From 248053852dcfbd4c73062ddccd95d773e5d56294 Mon Sep 17 00:00:00 2001 From: Chenjie Luo Date: Tue, 27 Oct 2026 05:36:23 +0000 Subject: [PATCH] Encode each GPU's imatrix weights in its own thread imatrix encodes every target weight once after its forward, one weight after another. In a single-pass run of a model that device_map spreads over several GPUs, each encode runs on its weight's GPU, so all but one GPU sat idle. Weights readable as they stand are now grouped by device and encoded by one thread per GPU, on the GPU they already live on, so the codes are the same. Weights that need a materialization window (offloaded, FSDP- or TP-sharded) are still encoded in turn inside it. The first encode runs before the workers start, so a format's CUDA extension loads once. Co-Authored-By: Claude Opus 5.5 (1M context) Signed-off-by: Chenjie Luo --- modelopt/torch/quantization/model_calib.py | 63 ++++++++++++-- .../quantization/test_ggml_imatrix_gpu.py | 83 +++++++++++++++++++ 2 files changed, 138 insertions(+), 8 deletions(-) create mode 100644 tests/gpu/torch/quantization/test_ggml_imatrix_gpu.py diff --git a/modelopt/torch/quantization/model_calib.py b/modelopt/torch/quantization/model_calib.py index 3de7e89786b..a784e63f57f 100644 --- a/modelopt/torch/quantization/model_calib.py +++ b/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. diff --git a/tests/gpu/torch/quantization/test_ggml_imatrix_gpu.py b/tests/gpu/torch/quantization/test_ggml_imatrix_gpu.py new file mode 100644 index 00000000000..31590a0e063 --- /dev/null +++ b/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 + )