From a53f43f645ed96ea1dfb52c082935348e0d0d0e2 Mon Sep 17 00:00:00 2001
From: Chenjie Luo <chenjiel@nvidia.com>
Date: Thu, 22 Oct 2026 17:28:24 +0000
Subject: [PATCH 1/6] Place GPTQ Hessians by measured headroom instead of a 65%
 usage cap

A Hessian stayed on a GPU only while the GPU was under 65% used, so on
large GPUs the Hessians of a big MoE went to the CPU with ~100 GB free:
on GLM-5.3-Flash over 4x GB300, half the MoE layers took 142 s instead
of 43 s. A GPU now takes a Hessian if, after it, the memory available to
an allocation (driver-free plus this process's cached blocks) still
covers 1.5x the most a GPTQ calibration forward has added on that GPU,
and at least 10% of it. Until a forward has been measured, the old 35%
reserve applies; each calibration run measures afresh. The Hessian log
line names where they went.

Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com>
Signed-off-by: Chenjie Luo <chenjiel@nvidia.com>
---
 modelopt/torch/quantization/mode.py           |  2 +
 modelopt/torch/quantization/model_calib.py    | 21 +++++-
 .../torch/quantization/utils/calib_utils.py   | 75 +++++++++++++++----
 tests/gpu/torch/quantization/test_gptq.py     | 68 +++++++++++++----
 tests/unit/torch/quantization/test_gptq.py    |  9 +++
 5 files changed, 141 insertions(+), 34 deletions(-)

diff --git a/modelopt/torch/quantization/mode.py b/modelopt/torch/quantization/mode.py
index 47b188ab9fc..fa36b566193 100644
--- a/modelopt/torch/quantization/mode.py
+++ b/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.
diff --git a/modelopt/torch/quantization/model_calib.py b/modelopt/torch/quantization/model_calib.py
index f02732a9334..b00c10a8ba5 100644
--- a/modelopt/torch/quantization/model_calib.py
+++ b/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:
diff --git a/modelopt/torch/quantization/utils/calib_utils.py b/modelopt/torch/quantization/utils/calib_utils.py
index 81f91d883e5..b30752fea84 100644
--- a/modelopt/torch/quantization/utils/calib_utils.py
+++ b/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.
diff --git a/tests/gpu/torch/quantization/test_gptq.py b/tests/gpu/torch/quantization/test_gptq.py
index ee6dbec46e1..d7faf3b41f9 100644
--- a/tests/gpu/torch/quantization/test_gptq.py
+++ b/tests/gpu/torch/quantization/test_gptq.py
@@ -39,13 +39,16 @@
 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 _report_available(monkeypatch, available):
+    """Make each GPU report ``available[gpu]`` bytes for a new allocation, with no forward measured."""
+    monkeypatch.setattr(calib_utils, "_available_bytes", lambda gpu: available[gpu])
+    monkeypatch.setattr(calib_utils, "_calib_forward_peak", {})
 
-    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)
+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."""
+    total = torch.cuda.get_device_properties(0).total_memory
+    _report_available(monkeypatch, {0: total // 10, 1: total * 9 // 10})
     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)
@@ -60,16 +63,15 @@ def used_fraction(device):
 @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":
+    total = torch.cuda.get_device_properties(0).total_memory
+    reserve = int(calib_utils._HESSIAN_UNMEASURED_RESERVE * total)
+    idle = {
+        "idle_gpu_full": reserve // 2,
+        "idle_gpu_too_small": reserve + 64 * 64 * 4 - 1,  # one byte short of reserve + Hessian
+        "multi_rank": total,
+    }[blocker]
+    _report_available(monkeypatch, {0: total // 10, 1: idle})
+    if 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")
@@ -78,6 +80,42 @@ def used_fraction(device):
     assert helper.hessian.device == torch.device("cpu")
 
 
+@pytest.mark.parametrize(
+    ("measured_fraction", "stays"),
+    [(None, False), (0.05, True), (0.25, False)],
+    ids=["unmeasured", "small_forward", "large_forward"],
+)
+def test_a_measured_calibration_forward_sets_the_hessian_reserve(
+    monkeypatch, measured_fraction, stays
+):
+    """At 30% available, a GPU keeps a Hessian only once a measured forward leaves it room."""
+    total = torch.cuda.get_device_properties(0).total_memory
+    _report_available(monkeypatch, {0: total * 3 // 10})
+    monkeypatch.setattr(calib_utils, "_process_gpus", lambda: [0])
+    if measured_fraction is not None:
+        calib_utils._calib_forward_peak[0] = int(measured_fraction * total)
+    linear = torch.nn.Linear(64, 8, bias=False, device="cuda:0")
+    helper = GPTQHelper(linear, "linear", offload_to_cpu=True)
+    assert helper.hessian.device == (torch.device("cuda", 0) if stays else torch.device("cpu"))
+
+
+def test_gptq_records_its_calibration_forward_memory(monkeypatch):
+    """The next layer's Hessian reserve covers what this layer's calibration forward added."""
+    monkeypatch.setattr(calib_utils, "_calib_forward_peak", {})
+    model = torch.nn.Linear(256, 256, device="cuda")
+    inputs = torch.randn(2, 16, 256, device="cuda")
+    mtq.quantize(model, mtq.NVFP4_DEFAULT_CFG, forward_loop=lambda m: m(inputs))
+    transient = 64 * 2**20
+
+    def forward_loop(m):
+        scratch = torch.empty(transient, dtype=torch.uint8, device="cuda")
+        m(inputs)
+        del scratch
+
+    gptq(model, forward_loop=forward_loop, perc_damp=0.1, block_size=16)
+    assert calib_utils._calib_forward_peak[torch.cuda.current_device()] >= transient
+
+
 @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 ee1a55a8554..5744348d67e 100644
--- a/tests/unit/torch/quantization/test_gptq.py
+++ b/tests/unit/torch/quantization/test_gptq.py
@@ -22,6 +22,7 @@
 from _test_utils.torch.distributed.utils import spawn_multiprocess_job
 
 import modelopt.torch.quantization as mtq
+from modelopt.torch.quantization.utils import calib_utils
 from modelopt.torch.quantization.utils.calib_utils import GPTQHelper, update_hessian
 from modelopt.torch.utils.distributed import ParallelState
 
@@ -172,3 +173,11 @@ def test_gptq_hessian_dp_sync(skip_on_windows):
     spawn_multiprocess_job(
         2, partial(_gptq_hessian_dp_sync_helper, token_counts=(5, 0)), backend="gloo"
     )
+
+
+def test_a_calibration_run_forgets_earlier_forward_measurements(monkeypatch):
+    """Another run, maybe of another model, must not size its Hessian reserve from this one."""
+    monkeypatch.setattr(calib_utils, "_calib_forward_peak", {0: 2**60})
+    model = torch.nn.Linear(16, 16)
+    mtq.quantize(model, mtq.INT8_DEFAULT_CFG, forward_loop=lambda m: m(torch.randn(4, 16)))
+    assert 2**60 not in calib_utils._calib_forward_peak.values()

From 919fa71805efd8c2ad951fdc6a4e8a5055427b65 Mon Sep 17 00:00:00 2001
From: Chenjie Luo <chenjiel@nvidia.com>
Date: Fri, 23 Oct 2026 17:43:04 +0000
Subject: [PATCH 2/6] Size the layerwise calibration batch by probing each kind
 of decoder layer

hf_ptq.py gave layerwise calibration a fixed 8,192-token batch, which can
run out of memory on a model or GPU whose decoder-layer forward is larger
than the ones it was tuned on. get_max_layerwise_batch_size runs the first
decoder layer of each kind (by module classes) on one and two dummy samples,
takes the growth of its peak memory per GPU, and keeps the largest batch
whose forward stays within half of the GPU's available memory after the
inputs layerwise calibration caches. The token budget stays the cap. All
ranks agree on the smallest batch.

Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com>
Signed-off-by: Chenjie Luo <chenjiel@nvidia.com>
---
 CHANGELOG.rst                                 |   2 +-
 .../announcements/20260928-single-gpu-ptq.rst |  15 +-
 examples/hf_ptq/hf_ptq.py                     |  19 ++-
 .../quantization/utils/layerwise_calib.py     | 131 ++++++++++++++++++
 tests/examples/hf_ptq/test_hf_ptq_args.py     |  24 ++--
 .../quantization/test_layerwise_calibrate.py  |  43 +++++-
 6 files changed, 210 insertions(+), 24 deletions(-)

diff --git a/CHANGELOG.rst b/CHANGELOG.rst
index 4f5b89098d4..bae8db22b40 100755
--- a/CHANGELOG.rst
+++ b/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 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.
 - 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 b5e4e7fbac2..f246bfdcf85 100644
--- a/docs/source/announcements/20260928-single-gpu-ptq.rst
+++ b/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 8,192 tokens' worth (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. 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:
 
diff --git a/examples/hf_ptq/hf_ptq.py b/examples/hf_ptq/hf_ptq.py
index 27e8fb3e8d5..5215de08cf3 100755
--- a/examples/hf_ptq/hf_ptq.py
+++ b/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).
+# Most 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."""
+    """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
diff --git a/modelopt/torch/quantization/utils/layerwise_calib.py b/modelopt/torch/quantization/utils/layerwise_calib.py
index 7d9609063b6..2a954607360 100644
--- a/modelopt/torch/quantization/utils/layerwise_calib.py
+++ b/modelopt/torch/quantization/utils/layerwise_calib.py
@@ -577,6 +577,137 @@ def cache_outputs_for_next_layer_calib(
         return self.get_input_activations(next_layer, forward_loop)
 
 
+def _used_gpus() -> list[int]:
+    """GPUs this process holds memory on; probing any other would create a context on it."""
+    return [i for i in range(torch.cuda.device_count()) if torch.cuda.memory_reserved(i)]
+
+
+def _available_gpu_bytes(gpu: int) -> int:
+    """Memory a new allocation on *gpu* can get: what is free plus what PyTorch caches unused."""
+    free, _ = torch.cuda.mem_get_info(gpu)
+    return free + torch.cuda.memory_reserved(gpu) - torch.cuda.memory_allocated(gpu)
+
+
+def _cuda_bytes_per_gpu(obj: Any, out: dict[int, int] | None = None) -> dict[int, int]:
+    """Bytes of the CUDA tensors in a nested args/kwargs structure, per GPU."""
+    out = {} if out is None else out
+    if isinstance(obj, torch.Tensor):
+        if obj.is_cuda:
+            out[obj.device.index] = out.get(obj.device.index, 0) + obj.nbytes
+    elif isinstance(obj, dict):
+        for value in obj.values():
+            _cuda_bytes_per_gpu(value, out)
+    elif isinstance(obj, (list, tuple)):
+        for value in obj:
+            _cuda_bytes_per_gpu(value, out)
+    return out
+
+
+def _capture_layer_inputs(
+    model: nn.Module, layer: nn.Module, batch_size: int, seq_len: int, device: torch.device
+):
+    """Run *model* on a dummy batch until *layer* and return the ``(args, kwargs)`` it receives."""
+    # Inline import: dataset_utils is not needed by the rest of this module.
+    from modelopt.torch.utils.dataset_utils import _disable_use_cache
+
+    input_ids = torch.zeros(batch_size, seq_len, dtype=torch.long, device=device)
+    # One left-pad position, like padded calibration rows, so the model builds the explicit
+    # attention mask that calibration batches get instead of a cheaper causal-only path.
+    attention_mask = torch.ones_like(input_ids)
+    attention_mask[:, 0] = 0
+    info = layer._layerwise_calib
+    info.mode = "capture"
+    try:
+        with _disable_use_cache(model):
+            model(input_ids, attention_mask=attention_mask)
+    finally:
+        info.mode = "original"
+    (inputs,) = info.collected_inputs
+    info.collected_inputs = []
+    return inputs
+
+
+def _forward_peak_per_gpu(layer: nn.Module, args: tuple, kwargs: dict) -> dict[int, int]:
+    """Peak memory one forward of *layer* adds on each GPU."""
+    base = {}
+    for gpu in _used_gpus():
+        torch.cuda.reset_peak_memory_stats(gpu)
+        base[gpu] = torch.cuda.memory_allocated(gpu)
+    layer(*args, **kwargs)
+    return {gpu: torch.cuda.max_memory_allocated(gpu) - base.get(gpu, 0) for gpu in _used_gpus()}
+
+
+def _layer_batch_fit(
+    model: nn.Module,
+    layer: nn.Module,
+    seq_len: int,
+    num_samples: int,
+    memory_fraction: float,
+    device: torch.device,
+) -> int | None:
+    """Largest batch whose forward through *layer* fits; ``None`` if no GPU memory grows with it."""
+    peaks = []
+    # The first forward only warms up: it also allocates one-off buffers like cuBLAS's workspace.
+    for batch_size in (1, 1, 2):
+        args, kwargs = _capture_layer_inputs(model, layer, batch_size, seq_len, device)
+        if batch_size == 1:
+            cached_per_sample = _cuda_bytes_per_gpu((args, kwargs))
+        peaks.append(_forward_peak_per_gpu(layer, args, kwargs))
+        del args, kwargs
+
+    _, one, two = peaks
+    fits = []
+    for gpu, peak in two.items():
+        per_sample = peak - one.get(gpu, 0)
+        if per_sample <= 0:
+            continue
+        # Layerwise calibration caches every sample's inputs for this layer and the next.
+        available = _available_gpu_bytes(gpu) - 2 * num_samples * cached_per_sample.get(gpu, 0)
+        fixed = one.get(gpu, 0) - per_sample
+        fits.append(int((memory_fraction * available - fixed) // per_sample))
+    return min(fits, default=None)
+
+
+@torch.no_grad()
+def get_max_layerwise_batch_size(
+    model: nn.Module,
+    seq_len: int,
+    max_batch_size: int,
+    num_samples: int,
+    memory_fraction: float = 0.5,
+    device: torch.device | None = None,
+) -> int:
+    """Largest calibration batch, up to *max_batch_size*, whose decoder-layer forwards fit.
+
+    Layerwise calibration runs one decoder layer at a time, so a full-model batch probe such as
+    :func:`get_max_batch_size <modelopt.torch.utils.dataset_utils.get_max_batch_size>` both
+    overstates the memory a batch needs and risks running out of it. Instead, the first layer
+    of each kind (by module classes: dense or MoE, full or linear attention) runs on one and
+    two samples of *seq_len* tokens, and the growth of its peak memory on each GPU is that GPU's
+    cost per sample. A batch fits when the layer's forward stays within *memory_fraction* of
+    the GPU's available memory, after the inputs layerwise calibration caches for all
+    *num_samples*; the rest is left for quantizers and Hessians. The dummy token ids go to
+    *device*, by default the model's. All ranks return the same, smallest batch.
+    """
+    batch_size = max_batch_size
+    layers = LayerActivationCollector.get_decoder_layers(model)
+    if layers and torch.cuda.is_available() and max_batch_size > 1:
+        kinds: dict[tuple[str, ...], nn.Module] = {}
+        for layer in layers:
+            kinds.setdefault(tuple(sorted({type(m).__name__ for m in layer.modules()})), layer)
+        device = device or get_module_device(model)
+        collector = LayerActivationCollector(model)
+        collector._patch_all_layers(layers)
+        try:
+            for layer in kinds.values():
+                fit = _layer_batch_fit(model, layer, seq_len, num_samples, memory_fraction, device)
+                if fit is not None:
+                    batch_size = min(batch_size, fit)
+        finally:
+            collector._unpatch_all_layers()
+    return max(1, min(dist.allgather(batch_size)))
+
+
 def _move_to_device(obj: Any, device: torch.device) -> Any:
     """Recursively move tensors to *device*. Non-tensors are returned as-is."""
     if isinstance(obj, torch.Tensor):
diff --git a/tests/examples/hf_ptq/test_hf_ptq_args.py b/tests/examples/hf_ptq/test_hf_ptq_args.py
index b441dcc3ce7..1f1020d41ab 100644
--- a/tests/examples/hf_ptq/test_hf_ptq_args.py
+++ b/tests/examples/hf_ptq/test_hf_ptq_args.py
@@ -436,17 +436,20 @@ def test_mlflow_is_off_by_default(monkeypatch):
 
 
 @pytest.mark.parametrize(
-    ("extra_args", "expected"),
+    ("extra_args", "memory_fit", "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
+        ((), 1000, 16),  # 8,192 tokens at --calib_seq 512
+        ((), 5, 5),  # what fits a decoder layer's forward
+        (("--calib_size", "4"), 1000, 4),  # never more than the samples
+        (("--calib_seq", "16384"), 1000, 1),  # at least one sample
+        (("--batch_size", "1"), 1000, 1),  # an explicit value is kept
     ],
-    ids=["token_budget", "sample_cap", "long_sequence", "explicit"],
+    ids=["token_budget", "memory", "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."""
+def test_layerwise_calibration_batch_reaches_the_dataloader(
+    monkeypatch, extra_args, memory_fit, expected
+):
+    """A layerwise recipe sizes the default calibration batch by tokens and layer memory."""
     hf_ptq, args = _parse_hf_ptq_args(
         monkeypatch,
         "--pyt_ckpt_path",
@@ -468,6 +471,11 @@ def make_calib_dataloader(args, *_args, **_kwargs):
         raise _DataloaderReachedError
 
     monkeypatch.setattr(hf_ptq, "make_calib_dataloader", make_calib_dataloader)
+    monkeypatch.setattr(
+        hf_ptq,
+        "get_max_layerwise_batch_size",
+        lambda model, seq_len, max_batch_size, num_samples, device: min(max_batch_size, memory_fit),
+    )
     with pytest.raises(_DataloaderReachedError):
         hf_ptq.quantize_main(args, None, None, None, False, None, None, torch.device("cpu"))
     assert batch_sizes == [expected]
diff --git a/tests/gpu/torch/quantization/test_layerwise_calibrate.py b/tests/gpu/torch/quantization/test_layerwise_calibrate.py
index 0e6751bfb79..00d4b2bb198 100644
--- a/tests/gpu/torch/quantization/test_layerwise_calibrate.py
+++ b/tests/gpu/torch/quantization/test_layerwise_calibrate.py
@@ -19,7 +19,11 @@
 import torch.nn as nn
 
 from modelopt.torch.quantization.model_calib import layerwise_calibrate
-from modelopt.torch.quantization.utils.layerwise_calib import LayerActivationCollector
+from modelopt.torch.quantization.utils import layerwise_calib
+from modelopt.torch.quantization.utils.layerwise_calib import (
+    LayerActivationCollector,
+    get_max_layerwise_batch_size,
+)
 
 
 class _DecoderBlock(nn.Module):
@@ -339,3 +343,40 @@ def test_skip_placeholder_uses_meta_device():
     assert out.device == torch.device("meta")
     assert out.shape == torch.Size([2, 3])
     assert out.dtype == torch.float32
+
+
+class _WideDecoderBlock(_DecoderBlock):
+    """A heavier kind of decoder layer: its FFN is 8x wider than _DecoderBlock's."""
+
+    def __init__(self, dim=16):
+        super().__init__(dim)
+        self.ffn = nn.Sequential(
+            nn.Linear(dim, dim * 32, bias=False),
+            nn.ReLU(),
+            nn.Linear(dim * 32, dim, bias=False),
+        )
+
+
+def test_layerwise_batch_probe_fits_the_heaviest_kind_of_layer(monkeypatch):
+    """The batch probe runs every kind of decoder layer, not just the first one."""
+    _register_test_discoverer(monkeypatch)
+    budget = 256 * 2**20
+    monkeypatch.setattr(layerwise_calib, "_available_gpu_bytes", lambda gpu: budget)
+    model = _SimpleTransformerModel(n_layers=3, dim=256).cuda()
+
+    def probe():
+        return get_max_layerwise_batch_size(model, seq_len=512, max_batch_size=4096, num_samples=1)
+
+    light = probe()
+    model.layers[2] = _WideDecoderBlock(256).cuda()
+    heavy = probe()
+    assert 1 < heavy < light < 4096
+    assert not any(hasattr(layer, "_layerwise_calib") for layer in model.layers)
+
+    # The heavy layer's forward at the probed batch takes at most half of the available memory.
+    x = torch.randn(heavy, 512, 256, device="cuda")
+    torch.cuda.reset_peak_memory_stats()
+    base = torch.cuda.memory_allocated()
+    with torch.no_grad():
+        model.layers[2](x)
+    assert torch.cuda.max_memory_allocated() - base <= budget / 2

From 2c8947374d5cbfc84bae87e91612c9a43b27b97b Mon Sep 17 00:00:00 2001
From: Chenjie Luo <chenjiel@nvidia.com>
Date: Fri, 23 Oct 2026 21:47:58 +0000
Subject: [PATCH 3/6] Raise the layerwise calibration token cap to 32,768

The batch probe now keeps a layerwise batch within GPU memory, so the
token budget no longer has to be the guard against running out of it.
At 32,768 tokens (64 samples at --calib_seq 512) a model whose decoder
layers fit the larger batch runs its MoE layers' per-expert loops a
quarter as often; one that does not, such as GLM-5.3-Flash with its
linear attention, keeps the batch the probe picks.

Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com>
Signed-off-by: Chenjie Luo <chenjiel@nvidia.com>
---
 CHANGELOG.rst                                         | 2 +-
 docs/source/announcements/20260928-single-gpu-ptq.rst | 2 +-
 examples/hf_ptq/hf_ptq.py                             | 4 ++--
 tests/examples/hf_ptq/test_hf_ptq_args.py             | 4 ++--
 4 files changed, 6 insertions(+), 6 deletions(-)

diff --git a/CHANGELOG.rst b/CHANGELOG.rst
index bae8db22b40..eee0e79bf80 100755
--- a/CHANGELOG.rst
+++ b/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 to what one decoder layer's forward fits in GPU memory, up to 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.
+- 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.
 
diff --git a/docs/source/announcements/20260928-single-gpu-ptq.rst b/docs/source/announcements/20260928-single-gpu-ptq.rst
index f246bfdcf85..58fcf3286e6 100644
--- a/docs/source/announcements/20260928-single-gpu-ptq.rst
+++ b/docs/source/announcements/20260928-single-gpu-ptq.rst
@@ -135,7 +135,7 @@ at sequence length 512, **batch size 8**.
 
    **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 8,192 tokens' worth (16 at sequence length 512). It was 1 when
+   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,
diff --git a/examples/hf_ptq/hf_ptq.py b/examples/hf_ptq/hf_ptq.py
index 5215de08cf3..3f41e64b83a 100755
--- a/examples/hf_ptq/hf_ptq.py
+++ b/examples/hf_ptq/hf_ptq.py
@@ -115,8 +115,8 @@
 from modelopt.torch.utils.vlm_dataset_utils import get_vlm_dataset_dataloader
 
 RAND_SEED = 1234
-# Most 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:
diff --git a/tests/examples/hf_ptq/test_hf_ptq_args.py b/tests/examples/hf_ptq/test_hf_ptq_args.py
index 1f1020d41ab..c8f7e1b667e 100644
--- a/tests/examples/hf_ptq/test_hf_ptq_args.py
+++ b/tests/examples/hf_ptq/test_hf_ptq_args.py
@@ -438,10 +438,10 @@ def test_mlflow_is_off_by_default(monkeypatch):
 @pytest.mark.parametrize(
     ("extra_args", "memory_fit", "expected"),
     [
-        ((), 1000, 16),  # 8,192 tokens at --calib_seq 512
+        ((), 1000, 64),  # 32,768 tokens at --calib_seq 512
         ((), 5, 5),  # what fits a decoder layer's forward
         (("--calib_size", "4"), 1000, 4),  # never more than the samples
-        (("--calib_seq", "16384"), 1000, 1),  # at least one sample
+        (("--calib_seq", "65536"), 1000, 1),  # at least one sample
         (("--batch_size", "1"), 1000, 1),  # an explicit value is kept
     ],
     ids=["token_budget", "memory", "sample_cap", "long_sequence", "explicit"],

From 84aec953b4ce9e033bb0e35e4cf3e918e4791ed2 Mon Sep 17 00:00:00 2001
From: Chenjie Luo <chenjiel@nvidia.com>
Date: Sat, 24 Oct 2026 07:16:21 +0000
Subject: [PATCH 4/6] Keep offloaded layers' quantizer state off the GPU and
 stop layerwise OOM retries

Accelerate's offload hooks only move the tensors a module had when it was
dispatched, so the quantizer state calibration adds (NVFP4's per-block
_amax and its calibrator's copy, GGML payload pins) stayed on the execution
device for every CPU- or disk-offloaded layer. On a single L40S that grew
by about 8 GB over Qwen3.6-35B-A3B and ran out of memory near the last
layers at any batch size. The first time a module's weights are
materialized for weight access, its offload hook now also carries its
quantizers' state: GPU copies for a forward, the CPU master copy after,
with anything the forward changed copied back.

During a layerwise capture forward, an out-of-memory error made the
calibration forward loop split the data batch and retry. The layer in
'run' mode replays its cached batches one per call at their cached size,
so the retry could not shrink the work and replayed the next batch's
inputs instead, misaligning the captured activations or failing on an
empty cache. It now stops with an error that asks for a smaller batch.

Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com>
Signed-off-by: Chenjie Luo <chenjiel@nvidia.com>
---
 CHANGELOG.rst                                 |  2 +
 .../torch/quantization/plugins/accelerate.py  | 72 +++++++++++++++++++
 .../quantization/utils/layerwise_calib.py     | 10 +++
 .../plugins/test_accelerate_gpu.py            | 15 ++++
 .../quantization/test_layerwise_calibrate.py  | 43 +++++++++++
 5 files changed, 142 insertions(+)

diff --git a/CHANGELOG.rst b/CHANGELOG.rst
index eee0e79bf80..44bca74e533 100755
--- a/CHANGELOG.rst
+++ b/CHANGELOG.rst
@@ -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.
diff --git a/modelopt/torch/quantization/plugins/accelerate.py b/modelopt/torch/quantization/plugins/accelerate.py
index eb206db2c5c..43deea1bbcd 100644
--- a/modelopt/torch/quantization/plugins/accelerate.py
+++ b/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))
diff --git a/modelopt/torch/quantization/utils/layerwise_calib.py b/modelopt/torch/quantization/utils/layerwise_calib.py
index 2a954607360..b37988fb521 100644
--- a/modelopt/torch/quantization/utils/layerwise_calib.py
+++ b/modelopt/torch/quantization/utils/layerwise_calib.py
@@ -361,6 +361,16 @@ def _early_stop_forward(module_self, *args, **kwargs):
                     if depth[0] > 1:
                         raise
                     return None
+                except torch.cuda.OutOfMemoryError as e:
+                    # The forward loop would split the data batch and retry, but the layer in 'run'
+                    # mode replays the batch it cached at full size, one cached batch per call: a
+                    # retry cannot shrink it and would replay the next batch's inputs instead.
+                    raise RuntimeError(
+                        "Layerwise calibration ran out of GPU memory while a decoder layer replayed "
+                        "a cached calibration batch. Splitting the batch now cannot help, since the "
+                        "layer replays it at the size it was cached with; rerun with a smaller "
+                        "calibration batch size."
+                    ) from e
                 except RuntimeError as e:
                     if "meta" not in str(e).lower():
                         raise
diff --git a/tests/gpu/torch/quantization/plugins/test_accelerate_gpu.py b/tests/gpu/torch/quantization/plugins/test_accelerate_gpu.py
index 6aab8fe48c8..c172cda3534 100644
--- a/tests/gpu/torch/quantization/plugins/test_accelerate_gpu.py
+++ b/tests/gpu/torch/quantization/plugins/test_accelerate_gpu.py
@@ -39,6 +39,7 @@
 
 import modelopt.torch.quantization as mtq
 from modelopt.torch.quantization.extensions import get_cuda_ext_mx
+from modelopt.torch.quantization.nn import TensorQuantizer
 from modelopt.torch.quantization.utils import (
     enable_weight_access_and_writeback,
     is_quantized_linear,
@@ -184,6 +185,20 @@ def test_layerwise_calibrate_cpu_offloaded(tmp_path, use_checkpoint):
 
     assert torch.allclose(output_ref.logits, output_test.logits)
 
+    # Between uses, the offloaded layer's quantizer state stays off the GPU like its weights.
+    offloaded_state = [
+        tensor
+        for quantizer in model.model.layers[0].modules()
+        if isinstance(quantizer, TensorQuantizer)
+        for tensor in [
+            *quantizer.buffers(recurse=False),
+            *(vars(quantizer._calibrator).values() if quantizer._calibrator is not None else ()),
+        ]
+        if isinstance(tensor, torch.Tensor)
+    ]
+    assert offloaded_state
+    assert all(tensor.device.type == "cpu" for tensor in offloaded_state)
+
     if use_checkpoint:
         manifest_path = os.path.join(ckpt_dir, "manifest.json")
         assert os.path.isfile(manifest_path)
diff --git a/tests/unit/torch/quantization/test_layerwise_calibrate.py b/tests/unit/torch/quantization/test_layerwise_calibrate.py
index 07e14d50f2a..5e7ee03307e 100644
--- a/tests/unit/torch/quantization/test_layerwise_calibrate.py
+++ b/tests/unit/torch/quantization/test_layerwise_calibrate.py
@@ -31,6 +31,7 @@
     _OutsideQuantizerCalibrator,
     _SkipLayer,
 )
+from modelopt.torch.utils.dataset_utils import create_forward_loop
 
 
 class _DecoderBlock(nn.Module):
@@ -681,6 +682,48 @@ def test_set_layer_states_raises_on_empty_collected_inputs(monkeypatch):
         collector._unpatch_all_layers()
 
 
+class _OutOfMemoryOnceBlock(_DecoderBlock):
+    """Decoder block whose next forward runs out of memory when ``oom`` is set."""
+
+    oom = False
+
+    def forward(self, x, **kwargs):
+        if self.oom:
+            self.oom = False
+            raise torch.cuda.OutOfMemoryError("out of memory")
+        return super().forward(x, **kwargs)
+
+
+class _InputIdsModel(_SimpleTransformerModel):
+    """Takes ``input_ids`` as the calibration forward loop passes it."""
+
+    def __init__(self):
+        super().__init__(n_layers=3, dim=16)
+        self.layers[0] = _OutOfMemoryOnceBlock(16)
+
+    def forward(self, input_ids, **kwargs):
+        return super().forward(input_ids)
+
+
+def test_out_of_memory_in_a_replayed_layer_is_not_split(monkeypatch):
+    """A batch split on OOM would replay the next cached batch in its place, so the OOM must stop."""
+    _register_test_discoverer(monkeypatch)
+    model = _InputIdsModel()
+    batches = [{"input_ids": torch.randint(0, 32, (4, 8))} for _ in range(3)]
+    forward_loop = create_forward_loop(dataloader=batches)
+
+    collector = LayerActivationCollector(model)
+    collector._patch_all_layers()
+    try:
+        collector.get_input_activations(model.layers[0], forward_loop)
+        model.layers[0].oom = True  # layer 0 replays its cached batches while layer 1 is captured
+        with pytest.raises(RuntimeError) as error:
+            collector.get_input_activations(model.layers[1], forward_loop)
+        assert not isinstance(error.value, torch.cuda.OutOfMemoryError)
+    finally:
+        collector._unpatch_all_layers()
+
+
 def test_run_asserts_on_empty_cached_inputs(monkeypatch):
     """A layer in 'run' mode with no cached inputs must raise AssertionError."""
     _register_test_discoverer(monkeypatch)

From 347289628e39ab2d51485101242eadc41f92b081 Mon Sep 17 00:00:00 2001
From: Chenjie Luo <chenjiel@nvidia.com>
Date: Mon, 26 Oct 2026 19:03:51 +0000
Subject: [PATCH 5/6] Address review: offload writeback, probe OOM, forward
 peak, shared helpers

- The offload hook now always copies a quantizer's own tensors back into
  their CPU master after a forward. Quantizer setters (amax,
  pre_quant_scale, bias_value, global_amax) write through .data, which
  leaves _version unchanged, so the version check restored the stale
  master and dropped the update. Copying into the existing master costs
  9-11 ms per offloaded forward for a Qwen3-30B-A3B MoE layer's NVFP4
  per-block amax (onload about 60 ms); a new CPU tensor cost 76 ms.
  Calibrator tensors, which are never written through .data, keep the
  version check.
- The batch probe returns 1 when one sample fits but two run out of
  memory, in the layer or while capturing its inputs, instead of
  stopping the run (and leaving other ranks waiting in the allgather).
- track_calib_forward_memory resets the peak counters first. Nothing
  reset them between layers, so from the second layer on the measured
  forward included the previous layer's Hessians and GPTQ updates and
  the reserve sent Hessians to the CPU for nothing.
- The probe groups decoder layers by parameter shapes as well as module
  classes, so layers of one class with different widths, like
  Nemotron-NAS's, are each probed.
- The GPU enumeration and available-memory helpers live once, in
  modelopt.torch.utils.perf; the dataset_utils and TensorQuantizer
  imports move to module scope.

Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com>
Signed-off-by: Chenjie Luo <chenjiel@nvidia.com>
---
 .../torch/quantization/plugins/accelerate.py  | 24 ++++--
 .../torch/quantization/utils/calib_utils.py   | 30 +++----
 .../quantization/utils/layerwise_calib.py     | 81 ++++++++++---------
 modelopt/torch/utils/perf.py                  | 13 +++
 .../plugins/test_accelerate_gpu.py            | 17 ++++
 tests/gpu/torch/quantization/test_gptq.py     | 11 ++-
 .../quantization/test_layerwise_calibrate.py  | 42 +++++++++-
 7 files changed, 146 insertions(+), 72 deletions(-)

diff --git a/modelopt/torch/quantization/plugins/accelerate.py b/modelopt/torch/quantization/plugins/accelerate.py
index 43deea1bbcd..aa57f759590 100644
--- a/modelopt/torch/quantization/plugins/accelerate.py
+++ b/modelopt/torch/quantization/plugins/accelerate.py
@@ -28,6 +28,8 @@
 
 import modelopt.torch.quantization as mtq
 
+from ..nn import TensorQuantizer
+
 __all__ = ["init_quantized_weights"]
 
 
@@ -81,8 +83,6 @@ def _writeback_params_to_weights_map(module, align_hook):
 
 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
@@ -111,9 +111,10 @@ def _carry_quantizers_with_offload(hook):
     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.
+    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
@@ -138,10 +139,17 @@ def carrying_post_forward(module, output):
                 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:
+                if entry is None or entry[0] is not tensor:
                     tensor.data = tensor.data.cpu()
+                    continue
+                master = entry[1]
+                if key[1] != "_calibrator" 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
diff --git a/modelopt/torch/quantization/utils/calib_utils.py b/modelopt/torch/quantization/utils/calib_utils.py
index b30752fea84..4895a31103b 100644
--- a/modelopt/torch/quantization/utils/calib_utils.py
+++ b/modelopt/torch/quantization/utils/calib_utils.py
@@ -43,7 +43,7 @@
 import torch
 import torch.distributed as dist
 
-from modelopt.torch.utils import print_rank_0
+from modelopt.torch.utils import get_available_gpu_memory, get_process_gpus, print_rank_0
 from modelopt.torch.utils.network import bind_forward_method, unpatch_forward_method
 
 # A GPU takes a new Hessian only if enough memory stays available after it for the calibration
@@ -57,17 +57,6 @@
 _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
@@ -76,7 +65,7 @@ def _gpu_has_room(gpu: int, nbytes: int) -> bool:
         reserve = _HESSIAN_UNMEASURED_RESERVE * total
     else:
         reserve = max(_HESSIAN_MIN_RESERVE * total, _HESSIAN_RESERVE_MARGIN * peak)
-    return _available_bytes(gpu) - nbytes >= reserve
+    return get_available_gpu_memory(gpu) - nbytes >= reserve
 
 
 def reset_calib_forward_memory():
@@ -88,12 +77,15 @@ def reset_calib_forward_memory():
 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.
+    PyTorch's peak counters are reset first: otherwise an earlier peak, such as the previous
+    layer's Hessians and GPTQ updates, would count as forward memory and send Hessians to the CPU
+    for nothing.
     """
-    gpus = _process_gpus() if torch.cuda.is_available() else []
-    allocated = {gpu: torch.cuda.memory_allocated(gpu) for gpu in gpus}
+    gpus = get_process_gpus() if torch.cuda.is_available() else []
+    allocated = {}
+    for gpu in gpus:
+        torch.cuda.reset_peak_memory_stats(gpu)
+        allocated[gpu] = torch.cuda.memory_allocated(gpu)
     yield
     for gpu in gpus:
         added = torch.cuda.max_memory_allocated(gpu) - allocated[gpu]
@@ -109,7 +101,7 @@ def _hessian_offload_device(nbytes: int) -> torch.device:
     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):
-        for gpu in sorted(_process_gpus(), key=_available_bytes, reverse=True):
+        for gpu in sorted(get_process_gpus(), key=get_available_gpu_memory, reverse=True):
             if _gpu_has_room(gpu, nbytes):
                 return torch.device("cuda", gpu)
     return torch.device("cpu")
diff --git a/modelopt/torch/quantization/utils/layerwise_calib.py b/modelopt/torch/quantization/utils/layerwise_calib.py
index b37988fb521..1fc8a37e6b7 100644
--- a/modelopt/torch/quantization/utils/layerwise_calib.py
+++ b/modelopt/torch/quantization/utils/layerwise_calib.py
@@ -35,7 +35,13 @@
 import torch.nn as nn
 
 from modelopt.torch.utils import distributed as dist
-from modelopt.torch.utils import print_rank_0, warn_rank_0
+from modelopt.torch.utils import (
+    get_available_gpu_memory,
+    get_process_gpus,
+    print_rank_0,
+    warn_rank_0,
+)
+from modelopt.torch.utils.dataset_utils import _disable_use_cache
 from modelopt.torch.utils.network import (
     bind_forward_method,
     get_module_device,
@@ -366,10 +372,9 @@ def _early_stop_forward(module_self, *args, **kwargs):
                     # mode replays the batch it cached at full size, one cached batch per call: a
                     # retry cannot shrink it and would replay the next batch's inputs instead.
                     raise RuntimeError(
-                        "Layerwise calibration ran out of GPU memory while a decoder layer replayed "
-                        "a cached calibration batch. Splitting the batch now cannot help, since the "
-                        "layer replays it at the size it was cached with; rerun with a smaller "
-                        "calibration batch size."
+                        "Layerwise calibration ran out of GPU memory. Splitting the batch now "
+                        "cannot help, since each decoder layer replays its cached inputs at the "
+                        "size they were cached with; rerun with a smaller calibration batch size."
                     ) from e
                 except RuntimeError as e:
                     if "meta" not in str(e).lower():
@@ -587,17 +592,6 @@ def cache_outputs_for_next_layer_calib(
         return self.get_input_activations(next_layer, forward_loop)
 
 
-def _used_gpus() -> list[int]:
-    """GPUs this process holds memory on; probing any other would create a context on it."""
-    return [i for i in range(torch.cuda.device_count()) if torch.cuda.memory_reserved(i)]
-
-
-def _available_gpu_bytes(gpu: int) -> int:
-    """Memory a new allocation on *gpu* can get: what is free plus what PyTorch caches unused."""
-    free, _ = torch.cuda.mem_get_info(gpu)
-    return free + torch.cuda.memory_reserved(gpu) - torch.cuda.memory_allocated(gpu)
-
-
 def _cuda_bytes_per_gpu(obj: Any, out: dict[int, int] | None = None) -> dict[int, int]:
     """Bytes of the CUDA tensors in a nested args/kwargs structure, per GPU."""
     out = {} if out is None else out
@@ -617,9 +611,6 @@ def _capture_layer_inputs(
     model: nn.Module, layer: nn.Module, batch_size: int, seq_len: int, device: torch.device
 ):
     """Run *model* on a dummy batch until *layer* and return the ``(args, kwargs)`` it receives."""
-    # Inline import: dataset_utils is not needed by the rest of this module.
-    from modelopt.torch.utils.dataset_utils import _disable_use_cache
-
     input_ids = torch.zeros(batch_size, seq_len, dtype=torch.long, device=device)
     # One left-pad position, like padded calibration rows, so the model builds the explicit
     # attention mask that calibration batches get instead of a cheaper causal-only path.
@@ -640,11 +631,13 @@ def _capture_layer_inputs(
 def _forward_peak_per_gpu(layer: nn.Module, args: tuple, kwargs: dict) -> dict[int, int]:
     """Peak memory one forward of *layer* adds on each GPU."""
     base = {}
-    for gpu in _used_gpus():
+    for gpu in get_process_gpus():
         torch.cuda.reset_peak_memory_stats(gpu)
         base[gpu] = torch.cuda.memory_allocated(gpu)
     layer(*args, **kwargs)
-    return {gpu: torch.cuda.max_memory_allocated(gpu) - base.get(gpu, 0) for gpu in _used_gpus()}
+    return {
+        gpu: torch.cuda.max_memory_allocated(gpu) - base.get(gpu, 0) for gpu in get_process_gpus()
+    }
 
 
 def _layer_batch_fit(
@@ -656,23 +649,34 @@ def _layer_batch_fit(
     device: torch.device,
 ) -> int | None:
     """Largest batch whose forward through *layer* fits; ``None`` if no GPU memory grows with it."""
-    peaks = []
     # The first forward only warms up: it also allocates one-off buffers like cuBLAS's workspace.
-    for batch_size in (1, 1, 2):
-        args, kwargs = _capture_layer_inputs(model, layer, batch_size, seq_len, device)
-        if batch_size == 1:
-            cached_per_sample = _cuda_bytes_per_gpu((args, kwargs))
-        peaks.append(_forward_peak_per_gpu(layer, args, kwargs))
+    for _ in range(2):
+        args, kwargs = _capture_layer_inputs(model, layer, 1, seq_len, device)
+        cached_per_sample = _cuda_bytes_per_gpu((args, kwargs))
+        one = _forward_peak_per_gpu(layer, args, kwargs)
         del args, kwargs
+    try:
+        args, kwargs = _capture_layer_inputs(model, layer, 2, seq_len, device)
+        two = _forward_peak_per_gpu(layer, args, kwargs)
+    except RuntimeError as e:
+        # Capturing inputs reports running out of memory as a RuntimeError raised from it.
+        if not any(isinstance(x, torch.cuda.OutOfMemoryError) for x in (e, e.__cause__)):
+            raise
+        two = None
+    finally:
+        args = kwargs = None
+    if two is None:
+        # One sample fits but two do not: release what the failed forward left cached.
+        torch.cuda.empty_cache()
+        return 1
 
-    _, one, two = peaks
     fits = []
     for gpu, peak in two.items():
         per_sample = peak - one.get(gpu, 0)
         if per_sample <= 0:
             continue
         # Layerwise calibration caches every sample's inputs for this layer and the next.
-        available = _available_gpu_bytes(gpu) - 2 * num_samples * cached_per_sample.get(gpu, 0)
+        available = get_available_gpu_memory(gpu) - 2 * num_samples * cached_per_sample.get(gpu, 0)
         fixed = one.get(gpu, 0) - per_sample
         fits.append(int((memory_fraction * available - fixed) // per_sample))
     return min(fits, default=None)
@@ -692,19 +696,22 @@ def get_max_layerwise_batch_size(
     Layerwise calibration runs one decoder layer at a time, so a full-model batch probe such as
     :func:`get_max_batch_size <modelopt.torch.utils.dataset_utils.get_max_batch_size>` both
     overstates the memory a batch needs and risks running out of it. Instead, the first layer
-    of each kind (by module classes: dense or MoE, full or linear attention) runs on one and
-    two samples of *seq_len* tokens, and the growth of its peak memory on each GPU is that GPU's
-    cost per sample. A batch fits when the layer's forward stays within *memory_fraction* of
-    the GPU's available memory, after the inputs layerwise calibration caches for all
-    *num_samples*; the rest is left for quantizers and Hessians. The dummy token ids go to
-    *device*, by default the model's. All ranks return the same, smallest batch.
+    of each kind (by module classes and parameter shapes: dense or MoE, full or linear attention,
+    and per-layer widths such as Nemotron-NAS's) runs on one and two samples of *seq_len* tokens,
+    and the growth of its peak memory on each GPU is that GPU's cost per sample. A batch fits
+    when the layer's forward stays within *memory_fraction* of the GPU's available memory, after
+    the inputs layerwise calibration caches for all *num_samples*; the rest is left for
+    quantizers and Hessians. A layer that fits one sample but not two gives a batch of one. The
+    dummy token ids go to *device*, by default the model's. All ranks return the same, smallest
+    batch.
     """
     batch_size = max_batch_size
     layers = LayerActivationCollector.get_decoder_layers(model)
     if layers and torch.cuda.is_available() and max_batch_size > 1:
-        kinds: dict[tuple[str, ...], nn.Module] = {}
+        kinds: dict[tuple, nn.Module] = {}
         for layer in layers:
-            kinds.setdefault(tuple(sorted({type(m).__name__ for m in layer.modules()})), layer)
+            classes = tuple(sorted({type(m).__name__ for m in layer.modules()}))
+            kinds.setdefault((classes, tuple(tuple(p.shape) for p in layer.parameters())), layer)
         device = device or get_module_device(model)
         collector = LayerActivationCollector(model)
         collector._patch_all_layers(layers)
diff --git a/modelopt/torch/utils/perf.py b/modelopt/torch/utils/perf.py
index 93bd4839c15..a9ee748659b 100644
--- a/modelopt/torch/utils/perf.py
+++ b/modelopt/torch/utils/perf.py
@@ -27,7 +27,9 @@
     "AccumulatingTimer",
     "Timer",
     "clear_cuda_cache",
+    "get_available_gpu_memory",
     "get_cuda_memory_stats",
+    "get_process_gpus",
     "get_used_gpu_mem_fraction",
     "maybe_clear_cuda_cache",
     "report_memory",
@@ -91,6 +93,17 @@ def get_used_gpu_mem_fraction(device="cuda:0"):
     return (total_memory - free_memory) / total_memory
 
 
+def get_process_gpus() -> list[int]:
+    """GPUs this process already holds memory on; querying any other would create a context on it."""
+    return [i for i in range(torch.cuda.device_count()) if torch.cuda.memory_reserved(i)]
+
+
+def get_available_gpu_memory(gpu: int) -> int:
+    """Bytes a new 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 report_memory(name="", rank=0, device=None):
     """Simple GPU memory report."""
     memory_stats = get_cuda_memory_stats(device)
diff --git a/tests/gpu/torch/quantization/plugins/test_accelerate_gpu.py b/tests/gpu/torch/quantization/plugins/test_accelerate_gpu.py
index c172cda3534..4808377ace7 100644
--- a/tests/gpu/torch/quantization/plugins/test_accelerate_gpu.py
+++ b/tests/gpu/torch/quantization/plugins/test_accelerate_gpu.py
@@ -141,6 +141,23 @@ def test_cpu_offloaded_tinyllama(tmp_path):
     assert torch.allclose(output_ref.logits, output_test.logits)
 
 
+def test_offloaded_quantizer_keeps_what_its_setter_wrote(tmp_path):
+    """A setter writes through ``.data``, which leaves ``_version`` alone; the offload must keep it."""
+    model, _, _, inputs = make_cpu_offloaded_model(tmp_path)
+    mtq.quantize(model, mtq.INT8_DEFAULT_CFG, lambda model: model(inputs))
+    linear = model.model.layers[0].self_attn.q_proj
+    quantizer = linear.weight_quantizer
+    with enable_weight_access_and_writeback(linear, model):
+        pass
+    assert quantizer._amax.device.type == "cpu"
+
+    amax = quantizer._amax * 2
+    with enable_weight_access_and_writeback(linear, model):
+        quantizer.amax = amax
+    assert quantizer._amax.device.type == "cpu"
+    assert torch.equal(quantizer._amax, amax)
+
+
 @pytest.mark.parametrize("use_checkpoint", [False, True], ids=["no_ckpt", "ckpt"])
 def test_layerwise_calibrate_cpu_offloaded(tmp_path, use_checkpoint):
     """Layerwise calibration on CPU-offloaded model matches GPU-only reference."""
diff --git a/tests/gpu/torch/quantization/test_gptq.py b/tests/gpu/torch/quantization/test_gptq.py
index d7faf3b41f9..5ea632a13e3 100644
--- a/tests/gpu/torch/quantization/test_gptq.py
+++ b/tests/gpu/torch/quantization/test_gptq.py
@@ -41,7 +41,7 @@
 
 def _report_available(monkeypatch, available):
     """Make each GPU report ``available[gpu]`` bytes for a new allocation, with no forward measured."""
-    monkeypatch.setattr(calib_utils, "_available_bytes", lambda gpu: available[gpu])
+    monkeypatch.setattr(calib_utils, "get_available_gpu_memory", lambda gpu: available[gpu])
     monkeypatch.setattr(calib_utils, "_calib_forward_peak", {})
 
 
@@ -91,7 +91,7 @@ def test_a_measured_calibration_forward_sets_the_hessian_reserve(
     """At 30% available, a GPU keeps a Hessian only once a measured forward leaves it room."""
     total = torch.cuda.get_device_properties(0).total_memory
     _report_available(monkeypatch, {0: total * 3 // 10})
-    monkeypatch.setattr(calib_utils, "_process_gpus", lambda: [0])
+    monkeypatch.setattr(calib_utils, "get_process_gpus", lambda: [0])
     if measured_fraction is not None:
         calib_utils._calib_forward_peak[0] = int(measured_fraction * total)
     linear = torch.nn.Linear(64, 8, bias=False, device="cuda:0")
@@ -100,7 +100,8 @@ def test_a_measured_calibration_forward_sets_the_hessian_reserve(
 
 
 def test_gptq_records_its_calibration_forward_memory(monkeypatch):
-    """The next layer's Hessian reserve covers what this layer's calibration forward added."""
+    """The next layer's Hessian reserve covers what this layer's calibration forward added, and
+    not an earlier peak such as the previous layer's Hessians."""
     monkeypatch.setattr(calib_utils, "_calib_forward_peak", {})
     model = torch.nn.Linear(256, 256, device="cuda")
     inputs = torch.randn(2, 16, 256, device="cuda")
@@ -112,8 +113,10 @@ def forward_loop(m):
         m(inputs)
         del scratch
 
+    earlier = 16 * transient
+    torch.empty(earlier, dtype=torch.uint8, device="cuda")  # freed at once; only the peak stays
     gptq(model, forward_loop=forward_loop, perc_damp=0.1, block_size=16)
-    assert calib_utils._calib_forward_peak[torch.cuda.current_device()] >= transient
+    assert transient <= calib_utils._calib_forward_peak[torch.cuda.current_device()] < earlier // 2
 
 
 @pytest.mark.parametrize(
diff --git a/tests/gpu/torch/quantization/test_layerwise_calibrate.py b/tests/gpu/torch/quantization/test_layerwise_calibrate.py
index 00d4b2bb198..3da2b016de7 100644
--- a/tests/gpu/torch/quantization/test_layerwise_calibrate.py
+++ b/tests/gpu/torch/quantization/test_layerwise_calibrate.py
@@ -15,6 +15,7 @@
 
 """Integration tests for layerwise_calibrate and LayerActivationCollector."""
 
+import pytest
 import torch
 import torch.nn as nn
 
@@ -357,18 +358,23 @@ def __init__(self, dim=16):
         )
 
 
-def test_layerwise_batch_probe_fits_the_heaviest_kind_of_layer(monkeypatch):
-    """The batch probe runs every kind of decoder layer, not just the first one."""
+@pytest.mark.parametrize("same_class", [False, True], ids=["new_class", "same_class"])
+def test_layerwise_batch_probe_fits_the_heaviest_kind_of_layer(monkeypatch, same_class):
+    """The batch probe runs every kind of decoder layer, not just the first one, including layers
+    of one class that differ in width, like Nemotron-NAS's."""
     _register_test_discoverer(monkeypatch)
     budget = 256 * 2**20
-    monkeypatch.setattr(layerwise_calib, "_available_gpu_bytes", lambda gpu: budget)
+    monkeypatch.setattr(layerwise_calib, "get_available_gpu_memory", lambda gpu: budget)
     model = _SimpleTransformerModel(n_layers=3, dim=256).cuda()
 
     def probe():
         return get_max_layerwise_batch_size(model, seq_len=512, max_batch_size=4096, num_samples=1)
 
     light = probe()
-    model.layers[2] = _WideDecoderBlock(256).cuda()
+    if same_class:
+        model.layers[2].ffn = _WideDecoderBlock(256).ffn.cuda()
+    else:
+        model.layers[2] = _WideDecoderBlock(256).cuda()
     heavy = probe()
     assert 1 < heavy < light < 4096
     assert not any(hasattr(layer, "_layerwise_calib") for layer in model.layers)
@@ -380,3 +386,31 @@ def probe():
     with torch.no_grad():
         model.layers[2](x)
     assert torch.cuda.max_memory_allocated() - base <= budget / 2
+
+
+class _OutOfMemoryOnTwoSamples(nn.Module):
+    """Runs *module*, but runs out of GPU memory on a batch of more than one sample."""
+
+    def __init__(self, module):
+        super().__init__()
+        self.module = module
+
+    def forward(self, x, **kwargs):
+        if x.shape[0] > 1:
+            raise torch.cuda.OutOfMemoryError("CUDA out of memory.")
+        return self.module(x, **kwargs)
+
+
+@pytest.mark.parametrize("where", ["layer_forward", "input_capture"])
+def test_layerwise_batch_probe_falls_back_to_one_sample(monkeypatch, where):
+    """A decoder layer that fits one sample but not two, or whose inputs for two samples cannot
+    be captured, gets a batch of one instead of stopping the run."""
+    _register_test_discoverer(monkeypatch)
+    model = _SimpleTransformerModel(n_layers=2, dim=64).cuda()
+    if where == "layer_forward":
+        model.layers[1] = _OutOfMemoryOnTwoSamples(model.layers[1])
+    else:
+        model.embed = _OutOfMemoryOnTwoSamples(model.embed)
+
+    assert get_max_layerwise_batch_size(model, seq_len=8, max_batch_size=64, num_samples=1) == 1
+    assert not any(hasattr(layer, "_layerwise_calib") for layer in model.layers)

From 8c7a66854711950258137017715c0b8315a16c43 Mon Sep 17 00:00:00 2001
From: Chenjie Luo <chenjiel@nvidia.com>
Date: Mon, 26 Oct 2026 21:02:51 +0000
Subject: [PATCH 6/6] Probe every GPU a layer kind runs on; count only
 allocatable cache

- The batch probe runs only the first layer of each kind, so with a
  device_map spreading identical layers over GPUs, only that layer's GPU
  limited the batch. Each kind's measured cost now applies to every GPU
  that runs a layer of that kind, without extra forwards.
- get_available_gpu_memory leaves out inactive split blocks, the free
  pieces of partly used cached segments: a large allocation such as a
  Hessian cannot use them, so counting them could place a Hessian on a
  GPU where allocating it runs out of memory instead of falling back to
  the CPU. One memory_stats call replaces memory_reserved plus
  memory_allocated, which is cheaper too.
- The offload hook also carries a static-bias quantizer's bias
  calibrator, which otherwise stayed on the GPU for offloaded layers.

Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com>
Signed-off-by: Chenjie Luo <chenjiel@nvidia.com>
---
 .../torch/quantization/plugins/accelerate.py  | 16 ++++---
 .../quantization/utils/layerwise_calib.py     | 45 ++++++++++++++-----
 modelopt/torch/utils/perf.py                  | 12 ++++-
 .../plugins/test_accelerate_gpu.py            | 23 ++++++++++
 tests/gpu/torch/quantization/test_gptq.py     | 15 +++++++
 .../quantization/test_layerwise_calibrate.py  | 25 +++++++++++
 6 files changed, 116 insertions(+), 20 deletions(-)

diff --git a/modelopt/torch/quantization/plugins/accelerate.py b/modelopt/torch/quantization/plugins/accelerate.py
index aa57f759590..fc2c37b9d94 100644
--- a/modelopt/torch/quantization/plugins/accelerate.py
+++ b/modelopt/torch/quantization/plugins/accelerate.py
@@ -90,15 +90,17 @@ def _own_quantizers(module):
             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):
