Add GLM-5.3-Flash and GLM-5.2 support to Megatron-Bridge PTQ and HF export

CHANGELOG.rst

@@ -26,6 +26,9 @@ 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 ``examples/megatron_bridge/quantize.py`` (MLflow's own ``MLFLOW_TRACKING_URI`` is honoured too), so a Megatron-Bridge PTQ run records the invocation, every argument as a searchable param, the resolved recipe, the master rank's log and the quantizer summary, and writes ``.experiment.json`` into ``--export_megatron_path``. The experiment defaults to ``$USER/megatron_bridge_quantize/<model basename>-<recipe name or --quant_cfg>`` and can be overridden with ``--mlflow_experiment`` / ``--mlflow_run_name``.
+- 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.
+- Add GLM-5.3-Flash (``glm5_next``) support to ``examples/megatron_bridge`` PTQ and unified HF export with the ``models/zai-org/GLM-5.3-Flash/ptq/nvfp4_experts_dense_mlp-kv_fp8_cast`` recipe, matching ``nvidia/GLM-5.3-Flash-NVFP4``. Requires a Megatron-Bridge and Megatron-Core with GLM-5.3-Flash support. DSA sparse-attention models also gain FP8 KV-cache quantization on Megatron.
+- 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``.
 
 *Misc*
 
@@ -77,6 +80,7 @@ Changelog
 
 **Bug Fixes**
 
+- Fix Megatron unified HF export of MoE models with grouped-GEMM experts when only the experts are quantized (e.g. ``nvfp4_experts_only-*`` recipes): ``hf_quant_config.json`` and the ``quantization_config`` in ``config.json`` were not written, so the quantized experts were served as unquantized weights. Re-export such checkpoints.
 - Fix shared ONNX export metadata and Diffusers attention policy: every ``NVFP4QuantExporter`` post-process now upgrades the default-domain opset to at least 23, all FP8 custom-op exports re-run ONNX shape/type inference after setting output metadata, and quantized SDPA derives FP8 MHA enablement from the live Q/K/V quantizers instead of honoring a caller-set ``_disable_fp8_mha`` attribute.
 - Fix ONNX FP16 conversion failing to preserve public output types when type inference changes a graph output declaration before output casts are inserted.
 - Fix ``examples/hf_ptq/hf_ptq.py`` discarding a completed PTQ run (no checkpoint exported) when the optional post-quantization sanity-check ``generate()`` call raised, for example because ``device_map="auto"`` placed part of the model on CPU. That failure is now caught and only skips the sanity check; export proceeds regardless.

examples/megatron_bridge/export_quantized_megatron_to_hf.py

@@ -22,7 +22,8 @@
 
 The HuggingFace unified exporter does not gather tensor-parallel-sharded weights, so this script
 always loads the checkpoint at tensor_model_parallel_size=1 (re-sharding from whatever TP was used
-during quantization). Use --pp_size to shard a large model across GPUs for export.
+during quantization). Use --pp_size and/or --ep_size to shard a large model across GPUs for
+export; --ep_size needs grouped-GEMM experts (not --no_moe_grouped_gemm).
 
 Example usage to export an FP8 checkpoint produced by quantize.py:
 
@@ -87,8 +88,14 @@ def get_args() -> argparse.Namespace:
         help="Export extra modules such as Medusa heads, EAGLE, or MTP.",
     )
 
-    # Only Pipeline parallelism is supported for export
+    # Only pipeline and expert parallelism are supported for export
     parser.add_argument("--pp_size", type=int, default=1, help="Pipeline parallel size")
