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."""