From 0101143bf6f1c3e0b15ef9a98ed0bf1e7c06550c Mon Sep 17 00:00:00 2001 From: Chenjie Luo Date: Sat, 17 Oct 2026 16:19:38 +0000 Subject: [PATCH 1/5] Speed up GPTQ calibration on MoE models GPTQ ran a max-calibration forward before its own Hessian forward even when no activation quantizer needed it; weight amax comes from the weights. Skip it, as max and imatrix already do. Fold the Hessian's incremental average into one addmm, which runs once per expert per forward on HF fused-MoE layers. A Hessian whose GPU was over 65% used went to the CPU, so every update copied the whole matrix to the host and rescaled it there. A model sharded over several GPUs fills the one holding the current layer while others sit idle; a single process now puts the Hessian on the emptiest GPU it already uses that is under the limit, and the CPU only when none is. hf_ptq's layerwise default batch size becomes a token budget (8192 tokens, 16 x 512) instead of 1: a batch holds one layer's activations, and batch 1 made each MoE layer run its per-expert loop once per calibration sample. Co-Authored-By: Claude Opus 5.5 (1M context) Signed-off-by: Chenjie Luo --- CHANGELOG.rst | 1 + .../announcements/20260928-single-gpu-ptq.rst | 7 +-- examples/hf_ptq/hf_ptq.py | 15 ++++-- modelopt/torch/quantization/model_calib.py | 16 +++--- .../torch/quantization/utils/calib_utils.py | 52 ++++++++++++++----- tests/gpu/torch/quantization/test_gptq.py | 20 +++++++ tests/unit/torch/quantization/test_gptq.py | 29 +++++++++++ 7 files changed, 114 insertions(+), 26 deletions(-) diff --git a/CHANGELOG.rst b/CHANGELOG.rst index 1fc54679a9a..7ac106c7ee7 100755 --- a/CHANGELOG.rst +++ b/CHANGELOG.rst @@ -11,6 +11,7 @@ Changelog - On Linux x86-64 with Python 3.11 and later, the ONNX extra now selects ONNX Runtime 1.30.x and CUDA 13 packages by default. GPU calibration requires a compatible CUDA 13 runtime and driver; Python 3.10 and other platforms retain their existing dependency selections. - Add imatrix calibration (``algorithm: {method: imatrix}``) for the GGML formats with a weighted search: each weight is encoded once, its search weighted by the square root of its inputs' per-column ``E[x^2]`` (llama.cpp's imatrix statistic), and HF fused-MoE experts get one statistic per expert. A weight that already holds final codes is left alone, so ``[{method: gptq}, {method: imatrix}]`` runs GPTQ where it applies and imatrix on the rest; the built-in IQ1 and IQ2 recipes now do this, and new ``_experts_only`` recipes apply imatrix to the routed MoE experts alone. - GPTQ now also quantizes HuggingFace fused MoE experts, each with its own Hessian from the tokens routed to it; they were previously skipped. The built-in IQ1 and IQ2 recipes therefore run GPTQ on the experts too, and large MoEs need layerwise calibration to hold one layer's expert Hessians at a time. +- GPTQ no longer runs a max-calibration forward when no activation quantizer needs one, and ``hf_ptq.py`` now sizes the default layerwise calibration batch at 8,192 tokens (16 samples at ``--calib_seq 512``) instead of 1. Layerwise GPTQ on HF fused-MoE models runs over 10x faster; pass ``--batch_size 1`` for the old batching, which pads no sample. - Add a checkpoint-specific NVIDIA Nemotron 3.5 Super VL AutoQuantize recipe with a 4.9-effective-bit NVFP4-MSE/FP8/BF16 search and fixed NVFP4 MTP experts. Select it under ``models/nvidia/NVIDIA-Nemotron-3.5-Super-VL-120B-A12B-BF16/auto_quantize/``. - Add ``examples/hf_ptq/convert_ggml_to_int8.py`` to expand a checkpoint exported with a GGML recipe (IQ formats, ``q8_0``) into an INT8 weight-only compressed-tensors checkpoint. Stock vLLM serves it with its integer kernels at about half the BF16 weight memory, reproducing the GGML weights up to rounding their scales to the model dtype. diff --git a/docs/source/announcements/20260928-single-gpu-ptq.rst b/docs/source/announcements/20260928-single-gpu-ptq.rst index 9ea16879777..fd367eed94b 100644 --- a/docs/source/announcements/20260928-single-gpu-ptq.rst +++ b/docs/source/announcements/20260928-single-gpu-ptq.rst @@ -133,9 +133,10 @@ at sequence length 512, **batch size 8**. .. note:: - **Pick the batch size on purpose.** For layerwise runs, ``hf_ptq.py`` defaults to - ``--batch_size 1``, but only if you leave the flag at its auto-probe default. An explicit - value passes straight through. Going to 8 cuts wall clock roughly in half. It does move the + **Pick the batch size on purpose.** For layerwise runs, ``hf_ptq.py``'s auto-probe default + picks 8,192 tokens per batch (16 samples at sequence length 512); it picked 1 when this run + was made, which is why it passed ``--batch_size 8``. An explicit value passes straight + through. Going from 1 to 8 cut wall clock roughly in half. It does move the numbers a little. Weights, ``weight_scale`` and ``weight_scale_2`` come out bit-identical, but activation ``input_scale`` values shift slightly, because batching pads each sequence to the longest one in the batch and those pad positions feed the activation amax. diff --git a/examples/hf_ptq/hf_ptq.py b/examples/hf_ptq/hf_ptq.py index 68f892f8e9b..a1643b9d531 100755 --- a/examples/hf_ptq/hf_ptq.py +++ b/examples/hf_ptq/hf_ptq.py @@ -114,6 +114,8 @@ from modelopt.torch.utils.vlm_dataset_utils import get_vlm_dataset_dataloader RAND_SEED = 1234 +# Tokens per calibration batch when layerwise calibration picks the batch size (16 x 512). +LAYERWISE_CALIB_TOKENS = 8192 mto.enable_huggingface_checkpointing() @@ -989,10 +991,17 @@ def quantize_main( ) args.batch_size = 1 # Layerwise calibration processes one layer at a time; auto batch-size probing runs a - # full-model forward which defeats the point and can OOM on very large models. + # full-model forward which defeats the point and can OOM on very large models. A batch + # then only holds one layer's activations, so size it by tokens: batch 1 makes an MoE + # layer run its per-expert loop once per sample. elif is_layerwise: - print("Layerwise calibration enabled. Using default batch_size=1 for calibration.") - args.batch_size = 1 + args.batch_size = max( + 1, min(LAYERWISE_CALIB_TOKENS // args.calib_seq, sum(args.calib_size)) + ) + print( + f"Layerwise calibration enabled. Using default batch_size={args.batch_size} " + f"({LAYERWISE_CALIB_TOKENS} tokens per batch) for calibration." + ) else: # Calibration/sparsification will actually take much more memory than regular inference # due to intermediate tensors for fake quantization. Setting sample_memory_usage_ratio diff --git a/modelopt/torch/quantization/model_calib.py b/modelopt/torch/quantization/model_calib.py index 5e7e62c4189..e50f3703f95 100644 --- a/modelopt/torch/quantization/model_calib.py +++ b/modelopt/torch/quantization/model_calib.py @@ -339,8 +339,9 @@ def max_calibrate( skip_forward_without_activation_calib: If True, skip the (potentially expensive) ``forward_loop`` when no enabled quantizer needs data-driven activation statistics (see :func:`_needs_activation_forward_for_max_calib`). Weight calibration still runs. - Only opt-in for the top-level ``max`` path; algorithms that always need activations - (MSE, local Hessian, SmoothQuant, SVDQuant, GPTQ) call this with the default False. + Opted into by the top-level ``max`` path and by GPTQ and imatrix, which run their own + forward; algorithms that always need activations (MSE, local Hessian, SmoothQuant, + SVDQuant) call this with the default False. See :class:`MaxCalibConfig ` for details on the remaining arguments. @@ -2320,12 +2321,14 @@ def gptq( * **Non-layerwise** (``layerwise.enable=False``): called once on the full model. All layers are quantized in parallel from the original activations. - We recommend enabling layerwise calibration - (``"layerwise": {"enable": True}``) and using a calibration batch size of 1. + We recommend enabling layerwise calibration (``"layerwise": {"enable": True}``). The + calibration batch size only trades memory for speed: each forward of an HF fused-MoE layer + loops over its experts, so batch 1 pays that loop once per sample. Per-module steps: - 1. ``max_calibrate`` to set amax values from the current activations. + 1. ``max_calibrate`` to set weight amax, running a forward only when an enabled activation + quantizer needs data-driven amax. 2. Promote eligible quantizers to ``StaticBlockScaleQuantizer`` (two-level scaling). 3. Collect Hessian matrices via forward hooks: one per quantized linear layer, and one per expert and projection of HF fused-MoE experts, over the tokens routed to that expert. @@ -2353,7 +2356,8 @@ def gptq( _refuse_packed_fused_expert_weights(name, module, "GPTQ", lambda q: q.is_enabled) # TODO: Add support for other scale setting strateiges like weight-mse or local-hessian - max_calibrate(model, forward_loop=forward_loop) + # The Hessian pass below is GPTQ's own forward; max only runs one if activations need amax. + max_calibrate(model, forward_loop=forward_loop, skip_forward_without_activation_calib=True) quantized_layers = [ (n, m) diff --git a/modelopt/torch/quantization/utils/calib_utils.py b/modelopt/torch/quantization/utils/calib_utils.py index c0c7b557fc8..fa7dc81c874 100644 --- a/modelopt/torch/quantization/utils/calib_utils.py +++ b/modelopt/torch/quantization/utils/calib_utils.py @@ -38,8 +38,6 @@ """GPTQ helper and Hessian utilities for calibration.""" -import math - import torch import torch.distributed as dist @@ -47,6 +45,29 @@ from modelopt.torch.utils.network import bind_forward_method, unpatch_forward_method from modelopt.torch.utils.perf import get_used_gpu_mem_fraction +# Used fraction of a GPU's memory above which GPTQ keeps new Hessians off it. +_HESSIAN_GPU_MEM_LIMIT = 0.65 + + +def _gpu_is_full(device) -> bool: + return get_used_gpu_mem_fraction(device) > _HESSIAN_GPU_MEM_LIMIT + + +def _hessian_offload_device() -> torch.device: + """Where a Hessian goes when its weight's GPU is full: the emptiest GPU with room, else CPU. + + A model sharded over several GPUs leaves the ones not holding the current layer idle, and a + Hessian there costs a device-to-device copy per update instead of a host copy and a CPU + rescale of the whole matrix. Only GPUs this process already holds memory on are candidates, + and only outside distributed runs, so no other job's or rank's GPU is touched. + """ + if not (dist.is_available() and dist.is_initialized() and dist.get_world_size() > 1): + gpus = [i for i in range(torch.cuda.device_count()) if torch.cuda.memory_reserved(i)] + gpus.sort(key=get_used_gpu_mem_fraction) + if gpus and not _gpu_is_full(gpus[0]): + return torch.device("cuda", gpus[0]) + return torch.device("cpu") + def update_hessian(input, hessian, n_samples): """Update hessian matrix with new input samples using incremental formula. @@ -61,19 +82,22 @@ def update_hessian(input, hessian, n_samples): Note: input must be non-empty (batch_size > 0); a zero-sized input causes division by zero. """ # Flatten to 2D (total_tokens, features) first, so batch_size counts tokens - input_flat = input.reshape(-1, input.shape[-1]).t().float() - batch_size = input_flat.shape[1] + input_flat = input.reshape(-1, input.shape[-1]).float() + batch_size = input_flat.shape[0] if batch_size == 0: # in MOEs some experts receive no tokens return hessian, n_samples - # Incremental averaging: scale down old hessian - hessian *= n_samples / (n_samples + batch_size) + # Incremental averaging, H = H * n / (n + b) + (2 / (n + b)) * X^T X, in one GEMM: this runs + # once per expert per forward, and separate scale and add passes each rewrite all of H. + scale = n_samples / (n_samples + batch_size) n_samples += batch_size - - # Compute outer product: H += (2/n_samples) * X @ X^T - scaled_input = math.sqrt(2 / n_samples) * input_flat - hessian.add_((scaled_input @ scaled_input.t()).to(hessian.device)) + if hessian.device == input_flat.device: + hessian.addmm_(input_flat.t(), input_flat, beta=scale, alpha=2 / n_samples) + else: + hessian.mul_(scale).add_( + (input_flat.t() @ input_flat).to(hessian.device), alpha=2 / n_samples + ) return hessian, n_samples @@ -154,10 +178,10 @@ def __init__( weight = module.weight if weight is None else weight in_features = weight.shape[-1] device = weight.device - if device.type == "meta" or ( - offload_to_cpu and device.type == "cuda" and get_used_gpu_mem_fraction(device) > 0.65 - ): - device = "cpu" + if device.type == "meta": + device = torch.device("cpu") + elif offload_to_cpu and device.type == "cuda" and _gpu_is_full(device): + device = _hessian_offload_device() self.hessian = torch.zeros(in_features, in_features, dtype=torch.float32, device=device) self.n_samples = 0 # Set by update_weights(); listed here for documentation. diff --git a/tests/gpu/torch/quantization/test_gptq.py b/tests/gpu/torch/quantization/test_gptq.py index 5388d8c6517..f4d484c7415 100644 --- a/tests/gpu/torch/quantization/test_gptq.py +++ b/tests/gpu/torch/quantization/test_gptq.py @@ -25,7 +25,9 @@ from modelopt.torch.export.unified_export_hf import _export_quantized_weight from modelopt.torch.quantization.model_calib import gptq from modelopt.torch.quantization.qtensor.nvfp4_tensor import NVFP4QTensor +from modelopt.torch.quantization.utils import calib_utils from modelopt.torch.quantization.utils.calib_utils import ( + GPTQHelper, compute_hessian_inverse, gptq_blockwise_update, gptq_blockwise_update_fused_scalar, @@ -37,6 +39,24 @@ torch.manual_seed(RAND_SEED) +def test_hessian_of_a_full_gpu_goes_to_an_idle_gpu(need_2_gpus, monkeypatch): + """A Hessian whose GPU is full lands on an idle GPU rather than the CPU, with the same values.""" + + def used_fraction(device): + return 0.9 if torch.device(device).index in (0, None) else 0.1 + + monkeypatch.setattr(calib_utils, "get_used_gpu_mem_fraction", used_fraction) + linear = torch.nn.Linear(64, 8, bias=False, device="cuda:0") + shard = torch.zeros(1, device="cuda:1") # the model also lives on cuda:1 + helper = GPTQHelper(linear, "linear", offload_to_cpu=True) + assert helper.hessian.device == shard.device + + inputs = torch.randn(32, 64, device="cuda:0") + helper.accumulate(inputs) + reference, _ = update_hessian(inputs, torch.zeros(64, 64, device="cuda:0"), 0) + torch.testing.assert_close(helper.hessian.to("cuda:0"), reference) + + @pytest.mark.parametrize( ("block_size", "dim", "model_weight", "expect_weight_change"), [ diff --git a/tests/unit/torch/quantization/test_gptq.py b/tests/unit/torch/quantization/test_gptq.py index 1c55c089943..ee1a55a8554 100644 --- a/tests/unit/torch/quantization/test_gptq.py +++ b/tests/unit/torch/quantization/test_gptq.py @@ -17,9 +17,11 @@ from functools import partial +import pytest import torch from _test_utils.torch.distributed.utils import spawn_multiprocess_job +import modelopt.torch.quantization as mtq from modelopt.torch.quantization.utils.calib_utils import GPTQHelper, update_hessian from modelopt.torch.utils.distributed import ParallelState @@ -123,6 +125,33 @@ def test_update_hessian_zero_token_input_noops(): torch.testing.assert_close(hessian, expected_hessian) +@pytest.mark.parametrize(("input_quantized", "forwards"), [(False, 1), (True, 2)]) +def test_gptq_runs_a_max_calibration_forward_only_for_activation_amax(input_quantized, forwards): + """Weight amax comes from the weights; only an input quantizer needs a forward before GPTQ's.""" + quant_cfg = [ + {"quantizer_name": "*", "enable": False}, + {"quantizer_name": "*weight_quantizer", "cfg": {"num_bits": 8, "axis": 0}, "enable": True}, + ] + if input_quantized: + quant_cfg.append( + { + "quantizer_name": "*input_quantizer", + "cfg": {"num_bits": 8, "axis": None}, + "enable": True, + } + ) + calls = [] + + def forward_loop(model): + calls.append(model) + model(torch.randn(4, 32)) + + config = {"quant_cfg": quant_cfg, "algorithm": {"method": "gptq", "block_size": 16}} + mtq.quantize(torch.nn.Linear(32, 8, bias=False), config, forward_loop) + + assert len(calls) == forwards + + def _gptq_hessian_dp_sync_helper(rank, size, token_counts): torch.manual_seed(0) inputs = [torch.randn(n, 8) for n in token_counts] From a1ff4c7d3d19947fa7e742e31b1f162e0c832d64 Mon Sep 17 00:00:00 2001 From: Chenjie Luo Date: Thu, 22 Oct 2026 04:23:28 +0000 Subject: [PATCH 2/5] Address review: Hessian GPU capacity check and layerwise batch tests - A Hessian leaving a full GPU now goes only to a GPU that is under the usage limit and has the Hessian's bytes free, else to the CPU; a GPU that passes the usage limit but is too small would otherwise OOM. - Test the CPU fallback on two GPUs: the idle GPU full, the idle GPU too small, and any rank of a multi-rank run. - Move hf_ptq's default layerwise batch size into a helper and test its token budget, sample cap and long-sequence floor. - Describe that default in the single-GPU PTQ note as a token budget, not a memory probe. Co-Authored-By: Claude Opus 5.5 (1M context) Signed-off-by: Chenjie Luo --- .../announcements/20260928-single-gpu-ptq.rst | 9 ++++---- examples/hf_ptq/hf_ptq.py | 10 ++++++--- .../torch/quantization/utils/calib_utils.py | 15 +++++++------ tests/examples/hf_ptq/test_hf_ptq_args.py | 10 +++++++++ tests/gpu/torch/quantization/test_gptq.py | 21 +++++++++++++++++++ 5 files changed, 52 insertions(+), 13 deletions(-) diff --git a/docs/source/announcements/20260928-single-gpu-ptq.rst b/docs/source/announcements/20260928-single-gpu-ptq.rst index fd367eed94b..b5e4e7fbac2 100644 --- a/docs/source/announcements/20260928-single-gpu-ptq.rst +++ b/docs/source/announcements/20260928-single-gpu-ptq.rst @@ -133,10 +133,11 @@ at sequence length 512, **batch size 8**. .. note:: - **Pick the batch size on purpose.** For layerwise runs, ``hf_ptq.py``'s auto-probe default - picks 8,192 tokens per batch (16 samples at sequence length 512); it picked 1 when this run - was made, which is why it passed ``--batch_size 8``. An explicit value passes straight - through. Going from 1 to 8 cut wall clock roughly in half. It does move the + **Pick the batch size on purpose.** For layerwise runs, ``hf_ptq.py`` does not probe memory: + left at its default, ``--batch_size`` becomes 8,192 tokens' worth of samples (16 at sequence + length 512). It was 1 when this run was made, which is why it passed ``--batch_size 8``. An + explicit value passes straight through; set a smaller one if a layer's activations do not + fit your GPU. Going from 1 to 8 cut wall clock roughly in half. It does move the numbers a little. Weights, ``weight_scale`` and ``weight_scale_2`` come out bit-identical, but activation ``input_scale`` values shift slightly, because batching pads each sequence to the longest one in the batch and those pad positions feed the activation amax. diff --git a/examples/hf_ptq/hf_ptq.py b/examples/hf_ptq/hf_ptq.py index a1643b9d531..6d2be697858 100755 --- a/examples/hf_ptq/hf_ptq.py +++ b/examples/hf_ptq/hf_ptq.py @@ -117,6 +117,12 @@ # Tokens per calibration batch when layerwise calibration picks the batch size (16 x 512). LAYERWISE_CALIB_TOKENS = 8192 + +def _layerwise_calib_batch_size(calib_seq: int, num_samples: int) -> int: + """Default layerwise batch size: LAYERWISE_CALIB_TOKENS tokens, but 1 to all of the samples.""" + return max(1, min(LAYERWISE_CALIB_TOKENS // calib_seq, num_samples)) + + mto.enable_huggingface_checkpointing() @@ -995,9 +1001,7 @@ def quantize_main( # then only holds one layer's activations, so size it by tokens: batch 1 makes an MoE # layer run its per-expert loop once per sample. elif is_layerwise: - args.batch_size = max( - 1, min(LAYERWISE_CALIB_TOKENS // args.calib_seq, sum(args.calib_size)) - ) + args.batch_size = _layerwise_calib_batch_size(args.calib_seq, sum(args.calib_size)) print( f"Layerwise calibration enabled. Using default batch_size={args.batch_size} " f"({LAYERWISE_CALIB_TOKENS} tokens per batch) for calibration." diff --git a/modelopt/torch/quantization/utils/calib_utils.py b/modelopt/torch/quantization/utils/calib_utils.py index fa7dc81c874..74e5ecb5349 100644 --- a/modelopt/torch/quantization/utils/calib_utils.py +++ b/modelopt/torch/quantization/utils/calib_utils.py @@ -53,19 +53,22 @@ def _gpu_is_full(device) -> bool: return get_used_gpu_mem_fraction(device) > _HESSIAN_GPU_MEM_LIMIT -def _hessian_offload_device() -> torch.device: +def _hessian_offload_device(nbytes: int) -> torch.device: """Where a Hessian goes when its weight's GPU is full: the emptiest GPU with room, else CPU. A model sharded over several GPUs leaves the ones not holding the current layer idle, and a Hessian there costs a device-to-device copy per update instead of a host copy and a CPU rescale of the whole matrix. Only GPUs this process already holds memory on are candidates, - and only outside distributed runs, so no other job's or rank's GPU is touched. + and only outside distributed runs, so no other job's or rank's GPU is touched. A candidate + must also have the Hessian's ``nbytes`` free. """ if not (dist.is_available() and dist.is_initialized() and dist.get_world_size() > 1): gpus = [i for i in range(torch.cuda.device_count()) if torch.cuda.memory_reserved(i)] - gpus.sort(key=get_used_gpu_mem_fraction) - if gpus and not _gpu_is_full(gpus[0]): - return torch.device("cuda", gpus[0]) + for gpu in sorted(gpus, key=get_used_gpu_mem_fraction): + if _gpu_is_full(gpu): + break + if torch.cuda.mem_get_info(gpu)[0] > nbytes: + return torch.device("cuda", gpu) return torch.device("cpu") @@ -181,7 +184,7 @@ def __init__( if device.type == "meta": device = torch.device("cpu") elif offload_to_cpu and device.type == "cuda" and _gpu_is_full(device): - device = _hessian_offload_device() + device = _hessian_offload_device(4 * in_features**2) self.hessian = torch.zeros(in_features, in_features, dtype=torch.float32, device=device) self.n_samples = 0 # Set by update_weights(); listed here for documentation. diff --git a/tests/examples/hf_ptq/test_hf_ptq_args.py b/tests/examples/hf_ptq/test_hf_ptq_args.py index 0854725ff90..c869ab46381 100644 --- a/tests/examples/hf_ptq/test_hf_ptq_args.py +++ b/tests/examples/hf_ptq/test_hf_ptq_args.py @@ -67,6 +67,16 @@ def _parse_hf_ptq_args(monkeypatch, *args): return hf_ptq, parsed_args +@pytest.mark.parametrize( + ("calib_seq", "num_samples", "expected"), + [(512, 512, 16), (512, 4, 4), (16384, 512, 1)], + ids=["token_budget", "sample_cap", "long_sequence"], +) +def test_layerwise_default_batch_size(monkeypatch, calib_seq, num_samples, expected): + hf_ptq = _import_hf_ptq(monkeypatch) + assert hf_ptq._layerwise_calib_batch_size(calib_seq, num_samples) == expected + + def test_recipe_help_distinguishes_weight_and_kv_autoquant(monkeypatch, capsys): hf_ptq = _import_hf_ptq(monkeypatch) monkeypatch.setattr(sys, "argv", ["hf_ptq.py", "--help"]) diff --git a/tests/gpu/torch/quantization/test_gptq.py b/tests/gpu/torch/quantization/test_gptq.py index f4d484c7415..ee6dbec46e1 100644 --- a/tests/gpu/torch/quantization/test_gptq.py +++ b/tests/gpu/torch/quantization/test_gptq.py @@ -57,6 +57,27 @@ def used_fraction(device): torch.testing.assert_close(helper.hessian.to("cuda:0"), reference) +@pytest.mark.parametrize("blocker", ["idle_gpu_full", "idle_gpu_too_small", "multi_rank"]) +def test_hessian_of_a_full_gpu_falls_back_to_the_cpu(need_2_gpus, monkeypatch, blocker): + """With no other GPU that can take it, or on any rank of a multi-rank run, it goes to the CPU.""" + idle_fraction = 0.9 if blocker == "idle_gpu_full" else 0.1 + + def used_fraction(device): + return 0.9 if torch.device(device).index in (0, None) else idle_fraction + + monkeypatch.setattr(calib_utils, "get_used_gpu_mem_fraction", used_fraction) + if blocker == "idle_gpu_too_small": + # 64 x 64 float32 is 16 KiB; leave 1 KiB free. + monkeypatch.setattr(torch.cuda, "mem_get_info", lambda device=None: (1024, 2**40)) + elif blocker == "multi_rank": + monkeypatch.setattr(calib_utils.dist, "is_initialized", lambda: True) + monkeypatch.setattr(calib_utils.dist, "get_world_size", lambda group=None: 2) + linear = torch.nn.Linear(64, 8, bias=False, device="cuda:0") + shard = torch.zeros(1, device="cuda:1") # noqa: F841 # the model also lives on cuda:1 + helper = GPTQHelper(linear, "linear", offload_to_cpu=True) + assert helper.hessian.device == torch.device("cpu") + + @pytest.mark.parametrize( ("block_size", "dim", "model_weight", "expect_weight_change"), [ From bf99bc3fd21e528fac2315a13fe70b5dd12edb8b Mon Sep 17 00:00:00 2001 From: Chenjie Luo Date: Thu, 22 Oct 2026 05:03:14 +0000 Subject: [PATCH 3/5] Skip per-expert GPTQ error logs Fused-expert GPTQ helpers no longer log a Hessian-weighted error each: it cost two weight-by-Hessian matmuls and a host sync per expert. gptq() logs one line per experts module instead. Co-Authored-By: Claude Opus 5.5 (1M context) Signed-off-by: Chenjie Luo --- modelopt/torch/quantization/model_calib.py | 18 +++++++++++++++--- .../torch/quantization/utils/calib_utils.py | 8 ++++++-- 2 files changed, 21 insertions(+), 5 deletions(-) diff --git a/modelopt/torch/quantization/model_calib.py b/modelopt/torch/quantization/model_calib.py index e50f3703f95..cf14a5b8c3c 100644 --- a/modelopt/torch/quantization/model_calib.py +++ b/modelopt/torch/quantization/model_calib.py @@ -2385,7 +2385,7 @@ def gptq( "exhaust it. Enable layerwise calibration to hold one layer's Hessians at a time." ) - def _make_gptq_handle(name, m, weight_quantizer=None, weight=None): + def _make_gptq_handle(name, m, weight_quantizer=None, weight=None, log_mse=True): quantizer = m.weight_quantizer if weight_quantizer is None else weight_quantizer backend = getattr(quantizer, "backend", None) if backend is None: @@ -2400,6 +2400,7 @@ def _make_gptq_handle(name, m, weight_quantizer=None, weight=None): importance_weighted=importance_weighted, weight_quantizer=weight_quantizer, weight=weight, + log_mse=log_mse, ) gptq_handles = {name: _make_gptq_handle(name, m) for name, m in quantized_layers} @@ -2426,8 +2427,11 @@ def _expert_hook(_input_quantizer, _args, output): for weight_name, weight, quantizers, input_quantizer in _fused_expert_projections( module ): + # An expert's logged error would cost two full matmuls; gptq() logs a summary. per_expert = { - idx: _make_gptq_handle(f"{name}.{weight_name}.{idx}", module, q, weight[idx]) + idx: _make_gptq_handle( + f"{name}.{weight_name}.{idx}", module, q, weight[idx], log_mse=False + ) for idx, q in enumerate(quantizers) if q.is_enabled } @@ -2461,13 +2465,21 @@ def _expert_hook(_input_quantizer, _args, output): with enable_weight_access_and_writeback(handle.module, model, names): handle.update_weights(block_size, perc_damp) handle.free() - for _, module in fused_experts: + for name, module in fused_experts: + updated = unreached = 0 with enable_weight_access_and_writeback(module, model, names): for weight, quantizer in module.iter_weights_for_calibration(): handle = expert_handles.pop(id(quantizer), None) if handle is not None: + unreached += handle.n_samples == 0 handle.update_weights(block_size, perc_damp, weight) handle.free() + updated += 1 + if updated: + print_rank_0( + f"[{name}] GPTQ updated {updated} expert weights, {unreached} of them never " + "reached by calibration" + ) del gptq_handles, expert_handles if torch.cuda.is_available(): diff --git a/modelopt/torch/quantization/utils/calib_utils.py b/modelopt/torch/quantization/utils/calib_utils.py index 74e5ecb5349..81f91d883e5 100644 --- a/modelopt/torch/quantization/utils/calib_utils.py +++ b/modelopt/torch/quantization/utils/calib_utils.py @@ -166,18 +166,21 @@ def __init__( importance_weighted=False, weight_quantizer=None, weight=None, + log_mse=True, ): """Initialize GPTQHelper with module state and Hessian storage. ``weight_quantizer`` and ``weight`` select one of several weights a module holds, such as one expert's slice of an HF fused-MoE parameter; by default the helper owns - ``module.weight`` and ``module.weight_quantizer``. + ``module.weight`` and ``module.weight_quantizer``. ``log_mse=False`` skips the logged + Hessian-weighted error, which costs two weight-by-Hessian matmuls. """ self.module = module self.name = name self.fused = fused self.importance_weighted = importance_weighted self._weight_quantizer = weight_quantizer + self.log_mse = log_mse weight = module.weight if weight is None else weight in_features = weight.shape[-1] device = weight.device @@ -262,7 +265,8 @@ def update_weights(self, block_size, perc_damp, weight=None): self.weight = weight.data.float().clone() self._prepare_hessian_inverse(hessian, perc_damp) self._blockwise_update(block_size) - self._print_mse_error(hessian, weight) + if self.log_mse: + self._print_mse_error(hessian, weight) weight.data.copy_(self.weight.reshape(weight.shape)) # ------------------------------------------------------------------ From 6cd5d46fcd8cdcd6a05d7004bfc02e7f759bd119 Mon Sep 17 00:00:00 2001 From: Chenjie Luo Date: Thu, 22 Oct 2026 14:58:19 +0000 Subject: [PATCH 4/5] Reject a non-positive --calib_seq in hf_ptq The default layerwise batch divides the token budget by --calib_seq, so 0 raised ZeroDivisionError; argument parsing now rejects it. Co-Authored-By: Claude Opus 5.5 (1M context) Signed-off-by: Chenjie Luo --- examples/hf_ptq/hf_ptq.py | 3 +++ tests/examples/hf_ptq/test_hf_ptq_args.py | 5 +++++ 2 files changed, 8 insertions(+) diff --git a/examples/hf_ptq/hf_ptq.py b/examples/hf_ptq/hf_ptq.py index 6d2be697858..27e8fb3e8d5 100755 --- a/examples/hf_ptq/hf_ptq.py +++ b/examples/hf_ptq/hf_ptq.py @@ -1433,6 +1433,9 @@ def parse_args() -> argparse.Namespace: if args.moe_calib_experts_ratio is not None and not (0.0 < args.moe_calib_experts_ratio <= 1.0): parser.error("--moe_calib_experts_ratio must be in the range (0.0, 1.0].") + if args.calib_seq <= 0: + parser.error("--calib_seq must be positive.") + if args.specdec_offline_dataset is not None and args.sparsity_fmt != "dense": parser.error("--specdec_offline_dataset is only supported with --sparsity_fmt dense (PTQ).") diff --git a/tests/examples/hf_ptq/test_hf_ptq_args.py b/tests/examples/hf_ptq/test_hf_ptq_args.py index c869ab46381..1e4596f9ba9 100644 --- a/tests/examples/hf_ptq/test_hf_ptq_args.py +++ b/tests/examples/hf_ptq/test_hf_ptq_args.py @@ -445,6 +445,11 @@ def test_mlflow_is_off_by_default(monkeypatch): assert args.mlflow_experiment is None +def test_calib_seq_must_be_positive(monkeypatch): + with pytest.raises(SystemExit): + _parse_hf_ptq_args(monkeypatch, "--pyt_ckpt_path", "/models/Qwen3-0.6B", "--calib_seq", "0") + + def test_mlflow_rejects_a_bad_tracking_uri(monkeypatch): with pytest.raises(SystemExit): _parse_hf_ptq_args( From 5487343c075c40244a27658f012ed37d8df38c5d Mon Sep 17 00:00:00 2001 From: Chenjie Luo Date: Thu, 22 Oct 2026 16:00:33 +0000 Subject: [PATCH 5/5] Test hf_ptq's layerwise calibration batch at the dataloader Drive quantize_main with a layerwise recipe up to the calibration dataloader and check the batch it gets: the 8,192-token default, its sample cap and floor, and an explicit --batch_size kept as is. This replaces the helper-only test. Co-Authored-By: Claude Opus 5.5 (1M context) Signed-off-by: Chenjie Luo --- tests/examples/hf_ptq/test_hf_ptq_args.py | 48 ++++++++++++++++++----- 1 file changed, 38 insertions(+), 10 deletions(-) diff --git a/tests/examples/hf_ptq/test_hf_ptq_args.py b/tests/examples/hf_ptq/test_hf_ptq_args.py index 1e4596f9ba9..b441dcc3ce7 100644 --- a/tests/examples/hf_ptq/test_hf_ptq_args.py +++ b/tests/examples/hf_ptq/test_hf_ptq_args.py @@ -67,16 +67,6 @@ def _parse_hf_ptq_args(monkeypatch, *args): return hf_ptq, parsed_args -@pytest.mark.parametrize( - ("calib_seq", "num_samples", "expected"), - [(512, 512, 16), (512, 4, 4), (16384, 512, 1)], - ids=["token_budget", "sample_cap", "long_sequence"], -) -def test_layerwise_default_batch_size(monkeypatch, calib_seq, num_samples, expected): - hf_ptq = _import_hf_ptq(monkeypatch) - assert hf_ptq._layerwise_calib_batch_size(calib_seq, num_samples) == expected - - def test_recipe_help_distinguishes_weight_and_kv_autoquant(monkeypatch, capsys): hf_ptq = _import_hf_ptq(monkeypatch) monkeypatch.setattr(sys, "argv", ["hf_ptq.py", "--help"]) @@ -445,6 +435,44 @@ def test_mlflow_is_off_by_default(monkeypatch): assert args.mlflow_experiment is None +@pytest.mark.parametrize( + ("extra_args", "expected"), + [ + ((), 16), # 8,192 tokens at --calib_seq 512 + (("--calib_size", "4"), 4), # never more than the samples + (("--calib_seq", "16384"), 1), # at least one sample + (("--batch_size", "1"), 1), # an explicit value is kept + ], + ids=["token_budget", "sample_cap", "long_sequence", "explicit"], +) +def test_layerwise_calibration_batch_reaches_the_dataloader(monkeypatch, extra_args, expected): + """A layerwise recipe sizes the default calibration batch by tokens.""" + hf_ptq, args = _parse_hf_ptq_args( + monkeypatch, + "--pyt_ckpt_path", + "/models/Qwen3-0.6B", + "--recipe", + "general/ptq/iq2_xxs", + "--calib_size", + "64", + *extra_args, + ) + + class _DataloaderReachedError(Exception): + pass + + batch_sizes = [] + + def make_calib_dataloader(args, *_args, **_kwargs): + batch_sizes.append(args.batch_size) + raise _DataloaderReachedError + + monkeypatch.setattr(hf_ptq, "make_calib_dataloader", make_calib_dataloader) + with pytest.raises(_DataloaderReachedError): + hf_ptq.quantize_main(args, None, None, None, False, None, None, torch.device("cpu")) + assert batch_sizes == [expected] + + def test_calib_seq_must_be_positive(monkeypatch): with pytest.raises(SystemExit): _parse_hf_ptq_args(monkeypatch, "--pyt_ckpt_path", "/models/Qwen3-0.6B", "--calib_seq", "0")