Encode each GPU's imatrix weights in its own thread

modelopt/torch/quantization/model_calib.py

@@ -19,8 +19,10 @@
 import math
 import time
 import warnings
-from collections import Counter
+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
 
@@ -65,10 +67,12 @@
     enable_weight_access_and_writeback,
     is_quantized_column_parallel_linear,
     is_quantized_linear,
+    is_quantized_parallel_linear,
     is_quantized_row_parallel_linear,
     module_name_maps,
     persistent_materialization,
     promote_static_block_weight_quantizers,
+    requires_weight_materialization,
 )
 from .utils.calib_utils import (
     _GPTQ_HELPER_REGISTRY,
@@ -2716,21 +2720,64 @@ def capture(quantizer, _weight, inputs):
         for handle in handles:
             handle.remove()
 
+    def encode(weight, quantizer):
+        sum_sq, count = targets[id(quantizer)]
+        second_moments = None if count == 0 else sum_sq.to(weight.device) / count
+        encoder = _IMATRIX_ENCODER_REGISTRY[quantizer.backend]
+        weight.copy_(encoder.encode(weight, quantizer, second_moments))
+
+    # Weights readable as they stand are grouped by device and encoded by one thread per GPU: a
+    # model that device_map spreads over GPUs would otherwise keep all but one of them idle.
+    # Offloaded or sharded weights are encoded in turn, each inside its materialization window.
+    by_device: dict[torch.device, list] = defaultdict(list)
     encoded = 0
     for _, module in target_modules:
-        with enable_weight_access_and_writeback(module, model, names):
+        in_window = requires_weight_materialization(
+            module, model, names
+        ) or is_quantized_parallel_linear(module)
+        with (
+            enable_weight_access_and_writeback(module, model, names) if in_window else nullcontext()
+        ):
             for weight, quantizer in module.iter_weights_for_calibration():
-                entry = targets.get(id(quantizer))
-                if entry is None:
+                if id(quantizer) not in targets:
                     continue
-                sum_sq, count = entry
-                second_moments = None if count == 0 else sum_sq.to(weight.device) / count
-                encoder = _IMATRIX_ENCODER_REGISTRY[quantizer.backend]
-                weight.copy_(encoder.encode(weight, quantizer, second_moments))
+                if in_window:
+                    encode(weight, quantizer)
+                else:
+                    by_device[weight.device].append((weight, quantizer))
                 encoded += 1
+    _encode_on_each_gpu(by_device, encode)
     print_rank_0(f"imatrix: encoded {encoded} weights")
 
 
+def _encode_on_each_gpu(by_device, encode):
+    """Call ``encode(weight, quantizer)`` on every item, one thread per GPU when several hold items."""
+    gpus = [device for device in by_device if device.type == "cuda"]
+    if len(gpus) < 2:
+        gpus = []
+    for device, items in by_device.items():
+        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)
+
+    def worker(device):
+        torch.cuda.set_device(device)  # per thread, like grad mode
+        with 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()
+
+
 def _run_weight_scale_calibration(model, forward_loop, scale_algorithm):
     """Run the weight-scale calibration algorithm.
 

tests/gpu/torch/quantization/test_ggml_imatrix_gpu.py

@@ -0,0 +1,83 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""GPU tests for imatrix calibration of the GGML formats."""
+
+import copy
+import time
+
+import torch
+from torch import nn
+
+import modelopt.torch.quantization as mtq
+from modelopt.torch.quantization.ggml.imatrix import GGMLImatrixEncoder
+
+IQ2_XXS_IMATRIX = {
+    "quant_cfg": [
+        {"quantizer_name": "*", "enable": False},
+        {
+            "quantizer_name": "*weight_quantizer",
+            "cfg": {"num_bits": "iq2_xxs", "block_sizes": {-1: 256}, "backend": "ggml"},
+        },
+    ],
+    "algorithm": {"method": "imatrix"},
+}
+
+
+class _Layers(nn.Module):
+    """Linear layers fed on whichever GPU each lives on, as a ``device_map`` places them."""
+
+    def __init__(self):
+        super().__init__()
+        self.layers = nn.ModuleList(nn.Linear(512, 512, bias=False) for _ in range(4))
+
+    def forward(self, x):
+        for layer in self.layers:
+            x = layer(x.to(layer.weight.device))
+        return x
+
+
+def test_imatrix_encodes_each_gpus_weights_at_once(need_2_gpus, monkeypatch):
+    """Weights on two GPUs are encoded concurrently, to the same codes as on one GPU."""
+    spans = []
+    encode = GGMLImatrixEncoder.encode
+
+    def timed_encode(self, weight, quantizer, second_moments):
+        start = time.perf_counter()
+        decoded = encode(self, weight, quantizer, second_moments)
+        torch.cuda.synchronize(weight.device)
+        time.sleep(0.2)  # long enough that encodes on different GPUs must overlap to save time
+        spans.append((weight.device.index, start, time.perf_counter()))
+        return decoded
+
+    monkeypatch.setattr(GGMLImatrixEncoder, "encode", timed_encode)
+    torch.manual_seed(0)
+    reference = _Layers().cuda()
+    model = copy.deepcopy(reference)
+    model.layers[2:].to("cuda:1")
+    inputs = torch.randn(8, 128, 512, device="cuda")
+
+    mtq.quantize(reference, IQ2_XXS_IMATRIX, forward_loop=lambda m: m(inputs))
+    spans.clear()
+    mtq.quantize(model, IQ2_XXS_IMATRIX, forward_loop=lambda m: m(inputs))
+
+    for expected, layer in zip(reference.layers, model.layers):
+        assert torch.equal(layer.weight.cuda(0), expected.weight)
+    assert any(
+        a_start < b_end and b_start < a_end
+        for gpu_a, a_start, a_end in spans
+        for gpu_b, b_start, b_end in spans
+        if gpu_a != gpu_b
+    )