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 + )