Add composable scale calibration to GPTQ

CHANGELOG.rst

@@ -19,6 +19,10 @@ Changelog
 
 **New Features**
 
+- Add ``scale_algorithm`` to GPTQ, supporting MSE and local-Hessian scale calibration (including
+  ``activation_error_coupling``) while retaining max calibration as the default. Combining MSE
+  scale calibration with ``four_over_six`` weight block sizes enables GPTQ with NVFP4 4/6;
+  ``fused=True`` now raises an error for that unsupported combination.
 - Add the ``prepare_megatron_data_blend`` utility to prepare weighted Megatron data blends from YAML configs, including optional token-budgeted subsets for distillation workflows. See the `Megatron data preparation guide <https://github.com/NVIDIA/Model-Optimizer/tree/main/examples/dataset/MEGATRON_DATA_PREP.md#prepare-token-budgeted-data-blends>`_.
 - Add Learned Scale Quantization (LSQ) and Dual-LSQ support for quantization-aware distillation, including learnable ``amax`` parameters, tied-scale and pre-scale options, focused NVFP4 recipes, and scale-only training.
 - Add the **D-PACE** loss objective for DFlash speculative-decoding training (`arXiv:2605.18810 <https://arxiv.org/abs/2605.18810>`_) and make it the default (``dflash_loss_objective: dpace``). It replaces the static exponential position decay with dynamic, confidence-derived per-position weights that adapt to whichever block positions currently limit acceptance. Smoothing is controlled by ``dflash_dpace_alpha`` (default 0.5); set ``dflash_loss_objective: decay`` to restore the previous static schedule. Training-only and detached from the gradient (no architecture or inference change).

modelopt/torch/quantization/config.py

@@ -1087,6 +1087,16 @@ class LocalHessianCalibConfig(_SharedStatesConfig, QuantizeAlgorithmConfig):
     )
 
 
+_ScaleCalibConfig: TypeAlias = MaxCalibConfig | MseCalibConfig | LocalHessianCalibConfig
+
+
+def _serialize_scale_algorithm_sparse(value: _ScaleCalibConfig | None) -> dict | None:
+    """Preserve the sparse public dict shape accepted by ``scale_algorithm`` fields."""
+    if value is None:
+        return None
+    return {"method": value.method, **value.model_dump(exclude={"method"}, exclude_unset=True)}
+
+
 class SmoothQuantCalibConfig(QuantizeAlgorithmConfig):
     """The config for ``smoothquant`` algorithm (SmoothQuant).
 
@@ -1237,6 +1247,11 @@ class GPTQCalibConfig(QuantizeAlgorithmConfig):
     sequentially so that each layer's Hessian is computed from activations that already reflect
     the quantization of preceding layers.
 
+    Scale calibration (max by default) sets the amax values that GPTQ quantizes against. NVFP4
+    four-over-six composes through ``four_over_six`` weight ``block_sizes`` and MSE scale
+    calibration. In layerwise mode, scale calibration runs once per decoder layer. Fused GPTQ is
+    incompatible with four-over-six weight quantization.
+
     The default values are taken from the official GPTQ implementation:
     https://github.com/IST-DASLab/FP-Quant/blob/d2e3092f968262c4de5fb050e1aef568a280dadd/src/quantization/gptq.py#L35
     """
@@ -1261,6 +1276,20 @@ class GPTQCalibConfig(QuantizeAlgorithmConfig):
         description="""When True, use a fused Triton kernel that combines quantization and
         per-column error propagation into one launch per GPTQ block.""",
     )
+    scale_algorithm: _ScaleCalibConfig | None = ModeloptField(
+        default=None,
+        title="Scale calibration algorithm run before the GPTQ weight update.",
+        description=(
+            "Dict with 'method' key: 'max', 'mse', or 'local_hessian'. Extra keys are "
+            "forwarded to the calibration function, e.g. 'fp8_scale_sweep' (mse/local_hessian) "
+            "or 'activation_error_coupling' (local_hessian). Defaults to {'method': 'max'} "
+            "if None, preserving historical GPTQ behavior."
+        ),
+    )
+
+    @field_serializer("scale_algorithm")
+    def _serialize_scale_algorithm(self, value: _ScaleCalibConfig | None):
+        return _serialize_scale_algorithm_sparse(value)
 
     @model_validator(mode="after")
     def _gptq_qdq_default(self):
