diff --git a/megatron/core/optimizer/distrib_optimizer.py b/megatron/core/optimizer/distrib_optimizer.py index 798230871c2..374b1aab096 100644 --- a/megatron/core/optimizer/distrib_optimizer.py +++ b/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."): diff --git a/megatron/core/transformer/multi_latent_attention.py b/megatron/core/transformer/multi_latent_attention.py index 5db3154f552..176b8e8451d 100644 --- a/megatron/core/transformer/multi_latent_attention.py +++ b/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) diff --git a/tests/unit_tests/test_optimizer.py b/tests/unit_tests/test_optimizer.py index 3c445cd3633..94613b7096c 100644 --- a/tests/unit_tests/test_optimizer.py +++ b/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", diff --git a/tests/unit_tests/transformer/test_multi_latent_attention.py b/tests/unit_tests/transformer/test_multi_latent_attention.py index 863a4e23d6d..646c87f2839 100644 --- a/tests/unit_tests/transformer/test_multi_latent_attention.py +++ b/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"