[Draft] Fix dist-ckpt for LayerwiseDistOpt

megatron/core/dist_checkpointing/dict_utils.py

@@ -217,22 +217,25 @@ def dict_list_map_outplace(f: Callable[[U], V], x: Union[Dict, List, U]) -> Unio
         return f(x)
 
 
-def merge(x1: Union[dict, list], x2: Union[dict, list], key: Tuple[Union[str, int], ...] = ()):
+def merge(x1: Union[dict, list], x2: Union[dict, list], key: Tuple[Union[str, int], ...] = (), layerwise_dist_opt: bool = False):
     """Merges dicts and lists recursively."""
     if isinstance(x1, dict) and isinstance(x2, dict):
         for k, v2 in x2.items():
             if k not in x1:
                 x1[k] = v2
             else:
-                x1[k] = merge(x1[k], v2, key=key + (k,))
+                x1[k] = merge(x1[k], v2, key=key + (k,), layerwise_dist_opt=layerwise_dist_opt)
     elif isinstance(x1, list) and isinstance(x2, list):
-        if len(x1) != len(x2):
+        # for layerwise dist opt, if the common list is empty, return the loaded list
+        if layerwise_dist_opt and len(x1) == 0:
+            return x2
+        elif len(x1) != len(x2):
             raise ValueError(
                 f"Cannot merge two lists with different lengths ({len(x1)} and {len(x2)}, "
                 f"encountered at level {key})"
             )
         for i, v2 in enumerate(x2):
-            x1[i] = merge(x1[i], v2, key=key + (i,))
+            x1[i] = merge(x1[i], v2, key=key + (i,), layerwise_dist_opt=layerwise_dist_opt)
     else:
         raise ValueError(
             f"Duplicate non-dict and non-list values encountered: `{x1}` and `{x2}` "

megatron/core/dist_checkpointing/serialization.py

@@ -64,6 +64,7 @@ def load(
     common_strategy: Union[LoadCommonStrategy, Tuple[str, int], None] = None,
     validate_access_integrity: bool = True,
     strict: Union[str, StrictHandling] = StrictHandling.ASSUME_OK_UNEXPECTED,
+    layerwise_dist_opt: bool = False,
 ) -> Union[StateDict, Tuple[StateDict, Set[str], Set[str]]]:
     """Loading entrypoint.
 
@@ -160,7 +161,7 @@ def load(
 
     loaded_state_dict = sharded_strategy.load(sharded_state_dict, checkpoint_dir)
 
-    merge(common_state_dict, loaded_state_dict)
+    merge(common_state_dict, loaded_state_dict, layerwise_dist_opt=True)
 
     loaded_state_dict = apply_factory_merges(common_state_dict, sh_ten_factories)
 

megatron/core/extensions/transformer_engine.py

@@ -1090,6 +1090,7 @@ def forward(
                 **attention_bias_kwargs,
                 **packed_seq_kwargs,
             )
+
         else:
             core_attn_out = super().forward(
                 query, key, value, attention_mask, **attention_bias_kwargs, **packed_seq_kwargs

megatron/core/optimizer/optimizer.py

@@ -1222,7 +1222,14 @@ def load_state_dict(self, state_dict):
         if isinstance(state_dict, dict):
             state_dict = (v for k, v in sorted(state_dict.items()))
         for optimizer, state in zip(self.chained_optimizers, state_dict):
-            optimizer.load_state_dict(state)
+            # Mostly for Layerwise Dist Opt situation but it can be a general check
+            # if the state dict optimizer is in this rank, load the state dict
+            if 'optimizer' in state:
+                if 'state' in state['optimizer']:
+                    optimizer.load_state_dict(state)
+            # else:
+            #     print(f"state_dict: {torch.distributed.get_rank()} {state['optimizer']} {type(optimizer)} \n")
+            #     assert False
         self._synchronize_steps()
 
     @torch.no_grad()

megatron/training/checkpointing.py

@@ -1059,7 +1059,7 @@ def _load_global_dist_base_checkpoint(
         )
     if checkpointing_context is not None:
         checkpointing_context["load_strategy"] = load_strategy
-    state_dict = dist_checkpointing.load(sharded_state_dict, checkpoint_name, load_strategy, strict=args.dist_ckpt_strictness)
+    state_dict = dist_checkpointing.load(sharded_state_dict, checkpoint_name, load_strategy, strict=args.dist_ckpt_strictness, layerwise_dist_opt=args.optimizer == 'dist_muon')
     return state_dict, checkpoint_name, release, CheckpointType.GLOBAL