Skip to content

Fix prequant layernorm export without scales #1838

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
realAsma merged 2 commits into from
Jun 27, 2026
Merged

Fix prequant layernorm export without scales #1838

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
97bebd3
Fix prequant layernorm export without scales
realAsma Jun 26, 2026
8d4711c
fix(recipes): compose int4 blockwise recipe from shared units
realAsma Jun 26, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
7 changes: 4 additions & 3 deletions modelopt/torch/export/quant_utils.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -1257,9 +1257,10 @@ def fuse_prequant_layernorm(
fused_bias = bias * avg_pre_quant_scale
layernorm_output_scaled = (normalization(input) * fused_weight) + fused_bias
"""
pre_quant_scale = getattr(modules[0].input_quantizer, "_pre_quant_scale").to(
layernorm_module.weight.device
)
if not hasattr(modules[0].input_quantizer, "_pre_quant_scale"):
return

pre_quant_scale = modules[0].input_quantizer._pre_quant_scale.to(layernorm_module.weight.device)
Comment on lines +1260 to +1263

@coderabbitai coderabbitai Bot Jun 26, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

🎯 Functional Correctness | 🟠 Major | ⚡ Quick win

Use pre_quant_scale instead of the private _pre_quant_scale buffer.

TensorQuantizer.pre_quant_scale is also None when pre-quant scaling is disabled, not just when the buffer is missing. This guard bypasses that contract, so export can still fuse a stale scale into the LayerNorm even though the forward path no longer applies it.

Suggested fix 
-    if not hasattr(modules[0].input_quantizer, "_pre_quant_scale"):
+    pre_quant_scale = modules[0].input_quantizer.pre_quant_scale
+    if pre_quant_scale is None:
         return
-
-    pre_quant_scale = modules[0].input_quantizer._pre_quant_scale.to(layernorm_module.weight.device)
+    pre_quant_scale = pre_quant_scale.to(layernorm_module.weight.device)
📝 Committable suggestion

‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.

Suggested change
if not hasattr(modules[0].input_quantizer, "_pre_quant_scale"):
return
pre_quant_scale = modules[0].input_quantizer._pre_quant_scale.to(layernorm_module.weight.device)
pre_quant_scale = modules[0].input_quantizer.pre_quant_scale
if pre_quant_scale is None:
return
pre_quant_scale = pre_quant_scale.to(layernorm_module.weight.device)
🤖 Prompt for AI Agents 
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

In `@modelopt/torch/export/quant_utils.py` around lines 1260 - 1263, The LayerNorm
export logic is checking the private TensorQuantizer buffer `_pre_quant_scale`
directly, which can miss the case where `TensorQuantizer.pre_quant_scale` is
`None` and incorrectly fuse a stale scale. Update the guard in `quant_utils`
where the `modules[0].input_quantizer` is handled to use the public
`pre_quant_scale` property instead of `hasattr(..., "_pre_quant_scale")`, and
only proceed with the fusion path when that property is present and non-None
before moving it to `layernorm_module.weight.device`.

if _layernorm_uses_weight_plus_one(layernorm_module):
# For norms that use (1 + weight) in forward, fold pre_quant_scale into the effective weight.
fused_weight = (layernorm_module.weight + 1.0) * pre_quant_scale - 1.0
Expand Down
46 changes: 9 additions & 37 deletions modelopt_recipes/general/ptq/int4_blockwise_weight_only.yaml
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -16,47 +16,19 @@
metadata:

@meenchen meenchen Jun 26, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

@realAsma could you share why we need this int4 recipe in the modelopt_recipes/general/ptq/? I think we don't support export for int4 blockwise. cc @Edwardf0t1

recipe_type: ptq
description: INT4 blockwise weight-only (W4A16, block size 128), max calibration.

imports:
base_disable_all: configs/ptq/units/base_disable_all
default_disabled_quantizers: configs/ptq/units/default_disabled_quantizers
int4_per_block: configs/numerics/int4_per_block

quantize:
algorithm: max
quant_cfg:
- quantizer_name: '*'
enable: false
- $import: base_disable_all
- quantizer_name: '*weight_quantizer'
cfg:
num_bits: 4
block_sizes:
-1: 128
$import: int4_per_block
- quantizer_name: '*input_quantizer'
enable: false
- quantizer_name: '*block_sparse_moe.gate*'
enable: false
- quantizer_name: '*linear_attn.conv1d*'
enable: false
- quantizer_name: '*lm_head*'
enable: false
- quantizer_name: '*mixer.conv1d*'
enable: false
- quantizer_name: '*mlp.gate.*'
enable: false
- quantizer_name: '*mlp.shared_expert_gate.*'
enable: false
- quantizer_name: '*output_layer*'
enable: false
- quantizer_name: '*proj_out.*'
enable: false
- quantizer_name: '*router*'
enable: false
- quantizer_name: 'output.*'
enable: false
- parent_class: 'nn.BatchNorm1d'
quantizer_name: '*'
enable: false
- parent_class: 'nn.BatchNorm2d'
quantizer_name: '*'
enable: false
- parent_class: 'nn.BatchNorm3d'
quantizer_name: '*'
enable: false
- parent_class: 'nn.LeakyReLU'
quantizer_name: '*'
enable: false
- $import: default_disabled_quantizers
9 changes: 7 additions & 2 deletions tests/unit/recipe/test_loader.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -156,6 +156,7 @@ def test_load_recipe_builtin_description():
_BUILTIN_PTQ_RECIPES = [
"general/ptq/fp8_default-kv_fp8",
"general/ptq/fp8_default-kv_fp8_cast",
"general/ptq/int4_blockwise_weight_only",
"general/ptq/nvfp4_default-kv_fp8",
"general/ptq/nvfp4_default-kv_fp8_cast",
"general/ptq/nvfp4_default-kv_nvfp4_cast",
Expand Down Expand Up @@ -509,6 +510,7 @@ def test_load_recipe_dflash_field_validation_raises(tmp_path):
("yaml_path", "model_cfg_name", "kv_cfg_name"),
[
("general/ptq/fp8_default-kv_fp8.yaml", "FP8_DEFAULT_CFG", "FP8_KV_CFG"),
("general/ptq/int4_blockwise_weight_only.yaml", "INT4_BLOCKWISE_WEIGHT_ONLY_CFG", None),
("general/ptq/nvfp4_default-kv_fp8.yaml", "NVFP4_DEFAULT_CFG", "FP8_KV_CFG"),
("general/ptq/nvfp4_mlp_only-kv_fp8.yaml", "NVFP4_MLP_ONLY_CFG", "FP8_KV_CFG"),
("general/ptq/nvfp4_omlp_only-kv_fp8.yaml", "NVFP4_OMLP_ONLY_CFG", "FP8_KV_CFG"),
Expand All @@ -517,7 +519,7 @@ def test_load_recipe_dflash_field_validation_raises(tmp_path):
def test_general_ptq_yaml_matches_config_dicts(yaml_path, model_cfg_name, kv_cfg_name):
"""Each general/ptq YAML's quant_cfg list matches the merged Python config dicts."""
model_cfg = getattr(qcfg, model_cfg_name)
kv_cfg = getattr(qcfg, kv_cfg_name)
kv_cfg = getattr(qcfg, kv_cfg_name) if kv_cfg_name is not None else None
recipe = load_recipe(yaml_path)
yaml_data = {"quantize": recipe.quantize}

Expand Down Expand Up @@ -552,7 +554,10 @@ def _normalize_entries(raw_entries):
def _sort_key(entry):
return json.dumps(entry, sort_keys=True, default=str)

python_entries = _normalize_entries(model_cfg["quant_cfg"] + kv_cfg["quant_cfg"])
python_quant_cfg = model_cfg["quant_cfg"]
if kv_cfg is not None:
python_quant_cfg = python_quant_cfg + kv_cfg["quant_cfg"]
python_entries = _normalize_entries(python_quant_cfg)
yaml_entries = _normalize_entries(yaml_data["quantize"]["quant_cfg"])

assert sorted(python_entries, key=_sort_key) == sorted(yaml_entries, key=_sort_key)
Expand Down
37 changes: 36 additions & 1 deletion tests/unit/torch/export/test_unified_export_hf.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -28,8 +28,9 @@
_collect_canonical_tied_patterns,
_reorder_canonical_first,
)
from modelopt.torch.export.quant_utils import sync_tied_input_amax
from modelopt.torch.export.quant_utils import fuse_prequant_layernorm, sync_tied_input_amax
from modelopt.torch.export.unified_export_hf import _export_quantized_weight
from modelopt.torch.quantization.nn import TensorQuantizer


