diff --git a/CHANGELOG.rst b/CHANGELOG.rst
index 6fdb4141845..7d928552900 100755
--- a/CHANGELOG.rst
+++ b/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 `_ for details.
- Add ``--mlflow `` 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/-`` 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.
diff --git a/examples/megatron_bridge/export_quantized_megatron_to_hf.py b/examples/megatron_bridge/export_quantized_megatron_to_hf.py
index faf1245ecf4..712dd6d3f10 100644
--- a/examples/megatron_bridge/export_quantized_megatron_to_hf.py
+++ b/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,
diff --git a/modelopt/torch/export/plugins/mcore_common.py b/modelopt/torch/export/plugins/mcore_common.py
index d275aee4c0a..b5dd68a5ab7 100644
--- a/modelopt/torch/export/plugins/mcore_common.py
+++ b/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] = {
diff --git a/modelopt/torch/export/plugins/mcore_custom.py b/modelopt/torch/export/plugins/mcore_custom.py
index ed3e00fd962..b09d1f2931e 100644
--- a/modelopt/torch/export/plugins/mcore_custom.py
+++ b/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."""
diff --git a/modelopt/torch/export/plugins/mcore_glm.py b/modelopt/torch/export/plugins/mcore_glm.py
new file mode 100644
index 00000000000..11ee40d63ed
--- /dev/null
+++ b/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."),
+}
diff --git a/modelopt/torch/export/quant_utils.py b/modelopt/torch/export/quant_utils.py
index fb25d23c5e4..2e98f8636fd 100755
--- a/modelopt/torch/export/quant_utils.py
+++ b/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)
diff --git a/modelopt/torch/export/unified_export_megatron.py b/modelopt/torch/export/unified_export_megatron.py
index 67564853ce8..0a7b6b4c636 100644
--- a/modelopt/torch/export/unified_export_megatron.py
+++ b/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
@@ -105,6 +105,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__ = [
@@ -175,9 +182,12 @@ 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
- self._hf_text_config.head_dim = language_model.config.kv_channels
+ # MLA's kv_channels is the V head dim, not HF's head_dim (e.g. glm5_next derives it from RoPE).
+ if not getattr(language_model.config, "multi_latent_attention", False):
+ self._hf_text_config.head_dim = language_model.config.kv_channels
self._hf_text_config.num_attention_heads = language_model.config.num_attention_heads
self._hf_text_config.num_key_value_heads = language_model.config.num_query_groups
self.is_multimodal = isinstance(model, LLaVAModel)
@@ -201,7 +211,14 @@ def __init__(
del self._hf_config.quantization_config
self.all_rules = self._populate_rule_book()
self.rules = self.all_rules[self.arch]
- self.exclude_modules = []
+ self.all_mcore_mappings = all_mcore_hf_export_mapping[self.arch]
+ 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 ()
+ ]
self.layer_config_dict = {}
if not hasattr(model, "_modelopt_state"):
@@ -369,7 +386,11 @@ def save_pretrained(
self._hf_pretrained_model_name,
trust_remote_code=self.trust_remote_code,
)
- generation_config.save_pretrained(save_directory)
+ # Pass it through unvalidated: save_pretrained rejects some shipped configs
+ # (e.g. GLM-5.3-Flash sets top_p without do_sample) on newer transformers.
+ generation_config.to_json_file(
+ os.path.join(save_directory, "generation_config.json")
+ )
except OSError:
pass
# Hub-ID / None source: fetch tokenizer files via AutoTokenizer.
@@ -461,9 +482,13 @@ def save_pretrained(
json.dump(config_dict, f, indent=4)
torch.distributed.barrier()
- # save_safetensors(state_dict, save_directory)
+ # 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 = (
+ tp_rank == 0 and get_data_parallel_rank() == 0 and get_expert_model_parallel_rank() == 0
+ )
save_safetensors_by_layer_index(
- layer_state_dicts=layer_state_dicts,
+ layer_state_dicts=layer_state_dicts if writes_layers else {},
total_layers=self.model.config.num_layers,
save_directory=save_directory,
name_template="model-{:05d}-of-{:05d}",
@@ -525,7 +550,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()
@@ -537,12 +566,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:
@@ -587,12 +622,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:
@@ -626,6 +658,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)
@@ -668,6 +726,14 @@ 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 "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)
@@ -733,6 +799,15 @@ def _get_transformer_layer_state_dict(self, layer, layer_id, is_mtp=False):
layer.mlp.shared_experts.gate_weight, layer_id, is_mtp=is_mtp
)
if hasattr(layer.mlp.experts, "local_experts"):
+ # SequentialMLP rules index experts by local position, so EP>1 would collide.
+ if (
+ torch.distributed.is_initialized()
+ and get_expert_model_parallel_world_size() > 1
+ ):
+ raise NotImplementedError(
+ "Export at expert parallel size > 1 needs grouped-GEMM experts; "
+ "export SequentialMLP (--no_moe_grouped_gemm) checkpoints at EP=1."
+ )
if not self.rules.get("use_packed_local_experts", False):
for expert_id, expert in enumerate(layer.mlp.experts.local_experts):
self.rules["local_experts.linear_fc1"](
@@ -783,12 +858,20 @@ def _get_mtp_state_dict(self) -> dict[str, torch.Tensor]:
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:
@@ -796,19 +879,18 @@ def _get_mtp_state_dict(self) -> dict[str, torch.Tensor]:
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)
@@ -832,6 +914,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
@@ -885,6 +969,53 @@ 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}")
+ 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
@@ -895,6 +1026,32 @@ 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 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):
self.rules["norm"](layer.norm, layer_id, is_mtp=is_mtp)
@@ -1031,6 +1188,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,
@@ -1089,7 +1247,8 @@ def _get_weight_bias(
and module.expert_bias is not None
and module.expert_bias.numel() > 0
):
- name_to_value["expert_bias"] = module.expert_bias.to(dtype).cpu()
+ # FP32 like Megatron's buffer and HF's e_score_correction_bias: it decides expert routing.
+ name_to_value["expert_bias"] = module.expert_bias.float().cpu()
return name_to_value
@@ -1585,12 +1744,18 @@ def _grouped_mlp_slicing(
torch.save(local_expert_state, _buf)
local_bytes = _buf.getvalue()
del _buf
- gathered_bytes: list = [None] * ep_size
- torch.distributed.all_gather_object(
- gathered_bytes, local_bytes, group=get_expert_model_parallel_group()
+ # Gather to EP rank 0 only, which writes the shards: holding every expert on every EP
+ # rank multiplies host memory by EP (a full GLM-5.3-Flash export OOMs at EP4).
+ ep_group = get_expert_model_parallel_group()
+ gathered_bytes: list | None = [None] * ep_size if ep_rank == 0 else None
+ torch.distributed.gather_object(
+ local_bytes,
+ gathered_bytes,
+ dst=torch.distributed.get_global_rank(ep_group, 0),
+ group=ep_group,
)
del local_bytes
- for b in gathered_bytes:
+ for b in gathered_bytes or ():
# weights_only=False: our own torch.save output from a sibling EP rank
# in this job's collective, not user-supplied.
s_loaded = torch.load(io.BytesIO(b), map_location="cpu", weights_only=False)
@@ -1771,11 +1936,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",
@@ -1795,12 +1955,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:
diff --git a/modelopt/torch/quantization/plugins/megatron.py b/modelopt/torch/quantization/plugins/megatron.py
index ec1958649dd..abedaeb102b 100644
--- a/modelopt/torch/quantization/plugins/megatron.py
+++ b/modelopt/torch/quantization/plugins/megatron.py
@@ -30,6 +30,15 @@
import megatron.core.transformer.moe.experts as megatron_moe
import torch
from megatron.core.dist_checkpointing.utils import replace_prefix_for_sharding
+from megatron.core.extensions.transformer_engine import (
+ TEColumnParallelGroupedLinear,
+ TEColumnParallelLinear,
+ TEDotProductAttention,
+ TELayerNormColumnParallelLinear,
+ TELinear,
+ TERowParallelGroupedLinear,
+ TERowParallelLinear,
+)
from megatron.core.models.gpt import GPTModel
from megatron.core.parallel_state import get_data_parallel_group
from megatron.core.tensor_parallel.mappings import gather_from_sequence_parallel_region
@@ -61,23 +70,14 @@
from ..utils import sync_moe_expert_amax
from ..utils.layerwise_calib import LayerActivationCollector
from .custom import CUSTOM_MODEL_PLUGINS, _ParallelLinear
+from .transformer_engine import _QuantTEGroupedLinear, _QuantTELayerNormLinear, _QuantTELinear
try:
- from megatron.core.extensions.transformer_engine import (
- TEColumnParallelGroupedLinear,
- TEColumnParallelLinear,
- TEDotProductAttention,
- TELayerNormColumnParallelLinear,
- TELinear,
- TERowParallelGroupedLinear,
- TERowParallelLinear,
- )
-
- from .transformer_engine import _QuantTEGroupedLinear, _QuantTELayerNormLinear, _QuantTELinear
-
- HAS_TE = True
+ from megatron.core.transformer.experimental_attention_variant.dsa import DSAttention
+
+ HAS_DSA = True
except ImportError:
- HAS_TE = False
+ HAS_DSA = False
__all__ = []
@@ -827,246 +827,268 @@ def sharded_state_dict(self, prefix="", sharded_offsets=(), metadata=None):
return sharded_state_dict
-if HAS_TE:
+@QuantModuleRegistry.register({TERowParallelLinear: "te_mcore_RowParallelLinear"})
+class _QuantTEMCoreRowParallelLinear(_QuantTELinear, _MegatronRowParallelLinear):
+ pass
- @QuantModuleRegistry.register({TERowParallelLinear: "te_mcore_RowParallelLinear"})
- class _QuantTEMCoreRowParallelLinear(_QuantTELinear, _MegatronRowParallelLinear):
- pass
- @QuantModuleRegistry.register({TEColumnParallelLinear: "te_mcore_ColumnParallelLinear"})
- class _QuantTEMCoreColumnParallelLinear(_QuantTELinear, _MegatronColumnParallelLinear):
- pass
+@QuantModuleRegistry.register({TEColumnParallelLinear: "te_mcore_ColumnParallelLinear"})
+class _QuantTEMCoreColumnParallelLinear(_QuantTELinear, _MegatronColumnParallelLinear):
+ pass
- @QuantModuleRegistry.register({TELinear: "te_mcore_Linear"})
- class _QuantTEMCoreLinear(_QuantTELinear):
- pass
- @QuantModuleRegistry.register(
- {TELayerNormColumnParallelLinear: "te_mcore_LayerNormColumnParallelLinear"}
- )
- class _QuantTELayerNormColumnParallelLinear(
- _QuantTELayerNormLinear, _MegatronColumnParallelLinear
- ):
- pass
-
- # Quantized subclasses to support TEGroupedLinear quantization
- class _QuantMegatronTEGroupedLinear(_QuantTEGroupedLinear, _MegatronParallelLinear):
- def modelopt_post_load_extra_state(self):
- _initialize_grouped_weight_quantizer_state(self)
-
- def _load_from_state_dict(self, state_dict, prefix, *args, **kwargs):
- # _sharded_state_dict_grouped adds _extra_state{gemm_idx} for gemm_idx:[1, num_gemms] in
- # sharded_state_dict which is same as _extra_state. The _extra_state{gemm_idx} is used for
- # TE Fp8 checkpoint, we need to remove the _extra_state{gemm_idx} for gemm_idx:[1, num_gemms]
- # for modelopt checkpoint restore
- filtered_state_dict = {
- k: v
- for k, v in state_dict.items()
- if not any(k.endswith(f"_extra_state{num}") for num in range(1, self.num_gemms))
- }
- return super()._load_from_state_dict(filtered_state_dict, prefix, *args, **kwargs)
+@QuantModuleRegistry.register({TELinear: "te_mcore_Linear"})
+class _QuantTEMCoreLinear(_QuantTELinear):
+ pass
+
+
+@QuantModuleRegistry.register(
+ {TELayerNormColumnParallelLinear: "te_mcore_LayerNormColumnParallelLinear"}
+)
+class _QuantTELayerNormColumnParallelLinear(_QuantTELayerNormLinear, _MegatronColumnParallelLinear):
+ pass
+
+
+# Quantized subclasses to support TEGroupedLinear quantization
+class _QuantMegatronTEGroupedLinear(_QuantTEGroupedLinear, _MegatronParallelLinear):
+ def modelopt_post_load_extra_state(self):
+ _initialize_grouped_weight_quantizer_state(self)
+
+ def _load_from_state_dict(self, state_dict, prefix, *args, **kwargs):
+ # _sharded_state_dict_grouped adds _extra_state{gemm_idx} for gemm_idx:[1, num_gemms] in
+ # sharded_state_dict which is same as _extra_state. The _extra_state{gemm_idx} is used for
+ # TE Fp8 checkpoint, we need to remove the _extra_state{gemm_idx} for gemm_idx:[1, num_gemms]
+ # for modelopt checkpoint restore
+ filtered_state_dict = {
+ k: v
+ for k, v in state_dict.items()
+ if not any(k.endswith(f"_extra_state{num}") for num in range(1, self.num_gemms))
+ }
+ return super()._load_from_state_dict(filtered_state_dict, prefix, *args, **kwargs)
+
+ def _process_quantizer_amax(self, k, v, quantizer_state_dict):
+ # Per-expert quantizers have independent checkpoint keys. Preserve their native
+ # scalar, channel, or block shape instead of flattening them through the legacy
+ # single-quantizer path.
+ if re.match(r"weight_quantizer\.\d+\..+_amax$", k):
+ quantizer_state_dict[k] = v
+ else:
+ quantizer_state_dict[k] = v.view(-1) if v.numel() == 1 else v
+
+ def _expert_parallel_groups(self):
+ """Return the (ep, expt_dp) process groups used to place fused experts globally."""
+ pg_collection = getattr(self, "_pg_collection", None)
+ if pg_collection is not None:
+ return pg_collection.ep, pg_collection.expt_dp
+ return (
+ mcore_parallel.get_expert_model_parallel_group(),
+ mcore_parallel.get_expert_data_parallel_group(),
+ )
+
+ def sharded_state_dict(self, prefix="", sharded_offsets=(), metadata=None):
+ """Emit per-expert quantizer amax with the same global expert identity as the weights.
+
+ The base linear emits ``weight_quantizer.{local_i}._amax`` with the local index and no
+ expert offset, so every EP rank writes identical keys and ``torch_dist`` dedup keeps only
+ one rank's experts. Here we mirror Megatron ``TEGroupedLinear._sharded_state_dict_grouped``:
+ each fused expert comes with its ``global_expert_idx`` (baked into the key prefix under
+ ``singleton_local_shards``, otherwise an EP sharded-offset) so all ``num_global_experts``
+ persist and reshard to any EP. Shared, whole-linear quantizer buffers (e.g.
+ ``input_quantizer``) keep the plain replicated path.
+ """
+ metadata = ensure_metadata_has_dp_cp_group(metadata)
+ singleton_local_shards = bool((metadata or {}).get("singleton_local_shards", False))
+
+ # Weights/bias/_extra_state come from the wrapped TE grouped linear, which already
+ # assigns each expert its global identity. Skip _MegatronParallelLinear's local-index
+ # amax emission by starting from the base MCore module's sharded_state_dict.
+ sharded_state_dict = super(_MegatronParallelLinear, self).sharded_state_dict(
+ prefix, sharded_offsets, metadata
+ )
- def _process_quantizer_amax(self, k, v, quantizer_state_dict):
- # Per-expert quantizers have independent checkpoint keys. Preserve their native
- # scalar, channel, or block shape instead of flattening them through the legacy
- # single-quantizer path.
- if re.match(r"weight_quantizer\.\d+\..+_amax$", k):
+ # Collect the quantizer buffers exactly like _MegatronParallelLinear.sharded_state_dict.
+ quantizer_state_dict = {}
+ for k, v in self.state_dict(prefix="", keep_vars=True).items():
+ if "_quantizer" in k and "_amax" in k:
+ self._process_quantizer_amax(k, v, quantizer_state_dict)
+ elif k == "input_quantizer._pre_quant_scale":
+ self._process_activation_quantizer_pre_quant_scale(k, v, quantizer_state_dict)
+ elif self._parameter_to_keep_in_quantizer_state_dict(k):
quantizer_state_dict[k] = v
+ elif "quantizer" in k:
+ warn_rank_0(
+ f"Quantizer state {k} is not supported for sharded_state_dict. "
+ "Please use regular state_dict."
+ )
+
+ # Channel shard axes (per real key); _global_amax stays un-sharded along channels but
+ # still rides with the expert identity below.
+ shard_axis_dict = self._get_shard_axis_dict(quantizer_state_dict)
+
+ # Split per-expert weight_quantizer.{i}.* from shared (input/output) quantizer buffers.
+ expert_re = re.compile(r"^weight_quantizer\.(\d+)\.(.+)$")
+ per_expert_subs = [[] for _ in range(self.num_gemms)]
+ shared_state = {}
+ for k, v in quantizer_state_dict.items():
+ m = expert_re.match(k)
+ if m:
+ per_expert_subs[int(m.group(1))].append((m.group(2), v, shard_axis_dict.get(k)))
else:
- quantizer_state_dict[k] = v.view(-1) if v.numel() == 1 else v
-
- def _expert_parallel_groups(self):
- """Return the (ep, expt_dp) process groups used to place fused experts globally."""
- pg_collection = getattr(self, "_pg_collection", None)
- if pg_collection is not None:
- return pg_collection.ep, pg_collection.expt_dp
- return (
- mcore_parallel.get_expert_model_parallel_group(),
- mcore_parallel.get_expert_data_parallel_group(),
- )
+ shared_state[k] = v
- def sharded_state_dict(self, prefix="", sharded_offsets=(), metadata=None):
- """Emit per-expert quantizer amax with the same global expert identity as the weights.
-
- The base linear emits ``weight_quantizer.{local_i}._amax`` with the local index and no
- expert offset, so every EP rank writes identical keys and ``torch_dist`` dedup keeps only
- one rank's experts. Here we mirror Megatron ``TEGroupedLinear._sharded_state_dict_grouped``:
- each fused expert comes with its ``global_expert_idx`` (baked into the key prefix under
- ``singleton_local_shards``, otherwise an EP sharded-offset) so all ``num_global_experts``
- persist and reshard to any EP. Shared, whole-linear quantizer buffers (e.g.
- ``input_quantizer``) keep the plain replicated path.
- """
- metadata = ensure_metadata_has_dp_cp_group(metadata)
- singleton_local_shards = bool((metadata or {}).get("singleton_local_shards", False))
-
- # Weights/bias/_extra_state come from the wrapped TE grouped linear, which already
- # assigns each expert its global identity. Skip _MegatronParallelLinear's local-index
- # amax emission by starting from the base MCore module's sharded_state_dict.
- sharded_state_dict = super(_MegatronParallelLinear, self).sharded_state_dict(
- prefix, sharded_offsets, metadata
+ # Shared quantizer buffers: replicated across experts, plain base offsets.
+ shared_axis_dict = {k: shard_axis_dict[k] for k in shared_state if k in shard_axis_dict}
+ sharded_state_dict.update(
+ make_sharded_tensors_for_checkpoint(
+ shared_state, prefix, shared_axis_dict, sharded_offsets
)
+ )
- # Collect the quantizer buffers exactly like _MegatronParallelLinear.sharded_state_dict.
- quantizer_state_dict = {}
- for k, v in self.state_dict(prefix="", keep_vars=True).items():
- if "_quantizer" in k and "_amax" in k:
- self._process_quantizer_amax(k, v, quantizer_state_dict)
- elif k == "input_quantizer._pre_quant_scale":
- self._process_activation_quantizer_pre_quant_scale(k, v, quantizer_state_dict)
- elif self._parameter_to_keep_in_quantizer_state_dict(k):
- quantizer_state_dict[k] = v
- elif "quantizer" in k:
- warn_rank_0(
- f"Quantizer state {k} is not supported for sharded_state_dict. "
- "Please use regular state_dict."
- )
-
- # Channel shard axes (per real key); _global_amax stays un-sharded along channels but
- # still rides with the expert identity below.
- shard_axis_dict = self._get_shard_axis_dict(quantizer_state_dict)
-
- # Split per-expert weight_quantizer.{i}.* from shared (input/output) quantizer buffers.
- expert_re = re.compile(r"^weight_quantizer\.(\d+)\.(.+)$")
- per_expert_subs = [[] for _ in range(self.num_gemms)]
- shared_state = {}
- for k, v in quantizer_state_dict.items():
- m = expert_re.match(k)
- if m:
- per_expert_subs[int(m.group(1))].append((m.group(2), v, shard_axis_dict.get(k)))
- else:
- shared_state[k] = v
-
- # Shared quantizer buffers: replicated across experts, plain base offsets.
- shared_axis_dict = {k: shard_axis_dict[k] for k in shared_state if k in shard_axis_dict}
- sharded_state_dict.update(
- make_sharded_tensors_for_checkpoint(
- shared_state, prefix, shared_axis_dict, sharded_offsets
+ # Per-expert amax: assign the same global expert identity the weights use.
+ ep_group, expt_dp_group = self._expert_parallel_groups()
+ num_global_experts = get_pg_size(ep_group) * self.num_gemms
+ local_expert_indices_offset = get_pg_rank(ep_group) * self.num_gemms
+ edp_replica_id = get_pg_rank(expt_dp_group)
+ ep_axis = len(sharded_offsets)
+ for gemm_idx, subs in enumerate(per_expert_subs):
+ if not subs:
+ continue
+ global_expert_idx = local_expert_indices_offset + gemm_idx
+ if singleton_local_shards:
+ expert_prefix = f"{global_expert_idx}.{prefix}"
+ new_sharded_offsets = sharded_offsets
+ else:
+ expert_prefix = prefix
+ new_sharded_offsets = (
+ *sharded_offsets,
+ (ep_axis, global_expert_idx, num_global_experts),
)
+ expert_state = {f"{gemm_idx}.weight_quantizer.{sub}": v for sub, v, _ in subs}
+ expert_axis = {
+ f"{gemm_idx}.weight_quantizer.{sub}": axis
+ for sub, _, axis in subs
+ if axis is not None
+ }
+ sub_sd = make_sharded_tensors_for_checkpoint(
+ expert_state, "", expert_axis, new_sharded_offsets
)
+ # Rewrite each ShardedTensor.key to carry the global expert identity (dict keys,
+ # which map to the local buffers on restore, are left untouched).
+ replace_prefix_for_sharding(sub_sd, f"{gemm_idx}.", expert_prefix)
+ for sub, _, _ in subs:
+ sh_ten = sub_sd[f"{gemm_idx}.weight_quantizer.{sub}"]
+ replica_id = sh_ten.replica_id
+ if len(replica_id) == 3:
+ sh_ten.replica_id = (*replica_id[:2], edp_replica_id)
+ sharded_state_dict[f"{prefix}weight_quantizer.{gemm_idx}.{sub}"] = sh_ten
+ return sharded_state_dict
- # Per-expert amax: assign the same global expert identity the weights use.
- ep_group, expt_dp_group = self._expert_parallel_groups()
- num_global_experts = get_pg_size(ep_group) * self.num_gemms
- local_expert_indices_offset = get_pg_rank(ep_group) * self.num_gemms
- edp_replica_id = get_pg_rank(expt_dp_group)
- ep_axis = len(sharded_offsets)
- for gemm_idx, subs in enumerate(per_expert_subs):
- if not subs:
- continue
- global_expert_idx = local_expert_indices_offset + gemm_idx
- if singleton_local_shards:
- expert_prefix = f"{global_expert_idx}.{prefix}"
- new_sharded_offsets = sharded_offsets
- else:
- expert_prefix = prefix
- new_sharded_offsets = (
- *sharded_offsets,
- (ep_axis, global_expert_idx, num_global_experts),
- )
- expert_state = {f"{gemm_idx}.weight_quantizer.{sub}": v for sub, v, _ in subs}
- expert_axis = {
- f"{gemm_idx}.weight_quantizer.{sub}": axis
- for sub, _, axis in subs
- if axis is not None
- }
- sub_sd = make_sharded_tensors_for_checkpoint(
- expert_state, "", expert_axis, new_sharded_offsets
- )
- # Rewrite each ShardedTensor.key to carry the global expert identity (dict keys,
- # which map to the local buffers on restore, are left untouched).
- replace_prefix_for_sharding(sub_sd, f"{gemm_idx}.", expert_prefix)
- for sub, _, _ in subs:
- sh_ten = sub_sd[f"{gemm_idx}.weight_quantizer.{sub}"]
- replica_id = sh_ten.replica_id
- if len(replica_id) == 3:
- sh_ten.replica_id = (*replica_id[:2], edp_replica_id)
- sharded_state_dict[f"{prefix}weight_quantizer.{gemm_idx}.{sub}"] = sh_ten
- return sharded_state_dict
-
- @QuantModuleRegistry.register(
- {TEColumnParallelGroupedLinear: "megatron_TEColumnParallelGroupedLinear"}
- )
- class _MegatronTEGroupedColumnParallelLinear(
- _QuantMegatronTEGroupedLinear, _MegatronColumnParallelLinear
- ):
- pass
- @QuantModuleRegistry.register(
- {TERowParallelGroupedLinear: "megatron_TERowParallelGroupedLinear"}
- )
- class _MegatronTEGroupedRowParallelLinear(
- _QuantMegatronTEGroupedLinear, _MegatronRowParallelLinear
- ):
- pass
-
- @QuantModuleRegistry.register({megatron_moe.TEGroupedMLP: "megatron_moe_TEGroupedMLP"})
- class _MegatronTEGroupedMLP(_MegatronMLP):
- def _setup(self):
- if not hasattr(self, "parallel_state") or self.parallel_state is None:
- self.parallel_state = ParallelState(
- mcore_parallel.get_expert_data_parallel_group(),
- tensor_parallel_group=mcore_parallel.get_expert_tensor_parallel_group(),
- expert_model_parallel_group=mcore_parallel.get_expert_model_parallel_group(),
- )
- # These child linears are still native MCore modules here. Seed `_parallel_state`
- # directly so the later QuantModule conversion sees the intended parallel state.
- self.linear_fc1._parallel_state = self.parallel_state
- self.linear_fc2._parallel_state = self.parallel_state
-
- @QuantModuleRegistry.register({TEDotProductAttention: "TEDotProductAttention"})
- class _QuantTEDotProductAttention(QuantModule):
- """Quantized version of TEDotProductAttention for Megatron models with KV cache quantization.
-
- This class adds KV cache quantization support to Transformer Engine's TEDotProductAttention
- module used in Megatron-Core models. It introduces three quantizers (q_bmm_quantizer,
- k_bmm_quantizer, v_bmm_quantizer) that quantize the query, key, and value tensors after
- RoPE has been applied.
- """
+@QuantModuleRegistry.register(
+ {TEColumnParallelGroupedLinear: "megatron_TEColumnParallelGroupedLinear"}
+)
+class _MegatronTEGroupedColumnParallelLinear(
+ _QuantMegatronTEGroupedLinear, _MegatronColumnParallelLinear
+):
+ pass
- def _setup(self):
- """Initialize quantizers for Q, K, V tensors."""
- self.q_bmm_quantizer = TensorQuantizer()
- self.k_bmm_quantizer = TensorQuantizer()
- self.v_bmm_quantizer = TensorQuantizer()
- # Set parallel_state for distributed sync of BMM quantizers
- try:
- data_parallel_group = get_data_parallel_group(with_context_parallel=True)
- except AssertionError:
- data_parallel_group = get_data_parallel_group()
+@QuantModuleRegistry.register({TERowParallelGroupedLinear: "megatron_TERowParallelGroupedLinear"})
+class _MegatronTEGroupedRowParallelLinear(
+ _QuantMegatronTEGroupedLinear, _MegatronRowParallelLinear
+):
+ pass
+
+
+@QuantModuleRegistry.register({megatron_moe.TEGroupedMLP: "megatron_moe_TEGroupedMLP"})
+class _MegatronTEGroupedMLP(_MegatronMLP):
+ def _setup(self):
+ if not hasattr(self, "parallel_state") or self.parallel_state is None:
self.parallel_state = ParallelState(
- data_parallel_group,
- mcore_parallel.get_tensor_model_parallel_group(),
+ mcore_parallel.get_expert_data_parallel_group(),
+ tensor_parallel_group=mcore_parallel.get_expert_tensor_parallel_group(),
+ expert_model_parallel_group=mcore_parallel.get_expert_model_parallel_group(),
)
+ # These child linears are still native MCore modules here. Seed `_parallel_state`
+ # directly so the later QuantModule conversion sees the intended parallel state.
+ self.linear_fc1._parallel_state = self.parallel_state
+ self.linear_fc2._parallel_state = self.parallel_state
- def forward(self, query, key, value, *args, **kwargs):
- """Apply post-RoPE quantization to KV cache."""
- # Quantize Q, K, V
- query = self.q_bmm_quantizer(query)
- key = self.k_bmm_quantizer(key)
+
+@QuantModuleRegistry.register({TEDotProductAttention: "TEDotProductAttention"})
+class _QuantCoreAttention(QuantModule):
+ """Core attention (TEDotProductAttention, DSAttention) with KV cache quantization.
+
+ Adds q/k/v_bmm_quantizers that quantize the post-RoPE query, key and value tensors.
+ """
+
+ def _setup(self):
+ """Initialize quantizers for Q, K, V tensors."""
+ self.q_bmm_quantizer = TensorQuantizer()
+ self.k_bmm_quantizer = TensorQuantizer()
+ self.v_bmm_quantizer = TensorQuantizer()
+
+ # Set parallel_state for distributed sync of BMM quantizers
+ try:
+ data_parallel_group = get_data_parallel_group(with_context_parallel=True)
+ except AssertionError:
+ data_parallel_group = get_data_parallel_group()
+ self.parallel_state = ParallelState(
+ data_parallel_group,
+ mcore_parallel.get_tensor_model_parallel_group(),
+ )
+
+ def forward(self, query, key, value, *args, **kwargs):
+ """Apply post-RoPE quantization to KV cache."""
+ # Quantize Q, K, V
+ query = self.q_bmm_quantizer(query)
+ if value is None:
+ # Absorbed MLA (DSAttention) passes value=None: the key is the KV latent that both K
+ # and V are read from, so calibrate V on it too (output unused) to export a V scale.
+ self.v_bmm_quantizer(key)
+ else:
value = self.v_bmm_quantizer(value)
- return super().forward(query, key, value, *args, **kwargs)
-
- def modelopt_post_restore(self, name=""):
- """Restore quantizer states after model loading."""
- for tq in [self.q_bmm_quantizer, self.k_bmm_quantizer, self.v_bmm_quantizer]:
- # TODO: Add support for non-scalar states such as
- # Affine KVCache bias vector which is per head per channel
- if not all(v.numel() == 1 for v in tq.state_dict().values()):
- raise NotImplementedError(
- "Only scalar states are supported for KV Cache/BMM Quantizers"
- )
- # dtype and device should have been set in `megatron_replace_quant_module_hook`
- # via `_configure_attention_for_kv_cache_quant`
- assert hasattr(self, "device") and hasattr(self, "dtype")
- self.to(device=self.device, dtype=self.dtype)
-
- def sharded_state_dict(self, prefix="", sharded_offsets=(), metadata=None):
- # Currently we do not need sharded_state_dict for TEDotProductAttention since the amax are scalar values.
- # However we would need this in future to support non-scalar states such as
- # Affine KVCache Quant bias vector.
- state_dict = self.state_dict(prefix="", keep_vars=True)
- return make_sharded_tensors_for_checkpoint(state_dict, prefix, {}, sharded_offsets)
+ key = self.k_bmm_quantizer(key)
+ return super().forward(query, key, value, *args, **kwargs)
+
+ def modelopt_post_restore(self, name=""):
+ """Restore quantizer states after model loading."""
+ for tq in [self.q_bmm_quantizer, self.k_bmm_quantizer, self.v_bmm_quantizer]:
+ # TODO: Add support for non-scalar states such as
+ # Affine KVCache bias vector which is per head per channel
+ if not all(v.numel() == 1 for v in tq.state_dict().values()):
+ raise NotImplementedError(
+ "Only scalar states are supported for KV Cache/BMM Quantizers"
+ )
+ # dtype and device should have been set in `megatron_replace_quant_module_hook`
+ # via `_configure_attention_for_kv_cache_quant`
+ assert hasattr(self, "device") and hasattr(self, "dtype")
+ self.to(device=self.device, dtype=self.dtype)
+
+ def sharded_state_dict(self, prefix="", sharded_offsets=(), metadata=None):
+ # Currently we do not need sharded_state_dict for core attention since the amax are scalar values.
+ # However we would need this in future to support non-scalar states such as
+ # Affine KVCache Quant bias vector.
+ state_dict = self.state_dict(prefix="", keep_vars=True)
+ return make_sharded_tensors_for_checkpoint(state_dict, prefix, {}, sharded_offsets)
+
+
+if HAS_DSA:
+
+ @QuantModuleRegistry.register({DSAttention: "megatron_DSAttention"})
+ class _QuantDSAttention(_QuantCoreAttention):
+ """DSAttention with KV-cache quantization.
+
+ torch's state_dict() / load_state_dict() route ``_extra_state`` only through classes that
+ override these; TEDotProductAttention does, DSAttention does not, so without them the
+ quantizer state (including amax) was dropped from checkpoints.
+ """
+
+ def get_extra_state(self):
+ return quant_module_get_extra_state(self)
+
+ def set_extra_state(self, state):
+ quant_module_set_extra_state(self, state)
def _is_supported_megatron_model(model: torch.nn.Module) -> bool:
diff --git a/modelopt/torch/quantization/utils/core_utils.py b/modelopt/torch/quantization/utils/core_utils.py
index eca3302a425..885d6f9546a 100644
--- a/modelopt/torch/quantization/utils/core_utils.py
+++ b/modelopt/torch/quantization/utils/core_utils.py
@@ -247,18 +247,24 @@ def representative_weight_quantizer(module: nn.Module, weight_name: str = "weigh
def weight_attr_names(module: nn.Module) -> "Generator[str, None, None]":
"""Get the weight param attribute names in a converted module, non-recursive.
- Covers three layouts:
+ Covers four layouts:
- standard ``nn.Linear``: ``weight`` + ``weight_quantizer``.
- custom per-weight quantizer (e.g. ``Llama4TextExperts`` with ``gate_up_proj`` +
``gate_up_proj_weight_quantizer``).
- fused-experts ``nn.ModuleList`` quantizers (``_QuantFusedExperts`` with
```` + ``_weight_quantizers`` plural list).
+ - TEGroupedLinear: ``weight0..N`` behind one ``GroupedQuantizer``, reported as ``weight``.
"""
+ from ..nn import GroupedQuantizer # local: ..nn imports this utils package at module scope
+
# standard: "weight" + "weight_quantizer" (singular) or "weight_quantizers" (plural)
if getattr(module, "weight", None) is not None:
if representative_weight_quantizer(module, "weight") is not None:
yield "weight"
+ elif isinstance(getattr(module, "weight_quantizer", None), GroupedQuantizer):
+ # TEGroupedLinear: ``weight0..N`` share one GroupedQuantizer and there is no ``weight``.
+ yield "weight"
# per-parameter custom attr names
for name, _ in module.named_parameters(recurse=False):
diff --git a/modelopt_recipes/models/zai-org/GLM-5.3-Flash/ptq/nvfp4_experts_dense_mlp-kv_fp8_cast.yaml b/modelopt_recipes/models/zai-org/GLM-5.3-Flash/ptq/nvfp4_experts_dense_mlp-kv_fp8_cast.yaml
index 2954aa12abe..6a7ce4d7f14 100644
--- a/modelopt_recipes/models/zai-org/GLM-5.3-Flash/ptq/nvfp4_experts_dense_mlp-kv_fp8_cast.yaml
+++ b/modelopt_recipes/models/zai-org/GLM-5.3-Flash/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: 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/examples/megatron_bridge/test_quantize_export.py b/tests/examples/megatron_bridge/test_quantize_export.py
index 3c877b1a145..a43ee99834f 100644
--- a/tests/examples/megatron_bridge/test_quantize_export.py
+++ b/tests/examples/megatron_bridge/test_quantize_export.py
@@ -41,29 +41,33 @@
@pytest.mark.parametrize(
- ("create_model", "model_kwargs"),
+ ("create_model", "model_kwargs", "recipe", "export_parallelism"),
[
# MoE: routed experts used to be dropped silently from the export.
- (create_tiny_qwen3_moe_dir, _DENSE_KWARGS),
- # Dense VLM: only the language model is quantized, vision is copied through.
- (create_tiny_qwen3vl_dir, {}),
- # Mamba hybrid + MoE, and the one architecture that keeps grouped-GEMM experts.
- (create_tiny_nemotron_h_dir, {}),
+ (create_tiny_qwen3_moe_dir, _DENSE_KWARGS, "general/ptq/nvfp4_default-kv_fp8", "pp_size"),
+ # Dense VLM: only the language model is quantized, vision is copied through. Keeps the
+ # FP8 script-to-checkpoint coverage.
+ (create_tiny_qwen3vl_dir, {}, "general/ptq/fp8_default-kv_fp8", "pp_size"),
+ # Mamba hybrid + MoE with grouped-GEMM experts: exported with its experts sharded
+ # across ranks (EP), resharded from the TP-quantized checkpoint.
+ (create_tiny_nemotron_h_dir, {}, "general/ptq/nvfp4_default-kv_fp8", "ep_size"),
],
ids=["qwen3_moe", "qwen3vl", "nemotron_h"],
)
@pytest.mark.timeout(360) # quantize + export in one test; 1-gpu CI exceeds the default 300s
-def test_quantize_and_export(tmp_path: Path, num_gpus, create_model, model_kwargs):
+def test_quantize_and_export(
+ tmp_path: Path, num_gpus, create_model, model_kwargs, recipe, export_parallelism
+):
"""Quantize a tiny model via a YAML recipe and export it to a unified HF checkpoint."""
hf_model_path = create_model(tmp_path, with_tokenizer=True, **model_kwargs)
- megatron_path = tmp_path / "fp8_megatron"
- hf_export_path = tmp_path / "fp8_hf"
+ megatron_path = tmp_path / "quantized_megatron"
+ hf_export_path = tmp_path / "quantized_hf"
# Step 1: quantize and save a Megatron checkpoint
quantize_cmd = extend_cmd_parts(
["torchrun", f"--nproc_per_node={num_gpus}", "quantize.py", "--skip_generate"],
hf_model_name_or_path=hf_model_path,
- recipe="general/ptq/fp8_default-kv_fp8",
+ recipe=recipe,
tp_size=num_gpus,
calib_dataset_name="cnn_dailymail",
calib_num_samples=4,
@@ -81,7 +85,7 @@ def test_quantize_and_export(tmp_path: Path, num_gpus, create_model, model_kwarg
hf_model_name_or_path=hf_model_path,
megatron_path=megatron_path,
export_unified_hf_path=hf_export_path,
- pp_size=num_gpus,
+ **{export_parallelism: num_gpus},
)
run_example_command(export_cmd, example_path="megatron_bridge", setup_free_port=True)
assert (hf_export_path / "config.json").exists()
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 9debbc2231d..484d7732978 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
@@ -34,6 +34,8 @@
"Qwen3VLForConditionalGeneration": "model.language_model.layers.{}.self_attn.",
"Qwen3_5ForConditionalGeneration": "model.language_model.layers.{}.self_attn.",
"Qwen3_5MoeForConditionalGeneration": "model.language_model.layers.{}.self_attn.",
+ "Glm5NextForConditionalGeneration": "model.language_model.layers.{}.self_attn.",
+ "GlmMoeDsaForCausalLM": "model.layers.{}.self_attn.",
}
@@ -54,6 +56,8 @@ def test_export_mapping_emits_kv_cache_scales(arch, prefix):
"Qwen3_5MoeForConditionalGeneration": "model.language_model.layers.{}.mlp.experts.{}",
"Qwen3MoeForCausalLM": "model.layers.{}.mlp.experts.{}",
"NemotronHForCausalLM": "backbone.layers.{}.mixer.experts.{}",
+ "Glm5NextForConditionalGeneration": "model.language_model.layers.{}.mlp.experts.{}",
+ "GlmMoeDsaForCausalLM": "model.layers.{}.mlp.experts.{}",
}
@@ -82,3 +86,33 @@ def test_qwen3_5_moe_expert_names_match_released_checkpoint():
fc2 = mapping["experts.linear_fc2"].target_name_or_prefix.format(7).format(3) + "."
assert fc2 == "model.language_model.layers.7.mlp.experts.3.down_proj."
+
+
+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 d3f23c7f4f8..033aa80dc0c 100644
--- a/tests/gpu_megatron/torch/export/test_unified_export_megatron.py
+++ b/tests/gpu_megatron/torch/export/test_unified_export_megatron.py
@@ -46,6 +46,7 @@
import modelopt.torch.quantization.ggml as ggml
import modelopt.torch.speculative as mtsp
from modelopt.torch.export import KV_CACHE_FP8, export_mcore_gpt_to_hf, import_mcore_gpt_from_hf
+from modelopt.torch.export.plugins.mcore_common import all_mcore_hf_export_mapping
from modelopt.torch.export.quant_format import IQ_FORMATS
from modelopt.torch.export.unified_export_megatron import GPTModelExporter
from modelopt.torch.quantization.config import QuantizerAttributeConfig
@@ -287,6 +288,59 @@ def test_megatron_gated_delta_net_slicing_exports_iq_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_iq_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_iq_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", IQ_FORMAT_NAMES)
def test_megatron_packed_experts_reject_iq_without_deployment_loader(qformat):
experts = _make_iq_experts(qformat, "linear_fc2")
@@ -902,6 +956,7 @@ def _make_exporter_for_mtp(model_dir: Path) -> GPTModelExporter:
exporter._hf_pretrained_model_name = str(model_dir)
exporter._state_dict = {} # MTP keys are absent — they should be picked up
exporter.exclude_modules = []
+ exporter.rules = {} # a Qwen-style ``mtp.*`` checkpoint, not MTP stored as decoder layers
return exporter
@@ -975,6 +1030,80 @@ def test_mtp_state_dict_index_file(tmp_path):
assert "mtp*" in exporter.exclude_modules
+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"
+ model_dir.mkdir()
+ fp8 = torch.full((128, 128), 2.0).to(torch.float8_e4m3fn)
+ save_file(
+ {
+ "model.layers.1.input_layernorm.weight": torch.ones(8), # pruned decoder layer
+ "model.layers.2.enorm.weight": torch.full((8,), 3.0), # MTP layer of the source
+ "model.layers.2.eh_proj.weight": fp8,
+ "model.layers.2.eh_proj.weight_scale_inv": torch.full((1, 1), 0.5),
+ },
+ str(model_dir / "model.safetensors"),
+ )
+ exporter = _make_exporter_for_mtp(model_dir)
+ exporter.rules = {"mtp_in_decoder_layers": True}
+ exporter.all_mcore_mappings = all_mcore_hf_export_mapping["GlmMoeDsaForCausalLM"]
+ # Depth-pruned from 2 to 1 decoder layers: the source MTP (layer 2) lands at layer 1.
+ exporter._src_num_hidden_layers = 2
+ exporter._hf_text_config = SimpleNamespace(num_hidden_layers=1, num_nextn_predict_layers=1)
+
+ mtp_state_dict = exporter._get_mtp_state_dict()
+
+ assert sorted(mtp_state_dict) == [
+ "model.layers.1.eh_proj.weight",
+ "model.layers.1.enorm.weight",
+ ]
+ assert mtp_state_dict["model.layers.1.eh_proj.weight"].dtype == torch.bfloat16
+ torch.testing.assert_close(
+ mtp_state_dict["model.layers.1.eh_proj.weight"].float(), torch.ones(128, 128)
+ )
+ assert exporter.exclude_modules == ["model.layers.1.*"]
+
+
+def test_mtp_state_dict_copies_decoder_mtp_layers_from_hub(tmp_path, monkeypatch):
+ """A Hub-ID source downloads only the shards holding the MTP layer, then copies it."""
+ hub = tmp_path / "hub"
+ hub.mkdir()
+ shards = {
+ "model.layers.0.input_layernorm.weight": "model-00001-of-00002.safetensors",
+ "model.layers.1.enorm.weight": "model-00002-of-00002.safetensors",
+ }
+ save_file(
+ {"model.layers.0.input_layernorm.weight": torch.ones(8)},
+ str(hub / shards["model.layers.0.input_layernorm.weight"]),
+ )
+ save_file(
+ {"model.layers.1.enorm.weight": torch.full((8,), 3.0)},
+ str(hub / shards["model.layers.1.enorm.weight"]),
+ )
+ (hub / "model.safetensors.index.json").write_text(json.dumps({"weight_map": shards}))
+ requested = {}
+
+ def fake_snapshot_download(repo_id, allow_patterns):
+ requested[repo_id] = allow_patterns
+ return str(hub)
+
+ monkeypatch.setattr(uem, "hf_hub_download", lambda repo_id, filename: str(hub / filename))
+ monkeypatch.setattr(uem, "snapshot_download", fake_snapshot_download)
+ exporter = _make_exporter_for_mtp(Path("zai-org/GLM-5.2"))
+ exporter.rules = {"mtp_in_decoder_layers": True}
+ exporter.all_mcore_mappings = all_mcore_hf_export_mapping["GlmMoeDsaForCausalLM"]
+ exporter._src_num_hidden_layers = 1
+ exporter._hf_text_config = SimpleNamespace(num_hidden_layers=1, num_nextn_predict_layers=1)
+
+ mtp_state_dict = exporter._get_mtp_state_dict()
+
+ assert requested == {
+ "zai-org/GLM-5.2": ["model.safetensors.index.json", "model-00002-of-00002.safetensors"]
+ }
+ assert list(mtp_state_dict) == ["model.layers.1.enorm.weight"]
+ assert exporter.exclude_modules == ["model.layers.1.*"]
+
+
class _FakeTEGroupedMLP:
"""Minimal TEGroupedMLP stand-in exposing num_gemms, weight{i}, and state_dict()."""
@@ -1183,12 +1312,39 @@ def test_is_sidecar_writer_rank_pins_to_dp0_ep0(monkeypatch):
assert GPTModelExporter._is_sidecar_writer_rank(True) is False
-def _make_exporter_for_key_check(num_layers: int) -> GPTModelExporter:
+def _make_exporter_for_key_check(
+ num_layers: int, src_num_layers: int | None = None, num_mtp: int = 0
+) -> GPTModelExporter:
+ """``num_layers`` is the exported HF decoder depth; ``num_mtp`` MTP layers follow it."""
exporter = object.__new__(GPTModelExporter)
- exporter.model = SimpleNamespace(config=SimpleNamespace(num_layers=num_layers))
+ exporter._hf_text_config = SimpleNamespace(
+ num_hidden_layers=num_layers, num_nextn_predict_layers=num_mtp
+ )
+ exporter._src_num_hidden_layers = num_layers if src_num_layers is None else src_num_layers
+ exporter.rules = {"mtp_in_decoder_layers": num_mtp > 0}
return exporter
+@pytest.mark.parametrize(("export_mtp", "raises"), [(True, False), (False, True)])
+def test_verify_exported_keys_depth_pruned_with_decoder_mtp(tmp_path, export_mtp, raises):
+ """Pruned 3 -> 2 layers: source layer 2 is not required, and its MTP (layer 3) must land at 2."""
+ source, export = tmp_path / "src", tmp_path / "exp"
+ _write_index(
+ source,
+ [f"model.layers.{i}.input_layernorm.weight" for i in range(3)]
+ + ["model.layers.3.eh_proj.weight"],
+ )
+ exported = [f"model.layers.{i}.input_layernorm.weight" for i in range(2)]
+ _write_index(export, exported + (["model.layers.2.eh_proj.weight"] if export_mtp else []))
+ exporter = _make_exporter_for_key_check(num_layers=2, src_num_layers=3, num_mtp=1)
+
+ if raises:
+ with pytest.raises(RuntimeError, match=r"model\.layers\.2\.eh_proj\.weight"):
+ exporter._verify_exported_keys(str(export), str(source))
+ else:
+ exporter._verify_exported_keys(str(export), str(source))
+
+
def _write_index(dir_path: Path, keys) -> None:
dir_path.mkdir(parents=True, exist_ok=True)
(dir_path / "model.safetensors.index.json").write_text(
diff --git a/tests/gpu_megatron/torch/quantization/plugins/test_megatron.py b/tests/gpu_megatron/torch/quantization/plugins/test_megatron.py
index 041a0bd3fca..d636fa2ebdb 100644
--- a/tests/gpu_megatron/torch/quantization/plugins/test_megatron.py
+++ b/tests/gpu_megatron/torch/quantization/plugins/test_megatron.py
@@ -1813,6 +1813,86 @@ def test_kv_cache_quant(dist_workers_size_1, config):
dist_workers_size_1.run(partial(_test_kv_cache_quant_helper, config))
+def _get_tiny_dsa_gpt_model():
+ """Tiny GPT with DSA sparse attention (AbsorbedMLASelfAttention + DSAttention + indexer)."""
+ from megatron.core.models.gpt.experimental_attention_variant_module_specs import (
+ get_transformer_block_with_experimental_attention_variant_spec,
+ )
+ from megatron.core.transformer.transformer_config import MLATransformerConfig
+
+ config = MLATransformerConfig(
+ num_layers=2,
+ hidden_size=64,
+ num_attention_heads=4,
+ ffn_hidden_size=128,
+ q_lora_rank=32,
+ kv_lora_rank=16,
+ qk_head_dim=16,
+ qk_pos_emb_head_dim=16,
+ v_head_dim=16,
+ experimental_attention_variant="dsa",
+ dsa_indexer_n_heads=2,
+ dsa_indexer_head_dim=32,
+ dsa_indexer_topk=8,
+ normalization="RMSNorm",
+ add_bias_linear=False,
+ hidden_dropout=0.0, # deterministic forward for the checkpoint round-trip comparison
+ attention_dropout=0.0,
+ bf16=True,
+ params_dtype=torch.bfloat16,
+ )
+ spec = get_transformer_block_with_experimental_attention_variant_spec(config)
+ return GPTModel(
+ config=config,
+ transformer_layer_spec=spec,
+ vocab_size=64,
+ max_sequence_length=64,
+ position_embedding_type="rope",
+ # Tied so a checkpoint round-trip restores the output layer too.
+ share_embeddings_and_output_weights=True,
+ ).cuda()
+
+
+def _test_dsa_kv_cache_quant_helper(tmp_path, rank, size):
+ from megatron.core.transformer.experimental_attention_variant.dsa import DSAttention
+
+ initialize_for_megatron(tensor_model_parallel_size=1, pipeline_model_parallel_size=1, seed=SEED)
+ model, model_test = _get_tiny_dsa_gpt_model(), _get_tiny_dsa_gpt_model()
+ forward = get_forward(model)
+ # Calibrated (not constant-amax) FP8 KV cache: the case that needs a real V amax to export.
+ kv_fp8_calibrated = {
+ "quant_cfg": [
+ {"quantizer_name": "*", "enable": False},
+ {"quantizer_name": "*[kv]_bmm_quantizer", "cfg": {"num_bits": (4, 3), "axis": None}},
+ ],
+ "algorithm": "max",
+ }
+ model = mtq.quantize(model, kv_fp8_calibrated, forward)
+
+ dsa_modules = [m for m in model.modules() if isinstance(m, DSAttention)]
+ assert dsa_modules, "DSAttention was not converted for KV-cache quantization"
+ for module in dsa_modules:
+ assert module.k_bmm_quantizer.is_enabled and module.v_bmm_quantizer.is_enabled
+ # Absorbed MLA passes value=None; V is calibrated on the shared KV latent like K, so a
+ # calibrated FP8 KV cache still exports a v_scale.
+ assert module.k_bmm_quantizer.amax is not None
+ assert torch.equal(module.v_bmm_quantizer.amax, module.k_bmm_quantizer.amax)
+
+ # The KV quantizer state (incl. amax) survives a torch-dist round-trip alongside the indexer.
+ # DSA trains its indexer through a separate indexer loss, so the LM-loss backward in the helper
+ # gives it no gradient; exclude it from that check.
+ for name, param in model_test.named_parameters():
+ if ".indexer." in name:
+ param.requires_grad_(False)
+ sharded_state_dict_test_helper(tmp_path, model, model_test, forward)
+
+
+def test_dsa_kv_cache_quant(dist_workers_size_1, tmp_path):
+ """DSAttention gets calibrated K/V KV-cache quantizers that survive a checkpoint round-trip."""
+ pytest.importorskip("megatron.core.transformer.experimental_attention_variant.dsa")
+ dist_workers_size_1.run(partial(_test_dsa_kv_cache_quant_helper, tmp_path))
+
+
def _test_kv_cache_amax_sync_helper(config, rank, size, tensor_model_parallel_size=1):
"""Helper function for testing KV cache quantizer amax sync across distributed world."""
# Use rank in seed to produce different amax values across ranks
diff --git a/tests/unit/recipe/test_glm_5_3_recipe.py b/tests/unit/recipe/test_glm_5_3_recipe.py
index 980cadf681b..51028015ba2 100644
--- a/tests/unit/recipe/test_glm_5_3_recipe.py
+++ b/tests/unit/recipe/test_glm_5_3_recipe.py
@@ -22,7 +22,7 @@
*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``.
@@ -144,7 +144,7 @@ def test_glm_5_3_recipe_quantizer_precedence():
# 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.
+ # later `*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
@@ -170,3 +170,68 @@ def test_glm_5_3_recipe_quantizer_precedence():
# lm_head stays BF16.
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
diff --git a/tests/unit/torch/export/test_get_quantization.py b/tests/unit/torch/export/test_get_quantization.py
index 8a8a7ed93f2..40745966051 100644
--- a/tests/unit/torch/export/test_get_quantization.py
+++ b/tests/unit/torch/export/test_get_quantization.py
@@ -50,6 +50,7 @@
uses_iq_quantization,
)
from modelopt.torch.quantization.nn import (
+ GroupedQuantizer,
NVFP4StaticQuantizer,
SequentialQuantizer,
TensorQuantizer,
@@ -142,6 +143,34 @@ def test_uses_iq_quantization_false_without_iq_layers():
assert not uses_iq_quantization(model)
+class _FakeGroupedLinear(torch.nn.Module):
+ """TEGroupedLinear layout: ``weight0..N`` behind one GroupedQuantizer and no ``weight``."""
+
+ def __init__(self, weight_cfg):
+ super().__init__()
+ self.weight0 = torch.nn.Parameter(torch.randn(4, 4))
+ self.weight1 = torch.nn.Parameter(torch.randn(4, 4))
+ quantizers = [TensorQuantizer(), TensorQuantizer()]
+ for q in quantizers:
+ q.set_from_attribute_config(weight_cfg)
+ self.weight_quantizer = GroupedQuantizer(*quantizers)
+ self.input_quantizer = TensorQuantizer()
+ self.input_quantizer.set_from_attribute_config({"num_bits": (4, 3)})
+
+
+def test_grouped_experts_only_model_reports_its_format():
+ """An experts-only recipe leaves grouped experts as the only quantized modules."""
+ model = torch.nn.Sequential(torch.nn.Linear(4, 4), _FakeGroupedLinear({"num_bits": (4, 3)}))
+
+ assert get_quantization_format(model) == QUANTIZATION_FP8
+
+
+def test_uses_iq_quantization_sees_grouped_experts():
+ model = torch.nn.Sequential(_FakeGroupedLinear(_IQ_WEIGHT_CFG))
+
+ assert uses_iq_quantization(model)
+
+
def test_uses_iq_quantization_tolerates_sequential_quantizer():
"""A SequentialQuantizer has is_enabled but no num_bits, and is never IQ.