-        calibrator = getattr(quantizer, "_calibrator", None)
-        sources = {
-            "_parameters": quantizer._parameters,
-            "_buffers": quantizer._buffers,
-            "_calibrator": vars(calibrator) if calibrator is not None else {},
-        }
+        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):
@@ -143,7 +145,7 @@ def carrying_post_forward(module, output):
                     tensor.data = tensor.data.cpu()
                     continue
                 master = entry[1]
-                if key[1] != "_calibrator" or entry[2] != tensor._version:
+                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
diff --git a/modelopt/torch/quantization/utils/layerwise_calib.py b/modelopt/torch/quantization/utils/layerwise_calib.py
index 1fc8a37e6b7..82cb4669561 100644
--- a/modelopt/torch/quantization/utils/layerwise_calib.py
+++ b/modelopt/torch/quantization/utils/layerwise_calib.py
@@ -640,15 +640,27 @@ def _forward_peak_per_gpu(layer: nn.Module, args: tuple, kwargs: dict) -> dict[i
     }
 
 
+def _gpu_index(device: torch.device) -> int | None:
+    """CUDA device index of *device*; ``None`` if it is not a GPU."""
+    if device.type != "cuda":
+        return None
+    return torch.cuda.current_device() if device.index is None else device.index
+
+
 def _layer_batch_fit(
     model: nn.Module,
     layer: nn.Module,
+    gpus: set[int],
     seq_len: int,
     num_samples: int,
     memory_fraction: float,
     device: torch.device,
 ) -> int | None:
