Fix prequant layernorm export without scales (#1838)

- Skip pre-quant LayerNorm fusion when the representative input
quantizer has no `_pre_quant_scale` buffer.
- Add CPU unit coverage for the no-scale weight-only path and the
existing fusion/removal behavior when pre-quant scales are present.
- Compose the public `general/ptq/int4_blockwise_weight_only` recipe
from the shared recipe units so it stays aligned with
`INT4_BLOCKWISE_WEIGHT_ONLY_CFG`.

NVBug: https://nvbugspro.nvidia.com/bug/6311597

NVBug 6311597 reports a deterministic HF export crash for
`general/ptq/int4_blockwise_weight_only` on Llama-3.1-8B. That
weight-only recipe disables input quantization and skips calibration, so
`_pre_quant_scale` is not registered, but export still reaches
`fuse_prequant_layernorm` through the AWQ-like format path.

This PR also carries the recipe-sync diff that was previously in draft
PR #1836; PR #1836 has been closed after moving that diff here.

- `pre-commit run --files modelopt/torch/export/quant_utils.py
tests/unit/torch/export/test_unified_export_hf.py
modelopt_recipes/general/ptq/int4_blockwise_weight_only.yaml
tests/unit/recipe/test_loader.py`
- `pytest_pwd
tests/unit/torch/export/test_unified_export_hf.py::test_fuse_prequant_layernorm_skips_modules_without_pre_quant_scale
tests/unit/recipe/test_loader.py::test_load_recipe_all_builtins
tests/unit/recipe/test_loader.py::test_general_ptq_yaml_matches_config_dicts
tests/unit/recipe/test_presets.py -q` -> 29 passed

Full GB10/Llama repro was not run locally.

<!-- This is an auto-generated comment: release notes by coderabbit.ai
-->

* **Bug Fixes**
* Improved handling of layer norm fusion so it safely skips cases where
pre-quantization scale data is unavailable, avoiding unexpected errors.
* Updated the layer norm fusion flow to correctly apply scaling when
pre-quantization data is present.

* **Tests**
* Expanded coverage for export and recipe loading scenarios, including
cases with and without pre-quantization scale data.
  * Added validation for the new int4 blockwise weight-only recipe.

<!-- end of auto-generated comment: release notes by coderabbit.ai -->

---------

Signed-off-by: realAsma <akuriparambi@nvidia.com>
Signed-off-by: Keval Morabia <28916987+kevalmorabia97@users.noreply.github.com>

Author: realAsma

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

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

tests/unit/recipe/test_loader.py

@@ -157,6 +157,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",
@@ -538,6 +539,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"),
@@ -546,7 +548,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}
 
@@ -581,7 +583,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)