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

CHANGELOG.rst

@@ -12,7 +12,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.
 
@@ -60,6 +60,8 @@ Changelog
 - Fix KV-cache quantization through ``register_attention_for_kv_quant``, the fallback for Hugging Face attention classes that do not use the attention interface, attaching the ``q``, ``k`` and ``v`` bmm quantizers to the wrong tensors for most attention layouts. Recalibrate and re-export models quantized through this fallback. Attention classes that compute their scores with ``torch.baddbmm`` (e.g. ``BloomAttention``, ``GPTNeoXJapaneseAttention``) are no longer patched, since their operands were always wired wrong (or raised ``IndexError``), and now get the existing warning to register a custom quantized attention class.
 - 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.
 - Nested ``only_student_forward()`` and ``only_teacher_forward()`` contexts now preserve the enclosing execution mode, including when an inner context is disabled or raises an exception.
+- 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",
@@ -2480,14 +2486,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

@@ -28,6 +28,8 @@
 
 import modelopt.torch.quantization as mtq
 
+from ..nn import TensorQuantizer
+
 __all__ = ["init_quantized_weights"]
 
 
@@ -79,6 +81,85 @@ 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."""
+    for child in module.children():
+        if isinstance(child, TensorQuantizer):
+            yield child
+        elif not hasattr(child, "_hf_hook"):
+            yield from _own_quantizers(child)
+
+
+# A quantizer's calibrators: its amax calibrator, and its bias calibrator under static bias.
+_CALIBRATORS = ("_calibrator", "_bias_calibrator")
+
+
+def _quantizer_tensors(module):
+    """``(key, tensor)`` for the state of ``module``'s own quantizers and their calibrators."""
+    for quantizer in _own_quantizers(module):
+        sources = {"_parameters": quantizer._parameters, "_buffers": quantizer._buffers}
+        for source in _CALIBRATORS:
+            calibrator = getattr(quantizer, source, None)
+            sources[source] = 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. A quantizer's own tensors
+    are always copied back, since its setters write through ``.data``, which leaves ``_version``
+    unchanged; a calibrator's tensor is copied back only if the forward replaced or modified it.
+    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 None or entry[0] is not tensor:
+                    tensor.data = tensor.data.cpu()
+                    continue
+                master = entry[1]
+                if key[1] not in _CALIBRATORS or entry[2] != tensor._version:
+                    if master.shape != tensor.shape or master.dtype != tensor.dtype:
+                        tensor.data = tensor.data.cpu()
+                        continue
+                    # Into the existing master: a new CPU tensor costs 8x as much in page faults.
+                    master.copy_(tensor.data)
+                tensor.data = master
+            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 +183,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))