Keep GatedDeltaNet beta / alpha gates unquantized in Megatron fake quant

CHANGELOG.rst

@@ -63,6 +63,7 @@ Changelog
 
 **Bug Fixes**
 
+- Megatron-Core fake quantization of Gated DeltaNet (Qwen3.5 / Qwen3.6) with a Transformer Engine ``in_proj`` now keeps its beta / alpha gates in BF16 during evaluation, QAD, ``fold_weight`` and the vLLM fake-quant export, matching deployment. Other ``in_proj`` layers stay fully quantized and warn.
 - ``hf_quant_config.json`` now declares ``pre_quant_scale`` only for INT4 / W4A8 layers that have the tensor; loading INT4 / W4A8 fused-MoE checkpoints in TensorRT-LLM also needs a release whose MoE loader treats it as optional. AWQ, SmoothQuant and SVDQuant now warn when they skip quantized fused-MoE experts.
 - Hugging Face fused MoE experts stored transposed, ``(num_experts, in_dim, out_dim)``, without a quantized wrapper of their own (for example Aria's) are now left unquantized with a warning. They were quantized with gate and up split along the wrong axis, so their exported weights were wrong; the Hugging Face export now refuses such a model's MoE block instead.
 - Calibrating a Megatron-Core MLA model now gives ``linear_q_down_proj`` and ``linear_kv_down_proj`` one per-tensor input and weight range, as vLLM and TensorRT-LLM fuse them into one GEMM with one scale per tensor. Re-calibrate existing MLA checkpoints to get the shared ranges.

modelopt/torch/export/plugins/vllm_fakequant_megatron.py

@@ -33,6 +33,7 @@
 from modelopt.torch.export.quant_format import QUANTIZATION_NONE
 from modelopt.torch.export.unified_export_megatron import GPTModelExporter
 from modelopt.torch.quantization.nn import GroupedQuantizer, SequentialQuantizer, TensorQuantizer
+from modelopt.torch.quantization.utils.core_utils import apply_without_qdq
 from modelopt.torch.utils import get_unwrapped_name
 from modelopt.torch.utils.distributed import DistributedProcessGroup, is_master
 
@@ -405,8 +406,12 @@ def _self_attention_scaling(
         self._state_dict.update({prefix + name: value for name, value in state.items()})
 
     @staticmethod
-    def _fakequant_weight(weight, weight_quantizer, dtype):
-        """Return a CPU QDQ weight while preserving the quantizer's device."""
+    def _fakequant_weight(weight, weight_quantizer, dtype, unquantized_rows=None):
+        """Return a CPU QDQ weight while preserving the quantizer's device.
+
+        ``unquantized_rows`` skip the QDQ but keep the quantizer's scaling / rotation, as in the
+        fake-quant forward (GatedDeltaNet's beta / alpha gates).
+        """
         weight = weight.to(dtype)
         # Fold the weight_quantizer into the weight by applying fake-quantization
         # (quantize then dequantize). The weight_quantizer amax is not exported;
@@ -430,7 +435,17 @@ def _fakequant_weight(weight, weight_quantizer, dtype):
                 if need_move:
                     weight_quantizer.to(quant_device)
                 try:
-                    weight = weight_quantizer(weight.to(quant_device)).to(dtype)
+                    weight = weight.to(quant_device)
+                    qweight = weight_quantizer(weight)
+                    if unquantized_rows is not None:
+                        qweight = torch.cat(
+                            [
+                                qweight[: unquantized_rows.start],
+                                apply_without_qdq(weight_quantizer, weight[unquantized_rows]),
+                                qweight[unquantized_rows.stop :],
+                            ]
+                        )
+                    weight = qweight.to(dtype)
                 finally:
                     if need_move:
                         weight_quantizer.to(wq_dev)
@@ -478,7 +493,10 @@ def _get_quantized_state(
         # Grouped slicing reads QDQ weights from the hooked module.state_dict().
         if "weight" in name_to_value and module not in self._grouped_fq_modules:
             name_to_value["weight"] = self._fakequant_weight(
-                name_to_value["weight"], getattr(module, "weight_quantizer", None), dtype
+                name_to_value["weight"],
+                getattr(module, "weight_quantizer", None),
+                dtype,
+                getattr(module, "_modelopt_unquantized_output_rows", None),
             )
 
         # Save activation ranges; weight quantizers are folded into the weights above.

modelopt/torch/quantization/plugins/megatron.py

@@ -1141,6 +1141,21 @@ def _setup(self):
                 data_parallel_group,
                 mcore_parallel.get_tensor_model_parallel_group(),
             )
+            # The fused in_proj ends with the beta / alpha gates, which recipes and every export
+            # keep in BF16 (vLLM's in_proj_ba is too narrow for FP4 kernels), so leave them
+            # unquantized here too.
+            names = list(getattr(self, "in_proj_split_names", ()))
+            if not isinstance(self.in_proj, TELayerNormColumnParallelLinear):
+                warn_rank_0(
+                    f"GatedDeltaNet in_proj is a {type(self.in_proj).__name__}, so fake quantization"
+                    " also quantizes its beta / alpha gates, which exports keep in BF16."
+                )
+            elif names[-2:] == ["beta", "alpha"]:
+                local_rows = sum(self.in_proj_split_sections)
+                gate_rows = sum(self.in_proj_split_sections[-2:])
+                self.in_proj._modelopt_unquantized_output_rows = slice(
+                    local_rows - gate_rows, local_rows
+                )
 
         @contextmanager
         def _quantized_gdn_kernel(self):

modelopt/torch/quantization/plugins/transformer_engine.py

@@ -26,6 +26,7 @@
 import transformer_engine.pytorch.module.linear as te_linear
 
 from modelopt.torch.quantization.utils import replace_function
+from modelopt.torch.quantization.utils.core_utils import apply_without_qdq
 
 from ..nn import GroupedQuantizer, QuantModuleRegistry, SequentialQuantizer, TensorQuantizer
 from .custom import _ParallelLinear
@@ -315,19 +316,41 @@ class _QuantLayerNormLinearFunc(torch.autograd.Function):
 
     @staticmethod
     def forward(ctx, inp, ln_weight, ln_bias, weight, *args, **kwargs):
-        input_quantizer, weight_quantizer = _QuantLayerNormLinearFunc.modelopt_quantizers
+        input_quantizer, weight_quantizer, unquantized_rows = (
+            _QuantLayerNormLinearFunc.modelopt_quantizers
+        )
 
         qweight = weight_quantizer(weight)
+        if unquantized_rows is not None:
+            # Deployment keeps these rows in BF16: they skip QDQ but keep the quantizers' scaling
+            # and rotation. Not in place: the fake-quant weight cache may hold qweight.
+            row_weight = apply_without_qdq(weight_quantizer, weight[unquantized_rows])
+            qweight = torch.cat(
+                (
+                    qweight[: unquantized_rows.start],
+                    row_weight.to(qweight.dtype),
+                    qweight[unquantized_rows.stop :],
+                )
+            )
         qweight.requires_grad = weight.requires_grad
         if ctx is not None:
             # We need to recompute the quantized input for the backward pass, so we save the input_quantizer
             ctx.modelopt_input_quantizer = input_quantizer
+            ctx.modelopt_unquantized_rows = unquantized_rows
 
         original_gemm = te_layernorm_linear.general_gemm
 
-        def _patched_general_gemm(weight, input, *gemm_args, **gemm_kwargs):
+        def _patched_general_gemm(gemm_weight, input, *gemm_args, **gemm_kwargs):
             qinput = input_quantizer(input)
-            return original_gemm(weight, qinput, *gemm_args, **gemm_kwargs)
+            outputs = original_gemm(gemm_weight, qinput, *gemm_args, **gemm_kwargs)
+            if unquantized_rows is not None:
+                bias = gemm_kwargs.get("bias")
+                outputs[0][..., unquantized_rows] = torch.nn.functional.linear(
+                    apply_without_qdq(input_quantizer, input).to(row_weight.dtype),
+                    row_weight,
+                    None if bias is None else bias[unquantized_rows],
+                ).to(outputs[0].dtype)
+            return outputs
 
         with replace_function(
             te_layernorm_linear,
@@ -366,7 +389,16 @@ def _patched_general_gemm(a, b, *gemm_args, **gemm_kwargs):
             # but gemm is called with the unquantized input (a)
             # So lets first get the quantized input (qinput) and then call the gemm
             qinput = ctx.modelopt_input_quantizer(a)
-            return original_gemm(qinput, b, *gemm_args, **gemm_kwargs)
+            outputs = original_gemm(qinput, b, *gemm_args, **gemm_kwargs)
+            rows = ctx.modelopt_unquantized_rows
+            if rows is not None and outputs[0] is not None:
+                # The forward computed these rows from the unquantized input.
+                x = apply_without_qdq(ctx.modelopt_input_quantizer, a)
+                dy = b[..., rows].reshape(-1, rows.stop - rows.start)
+                outputs[0][rows] = (dy.t() @ x.reshape(dy.shape[0], -1).to(dy.dtype)).to(
+                    outputs[0].dtype
+                )
+            return outputs
 
         with replace_function(
             te_layernorm_linear,
@@ -378,6 +410,7 @@ def _patched_general_gemm(a, b, *gemm_args, **gemm_kwargs):
             outputs = te_layernorm_linear._LayerNormLinear.backward(ctx, *grad_outputs)
 
         delattr(ctx, "modelopt_input_quantizer")
+        delattr(ctx, "modelopt_unquantized_rows")
         return outputs
 
 
@@ -394,13 +427,26 @@ def _setup(self):
             )
             self.fuse_wgrad_accumulation = False
 
+    def fold_weight(self, keep_attrs: bool = False):
+        """Fold the weight quantizer, keeping the rows the forward pass leaves unquantized."""
+        rows = getattr(self, "_modelopt_unquantized_output_rows", None)
+        unquantized = None
+        if rows is not None:
+            with torch.no_grad():
+                # Cloned: without transforms this returns a view of the weight that folding overwrites.
+                unquantized = apply_without_qdq(self.weight_quantizer, self.weight[rows]).clone()
+        super().fold_weight(keep_attrs)
+        if unquantized is not None:
+            self.weight.data[rows] = unquantized.to(self.weight.dtype)
+
     def forward(self, *args, **kwargs):
         """Call ModelOpt patch for _LayerNormLinear functional."""
         _assert_te_fp8_enabled()
         # This is multi-process safe (such as in torch distributed jobs), not multi-thread safe
         _QuantLayerNormLinearFunc.modelopt_quantizers = (
             self.input_quantizer,
             self.weight_quantizer,
+            getattr(self, "_modelopt_unquantized_output_rows", None),
         )
         with replace_function(
             te_layernorm_linear,

modelopt/torch/quantization/utils/core_utils.py

@@ -1024,6 +1024,25 @@ def enable_quant(quantizer):
         quantizer._if_quant = original_if_quant
 
 
+def apply_without_qdq(quantizer, inputs: torch.Tensor) -> torch.Tensor:
+    """Apply ``quantizer``'s pre-quant scale and rotation to ``inputs``, without quantizing them.
+
+    For rows a quantized module computes in high precision, so they keep the transforms (e.g.
+    an AWQ / SmoothQuant scale that the stored weight compensates for) of the quantized rows.
+    """
+    from ..nn import SequentialQuantizer  # local import: ..nn imports this module
+
+    members = list(quantizer) if isinstance(quantizer, SequentialQuantizer) else [quantizer]
+    for member in members:
+        disabled = member._disabled
+        member._disabled = True
+        try:
+            inputs = member(inputs)
+        finally:
+            member._disabled = disabled
+    return inputs
+
+
 @contextmanager
 def disable_calib(quantizer):
     """Temporarily disable calibration for a quantizer.

tests/gpu_megatron/torch/export/test_vllm_fakequant_megatron_export.py

@@ -22,13 +22,15 @@
 
 import pytest
 import torch
+import transformers
 import yaml
 from _test_utils.torch.megatron.models import get_mcore_gpt_model, get_mcore_hybrid_model
 from _test_utils.torch.megatron.utils import initialize_for_megatron, run_mcore_inference
 from _test_utils.torch.transformers_models import (
     create_tiny_gpt_oss_dir,
     create_tiny_llama_dir,
     create_tiny_nemotron_h_dir,
+    get_tiny_qwen3_5_vl_offline,
 )
 from megatron.core.parallel_state import is_pipeline_last_stage
 from safetensors import safe_open
@@ -45,6 +47,7 @@
     gather_mcore_vllm_fq_quantizer_recipe,
 )
 from modelopt.torch.quantization.nn import GroupedQuantizer, SequentialQuantizer, TensorQuantizer
+from modelopt.torch.quantization.utils.core_utils import apply_without_qdq
 
 
 def _assert_exported_quantizers(export_dir, expected_names, amax=1.001, disabled_names=()):
@@ -414,6 +417,80 @@ def test_mcore_vllm_export_mtp(request, tmp_path):
     workers.run(partial(_test_mcore_vllm_export_mtp, tmp_path))
 
 
+def _test_mcore_vllm_export_gdn(source, tmp_path, rank, size):
+    text_cfg = transformers.AutoConfig.from_pretrained(source).text_config
+    model = get_mcore_gpt_model(
+        tensor_model_parallel_size=size,
+        initialize_megatron=True,
+        num_layers=text_cfg.num_hidden_layers,
+        hidden_size=text_cfg.hidden_size,
+        num_attention_heads=text_cfg.num_attention_heads,
+        num_query_groups=text_cfg.num_key_value_heads,
+        ffn_hidden_size=text_cfg.intermediate_size,
+        max_sequence_length=text_cfg.max_position_embeddings,
+        vocab_size=text_cfg.vocab_size,
+        normalization="RMSNorm",
+        transformer_impl="modelopt",
+        kv_channels=text_cfg.head_dim,
+        experimental_attention_variant="gated_delta_net",
+        linear_attention_freq=len(text_cfg.layer_types),
+    ).cuda()
+    model.eval()
+
+    def forward_loop(model):
+        with torch.no_grad():
+            run_mcore_inference(model, torch.randint(0, 32, (1, 32), device="cuda"))
+
+    model = mtq.quantize(model, mtq.FP8_DEFAULT_CFG, forward_loop)
+    gdn = model.decoder.layers[0].self_attention
+    weight_quantizer = gdn.in_proj.weight_quantizer
+    if find_spec("fast_hadamard_transform") is not None:
+        # A CUDA-only transform that the gate rows keep: the export must stage them onto CUDA.
+        weight_quantizer.set_from_attribute_config(
+            {"rotate": {"enable": True, "mode": "rotate_back"}}
+        )
+    with torch.no_grad():
+        weight = gdn.in_proj.weight.to(torch.bfloat16)
+        transformed = apply_without_qdq(weight_quantizer, weight).cpu()
+        weight = weight.cpu()
+    model.cpu()  # CPU-offloaded weights
+    export_dir = tmp_path / "gdn_export"
+    export_mcore_gpt_to_hf_vllm_fq(
+        model, str(source), dtype=torch.bfloat16, export_dir=str(export_dir)
+    )
+
+    sections = dict(zip(gdn.in_proj_split_names, gdn.in_proj_split_sections))
+    qkvz_rows = sum(sections[name] for name in ("query", "key", "value", "z"))
+    gate_rows = {
+        "in_proj_b": transformed[qkvz_rows : qkvz_rows + sections["beta"]],
+        "in_proj_a": transformed[qkvz_rows + sections["beta"] :],
+    }
+    prefix = "model.language_model.layers.0.linear_attn."
+    state = torch.load(export_dir / "quantizer_state.pth", weights_only=True, map_location="cpu")
+    recipe = yaml.safe_load((export_dir / "quant_recipe.yaml").read_text())["quantizers"]
+    weight_map = json.loads((export_dir / "model.safetensors.index.json").read_text())["weight_map"]
+    # The beta / alpha gates are written unquantized and get no quantizers; qkv / z keep theirs.
+    for name, expected in gate_rows.items():
+        key = f"{prefix}{name}.weight"
+        with safe_open(export_dir / weight_map[key], framework="pt") as f:
+            assert torch.equal(f.get_tensor(key), expected)
+        assert not any(f"{prefix}{name}." in key for key in state.keys() | recipe.keys())
+    key = f"{prefix}in_proj_qkv.weight"
+    qkv_rows = sum(sections[name] for name in ("query", "key", "value"))
+    with safe_open(export_dir / weight_map[key], framework="pt") as f:
+        assert not torch.equal(f.get_tensor(key), weight[:qkv_rows])
+    for name in ("in_proj_qkv", "in_proj_z"):
+        assert f"{prefix}{name}.input_quantizer" in recipe
+
+
+def test_mcore_vllm_export_gdn(dist_workers_size_1, tmp_path):
+    """GatedDeltaNet's beta / alpha gates export unquantized, with no quantizers of their own."""
+    pytest.importorskip("fla")  # Megatron-Core GatedDeltaNet needs the optional fla
+    source = tmp_path / "tiny_qwen3_5_vl"
+    get_tiny_qwen3_5_vl_offline().save_pretrained(source)
+    dist_workers_size_1.run(partial(_test_mcore_vllm_export_gdn, source, tmp_path))
+
+
 def _assert_unsupported_settings(model, source, export_dir, rank, size):
     attribute_cfgs = [
         (