From 2a0039f795e75d4717a2427e560053b2a1b91bc2 Mon Sep 17 00:00:00 2001
From: Chenjie Luo <chenjiel@nvidia.com>
Date: Thu, 22 Oct 2026 05:23:49 +0000
Subject: [PATCH 1/7] Run fused-MoE expert GPTQ updates on all of this
 process's GPUs

In a layerwise run of a model sharded over several GPUs, only the GPU
holding the current layer worked on its expert updates while the others
sat idle. The updates of a fused-MoE layer's experts are independent, so
helpers whose update reads only the weight, its Hessian and the format's
codec (the GGML formats) now run in one worker thread per GPU this process
uses that has room for an update, each writing its result back to its
weight. The first update runs in the caller's thread so lazy setup, such as
loading a format's CUDA extension, happens once before the workers start.
The results are bit-identical to running them in turn on one GPU.

Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com>
Signed-off-by: Chenjie Luo <chenjiel@nvidia.com>
---
 modelopt/torch/quantization/ggml/gptq.py      | 10 ++-
 modelopt/torch/quantization/model_calib.py    | 20 ++---
 .../torch/quantization/utils/calib_utils.py   | 82 +++++++++++++++++--
 tests/gpu/torch/quantization/test_gptq.py     | 53 ++++++++++++
 4 files changed, 143 insertions(+), 22 deletions(-)

diff --git a/modelopt/torch/quantization/ggml/gptq.py b/modelopt/torch/quantization/ggml/gptq.py
index 92c56ed5003..0ed4bc3055c 100644
--- a/modelopt/torch/quantization/ggml/gptq.py
+++ b/modelopt/torch/quantization/ggml/gptq.py
@@ -78,11 +78,15 @@ class GGMLGPTQHelper(GPTQHelper):
     those exact codes rather than encoding the GPTQ'd weight again, which would not return them.
     """
 
-    def update_weights(self, block_size, perc_damp, weight=None):
+    # The update reads only the weight, its Hessian and the format's codec.
+    UPDATES_ON_ANY_DEVICE = True
+
+    def update_weights(self, block_size, perc_damp, weight=None, device=None):
         """Run the GPTQ update, then pin the chosen payload to the weight quantizer."""
-        super().update_weights(block_size, perc_damp, weight)
+        super().update_weights(block_size, perc_damp, weight, device)
         quantizer = self.weight_quantizer
-        pin_packed_weight(quantizer, quantizer.num_bits, self._packed)
+        weight = self.module.weight if weight is None else weight
+        pin_packed_weight(quantizer, quantizer.num_bits, self._packed.to(weight.device))
         self._packed = None
 
     def _blockwise_update(self, block_size):
diff --git a/modelopt/torch/quantization/model_calib.py b/modelopt/torch/quantization/model_calib.py
index 3a210a3545a..725770c11e6 100644
--- a/modelopt/torch/quantization/model_calib.py
+++ b/modelopt/torch/quantization/model_calib.py
@@ -80,6 +80,7 @@
     GPTQHelper,
     bfloat16_hessian_gemm,
     track_calib_forward_memory,
+    update_expert_weights,
 )
 
 __all__ = [
@@ -2648,18 +2649,17 @@ def _expert_hook(_input_quantizer, _args, output):
             handle.update_weights(block_size, perc_damp)
         handle.free()
     for name, module in fused_experts:
-        updated = unreached = 0
         with enable_weight_access_and_writeback(module, model, names):
-            for weight, quantizer in module.iter_weights_for_calibration():
-                handle = expert_handles.pop(id(quantizer), None)
-                if handle is not None:
-                    unreached += handle.n_samples == 0
-                    handle.update_weights(block_size, perc_damp, weight)
-                    handle.free()
-                    updated += 1
-        if updated:
+            work = [
+                (expert_handles.pop(id(quantizer)), weight)
+                for weight, quantizer in module.iter_weights_for_calibration()
+                if id(quantizer) in expert_handles
+            ]
+            unreached = sum(handle.n_samples == 0 for handle, _ in work)
+            update_expert_weights(work, block_size, perc_damp)
+        if work:
             print_rank_0(
-                f"[{name}] GPTQ updated {updated} expert weights, {unreached} of them never "
+                f"[{name}] GPTQ updated {len(work)} expert weights, {unreached} of them never "
                 "reached by calibration"
             )
     del gptq_handles, expert_handles
diff --git a/modelopt/torch/quantization/utils/calib_utils.py b/modelopt/torch/quantization/utils/calib_utils.py
index 8b045de3f15..9a0927b9f7f 100644
--- a/modelopt/torch/quantization/utils/calib_utils.py
+++ b/modelopt/torch/quantization/utils/calib_utils.py
@@ -39,6 +39,8 @@
 """GPTQ helper and Hessian utilities for calibration."""
 
 import functools
+import queue
+from concurrent.futures import ThreadPoolExecutor
 from contextlib import contextmanager
 
 import torch
@@ -106,13 +108,70 @@ def _hessian_offload_device(nbytes: int) -> torch.device:
     rescale of the whole matrix. Only GPUs this process already holds memory on are candidates,
     and only outside distributed runs, so no other job's or rank's GPU is touched.
     """
-    if not (dist.is_available() and dist.is_initialized() and dist.get_world_size() > 1):
-        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)
+    for gpu in sorted(_own_gpus(), key=get_available_gpu_memory, reverse=True):
+        if _gpu_has_room(gpu, nbytes):
+            return torch.device("cuda", gpu)
     return torch.device("cpu")
 
 
