Merge branch 'consumed_tokens_restart_fix' into 'main'

found a bug in consumed tokens initialization

See merge request ADLR/megatron-lm!182

megatron/checkpointing.py

@@ -220,8 +220,10 @@ def load_checkpoint(model, optimizer, lr_scheduler, load_arg='load'):
     if 'args' in state_dict:
         checkpoint_args = state_dict['args']
         check_checkpoint_args(checkpoint_args)
-        args.consumed_train_samples = getattr(args, 'consumed_train_samples', 0)
-        args.consumed_valid_samples = getattr(args, 'consumed_valid_samples', 0)
+        args.consumed_train_samples = getattr(checkpoint_args,
+                                              'consumed_train_samples', 0)
+        args.consumed_valid_samples = getattr(checkpoint_args,
+                                              'consumed_valid_samples', 0)
     else:
         print_rank_0('could not find arguments in the checkpoint ...')