Support fused MLA QKV checkpoint reload

megatron/core/optimizer/distrib_optimizer.py

@@ -2884,6 +2884,47 @@ def _normalize_state_dict_for_grouped_params(state_dict_flat, model_chunk):
                     for i, tensor in enumerate(split_tensors):
                         state_dict_flat[f"{key_prefix}{indexed_suffixes[i]}"] = tensor
 
+    @staticmethod
+    def _synthesize_state_dict_params_for_model(state_dict_flat, model_chunk):
+        """Let modules materialize runtime params before optimizer state-dict matching.
+
+        This covers modules whose runtime parameter is assembled from multiple checkpoint
+        tensors. Normal model loading can handle this in _load_from_state_dict(), but
+        reload_model_params(state_dict=...) matches optimizer master params directly by name.
+        """
+        for module_name, module in model_chunk.named_modules():
+            synthesize = getattr(module, '_synthesize_fused_qkv_down_weight', None)
+            key_suffixes = getattr(module, '_synthetic_state_dict_key_suffixes', None)
+            if not callable(synthesize) or not callable(key_suffixes):
+                continue
+
+            # Optional compatibility hook contract:
+            # - key_suffixes() returns source checkpoint keys under this module, used only
+            #   to infer checkpoint prefixes despite wrapper-added path segments.
+            # - synthesize(state_dict_flat, module_prefix) mutates the flat checkpoint dict
+            #   in place, materializing runtime parameter names from checkpoint tensors.
+            for inner_key in key_suffixes():
+                for module_prefix in DistributedOptimizer._state_dict_module_prefixes(
+                    state_dict_flat, module_name, inner_key
+                ):
+                    synthesize(state_dict_flat, module_prefix)
+
+    @staticmethod
+    def _state_dict_module_prefixes(state_dict_flat, module_name, inner_key):
+        """Find checkpoint prefixes for a module by suffix-matching one inner key."""
+        module_parts = module_name.split(".") if module_name else []
+        for start_idx in range(len(module_parts) + 1):
+            module_suffix = ".".join(module_parts[start_idx:])
+            key_suffix = f"{module_suffix}.{inner_key}" if module_suffix else inner_key
+            prefixes = {
+                state_key[: len(state_key) - len(inner_key)]
+                for state_key in state_dict_flat
+                if state_key.endswith(key_suffix)
+            }
+            if prefixes:
+                return prefixes
+        return set()
+
     def _build_model_param_to_state_dict_param_map(self, state_dict):
         """Create a map from model params to tensors in state_dict based on their names."""
         state_dict_list = []
@@ -2906,6 +2947,7 @@ def _build_model_param_to_state_dict_param_map(self, state_dict):
         model_param_to_state_dict_param_map = {}
         for chunk_idx, model_chunk in enumerate(self.model_chunks):
             self._normalize_state_dict_for_grouped_params(state_dict_list[chunk_idx], model_chunk)
+            self._synthesize_state_dict_params_for_model(state_dict_list[chunk_idx], model_chunk)
             names_in_state_dict = set(state_dict_list[chunk_idx].keys())
             for name, model_param in model_chunk.named_parameters():
                 while name.startswith("module."):

megatron/core/transformer/multi_latent_attention.py

@@ -1384,8 +1384,11 @@ def _clone_sharded_object_with_key(obj: ShardedObject, new_key: str) -> ShardedO
                 sharded_state_dict[q_extra_key] = fused_obj
                 sharded_state_dict[kv_extra_key] = fused_obj
 
+        # Keep fused layernorm params so TransformerLayer's key map can load old
+        # input_layernorm checkpoints into the fused TE down-proj module.
         for key in list(sharded_state_dict.keys()):
-            if key.startswith(fused_prefix):
+            suffix = key[len(fused_prefix) :] if key.startswith(fused_prefix) else ""
+            if key.startswith(fused_prefix) and not suffix.startswith("layer_norm_"):
                 del sharded_state_dict[key]
 
         fused_weight = self.linear_qkv_down_proj.weight
@@ -1429,8 +1432,12 @@ def _clone_sharded_object_with_key(obj: ShardedObject, new_key: str) -> ShardedO
 
         return sharded_state_dict
 
-    def _load_from_state_dict(self, state_dict, prefix, *args, **kwargs):
-        """Load state dict with automatic unfused->fused conversion."""
+    def _synthetic_state_dict_key_suffixes(self):
+        """Return source checkpoint keys used to locate this module in a state dict."""
+        return ("linear_q_down_proj.weight",)
+
+    def _synthesize_fused_qkv_down_weight(self, state_dict, prefix):
+        """Materialize fused qkv-down weight from old separate q/kv checkpoint keys."""
         q_key = f"{prefix}linear_q_down_proj.weight"
         kv_key = f"{prefix}linear_kv_down_proj.weight"
         fused_key = f"{prefix}linear_qkv_down_proj.weight"
@@ -1447,4 +1454,8 @@ def _as_tensor(x):
             state_dict.pop(f"{prefix}linear_q_down_proj.bias", None)
             state_dict.pop(f"{prefix}linear_kv_down_proj.bias", None)
 
