Address review: warn when the GLM-5 MTP copy finds no tensors; note optional DSA test imports

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/unified_export_megatron.py

@@ -936,6 +936,8 @@ def _copy_decoder_mtp_layers_from_pretrained(self) -> dict[str, torch.Tensor]:
             self.exclude_modules.append(dst + "*")
         if mtp_state_dict:
             print(f"Copied {len(mtp_state_dict)} MTP tensors from {source}")
+        else:
+            warn_rank_0(f"No MTP tensors under {src_prefixes} in {source}; exporting without MTP.")
         return mtp_state_dict
 
     @staticmethod

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

@@ -1886,6 +1886,7 @@ def test_registered_megatron_quant_modules_checkpoint_quantizer_state():
 
 def _get_tiny_dsa_gpt_model():
     """Tiny GPT with DSA sparse attention (AbsorbedMLASelfAttention + DSAttention + indexer)."""
+    # Local: DSA is optional; test_dsa_kv_cache_quant importorskips it before calling this.
     from megatron.core.models.gpt.experimental_attention_variant_module_specs import (
         get_transformer_block_with_experimental_attention_variant_spec,
     )
@@ -1925,6 +1926,7 @@ def _get_tiny_dsa_gpt_model():
 
 
 def _test_dsa_kv_cache_quant_helper(tmp_path, rank, size):
+    # Local: DSA is optional; test_dsa_kv_cache_quant importorskips it before calling this.
     from megatron.core.transformer.experimental_attention_variant.dsa import DSAttention
 
     initialize_for_megatron(tensor_model_parallel_size=1, pipeline_model_parallel_size=1, seed=SEED)