def test_collect_canonical_tied_patterns_dict_style():
Expand Down Expand Up @@ -180,3 +181,37 @@ def test_export_quantized_weight_skips_alias_when_one_tied_side_is_unquantized()
assert enc.weight.data_ptr() != original_shared_data_ptr # encoder got fresh packed
assert dec.weight.data_ptr() == original_shared_data_ptr # decoder untouched
assert enc.weight.data_ptr() != dec.weight.data_ptr()


def _linear_with_input_quantizer():
linear = torch.nn.Linear(4, 4, bias=False)
linear.input_quantizer = TensorQuantizer()
return linear


def test_fuse_prequant_layernorm_skips_modules_without_pre_quant_scale():
layernorm = torch.nn.LayerNorm(4)
original_weight = layernorm.weight.detach().clone()
modules = [_linear_with_input_quantizer(), _linear_with_input_quantizer()]

fuse_prequant_layernorm(layernorm, modules)

assert torch.allclose(layernorm.weight, original_weight)
assert not hasattr(modules[0], "fused_with_prequant")
assert not hasattr(modules[1], "fused_with_prequant")


def test_fuse_prequant_layernorm_fuses_and_removes_pre_quant_scale():
layernorm = torch.nn.LayerNorm(4)
modules = [_linear_with_input_quantizer(), _linear_with_input_quantizer()]
pre_quant_scale = torch.tensor([1.0, 2.0, 3.0, 4.0])
for module in modules:
module.input_quantizer._pre_quant_scale = pre_quant_scale

fuse_prequant_layernorm(layernorm, modules)

assert torch.allclose(layernorm.weight, pre_quant_scale)
assert torch.allclose(layernorm.bias, torch.zeros_like(pre_quant_scale))
for module in modules:
assert not hasattr(module.input_quantizer, "_pre_quant_scale")
assert module.fused_with_prequant
Loading