+    def _load_from_state_dict(self, state_dict, prefix, *args, **kwargs):
+        """Load state dict with automatic unfused->fused conversion."""
+        self._synthesize_fused_qkv_down_weight(state_dict, prefix)
+
         return super()._load_from_state_dict(state_dict, prefix, *args, **kwargs)

tests/unit_tests/test_optimizer.py

@@ -13,6 +13,9 @@
 from transformer_engine.pytorch.fp8 import fp8_autocast
 
 from megatron.core.distributed import DistributedDataParallel, DistributedDataParallelConfig
+from megatron.core.models.gpt.gpt_layer_specs import (
+    get_gpt_layer_with_transformer_engine_submodules,
+)
 from megatron.core.optimizer import (
     ChainedOptimizer,
     OptimizerConfig,
@@ -23,9 +26,17 @@
     get_megatron_optimizer,
     get_standard_config_overrides,
 )
+from megatron.core.optimizer.distrib_optimizer import DistributedOptimizer
 from megatron.core.optimizer_param_scheduler import ParamGroupOverride
 from megatron.core.process_groups_config import ProcessGroupCollection
+from megatron.core.tensor_parallel.random import model_parallel_cuda_manual_seed
 from megatron.core.transformer import TransformerConfig
+from megatron.core.transformer.enums import AttnMaskType
+from megatron.core.transformer.multi_latent_attention import (
+    FusedMLASelfAttention,
+    MLASelfAttentionSubmodules,
+)
+from megatron.core.transformer.transformer_config import MLATransformerConfig
 from megatron.core.utils import is_te_min_version, is_torch_min_version
 from tests.unit_tests.test_utilities import Utils
 from tests.unit_tests.test_utils import _init_distributed
@@ -1065,6 +1076,59 @@ def test_optimizer_reload_model_params():
             )
 
 
+def test_distributed_optimizer_synthesizes_fused_qkv_down_weight_for_state_dict_matching():
+    if not is_te_min_version("1.10.0"):
+        pytest.skip("Requires TE >= 1.10.0")
+
+    Utils.initialize_model_parallel(1, 1)
+    model_parallel_cuda_manual_seed(123)
+    try:
+        transformer_config = MLATransformerConfig(
+            num_layers=2,
+            hidden_size=12,
+            num_attention_heads=4,
+            use_cpu_initialization=True,
+            q_lora_rank=32,
+            kv_lora_rank=32,
+            qk_head_dim=128,
+            v_head_dim=128,
+            qk_pos_emb_head_dim=64,
+            rope_type="rope",
+            rotary_base=10000,
+            original_max_position_embeddings=32,
+        )
+        submodules = get_gpt_layer_with_transformer_engine_submodules(
+            multi_latent_attention=True, mla_down_proj_fusion=True
+        ).self_attention.submodules
+        assert isinstance(submodules, MLASelfAttentionSubmodules)
+        fused_mla = FusedMLASelfAttention(
+            transformer_config, submodules, layer_number=1, attn_mask_type=AttnMaskType.causal
+        )
+
+        model = nn.Module()
+        model.decoder = nn.Module()
+        model.decoder.layers = nn.ModuleList([nn.Module()])
+        model.decoder.layers[0].self_attention = fused_mla
+
+        prefix = "module.decoder.layers.0.self_attention."
+        sharded_state_dict = fused_mla.sharded_state_dict(prefix=prefix)
+        q_key = next(k for k in sharded_state_dict if k.endswith("linear_q_down_proj.weight"))
+        kv_key = next(k for k in sharded_state_dict if k.endswith("linear_kv_down_proj.weight"))
+        fused_key = f"{prefix}{next(k for k in fused_mla.state_dict() if k.endswith('linear_qkv_down_proj.weight'))}"
+        q_weight = sharded_state_dict[q_key].data
+        kv_weight = sharded_state_dict[kv_key].data
+        state_dict = {q_key: q_weight, kv_key: kv_weight}
+
+        DistributedOptimizer._synthesize_state_dict_params_for_model(state_dict, model)
+
+        assert fused_key in state_dict
+        assert q_key not in state_dict
+        assert kv_key not in state_dict
+        torch.testing.assert_close(state_dict[fused_key], torch.cat([q_weight, kv_weight], dim=0))
+    finally:
+        Utils.destroy_model_parallel()
+
+
 @pytest.mark.skipif(
     not is_torch_min_version("2.4.0"),
     reason="torch.distributed.init_device_mesh requires torch >= 2.4.0",

tests/unit_tests/transformer/test_multi_latent_attention.py

@@ -1875,6 +1875,9 @@ def test_sharded_state_dict_splits_back(self):
         assert any(
             'linear_kv_down_proj.weight' in k for k in sharded_sd
         ), f"Expected linear_kv_down_proj.weight in sharded state dict, got keys: {list(sharded_sd.keys())}"
+        assert any(
+            'linear_qkv_down_proj.layer_norm_weight' in k for k in sharded_sd
+        ), f"Expected linear_qkv_down_proj.layer_norm_weight in sharded state dict, got keys: {list(sharded_sd.keys())}"
         assert not any(
             'linear_qkv_down_proj.weight' in k for k in sharded_sd
         ), f"Unexpected linear_qkv_down_proj.weight in sharded state dict"