From 37c5210cfc3299a07475d95b75996f32869d0da5 Mon Sep 17 00:00:00 2001 From: Chenjie Luo Date: Thu, 22 Oct 2026 05:23:49 +0000 Subject: [PATCH] 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) Signed-off-by: Chenjie Luo --- modelopt/torch/quantization/ggml/gptq.py | 10 ++- modelopt/torch/quantization/model_calib.py | 26 +++--- .../torch/quantization/utils/calib_utils.py | 88 ++++++++++++++++--- tests/gpu/torch/quantization/test_gptq.py | 53 +++++++++++ 4 files changed, 151 insertions(+), 26 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 cf14a5b8c3c..190057583b2 100644 --- a/modelopt/torch/quantization/model_calib.py +++ b/modelopt/torch/quantization/model_calib.py @@ -69,7 +69,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, + update_expert_weights, +) __all__ = [ "CalibratorFactory", @@ -2466,18 +2471,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 81f91d883e5..c363cefa45b 100644 --- a/modelopt/torch/quantization/utils/calib_utils.py +++ b/modelopt/torch/quantization/utils/calib_utils.py @@ -38,6 +38,9 @@ """GPTQ helper and Hessian utilities for calibration.""" +import queue +from concurrent.futures import ThreadPoolExecutor + import torch import torch.distributed as dist @@ -62,16 +65,72 @@ def _hessian_offload_device(nbytes: int) -> torch.device: 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. """ - 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: - return torch.device("cuda", gpu) + for gpu in sorted(_process_gpus(), key=get_used_gpu_mem_fraction): + if _gpu_is_full(gpu): + break + if torch.cuda.mem_get_info(gpu)[0] > nbytes: + return torch.device("cuda", gpu) return torch.device("cpu") +def _process_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 [i for i in range(torch.cuda.device_count()) if torch.cuda.memory_reserved(i)] + + +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 _process_gpus() if torch.cuda.mem_get_info(gpu)[0] > 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): """Update hessian matrix with new input samples using incremental formula. @@ -156,6 +215,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, @@ -253,20 +315,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 ee6dbec46e1..869ef2154d3 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 ( @@ -57,6 +61,55 @@ def used_fraction(device): 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, "_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."""