+def _own_gpus() -> list[int]:
+    """GPUs this process already holds memory on; none in a multi-rank run, where they may be another rank's."""
+    if dist.is_available() and dist.is_initialized() and dist.get_world_size() > 1:
+        return []
+    return get_process_gpus()
+
+
+def update_expert_weights(work, block_size, perc_damp):
+    """Run the GPTQ updates of ``(helper, weight)`` pairs, spread over this process's GPUs.
+
+    Expert updates are independent, and in a layerwise run of a model sharded over several GPUs
+    only the current layer's GPU works. Helpers whose update can run on any device run in one
+    worker thread per GPU this process uses that has room for an update, each writing its result
+    back to its weight; the others, or all of them when fewer than two GPUs qualify, run in turn on
+    their weight's device.
+    """
+    movable = [(h, w) for h, w in work if h.UPDATES_ON_ANY_DEVICE and w.is_cuda]
+    # A worker holds one expert at a time: its weight, Hessian, inverse and their temporaries.
+    room = 8 * max((4 * h.hessian.numel() for h, _ in movable), default=0)
+    gpus = [gpu for gpu in _own_gpus() if get_available_gpu_memory(gpu) > room]
+    if len(gpus) < 2 or len(movable) < 2:
+        movable = []
+    moved = {id(helper) for helper, _ in movable}
+    for helper, weight in work:
+        if id(helper) not in moved:
+            helper.update_weights(block_size, perc_damp, weight)
+            helper.free()
+    if not movable:
+        return
+
+    # The first update runs here, so lazy setup such as loading a format's CUDA extension, which is
+    # not thread-safe, happens once before the workers start.
+    helper, weight = movable[0]
+    helper.update_weights(block_size, perc_damp, weight)
+    helper.free()
+    pending: queue.SimpleQueue = queue.SimpleQueue()
+    for item in movable[1:]:
+        pending.put(item)
+
+    def worker(gpu):
+        torch.cuda.set_device(gpu)  # per thread
+        device = torch.device("cuda", gpu)
+        while True:
+            try:
+                helper, weight = pending.get_nowait()
+            except queue.Empty:
+                break
+            helper.update_weights(block_size, perc_damp, weight, device=device)
+            helper.free()
+        torch.cuda.synchronize(device)
+
+    with ThreadPoolExecutor(len(gpus)) as pool:
+        for future in [pool.submit(worker, gpu) for gpu in gpus]:
+            future.result()
+    for device in {weight.device for _, weight in movable}:
+        torch.cuda.synchronize(device)
+
+
 def update_hessian(input, hessian, n_samples, input_dtype=torch.float32):
     """Update hessian matrix with new input samples using incremental formula.
 
@@ -255,6 +314,9 @@ class GPTQHelper:
     """
 
     CACHE_NAME = "_forward_no_gptq_hessian"
