[ckpt] fix: support multi-layer MTP export filtering (#5495)

Signed-off-by: Yu Yao <yaoyu.094@gmail.com>

Author: yaoyu-33

src/megatron/bridge/models/conversion/auto_bridge.py

@@ -174,9 +174,10 @@ def _mtp_source_key_prefixes(source: Any, *configs: Any) -> tuple[str, ...]:
     differently:
 
     * DeepSeek-style: a dedicated ``mtp.*`` prefix.
-    * GLM-4.x ``glm4_moe_lite``: the nextn layer is stored as a regular decoder
-      layer at index ``num_hidden_layers`` (i.e. one past the last real layer),
-      so its tensors live under ``model.layers.{num_hidden_layers}.*``.
+    * GLM-4.x ``glm4_moe_lite`` and Step3.7: nextn layers are stored as regular
+      decoder layers starting at index ``num_hidden_layers`` (i.e. one past the
+      last real layer), so their tensors live under consecutive
+      ``model.layers.{layer_index}.*`` prefixes.
 
     When the megatron model is built without an MTP head the generator never
     yields these tensors. If they remain in the source sharding map, the shards
@@ -212,9 +213,20 @@ def _mtp_source_key_prefixes(source: Any, *configs: Any) -> tuple[str, ...]:
             break
 
     if num_hidden_layers is not _MISSING:
-        nextn_prefix = f"model.layers.{num_hidden_layers}."
-        if source.has_glob(f"{nextn_prefix}*"):
-            prefixes.append(nextn_prefix)
+        num_nextn_layers = 1
+        for config in configs:
+            text_config = _get_config_field(config, "text_config")
+            config_candidates = (config,) if text_config is _MISSING else (config, text_config)
+            for config_candidate in config_candidates:
+                for field in MTP_CONFIG_FIELDS:
+                    value = _get_config_field(config_candidate, field)
+                    if value is not _MISSING and value is not None:
+                        num_nextn_layers = max(num_nextn_layers, int(value))
+
+        for layer_index in range(num_hidden_layers, num_hidden_layers + num_nextn_layers):
+            nextn_prefix = f"model.layers.{layer_index}."
+            if source.has_glob(f"{nextn_prefix}*"):
+                prefixes.append(nextn_prefix)
 
     return tuple(prefixes)
 

tests/unit_tests/models/test_auto_bridge.py

@@ -311,6 +311,15 @@ def src(*present_globs):
         # Nested text_config carries num_hidden_layers.
         assert _mtp_source_key_prefixes(glm_src, {"text_config": {"num_hidden_layers": 47}}) == ("model.layers.47.",)
 
+        # Step3.7 stores multiple MTP layers after the regular decoder layers.
+        step37_src = src("model.layers.45.*", "model.layers.46.*", "model.layers.47.*")
+        step37_config = {"text_config": {"num_hidden_layers": 45, "num_nextn_predict_layers": 3}}
+        assert _mtp_source_key_prefixes(step37_src, step37_config) == (
+            "model.layers.45.",
+            "model.layers.46.",
+            "model.layers.47.",
+        )
+
         # No matching source keys -> nothing to strip.
         assert _mtp_source_key_prefixes(src(), {"num_hidden_layers": 47}) == ()