+    parser.add_argument(
+        "--ep_size",
+        type=int,
+        default=1,
+        help="Expert parallel size (grouped-GEMM experts only); shards the experts across ranks",
+    )
     parser.add_argument(
         "--num_layers_in_first_pipeline_stage",
         type=int,
@@ -126,7 +133,7 @@ def main(args: argparse.Namespace):
         provider_overrides={
             "tensor_model_parallel_size": 1,  # Tensor parallelism is not supported
             "pipeline_model_parallel_size": args.pp_size,
-            "expert_model_parallel_size": 1,  # Expert parallelism is not supported
+            "expert_model_parallel_size": args.ep_size,
             "expert_tensor_parallel_size": 1,  # Expert tensor parallelism is not supported
             "num_layers_in_first_pipeline_stage": args.num_layers_in_first_pipeline_stage,
             "num_layers_in_last_pipeline_stage": args.num_layers_in_last_pipeline_stage,

modelopt/torch/export/plugins/mcore_common.py

@@ -18,6 +18,11 @@
 from typing import Any
 
 from .mcore_deepseek import deepseek_causal_lm_export, deepseek_causal_lm_import
+from .mcore_glm import (
+    GLM5_NEXT_VISION_PREFIXES,
+    glm5_next_causal_lm_export,
+    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 +68,8 @@
     "Qwen3VLForConditionalGeneration": qwen3vl_causal_lm_export,
     "Qwen3_5ForConditionalGeneration": qwen3_5_vl_causal_lm_export,
     "Qwen3_5MoeForConditionalGeneration": qwen3_5_vl_causal_lm_export,
+    "Glm5NextForConditionalGeneration": glm5_next_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
@@ -72,6 +79,7 @@
     "Qwen3VLForConditionalGeneration": QWEN3VL_VISION_PREFIXES,
     "Qwen3_5ForConditionalGeneration": QWEN3_5_VL_VISION_PREFIXES,
     "Qwen3_5MoeForConditionalGeneration": QWEN3_5_VL_VISION_PREFIXES,
+    "Glm5NextForConditionalGeneration": GLM5_NEXT_VISION_PREFIXES,
 }
 
 all_mcore_hf_import_mapping: dict[str, Any] = {

modelopt/torch/export/plugins/mcore_custom.py

@@ -200,6 +200,18 @@ def __init__(self, target_name_or_prefix: str = "", func_kwargs: dict[str, Any]
         )
 
 
+class KimiDeltaAttentionSlicing(CustomModuleMapping):
+    """A custom module mapping that splits KDA's fused q|k|v ``in_proj`` and ``conv1d``."""
+
+    def __init__(self, target_name_or_prefix: str = "", func_kwargs: dict[str, Any] = {}):
+        """Create a custom module mapping that splits the fused KDA projections."""
+        super().__init__(
+            func_name="kda_slicing",
+            target_name_or_prefix=target_name_or_prefix,
+            func_kwargs=func_kwargs,
+        )
+
+
 class PackNameRemapping(CustomModuleMapping):
     """A custom module mapping that packs module after name remapping."""
 

modelopt/torch/export/plugins/mcore_glm.py

@@ -0,0 +1,121 @@
+# 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.
+
+For GLM-5.3-Flash (``glm5_next``), Megatron-Bridge builds each HF decoder layer as two physical
+layers wrapped in mHC hyper-connections: attention (KDA or NoPE-MLA DSA) at ``2N`` and the dense MLP /
+MoE at ``2N+1``. The vision tower is copied from HF.
+
+Both keep the MTP layer at HF index ``num_hidden_layers`` under the decoder's names.
+"""
+
+from .mcore_custom import (
+    GatedMLPSlicing,
+    GroupedGatedMLPSlicing,
+    GroupedMLPSlicing,
+    KimiDeltaAttentionSlicing,
+    NameRemapping,
+    SelfAttentionScaling,
+    with_language_model_prefix,
+)
+from .mcore_deepseek import deepseek_causal_lm_export
+
+# Vision-tower weights copied straight from the HF checkpoint (never quantized).
+GLM5_NEXT_VISION_PREFIXES = ("model.visual.",)
+
+_glm5_next_causal_lm_export: dict = {
+    # Layer-level flags read by the exporter (see the module docstring).
+    "fold_attn_mlp_layer_pairs": True,
+    "mtp_in_decoder_layers": True,
+    "word_embeddings": NameRemapping("model.embed_tokens."),
+    "final_norm": NameRemapping("model.norm."),
+    "output_layer": NameRemapping("lm_head."),
+    # mHC hyper-connections: formatted with (hf_layer_id, "attn" | "ffn").
+    "hc_fn": NameRemapping("model.layers.{}.hc_{}_fn"),
+    "hc_base": NameRemapping("model.layers.{}.hc_{}_base"),
+    "hc_scale": NameRemapping("model.layers.{}.hc_{}_scale"),
+    "input_layernorm": NameRemapping("model.layers.{}.input_layernorm."),
+    # KDA linear attention (fused q|k|v in_proj and conv1d are split by ``kda``).
+    "kda": KimiDeltaAttentionSlicing("model.layers.{}.self_attn."),
+    "kda.beta_proj": NameRemapping("model.layers.{}.self_attn.b_proj."),
+    "kda.f_a_proj": NameRemapping("model.layers.{}.self_attn.f_a_proj."),
+    "kda.f_b_proj": NameRemapping("model.layers.{}.self_attn.f_b_proj."),
+    "kda.g_a_proj": NameRemapping("model.layers.{}.self_attn.g_a_proj."),
+    "kda.g_b_proj": NameRemapping("model.layers.{}.self_attn.g_b_proj."),
+    "kda.A_log": NameRemapping("model.layers.{}.self_attn.A_log"),
+    "kda.dt_bias": NameRemapping("model.layers.{}.self_attn.dt_bias"),
+    "kda.out_norm": NameRemapping("model.layers.{}.self_attn.o_norm."),
+    "kda.out_proj": NameRemapping("model.layers.{}.self_attn.o_proj."),
+    # NoPE MLA with the DSA kpool indexer
+    "linear_q_down_proj": NameRemapping("model.layers.{}.self_attn.q_a_proj."),
+    "linear_q_layernorm": NameRemapping("model.layers.{}.self_attn.q_a_layernorm."),
+    "linear_q_up_proj": NameRemapping("model.layers.{}.self_attn.q_b_proj."),
+    "linear_kv_down_proj": NameRemapping("model.layers.{}.self_attn.kv_a_proj_with_mqa."),
+    "linear_kv_layernorm": NameRemapping("model.layers.{}.self_attn.kv_a_layernorm."),
+    "linear_kv_up_proj": NameRemapping("model.layers.{}.self_attn.kv_b_proj."),
+    "linear_proj": NameRemapping("model.layers.{}.self_attn.o_proj."),
+    "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."),
+    "indexer.index_kpool_compress_ape": NameRemapping(
+        "model.layers.{}.self_attn.indexer.index_kpool_compress_ape"
+    ),
+    "indexer.index_kpool_compress_gate": NameRemapping(
+        "model.layers.{}.self_attn.indexer.index_kpool_compress_gate"
+    ),
+    # Dense MLP (the pre-MLP norm is fused into linear_fc1)
+    "pre_mlp_layernorm": NameRemapping("model.layers.{}.post_attention_layernorm."),
+    "fused_pre_mlp_layernorm": NameRemapping("model.layers.{}.post_attention_layernorm.weight"),
+    "linear_fc1": GatedMLPSlicing("model.layers.{}.mlp."),
+    "linear_fc2": NameRemapping("model.layers.{}.mlp.down_proj."),
+    # MoE
+    "router": NameRemapping(
+        "model.layers.{}.mlp.gate.", {"mapping": {"expert_bias": "e_score_correction_bias"}}
+    ),
+    "shared_experts.linear_fc1": GatedMLPSlicing("model.layers.{}.mlp.shared_experts."),
+    "shared_experts.linear_fc2": NameRemapping("model.layers.{}.mlp.shared_experts.down_proj."),
+    "local_experts.linear_fc1": GatedMLPSlicing("model.layers.{}.mlp.experts.{}."),
+    "local_experts.linear_fc2": NameRemapping("model.layers.{}.mlp.experts.{}.down_proj."),
+    "experts.linear_fc1": GroupedGatedMLPSlicing("model.layers.{}.mlp.experts.{{}}"),
+    "experts.linear_fc2": GroupedMLPSlicing("model.layers.{}.mlp.experts.{{}}.down_proj"),
+    # MTP (split e_proj / h_proj are concatenated back into eh_proj)
+    "mtp.enorm": NameRemapping("model.layers.{}.enorm."),
+    "mtp.hnorm": NameRemapping("model.layers.{}.hnorm."),
+    "mtp.eh_proj": NameRemapping("model.layers.{}.eh_proj.weight"),
+    "mtp.final_layernorm": NameRemapping("model.layers.{}.shared_head.norm."),
+}
+
+glm5_next_causal_lm_export = with_language_model_prefix(_glm5_next_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"),
+    "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/quant_utils.py

@@ -452,12 +452,6 @@ def uses_iq_quantization(module) -> bool:
 
     This reads ``num_bits`` directly rather than resolving each layer's full format, so an
     unrelated unsupported quantizer elsewhere in the model cannot turn the check into an error.
-
-    Known gap, shared with ``get_quantization_format``: ``weight_attr_names`` yields nothing for
-    a TEGroupedLinear, whose parameters are ``weight0..N`` while its quantizer is a single
-    ``GroupedQuantizer`` under ``weight_quantizer``. Neither function sees such a module, so an
-    experts-only IQ model reports no format at all -- not just here. Closing it belongs in
-    ``weight_attr_names``, where it affects every format, rather than in this helper.
     """
     for weight_name in weight_attr_names(module):
         weight_quantizer = representative_weight_quantizer(module, weight_name)