diff --git a/CHANGELOG.rst b/CHANGELOG.rst index 8e6b14e2468..1028ef6363c 100755 --- a/CHANGELOG.rst +++ b/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 `_. - 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 `_) 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). diff --git a/modelopt/torch/quantization/config.py b/modelopt/torch/quantization/config.py index c8eb109ca0c..c1d9d6b4a2d 100644 --- a/modelopt/torch/quantization/config.py +++ b/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): diff --git a/modelopt/torch/quantization/model_calib.py b/modelopt/torch/quantization/model_calib.py index 21fd9429d14..73036889275 100644 --- a/modelopt/torch/quantization/model_calib.py +++ b/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) diff --git a/modelopt_recipes/ptq.md b/modelopt_recipes/ptq.md index 78d44ecbc5b..36135cabe66 100644 --- a/modelopt_recipes/ptq.md +++ b/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}}`. --- diff --git a/tests/gpu/torch/quantization/test_gptq.py b/tests/gpu/torch/quantization/test_gptq.py index 0a1849544b8..2b64b6becfb 100644 --- a/tests/gpu/torch/quantization/test_gptq.py +++ b/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() + ) # --------------------------------------------------------------------------- diff --git a/tests/unit/torch/quantization/test_gptq.py b/tests/unit/torch/quantization/test_gptq.py index 345bd167112..8f550ddd049 100644 --- a/tests/unit/torch/quantization/test_gptq.py +++ b/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