Improve layerwise GPTQ memory and speed: Hessian headroom, batch probe and offload fixes

CHANGELOG.rst

@@ -11,7 +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.
+- 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 to what one decoder layer's forward fits in GPU memory, up to 32,768 tokens (64 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.
 - 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.
 
@@ -49,6 +49,8 @@ Changelog
 - Fix ONNX quantization placing a per-tensor Q/DQ pair on a rank-0 (scalar) constant input, which TensorRT's ONNX parser rejects (``convertAxis`` asserts the quantization axis lies in ``[0, nbDims]`` but ``nbDims == 0``). Rank-0 constant operands are now kept unquantized in the int8, fp8, and autotune export paths; the existing single-element Mul/Div/Pad stripping stays gated under ``target_dla``.
 - Fix PTQ of native Transformers FP8 experts to apply source weight scales during calibration and export. Re-quantize affected checkpoints from the original FP8 source.
 - Fix GPTQ with ``fused: true`` failing with a CUDA illegal memory access, or a Triton pointer error, on a model that ``device_map="auto"`` spreads over several GPUs once calibration reaches a layer outside the current GPU.
+- Fix calibrating a CPU- or disk-offloaded model running out of GPU memory as it progresses: the quantizer state calibration adds to an offloaded layer now leaves the GPU with the layer's weights instead of accumulating there.
+- Fix layerwise calibration replaying the wrong cached activations, or failing with ``is in 'run' mode but has no cached inputs to replay``, after a GPU out-of-memory error. It now stops with an error asking for a smaller calibration batch size.
 - Fix ONNX preprocessing with custom TensorRT plugin libraries on TensorRT 11.
 - Fix ONNX remote AutoQDQ safety benchmarks returning infinite latency by running the generated engine with ``trtexec_safe`` on the configured target.
 - Fix ``INT8_DEFAULT_CFG``, ``FP8_DEFAULT_CFG``, ``NVFP4_DEFAULT_CFG`` and the other presets and recipes that import ``configs/ptq/units/default_disabled_quantizers`` and enable ``*input_quantizer``: since 0.44 they also quantized ``nn.LayerNorm`` inputs, which adds quantization error to the residual stream of models built on ``nn.LayerNorm`` such as BERT-style encoders. Re-quantize affected models; a recipe that wants ``nn.LayerNorm`` quantizers must enable them after importing that unit.

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

@@ -133,14 +133,15 @@ at sequence length 512, **batch size 8**.
 
 .. note::
 
