From 97bebd360c27e0a1a07e4a8a91390fa589099216 Mon Sep 17 00:00:00 2001 From: realAsma Date: Fri, 26 Jun 2026 22:18:58 +0000 Subject: [PATCH 1/2] Fix prequant layernorm export without scales Signed-off-by: realAsma --- modelopt/torch/export/quant_utils.py | 7 ++-- .../torch/export/test_unified_export_hf.py | 37 ++++++++++++++++++- 2 files changed, 40 insertions(+), 4 deletions(-) diff --git a/modelopt/torch/export/quant_utils.py b/modelopt/torch/export/quant_utils.py index 2af5f6eab0b..b5867729fbd 100755 --- a/modelopt/torch/export/quant_utils.py +++ b/modelopt/torch/export/quant_utils.py @@ -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) 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 diff --git a/tests/unit/torch/export/test_unified_export_hf.py b/tests/unit/torch/export/test_unified_export_hf.py index 3032353f914..118331ce3d9 100644 --- a/tests/unit/torch/export/test_unified_export_hf.py +++ b/tests/unit/torch/export/test_unified_export_hf.py @@ -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(): @@ -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 From 8d4711c8531a6cb1c34253c412eb7e9bdb60f573 Mon Sep 17 00:00:00 2001 From: realAsma Date: Fri, 26 Jun 2026 21:23:24 +0000 Subject: [PATCH 2/2] fix(recipes): compose int4 blockwise recipe from shared units Signed-off-by: realAsma --- .../ptq/int4_blockwise_weight_only.yaml | 46 ++++--------------- tests/unit/recipe/test_loader.py | 9 +++- 2 files changed, 16 insertions(+), 39 deletions(-) diff --git a/modelopt_recipes/general/ptq/int4_blockwise_weight_only.yaml b/modelopt_recipes/general/ptq/int4_blockwise_weight_only.yaml index 432b970339c..91bf57971de 100644 --- a/modelopt_recipes/general/ptq/int4_blockwise_weight_only.yaml +++ b/modelopt_recipes/general/ptq/int4_blockwise_weight_only.yaml @@ -16,47 +16,19 @@ metadata: 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 diff --git a/tests/unit/recipe/test_loader.py b/tests/unit/recipe/test_loader.py index 26e328426d0..829f4ec3a37 100644 --- a/tests/unit/recipe/test_loader.py +++ b/tests/unit/recipe/test_loader.py @@ -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", @@ -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"), @@ -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} @@ -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)