Address review: refuse a Megatron-built GLM-5 MTP, drop the KPool indexer export

- GLM-5's MTP is only copied from the source checkpoint; raise on every rank if Megatron built
  one instead of exporting it under mtp.* names, and drop the unused mtp.* rules.
- Leave the KPool indexer parameters to the GLM-5.3-Flash PR that adds their rules.
- Check the restored DSA K/V quantizers stay enabled (V is unused in the forward).

Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com>
Signed-off-by: Keval Morabia <28916987+kevalmorabia97@users.noreply.github.com>

Author: kevalmorabia97

modelopt/torch/export/plugins/mcore_glm.py

@@ -16,7 +16,8 @@
 """Custom mappings from Megatron Core models to GLM-5.x Hugging Face models.
 
 GLM-5 / GLM-5.2 (``glm_moe_dsa``) is DeepSeek-V3-style MLA with a DSA indexer. It keeps the MTP
-layer at HF index ``num_hidden_layers`` under the decoder's names.
+layer at HF index ``num_hidden_layers`` under the decoder's names; Megatron-Bridge does not build
+it, so the exporter copies it from the source checkpoint.
 """
 
 from .mcore_custom import (
@@ -37,8 +38,4 @@
     "indexer.linear_weights_proj": NameRemapping("model.layers.{}.self_attn.indexer.weights_proj."),
     "experts.linear_fc1": GroupedGatedMLPSlicing("model.layers.{}.mlp.experts.{{}}"),
     "experts.linear_fc2": GroupedMLPSlicing("model.layers.{}.mlp.experts.{{}}.down_proj"),
-    "mtp.enorm": NameRemapping("model.layers.{}.enorm."),
-    "mtp.hnorm": NameRemapping("model.layers.{}.hnorm."),
-    "mtp.eh_proj": NameRemapping("model.layers.{}.eh_proj."),
-    "mtp.final_layernorm": NameRemapping("model.layers.{}.shared_head.norm."),
 }

modelopt/torch/export/unified_export_megatron.py

@@ -205,6 +205,11 @@ def __init__(
         self.all_rules = self._populate_rule_book()
         self.rules = self.all_rules[self.arch]
         self.all_mcore_mappings = all_mcore_hf_export_mapping[self.arch]
+        if self.rules.get("mtp_in_decoder_layers", False) and language_model.config.mtp_num_layers:
+            # Only the copy from the source checkpoint is supported (the bridge builds no MTP).
+            raise NotImplementedError(
+                f"Exporting a Megatron-built MTP for {self.arch} is not supported yet."
+            )
         # The vision tower is copied through unquantized, so deployments must not treat it as such.
         self.exclude_modules = [
             prefix.removesuffix(".") + "*" for prefix in self.vision_passthrough_prefixes or ()
@@ -990,18 +995,11 @@ def _get_gated_delta_net_state_dict(self, layer, layer_id, is_mtp=False):
         self.rules["linear_attn.out_proj"](gdn.out_proj, layer_id, is_mtp=is_mtp)
 
     def _get_dsa_indexer_state_dict(self, indexer, layer_id, is_mtp=False):
-        """Export the DSA kpool indexer of a sparse MLA layer."""
+        """Export the DSA indexer of a sparse MLA layer."""
         if "indexer.linear_wq_b" not in self.rules:
             raise NotImplementedError(f"No export rule for the DSA indexer of {self.arch}.")
         for name in ("linear_wq_b", "linear_wk", "k_norm", "linear_weights_proj"):
             self.rules[f"indexer.{name}"](getattr(indexer, name), layer_id, is_mtp=is_mtp)
-        # KPool (GLM-5.3-Flash) only; plain DSA indexers (GLM-5.2) have no compression params.
-        for name in ("index_kpool_compress_ape", "index_kpool_compress_gate"):
-            param = getattr(indexer, name, None)
-            if param is not None:
-                self.rules[f"indexer.{name}"](
-                    param.detach().to(self.dtype), layer_id, is_mtp=is_mtp
-                )
 
     def _get_mamba_layer_state_dict(self, layer, layer_id, is_mtp=False):
         if not isinstance(layer.norm, IdentityOp):

tests/gpu_megatron/torch/export/test_unified_export_megatron.py

@@ -1054,6 +1054,30 @@ def test_mtp_state_dict_index_file(tmp_path):
     assert "mtp*" in exporter.exclude_modules
 
 
+def _test_live_decoder_mtp_export_rejected(model_dir, rank, size):
+    model = get_mcore_gpt_model(
+        tensor_model_parallel_size=size,
+        pipeline_model_parallel_size=1,
+        initialize_megatron=True,
+        num_layers=2,
+        hidden_size=64,
+        num_attention_heads=4,
+        vocab_size=128,
+        max_sequence_length=32,
+        mtp_num_layers=1,
+    ).cuda()
+    with pytest.raises(NotImplementedError, match="Megatron-built MTP"):
+        GPTModelExporter(model, str(model_dir))
+
+
+def test_live_decoder_mtp_export_rejected(dist_workers_size_1, tmp_path):
+    """GLM-5 MTP is only copied from the source; a Megatron-built one must not export misnamed."""
+    transformers.GlmMoeDsaConfig(
+        num_hidden_layers=2, architectures=["GlmMoeDsaForCausalLM"]
+    ).save_pretrained(tmp_path)
+    dist_workers_size_1.run(partial(_test_live_decoder_mtp_export_rejected, tmp_path))
+
+
 def test_mtp_state_dict_copies_decoder_mtp_layers(tmp_path):
     """GLM-5 keeps MTP as an extra decoder layer; copy it dequantized when Megatron did not build it."""
     model_dir = tmp_path / "fake_glm5"

tests/gpu_megatron/torch/quantization/plugins/test_megatron.py

@@ -2176,6 +2176,11 @@ def _test_dsa_kv_cache_quant_helper(tmp_path, rank, size):
         if ".indexer." in name:
             param.requires_grad_(False)
     sharded_state_dict_test_helper(tmp_path, model, model_test, forward)
+    # The helper compares every restored tensor (incl. K/V amax), but V is unused in the forward
+    # pass, so also check both quantizers came back enabled.
+    restored = [m for m in model_test.modules() if isinstance(m, DSAttention)]
+    assert len(restored) == len(dsa_modules)
+    assert all(m.k_bmm_quantizer.is_enabled and m.v_bmm_quantizer.is_enabled for m in restored)
 
 
 def test_dsa_kv_cache_quant(dist_workers_size_1, tmp_path):