-    """Largest batch whose forward through *layer* fits; ``None`` if no GPU memory grows with it."""
+    """Largest batch whose forward through *layer* fits; ``None`` if no GPU memory grows with it.
+
+    Layers of the same kind on other *gpus* cost what *layer* costs on its own GPU, so those GPUs'
+    memory limits the batch too.
+    """
     # The first forward only warms up: it also allocates one-off buffers like cuBLAS's workspace.
     for _ in range(2):
         args, kwargs = _capture_layer_inputs(model, layer, 1, seq_len, device)
@@ -670,14 +682,19 @@ def _layer_batch_fit(
         torch.cuda.empty_cache()
         return 1
 
-    fits = []
+    costs = {}
     for gpu, peak in two.items():
         per_sample = peak - one.get(gpu, 0)
-        if per_sample <= 0:
-            continue
+        if per_sample > 0:
+            costs[gpu] = (per_sample, one.get(gpu, 0) - per_sample, cached_per_sample.get(gpu, 0))
+    home = _gpu_index(get_module_device(layer))
+    if home in costs:
+        for gpu in gpus:
+            costs.setdefault(gpu, costs[home])
+    fits = []
+    for gpu, (per_sample, fixed, cached) in costs.items():
         # Layerwise calibration caches every sample's inputs for this layer and the next.
-        available = get_available_gpu_memory(gpu) - 2 * num_samples * cached_per_sample.get(gpu, 0)
-        fixed = one.get(gpu, 0) - per_sample
+        available = get_available_gpu_memory(gpu) - 2 * num_samples * cached
         fits.append(int((memory_fraction * available - fixed) // per_sample))
     return min(fits, default=None)
 
@@ -698,7 +715,8 @@ def get_max_layerwise_batch_size(
     overstates the memory a batch needs and risks running out of it. Instead, the first layer
     of each kind (by module classes and parameter shapes: dense or MoE, full or linear attention,
     and per-layer widths such as Nemotron-NAS's) runs on one and two samples of *seq_len* tokens,
-    and the growth of its peak memory on each GPU is that GPU's cost per sample. A batch fits
+    and the growth of its peak memory on each GPU is that GPU's cost per sample; every other GPU
+    that runs a layer of that kind is charged the same. A batch fits
     when the layer's forward stays within *memory_fraction* of the GPU's available memory, after
     the inputs layerwise calibration caches for all *num_samples*; the rest is left for
     quantizers and Hessians. A layer that fits one sample but not two gives a batch of one. The
@@ -708,16 +726,21 @@ def get_max_layerwise_batch_size(
     batch_size = max_batch_size
     layers = LayerActivationCollector.get_decoder_layers(model)
     if layers and torch.cuda.is_available() and max_batch_size > 1:
-        kinds: dict[tuple, nn.Module] = {}
+        kinds: dict[tuple, list[nn.Module]] = {}
         for layer in layers:
             classes = tuple(sorted({type(m).__name__ for m in layer.modules()}))
-            kinds.setdefault((classes, tuple(tuple(p.shape) for p in layer.parameters())), layer)
+            shapes = tuple(tuple(p.shape) for p in layer.parameters())
+            kinds.setdefault((classes, shapes), []).append(layer)
         device = device or get_module_device(model)
         collector = LayerActivationCollector(model)
         collector._patch_all_layers(layers)
         try:
-            for layer in kinds.values():
-                fit = _layer_batch_fit(model, layer, seq_len, num_samples, memory_fraction, device)
+            for same_kind in kinds.values():
+                indexes = (_gpu_index(get_module_device(layer)) for layer in same_kind)
+                gpus = {gpu for gpu in indexes if gpu is not None}
+                fit = _layer_batch_fit(
+                    model, same_kind[0], gpus, seq_len, num_samples, memory_fraction, device
+                )
                 if fit is not None:
                     batch_size = min(batch_size, fit)
         finally:
diff --git a/modelopt/torch/utils/perf.py b/modelopt/torch/utils/perf.py
index a9ee748659b..196aba04287 100644
--- a/modelopt/torch/utils/perf.py
+++ b/modelopt/torch/utils/perf.py
@@ -99,9 +99,17 @@ def get_process_gpus() -> list[int]:
 
 
 def get_available_gpu_memory(gpu: int) -> int:
-    """Bytes a new allocation on ``gpu`` can get: the driver's free memory and our cached blocks."""
+    """Bytes a new allocation on ``gpu`` can get: the driver's free memory and our cached blocks.
+
+    Free pieces of partly used cached blocks are left out: a large allocation cannot use them,
+    and PyTorch cannot hand them back to the driver.
+    """
     free, _ = torch.cuda.mem_get_info(gpu)
-    return free + torch.cuda.memory_reserved(gpu) - torch.cuda.memory_allocated(gpu)
+    stats = torch.cuda.memory_stats(gpu)
+    cached = stats.get("reserved_bytes.all.current", 0) - stats.get(
+        "allocated_bytes.all.current", 0
+    )
+    return free + cached - stats.get("inactive_split_bytes.all.current", 0)
 
 
 def report_memory(name="", rank=0, device=None):
diff --git a/tests/gpu/torch/quantization/plugins/test_accelerate_gpu.py b/tests/gpu/torch/quantization/plugins/test_accelerate_gpu.py
index 4808377ace7..c92c9b54c1a 100644
--- a/tests/gpu/torch/quantization/plugins/test_accelerate_gpu.py
+++ b/tests/gpu/torch/quantization/plugins/test_accelerate_gpu.py
@@ -158,6 +158,29 @@ def test_offloaded_quantizer_keeps_what_its_setter_wrote(tmp_path):
     assert torch.equal(quantizer._amax, amax)
 
 
+def test_offloaded_layer_keeps_bias_calibration_off_the_gpu(tmp_path):
+    """A static-bias quantizer's bias calibrator leaves the GPU with the offloaded layer's weights."""
+    model, _, _, inputs = make_cpu_offloaded_model(tmp_path)
+    input_cfg = {"num_bits": 8, "axis": None, "bias": {-1: None, "type": "static"}}
+    quant_cfg = {
+        "quant_cfg": [
+            {"quantizer_name": "*", "enable": False},
+            {"quantizer_name": "*input_quantizer", "cfg": input_cfg},
+        ],
+        "algorithm": "max",
+    }
+    mtq.quantize(model, make_layerwise_cfg(quant_cfg), lambda model: model(inputs))
+    bias_state = [
+        tensor
+        for quantizer in model.model.layers[0].modules()
+        if isinstance(quantizer, TensorQuantizer) and quantizer._bias_calibrator is not None
+        for tensor in vars(quantizer._bias_calibrator).values()
+        if isinstance(tensor, torch.Tensor)
+    ]
+    assert bias_state
+    assert all(tensor.device.type == "cpu" for tensor in bias_state)
+
+
 @pytest.mark.parametrize("use_checkpoint", [False, True], ids=["no_ckpt", "ckpt"])
 def test_layerwise_calibrate_cpu_offloaded(tmp_path, use_checkpoint):
     """Layerwise calibration on CPU-offloaded model matches GPU-only reference."""
diff --git a/tests/gpu/torch/quantization/test_gptq.py b/tests/gpu/torch/quantization/test_gptq.py
index 5ea632a13e3..286883969f9 100644
--- a/tests/gpu/torch/quantization/test_gptq.py
+++ b/tests/gpu/torch/quantization/test_gptq.py
@@ -33,6 +33,7 @@
     gptq_blockwise_update_fused_scalar,
     update_hessian,
 )
+from modelopt.torch.utils import get_available_gpu_memory
 from modelopt.torch.utils.dataset_utils import create_forward_loop, get_dataset_dataloader
 
 RAND_SEED = 42
@@ -99,6 +100,20 @@ def test_a_measured_calibration_forward_sets_the_hessian_reserve(
     assert helper.hessian.device == (torch.device("cuda", 0) if stays else torch.device("cpu"))
 
 
+def test_available_gpu_memory_is_what_one_allocation_can_get(monkeypatch):
+    """Free pieces of partly used cached blocks do not count: one allocation of all the memory
+    reported beyond the driver's free memory comes from the cache."""
+    torch.cuda.empty_cache()
+    whole = torch.empty(64 * 2**20, dtype=torch.uint8, device="cuda")
+    pieces = [torch.empty(3 * 2**20, dtype=torch.uint8, device="cuda") for _ in range(24)]
+    del whole, pieces[::2]  # a whole free block, and free pieces of partly used ones
+    monkeypatch.setattr(torch.cuda, "mem_get_info", lambda device=None: (0, 2**40))
+    available = get_available_gpu_memory(torch.cuda.current_device())
+    reserved = torch.cuda.memory_reserved()
+    block = torch.empty(available, dtype=torch.uint8, device="cuda")  # noqa: F841
+    assert torch.cuda.memory_reserved() == reserved
+
+
 def test_gptq_records_its_calibration_forward_memory(monkeypatch):
     """The next layer's Hessian reserve covers what this layer's calibration forward added, and
     not an earlier peak such as the previous layer's Hessians."""
diff --git a/tests/gpu/torch/quantization/test_layerwise_calibrate.py b/tests/gpu/torch/quantization/test_layerwise_calibrate.py
index 3da2b016de7..0ad88b1fa8a 100644
--- a/tests/gpu/torch/quantization/test_layerwise_calibrate.py
+++ b/tests/gpu/torch/quantization/test_layerwise_calibrate.py
@@ -414,3 +414,28 @@ def test_layerwise_batch_probe_falls_back_to_one_sample(monkeypatch, where):
 
     assert get_max_layerwise_batch_size(model, seq_len=8, max_batch_size=64, num_samples=1) == 1
     assert not any(hasattr(layer, "_layerwise_calib") for layer in model.layers)
+
+
+class _MultiGPUModel(_SimpleTransformerModel):
+    """Decoder layers spread over GPUs, as a ``device_map`` places them."""
+
+    def forward(self, x, **kwargs):
+        x = self.embed(x)
+        for layer in self.layers:
+            x = layer(x.to(layer.norm.weight.device))
+        return x
+
+
+def test_layerwise_batch_probe_fits_every_gpu_a_kind_of_layer_runs_on(need_2_gpus, monkeypatch):
+    """Of two identical layers, only the first is probed, but the second's GPU, with less memory,
+    sets the batch."""
+    _register_test_discoverer(monkeypatch)
+    model = _MultiGPUModel(n_layers=2, dim=256).cuda()
+    model.layers[1].to("cuda:1")
+
+    def probe(budgets):
+        monkeypatch.setattr(layerwise_calib, "get_available_gpu_memory", lambda gpu: budgets[gpu])
+        return get_max_layerwise_batch_size(model, seq_len=512, max_batch_size=4096, num_samples=1)
+
+    small, large = 256 * 2**20, 1024 * 2**20
+    assert probe({0: large, 1: small}) == probe({0: small, 1: small}) < probe({0: large, 1: large})