diff --git a/miles/backends/megatron_utils/megatron_to_hf/qwen3_5.py b/miles/backends/megatron_utils/megatron_to_hf/qwen3_5.py index 2aabd86eba..eb2f4ea291 100644 --- a/miles/backends/megatron_utils/megatron_to_hf/qwen3_5.py +++ b/miles/backends/megatron_utils/megatron_to_hf/qwen3_5.py @@ -14,8 +14,11 @@ def _convert_mtp_layer(args, name, param, layer_idx): if "eh_proj.weight" in name: return [("mtp.fc.weight", param)] - if "transformer_layer" in name: - proxy_name = name.replace(f"mtp.layers.{layer_idx}.transformer_layer", f"decoder.layers.{layer_idx}") + # Accept both MTP submodule names: mtp_model_layer (new) and transformer_layer (old). + for mtp_layer_attr in ("mtp_model_layer", "transformer_layer"): + if mtp_layer_attr not in name: + continue + proxy_name = name.replace(f"mtp.layers.{layer_idx}.{mtp_layer_attr}", f"decoder.layers.{layer_idx}") mapped_params = convert_qwen3_5_to_hf(args, proxy_name, param) final_params = [] diff --git a/miles/backends/megatron_utils/megatron_to_hf/qwen3_next.py b/miles/backends/megatron_utils/megatron_to_hf/qwen3_next.py index f248cc0fc3..85633cb234 100644 --- a/miles/backends/megatron_utils/megatron_to_hf/qwen3_next.py +++ b/miles/backends/megatron_utils/megatron_to_hf/qwen3_next.py @@ -155,13 +155,16 @@ def convert_qwen3_next_to_hf(args, name, param): elif rest == "final_layernorm.weight": return [("mtp.norm.weight", param)] - # transformer_layer components → reuse decoder conversion with mtp prefix - if rest.startswith("transformer_layer."): - transformer_rest = rest[len("transformer_layer.") :] - proxy_name = f"module.module.decoder.layers.{layer_idx}.{transformer_rest}" - results = convert_qwen3_next_to_hf(args, proxy_name, param) - return [ - (hf_name.replace(f"model.layers.{layer_idx}", f"mtp.layers.{layer_idx}"), p) for hf_name, p in results - ] + # MTP transformer-layer components reuse the decoder conversion with an mtp prefix. + # Accept both MTP submodule names: mtp_model_layer (new) and transformer_layer (old). + for mtp_layer_attr in ("mtp_model_layer.", "transformer_layer."): + if rest.startswith(mtp_layer_attr): + transformer_rest = rest[len(mtp_layer_attr) :] + proxy_name = f"module.module.decoder.layers.{layer_idx}.{transformer_rest}" + results = convert_qwen3_next_to_hf(args, proxy_name, param) + return [ + (hf_name.replace(f"model.layers.{layer_idx}", f"mtp.layers.{layer_idx}"), p) + for hf_name, p in results + ] raise ValueError(f"Unknown parameter name: {name}") diff --git a/miles/backends/megatron_utils/update_weight/common.py b/miles/backends/megatron_utils/update_weight/common.py index d1176e7b8b..5bfa3d6e03 100644 --- a/miles/backends/megatron_utils/update_weight/common.py +++ b/miles/backends/megatron_utils/update_weight/common.py @@ -322,15 +322,19 @@ def _named_params_and_buffers_global( # MTP layer indices start from 0 layer_idx, rest = match.groups() - expert_pattern = r"transformer_layer.mlp.experts\.(.+)\.weight(\d+)" + # Match both MTP submodule names and re-emit whichever the running Megatron uses. + expert_pattern = r"(mtp_model_layer|transformer_layer)\.mlp\.experts\.(.+)\.weight(\d+)" match = re.match(expert_pattern, rest) if not match: yield name, param continue - rest, expert_idx = match.groups() + mtp_layer_attr, rest, expert_idx = match.groups() expert_idx = int(expert_idx) + expert_offset - yield f"module.module.mtp.layers.{layer_idx}.transformer_layer.mlp.experts.{rest}.weight{expert_idx}", param + yield ( + f"module.module.mtp.layers.{layer_idx}.{mtp_layer_attr}.mlp.experts.{rest}.weight{expert_idx}", + param, + ) continue layer_idx, rest = match.groups() diff --git a/miles_plugins/mbridge/glm4moe_lite.py b/miles_plugins/mbridge/glm4moe_lite.py index c44888cd45..44238af639 100644 --- a/miles_plugins/mbridge/glm4moe_lite.py +++ b/miles_plugins/mbridge/glm4moe_lite.py @@ -130,8 +130,13 @@ def _convert_mtp_param(self, name: str) -> tuple[list[str]]: if name in direct_name_mapping: return [direct_name_mapping[name]] - assert "mtp.layers.0.transformer_layer" in name, "mtp not found" - proxy_name = name.replace("mtp.layers.0.transformer_layer", f"decoder.layers.{mtp_layer_id}") + # Accept both MTP submodule names: mtp_model_layer (new) and transformer_layer (old). + for mtp_layer_attr in ("mtp_model_layer", "transformer_layer"): + if f"mtp.layers.0.{mtp_layer_attr}" in name: + proxy_name = name.replace(f"mtp.layers.0.{mtp_layer_attr}", f"decoder.layers.{mtp_layer_id}") + break + else: + raise AssertionError(f"mtp not found in parameter name: {name}") if "self_attention" in proxy_name or "input_layernorm.weight" in proxy_name: return self._weight_name_mapping_attention(proxy_name) if "mlp" in proxy_name: diff --git a/miles_plugins/mbridge/qwen3_5.py b/miles_plugins/mbridge/qwen3_5.py index fe1556ef7d..9bc15de90b 100644 --- a/miles_plugins/mbridge/qwen3_5.py +++ b/miles_plugins/mbridge/qwen3_5.py @@ -269,9 +269,12 @@ def _convert_mtp_param(self, name: str) -> list[str]: if name in direct_name_mapping: return [direct_name_mapping[name]] - if "transformer_layer" in name: + # Accept both MTP submodule names: mtp_model_layer (new) and transformer_layer (old). + for mtp_layer_attr in ("mtp_model_layer", "transformer_layer"): + if f".{mtp_layer_attr}." not in name: + continue proxy_name = name.replace( - f"mtp.layers.{mtp_layer_idx}.transformer_layer", + f"mtp.layers.{mtp_layer_idx}.{mtp_layer_attr}", f"decoder.layers.{mtp_layer_idx}", ) diff --git a/miles_plugins/mbridge/qwen3_next.py b/miles_plugins/mbridge/qwen3_next.py index 953bd29027..00c3125137 100644 --- a/miles_plugins/mbridge/qwen3_next.py +++ b/miles_plugins/mbridge/qwen3_next.py @@ -130,9 +130,12 @@ def _convert_mtp_param(self, name: str) -> list[str]: if name in direct_mappings: return [direct_mappings[name]] - if "transformer_layer" in name: + # Accept both MTP submodule names: mtp_model_layer (new) and transformer_layer (old). + for mtp_layer_attr in ("mtp_model_layer", "transformer_layer"): + if f".{mtp_layer_attr}." not in name: + continue proxy_name = name.replace( - f"mtp.layers.{mtp_layer_idx}.transformer_layer", + f"mtp.layers.{mtp_layer_idx}.{mtp_layer_attr}", f"decoder.layers.{mtp_layer_idx}", ) diff --git a/tests/fast/backends/megatron_utils/test_qwen3_5_mtp_bridge_mapping.py b/tests/fast/backends/megatron_utils/test_qwen3_5_mtp_bridge_mapping.py index af1bd73d54..5f1f9cf538 100644 --- a/tests/fast/backends/megatron_utils/test_qwen3_5_mtp_bridge_mapping.py +++ b/tests/fast/backends/megatron_utils/test_qwen3_5_mtp_bridge_mapping.py @@ -3,8 +3,12 @@ import types from pathlib import Path +import pytest import torch +# The converters accept both the old and new MTP submodule names; cover both. +MTP_LAYER_ATTRS = ("mtp_model_layer", "transformer_layer") + def install_bridge_stubs(): megatron_mod = types.ModuleType("megatron") @@ -126,12 +130,13 @@ def load_raw_export_module(): return module -def test_mtp_moe_expert_mapping_uses_individual_hf_weights(): +@pytest.mark.parametrize("mtp_layer_attr", MTP_LAYER_ATTRS) +def test_mtp_moe_expert_mapping_uses_individual_hf_weights(mtp_layer_attr): module = load_bridge_module() bridge = module.Qwen3_5Bridge.__new__(module.Qwen3_5Bridge) - fc1_names = bridge._convert_mtp_param("mtp.layers.0.transformer_layer.mlp.experts.linear_fc1.weight42") - fc2_names = bridge._convert_mtp_param("mtp.layers.0.transformer_layer.mlp.experts.linear_fc2.weight42") + fc1_names = bridge._convert_mtp_param(f"mtp.layers.0.{mtp_layer_attr}.mlp.experts.linear_fc1.weight42") + fc2_names = bridge._convert_mtp_param(f"mtp.layers.0.{mtp_layer_attr}.mlp.experts.linear_fc2.weight42") assert fc1_names == [ "mtp.layers.0.mlp.experts.42.gate_proj.weight", @@ -140,17 +145,35 @@ def test_mtp_moe_expert_mapping_uses_individual_hf_weights(): assert fc2_names == ["mtp.layers.0.mlp.experts.42.down_proj.weight"] -def test_mtp_dense_mlp_mapping_still_uses_dense_hf_weights(): +@pytest.mark.parametrize("mtp_layer_attr", MTP_LAYER_ATTRS) +def test_mtp_dense_mlp_mapping_still_uses_dense_hf_weights(mtp_layer_attr): module = load_bridge_module() bridge = module.Qwen3_5Bridge.__new__(module.Qwen3_5Bridge) - fc1_names = bridge._convert_mtp_param("mtp.layers.0.transformer_layer.mlp.linear_fc1.weight") - fc2_names = bridge._convert_mtp_param("mtp.layers.0.transformer_layer.mlp.linear_fc2.weight") + fc1_names = bridge._convert_mtp_param(f"mtp.layers.0.{mtp_layer_attr}.mlp.linear_fc1.weight") + fc2_names = bridge._convert_mtp_param(f"mtp.layers.0.{mtp_layer_attr}.mlp.linear_fc2.weight") assert fc1_names == ["mtp.layers.0.mlp.gate_proj.weight", "mtp.layers.0.mlp.up_proj.weight"] assert fc2_names == ["mtp.layers.0.mlp.down_proj.weight"] +@pytest.mark.parametrize("mtp_layer_attr", MTP_LAYER_ATTRS) +def test_raw_qwen3_5_mtp_export_handles_both_mtp_layer_attrs(mtp_layer_attr): + module = load_raw_export_module() + + param = torch.zeros(4) + # convert_qwen3_5_to_hf derives head_dim at entry, so the stub needs attention dims. + args = types.SimpleNamespace(kv_channels=None, hidden_size=64, num_attention_heads=4, num_query_groups=2) + result = module._convert_mtp_layer( + args, + f"module.module.mtp.layers.0.{mtp_layer_attr}.mlp.experts.linear_fc1", + param, + 0, + ) + + assert result == [("mtp.layers.0.mlp.experts.gate_up_proj", param)] + + def test_mtp_block_spec_uses_current_transformer_layer_spec(): module = load_bridge_module() bridge = module.Qwen3_5Bridge.__new__(module.Qwen3_5Bridge)