[3/3] Add GLM-5.3-Flash support to Megatron-Bridge PTQ and HF export

CHANGELOG.rst

@@ -21,6 +21,7 @@ Changelog
 - Add ``--sft_in_batch_packing``, ``--moe_aux_loss_coeff`` and ``--dsa_kernel_backend`` to ``examples/megatron_bridge/distill.py`` (``--dsa_kernel_backend`` also to ``quantize.py``): pack ``--sft_hf_dataset`` conversations per micro-batch instead of padding each to ``--seq_length``, override the MoE load-balancing loss weight (e.g. 0 for padded SFT data), and select fused DeepSeek sparse-attention kernels for long sequences.
 - Add ``--vllm_fakequant_export`` to ``examples/megatron_bridge/export_quantized_megatron_to_hf.py`` to evaluate formats vLLM has no kernels for (e.g. GGML IQ) as BF16 fake-quant checkpoints, and allow ``--ep_size`` > 1 for checkpoints built with ``--no_moe_grouped_gemm``.
 - GGML GPTQ checkpoints from ``examples/megatron_bridge`` keep their calibrated codes when the Megatron checkpoint is reloaded for QAD or export, and layerwise calibration now runs on Megatron vision-language models.
+- 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-BF16/ptq/nvfp4_experts_dense_mlp-kv_fp8_cast`` recipe, matching ``nvidia/GLM-5.3-Flash-NVFP4``, and two GGML IQ2_S recipes for its routed experts: ``models/zai-org/GLM-5.3-Flash-BF16/ptq/iq2_s_experts-q8_0_shared_dense-gptq`` and, with weight-only NVFP4 on the experts' down projections, ``iq2_s_experts-nvfp4_down-q8_0_shared_dense-gptq``. Requires a Megatron-Bridge and Megatron-Core with GLM-5.3-Flash support.
 
 *Misc*
 

modelopt/torch/export/plugins/mcore_common.py

@@ -18,7 +18,11 @@
 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_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,
@@ -66,6 +70,7 @@
     "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,
 }
 
@@ -77,6 +82,7 @@
     "Qwen3_5ForConditionalGeneration": QWEN3_5_VL_VISION_PREFIXES,
     "Qwen3_5MoeForConditionalGeneration": QWEN3_5_VL_VISION_PREFIXES,
     "NemotronH_Omni_Reasoning_V3": NEMOTRON_H_OMNI_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

@@ -15,19 +15,95 @@
 
 """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.
+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,
@@ -38,4 +114,8 @@
     "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."),
 }