[2/3] Quantize the DSA KV cache and export GLM-5 / GLM-5.2 from Megatron-Bridge

CHANGELOG.rst

@@ -40,6 +40,7 @@ Changelog
 - Add an end-to-end W4A4 NVFP4 PTQ and QAD tutorial for Qwen3.6-35B-A3B also covering evaluation and vLLM throughput benchmarking. See `examples/megatron_bridge/tutorials/Qwen3.6-35B-A3B/README.md <https://github.com/NVIDIA/Model-Optimizer/tree/main/examples/megatron_bridge/tutorials/Qwen3.6-35B-A3B/>`_ for details.
 - Add ``--mlflow <tracking-uri>`` to the ``examples/megatron_bridge`` scripts that write a checkpoint -- ``prune_minitron.py``, ``quantize.py``, ``distill.py``, ``export_quantized_megatron_to_hf.py`` and ``export_distilled_megatron_to_hf.py`` (MLflow's own ``MLFLOW_TRACKING_URI`` is honoured too). Each run records the invocation, its arguments as searchable params and its log -- a distillation records its training metrics instead of the rank-0 log -- and writes ``.experiment.json`` into the checkpoint it produced, so a pruning, a quantization, the distillation that refines its checkpoint and the export that deploys it can be traced to one another; uploading the checkpoints themselves stays off unless ``--mlflow_log_checkpoints`` is passed.
 - Add ``--ep_size`` to ``examples/megatron_bridge/export_quantized_megatron_to_hf.py`` so large MoE models with grouped-GEMM experts can be exported with their experts sharded across GPUs. Checkpoints built with ``--no_moe_grouped_gemm`` must still be exported at ``--ep_size 1``.
+- Add unified HF export of GLM-5 / GLM-5.2 (``glm_moe_dsa``) checkpoints quantized with ``examples/megatron_bridge``. The MTP layer, which Megatron-Bridge does not build for these models, is copied through unquantized from the source checkpoint. DSA sparse-attention models also gain FP8 KV-cache quantization on Megatron.
 
 *Misc*
 

modelopt/torch/distill/plugins/megatron.py

@@ -40,6 +40,7 @@
 import modelopt.torch.distill as mtd
 from modelopt.torch.distill.config import Criterion
 from modelopt.torch.quantization.nn import TensorQuantizer
+from modelopt.torch.utils import warn_rank_0
 
 if TYPE_CHECKING:
     from megatron.core.dist_checkpointing.mapping import ShardedStateDict
@@ -592,9 +593,11 @@ def _sharded_state_dict(self, *args, **kwargs) -> "ShardedStateDict":
     skip_mtp_loss = _mtp_excluded_from_quantization(model)
     if distill_cfg.skip_lm_loss and skip_mtp_loss:
         # Freeze the untrained MTP head: DDP's overlapped grad reduce asserts on params with no grad.
-        for name, param in model.named_parameters():
-            if "mtp" in name.split("."):
-                param.requires_grad_(False)
+        warn_rank_0("MTP head is outside quantization and its loss is skipped: freezing it.")
+        with model.hide_teacher_model():
+            for name, param in model.named_parameters():
+                if "mtp" in name.split("."):
+                    param.requires_grad_(False)
 
     def _compute_student_lm_loss(self, labels, logits) -> Tensor:
         self._lm_loss_call_count += 1

modelopt/torch/export/plugins/mcore_common.py

@@ -18,6 +18,7 @@
 from typing import Any
 
 from .mcore_deepseek import deepseek_causal_lm_export, deepseek_causal_lm_import
+from .mcore_glm import glm_moe_dsa_causal_lm_export
 from .mcore_gptoss import gptoss_causal_lm_export, gptoss_causal_lm_import
 from .mcore_llama import (
     eagle3_deep_llama_causal_lm_export,
@@ -63,6 +64,7 @@
     "Qwen3VLForConditionalGeneration": qwen3vl_causal_lm_export,
     "Qwen3_5ForConditionalGeneration": qwen3_5_vl_causal_lm_export,
     "Qwen3_5MoeForConditionalGeneration": qwen3_5_vl_causal_lm_export,
+    "GlmMoeDsaForCausalLM": glm_moe_dsa_causal_lm_export,
 }
 
 # VLM architectures whose Megatron export covers the language model only: the vision tower is copied

modelopt/torch/export/plugins/mcore_glm.py

@@ -0,0 +1,41 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""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; Megatron-Bridge does not build
+it, so the exporter copies it from the source checkpoint.
+"""
+
+from .mcore_custom import (
+    GroupedGatedMLPSlicing,
+    GroupedMLPSlicing,
+    NameRemapping,
+    SelfAttentionScaling,
+)
+from .mcore_deepseek import deepseek_causal_lm_export
+
+glm_moe_dsa_causal_lm_export: dict = {
+    **deepseek_causal_lm_export,
+    "mtp_in_decoder_layers": True,
+    "core_attention": SelfAttentionScaling("model.layers.{}.self_attn."),
+    "indexer.linear_wq_b": NameRemapping("model.layers.{}.self_attn.indexer.wq_b."),
+    "indexer.linear_wk": NameRemapping("model.layers.{}.self_attn.indexer.wk."),
+    "indexer.k_norm": NameRemapping("model.layers.{}.self_attn.indexer.k_norm."),
+    "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"),
+}

modelopt/torch/export/unified_export_megatron.py

@@ -29,7 +29,7 @@
 
 import torch
 import torch.distributed
-from huggingface_hub import get_safetensors_metadata, hf_hub_download
+from huggingface_hub import get_safetensors_metadata, hf_hub_download, snapshot_download
 from huggingface_hub.errors import EntryNotFoundError
 from safetensors import safe_open
 from safetensors.torch import save_file
@@ -175,6 +175,7 @@ def __init__(
         self._hf_text_config = getattr(self._hf_config, "text_config", self._hf_config)
 
         # Update hf_config
+        self._src_num_hidden_layers = self._hf_text_config.num_hidden_layers
         self._hf_text_config.num_hidden_layers = language_model.config.num_layers
         self._hf_text_config.hidden_size = language_model.config.hidden_size
         # MLA's kv_channels is the V head dim, not HF's head_dim (e.g. glm5_next derives it from RoPE).
@@ -203,6 +204,12 @@ def __init__(
                 del self._hf_config.quantization_config
         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 ()
@@ -351,9 +358,11 @@ def save_pretrained(
                 "Export at expert parallel size > 1 needs grouped-GEMM experts; "
                 "export SequentialMLP (--no_moe_grouped_gemm) checkpoints at EP=1."
             )
-        # One writer per pipeline stage: other TP / DP / EP ranks hold the same layers (EP>1 ranks
-        # hold no gathered experts at all), and writing them too would race on the same files.
-        writes_layers = (
+        # One layer-shard writer per pipeline stage: other TP / DP / EP ranks hold the same layers
+        # (EP>1 ranks hold no gathered experts at all), and writing them too would race on the same
+        # files. Unlike is_writer_rank (one global writer of save_directory metadata), this has no
+        # PP term on purpose: each stage owns a disjoint slice of the layer shards.
+        is_stage_layer_writer = (
             tp_rank == 0 and get_data_parallel_rank() == 0 and get_expert_model_parallel_rank() == 0
         )
 
@@ -421,7 +430,7 @@ def save_pretrained(
 
             # The live MTP export runs EP collectives, so every last-stage main rank joins it; the
             # collective-free copy from the source checkpoint only runs on the writer.
-            mtp_state_dict = self._get_mtp_state_dict(copy_from_pretrained=writes_layers)
+            mtp_state_dict = self._get_mtp_state_dict(copy_from_pretrained=is_stage_layer_writer)
             if len(mtp_state_dict) > 0:
                 layer_state_dicts[self.model.config.num_layers].update(mtp_state_dict)
                 print(f"Successfully loaded {len(mtp_state_dict)} MTP tensors")
@@ -459,7 +468,7 @@ def save_pretrained(
         # Add multimodal components to state_dict. Since only support decoder model quantization,
         # no changes will be made to the multimodal components. We copy the multimodal components
         # from the pretrained model directly to the state_dict to avoid implementing the export logic.
-        if is_first_stage_main_rank and writes_layers:
+        if is_first_stage_main_rank and is_stage_layer_writer:
             # layer_state_dicts is keyed by layer_number (1-indexed), so the first
             # decoder layer on this (first) PP stage is the smallest key, not 0.
             # Merge the multimodal components into that shard so they land in a file
@@ -488,7 +497,7 @@ def save_pretrained(
         torch.distributed.barrier()
 
         save_safetensors_by_layer_index(
-            layer_state_dicts=layer_state_dicts if writes_layers else {},
+            layer_state_dicts=layer_state_dicts if is_stage_layer_writer else {},
             total_layers=self.model.config.num_layers,
             save_directory=save_directory,
             name_template="model-{:05d}-of-{:05d}",
@@ -550,7 +559,11 @@ def _verify_exported_keys(self, save_directory, pretrained_model_name_or_path) -
         # Narrow on purpose: compare module prefixes, not tensor names, since a quantized source
         # carries extras with no export counterpart, and only inside decoder layers, whose naming
         # is stable. A dropped decoder module is the case that loads fine and produces garbage.
-        num_layers = self.model.config.num_layers
+        # HF decoder depth of this export: Megatron may build several physical layers per HF layer.
+        hf_depth = self._hf_text_config.num_hidden_layers
+        num_mtp = 0
+        if self.rules.get("mtp_in_decoder_layers", False):
+            num_mtp = getattr(self._hf_text_config, "num_nextn_predict_layers", 0) or 0
         # Ancestors too: an export may expand one source module into several (Qwen3.5 packs
         # routed experts; the quantized export writes them per expert). Expansion is not a drop.
         exported_modules = set()
@@ -562,12 +575,18 @@ def _verify_exported_keys(self, save_directory, pretrained_model_name_or_path) -
                     break
                 exported_modules.add(prefix)
         missing = set()
-        for key in source - exported:
+        for key in source:
             layer = re.search(r"\.layers\.(\d+)\.", key)
             if layer is None:
                 continue  # see the note above: decoder layers only
-            if int(layer.group(1)) >= num_layers:
-                continue  # depth-pruned model: the source has layers this export does not
+            if int(layer.group(1)) >= hf_depth:
+                mtp_id = int(layer.group(1)) - self._src_num_hidden_layers
+                if not 0 <= mtp_id < num_mtp:
+                    continue  # depth-pruned model: the source has layers this export does not
+                # An MTP stored as extra decoder layers follows the exported decoder layers.
+                key = f"{key[: layer.start(1)]}{hf_depth + mtp_id}{key[layer.end(1) :]}"
+            if key in exported:
+                continue
             if key.rsplit(".", 1)[0] in exported_modules:
                 continue  # module is exported; this name is a source-side quantization artifact
             if "rotary_emb" in key:
@@ -693,6 +712,12 @@ def _get_transformer_layer_state_dict(self, layer, layer_id, is_mtp=False):
                     layer.self_attention.linear_kv_up_proj, layer_id, is_mtp=is_mtp
                 )
                 self.rules["linear_proj"](layer.self_attention.linear_proj, layer_id, is_mtp=is_mtp)
+                core_attention = getattr(layer.self_attention, "core_attention", None)
+                if core_attention is not None and "core_attention" in self.rules:
+                    self.rules["core_attention"](core_attention, layer_id, is_mtp=is_mtp)
+                indexer = getattr(core_attention, "indexer", None)
+                if indexer is not None:
+                    self._get_dsa_indexer_state_dict(indexer, layer_id, is_mtp)
             elif "linear_attn" in self.rules and hasattr(layer.self_attention, "in_proj"):
                 # GatedDeltaNet (Qwen3.5 linear attention): no q/k layernorm, no core_attention.
                 self._get_gated_delta_net_state_dict(layer, layer_id, is_mtp=is_mtp)
@@ -857,6 +882,8 @@ def _copy_mtp_state_dict_from_pretrained(self) -> dict[str, torch.Tensor]:
         mtp_state_dict = {}
         if not self._hf_pretrained_model_name:
             return mtp_state_dict
+        if self.rules.get("mtp_in_decoder_layers", False):
+            return self._copy_decoder_mtp_layers_from_pretrained()
 
         mtp_exists = False
 
@@ -910,6 +937,55 @@ def _copy_mtp_state_dict_from_pretrained(self) -> dict[str, torch.Tensor]:
             self.exclude_modules.append("mtp*")
         return mtp_state_dict
 
+    def _copy_decoder_mtp_layers_from_pretrained(self) -> dict[str, torch.Tensor]:
+        """Copy MTP layers stored as extra decoder layers (GLM-5.x) from the source, dequantized.
+
+        Used when Megatron did not build the MTP (e.g. Megatron-Bridge's GLM-5 bridge); the copies
+        stay BF16 and are excluded from quantization, like the released NVFP4 checkpoints.
+        """
+        num_mtp = getattr(self._hf_text_config, "num_nextn_predict_layers", 0) or 0
+        source = self._hf_pretrained_model_name
+        if num_mtp == 0 or source is None:
+            return {}
+        layers_prefix = self.all_mcore_mappings["input_layernorm"].target_name_or_prefix
+        layers_prefix = layers_prefix.split("{}")[0]  # e.g. "model.layers."
+        src_prefixes = [
+            f"{layers_prefix}{self._src_num_hidden_layers + i}." for i in range(num_mtp)
+        ]
+        if not os.path.isdir(source):
+            source = self._download_hub_shards(str(source), tuple(src_prefixes))
+        keys = _read_checkpoint_keys(source)
+        mtp_state_dict = {}
+        for i in range(num_mtp):
+            src = src_prefixes[i]
+            dst = f"{layers_prefix}{self._hf_text_config.num_hidden_layers + i}."
+            for key in sorted(
+                k for k in keys if k.startswith(src) and not k.endswith("_scale_inv")
+            ):
+                mtp_state_dict[dst + key[len(src) :]] = get_safetensor(
+                    str(source), key, dequantize=True
+                )
+            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
+    def _download_hub_shards(repo_id: str, key_prefixes: tuple[str, ...]) -> str:
+        """Download only the Hub shards holding tensors under ``key_prefixes``; return the local dir."""
+        try:
+            index_file = hf_hub_download(repo_id, "model.safetensors.index.json")
+        except EntryNotFoundError:  # unsharded checkpoint
+            return snapshot_download(repo_id, allow_patterns=["model.safetensors"])
+        with open(index_file) as f:
+            weight_map = json.load(f)["weight_map"]
+        shards = sorted(
+            {shard for key, shard in weight_map.items() if key.startswith(key_prefixes)}
+        )
+        return snapshot_download(repo_id, allow_patterns=["model.safetensors.index.json", *shards])
+
     def _get_gated_delta_net_state_dict(self, layer, layer_id, is_mtp=False):
         """Export a GatedDeltaNet (Qwen3.5 linear-attention) layer's ``self_attention``."""
         gdn = layer.self_attention
@@ -920,6 +996,13 @@ def _get_gated_delta_net_state_dict(self, layer, layer_id, is_mtp=False):
         self.rules["linear_attn.out_norm"](gdn.out_norm, layer_id, is_mtp=is_mtp)
         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 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)
+
     def _get_mamba_layer_state_dict(self, layer, layer_id, is_mtp=False):
         if not isinstance(layer.norm, IdentityOp):
             self.rules["norm"](layer.norm, layer_id, is_mtp=is_mtp)