@@ -1278,9 +1307,6 @@ def _gptq_qdq_default(self):
         return self
 
 
-_ScaleCalibConfig: TypeAlias = MaxCalibConfig | MseCalibConfig | LocalHessianCalibConfig
-
-
 class LSQConfig(QuantizeAlgorithmConfig):
     """Config for LSQ (Learnt Scale Quantization) and Dual-LSQ algorithms.
 
@@ -1356,9 +1382,7 @@ class LSQConfig(QuantizeAlgorithmConfig):
     @field_serializer("scale_algorithm")
     def _serialize_scale_algorithm(self, value: _ScaleCalibConfig | None):
         """Preserve the sparse public dict shape accepted by this field."""
-        if value is None:
-            return None
-        return {"method": value.method, **value.model_dump(exclude={"method"}, exclude_unset=True)}
+        return _serialize_scale_algorithm_sparse(value)
 
     @model_validator(mode="after")
     def _validate_tied_amax(self):

modelopt/torch/quantization/model_calib.py

@@ -2184,6 +2184,7 @@ def gptq(
     perc_damp: float = 0.01,
     block_size: int = 128,
     fused: bool = False,
+    scale_algorithm: dict | None = None,
 ):
     """GPTQ quantization.
 
@@ -2197,7 +2198,8 @@ def gptq(
 
     Per-module steps:
 
-    1. ``max_calibrate`` to set amax values from the current activations.
+    1. Scale calibration (``scale_algorithm``; max by default) to set amax values
+       from the current activations.
     2. Promote eligible quantizers to ``StaticBlockScaleQuantizer`` (two-level scaling).
     3. Collect per-linear-layer Hessian matrices via forward hooks.
     4. Blockwise weight updates using the inverse Hessian to compensate for
@@ -2210,11 +2212,31 @@ def gptq(
         perc_damp: Percentage of avg Hessian diagonal for damping (default: 0.01).
         block_size: Block size for GPTQ weight update.
         fused: If True, use fused Triton kernel for NVFP4 static quantization.
+        scale_algorithm: Calibration algorithm config to run first. Dict with
+            'method' key: 'mse', 'local_hessian', or 'max'. Defaults to
+            {'method': 'max'} if None. Under layerwise mode, calibration runs
+            once per decoder layer.
     """
     total_start = time.time()
 
-    # TODO: Add support for other scale setting strateiges like weight-mse or local-hessian
-    max_calibrate(model, forward_loop=forward_loop)
+    if fused and any(
+        is_quantized_linear(m)
+        and m.weight_quantizer.is_enabled
+        and any(
+            isinstance(q, TensorQuantizer) and (q.block_sizes or {}).get("four_over_six")
+            for q in m.weight_quantizer.modules()
+        )
+        for m in model.modules()
+    ):
+        raise ValueError(
+            "GPTQ fused=True is incompatible with four_over_six weight block_sizes: the fused "
+            "NVFP4 kernel assumes a 448-normalized global scale, while 4/6 normalizes by 256. "
+            "Use fused=False."
+        )
+
+    if scale_algorithm is None:
+        scale_algorithm = {"method": "max"}
+    _run_scale_calibration(model, forward_loop, scale_algorithm)
 
     quantized_layers = [
         (n, m)

modelopt_recipes/ptq.md

@@ -187,8 +187,8 @@ How the quantization scales are searched. The default (no suffix) is `max`.
   one decoder layer at a time to **lower peak memory**; same numerics as the
   non-layerwise variant.
 
-These can also be **stacked** when a single method isn't enough — e.g. `mse` +
-`gptq` combines an MSE-searched weight scale with GPTQ's layerwise update.
+To combine a searched weight scale with GPTQ, configure the scale search inside GPTQ:
+`algorithm: {method: gptq, layerwise: true, scale_algorithm: {method: mse}}`.
 
 ---
 

tests/gpu/torch/quantization/test_gptq.py

@@ -127,18 +127,15 @@ def test_gptq_export_roundtrip():
     )
 
 
-@pytest.mark.parametrize(
-    "quant_cfg", [mtq.NVFP4_DEFAULT_CFG, mtq.FP8_DEFAULT_CFG, mtq.INT4_BLOCKWISE_WEIGHT_ONLY_CFG]
-)
-def test_gptq_e2e_flow(quant_cfg, tiny_tokenizer):
-    model = get_tiny_llama(vocab_size=tiny_tokenizer.vocab_size).to("cuda")
+def _run_gptq_e2e(quant_cfg, algorithm, tokenizer):
+    model = get_tiny_llama(vocab_size=tokenizer.vocab_size).to("cuda")
     model.eval()
 
     quant_cfg = copy.deepcopy(quant_cfg)
-    quant_cfg["algorithm"] = {"method": "gptq", "layerwise": True}
+    quant_cfg["algorithm"] = algorithm
     calib_dataloader = get_dataset_dataloader(
         dataset_name="cnn_dailymail",
-        tokenizer=tiny_tokenizer,
+        tokenizer=tokenizer,
         batch_size=2,
         num_samples=8,
         device="cuda",
@@ -147,6 +144,49 @@ def test_gptq_e2e_flow(quant_cfg, tiny_tokenizer):
 
     calibrate_loop = create_forward_loop(dataloader=calib_dataloader)
     model = mtq.quantize(model, quant_cfg, forward_loop=calibrate_loop)
+    return model
+
+
+@pytest.mark.parametrize(
+    "quant_cfg", [mtq.NVFP4_DEFAULT_CFG, mtq.FP8_DEFAULT_CFG, mtq.INT4_BLOCKWISE_WEIGHT_ONLY_CFG]
+)
+def test_gptq_e2e_flow(quant_cfg, tiny_tokenizer):
+    _run_gptq_e2e(quant_cfg, {"method": "gptq", "layerwise": True}, tiny_tokenizer)
+
+
+@pytest.mark.parametrize(
+    ("quant_cfg", "scale_algorithm"),
+    [
+        (mtq.NVFP4_DEFAULT_CFG, {"method": "mse", "fp8_scale_sweep": True}),
+        (mtq.NVFP4_DEFAULT_CFG, {"method": "local_hessian"}),
+        (
+            mtq.NVFP4_DEFAULT_CFG,
+            {"method": "local_hessian", "activation_error_coupling": True},
+        ),
+        (
+            mtq.NVFP4_FOUR_OVER_SIX_CFG,
+            {
+                "method": "mse",
+                "fp8_scale_sweep": False,
+                "start_multiplier": 1.0,
+                "stop_multiplier": 1.5,
+                "step_size": 0.5,
+            },
+        ),
+    ],
+    ids=["mse", "local_hessian", "local_hessian_coupled", "four_over_six"],
+)
+def test_gptq_e2e_scale_algorithms(quant_cfg, scale_algorithm, tiny_tokenizer):
+    model = _run_gptq_e2e(
+        quant_cfg,
+        {"method": "gptq", "layerwise": True, "scale_algorithm": scale_algorithm},
+        tiny_tokenizer,
+    )
+
+    assert any(
+        hasattr(module, "weight_quantizer") and module.weight_quantizer.amax is not None
+        for module in model.modules()
+    )
 
 
 # ---------------------------------------------------------------------------

tests/unit/torch/quantization/test_gptq.py

@@ -15,11 +15,141 @@
 
 """CPU unit tests for GPTQ utilities."""
 
-import torch
+from unittest.mock import Mock, create_autospec
 
+import pytest
+import torch
+from torch import nn
+
+import modelopt.torch.quantization as mtq
+import modelopt.torch.quantization.model_calib as model_calib_module
+from modelopt.torch.quantization.config import (
+    GPTQCalibConfig,
+    LocalHessianCalibConfig,
+    MaxCalibConfig,
+    MseCalibConfig,
+)
+from modelopt.torch.quantization.model_calib import gptq
 from modelopt.torch.quantization.utils.calib_utils import update_hessian
 
 
+class TestGPTQScaleAlgorithmConfig:
+    """Tests for GPTQ scale-calibration configuration."""
+
+    def test_default_config(self):
+        cfg = GPTQCalibConfig()
+        assert cfg.scale_algorithm is None
+        assert cfg.model_dump()["scale_algorithm"] is None
+        assert cfg.perc_damp == 0.01
+        assert cfg.block_size == 128
+        assert cfg.fused is False
+
+    def test_old_style_config_dict_restores(self):
+        cfg = GPTQCalibConfig(method="gptq", perc_damp=0.02)
+        assert cfg.scale_algorithm is None
+
+    @pytest.mark.parametrize(
+        ("method", "config_type"),
+        [
+            ("max", MaxCalibConfig),
+            ("mse", MseCalibConfig),
+            ("local_hessian", LocalHessianCalibConfig),
+        ],
+    )
+    def test_scale_algorithm(self, method, config_type):
+        cfg = GPTQCalibConfig(scale_algorithm={"method": method})
+        assert isinstance(cfg.scale_algorithm, config_type)
+
+    def test_unsupported_scale_algorithm(self):
+        with pytest.raises(ValueError):
+            GPTQCalibConfig(scale_algorithm={"method": "smoothquant"})
+
+    def test_local_hessian_activation_error_coupling(self):
+        cfg = GPTQCalibConfig(
+            scale_algorithm={"method": "local_hessian", "activation_error_coupling": True}
+        )
+        assert cfg.model_dump()["scale_algorithm"]["activation_error_coupling"] is True
+
+    def test_scale_algorithm_preserves_sparse_dict(self, monkeypatch):
+        cfg = GPTQCalibConfig(scale_algorithm={"method": "mse", "fp8_scale_sweep": True})
+        assert cfg.model_dump()["scale_algorithm"] == {
+            "method": "mse",
+            "fp8_scale_sweep": True,
+        }
+
+        calibrate = create_autospec(model_calib_module.mse_calibrate)
+        monkeypatch.setattr(model_calib_module, "mse_calibrate", calibrate)
+        model = Mock()
+        model_calib_module._run_scale_calibration(model, None, cfg.scale_algorithm)
+        calibrate.assert_called_once_with(model, forward_loop=None, fp8_scale_sweep=True)
+
+
+@pytest.mark.parametrize(
+    ("scale_algorithm", "expected_func", "expected_kwargs"),
+    [
+        (None, "max", {}),
+        ({"method": "mse", "fp8_scale_sweep": True}, "mse", {"fp8_scale_sweep": True}),
+        (
+            {"method": "local_hessian", "activation_error_coupling": True},
+            "local_hessian",
+            {"activation_error_coupling": True},
+        ),
+    ],
+)
+def test_gptq_scale_algorithm_dispatch(
+    monkeypatch, scale_algorithm, expected_func, expected_kwargs
+):
+    calls = []
+
+    def record(name):
+        def calibrate(model, forward_loop, **kwargs):
+            calls.append((name, model, forward_loop, kwargs))
+
+        return calibrate
+
+    monkeypatch.setattr(model_calib_module, "max_calibrate", record("max"))
+    monkeypatch.setattr(model_calib_module, "mse_calibrate", record("mse"))
+    monkeypatch.setattr(model_calib_module, "local_hessian_calibrate", record("local_hessian"))
+    model = nn.Linear(4, 4)
+
+    def forward_loop(model):
+        pass
+
+    gptq(model, forward_loop=forward_loop, scale_algorithm=scale_algorithm)
+
+    assert calls == [(expected_func, model, forward_loop, expected_kwargs)]
+
+
+@pytest.mark.parametrize("sequential", [False, True])
+def test_gptq_fused_rejects_four_over_six(sequential):
+    model = nn.Linear(32, 32)
+    weight_quantizer_cfg = {
+        "num_bits": (2, 1),
+        "block_sizes": {
+            -1: 16,
+            "type": "static",
+            "scale_bits": (4, 3),
+            "four_over_six": True,
+        },
+    }
+    if sequential:
+        weight_quantizer_cfg = [weight_quantizer_cfg, {"num_bits": (4, 3), "axis": None}]
+    quant_cfg = {
+        "quant_cfg": [
+            {"quantizer_name": "*", "enable": False},
+            {
+                "quantizer_name": "*weight_quantizer",
+                "cfg": weight_quantizer_cfg,
+            },
+        ],
+        "algorithm": None,
+    }
+    mtq.quantize(model, quant_cfg)
+
+    with pytest.raises(ValueError, match="four_over_six"):
+        gptq(model, forward_loop=lambda m: None, fused=True)
+
+
 def test_update_hessian():
     """Test for update_hessian function with both random and known inputs."""
     # Test 1: Random input - general functionality test