[2/3] Speed up layerwise GPTQ on MoE models

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 ``<format>_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.
 

docs/source/announcements/20260928-single-gpu-ptq.rst

@@ -133,9 +133,11 @@ 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`` 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.

examples/hf_ptq/hf_ptq.py

@@ -114,6 +114,14 @@
 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
+
+
+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()
 
@@ -989,10 +997,15 @@ 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 = _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."
+            )
         else:
             # Calibration/sparsification will actually take much more memory than regular inference
             # due to intermediate tensors for fake quantization. Setting sample_memory_usage_ratio
@@ -1420,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).")
 

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 <modelopt.torch.quantization.config.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)
@@ -2381,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:
@@ -2396,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}
@@ -2422,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
                 }
@@ -2457,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():

modelopt/torch/quantization/utils/calib_utils.py

@@ -38,15 +38,39 @@
 
 """GPTQ helper and Hessian utilities for calibration."""
 
-import math
-
 import torch
 import torch.distributed as dist
 
 from modelopt.torch.utils import print_rank_0
 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(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. 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)]
+        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")
+
 
 def update_hessian(input, hessian, n_samples):
     """Update hessian matrix with new input samples using incremental formula.
@@ -61,19 +85,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
 
@@ -139,25 +166,28 @@ 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
-        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(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.
@@ -235,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))
 
     # ------------------------------------------------------------------

tests/examples/hf_ptq/test_hf_ptq_args.py

@@ -435,6 +435,49 @@ 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")
+
+
 def test_mlflow_rejects_a_bad_tracking_uri(monkeypatch):
     with pytest.raises(SystemExit):
         _parse_hf_ptq_args(

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,45 @@
 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("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"),
     [