fix(mtp): track megatron mtp_model_layer rename in raw converters

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 = []

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}")

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()

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:

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}",
             )
 

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}",
             )
 

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)