diff --git a/modelopt/torch/quantization/ggml/gptq.py b/modelopt/torch/quantization/ggml/gptq.py index 92c56ed5003..7b5cbe91295 100644 --- a/modelopt/torch/quantization/ggml/gptq.py +++ b/modelopt/torch/quantization/ggml/gptq.py @@ -19,6 +19,7 @@ import torch +from ..extensions import get_cuda_ext_ggml from ..utils.calib_utils import GPTQHelper, register_gptq_helper from .common import floored_second_moments, pin_packed_weight from .registry import GGML_FORMAT_REGISTRY @@ -78,11 +79,20 @@ 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): + def updates_on_any_device(self): + """Whether the CUDA encoders are available. + + The update reads only the weight, its Hessian and the format's codec, but the PyTorch + fallback's scratch grows with ``block_chunk_size``, past the room the update is checked for. + """ + return get_cuda_ext_ggml() is not None + + 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 3a210a3545a..a388c1d319e 100644 --- a/modelopt/torch/quantization/model_calib.py +++ b/modelopt/torch/quantization/model_calib.py @@ -21,7 +21,6 @@ import warnings 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 @@ -79,7 +78,9 @@ _IMATRIX_ENCODER_REGISTRY, GPTQHelper, bfloat16_hessian_gemm, + run_on_each_gpu, track_calib_forward_memory, + update_expert_weights, ) __all__ = [ @@ -2648,15 +2649,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 + work = [ + (expert_handles.pop(id(quantizer)), weight) + for weight, quantizer in module.iter_weights_for_calibration() + if id(quantizer) in expert_handles + ] + updated = len(work) + unreached = sum(handle.n_samples == 0 for handle, _ in work) + update_expert_weights(work, block_size, perc_damp) + # Its weight slices would keep this module's onloaded weights alive past the window. + del work if updated: print_rank_0( f"[{name}] GPTQ updated {updated} expert weights, {unreached} of them never " @@ -2794,25 +2797,11 @@ def _encode_on_each_gpu(by_device, encode): 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) - inference = torch.is_inference_mode_enabled() - - def worker(device): - # The device, grad mode and inference mode are per thread: take the caller's. - torch.cuda.set_device(device) - with torch.inference_mode() if inference else 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() + if gpus: + pending = {device: iter(by_device[device]) for device in gpus} + run_on_each_gpu( + gpus, lambda device: next(pending[device], None), lambda item, _: encode(*item) + ) def _run_weight_scale_calibration(model, forward_loop, scale_algorithm): diff --git a/modelopt/torch/quantization/utils/calib_utils.py b/modelopt/torch/quantization/utils/calib_utils.py index 8b045de3f15..33f075a4e7d 100644 --- a/modelopt/torch/quantization/utils/calib_utils.py +++ b/modelopt/torch/quantization/utils/calib_utils.py @@ -39,6 +39,8 @@ """GPTQ helper and Hessian utilities for calibration.""" import functools +import queue +from concurrent.futures import ThreadPoolExecutor from contextlib import contextmanager import torch @@ -106,13 +108,103 @@ def _hessian_offload_device(nbytes: int) -> torch.device: rescale of the whole matrix. Only GPUs this process already holds memory on are candidates, and only outside distributed runs, so no other job's or rank's GPU is touched. """ - if not (dist.is_available() and dist.is_initialized() and dist.get_world_size() > 1): - for gpu in sorted(get_process_gpus(), key=get_available_gpu_memory, reverse=True): - if _gpu_has_room(gpu, nbytes): - return torch.device("cuda", gpu) + for gpu in sorted(_own_gpus(), key=get_available_gpu_memory, reverse=True): + if _gpu_has_room(gpu, nbytes): + return torch.device("cuda", gpu) return torch.device("cpu") +def _own_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 get_process_gpus() + + +def _update_room(helper, weight, block_size) -> int: + """Bytes a GPU needs to run ``helper``'s update of ``weight``, as measured with the CUDA codecs. + + The Hessian, the inverse and their temporaries; the float32 copy of the weight with a + dtype-conversion temporary (a projection much wider than its input makes the weight the larger + part); the block's errors and up to four temporaries of a group no wider than the block, each + as large as the weight when the input is no wider than a block; and the codec's tables and + scratch, about 35 MiB on a GPU's first update. + """ + columns = weight.shape[-1] + rows = weight.numel() // columns + block = min(block_size, columns) + return 2**26 + 4 * (8 * helper.hessian.numel() + rows * (2 * columns + 5 * block)) + + +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 w.is_cuda and h.updates_on_any_device()] + # A worker holds one expert at a time. + room = max((_update_room(h, w, block_size) for h, w in movable), default=0) + gpus = [torch.device("cuda", g) for g in _own_gpus() if get_available_gpu_memory(g) > room] + if len(gpus) < 2 or len(movable) < 2: + movable = [] + + def update(item, device): + helper, weight = item + helper.update_weights(block_size, perc_damp, weight, device=device) + helper.free() + + moved = {id(helper) for helper, _ in movable} + for item in work: + if id(item[0]) not in moved: + update(item, None) + if not movable: + return + + pending: queue.SimpleQueue = queue.SimpleQueue() + for item in movable: + pending.put(item) + + def next_item(device): + try: + return pending.get_nowait() + except queue.Empty: + return None + + run_on_each_gpu(gpus, next_item, update) + for device in {weight.device for _, weight in movable}: + torch.cuda.synchronize(device) + + +def run_on_each_gpu(gpus, next_item, run): + """Call ``run(item, device)`` on the items ``next_item(device)`` hands out, one thread per GPU. + + ``next_item`` returns the next item for that GPU, or ``None`` when there is none left; each GPU's + items can be its own or come from a queue all of them share. The first item runs in the caller's + thread with ``device=None``, on the device its data is on, so lazy setup that is not thread-safe, + such as loading a format's CUDA extension, happens once before the workers start. + """ + first = next_item(gpus[0]) + if first is not None: + run(first, None) + inference = torch.is_inference_mode_enabled() + + def worker(device): + # The device, grad mode and inference mode are per thread: take the caller's. + torch.cuda.set_device(device) + with torch.inference_mode() if inference else torch.no_grad(): + while (item := next_item(device)) is not None: + run(item, device) + torch.cuda.synchronize(device) + + with ThreadPoolExecutor(len(gpus)) as pool: + for future in [pool.submit(worker, device) for device in gpus]: + future.result() + + def update_hessian(input, hessian, n_samples, input_dtype=torch.float32): """Update hessian matrix with new input samples using incremental formula. @@ -359,20 +451,31 @@ def free(self): self.weight = None self.h_inv = None - def update_weights(self, block_size, perc_damp, weight=None): + def updates_on_any_device(self) -> bool: + """Whether :meth:`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, and fits the room :func:`update_expert_weights` checks for. + """ + return False + + 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 whose :meth:`updates_on_any_device` is true 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/_test_utils/torch/quantization/models.py b/tests/_test_utils/torch/quantization/models.py index 2e1875337df..b8542545895 100644 --- a/tests/_test_utils/torch/quantization/models.py +++ b/tests/_test_utils/torch/quantization/models.py @@ -174,15 +174,13 @@ def get_input(self, batch_size=2): class ToyFusedExperts(nn.Module): """HF transformers 5 fused-experts layout: one ``F.linear`` pair per routed expert.""" - def __init__(self): + def __init__(self, hidden=TOY_MOE_HIDDEN, intermediate=TOY_MOE_INTERMEDIATE): super().__init__() self.num_experts = TOY_MOE_EXPERTS self.gate_up_proj = nn.Parameter( - torch.randn(TOY_MOE_EXPERTS, 2 * TOY_MOE_INTERMEDIATE, TOY_MOE_HIDDEN) * 0.02 - ) - self.down_proj = nn.Parameter( - torch.randn(TOY_MOE_EXPERTS, TOY_MOE_HIDDEN, TOY_MOE_INTERMEDIATE) * 0.02 + torch.randn(TOY_MOE_EXPERTS, 2 * intermediate, hidden) * 0.02 ) + self.down_proj = nn.Parameter(torch.randn(TOY_MOE_EXPERTS, hidden, intermediate) * 0.02) def forward(self, hidden_states, expert_of_token): output = torch.zeros_like(hidden_states) @@ -198,10 +196,10 @@ def forward(self, hidden_states, expert_of_token): class ToyMoEBlock(nn.Module): """Routes token ``t`` to expert ``t % 3``, so the last expert receives no token.""" - def __init__(self): + def __init__(self, hidden=TOY_MOE_HIDDEN, intermediate=TOY_MOE_INTERMEDIATE): super().__init__() - self.experts = ToyFusedExperts() - self.shared_expert = nn.Linear(TOY_MOE_HIDDEN, TOY_MOE_HIDDEN, bias=False) + self.experts = ToyFusedExperts(hidden, intermediate) + self.shared_expert = nn.Linear(hidden, hidden, bias=False) def forward(self, hidden_states): routed = self.experts(hidden_states, torch.arange(len(hidden_states)) % 3) @@ -209,9 +207,9 @@ def forward(self, hidden_states): class ToyMoE(nn.Module): - def __init__(self): + def __init__(self, hidden=TOY_MOE_HIDDEN, intermediate=TOY_MOE_INTERMEDIATE): super().__init__() - self.mlp = ToyMoEBlock() + self.mlp = ToyMoEBlock(hidden, intermediate) def forward(self, hidden_states): return self.mlp(hidden_states) diff --git a/tests/gpu/torch/quantization/test_gptq.py b/tests/gpu/torch/quantization/test_gptq.py index d2733293234..a7f441b243a 100644 --- a/tests/gpu/torch/quantization/test_gptq.py +++ b/tests/gpu/torch/quantization/test_gptq.py @@ -18,12 +18,17 @@ import pytest import torch +from _test_utils.torch.quantization.models import 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.extensions import get_cuda_ext_ggml +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 ( @@ -61,6 +66,111 @@ def test_hessian_of_a_full_gpu_goes_to_an_idle_gpu(need_2_gpus, monkeypatch): torch.testing.assert_close(helper.hessian.to("cuda:0"), reference) +@pytest.mark.parametrize("inference_mode", [False, True]) +def test_expert_updates_on_another_gpu_match_one_gpu(need_2_gpus, monkeypatch, inference_mode): + """GGML expert updates spread over two GPUs give the one-GPU weights and codes, pinned at home, + also on inference tensors under the caller's inference mode, which worker threads do not inherit. + """ + pytest.importorskip("transformers") # the fused-experts quant wrapper is the HF plugin's + get_cuda_ext_ggml(raise_if_failed=True) # updates move only with the CUDA encoders + # Two GGML groups per expert input, so the Hessian carries the first group's error to the next. + width = 512 + 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: + with torch.inference_mode(inference_mode): + inputs = torch.randn(96, width, device="cuda") + for gpus in ([0], [0, 1]): + monkeypatch.setattr(calib_utils, "get_process_gpus", lambda gpus=gpus: gpus) + torch.manual_seed(RAND_SEED) + model = ToyMoE(width, width).cuda() + 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( + ("num_bits", "shape", "cuda_extension"), + [ + ("iq2_xxs", (16384, 2048), True), # a projection 8x wider than its input + ("iq2_xxs", (65536, 256), True), # one GPTQ block spans the input + ("iq1_s", (1024, 512), True), + # The PyTorch encoder's scratch grows with block_chunk_size, past the room check. + ("iq1_s", (1024, 512), False), + ], +) +def test_an_update_that_may_move_fits_its_room( + need_2_gpus, monkeypatch, num_bits, shape, cuda_extension +): + """An expert update that may run on another GPU peaks there within the room it is admitted with.""" + if cuda_extension: + get_cuda_ext_ggml(raise_if_failed=True) # without it, no update may move to check + else: + monkeypatch.setattr(get_cuda_ext_ggml, "extension", None, raising=False) + out_features, in_features = shape + torch.manual_seed(RAND_SEED) + linear = torch.nn.Linear( + in_features, out_features, bias=False, dtype=torch.bfloat16, device="cuda:0" + ) + inputs = torch.randn(512, in_features, dtype=torch.bfloat16, device="cuda:0") + config = { + "quant_cfg": [ + {"quantizer_name": "*", "enable": False}, + { + "quantizer_name": "*weight_quantizer", + "cfg": {"num_bits": num_bits, "block_sizes": {-1: 256}, "backend": "ggml"}, + }, + ], + "algorithm": "max", + } + mtq.quantize(linear, config, forward_loop=lambda m: m(inputs)) + helper = GGMLGPTQHelper(linear, "linear", log_mse=False) + helper.accumulate(inputs) + block_size = 256 + worker = torch.device("cuda", 1) + torch.cuda.synchronize(worker) + torch.cuda.reset_peak_memory_stats(worker) + allocated = torch.cuda.memory_allocated(worker) + with torch.cuda.device(worker): + helper.update_weights(block_size, 0.3, linear.weight, device=worker) + peak = torch.cuda.max_memory_allocated(worker) - allocated + + if helper.updates_on_any_device(): + assert peak <= calib_utils._update_room(helper, linear.weight, block_size) + + @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.""" diff --git a/tests/unit/torch/quantization/test_ggml_gptq.py b/tests/unit/torch/quantization/test_ggml_gptq.py index fc5a01ca080..a0c0f480a43 100644 --- a/tests/unit/torch/quantization/test_ggml_gptq.py +++ b/tests/unit/torch/quantization/test_ggml_gptq.py @@ -14,9 +14,11 @@ # limitations under the License. """GPTQ for GGML block formats: the group update, the helper, and the payload pin.""" +import contextlib import copy import dataclasses import io +import weakref from types import SimpleNamespace import pytest @@ -27,6 +29,7 @@ import modelopt.torch.opt as mto import modelopt.torch.quantization as mtq from modelopt.torch.export.unified_export_megatron import GPTModelExporter +from modelopt.torch.quantization import model_calib from modelopt.torch.quantization.ggml import GGML_FORMAT_REGISTRY from modelopt.torch.quantization.ggml.common import _row_keys, pinned_packed_weight from modelopt.torch.quantization.ggml.gptq import gptq_group_update @@ -291,6 +294,31 @@ def same(a, b): assert not same(single_pass[1], layerwise[1]) +def test_gptq_drops_each_fused_experts_weights_before_the_next_window(monkeypatch, moe): + """No expert weight GPTQ updates outlives its module's weight window: under offload it would + keep that module's onloaded weights on the GPU while the next module's are onloaded.""" + updated = [] + update_expert_weights = model_calib.update_expert_weights + window = model_calib.enable_weight_access_and_writeback + + def recording_update_expert_weights(work, *args): + updated.extend(weakref.ref(weight) for _, weight in work) + update_expert_weights(work, *args) + + @contextlib.contextmanager + def checked_window(*args, **kwargs): + assert all(ref() is None for ref in updated) + with window(*args, **kwargs): + yield + + monkeypatch.setattr(model_calib, "update_expert_weights", recording_update_expert_weights) + monkeypatch.setattr(model_calib, "enable_weight_access_and_writeback", checked_window) + inputs = _inputs()[:96, :TOY_MOE_HIDDEN] + mtq.quantize(_MoEStack(moe, ToyMoE()), _int4_config(GPTQ), forward_loop=lambda m: m(inputs)) + + assert updated + + def test_pin_outlives_the_weight_tensor(): model, _ = _gptq_model() # Same values in a new tensor, as an offload round trip or a layerwise resume leaves them.