-   **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
+   **Pick the batch size on purpose.** For layerwise runs, ``hf_ptq.py`` left at its default
+   probes one decoder layer of each kind and takes the most samples whose forward fits in half
+   the GPU's free memory, up to 32,768 tokens' worth (64 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. 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.
+   but activation ``input_scale`` values shift slightly, because a different batch shape
+   rounds the activations differently. Padding does not change: ``hf_ptq.py`` pads every
+   calibration sample to the same length whatever the batch size.
 
 Then launch the run:
 

examples/hf_ptq/hf_ptq.py

@@ -89,6 +89,7 @@
 from modelopt.torch.quantization.config import need_calibration
 from modelopt.torch.quantization.plugins.accelerate import init_quantized_weights
 from modelopt.torch.quantization.utils import is_quantized
+from modelopt.torch.quantization.utils.layerwise_calib import get_max_layerwise_batch_size
 from modelopt.torch.speculative.eagle.utils import (
     EagleOfflineDataCollator,
     OfflineSupervisedDataset,
@@ -114,12 +115,12 @@
 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
+# Most tokens per calibration batch when layerwise calibration picks the batch size (64 x 512).
+LAYERWISE_CALIB_TOKENS = 32768
 
 
 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."""
+    """Largest default layerwise batch: LAYERWISE_CALIB_TOKENS tokens, but 1 to all of the samples."""
     return max(1, min(LAYERWISE_CALIB_TOKENS // calib_seq, num_samples))
 
 
@@ -998,13 +999,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. 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.
+        # then only holds one layer's activations, so probe one decoder layer of each kind, up
+        # to a token budget: batch 1 makes an MoE layer run its per-expert loop once per sample.
         elif is_layerwise:
-            args.batch_size = _layerwise_calib_batch_size(args.calib_seq, sum(args.calib_size))
+            max_batch_size = _layerwise_calib_batch_size(args.calib_seq, sum(args.calib_size))
+            args.batch_size = get_max_layerwise_batch_size(
+                language_model, args.calib_seq, max_batch_size, sum(args.calib_size), device=device
+            )
             print(
-                f"Layerwise calibration enabled. Using default batch_size={args.batch_size} "
-                f"({LAYERWISE_CALIB_TOKENS} tokens per batch) for calibration."
+                f"Layerwise calibration enabled. Using batch_size={args.batch_size}, the most "
+                f"that fits a decoder layer's forward, up to {max_batch_size} "
+                f"({LAYERWISE_CALIB_TOKENS} tokens)."
             )
         else:
             # Calibration/sparsification will actually take much more memory than regular inference

modelopt/torch/quantization/mode.py

@@ -73,6 +73,7 @@
     smoothquant,
     svdquant,
 )
+from .utils.calib_utils import reset_calib_forward_memory
 
 __all__ = ["BaseCalibrateModeDescriptor"]
 
@@ -250,6 +251,7 @@ def wrapped_calib_func(
                 module._moe_calib_experts_ratio = moe_calib_experts_ratio
 
     if func is not None:
+        reset_calib_forward_memory()
         if layerwise:
             # All currently implemented PTQ algorithms support layerwise calibration;
             # future algorithms that need full-model context must add a guard here.

modelopt/torch/quantization/model_calib.py

@@ -19,6 +19,7 @@
 import math
 import time
 import warnings
+from collections import Counter
 from collections.abc import Callable, Mapping, Sequence
 from functools import partial
 from typing import Any, TypeAlias
@@ -69,7 +70,12 @@
     persistent_materialization,
     promote_static_block_weight_quantizers,
 )
-from .utils.calib_utils import _GPTQ_HELPER_REGISTRY, _IMATRIX_ENCODER_REGISTRY, GPTQHelper
+from .utils.calib_utils import (
+    _GPTQ_HELPER_REGISTRY,
+    _IMATRIX_ENCODER_REGISTRY,
+    GPTQHelper,
+    track_calib_forward_memory,
+)
 
 __all__ = [
     "CalibratorFactory",
@@ -2441,14 +2447,21 @@ def _expert_hook(_input_quantizer, _args, output):
                         input_quantizer.register_forward_hook(_make_expert_hook(module, per_expert))
                     )
 
+    placement = Counter(
+        str(h.hessian.device) for h in [*gptq_handles.values(), *expert_handles.values()]
+    )
     print_rank_0(
         f"Computing Hessians for {len(gptq_handles)} linear layers and "
-        f"{len(expert_handles)} fused-expert weights ({expert_hessian_bytes / 2**30:.1f} GiB)..."
+        f"{len(expert_handles)} fused-expert weights ({expert_hessian_bytes / 2**30:.1f} GiB, "
+        f"on {', '.join(f'{d} x{n}' for d, n in sorted(placement.items()))})..."
     )
 
     try:
-        with set_quantizer_by_cfg_context(
-            model, [{"quantizer_name": "*weight_quantizer", "enable": False}]
+        with (
+            set_quantizer_by_cfg_context(
+                model, [{"quantizer_name": "*weight_quantizer", "enable": False}]
+            ),
+            track_calib_forward_memory(),
         ):
             forward_loop(model)
     finally:

modelopt/torch/quantization/plugins/accelerate.py

@@ -79,6 +79,77 @@ def _writeback_params_to_weights_map(module, align_hook):
             w_map[key] = tensor.detach().cpu()
 
 
+def _own_quantizers(module):
+    """Quantizers whose state belongs with ``module``'s weights, not with a hooked descendant's."""
+    from ..nn import TensorQuantizer
+
+    for child in module.children():
+        if isinstance(child, TensorQuantizer):
+            yield child
+        elif not hasattr(child, "_hf_hook"):
+            yield from _own_quantizers(child)
+
+
+def _quantizer_tensors(module):
+    """``(key, tensor)`` for the state of ``module``'s own quantizers and their calibrators."""
+    for quantizer in _own_quantizers(module):
+        calibrator = getattr(quantizer, "_calibrator", None)
+        sources = {
+            "_parameters": quantizer._parameters,
+            "_buffers": quantizer._buffers,
+            "_calibrator": vars(calibrator) if calibrator is not None else {},
+        }
+        for source, tensors in sources.items():
+            for name, tensor in tensors.items():
+                if isinstance(tensor, torch.Tensor):
+                    yield (id(quantizer), source, name), tensor
+
+
+def _carry_quantizers_with_offload(hook):
+    """Make an offload ``hook`` move its module's quantizer state on and off the GPU with the weights.
+
+    Accelerate only manages the tensors a module had when it was dispatched, so calibration state
+    such as NVFP4's per-block ``_amax``, and its calibrator's copy, stayed on the execution device
+    for every offloaded layer and added up over the model. The CPU tensors stay the master copy:
+    a forward gets GPU copies, and after it the CPU tensors come back, once any tensor the forward
+    modified or replaced is copied back. The packed-weight cache is keyed to the weight just
+    released, so it is dropped.
+    """
+    if getattr(hook, "_modelopt_carries_quantizers", False):
+        return
+    pre_forward, post_forward = hook.pre_forward, hook.post_forward
+    onloaded: dict[tuple, tuple[torch.Tensor, torch.Tensor, int]] = {}
+
+    def carrying_pre_forward(module, *args, **kwargs):
+        result = pre_forward(module, *args, **kwargs)
+        if hook.offload:
+            for key, tensor in _quantizer_tensors(module):
+                if tensor.device.type == "cpu":
+                    master = tensor.data
+                    tensor.data = master.to(hook.execution_device)
+                    onloaded[key] = (tensor, master, tensor._version)
+        return result
+
+    def carrying_post_forward(module, output):
+        output = post_forward(module, output)
+        if hook.offload:
+            for key, tensor in _quantizer_tensors(module):
+                # meta: a hook that places submodules offloaded it already.
+                if tensor.device.type in ("cpu", "meta"):
+                    continue
+                entry = onloaded.pop(key, None)
+                if entry is not None and entry[0] is tensor and entry[2] == tensor._version:
+                    tensor.data = entry[1]
+                else:
+                    tensor.data = tensor.data.cpu()
+            for quantizer in _own_quantizers(module):
+                quantizer._quantizer_cache = None
+        return output
+
+    hook.pre_forward, hook.post_forward = carrying_pre_forward, carrying_post_forward
+    hook._modelopt_carries_quantizers = True
+
+
 @contextmanager
 def weight_access_and_writeback_context(module, writeback: bool = True):
     """Context manager for weight access and writeback for modules managed by accelerate.
@@ -102,6 +173,7 @@ def weight_access_and_writeback_context(module, writeback: bool = True):
         # weights would be overwritten with stale CPU state_dict values.
         needs_materialize = any(p.device.type == "meta" for p in mod.parameters())
         if needs_materialize:
+            _carry_quantizers_with_offload(hook)
             hook.pre_forward(mod)
         hook.offload = False
         materialized.append((mod, hook, needs_materialize))

modelopt/torch/quantization/utils/calib_utils.py

@@ -38,36 +38,79 @@
 
 """GPTQ helper and Hessian utilities for calibration."""
 
+from contextlib import contextmanager
+
 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
+# 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
+# Once measured, keep this multiple of the most a calibration forward added on the GPU, and at
+# least this share of the GPU.
+_HESSIAN_RESERVE_MARGIN = 1.5
+_HESSIAN_MIN_RESERVE = 0.1
+# The most memory a GPTQ calibration forward has added on each GPU, by device index.
+_calib_forward_peak: dict[int, int] = {}
+
+
+def _process_gpus() -> list[int]:
+    """GPUs this process already holds memory on."""
+    return [i for i in range(torch.cuda.device_count()) if torch.cuda.memory_reserved(i)]
+
+
+def _available_bytes(gpu: int) -> int:
+    """Memory an allocation on ``gpu`` can get: the driver's free memory and our cached blocks."""
+    free, _ = torch.cuda.mem_get_info(gpu)
+    return free + torch.cuda.memory_reserved(gpu) - torch.cuda.memory_allocated(gpu)
+
 
+def _gpu_has_room(gpu: int, nbytes: int) -> bool:
+    """Whether ``nbytes`` more on ``gpu`` leaves the calibration forward the memory it needs."""
+    total = torch.cuda.get_device_properties(gpu).total_memory
+    peak = _calib_forward_peak.get(gpu)
+    if peak is None:
+        reserve = _HESSIAN_UNMEASURED_RESERVE * total
+    else:
+        reserve = max(_HESSIAN_MIN_RESERVE * total, _HESSIAN_RESERVE_MARGIN * peak)
+    return _available_bytes(gpu) - nbytes >= reserve
+
+
+def reset_calib_forward_memory():
+    """Forget measured calibration forwards; a new run, maybe of another model, measures afresh."""
+    _calib_forward_peak.clear()
 
-def _gpu_is_full(device) -> bool:
-    return get_used_gpu_mem_fraction(device) > _HESSIAN_GPU_MEM_LIMIT
+
+@contextmanager
+def track_calib_forward_memory():
+    """Record the memory the enclosed calibration forward adds on each GPU, for Hessian placement.
+
+    The peak is read without resetting PyTorch's peak counter, which callers may be measuring
+    with, so an earlier, higher peak can only overstate the need: the placement errs toward the
+    CPU, never toward running the next forward out of memory.
+    """
+    gpus = _process_gpus() if torch.cuda.is_available() else []
+    allocated = {gpu: torch.cuda.memory_allocated(gpu) for gpu in gpus}
+    yield
+    for gpu in gpus:
+        added = torch.cuda.max_memory_allocated(gpu) - allocated[gpu]
+        _calib_forward_peak[gpu] = max(_calib_forward_peak.get(gpu, 0), added)
 
 
 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.
+    """Where a Hessian goes when its weight's GPU lacks room: 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.
+    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)]
-        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:
+        for gpu in sorted(_process_gpus(), key=_available_bytes, reverse=True):
+            if _gpu_has_room(gpu, nbytes):
                 return torch.device("cuda", gpu)
     return torch.device("cpu")
 
@@ -186,8 +229,10 @@ def __init__(
         device = weight.device
         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)
+        elif offload_to_cpu and device.type == "cuda":
+            index = torch.cuda.current_device() if device.index is None else device.index
+            if not _gpu_has_room(index, 4 * in_features**2):
+                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.