diff --git a/CHANGELOG.rst b/CHANGELOG.rst
index 822768dabbf..706c64fc2de 100755
--- a/CHANGELOG.rst
+++ b/CHANGELOG.rst
@@ -35,6 +35,7 @@ Changelog
- 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.
- The Megatron unified HF export (``export_mcore_gpt_to_hf``) now copies the source checkpoint's non-model files (tokenizer, processor, remote code, chat templates, ...) into the export unchanged, from local disk only. The ``examples/megatron_bridge`` scripts make a Hugging Face Hub model local first with ``modelopt.torch.export.ensure_local_checkpoint``.
- Add an end-to-end W4A4 NVFP4 PTQ and QAD tutorial for NVIDIA Nemotron 3.5 Super VL, covering chat-data blend preparation, assistant-only distillation, Hugging Face export and vLLM serving. See `examples/megatron_bridge/tutorials/NVIDIA-Nemotron-3.5-Super-VL-120B-A12B-BF16/README.md `_ for details.
+- 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*
diff --git a/modelopt/torch/export/plugins/mcore_common.py b/modelopt/torch/export/plugins/mcore_common.py
index fbfbea57133..2066ff5c71f 100644
--- a/modelopt/torch/export/plugins/mcore_common.py
+++ b/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] = {
diff --git a/modelopt/torch/export/plugins/mcore_custom.py b/modelopt/torch/export/plugins/mcore_custom.py
index 89f41d03016..c7271ee08ea 100644
--- a/modelopt/torch/export/plugins/mcore_custom.py
+++ b/modelopt/torch/export/plugins/mcore_custom.py
@@ -207,6 +207,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."""
diff --git a/modelopt/torch/export/plugins/mcore_glm.py b/modelopt/torch/export/plugins/mcore_glm.py
index 25a967860ac..11ee40d63ed 100644
--- a/modelopt/torch/export/plugins/mcore_glm.py
+++ b/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."),
}
diff --git a/modelopt/torch/export/unified_export_megatron.py b/modelopt/torch/export/unified_export_megatron.py
index ee2dfcb24c2..304a7f4d2e4 100644
--- a/modelopt/torch/export/unified_export_megatron.py
+++ b/modelopt/torch/export/unified_export_megatron.py
@@ -106,6 +106,13 @@
from megatron.core.transformer.torch_norm import L2Norm
from megatron.core.transformer.transformer_layer import TransformerLayer
+ try:
+ from megatron.core.models.hybrid.layers.hybrid_hyper_connection import (
+ HyperConnectionHybridLayer,
+ )
+ except ImportError: # older Megatron-Core without mHC
+ HyperConnectionHybridLayer = None
+
has_mcore = True
__all__ = [
@@ -212,11 +219,9 @@ 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."
- )
+ if self.rules.get("fold_attn_mlp_layer_pairs", False):
+ # Each HF decoder layer is an attention and an MLP physical layer in Megatron.
+ self._hf_text_config.num_hidden_layers = language_model.config.num_layers // 2
# 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 ()
@@ -634,12 +639,9 @@ def _get_state_dict(self):
# Decoder layers
for layer in model.decoder.layers:
layer_id = layer.layer_number - 1
- if isinstance(layer, MambaLayer):
- self._get_mamba_layer_state_dict(layer, layer_id)
- elif isinstance(layer, TransformerLayer):
- self._get_transformer_layer_state_dict(layer, layer_id)
- else:
- raise ValueError("Only TransformerLayer or MambaLayer are supported.")
+ if self.rules.get("fold_attn_mlp_layer_pairs", False):
+ layer_id //= 2
+ self._get_decoder_layer_state_dict(layer, layer_id)
self._layer_state_dicts[layer.layer_number] = self._state_dict
if layer.layer_number != self.model.config.num_layers:
@@ -673,6 +675,32 @@ def _get_fused_norm_weight(self, module, primary_key: str = "fused_norm"):
return None, None
return fused_key, weight
+ def _get_decoder_layer_state_dict(
+ self, layer, layer_id, is_mtp=False, export_hyper_connection=True
+ ):
+ """Export one decoder layer, unwrapping an mHC ``HyperConnectionHybridLayer`` first."""
+ if HyperConnectionHybridLayer is not None and isinstance(layer, HyperConnectionHybridLayer):
+ if export_hyper_connection:
+ self._get_hyper_connection_state_dict(layer, layer_id)
+ layer = layer.inner_layer
+ if isinstance(layer, MambaLayer):
+ self._get_mamba_layer_state_dict(layer, layer_id, is_mtp=is_mtp)
+ elif isinstance(layer, TransformerLayer):
+ self._get_transformer_layer_state_dict(layer, layer_id, is_mtp=is_mtp)
+ else:
+ raise ValueError("Only TransformerLayer or MambaLayer are supported.")
+
+ def _get_hyper_connection_state_dict(self, layer, layer_id):
+ """Export an mHC wrapper as HF's ``hc_{attn,ffn}_{fn,base,scale}`` of its HF layer."""
+ attention = getattr(layer.inner_layer, "self_attention", None)
+ kind = "ffn" if attention is None or isinstance(attention, IdentityOp) else "attn"
+ hc = layer.hyper_connection
+ self.rules["hc_fn"](hc.mapping_proj.weight.detach().to(self.dtype), layer_id, kind)
+ # The HF checkpoint keeps the mHC bias and alpha scales in FP32.
+ self.rules["hc_base"](hc.bias.detach().float(), layer_id, kind)
+ alphas = torch.cat([hc.alpha_pre, hc.alpha_post, hc.alpha_res]).detach().float()
+ self.rules["hc_scale"](alphas, layer_id, kind)
+
def _get_transformer_layer_state_dict(self, layer, layer_id, is_mtp=False):
if not isinstance(layer.input_layernorm, IdentityOp):
self.rules["input_layernorm"](layer.input_layernorm, layer_id, is_mtp=is_mtp)
@@ -721,6 +749,8 @@ def _get_transformer_layer_state_dict(self, layer, layer_id, is_mtp=False):
indexer = getattr(core_attention, "indexer", None)
if indexer is not None:
self._get_dsa_indexer_state_dict(indexer, layer_id, is_mtp)
+ elif "kda" in self.rules and hasattr(layer.self_attention, "in_proj"):
+ self._get_kda_state_dict(layer, layer_id, is_mtp=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)
@@ -841,12 +871,20 @@ def _get_mtp_state_dict(self, copy_from_pretrained: bool = True) -> dict[str, to
saved_state_dict = self._state_dict
self._state_dict = OrderedDict()
try:
- for mtp_layer in mtp.layers:
+ for mtp_idx, mtp_layer in enumerate(mtp.layers):
# Some architectures (Qwen3.5) put a single TransformerLayer here, not a container.
inner = mtp_layer.mtp_model_layer
inner_layers = getattr(inner, "layers", None) or [inner]
- first_id = inner_layers[0].layer_number - 1
- last_id = inner_layers[-1].layer_number - 1
+ if self.rules.get("mtp_in_decoder_layers", False):
+ # HF stores the MTP layer after the decoder layers, under the decoder's names.
+ first_id = last_id = self._hf_text_config.num_hidden_layers + mtp_idx
+ inner_ids = [first_id] * len(inner_layers)
+ inner_is_mtp = False
+ else:
+ first_id = inner_layers[0].layer_number - 1
+ last_id = inner_layers[-1].layer_number - 1
+ inner_ids = [layer.layer_number - 1 for layer in inner_layers]
+ inner_is_mtp = True
# Outer predictor projections attach to the first inner HF index.
if "mtp.enorm" in self.rules:
@@ -854,19 +892,18 @@ def _get_mtp_state_dict(self, copy_from_pretrained: bool = True) -> dict[str, to
if "mtp.hnorm" in self.rules:
self.rules["mtp.hnorm"](mtp_layer.hnorm, first_id)
if "mtp.eh_proj" in self.rules:
- self.rules["mtp.eh_proj"](mtp_layer.eh_proj, first_id)
-
- # Inner layers reuse the base decoder walker (is_mtp=True).
- for inner in inner_layers:
- hf_layer_id = inner.layer_number - 1
- if isinstance(inner, MambaLayer):
- self._get_mamba_layer_state_dict(inner, hf_layer_id, is_mtp=True)
- elif isinstance(inner, TransformerLayer):
- self._get_transformer_layer_state_dict(inner, hf_layer_id, is_mtp=True)
- else:
- raise ValueError(
- "Only TransformerLayer or MambaLayer are supported in the MTP block."
- )
+ if getattr(mtp_layer, "eh_proj", None) is not None:
+ self.rules["mtp.eh_proj"](mtp_layer.eh_proj, first_id)
+ else: # split projections (GLM-5.3-Flash): HF eh_proj is [e_proj | h_proj]
+ eh_proj = torch.cat([mtp_layer.e_proj.weight, mtp_layer.h_proj.weight], 1)
+ self.rules["mtp.eh_proj"](eh_proj.detach().to(self.dtype), first_id)
+
+ # Inner layers reuse the base decoder walker. The MTP mHC weights have no HF
+ # counterpart (Megatron-Bridge initializes them on import), so they are dropped.
+ for inner, hf_layer_id in zip(inner_layers, inner_ids):
+ self._get_decoder_layer_state_dict(
+ inner, hf_layer_id, is_mtp=inner_is_mtp, export_hyper_connection=False
+ )
# The MTP block's own final layernorm attaches to the last inner HF index.
final_layernorm = getattr(mtp_layer, "final_layernorm", None)
@@ -1008,12 +1045,31 @@ 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_kda_state_dict(self, layer, layer_id, is_mtp=False):
+ """Export a KDA (Kimi Delta Attention) layer's ``self_attention``."""
+ kda = layer.self_attention
+ self.rules["kda"](kda, layer_id, is_mtp=is_mtp)
+ for name in ("beta_proj", "f_a_proj", "f_b_proj", "g_a_proj", "g_b_proj"):
+ self.rules[f"kda.{name}"](getattr(kda, name), layer_id, is_mtp=is_mtp)
+ self.rules["kda.out_norm"](kda.out_norm, layer_id, is_mtp=is_mtp)
+ self.rules["kda.out_proj"](kda.out_proj, layer_id, is_mtp=is_mtp)
+ # The HF checkpoint keeps the kernel parameters in FP32.
+ self.rules["kda.A_log"](kda.A_log.detach().float(), layer_id, is_mtp=is_mtp)
+ self.rules["kda.dt_bias"](kda.dt_bias.detach().float(), 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."""
+ """Export the DSA kpool 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):
@@ -1151,6 +1207,7 @@ def _custom_mapping_to_lambda(mapping):
"self_attention_scaling": self._self_attention_scaling,
"gated_mlp_slicing": self._gated_mlp_slicing,
"gated_delta_net_slicing": self._gated_delta_net_slicing,
+ "kda_slicing": self._kda_slicing,
"grouped_mlp_slicing": self._grouped_mlp_slicing,
"pack_name_remapping": self._pack_name_remapping,
"pack_name_remapping_gpt_oss": self._pack_name_remapping_gpt_oss,
@@ -1910,11 +1967,6 @@ def _gated_delta_net_slicing(self, module, prefix, is_mtp=False):
"""
if is_mtp:
prefix = self._mtp_prefix(prefix)
- in_proj = module.in_proj
- name_to_value, qformat, block_size = self._get_quantized_state(
- in_proj, self.dtype, prefix=prefix
- )
-
assert tuple(module.in_proj_split_names) == (
"query",
"key",
@@ -1934,12 +1986,39 @@ def _gated_delta_net_slicing(self, module, prefix, is_mtp=False):
sections["alpha"],
]
proj_names = ("in_proj_qkv", "in_proj_z", "in_proj_b", "in_proj_a")
- proj_prefixes = [prefix + name + "." for name in proj_names]
# The recipes keep the alpha / beta gates in BF16, but Megatron fuses all six sections
# behind one quantizer, so they can only be dropped here rather than by a quantizer_name.
- keep_bf16 = {
- p for p, n in zip(proj_prefixes, proj_names) if n in ("in_proj_a", "in_proj_b")
- }
+ self._split_fused_projection(
+ module.in_proj,
+ prefix,
+ proj_names,
+ split_sizes,
+ keep_bf16_names=("in_proj_a", "in_proj_b"),
+ )
+
+ def _kda_slicing(self, module, prefix, is_mtp=False):
+ """Split KDA's fused q|k|v ``in_proj`` and depthwise ``conv1d`` into HF's q/k/v tensors."""
+ if is_mtp:
+ prefix = self._mtp_prefix(prefix)
+ assert tuple(module.in_proj_split_names) == ("query", "key", "value"), (
+ f"Unexpected KDA in_proj layout {tuple(module.in_proj_split_names)}; only the "
+ "two-stage-gate layout [query, key, value] is supported"
+ )
+ split_sizes = list(module.in_proj_split_sections)
+ self._split_fused_projection(
+ module.in_proj, prefix, ("q_proj", "k_proj", "v_proj"), split_sizes
+ )
+ conv_weights = torch.split(module.conv1d.weight.detach().to(self.dtype), split_sizes)
+ for name, weight in zip(("q_conv1d", "k_conv1d", "v_conv1d"), conv_weights):
+ self._state_dict[prefix + name + ".weight"] = weight
+
+ def _split_fused_projection(self, fused, prefix, proj_names, split_sizes, keep_bf16_names=()):
+ """Export a row-fused projection as the HF projections ``prefix + proj_names[i]``."""
+ name_to_value, qformat, block_size = self._get_quantized_state(
+ fused, self.dtype, prefix=prefix
+ )
+ proj_prefixes = [prefix + name + "." for name in proj_names]
+ keep_bf16 = {p for p, n in zip(proj_prefixes, proj_names) if n in keep_bf16_names}
for proj_prefix in proj_prefixes:
if proj_prefix in keep_bf16:
@@ -1969,7 +2048,7 @@ def _gated_delta_net_slicing(self, module, prefix, is_mtp=False):
proj_prefix + "weight",
proj_weight,
qformat,
- getattr(in_proj, "weight_quantizer", None),
+ getattr(fused, "weight_quantizer", None),
)
)
elif weight_scale is None:
diff --git a/modelopt_recipes/models/zai-org/GLM-5.3-Flash-BF16/ptq/iq2_s_experts-nvfp4_down-q8_0_shared_dense-gptq.yaml b/modelopt_recipes/models/zai-org/GLM-5.3-Flash-BF16/ptq/iq2_s_experts-nvfp4_down-q8_0_shared_dense-gptq.yaml
new file mode 100644
index 00000000000..3c8e6ab1497
--- /dev/null
+++ b/modelopt_recipes/models/zai-org/GLM-5.3-Flash-BF16/ptq/iq2_s_experts-nvfp4_down-q8_0_shared_dense-gptq.yaml
@@ -0,0 +1,68 @@
+# 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.
+
+# Recipe for zai-org/GLM-5.3-Flash-BF16 with Megatron-Bridge: GGML IQ2_S on the routed experts' gate/up projections,
+# NVFP4 (weight-only) on their down projections, and Q8_0 on the shared experts and the dense MLP of layers 0-2,
+# calibrated with layerwise GPTQ. Both attention families (KDA and sparse MLA), the router, the vision tower, the
+# MTP layer, embeddings and lm_head stay BF16.
+#
+# Megatron-Bridge only (examples/megatron_bridge): the patterns use Megatron module names, and
+# GPTQ skips the fused routed experts of the HF model. Run quantize.py, distill.py and export
+# with --no_moe_grouped_gemm: GPTQ and the GGML export only handle per-expert linears. vLLM has no IQ2 kernels, so evaluate an export made with
+# --vllm_fakequant_export, which folds the IQ2 and NVFP4 weights into BF16.
+#
+# Every quantized linear's input width is a multiple of the 256-weight GGML block
+# (hidden 4096, dense MLP 12288, experts 2048).
+
+# modelopt-schema: modelopt.recipe.config.ModelOptPTQRecipe
+imports:
+ base_disable_all: configs/ptq/units/base_disable_all
+ iq2_s: configs/numerics/iq2_s
+ nvfp4_static: configs/numerics/nvfp4_static
+ q8_0: configs/numerics/q8_0
+
+metadata:
+ description: >-
+ GLM-5.3-Flash-BF16 (Megatron-Bridge): GGML IQ2_S on the routed experts' gate/up, weight-only NVFP4 on their
+ down projections and Q8_0 on the shared experts and dense MLP, calibrated with layerwise GPTQ. Attention,
+ router, vision tower, MTP layer, embeddings and lm_head stay BF16. Requires calibration data.
+quantize:
+ algorithm:
+ method: gptq
+ # GGML encodes 256-weight blocks jointly, so GPTQ updates one block at a time.
+ block_size: 256
+ perc_damp: 0.3
+ importance_weighted: true
+ layerwise: {enable: true}
+ quant_cfg:
+ - $import: base_disable_all
+ - quantizer_name: '*mlp*weight_quantizer'
+ cfg:
+ $import: iq2_s
+ # The routed experts' down projections, which take the gated activation, get 4.5 instead of 2.56 bits.
+ - quantizer_name: '*mlp.experts.local_experts.*.linear_fc2.weight_quantizer'
+ cfg:
+ $import: nvfp4_static
+ # Every token runs through the shared experts and the dense MLP, and they are ~0.5% of the MLP
+ # weights, so Q8_0 there costs almost nothing in size.
+ - quantizer_name: '*mlp.shared_experts*weight_quantizer'
+ cfg:
+ $import: q8_0
+ - quantizer_name: '*mlp.linear_fc*weight_quantizer'
+ cfg:
+ $import: q8_0
+ - {quantizer_name: '*visual*', enable: false}
+ # Megatron-Bridge builds the MTP layer (HF does not); keep it BF16 like the HF export.
+ - {quantizer_name: '*mtp*', enable: false}
diff --git a/modelopt_recipes/models/zai-org/GLM-5.3-Flash-BF16/ptq/iq2_s_experts-q8_0_shared_dense-gptq.yaml b/modelopt_recipes/models/zai-org/GLM-5.3-Flash-BF16/ptq/iq2_s_experts-q8_0_shared_dense-gptq.yaml
new file mode 100644
index 00000000000..82b37f54fba
--- /dev/null
+++ b/modelopt_recipes/models/zai-org/GLM-5.3-Flash-BF16/ptq/iq2_s_experts-q8_0_shared_dense-gptq.yaml
@@ -0,0 +1,62 @@
+# 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.
+
+# Recipe for zai-org/GLM-5.3-Flash-BF16 with Megatron-Bridge: GGML IQ2_S on all routed-expert projections and
+# Q8_0 on the shared experts and the dense MLP of layers 0-2, calibrated with layerwise GPTQ. Both attention
+# families (KDA and sparse MLA), the router, the vision tower, the MTP layer, embeddings and lm_head stay BF16.
+#
+# Megatron-Bridge only (examples/megatron_bridge): the patterns use Megatron module names, and
+# GPTQ skips the fused routed experts of the HF model. Run quantize.py, distill.py and export
+# with --no_moe_grouped_gemm: GPTQ and the GGML export only handle per-expert linears. vLLM has no IQ2 kernels, so evaluate an export made with
+# --vllm_fakequant_export, which folds the IQ2 weights into BF16.
+#
+# Every quantized linear's input width is a multiple of the 256-weight GGML block
+# (hidden 4096, dense MLP 12288, experts 2048).
+
+# modelopt-schema: modelopt.recipe.config.ModelOptPTQRecipe
+imports:
+ base_disable_all: configs/ptq/units/base_disable_all
+ iq2_s: configs/numerics/iq2_s
+ q8_0: configs/numerics/q8_0
+
+metadata:
+ description: >-
+ GLM-5.3-Flash-BF16 (Megatron-Bridge): GGML IQ2_S on the routed experts and Q8_0 on the shared experts and
+ dense MLP, calibrated with layerwise GPTQ. Attention, router, vision tower, MTP layer, embeddings and
+ lm_head stay BF16. Requires calibration data.
+quantize:
+ algorithm:
+ method: gptq
+ # GGML encodes 256-weight blocks jointly, so GPTQ updates one block at a time.
+ block_size: 256
+ perc_damp: 0.3
+ importance_weighted: true
+ layerwise: {enable: true}
+ quant_cfg:
+ - $import: base_disable_all
+ - quantizer_name: '*mlp*weight_quantizer'
+ cfg:
+ $import: iq2_s
+ # Every token runs through the shared experts and the dense MLP, and they are ~0.5% of the MLP
+ # weights, so Q8_0 there costs almost nothing in size.
+ - quantizer_name: '*mlp.shared_experts*weight_quantizer'
+ cfg:
+ $import: q8_0
+ - quantizer_name: '*mlp.linear_fc*weight_quantizer'
+ cfg:
+ $import: q8_0
+ - {quantizer_name: '*visual*', enable: false}
+ # Megatron-Bridge builds the MTP layer (HF does not); keep it BF16 like the HF export.
+ - {quantizer_name: '*mtp*', enable: false}
diff --git a/modelopt_recipes/models/zai-org/GLM-5.3-Flash-BF16/ptq/nvfp4_experts_dense_mlp-kv_fp8_cast.yaml b/modelopt_recipes/models/zai-org/GLM-5.3-Flash-BF16/ptq/nvfp4_experts_dense_mlp-kv_fp8_cast.yaml
index d897eb79f27..59cb92c7715 100644
--- a/modelopt_recipes/models/zai-org/GLM-5.3-Flash-BF16/ptq/nvfp4_experts_dense_mlp-kv_fp8_cast.yaml
+++ b/modelopt_recipes/models/zai-org/GLM-5.3-Flash-BF16/ptq/nvfp4_experts_dense_mlp-kv_fp8_cast.yaml
@@ -29,10 +29,10 @@
# so only 3 of the 45 layers have a plain MLP -- 9 modules in total (`mlp.gate_proj`,
# `mlp.up_proj`, `mlp.down_proj` each). The other 42 layers carry `mlp.experts..*`.
#
-# The one rule that is not obvious is the trailing `*visual*` disable, and it is load-bearing.
+# The one rule that is not obvious is the `*visual*` disable, and it is load-bearing.
# The vision tower reuses the same leaf names -- `model.visual.blocks..mlp.gate_proj` and
# friends, across 24 blocks -- so the dense-MLP patterns match 144 modules inside it. Entries
-# apply in order, so the disable has to come LAST or the vision tower is quantized by accident.
+# apply in order, so the disable has to come AFTER them or the vision tower is quantized by accident.
# (`*.experts.*` matches nothing under `model.visual.*`, so only the dense-MLP patterns reach it.)
#
# Two model-specific notes:
@@ -41,10 +41,12 @@
# layers nest under `model.language_model.layers` and layerwise_calibrate cannot find
# them.
#
-# * The MTP (next-token-prediction) head is not quantized because it is not built. The config
-# declares `num_hidden_layers: 45` (with `num_nextn_predict_layers: 1`), so the HF model class
-# instantiates decoder layers 0-44 only and does not construct the MTP layer -- it is therefore
-# neither quantized nor carried into the exported checkpoint.
+# * The MTP (next-token-prediction) head stays BF16. The HF model class does not build it
+# (decoder layers 0-44 only); Megatron-Bridge does, so the trailing `*mtp*` disable keeps it off.
+#
+# The same recipe also drives Megatron-Bridge `quantize.py`, where the dense MLP is
+# `mlp.linear_fc1` / `mlp.linear_fc2`. Those patterns need a literal `mlp.linear_fc`, so they
+# skip `mlp.shared_experts.linear_fc*` and `mlp.experts.local_experts..linear_fc*`.
#
# The shared `default_disabled_quantizers` unit is deliberately not imported.
# Of its patterns only `*visual*` changes anything here; every other one either matches no
@@ -63,8 +65,8 @@ metadata:
description: >-
GLM-5.3-Flash-BF16: NVFP4 (W4A4) on the routed experts and on the dense MLP of layers 0-2,
plus an FP8 KV cache in cast mode using constant amax; max calibration. Shared experts,
- router gate, KDA and MLA attention, vision tower, embeddings and lm_head stay BF16. The
- MTP layer is not built at num_hidden_layers=45, so it is neither quantized nor exported.
+ router gate, KDA and MLA attention, vision tower, MTP layer, embeddings and lm_head
+ stay BF16. Works for both HF (hf_ptq.py) and Megatron-Bridge (quantize.py) module names.
quantize:
algorithm:
method: max
@@ -100,6 +102,21 @@ quantize:
- quantizer_name: '*mlp.down_proj*input_quantizer'
cfg:
$import: nvfp4
+ # Dense MLP, Megatron-Bridge names (fc1 = fused gate_proj + up_proj).
+ - quantizer_name: '*mlp.linear_fc1*weight_quantizer'
+ cfg:
+ $import: nvfp4
+ - quantizer_name: '*mlp.linear_fc1*input_quantizer'
+ cfg:
+ $import: nvfp4
+ - quantizer_name: '*mlp.linear_fc2*weight_quantizer'
+ cfg:
+ $import: nvfp4
+ - quantizer_name: '*mlp.linear_fc2*input_quantizer'
+ cfg:
+ $import: nvfp4
- $import: kv_fp8_cast
- # MUST stay last: keeps the vision tower BF16 after the dense-MLP patterns matched it.
+ # MUST follow the dense-MLP patterns: keeps the vision tower BF16 after they matched it.
- {quantizer_name: '*visual*', enable: false}
+ # Megatron-Bridge builds the MTP layer (HF does not); keep it BF16 like the HF export.
+ - {quantizer_name: '*mtp*', enable: false}
diff --git a/tests/_test_utils/torch/transformers_models.py b/tests/_test_utils/torch/transformers_models.py
index bdf910981d2..f1b51a06402 100644
--- a/tests/_test_utils/torch/transformers_models.py
+++ b/tests/_test_utils/torch/transformers_models.py
@@ -639,6 +639,56 @@ def create_tiny_deepseek_v3_dir(
)
+##### GLM-5.3-FLASH #####
+def get_tiny_glm5_next(**text_config_kwargs) -> PreTrainedModel:
+ """Tiny GLM-5.3-Flash VLM: one KDA + dense-MLP layer, one sparse-MLA + MoE layer, one vision block."""
+ set_seed(SEED)
+
+ # Lazy import — glm5_next is native only from transformers 5.16.1.
+ from transformers import Glm5NextConfig
+
+ text_kwargs = {
+ "vocab_size": 128,
+ "hidden_size": 64,
+ "intermediate_size": 128,
+ "moe_intermediate_size": 32,
+ "num_hidden_layers": 2,
+ "layer_types": ["linear_attention", "indexed_attention"],
+ "mlp_layer_types": ["dense", "sparse"],
+ "indexer_types": ["full", "full"],
+ "num_attention_heads": 4,
+ "num_key_value_heads": 4,
+ "q_lora_rank": 32,
+ "kv_lora_rank": 16,
+ "qk_nope_head_dim": 16,
+ "v_head_dim": 16,
+ "linear_num_heads": 4,
+ "linear_head_dim": 16,
+ "index_n_heads": 2,
+ "index_head_dim": 16,
+ "index_topk": 8,
+ "index_kpool": 4,
+ "n_routed_experts": 4,
+ "num_experts_per_tok": 2,
+ "n_shared_experts": 1,
+ "max_position_embeddings": 64,
+ "pad_token_id": 0,
+ }
+ text_kwargs.update(text_config_kwargs)
+ vision_kwargs = {
+ "depth": 1,
+ "hidden_size": 32,
+ "intermediate_size": 64,
+ "num_heads": 2,
+ "image_size": 28,
+ "patch_size": 14,
+ "out_hidden_size": text_kwargs["hidden_size"],
+ "projection_intermediate_size": 64,
+ }
+ cfg = Glm5NextConfig(text_config=text_kwargs, vision_config=vision_kwargs)
+ return AutoModelForImageTextToText.from_config(cfg, dtype=torch.bfloat16)
+
+
##### GLM-5.3-FLASH / DEEPSEEK-V4 (config-only, for vLLM ``load_format="dummy"``) #####
def _write_config_dir(dir_path: Path, config: dict) -> Path:
dir_path.mkdir(parents=True, exist_ok=True)
diff --git a/tests/gpu_megatron/torch/export/plugins/test_mcore_export_mappings.py b/tests/gpu_megatron/torch/export/plugins/test_mcore_export_mappings.py
index 965c44d7c29..f8aa35ca484 100644
--- a/tests/gpu_megatron/torch/export/plugins/test_mcore_export_mappings.py
+++ b/tests/gpu_megatron/torch/export/plugins/test_mcore_export_mappings.py
@@ -38,6 +38,7 @@
"Qwen3_5ForConditionalGeneration": "model.language_model.layers.{}.self_attn.",
"Qwen3_5MoeForConditionalGeneration": "model.language_model.layers.{}.self_attn.",
"NemotronH_Omni_Reasoning_V3": "language_model.backbone.layers.{}.mixer.",
+ "Glm5NextForConditionalGeneration": "model.language_model.layers.{}.self_attn.",
"GlmMoeDsaForCausalLM": "model.layers.{}.self_attn.",
}
@@ -60,6 +61,7 @@ def test_export_mapping_emits_kv_cache_scales(arch, prefix):
"Qwen3MoeForCausalLM": "model.layers.{}.mlp.experts.{}",
"NemotronHForCausalLM": "backbone.layers.{}.mixer.experts.{}",
"NemotronH_Omni_Reasoning_V3": "language_model.backbone.layers.{}.mixer.experts.{}",
+ "Glm5NextForConditionalGeneration": "model.language_model.layers.{}.mlp.experts.{}",
"GlmMoeDsaForCausalLM": "model.layers.{}.mlp.experts.{}",
}
@@ -105,3 +107,33 @@ def test_new_model_nests_every_target_under_language_model():
assert vl["mtp.eh_proj"].target_name_or_prefix.startswith("language_model.mtp.")
passthrough = all_mcore_hf_vision_passthrough_mapping["NemotronH_Omni_Reasoning_V3"]
assert not any(p.startswith("language_model.") for p in passthrough)
+
+
+def test_glm5_next_names_match_released_checkpoint():
+ """Rule targets must format to tensor names of the released zai-org/GLM-5.3-Flash checkpoint."""
+ mapping = all_mcore_hf_export_mapping["Glm5NextForConditionalGeneration"]
+ assert mapping["fold_attn_mlp_layer_pairs"] is True
+ assert mapping["mtp_in_decoder_layers"] is True
+
+ layer = "model.language_model.layers.3."
+ expected = {
+ ("hc_fn", 3, "attn"): layer + "hc_attn_fn",
+ ("hc_base", 3, "ffn"): layer + "hc_ffn_base",
+ ("hc_scale", 3, "ffn"): layer + "hc_ffn_scale",
+ ("kda", 3): layer + "self_attn.",
+ ("kda.beta_proj", 3): layer + "self_attn.b_proj.",
+ ("kda.out_norm", 3): layer + "self_attn.o_norm.",
+ ("kda.A_log", 3): layer + "self_attn.A_log",
+ ("linear_kv_down_proj", 3): layer + "self_attn.kv_a_proj_with_mqa.",
+ ("indexer.linear_wq_b", 3): layer + "self_attn.indexer.wq_b.",
+ ("indexer.index_kpool_compress_ape", 3): layer
+ + "self_attn.indexer.index_kpool_compress_ape",
+ ("fused_pre_mlp_layernorm", 3): layer + "post_attention_layernorm.weight",
+ ("shared_experts.linear_fc2", 3): layer + "mlp.shared_experts.down_proj.",
+ ("mtp.eh_proj", 3): layer + "eh_proj.weight",
+ ("mtp.final_layernorm", 3): layer + "shared_head.norm.",
+ }
+ for (rule, *args), name in expected.items():
+ assert mapping[rule].target_name_or_prefix.format(*args) == name, rule
+ assert mapping["kda"].func_name == "kda_slicing"
+ assert mapping["router"].func_kwargs["mapping"] == {"expert_bias": "e_score_correction_bias"}
diff --git a/tests/gpu_megatron/torch/export/test_unified_export_megatron.py b/tests/gpu_megatron/torch/export/test_unified_export_megatron.py
index c70e5bb432f..b6b625c84b0 100644
--- a/tests/gpu_megatron/torch/export/test_unified_export_megatron.py
+++ b/tests/gpu_megatron/torch/export/test_unified_export_megatron.py
@@ -307,6 +307,59 @@ def test_megatron_gated_delta_net_slicing_exports_ggml_payloads(qformat):
)
+def test_megatron_kda_slicing_splits_fused_qkv_and_conv1d():
+ weight = torch.randn(7, 4).bfloat16()
+ conv = torch.randn(7, 1, 4).bfloat16()
+ module = SimpleNamespace(
+ in_proj=object(),
+ in_proj_split_names=("query", "key", "value"),
+ in_proj_split_sections=(2, 2, 3),
+ conv1d=SimpleNamespace(weight=conv),
+ )
+ exporter = _make_ggml_exporter()
+ exporter._get_quantized_state = lambda *a, **k: ({"weight": weight}, None, None)
+
+ exporter._kda_slicing(module, "model.layers.0.self_attn.")
+
+ for name, rows in (("q", slice(0, 2)), ("k", slice(2, 4)), ("v", slice(4, 7))):
+ torch.testing.assert_close(
+ exporter._state_dict[f"model.layers.0.self_attn.{name}_proj.weight"], weight[rows]
+ )
+ torch.testing.assert_close(
+ exporter._state_dict[f"model.layers.0.self_attn.{name}_conv1d.weight"], conv[rows]
+ )
+ assert exporter.exclude_modules == [
+ f"model.layers.0.self_attn.{name}_proj" for name in ("q", "k", "v")
+ ]
+
+
+@pytest.mark.parametrize(("self_attention", "kind"), [(object(), "attn"), (None, "ffn")])
+def test_megatron_hyper_connection_exports_hf_hc_tensors(self_attention, kind):
+ hc = SimpleNamespace(
+ mapping_proj=SimpleNamespace(weight=torch.randn(24, 64)),
+ bias=torch.randn(24),
+ alpha_pre=torch.tensor([0.1]),
+ alpha_post=torch.tensor([0.2]),
+ alpha_res=torch.tensor([0.3]),
+ )
+ layer = SimpleNamespace(
+ hyper_connection=hc, inner_layer=SimpleNamespace(self_attention=self_attention)
+ )
+ exporter = _make_ggml_exporter()
+ exporter.rules = exporter._populate_rule_book()["Glm5NextForConditionalGeneration"]
+
+ exporter._get_hyper_connection_state_dict(layer, 2)
+
+ prefix = f"model.language_model.layers.2.hc_{kind}_"
+ assert sorted(exporter._state_dict) == [prefix + n for n in ("base", "fn", "scale")]
+ assert exporter._state_dict[prefix + "fn"].dtype == torch.bfloat16
+ # The released checkpoint keeps the mHC bias and alpha scales in FP32.
+ assert exporter._state_dict[prefix + "base"].dtype == torch.float32
+ torch.testing.assert_close(
+ exporter._state_dict[prefix + "scale"], torch.tensor([0.1, 0.2, 0.3])
+ )
+
+
@pytest.mark.parametrize("qformat", GGML_FORMAT_NAMES)
def test_megatron_packed_experts_reject_ggml_without_deployment_loader(qformat):
experts = _make_ggml_experts(qformat, "linear_fc2")
@@ -1173,30 +1226,6 @@ 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"
@@ -1218,6 +1247,8 @@ def test_mtp_state_dict_copies_decoder_mtp_layers(tmp_path):
exporter._src_num_hidden_layers = 2
exporter._hf_text_config = SimpleNamespace(num_hidden_layers=1, num_nextn_predict_layers=1)
+ # Non-writer ranks skip the copy.
+ assert exporter._get_mtp_state_dict(copy_from_pretrained=False) == {}
mtp_state_dict = exporter._get_mtp_state_dict()
assert sorted(mtp_state_dict) == [
diff --git a/tests/unit/recipe/test_glm_5_3_recipe.py b/tests/unit/recipe/test_glm_5_3_recipe.py
index de0c7c73942..cd80a639403 100644
--- a/tests/unit/recipe/test_glm_5_3_recipe.py
+++ b/tests/unit/recipe/test_glm_5_3_recipe.py
@@ -22,24 +22,29 @@
*not* matched and the shared experts stay BF16.
* The vision tower reuses the language MLP's leaf names (``mlp.gate_proj`` /
``up_proj`` / ``down_proj``), so the dense-MLP patterns match ``model.visual.*``
- too -- only the trailing ``*visual*`` disable (which must stay last) keeps the
+ too -- only the ``*visual*`` disable (which must follow them) keeps the
vision tower in BF16.
-* ``*mlp.gate_proj*`` must not catch the router ``mlp.gate``.
This pins that behaviour so it can't silently drift.
"""
+import pytest
import torch.nn as nn
+from _test_utils.torch import transformers_models
import modelopt.torch.quantization as mtq
from modelopt.recipe import load_recipe
_RECIPE = "models/zai-org/GLM-5.3-Flash-BF16/ptq/nvfp4_experts_dense_mlp-kv_fp8_cast"
+_IQ2_RECIPE = "models/zai-org/GLM-5.3-Flash-BF16/ptq/iq2_s_experts-q8_0_shared_dense-gptq"
+_IQ2_NVFP4_DOWN_RECIPE = (
+ "models/zai-org/GLM-5.3-Flash-BF16/ptq/iq2_s_experts-nvfp4_down-q8_0_shared_dense-gptq"
+)
_H = 32
class _MLP(nn.Module):
- """Plain MLP leaf names, shared by the dense MLP, each routed/shared expert, and vision."""
+ """Plain MLP leaf names, shared by the dense MLP, the shared experts, and vision."""
def __init__(self):
super().__init__()
@@ -48,78 +53,31 @@ def __init__(self):
self.down_proj = nn.Linear(_H, _H, bias=False)
-class _MoE(nn.Module):
+class _VisionAttention(nn.Module):
def __init__(self):
super().__init__()
- self.experts = nn.ModuleList([_MLP(), _MLP()])
- self.shared_experts = _MLP()
- self.gate = nn.Linear(_H, 2, bias=False) # router
-
-
-class _KDA(nn.Module):
- """KDA linear attention (projections plus a depthwise causal conv1d)."""
-
- def __init__(self):
- super().__init__()
- self.in_proj_qkvz = nn.Linear(_H, _H, bias=False)
- self.conv1d = nn.Conv1d(_H, _H, kernel_size=3, groups=_H, bias=False)
- self.out_proj = nn.Linear(_H, _H, bias=False)
-
-
-class _MLA(nn.Module):
- """NoPE sparse-MLA attention (and the vision attention, which reuses the leaf names)."""
-
- def __init__(self):
- super().__init__()
- self.q_proj = nn.Linear(_H, _H, bias=False)
- self.kv_proj = nn.Linear(_H, _H, bias=False)
- self.o_proj = nn.Linear(_H, _H, bias=False)
-
-
-class _DenseLayer(nn.Module):
- """Layers 0-2: KDA attention + a plain (dense) MLP."""
-
- def __init__(self):
- super().__init__()
- self.linear_attn = _KDA()
- self.mlp = _MLP()
-
-
-class _SparseLayer(nn.Module):
- """Layers 3-44: MLA attention + an MoE block."""
-
- def __init__(self):
- super().__init__()
- self.self_attn = _MLA()
- self.mlp = _MoE()
+ self.qkv = nn.Linear(_H, 3 * _H)
+ self.proj = nn.Linear(_H, _H)
class _VisionBlock(nn.Module):
def __init__(self):
super().__init__()
self.mlp = _MLP() # same gate_proj / up_proj / down_proj leaf names as the language MLP
- self.attn = _MLA()
-
-
-class _GLM53Flash(nn.Module):
- """Tiny stand-in for the ``glm5_next`` VLM MoE (one dense + one sparse layer + vision)."""
-
- def __init__(self):
- super().__init__()
- self.model = nn.Module()
- self.model.language_model = nn.Module()
- self.model.language_model.layers = nn.ModuleList([_DenseLayer(), _SparseLayer()])
- self.model.visual = nn.Module()
- self.model.visual.blocks = nn.ModuleList([_VisionBlock()])
- self.lm_head = nn.Linear(_H, _H, bias=False)
+ self.attn = _VisionAttention()
def _nvfp4(quantizer):
return quantizer.is_enabled and quantizer.num_bits == (2, 1)
+def _fp8(quantizer):
+ return quantizer.is_enabled and quantizer.num_bits == (4, 3)
+
+
def test_glm_5_3_recipe_quantizer_precedence():
- model = _GLM53Flash()
+ pytest.importorskip("transformers.models.glm5_next", reason="needs transformers>=5.16.1")
+ model = transformers_models.get_tiny_glm5_next()
config = load_recipe(_RECIPE).quantize.model_dump()
# The recipe uses plain max calibration; here we only assert quantizer placement,
@@ -128,45 +86,145 @@ def test_glm_5_3_recipe_quantizer_precedence():
config["algorithm"] = None
mtq.quantize(model, config)
- dense = model.model.language_model.layers[0]
- sparse = model.model.language_model.layers[1]
+ dense, sparse = model.model.language_model.layers
- # Routed experts -> NVFP4 W4A4.
- for expert in sparse.mlp.experts:
- for proj in (expert.gate_proj, expert.up_proj, expert.down_proj):
- assert _nvfp4(proj.weight_quantizer)
- assert _nvfp4(proj.input_quantizer)
+ # Routed experts (fused into 3D params, one weight quantizer per expert) -> NVFP4 W4A4.
+ experts = sparse.mlp.experts
+ for name in ("gate_up_proj", "down_proj"):
+ assert _nvfp4(getattr(experts, f"{name}_input_quantizer"))
+ assert all(_nvfp4(q) for q in getattr(experts, f"{name}_weight_quantizers"))
# Dense MLP (layers 0-2) -> NVFP4.
for proj in (dense.mlp.gate_proj, dense.mlp.up_proj, dense.mlp.down_proj):
assert _nvfp4(proj.weight_quantizer)
assert _nvfp4(proj.input_quantizer)
- # Vision tower stays BF16 -- the load-bearing case: the vision MLP reuses
- # gate_proj/up_proj/down_proj, so the dense-MLP patterns match it and only the
- # trailing `*visual*` disable keeps it off.
- vblock = model.model.visual.blocks[0]
- for proj in (vblock.mlp.gate_proj, vblock.mlp.up_proj, vblock.mlp.down_proj):
- assert proj.weight_quantizer.is_enabled is False
- assert proj.input_quantizer.is_enabled is False
+ # Sparse-MLA KV cache -> FP8.
+ assert _fp8(sparse.self_attn.k_bmm_quantizer)
+ assert _fp8(sparse.self_attn.v_bmm_quantizer)
- # Shared experts and the router gate stay BF16: `*.experts.*` needs a literal
- # `.experts.` (so `shared_experts` is skipped), and `*mlp.gate_proj*` doesn't match
- # the router `mlp.gate`.
+ # The whole vision tower stays BF16 -- the load-bearing case: its blocks and merger reuse
+ # gate_proj/up_proj/down_proj, so the dense-MLP patterns match them and only the later
+ # `*visual*` disable keeps them off.
+ for name, module in model.model.visual.named_modules():
+ if name.endswith("quantizer"):
+ assert module.is_enabled is False, name
+
+ # Shared experts stay BF16: `*.experts.*` needs a literal `.experts.`, so `shared_experts`
+ # is skipped. (The router `mlp.gate` is not an nn.Linear, so it gets no quantizer at all.)
for proj in (
sparse.mlp.shared_experts.gate_proj,
sparse.mlp.shared_experts.up_proj,
sparse.mlp.shared_experts.down_proj,
):
assert proj.weight_quantizer.is_enabled is False
- assert sparse.mlp.gate.weight_quantizer.is_enabled is False
+ assert proj.input_quantizer.is_enabled is False
- # Both attention families stay BF16, including the KDA conv1d.
- assert dense.linear_attn.conv1d.weight_quantizer.is_enabled is False
- assert dense.linear_attn.in_proj_qkvz.weight_quantizer.is_enabled is False
- assert dense.linear_attn.out_proj.weight_quantizer.is_enabled is False
- for proj in (sparse.self_attn.q_proj, sparse.self_attn.kv_proj, sparse.self_attn.o_proj):
- assert proj.weight_quantizer.is_enabled is False
+ # Both attention families' projections stay BF16, including the KDA conv1d and the indexer.
+ for attn in (dense.self_attn, sparse.self_attn):
+ for name, module in attn.named_modules():
+ if name.endswith(("weight_quantizer", "input_quantizer")):
+ assert module.is_enabled is False, name
- # lm_head stays BF16.
+ # Embeddings and lm_head stay BF16.
+ assert model.model.language_model.embed_tokens.weight_quantizer.is_enabled is False
assert model.lm_head.weight_quantizer.is_enabled is False
+
+
+class _McoreMLP(nn.Module):
+ """Megatron-Bridge MLP leaf names (dense MLP, each local expert, and the shared experts)."""
+
+ def __init__(self):
+ super().__init__()
+ self.linear_fc1 = nn.Linear(_H, 2 * _H, bias=False)
+ self.linear_fc2 = nn.Linear(_H, _H, bias=False)
+
+
+class _McoreMoE(nn.Module):
+ def __init__(self):
+ super().__init__()
+ self.experts = nn.Module()
+ self.experts.local_experts = nn.ModuleList([_McoreMLP(), _McoreMLP()])
+ self.shared_experts = _McoreMLP()
+
+
+class _McoreLayer(nn.Module):
+ def __init__(self, mlp):
+ super().__init__()
+ self.inner_layer = nn.Module() # mHC wraps each block as `.inner_layer`
+ self.inner_layer.mlp = mlp
+
+
+class _McoreGLM53Flash(nn.Module):
+ """Megatron-Bridge naming: a dense and an MoE decoder layer, the MTP layer, and vision."""
+
+ def __init__(self):
+ super().__init__()
+ self.language_model = nn.Module()
+ self.language_model.decoder = nn.Module()
+ self.language_model.decoder.layers = nn.ModuleList(
+ [_McoreLayer(_McoreMLP()), _McoreLayer(_McoreMoE())]
+ )
+ self.language_model.mtp = nn.Module()
+ self.language_model.mtp.layers = nn.ModuleList([_McoreLayer(_McoreMoE())])
+ self.visual = nn.Module()
+ self.visual.blocks = nn.ModuleList([_VisionBlock()])
+
+
+def test_glm_5_3_recipe_megatron_names():
+ model = _McoreGLM53Flash()
+ config = load_recipe(_RECIPE).quantize.model_dump()
+ config["algorithm"] = None
+ mtq.quantize(model, config)
+
+ dense, sparse = (layer.inner_layer.mlp for layer in model.language_model.decoder.layers)
+ mtp = model.language_model.mtp.layers[0].inner_layer.mlp
+
+ # Dense MLP and routed experts -> NVFP4 W4A4.
+ for mlp in (dense, *sparse.experts.local_experts):
+ for proj in (mlp.linear_fc1, mlp.linear_fc2):
+ assert _nvfp4(proj.weight_quantizer)
+ assert _nvfp4(proj.input_quantizer)
+
+ # Shared experts, the whole MTP layer, and the vision tower stay BF16.
+ for mlp in (sparse.shared_experts, mtp.shared_experts, *mtp.experts.local_experts):
+ for proj in (mlp.linear_fc1, mlp.linear_fc2):
+ assert proj.weight_quantizer.is_enabled is False
+ assert proj.input_quantizer.is_enabled is False
+ vmlp = model.visual.blocks[0].mlp
+ for proj in (vmlp.gate_proj, vmlp.up_proj, vmlp.down_proj):
+ assert proj.weight_quantizer.is_enabled is False
+
+
+@pytest.mark.parametrize(
+ ("recipe", "expert_fc2_bits"),
+ [(_IQ2_RECIPE, "iq2_s"), (_IQ2_NVFP4_DOWN_RECIPE, (2, 1))],
+)
+def test_glm_5_3_iq2_recipe_megatron_names(recipe, expert_fc2_bits):
+ model = _McoreGLM53Flash()
+ config = load_recipe(recipe).quantize.model_dump()
+ config["algorithm"] = None
+ mtq.quantize(model, config)
+
+ dense, sparse = (layer.inner_layer.mlp for layer in model.language_model.decoder.layers)
+ mtp = model.language_model.mtp.layers[0].inner_layer.mlp
+
+ # Routed experts -> IQ2_S (the NVFP4-down variant: NVFP4 on linear_fc2); dense MLP and shared
+ # experts -> Q8_0; activations stay BF16.
+ for mlp, fc1_bits, fc2_bits in [
+ *((expert, "iq2_s", expert_fc2_bits) for expert in sparse.experts.local_experts),
+ (dense, "q8_0", "q8_0"),
+ (sparse.shared_experts, "q8_0", "q8_0"),
+ ]:
+ for proj, num_bits in ((mlp.linear_fc1, fc1_bits), (mlp.linear_fc2, fc2_bits)):
+ assert proj.weight_quantizer.is_enabled
+ assert proj.weight_quantizer.num_bits == num_bits
+ assert proj.input_quantizer.is_enabled is False
+
+ # The MTP layer and the vision tower stay BF16.
+ for mlp in (mtp.shared_experts, *mtp.experts.local_experts):
+ for proj in (mlp.linear_fc1, mlp.linear_fc2):
+ assert proj.weight_quantizer.is_enabled is False
+ vmlp = model.visual.blocks[0].mlp
+ for proj in (vmlp.gate_proj, vmlp.up_proj, vmlp.down_proj):
+ assert proj.weight_quantizer.is_enabled is False