diff --git a/modelopt/torch/quantization/model_calib.py b/modelopt/torch/quantization/model_calib.py index 40c86bf7bd0..bdba6d82fb8 100644 --- a/modelopt/torch/quantization/model_calib.py +++ b/modelopt/torch/quantization/model_calib.py @@ -2662,6 +2662,31 @@ def _expert_hook(_input_quantizer, _args, output): # Calibration rows squared at a time, bounding imatrix's float32 temporaries to this slice of an # input rather than twice the whole of it. _IMATRIX_ROW_CHUNK = 1024 +# CUDA reductions over about 1,024 rows or more stage partial sums in a global buffer that grows +# with the input, so longer GPU inputs are reduced in groups of this many rows. +_IMATRIX_GPU_GROUP = 512 + + +def _imatrix_sum_of_squares(x: torch.Tensor) -> torch.Tensor: + """Per-column sum of squares of the rows of ``x``, in float32.""" + if not (x.is_cuda and x.dtype in (torch.bfloat16, torch.float16, torch.float32)): + # On a CPU vector_norm copies the whole input to float32 first, and it rejects float64, + # so square float32 slices of the input instead. + return sum(rows.float().square().sum(dim=0) for rows in x.split(_IMATRIX_ROW_CHUNK)) + if x.shape[0] > _IMATRIX_GPU_GROUP**2: + # Bound the group sums, and the reduction over them, to _IMATRIX_GPU_GROUP rows too. + return sum(_imatrix_sum_of_squares(rows) for rows in x.split(_IMATRIX_GPU_GROUP**2)) + # On a GPU vector_norm squares and sums in float32 inside the kernel, without a float32 copy. + if x.shape[0] <= 3 * _IMATRIX_GPU_GROUP: + # Launches dominate at this size, so one reduction is fastest, and up to here it stages no + # more than squaring float32 slices would. + return torch.linalg.vector_norm(x, 2, dim=0, dtype=torch.float32).square_() + full = x.shape[0] // _IMATRIX_GPU_GROUP * _IMATRIX_GPU_GROUP + groups = x[:full].reshape(-1, _IMATRIX_GPU_GROUP, x.shape[1]) + sum_sq = torch.linalg.vector_norm(groups, 2, dim=1, dtype=torch.float32).square_().sum(dim=0) + if full < x.shape[0]: + sum_sq += torch.linalg.vector_norm(x[full:], 2, dim=0, dtype=torch.float32).square_() + return sum_sq @torch.no_grad() @@ -2731,7 +2756,7 @@ def capture(quantizer, _weight, inputs): x = inputs.detach().reshape(-1, inputs.shape[-1]) if x.shape[0] == 0: return - sum_sq = sum(rows.float().square().sum(dim=0) for rows in x.split(_IMATRIX_ROW_CHUNK)) + sum_sq = _imatrix_sum_of_squares(x) entry[0] = sum_sq if entry[0] is None else entry[0] + sum_sq entry[1] += x.shape[0] diff --git a/tests/gpu/torch/quantization/test_ggml_imatrix_gpu.py b/tests/gpu/torch/quantization/test_ggml_imatrix_gpu.py index 1ab9812762f..4b93df1ee18 100644 --- a/tests/gpu/torch/quantization/test_ggml_imatrix_gpu.py +++ b/tests/gpu/torch/quantization/test_ggml_imatrix_gpu.py @@ -15,6 +15,7 @@ """GPU tests for imatrix calibration of the GGML formats.""" +import dataclasses import time import pytest @@ -22,7 +23,9 @@ from torch import nn import modelopt.torch.quantization as mtq +from modelopt.torch.quantization.ggml import GGML_FORMAT_REGISTRY from modelopt.torch.quantization.ggml.imatrix import GGMLImatrixEncoder +from modelopt.torch.quantization.model_calib import _imatrix_sum_of_squares IQ2_XXS_IMATRIX = { "quant_cfg": [ @@ -86,3 +89,55 @@ def timed_encode(self, weight, quantizer, second_moments): for gpu_b, b_start, b_end in spans if gpu_a != gpu_b ) + + +@pytest.mark.parametrize( + "rows", [1500, 2048, 2500, 512 * 512 + 700], ids=["one", "groups", "leftover", "slices"] +) +def test_imatrix_second_moments_of_bf16_inputs_on_gpu(monkeypatch, rows): + """On a GPU, bf16 inputs take the fused float32 reduction: in one call up to 1,536 rows, else + over 512-row groups and the rows left over, in slices of 512 groups. Each gives the float64 + statistic.""" + calls = [] + iq2_xxs = GGML_FORMAT_REGISTRY["iq2_xxs"] + + def recording_quantize(weight, **kwargs): + calls.append(kwargs.get("importance")) + return iq2_xxs.quantize(weight, **kwargs) + + monkeypatch.setitem( + GGML_FORMAT_REGISTRY, "iq2_xxs", dataclasses.replace(iq2_xxs, quantize=recording_quantize) + ) + torch.manual_seed(0) + scale = torch.randn(512).mul(1.5).exp() # columns of very different scale + inputs = (torch.randn(rows, 512) * scale).to("cuda", torch.bfloat16) + model = nn.Linear(512, 8, bias=False, device="cuda", dtype=torch.bfloat16) + mtq.quantize(model, IQ2_XXS_IMATRIX, forward_loop=lambda m: m(inputs)) + + second_moments = inputs.double().square().mean(dim=0) + expected = second_moments.clamp_min(1e-4 * second_moments.mean()).sqrt().float() + torch.testing.assert_close(calls[-1], expected, rtol=1e-5, atol=0) + + +@pytest.mark.parametrize( + ("rows", "columns"), + [(1536, 4096), (2048, 4096), (32768 + 300, 4096), (4 * 512 * 512 + 300, 512)], + ids=["1.5K", "2K", "32K", "1M"], +) +def test_imatrix_statistic_takes_no_more_gpu_memory_than_float32_slices(rows, columns): + """The fused reduction's temporaries peak no higher than squaring float32 slices of 1,024 rows, + the statistic's previous path, so a calibration that fit before still fits: one reduction over + all rows exceeds them at 32K rows, and one over all 512-row groups past 512K rows.""" + x = torch.randn(rows, columns, device="cuda", dtype=torch.bfloat16) + + def peak_memory(sum_of_squares): + torch.cuda.synchronize() + before = torch.cuda.memory_allocated() + torch.cuda.reset_peak_memory_stats() + sum_of_squares(x) + return torch.cuda.max_memory_allocated() - before + + float32_slices = peak_memory( + lambda x: sum(r.float().square().sum(dim=0) for r in x.split(1024)) + ) + assert peak_memory(_imatrix_sum_of_squares) <= float32_slices diff --git a/tests/unit/torch/quantization/test_ggml_imatrix.py b/tests/unit/torch/quantization/test_ggml_imatrix.py index 3955dd222b9..4de4f98c3dc 100644 --- a/tests/unit/torch/quantization/test_ggml_imatrix.py +++ b/tests/unit/torch/quantization/test_ggml_imatrix.py @@ -81,12 +81,14 @@ def _decoded(quantizer, shape): return IQ2_XXS.dequantize(pinned, torch.tensor(shape), dtype=torch.float32) -def test_imatrix_weights_the_search_by_the_root_of_the_inputs_second_moments(monkeypatch): +@pytest.mark.parametrize("dtype", [torch.float32, torch.float64]) +def test_imatrix_weights_the_search_by_the_root_of_the_inputs_second_moments(monkeypatch, dtype): calls = _record_importance(monkeypatch) - inputs = _inputs(rows=2500) # more rows than the hook squares at a time - mtq.quantize(nn.Linear(512, 8, bias=False), _config(IMATRIX), forward_loop=lambda m: m(inputs)) + inputs = _inputs(rows=2500).to(dtype) # more rows than the hook squares at a time + model = nn.Linear(512, 8, bias=False, dtype=dtype) + mtq.quantize(model, _config(IMATRIX), forward_loop=lambda m: m(inputs)) - torch.testing.assert_close(calls[-1], _importance(inputs.square().mean(dim=0))) + torch.testing.assert_close(calls[-1], _importance(inputs.float().square().mean(dim=0))) def test_imatrix_pins_the_codes_it_chose():