Merge branch 'skierat/avoid_prefix' into 'main'

Fix reshardable checkpoint format by removing 'chained_*.' prefixes

Closes dl/JoC/nemo-ci#1995

See merge request ADLR/megatron-lm!3455

Author: ko3n1g

megatron/core/optimizer/optimizer.py

@@ -1136,6 +1136,15 @@ def state_dict(self):
     def sharded_state_dict(
         self, model_sharded_state_dict: ShardedStateDict, is_loading: bool = False, **kwargs
     ):
+        metadata = kwargs.get('metadata') or {}
+        should_add_prefix = True  # Backward-compatibility
+        if (
+            metadata.get('chained_optim_avoid_prefix', False)
+            # This condition should be True if a distributed optimizer isn't used
+            and metadata.get('distrib_optim_sharding_type') != 'dp_zero_gather_scatter'
+        ):
+            should_add_prefix = False
+
         if len(self.chained_optimizers) == 1:
             return self.chained_optimizers[0].sharded_state_dict(
                 model_sharded_state_dict, is_loading, **kwargs
@@ -1147,7 +1156,8 @@ def sharded_state_dict(
                 optim_state_dict = optimizer.sharded_state_dict(
                     model_sharded_state_dict, is_loading, **kwargs
                 )
-                add_prefix_for_sharding(optim_state_dict, f'chained_{optimizer_idx}.')
+                if should_add_prefix:
+                    add_prefix_for_sharding(optim_state_dict, f'chained_{optimizer_idx}.')
                 sharded_state_dict[optimizer_idx] = optim_state_dict
             return sharded_state_dict
 

megatron/training/checkpointing.py

@@ -336,6 +336,7 @@ def _build_sharded_state_dict_metadata(args: Namespace) -> dict:
             metadata['distrib_optim_sharding_type'] = 'fully_sharded_model_space'
         else:
             metadata['distrib_optim_sharding_type'] = 'dp_zero_gather_scatter'
+    metadata['chained_optim_avoid_prefix'] = True
     return metadata
 
 def save_checkpoint(iteration, model, optimizer, opt_param_scheduler, num_floating_point_operations_so_far,