diff --git a/CHANGELOG.rst b/CHANGELOG.rst index 89fb2a8b8d7..d3909219044 100755 --- a/CHANGELOG.rst +++ b/CHANGELOG.rst @@ -24,6 +24,7 @@ Changelog - Speed up NVFP4 QAT and QAD of Hugging Face fused MoE experts: all experts of a projection are now fake-quantized in one kernel launch instead of one per expert, about 2x faster per MoE layer on Qwen3-30B-A3B with FSDP2. - ``hf_ptq.py`` calibrates Qwen3.8 (``qwen4_exp``) without running its QSA indexer's per-query Python loop when every visible key fits the indexer's token budget, with bit-identical results. On Qwen3.8-Flash-Next this saves about two hours of layerwise calibration. - Speed up Hugging Face QAT and QAD with ``QATTrainer``, ``QADTrainer`` or ``--manual_gc``: the ``gc.collect()`` run before each training step now skips the model and optimizer state, saving about 1 s per step (30%) on Qwen3-30B-A3B. +- GPTQ now defaults to ``fused=True``, running the fused Triton update for static NVFP4 weights on SM89+ (4/6 scaling and other formats keep the unfused loop). Set ``fused: false`` to restore the previous behavior. *Speculative Decoding* diff --git a/modelopt/torch/kernels/quantization/gemm/gptq_fused_kernel.py b/modelopt/torch/kernels/quantization/gemm/gptq_fused_kernel.py index 3b410e97bae..e5b8721d7f6 100644 --- a/modelopt/torch/kernels/quantization/gemm/gptq_fused_kernel.py +++ b/modelopt/torch/kernels/quantization/gemm/gptq_fused_kernel.py @@ -15,22 +15,22 @@ """Fused Triton kernels for GPTQ blockwise weight-update. -A kernel for scalar (NVFP4) quantization with inline two-level scale computation. -Fuses scale computation + quantization + per-column GPTQ error propagation into -one launch per GPTQ block, avoiding the Python-level per-column loop. +A kernel for scalar (NVFP4) quantization with precomputed FP8 block scales. +Fuses quantization + per-column GPTQ error propagation into one launch per GPTQ +block, avoiding the Python-level per-column loop. Architecture: - One Triton program per output row. - ``w_full [BLOCK_SIZE]`` register tensor holds working weights. - - Per-column: calls ``nvfp4_scalar_qdq()`` for FP4 QDQ with inline scale - computation, then propagates error via ``w_full -= err * h_inv_row``. + - Per-column: calls ``nvfp4_scalar_quant()`` for FP4 QDQ with the column's block + scale, then propagates error via ``w_full -= err * h_inv_row``. """ import torch import triton import triton.language as tl -from ..common.nvfp4_quant import nvfp4_scalar_qdq +from ..common.nvfp4_quant import nvfp4_scalar_quant __all__ = ["gptq_fused_block_scalar"] @@ -45,68 +45,62 @@ def _gptq_scalar_kernel( w_ptr, qw_ptr, err_ptr, - amax_ptr, - global_scale, + scale_ptr, hinv_ptr, num_rows, - n_amax_blocks, + n_scale_blocks, quant_block_size, block_start, + n_cols, BLOCK_SIZE: tl.constexpr, ): row = tl.program_id(0) if row >= num_rows: return - w_base = w_ptr + row * BLOCK_SIZE - qw_base = qw_ptr + row * BLOCK_SIZE - err_base = err_ptr + row * BLOCK_SIZE - amax_base = amax_ptr + row * n_amax_blocks + w_base = w_ptr + row * n_cols + qw_base = qw_ptr + row * n_cols + err_base = err_ptr + row * n_cols + scale_base = scale_ptr + row * n_scale_blocks + # BLOCK_SIZE is n_cols rounded up to a power of two, as tl.arange requires. j_range = tl.arange(0, BLOCK_SIZE) - w_full = tl.load(w_base + j_range) + in_block = j_range < n_cols + w_full = tl.load(w_base + j_range, mask=in_block, other=0.0) - for col in range(0, BLOCK_SIZE, 1): - block_amax = tl.load(amax_base + (block_start + col) // quant_block_size) + for col in range(0, n_cols, 1): + block_scale = tl.load(scale_base + (block_start + col) // quant_block_size) w_scalar = tl.sum(tl.where(j_range == col, w_full, 0.0)) q_scalar = tl.sum( - nvfp4_scalar_qdq( - tl.full([1], w_scalar, dtype=tl.float32), - block_amax, - global_scale, - 1, - ) + nvfp4_scalar_quant(tl.full([1], w_scalar, dtype=tl.float32), block_scale, 1) ) - d_val = tl.load(hinv_ptr + col * BLOCK_SIZE + col) + d_val = tl.load(hinv_ptr + col * n_cols + col) err_val = (w_scalar - q_scalar) / d_val tl.store(err_base + col, err_val) tl.store(qw_base + col, q_scalar) - remaining = j_range > col - hinv_row = tl.load(hinv_ptr + col * BLOCK_SIZE + j_range, mask=remaining, other=0.0) + remaining = (j_range > col) & in_block + hinv_row = tl.load(hinv_ptr + col * n_cols + j_range, mask=remaining, other=0.0) w_full = w_full - err_val * hinv_row def gptq_fused_block_scalar( w_block: torch.Tensor, - block_amax: torch.Tensor, - global_scale: float, + block_scale: torch.Tensor, h_inv_cho_blk: torch.Tensor, quant_block_size: int, block_start: int, ) -> tuple[torch.Tensor, torch.Tensor]: """Run scalar GPTQ (NVFP4) column loop for one block in a single Triton kernel launch. - Computes FP8-quantized scales from per-block amax inline via - :func:`nvfp4_scalar_qdq`, then performs NVFP4 fake quantization and + Performs NVFP4 fake quantization with the given FP8-quantized block scales and GPTQ error propagation per column. Args: w_block: Working weights ``[num_rows, block_size]`` (float32). - block_amax: Per-block amax values ``[num_rows, n_amax_blocks]`` (float32). - global_scale: Pre-computed ``global_amax / (6.0 * 448.0)`` (scalar). + block_scale: FP8-quantized per-block scales ``[num_rows, n_scale_blocks]`` (float32). h_inv_cho_blk: Block of upper-Cholesky inverse Hessian ``[block_size, block_size]``. quant_block_size: Number of elements sharing one scale factor. block_start: Column offset of this block in the full weight matrix. @@ -126,14 +120,14 @@ def gptq_fused_block_scalar( w_block.contiguous(), qw_block, err_block, - block_amax.contiguous(), - global_scale, + block_scale.contiguous(), h_inv_cho_blk.contiguous(), num_rows, - block_amax.shape[1], + block_scale.shape[1], quant_block_size, block_start, - BLOCK_SIZE=block_size, + block_size, + BLOCK_SIZE=triton.next_power_of_2(block_size), ) return qw_block, err_block diff --git a/modelopt/torch/quantization/config.py b/modelopt/torch/quantization/config.py index e5700be1a7f..6f097c3efb0 100644 --- a/modelopt/torch/quantization/config.py +++ b/modelopt/torch/quantization/config.py @@ -1266,10 +1266,11 @@ class GPTQCalibConfig(_HessianInputDtypeConfig, QuantizeAlgorithmConfig): of the group_size used in the quantization.""", ) fused: bool = ModeloptField( - default=False, + default=True, title="Use fused Triton kernel for GPTQ.", description="""When True, use a fused Triton kernel that combines quantization and - per-column error propagation into one launch per GPTQ block.""", + per-column error propagation into one launch per GPTQ block. Applies to static NVFP4 + weights without 4/6 scaling on SM89+; other cases use the unfused loop.""", ) importance_weighted: bool = ModeloptField( default=False, diff --git a/modelopt/torch/quantization/model_calib.py b/modelopt/torch/quantization/model_calib.py index 7397d79eb56..e6b60e02dca 100644 --- a/modelopt/torch/quantization/model_calib.py +++ b/modelopt/torch/quantization/model_calib.py @@ -2520,7 +2520,7 @@ def gptq( forward_loop: ForwardLoop, perc_damp: float = 0.01, block_size: int = 128, - fused: bool = False, + fused: bool = True, importance_weighted: bool = False, hessian_input_dtype: str = "float32", ): diff --git a/modelopt/torch/quantization/utils/calib_utils.py b/modelopt/torch/quantization/utils/calib_utils.py index 33f075a4e7d..691494279c4 100644 --- a/modelopt/torch/quantization/utils/calib_utils.py +++ b/modelopt/torch/quantization/utils/calib_utils.py @@ -46,6 +46,7 @@ import torch import torch.distributed as dist +import modelopt.torch.kernels.quantization.gemm as triton_kernel from modelopt.torch.utils import ( get_available_gpu_memory, get_process_gpus, @@ -54,6 +55,8 @@ ) from modelopt.torch.utils.network import bind_forward_method, unpatch_forward_method +from .numeric_utils import E4M3_MAX, fp8_max_for_normalization + # A GPU takes a new Hessian only if enough memory stays available after it for the calibration # forward that follows. Until GPTQ has measured that forward on the GPU, keep this share free. _HESSIAN_UNMEASURED_RESERVE = 0.35 @@ -490,9 +493,8 @@ def _prepare_hessian_inverse(self, hessian, perc_damp): def _blockwise_update(self, block_size): """Column-wise GPTQ update. - When ``self.fused`` is True and the weight quantizer is an - ``NVFP4StaticQuantizer``, uses :func:`gptq_blockwise_update_fused_scalar` - (a fused Triton kernel). Otherwise falls back to + When ``self.fused`` is True and :meth:`_fused_supported` holds, uses + :func:`gptq_blockwise_update_fused_scalar` (a fused Triton kernel). Otherwise falls back to :func:`gptq_blockwise_update` (unfused column-by-column loop). """ assert self.weight is not None and self.h_inv is not None, ( @@ -500,24 +502,43 @@ def _blockwise_update(self, block_size): ) quantizer = self.weight_quantizer - if self.fused and getattr(quantizer, "_is_nvfp4_static_quantizer", False): + if self.fused and self._fused_supported(quantizer, self.weight): block_sizes = quantizer.block_sizes + # The kernel looks up each column's scale by absolute index, so GPTQ blocks need not + # align with quantizer groups. quant_block_size = block_sizes.get(-1) or block_sizes.get(1) - if quant_block_size is not None and block_size % quant_block_size != 0: - raise ValueError( - f"GPTQ block_size ({block_size}) must be divisible by the quantizer" - f" group_size ({quant_block_size})" - ) out_features, num_cols = self.weight.shape - n_blocks = num_cols // quant_block_size - block_amax = quantizer.amax.reshape(out_features, n_blocks).float() - global_scale = quantizer.global_amax.float().item() / (6.0 * 448.0) + # A width the group size does not divide carries a scale for its padded last group. + n_blocks = -(-num_cols // quant_block_size) + # The quantizer's own FP8 scales, so the fused update rounds onto the same grid. + block_scale = triton_kernel.compute_fp4_scales(quantizer.amax, quantizer.global_amax) gptq_blockwise_update_fused_scalar( - self.weight, block_amax, global_scale, self.h_inv, block_size, quant_block_size + self.weight, + block_scale.reshape(out_features, n_blocks), + self.h_inv, + block_size, + quant_block_size, ) else: gptq_blockwise_update(self.weight, self.h_inv, block_size, quantizer) + def _fused_supported(self, quantizer, weight): + # The fused kernel stands in for the quantizer's forward with a plain static QDQ over + # last-dim groups, so anything else that forward does (pre-quant scale, rotation, LSQ, + # 2-D blocks) takes the unfused loop. It also needs FP8 Triton casts (SM89+) and the + # default FP8 scale range (no 4/6). + return ( + getattr(quantizer, "_is_nvfp4_static_quantizer", False) + and {k for k in quantizer.block_sizes if isinstance(k, int)} in ({-1}, {1}) + and quantizer.pre_quant_scale is None + and not quantizer.rotate_is_enabled + and not quantizer._lsq + and fp8_max_for_normalization(quantizer) == E4M3_MAX + and triton_kernel.IS_AVAILABLE + and weight.is_cuda + and torch.cuda.get_device_capability(weight.device) >= (8, 9) + ) + def _print_mse_error(self, hessian, weight): """Log Hessian-weighted relative MSE between ``self.weight`` and original ``weight``.""" w_orig = weight.float() @@ -565,20 +586,19 @@ def gptq_blockwise_update(weight, h_inv, block_size, quantize_fn): weight[:, block_end:].addmm_(errs, h_inv[block_start:block_end, block_end:], alpha=-1) -def gptq_blockwise_update_fused_scalar( - weight, block_amax, global_scale, h_inv, block_size, quant_block_size -): +def gptq_blockwise_update_fused_scalar(weight, block_scale, h_inv, block_size, quant_block_size): """Fused GPTQ blockwise update for NVFP4 scalar quantization. - Uses a fused Triton kernel that combines scale computation, quantization, - and per-column error propagation into one launch per GPTQ block, avoiding - the Python-level per-column loop in :func:`gptq_blockwise_update`. + Uses a fused Triton kernel that combines quantization and per-column error propagation + into one launch per GPTQ block, avoiding the Python-level per-column loop in + :func:`gptq_blockwise_update`. Args: weight: Weight tensor ``[out_features, in_features]``, modified **in-place** with fake-quantized values. - block_amax: Per-block amax values ``[out_features, n_amax_blocks]``. - global_scale: Pre-computed ``global_amax / (6.0 * 448.0)`` (scalar). + block_scale: FP8-quantized per-block scales ``[out_features, n_scale_blocks]``, as + :func:`compute_fp4_scales ` + returns them. h_inv: Upper-triangular Cholesky factor of the damped inverse Hessian ``[in_features, in_features]``. block_size: Number of columns to process per GPTQ block. @@ -591,8 +611,7 @@ def gptq_blockwise_update_fused_scalar( be = min(bs + block_size, num_cols) qw, err = gptq_fused_block_scalar( weight[:, bs:be].clone().contiguous(), - block_amax, - global_scale, + block_scale, h_inv[bs:be, bs:be].contiguous(), quant_block_size, bs, diff --git a/tests/gpu/torch/quantization/test_gptq.py b/tests/gpu/torch/quantization/test_gptq.py index a7f441b243a..408dba33f23 100644 --- a/tests/gpu/torch/quantization/test_gptq.py +++ b/tests/gpu/torch/quantization/test_gptq.py @@ -24,6 +24,7 @@ import modelopt.torch.quantization as mtq from modelopt.torch.export.unified_export_hf import _export_quantized_weight +from modelopt.torch.kernels.quantization.gemm import compute_fp4_scales from modelopt.torch.quantization.extensions import get_cuda_ext_ggml from modelopt.torch.quantization.ggml.common import pinned_packed_weight from modelopt.torch.quantization.ggml.gptq import GGMLGPTQHelper @@ -408,7 +409,7 @@ def test_gptq_e2e_flow(quant_cfg, tiny_tokenizer, calib_jsonl): def _make_nvfp4_test_data(quant_block_size, out_features, dim): - """Create weight, weight_quantizer, block_amax, global_scale, and h_inv for NVFP4 GPTQ tests.""" + """Create weight, weight_quantizer, block_scale, and h_inv for NVFP4 GPTQ tests.""" # Build a quantized Linear with NVFP4 static config at the desired block size model = torch.nn.Linear(dim, out_features, bias=False, device="cuda") weight = model.weight.data.clone() @@ -431,8 +432,9 @@ def _make_nvfp4_test_data(quant_block_size, out_features, dim): model.weight.data = weight.clone() weight_quantizer = model.weight_quantizer - block_amax = weight_quantizer.amax.reshape(out_features, -1).float() - global_scale = weight_quantizer.global_amax.float().item() / (6.0 * 448.0) + block_scale = compute_fp4_scales(weight_quantizer.amax, weight_quantizer.global_amax).reshape( + out_features, -1 + ) # Compute Hessian hessian = torch.zeros(dim, dim, dtype=torch.float32) @@ -440,7 +442,7 @@ def _make_nvfp4_test_data(quant_block_size, out_features, dim): hessian = hessian.to("cuda") h_inv = compute_hessian_inverse(hessian, weight, perc_damp=0.01) - return weight, weight_quantizer, block_amax, global_scale, h_inv + return weight, weight_quantizer, block_scale, h_inv def _run_unfused_gptq_nvfp4(weight, weight_quantizer, h_inv, gptq_block_size): @@ -450,14 +452,10 @@ def _run_unfused_gptq_nvfp4(weight, weight_quantizer, h_inv, gptq_block_size): return w -def _run_fused_gptq_nvfp4( - weight, block_amax, global_scale, h_inv, gptq_block_size, quant_block_size -): +def _run_fused_gptq_nvfp4(weight, block_scale, h_inv, gptq_block_size, quant_block_size): """Fused Triton GPTQ for NVFP4 using the production fused update.""" w = weight.float().clone() - gptq_blockwise_update_fused_scalar( - w, block_amax, global_scale, h_inv, gptq_block_size, quant_block_size - ) + gptq_blockwise_update_fused_scalar(w, block_scale, h_inv, gptq_block_size, quant_block_size) return w @@ -465,11 +463,11 @@ def _run_fused_gptq_nvfp4( def test_fused_nvfp4_gptq_runs_on_the_weights_gpu(need_2_gpus): """A layer device_map placed on another GPU is updated there, not from the current GPU.""" torch.manual_seed(42) - weight, _, block_amax, global_scale, h_inv = _make_nvfp4_test_data(16, 64, 256) - expected = _run_fused_gptq_nvfp4(weight, block_amax, global_scale, h_inv, 128, 16) + weight, _, block_scale, h_inv = _make_nvfp4_test_data(16, 64, 256) + expected = _run_fused_gptq_nvfp4(weight, block_scale, h_inv, 128, 16) - weight, block_amax, h_inv = (t.to("cuda:1") for t in (weight, block_amax, h_inv)) - got = _run_fused_gptq_nvfp4(weight, block_amax, global_scale, h_inv, 128, 16) + weight, block_scale, h_inv = (t.to("cuda:1") for t in (weight, block_scale, h_inv)) + got = _run_fused_gptq_nvfp4(weight, block_scale, h_inv, 128, 16) assert torch.cuda.current_device() == 0 torch.testing.assert_close(got.cpu(), expected.cpu(), rtol=0, atol=0) @@ -488,7 +486,7 @@ def test_fused_vs_unfused_nvfp4(quant_block_size, gptq_block_size): dim = max(256, quant_block_size * 4) out_features = 64 - weight, weight_quantizer, block_amax, global_scale, h_inv = _make_nvfp4_test_data( + weight, weight_quantizer, block_scale, h_inv = _make_nvfp4_test_data( quant_block_size, out_features, dim, @@ -496,8 +494,7 @@ def test_fused_vs_unfused_nvfp4(quant_block_size, gptq_block_size): weight_fused = _run_fused_gptq_nvfp4( weight, - block_amax, - global_scale, + block_scale, h_inv, gptq_block_size, quant_block_size, @@ -526,6 +523,56 @@ def test_fused_vs_unfused_nvfp4(quant_block_size, gptq_block_size): torch.testing.assert_close(weight_fused, weight_unfused, atol=1e-4, rtol=1e-4) +@requires_triton +@pytest.mark.parametrize( + ("in_features", "block_size", "cfg_update", "takes_fused"), + [ + (96, 48, {}, True), # non-power-of-two GPTQ block + (208, 128, {}, True), # non-power-of-two final block + (70, 128, {}, True), # width the NVFP4 group size does not divide + (96, 40, {}, True), # GPTQ block not a multiple of the group size + # The quantizer forward does more than a per-row static QDQ, so GPTQ stays unfused. + (64, 128, {"block_sizes": {"four_over_six": True}}, False), + (64, 128, {"block_sizes": {-2: 16}}, False), + (64, 128, {"rotate": {"enable": True, "mode": "rotate_back"}}, False), + ], +) +def test_gptq_fused_matches_unfused(in_features, block_size, cfg_update, takes_fused, monkeypatch): + if torch.cuda.get_device_capability() < (8, 9): + pytest.skip("fused GPTQ needs SM89+") + if "rotate" in cfg_update: + pytest.importorskip("fast_hadamard_transform") + fused_calls = [] + fused_update = calib_utils.gptq_blockwise_update_fused_scalar + monkeypatch.setattr( + calib_utils, + "gptq_blockwise_update_fused_scalar", + lambda *args: fused_calls.append(1) or fused_update(*args), + ) + torch.manual_seed(0) + model = torch.nn.Linear(in_features, 32, bias=False, device="cuda") + inp = torch.randn(4, 16, in_features, device="cuda") + block_sizes = {-1: 16, "type": "static", "scale_bits": (4, 3)} + weight_cfg = { + "num_bits": (2, 1), + **cfg_update, + "block_sizes": {**block_sizes, **cfg_update.get("block_sizes", {})}, + } + weights = [] + for fused in (True, False): + quant_cfg = { + "quant_cfg": [ + {"quantizer_name": "*", "enable": False}, + {"quantizer_name": "*weight_quantizer", "cfg": weight_cfg}, + ], + "algorithm": {"method": "gptq", "block_size": block_size, "fused": fused}, + } + qmodel = mtq.quantize(copy.deepcopy(model), quant_cfg, forward_loop=lambda m: m(inp)) + weights.append(qmodel.weight.data) + assert len(fused_calls) == int(takes_fused) + torch.testing.assert_close(weights[0], weights[1], atol=1e-4, rtol=1e-4) + + _NVFP4_BENCH_CONFIGS = [ (16, 128, 256, 512), (16, 128, 256, 2048), @@ -557,13 +604,13 @@ def _bench(fn, n_warmup=2, n_iters=5): for quant_block_size, gptq_block_size, out_features, dim in _NVFP4_BENCH_CONFIGS: torch.manual_seed(42) - weight, weight_quantizer, block_amax, global_scale, h_inv = _make_nvfp4_test_data( + weight, weight_quantizer, block_scale, h_inv = _make_nvfp4_test_data( quant_block_size, out_features, dim ) def run_fused(): return _run_fused_gptq_nvfp4( - weight, block_amax, global_scale, h_inv, gptq_block_size, quant_block_size + weight, block_scale, h_inv, gptq_block_size, quant_block_size ) def run_unfused():