rename var

Signed-off-by: Maanu Grover <[email protected]>

Author: maanug-nv

megatron/training/checkpointing.py

@@ -1413,35 +1413,35 @@ def load_checkpoint(ddp_model, optimizer, opt_param_scheduler, load_arg='load',
     ignore_rng_state = False
     ignore_rerun_state = True
     if ckpt_format == "torch_dist":
-        state_dict_args = types.SimpleNamespace()
+        ckpt_args = types.SimpleNamespace()
         if state_dict is not None and "args" in state_dict:
-            state_dict_args = state_dict.get("args")
+            ckpt_args = state_dict.get("args")
 
-        if not hasattr(state_dict_args, "tensor_model_parallel_size"):
+        if not hasattr(ckpt_args, "tensor_model_parallel_size"):
             print_rank_0("WARNING: TP size not found in checkpoint args, using 1 as default.")
-        if not hasattr(state_dict_args, "pipeline_model_parallel_size"):
+        if not hasattr(ckpt_args, "pipeline_model_parallel_size"):
             print_rank_0("WARNING: PP size not found in checkpoint args, using 1 as default.")
 
         ckpt_tp_pp = (
-            getattr(state_dict_args, "tensor_model_parallel_size", 1),
-            getattr(state_dict_args, "pipeline_model_parallel_size", 1),
+            getattr(ckpt_args, "tensor_model_parallel_size", 1),
+            getattr(ckpt_args, "pipeline_model_parallel_size", 1),
         )
         run_tp_pp = (
             args.tensor_model_parallel_size,
             args.pipeline_model_parallel_size,
         )
 
-        ckpt_world_size = getattr(state_dict_args, 'world_size', 0)
+        ckpt_world_size = getattr(ckpt_args, 'world_size', 0)
         run_world_size = getattr(args, 'world_size', 0)
-        ckpt_dp = getattr(state_dict_args, 'data_parallel_size', 0)
+        ckpt_dp = getattr(ckpt_args, 'data_parallel_size', 0)
         run_dp = getattr(args, 'data_parallel_size', 0)
         mismatch_msg = "(TP, PP) mismatch after resume ({} vs {} from checkpoint)".format(
             run_tp_pp, ckpt_tp_pp
         )
 
         # Determine if RNG state will be loaded
         if (ckpt_tp_pp == run_tp_pp and not release and not args.finetune and not args.no_load_rng
-                and not getattr(state_dict_args, 'no_save_rng', False)):
+                and not getattr(ckpt_args, 'no_save_rng', False)):
             gen_sd_rng_state = get_rng_state(args.ckpt_format)  # we can load the rng state
         else:
             ignore_rng_state = True
@@ -1456,7 +1456,7 @@ def load_checkpoint(ddp_model, optimizer, opt_param_scheduler, load_arg='load',
         print_rank_0(f'sharded_state_dict metadata loaded from the checkpoint: {sharded_sd_metadata}')
         # Determine if optimizer state will be loaded
         if (not release and not args.finetune and not args.no_load_optim
-                and not getattr(state_dict_args, 'no_save_optim', False)):
+                and not getattr(ckpt_args, 'no_save_optim', False)):
             gen_sd_optim = optimizer
             gen_sd_opt_param_scheduler = opt_param_scheduler
 
@@ -1467,7 +1467,7 @@ def load_checkpoint(ddp_model, optimizer, opt_param_scheduler, load_arg='load',
                     # (for MCore v0.13+ checkpoints `sharded_sd_metadata is not None`)
                     sharded_sd_metadata = {
                         'distrib_optim_sharding_type': ('fully_sharded_model_space'
-                                                        if getattr(state_dict_args, 'ckpt_fully_parallel_save', False)
+                                                        if getattr(ckpt_args, 'ckpt_fully_parallel_save', False)
                                                         else 'dp_zero_gather_scatter'),
                     }
                 if (