+    # Whether update_weights() may run on a device other than its weight's: true only for helpers
+    # whose update reads no quantizer state, which lives on the weight's device.
+    UPDATES_ON_ANY_DEVICE = False
 
     def __init__(
         self,
@@ -359,20 +421,22 @@ def free(self):
         self.weight = None
         self.h_inv = None
 
-    def update_weights(self, block_size, perc_damp, weight=None):
+    def update_weights(self, block_size, perc_damp, weight=None, device=None):
         """Run GPTQ blockwise weight update on this helper's weight.
 
         Populates ``self.weight`` and ``self.h_inv``, runs the blockwise update, logs MSE, and
         overwrites ``weight`` (default ``module.weight``) in place, so an expert's slice of a
-        fused 3-D parameter is written back into that parameter.
+        fused 3-D parameter is written back into that parameter. The update runs on ``device``,
+        by default the weight's; only a helper with ``UPDATES_ON_ANY_DEVICE`` may pass another.
         """
         weight = self.module.weight if weight is None else weight
-        hessian = self.hessian.to(weight.device)
-        self.weight = weight.data.float().clone()
+        device = weight.device if device is None else device
+        hessian = self.hessian.to(device)
+        self.weight = weight.data.to(device, torch.float32, copy=True)
         self._prepare_hessian_inverse(hessian, perc_damp)
         self._blockwise_update(block_size)
         if self.log_mse:
-            self._print_mse_error(hessian, weight)
+            self._print_mse_error(hessian, weight.data.to(device))
         weight.data.copy_(self.weight.reshape(weight.shape))
 
     # ------------------------------------------------------------------
diff --git a/tests/gpu/torch/quantization/test_gptq.py b/tests/gpu/torch/quantization/test_gptq.py
index d2733293234..667f7e19d47 100644
--- a/tests/gpu/torch/quantization/test_gptq.py
+++ b/tests/gpu/torch/quantization/test_gptq.py
@@ -18,12 +18,16 @@
 
 import pytest
 import torch
+from _test_utils.torch.quantization.models import TOY_MOE_HIDDEN, ToyFusedExperts, ToyMoE
 from _test_utils.torch.transformers_models import get_tiny_llama
 from conftest import requires_triton
 
 import modelopt.torch.quantization as mtq
 from modelopt.torch.export.unified_export_hf import _export_quantized_weight
+from modelopt.torch.quantization.ggml.common import pinned_packed_weight
+from modelopt.torch.quantization.ggml.gptq import GGMLGPTQHelper
 from modelopt.torch.quantization.model_calib import gptq
+from modelopt.torch.quantization.nn import QuantModuleRegistry
 from modelopt.torch.quantization.qtensor.nvfp4_tensor import NVFP4QTensor
 from modelopt.torch.quantization.utils import calib_utils
 from modelopt.torch.quantization.utils.calib_utils import (
@@ -61,6 +65,55 @@ def test_hessian_of_a_full_gpu_goes_to_an_idle_gpu(need_2_gpus, monkeypatch):
     torch.testing.assert_close(helper.hessian.to("cuda:0"), reference)
 
 
+def test_expert_updates_on_another_gpu_match_one_gpu(need_2_gpus, monkeypatch):
+    """GGML expert updates spread over two GPUs give the one-GPU weights and codes, pinned at home."""
+    pytest.importorskip("transformers")  # the fused-experts quant wrapper is the HF plugin's
+    torch.manual_seed(RAND_SEED)
+    reference = ToyMoE().cuda()
+    inputs = torch.randn(96, TOY_MOE_HIDDEN, device="cuda")
+    shard = torch.zeros(1, device="cuda:1")  # noqa: F841  # the model also lives on cuda:1
+    config = {
+        "quant_cfg": [
+            {"quantizer_name": "*", "enable": False},
+            {
+                "quantizer_name": "*weight_quantizer",
+                "cfg": {"num_bits": "iq2_xxs", "block_sizes": {-1: 256}, "backend": "ggml"},
+            },
+        ],
+        "algorithm": {"method": "gptq", "block_size": 256, "perc_damp": 0.3},
+    }
+    devices = []
+    update_weights = GGMLGPTQHelper.update_weights
+
+    def recording_update_weights(self, *args, device=None, **kwargs):
+        devices.append(device)
+        return update_weights(self, *args, device=device, **kwargs)
+
+    monkeypatch.setattr(GGMLGPTQHelper, "update_weights", recording_update_weights)
+    experts = []
+    try:
+        for gpus in ([0], [0, 1]):
+            monkeypatch.setattr(calib_utils, "get_process_gpus", lambda gpus=gpus: gpus)
+            model = copy.deepcopy(reference)
+            mtq.quantize(model, config, forward_loop=lambda m: m(inputs))
+            experts.append(model.mlp.experts)
+    finally:
+        QuantModuleRegistry.unregister(ToyFusedExperts)
+
+    assert torch.device("cuda", 1) in devices
+    one_gpu, two_gpus = experts
+    for name in ("gate_up_proj", "down_proj"):
+        assert torch.equal(getattr(one_gpu, name), getattr(two_gpus, name))
+        for q1, q2 in zip(
+            getattr(one_gpu, f"{name}_weight_quantizers"),
+            getattr(two_gpus, f"{name}_weight_quantizers"),
+        ):
+            assert pinned_packed_weight(q2, "iq2_xxs").device == getattr(two_gpus, name).device
+            assert torch.equal(
+                pinned_packed_weight(q1, "iq2_xxs"), pinned_packed_weight(q2, "iq2_xxs")
+            )
+
+
 @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."""

From 46262c8d7ea60e6f16f8d9d752bb92903c31b3c8 Mon Sep 17 00:00:00 2001
From: Chenjie Luo <chenjiel@nvidia.com>
Date: Fri, 30 Oct 2026 16:56:47 +0000
Subject: [PATCH 2/7] Disable autograd in the GPTQ expert-update worker threads

Grad mode is per thread, so the worker threads did not inherit the
no_grad that gptq() runs under, unlike the imatrix encode workers.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Signed-off-by: Chenjie Luo <chenjiel@nvidia.com>
---
 modelopt/torch/quantization/utils/calib_utils.py | 3 ++-
 1 file changed, 2 insertions(+), 1 deletion(-)

diff --git a/modelopt/torch/quantization/utils/calib_utils.py b/modelopt/torch/quantization/utils/calib_utils.py
index 9a0927b9f7f..b4449b4cdcf 100644
--- a/modelopt/torch/quantization/utils/calib_utils.py
+++ b/modelopt/torch/quantization/utils/calib_utils.py
@@ -153,8 +153,9 @@ def update_expert_weights(work, block_size, perc_damp):
     for item in movable[1:]:
         pending.put(item)
 
+    @torch.no_grad()  # grad mode is per thread, like the current device
     def worker(gpu):
-        torch.cuda.set_device(gpu)  # per thread
+        torch.cuda.set_device(gpu)
         device = torch.device("cuda", gpu)
         while True:
             try:

From 88de90b1ad28cd58c883c89a50647054c4d34a2e Mon Sep 17 00:00:00 2001
From: Chenjie Luo <chenjiel@nvidia.com>
Date: Fri, 30 Oct 2026 17:46:18 +0000
Subject: [PATCH 3/7] Keep GPTQ expert updates' modes, memory and weight
 lifetime in check

- The worker threads take the caller's inference mode as well as grad
  mode, as the imatrix encode workers do, so the threaded path creates its
  tensors in the same mode as the sequential one.
- The room check counts each expert's float32 weight copy and a conversion
  temporary, and the codec's tables and scratch, not only its Hessian. A
  projection wider than its input, such as a gate/up projection whose down
  projection stays on its own GPU, needed more than the eight Hessians the
  check allowed.
- gptq() drops its expert weight slices before leaving each module's
  weight window. Under offload they kept that module's onloaded weights on
  the GPU while the next module's were onloaded.
- The two-GPU test feeds each expert 512 columns, two GGML groups, so its
  Hessian shapes the weights and codes it compares, and also runs under
  inference mode.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Signed-off-by: Chenjie Luo <chenjiel@nvidia.com>
---
 modelopt/torch/quantization/model_calib.py    |  7 +++--
 .../torch/quantization/utils/calib_utils.py   | 25 ++++++++++-------
 .../_test_utils/torch/quantization/models.py  | 18 ++++++------
 tests/gpu/torch/quantization/test_gptq.py     | 27 ++++++++++--------
 .../unit/torch/quantization/test_ggml_gptq.py | 28 +++++++++++++++++++
 5 files changed, 72 insertions(+), 33 deletions(-)

diff --git a/modelopt/torch/quantization/model_calib.py b/modelopt/torch/quantization/model_calib.py
index 725770c11e6..e570fa61128 100644
--- a/modelopt/torch/quantization/model_calib.py
+++ b/modelopt/torch/quantization/model_calib.py
@@ -2655,11 +2655,14 @@ def _expert_hook(_input_quantizer, _args, output):
                 for weight, quantizer in module.iter_weights_for_calibration()
                 if id(quantizer) in expert_handles
             ]
+            updated = len(work)
             unreached = sum(handle.n_samples == 0 for handle, _ in work)
             update_expert_weights(work, block_size, perc_damp)
-        if work:
+            # Its weight slices would keep this module's onloaded weights alive past the window.
+            del work
+        if updated:
             print_rank_0(
-                f"[{name}] GPTQ updated {len(work)} expert weights, {unreached} of them never "
+                f"[{name}] GPTQ updated {updated} expert weights, {unreached} of them never "
                 "reached by calibration"
             )
     del gptq_handles, expert_handles
diff --git a/modelopt/torch/quantization/utils/calib_utils.py b/modelopt/torch/quantization/utils/calib_utils.py
index b4449b4cdcf..00bd00c5745 100644
--- a/modelopt/torch/quantization/utils/calib_utils.py
+++ b/modelopt/torch/quantization/utils/calib_utils.py
@@ -131,8 +131,11 @@ def update_expert_weights(work, block_size, perc_damp):
     their weight's device.
     """
     movable = [(h, w) for h, w in work if h.UPDATES_ON_ANY_DEVICE and w.is_cuda]
-    # A worker holds one expert at a time: its weight, Hessian, inverse and their temporaries.
-    room = 8 * max((4 * h.hessian.numel() for h, _ in movable), default=0)
+    # A worker holds one expert at a time: its Hessian, the inverse and their temporaries, the
+    # float32 copy of its weight with a dtype-conversion temporary (a projection much wider than
+    # its input makes the weight the larger part), and the codec's tables and scratch, about
+    # 35 MiB on a GPU's first update.
+    room = 2**26 + max((4 * (8 * h.hessian.numel() + 2 * w.numel()) for h, w in movable), default=0)
     gpus = [gpu for gpu in _own_gpus() if get_available_gpu_memory(gpu) > room]
     if len(gpus) < 2 or len(movable) < 2:
         movable = []
@@ -152,18 +155,20 @@ def update_expert_weights(work, block_size, perc_damp):
     pending: queue.SimpleQueue = queue.SimpleQueue()
     for item in movable[1:]:
         pending.put(item)
+    inference = torch.is_inference_mode_enabled()
 
-    @torch.no_grad()  # grad mode is per thread, like the current device
     def worker(gpu):
+        # The device, grad mode and inference mode are per thread: take the caller's.
         torch.cuda.set_device(gpu)
         device = torch.device("cuda", gpu)
-        while True:
-            try:
-                helper, weight = pending.get_nowait()
-            except queue.Empty:
-                break
-            helper.update_weights(block_size, perc_damp, weight, device=device)
-            helper.free()
+        with torch.inference_mode() if inference else torch.no_grad():
+            while True:
+                try:
+                    helper, weight = pending.get_nowait()
+                except queue.Empty:
+                    break
+                helper.update_weights(block_size, perc_damp, weight, device=device)
+                helper.free()
         torch.cuda.synchronize(device)
 
     with ThreadPoolExecutor(len(gpus)) as pool:
diff --git a/tests/_test_utils/torch/quantization/models.py b/tests/_test_utils/torch/quantization/models.py
index 2e1875337df..b8542545895 100644
--- a/tests/_test_utils/torch/quantization/models.py
+++ b/tests/_test_utils/torch/quantization/models.py
@@ -174,15 +174,13 @@ def get_input(self, batch_size=2):
 class ToyFusedExperts(nn.Module):
     """HF transformers 5 fused-experts layout: one ``F.linear`` pair per routed expert."""
 
-    def __init__(self):
+    def __init__(self, hidden=TOY_MOE_HIDDEN, intermediate=TOY_MOE_INTERMEDIATE):
         super().__init__()
         self.num_experts = TOY_MOE_EXPERTS
         self.gate_up_proj = nn.Parameter(
-            torch.randn(TOY_MOE_EXPERTS, 2 * TOY_MOE_INTERMEDIATE, TOY_MOE_HIDDEN) * 0.02
-        )
-        self.down_proj = nn.Parameter(
-            torch.randn(TOY_MOE_EXPERTS, TOY_MOE_HIDDEN, TOY_MOE_INTERMEDIATE) * 0.02
+            torch.randn(TOY_MOE_EXPERTS, 2 * intermediate, hidden) * 0.02
         )
+        self.down_proj = nn.Parameter(torch.randn(TOY_MOE_EXPERTS, hidden, intermediate) * 0.02)
 
     def forward(self, hidden_states, expert_of_token):
         output = torch.zeros_like(hidden_states)
@@ -198,10 +196,10 @@ def forward(self, hidden_states, expert_of_token):
 class ToyMoEBlock(nn.Module):
     """Routes token ``t`` to expert ``t % 3``, so the last expert receives no token."""
 
-    def __init__(self):
+    def __init__(self, hidden=TOY_MOE_HIDDEN, intermediate=TOY_MOE_INTERMEDIATE):
         super().__init__()
-        self.experts = ToyFusedExperts()
-        self.shared_expert = nn.Linear(TOY_MOE_HIDDEN, TOY_MOE_HIDDEN, bias=False)
+        self.experts = ToyFusedExperts(hidden, intermediate)
+        self.shared_expert = nn.Linear(hidden, hidden, bias=False)
 
     def forward(self, hidden_states):
         routed = self.experts(hidden_states, torch.arange(len(hidden_states)) % 3)
@@ -209,9 +207,9 @@ def forward(self, hidden_states):
 
 
 class ToyMoE(nn.Module):
-    def __init__(self):
+    def __init__(self, hidden=TOY_MOE_HIDDEN, intermediate=TOY_MOE_INTERMEDIATE):
         super().__init__()
-        self.mlp = ToyMoEBlock()
+        self.mlp = ToyMoEBlock(hidden, intermediate)
 
     def forward(self, hidden_states):
         return self.mlp(hidden_states)
diff --git a/tests/gpu/torch/quantization/test_gptq.py b/tests/gpu/torch/quantization/test_gptq.py
index 667f7e19d47..f86e015f803 100644
--- a/tests/gpu/torch/quantization/test_gptq.py
+++ b/tests/gpu/torch/quantization/test_gptq.py
@@ -18,7 +18,7 @@
 
 import pytest
 import torch
-from _test_utils.torch.quantization.models import TOY_MOE_HIDDEN, ToyFusedExperts, ToyMoE
+from _test_utils.torch.quantization.models import ToyFusedExperts, ToyMoE
 from _test_utils.torch.transformers_models import get_tiny_llama
 from conftest import requires_triton
 
@@ -65,12 +65,14 @@ def test_hessian_of_a_full_gpu_goes_to_an_idle_gpu(need_2_gpus, monkeypatch):
     torch.testing.assert_close(helper.hessian.to("cuda:0"), reference)
 
 
-def test_expert_updates_on_another_gpu_match_one_gpu(need_2_gpus, monkeypatch):
-    """GGML expert updates spread over two GPUs give the one-GPU weights and codes, pinned at home."""
+@pytest.mark.parametrize("inference_mode", [False, True])
+def test_expert_updates_on_another_gpu_match_one_gpu(need_2_gpus, monkeypatch, inference_mode):
+    """GGML expert updates spread over two GPUs give the one-GPU weights and codes, pinned at home,
+    also on inference tensors under the caller's inference mode, which worker threads do not inherit.
+    """
     pytest.importorskip("transformers")  # the fused-experts quant wrapper is the HF plugin's
-    torch.manual_seed(RAND_SEED)
-    reference = ToyMoE().cuda()
-    inputs = torch.randn(96, TOY_MOE_HIDDEN, device="cuda")
+    # Two GGML groups per expert input, so the Hessian carries the first group's error to the next.
+    width = 512
     shard = torch.zeros(1, device="cuda:1")  # noqa: F841  # the model also lives on cuda:1
     config = {
         "quant_cfg": [
@@ -92,11 +94,14 @@ def recording_update_weights(self, *args, device=None, **kwargs):
     monkeypatch.setattr(GGMLGPTQHelper, "update_weights", recording_update_weights)
     experts = []
     try:
-        for gpus in ([0], [0, 1]):
-            monkeypatch.setattr(calib_utils, "get_process_gpus", lambda gpus=gpus: gpus)
-            model = copy.deepcopy(reference)
-            mtq.quantize(model, config, forward_loop=lambda m: m(inputs))
-            experts.append(model.mlp.experts)
+        with torch.inference_mode(inference_mode):
+            inputs = torch.randn(96, width, device="cuda")
+            for gpus in ([0], [0, 1]):
+                monkeypatch.setattr(calib_utils, "get_process_gpus", lambda gpus=gpus: gpus)
+                torch.manual_seed(RAND_SEED)
+                model = ToyMoE(width, width).cuda()
+                mtq.quantize(model, config, forward_loop=lambda m: m(inputs))
+                experts.append(model.mlp.experts)
     finally:
         QuantModuleRegistry.unregister(ToyFusedExperts)
 
diff --git a/tests/unit/torch/quantization/test_ggml_gptq.py b/tests/unit/torch/quantization/test_ggml_gptq.py
index fc5a01ca080..a0c0f480a43 100644
--- a/tests/unit/torch/quantization/test_ggml_gptq.py
+++ b/tests/unit/torch/quantization/test_ggml_gptq.py
@@ -14,9 +14,11 @@
 # limitations under the License.
 """GPTQ for GGML block formats: the group update, the helper, and the payload pin."""
 
+import contextlib
 import copy
 import dataclasses
 import io
+import weakref
 from types import SimpleNamespace
 
 import pytest
@@ -27,6 +29,7 @@
 import modelopt.torch.opt as mto
 import modelopt.torch.quantization as mtq
 from modelopt.torch.export.unified_export_megatron import GPTModelExporter
+from modelopt.torch.quantization import model_calib
 from modelopt.torch.quantization.ggml import GGML_FORMAT_REGISTRY
 from modelopt.torch.quantization.ggml.common import _row_keys, pinned_packed_weight
 from modelopt.torch.quantization.ggml.gptq import gptq_group_update
@@ -291,6 +294,31 @@ def same(a, b):
     assert not same(single_pass[1], layerwise[1])
 
 
+def test_gptq_drops_each_fused_experts_weights_before_the_next_window(monkeypatch, moe):
+    """No expert weight GPTQ updates outlives its module's weight window: under offload it would
+    keep that module's onloaded weights on the GPU while the next module's are onloaded."""
+    updated = []
+    update_expert_weights = model_calib.update_expert_weights
+    window = model_calib.enable_weight_access_and_writeback
+
+    def recording_update_expert_weights(work, *args):
+        updated.extend(weakref.ref(weight) for _, weight in work)
+        update_expert_weights(work, *args)
+
+    @contextlib.contextmanager
+    def checked_window(*args, **kwargs):
+        assert all(ref() is None for ref in updated)
+        with window(*args, **kwargs):
+            yield
+
+    monkeypatch.setattr(model_calib, "update_expert_weights", recording_update_expert_weights)
+    monkeypatch.setattr(model_calib, "enable_weight_access_and_writeback", checked_window)
+    inputs = _inputs()[:96, :TOY_MOE_HIDDEN]
+    mtq.quantize(_MoEStack(moe, ToyMoE()), _int4_config(GPTQ), forward_loop=lambda m: m(inputs))
+
+    assert updated
+
+
 def test_pin_outlives_the_weight_tensor():
     model, _ = _gptq_model()
     # Same values in a new tensor, as an offload round trip or a layerwise resume leaves them.

From df3f78f70c9687d0f3505e6fe8899425a8d27008 Mon Sep 17 00:00:00 2001
From: Chenjie Luo <chenjiel@nvidia.com>
Date: Fri, 30 Oct 2026 18:55:14 +0000
Subject: [PATCH 4/7] Move GPTQ expert updates only where the room check covers
 the encoder

The room an expert update is checked for was measured with the GGML CUDA
encoders. Without the extension, the PyTorch fallback's scratch grows with
block_chunk_size and can exceed it (IQ1_S at 1024 x 512 peaks at 96 MiB
against a 76 MiB check), so GGML updates move to another GPU only when the
extension is loaded. UPDATES_ON_ANY_DEVICE becomes updates_on_any_device(),
since the answer now depends on the run. A GPU test checks that an update a
helper lets move peaks within the room it is admitted with.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Signed-off-by: Chenjie Luo <chenjiel@nvidia.com>
---
 modelopt/torch/quantization/ggml/gptq.py      | 10 +++-
 .../torch/quantization/utils/calib_utils.py   | 33 +++++++++----
 tests/gpu/torch/quantization/test_gptq.py     | 47 +++++++++++++++++++
 3 files changed, 78 insertions(+), 12 deletions(-)

diff --git a/modelopt/torch/quantization/ggml/gptq.py b/modelopt/torch/quantization/ggml/gptq.py
index 0ed4bc3055c..7b5cbe91295 100644
--- a/modelopt/torch/quantization/ggml/gptq.py
+++ b/modelopt/torch/quantization/ggml/gptq.py
@@ -19,6 +19,7 @@
 
 import torch
 
+from ..extensions import get_cuda_ext_ggml
 from ..utils.calib_utils import GPTQHelper, register_gptq_helper
 from .common import floored_second_moments, pin_packed_weight
 from .registry import GGML_FORMAT_REGISTRY
@@ -78,8 +79,13 @@ class GGMLGPTQHelper(GPTQHelper):
     those exact codes rather than encoding the GPTQ'd weight again, which would not return them.
     """
 
-    # The update reads only the weight, its Hessian and the format's codec.
-    UPDATES_ON_ANY_DEVICE = True
+    def updates_on_any_device(self):
+        """Whether the CUDA encoders are available.
+
+        The update reads only the weight, its Hessian and the format's codec, but the PyTorch
+        fallback's scratch grows with ``block_chunk_size``, past the room the update is checked for.
+        """
+        return get_cuda_ext_ggml() is not None
 
     def update_weights(self, block_size, perc_damp, weight=None, device=None):
         """Run the GPTQ update, then pin the chosen payload to the weight quantizer."""
diff --git a/modelopt/torch/quantization/utils/calib_utils.py b/modelopt/torch/quantization/utils/calib_utils.py
index 00bd00c5745..d334584373d 100644
--- a/modelopt/torch/quantization/utils/calib_utils.py
+++ b/modelopt/torch/quantization/utils/calib_utils.py
@@ -121,6 +121,16 @@ def _own_gpus() -> list[int]:
     return get_process_gpus()
 
 
+def _update_room(helper, weight) -> int:
+    """Bytes a GPU needs to run ``helper``'s update of ``weight``, as measured with the CUDA codecs.
+
+    The Hessian, the inverse and their temporaries; the float32 copy of the weight with a
+    dtype-conversion temporary (a projection much wider than its input makes the weight the larger
+    part); and the codec's tables and scratch, about 35 MiB on a GPU's first update.
+    """
+    return 2**26 + 4 * (8 * helper.hessian.numel() + 2 * weight.numel())
+
+
 def update_expert_weights(work, block_size, perc_damp):
     """Run the GPTQ updates of ``(helper, weight)`` pairs, spread over this process's GPUs.
 
@@ -130,12 +140,9 @@ def update_expert_weights(work, block_size, perc_damp):
     back to its weight; the others, or all of them when fewer than two GPUs qualify, run in turn on
     their weight's device.
     """
-    movable = [(h, w) for h, w in work if h.UPDATES_ON_ANY_DEVICE and w.is_cuda]
-    # A worker holds one expert at a time: its Hessian, the inverse and their temporaries, the
-    # float32 copy of its weight with a dtype-conversion temporary (a projection much wider than
-    # its input makes the weight the larger part), and the codec's tables and scratch, about
-    # 35 MiB on a GPU's first update.
-    room = 2**26 + max((4 * (8 * h.hessian.numel() + 2 * w.numel()) for h, w in movable), default=0)
+    movable = [(h, w) for h, w in work if w.is_cuda and h.updates_on_any_device()]
+    # A worker holds one expert at a time.
+    room = max((_update_room(h, w) for h, w in movable), default=0)
     gpus = [gpu for gpu in _own_gpus() if get_available_gpu_memory(gpu) > room]
     if len(gpus) < 2 or len(movable) < 2:
         movable = []
@@ -320,9 +327,6 @@ class GPTQHelper:
     """
 
     CACHE_NAME = "_forward_no_gptq_hessian"
-    # Whether update_weights() may run on a device other than its weight's: true only for helpers
-    # whose update reads no quantizer state, which lives on the weight's device.
-    UPDATES_ON_ANY_DEVICE = False
 
     def __init__(
         self,
@@ -427,13 +431,22 @@ def free(self):
         self.weight = None
         self.h_inv = None
 
+    def updates_on_any_device(self) -> bool:
+        """Whether :meth:`update_weights` may run on a device other than its weight's.
+
+        True only for helpers whose update reads no quantizer state, which lives on the weight's
+        device, and fits the room :func:`update_expert_weights` checks for.
+        """
+        return False
+
     def update_weights(self, block_size, perc_damp, weight=None, device=None):
         """Run GPTQ blockwise weight update on this helper's weight.
 
         Populates ``self.weight`` and ``self.h_inv``, runs the blockwise update, logs MSE, and
         overwrites ``weight`` (default ``module.weight``) in place, so an expert's slice of a
         fused 3-D parameter is written back into that parameter. The update runs on ``device``,
-        by default the weight's; only a helper with ``UPDATES_ON_ANY_DEVICE`` may pass another.
+        by default the weight's; only a helper whose :meth:`updates_on_any_device` is true may
+        pass another.
         """
         weight = self.module.weight if weight is None else weight
         device = weight.device if device is None else device
diff --git a/tests/gpu/torch/quantization/test_gptq.py b/tests/gpu/torch/quantization/test_gptq.py
index f86e015f803..552d9da117c 100644
--- a/tests/gpu/torch/quantization/test_gptq.py
+++ b/tests/gpu/torch/quantization/test_gptq.py
@@ -24,6 +24,7 @@
 
 import modelopt.torch.quantization as mtq
 from modelopt.torch.export.unified_export_hf import _export_quantized_weight
+from modelopt.torch.quantization.extensions import get_cuda_ext_ggml
 from modelopt.torch.quantization.ggml.common import pinned_packed_weight
 from modelopt.torch.quantization.ggml.gptq import GGMLGPTQHelper
 from modelopt.torch.quantization.model_calib import gptq
@@ -119,6 +120,52 @@ def recording_update_weights(self, *args, device=None, **kwargs):
             )
 
 
+@pytest.mark.parametrize(
+    ("num_bits", "shape", "cuda_extension"),
+    [
+        ("iq2_xxs", (16384, 2048), True),  # a projection 8x wider than its input
+        ("iq1_s", (1024, 512), True),
+        # The PyTorch encoder's scratch grows with block_chunk_size, past the room check.
+        ("iq1_s", (1024, 512), False),
+    ],
+)
+def test_an_update_that_may_move_fits_its_room(
+    need_2_gpus, monkeypatch, num_bits, shape, cuda_extension
+):
+    """An expert update that may run on another GPU peaks there within the room it is admitted with."""
+    if not cuda_extension:
+        monkeypatch.setattr(get_cuda_ext_ggml, "extension", None, raising=False)
+    out_features, in_features = shape
+    torch.manual_seed(RAND_SEED)
+    linear = torch.nn.Linear(
+        in_features, out_features, bias=False, dtype=torch.bfloat16, device="cuda:0"
+    )
+    inputs = torch.randn(512, in_features, dtype=torch.bfloat16, device="cuda:0")
+    config = {
+        "quant_cfg": [
+            {"quantizer_name": "*", "enable": False},
+            {
+                "quantizer_name": "*weight_quantizer",
+                "cfg": {"num_bits": num_bits, "block_sizes": {-1: 256}, "backend": "ggml"},
+            },
+        ],
+        "algorithm": "max",
+    }
+    mtq.quantize(linear, config, forward_loop=lambda m: m(inputs))
+    helper = GGMLGPTQHelper(linear, "linear", log_mse=False)
+    helper.accumulate(inputs)
+    worker = torch.device("cuda", 1)
+    torch.cuda.synchronize(worker)
+    torch.cuda.reset_peak_memory_stats(worker)
+    allocated = torch.cuda.memory_allocated(worker)
+    with torch.cuda.device(worker):
+        helper.update_weights(256, 0.3, linear.weight, device=worker)
+    peak = torch.cuda.max_memory_allocated(worker) - allocated
+
+    if helper.updates_on_any_device():
+        assert peak <= calib_utils._update_room(helper, linear.weight)
+
+
 @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."""

From 958a02de551d4508d5fa78e3a63a0a6d809c4f66 Mon Sep 17 00:00:00 2001
From: Chenjie Luo <chenjiel@nvidia.com>
Date: Fri, 30 Oct 2026 21:56:26 +0000
Subject: [PATCH 5/7] Fail the GPTQ expert-update GPU tests when the GGML CUDA
 extension is missing

Expert updates move to another GPU only with the GGML CUDA encoders. Without
the extension, the room test's CUDA cases passed without checking anything
and the two-GPU test failed on an unclear device assertion. Both now load
the extension first, failing with its build error, as the format tests do.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Signed-off-by: Chenjie Luo <chenjiel@nvidia.com>
---
 tests/gpu/torch/quantization/test_gptq.py | 5 ++++-
 1 file changed, 4 insertions(+), 1 deletion(-)

diff --git a/tests/gpu/torch/quantization/test_gptq.py b/tests/gpu/torch/quantization/test_gptq.py
index 552d9da117c..a2caa929894 100644
--- a/tests/gpu/torch/quantization/test_gptq.py
+++ b/tests/gpu/torch/quantization/test_gptq.py
@@ -72,6 +72,7 @@ def test_expert_updates_on_another_gpu_match_one_gpu(need_2_gpus, monkeypatch, i
     also on inference tensors under the caller's inference mode, which worker threads do not inherit.
     """
     pytest.importorskip("transformers")  # the fused-experts quant wrapper is the HF plugin's
+    get_cuda_ext_ggml(raise_if_failed=True)  # updates move only with the CUDA encoders
     # Two GGML groups per expert input, so the Hessian carries the first group's error to the next.
     width = 512
     shard = torch.zeros(1, device="cuda:1")  # noqa: F841  # the model also lives on cuda:1
@@ -133,7 +134,9 @@ def test_an_update_that_may_move_fits_its_room(
     need_2_gpus, monkeypatch, num_bits, shape, cuda_extension
 ):
     """An expert update that may run on another GPU peaks there within the room it is admitted with."""
-    if not cuda_extension:
+    if cuda_extension:
+        get_cuda_ext_ggml(raise_if_failed=True)  # without it, no update may move to check
+    else:
         monkeypatch.setattr(get_cuda_ext_ggml, "extension", None, raising=False)
     out_features, in_features = shape
     torch.manual_seed(RAND_SEED)

From 3d90e69f63a9b9e3257051a43384a2c867cedaf2 Mon Sep 17 00:00:00 2001
From: Chenjie Luo <chenjiel@nvidia.com>
Date: Fri, 30 Oct 2026 22:32:48 +0000
Subject: [PATCH 6/7] Count GPTQ's block and group buffers in the expert-update
 room check

The room an expert update is admitted with counted the weight, its Hessian
and the codec, but not the update's block of errors and the group-sized
temporaries around each triangular solve. For an input no wider than a
GPTQ block, each is as large as the weight: a 65536 x 256 IQ2_XXS expert
peaked at 325 MiB against a 194 MiB check. The room now adds up to five
blocks of columns, so it takes block_size, and the GPU test covers that
shape.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Signed-off-by: Chenjie Luo <chenjiel@nvidia.com>
---
 modelopt/torch/quantization/utils/calib_utils.py | 13 +++++++++----
 tests/gpu/torch/quantization/test_gptq.py        |  6 ++++--
 2 files changed, 13 insertions(+), 6 deletions(-)

diff --git a/modelopt/torch/quantization/utils/calib_utils.py b/modelopt/torch/quantization/utils/calib_utils.py
index d334584373d..e3630c69eb5 100644
--- a/modelopt/torch/quantization/utils/calib_utils.py
+++ b/modelopt/torch/quantization/utils/calib_utils.py
@@ -121,14 +121,19 @@ def _own_gpus() -> list[int]:
     return get_process_gpus()
 
 
-def _update_room(helper, weight) -> int:
+def _update_room(helper, weight, block_size) -> int:
     """Bytes a GPU needs to run ``helper``'s update of ``weight``, as measured with the CUDA codecs.
 
     The Hessian, the inverse and their temporaries; the float32 copy of the weight with a
     dtype-conversion temporary (a projection much wider than its input makes the weight the larger
-    part); and the codec's tables and scratch, about 35 MiB on a GPU's first update.
+    part); the block's errors and up to four temporaries of a group no wider than the block, each
+    as large as the weight when the input is no wider than a block; and the codec's tables and
+    scratch, about 35 MiB on a GPU's first update.
     """
-    return 2**26 + 4 * (8 * helper.hessian.numel() + 2 * weight.numel())
+    columns = weight.shape[-1]
+    rows = weight.numel() // columns
+    block = min(block_size, columns)
+    return 2**26 + 4 * (8 * helper.hessian.numel() + rows * (2 * columns + 5 * block))
 
 
 def update_expert_weights(work, block_size, perc_damp):
@@ -142,7 +147,7 @@ def update_expert_weights(work, block_size, perc_damp):
     """
     movable = [(h, w) for h, w in work if w.is_cuda and h.updates_on_any_device()]
     # A worker holds one expert at a time.
-    room = max((_update_room(h, w) for h, w in movable), default=0)
+    room = max((_update_room(h, w, block_size) for h, w in movable), default=0)
     gpus = [gpu for gpu in _own_gpus() if get_available_gpu_memory(gpu) > room]
     if len(gpus) < 2 or len(movable) < 2:
         movable = []
diff --git a/tests/gpu/torch/quantization/test_gptq.py b/tests/gpu/torch/quantization/test_gptq.py
index a2caa929894..a7f441b243a 100644
--- a/tests/gpu/torch/quantization/test_gptq.py
+++ b/tests/gpu/torch/quantization/test_gptq.py
@@ -125,6 +125,7 @@ def recording_update_weights(self, *args, device=None, **kwargs):
     ("num_bits", "shape", "cuda_extension"),
     [
         ("iq2_xxs", (16384, 2048), True),  # a projection 8x wider than its input
+        ("iq2_xxs", (65536, 256), True),  # one GPTQ block spans the input
         ("iq1_s", (1024, 512), True),
         # The PyTorch encoder's scratch grows with block_chunk_size, past the room check.
         ("iq1_s", (1024, 512), False),
@@ -157,16 +158,17 @@ def test_an_update_that_may_move_fits_its_room(
     mtq.quantize(linear, config, forward_loop=lambda m: m(inputs))
     helper = GGMLGPTQHelper(linear, "linear", log_mse=False)
     helper.accumulate(inputs)
+    block_size = 256
     worker = torch.device("cuda", 1)
     torch.cuda.synchronize(worker)
     torch.cuda.reset_peak_memory_stats(worker)
     allocated = torch.cuda.memory_allocated(worker)
     with torch.cuda.device(worker):
-        helper.update_weights(256, 0.3, linear.weight, device=worker)
+        helper.update_weights(block_size, 0.3, linear.weight, device=worker)
     peak = torch.cuda.max_memory_allocated(worker) - allocated
 
     if helper.updates_on_any_device():
-        assert peak <= calib_utils._update_room(helper, linear.weight)
+        assert peak <= calib_utils._update_room(helper, linear.weight, block_size)
 
 
 @pytest.mark.parametrize("blocker", ["idle_gpu_full", "idle_gpu_too_small", "multi_rank"])

From e3123f97ca8d7ffb303b597196bb3397a1aec1ac Mon Sep 17 00:00:00 2001
From: Chenjie Luo <chenjiel@nvidia.com>
Date: Sat, 31 Oct 2026 03:36:47 +0000
Subject: [PATCH 7/7] Share one per-GPU thread runner between imatrix encodes
 and GPTQ updates

The imatrix encoder and the GPTQ expert updates each ran their own copy of
the same scaffolding: the first item in the caller's thread for lazy setup,
one thread per GPU under the caller's grad and inference mode, and a sync.
The copies had already drifted once, when the GPTQ workers lacked the
inference mode. run_on_each_gpu now holds those rules; the callers only say
where items come from, a list per GPU or a queue all GPUs share.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Signed-off-by: Chenjie Luo <chenjiel@nvidia.com>
---
 modelopt/torch/quantization/model_calib.py    | 26 ++------
 .../torch/quantization/utils/calib_utils.py   | 63 ++++++++++++-------
 2 files changed, 45 insertions(+), 44 deletions(-)

diff --git a/modelopt/torch/quantization/model_calib.py b/modelopt/torch/quantization/model_calib.py
index e570fa61128..a388c1d319e 100644
--- a/modelopt/torch/quantization/model_calib.py
+++ b/modelopt/torch/quantization/model_calib.py
@@ -21,7 +21,6 @@
 import warnings
 from collections import Counter, defaultdict
 from collections.abc import Callable, Mapping, Sequence
-from concurrent.futures import ThreadPoolExecutor
 from contextlib import nullcontext
 from functools import partial
 from typing import Any, TypeAlias
@@ -79,6 +78,7 @@
     _IMATRIX_ENCODER_REGISTRY,
     GPTQHelper,
     bfloat16_hessian_gemm,
+    run_on_each_gpu,
     track_calib_forward_memory,
     update_expert_weights,
 )
@@ -2797,25 +2797,11 @@ def _encode_on_each_gpu(by_device, encode):
         if device not in gpus:
             for item in items:
                 encode(*item)
-    if not gpus:
-        return
-    # The first encode runs here, so lazy setup that is not thread-safe, such as loading a format's
-    # CUDA extension, happens once before the workers start.
-    first, *by_device[gpus[0]] = by_device[gpus[0]]
-    encode(*first)
-    inference = torch.is_inference_mode_enabled()
-
-    def worker(device):
-        # The device, grad mode and inference mode are per thread: take the caller's.
-        torch.cuda.set_device(device)
-        with torch.inference_mode() if inference else torch.no_grad():
-            for item in by_device[device]:
-                encode(*item)
-        torch.cuda.synchronize(device)
-
-    with ThreadPoolExecutor(len(gpus)) as pool:
-        for future in [pool.submit(worker, gpu) for gpu in gpus]:
-            future.result()
+    if gpus:
+        pending = {device: iter(by_device[device]) for device in gpus}
+        run_on_each_gpu(
+            gpus, lambda device: next(pending[device], None), lambda item, _: encode(*item)
+        )
 
 
 def _run_weight_scale_calibration(model, forward_loop, scale_algorithm):
diff --git a/modelopt/torch/quantization/utils/calib_utils.py b/modelopt/torch/quantization/utils/calib_utils.py
index e3630c69eb5..33f075a4e7d 100644
--- a/modelopt/torch/quantization/utils/calib_utils.py
+++ b/modelopt/torch/quantization/utils/calib_utils.py
@@ -148,46 +148,61 @@ def update_expert_weights(work, block_size, perc_damp):
     movable = [(h, w) for h, w in work if w.is_cuda and h.updates_on_any_device()]
     # A worker holds one expert at a time.
     room = max((_update_room(h, w, block_size) for h, w in movable), default=0)
-    gpus = [gpu for gpu in _own_gpus() if get_available_gpu_memory(gpu) > room]
+    gpus = [torch.device("cuda", g) for g in _own_gpus() if get_available_gpu_memory(g) > room]
     if len(gpus) < 2 or len(movable) < 2:
         movable = []
+
+    def update(item, device):
+        helper, weight = item
+        helper.update_weights(block_size, perc_damp, weight, device=device)
+        helper.free()
+
     moved = {id(helper) for helper, _ in movable}
-    for helper, weight in work:
-        if id(helper) not in moved:
-            helper.update_weights(block_size, perc_damp, weight)
-            helper.free()
+    for item in work:
+        if id(item[0]) not in moved:
+            update(item, None)
     if not movable:
         return
 
-    # The first update runs here, so lazy setup such as loading a format's CUDA extension, which is
-    # not thread-safe, happens once before the workers start.
-    helper, weight = movable[0]
-    helper.update_weights(block_size, perc_damp, weight)
-    helper.free()
     pending: queue.SimpleQueue = queue.SimpleQueue()
-    for item in movable[1:]:
+    for item in movable:
         pending.put(item)
+
+    def next_item(device):
+        try:
+            return pending.get_nowait()
+        except queue.Empty:
+            return None
+
+    run_on_each_gpu(gpus, next_item, update)
+    for device in {weight.device for _, weight in movable}:
+        torch.cuda.synchronize(device)
+
+
+def run_on_each_gpu(gpus, next_item, run):
+    """Call ``run(item, device)`` on the items ``next_item(device)`` hands out, one thread per GPU.
+
+    ``next_item`` returns the next item for that GPU, or ``None`` when there is none left; each GPU's
+    items can be its own or come from a queue all of them share. The first item runs in the caller's
+    thread with ``device=None``, on the device its data is on, so lazy setup that is not thread-safe,
+    such as loading a format's CUDA extension, happens once before the workers start.
+    """
+    first = next_item(gpus[0])
+    if first is not None:
+        run(first, None)
     inference = torch.is_inference_mode_enabled()
 
-    def worker(gpu):
+    def worker(device):
         # The device, grad mode and inference mode are per thread: take the caller's.
-        torch.cuda.set_device(gpu)
-        device = torch.device("cuda", gpu)
+        torch.cuda.set_device(device)
         with torch.inference_mode() if inference else torch.no_grad():
-            while True:
-                try:
-                    helper, weight = pending.get_nowait()
-                except queue.Empty:
-                    break
-                helper.update_weights(block_size, perc_damp, weight, device=device)
-                helper.free()
+            while (item := next_item(device)) is not None:
+                run(item, device)
         torch.cuda.synchronize(device)
 
     with ThreadPoolExecutor(len(gpus)) as pool:
-        for future in [pool.submit(worker, gpu) for gpu in gpus]:
+        for future in [pool.submit(worker, device) for device in gpus]:
             future.result()
-    for device in {weight.device for _, weight in movable}:
-        torch.cuda.synchronize(device)
 
 
 def update_hessian(input, hessian, n_samples, input_dtype=torch.float32):