From aad969172b4871fa5855086c1256577cf32a4909 Mon Sep 17 00:00:00 2001 From: Chenjie Luo Date: Tue, 27 Oct 2026 05:36:49 +0000 Subject: [PATCH 1/5] Collect imatrix statistics with one fused float32 reduction imatrix summed each input's per-column squares by copying 1,024-row slices to float32, squaring and summing them. torch.linalg.vector_norm with dtype=float32 squares and accumulates in float32 inside one kernel, with no float32 copy: 3-10x faster on an RTX PRO 6000 (1.19 -> 0.18 ms for 32K x 4,096 bf16, 0.050 -> 0.016 ms for a small per-expert batch) and as precise (both within 3e-7 of float64). Co-Authored-By: Claude Opus 5.5 (1M context) Signed-off-by: Chenjie Luo --- modelopt/torch/quantization/model_calib.py | 9 +++------ 1 file changed, 3 insertions(+), 6 deletions(-) diff --git a/modelopt/torch/quantization/model_calib.py b/modelopt/torch/quantization/model_calib.py index 40c86bf7bd0..5e992026c21 100644 --- a/modelopt/torch/quantization/model_calib.py +++ b/modelopt/torch/quantization/model_calib.py @@ -2659,11 +2659,6 @@ def _expert_hook(_input_quantizer, _args, output): print_rank_0(f"GPTQ time: {time.time() - total_start:.2f}s") -# 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 - - @torch.no_grad() def imatrix_calibrate(model: nn.Module, forward_loop: ForwardLoop | None = None): """Encode weights once with their inputs' per-column second moments, llama.cpp's imatrix. @@ -2731,7 +2726,9 @@ 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)) + # One reduction that squares and sums in float32 without copying the inputs to float32: + # 3-6x faster than squaring float32 slices of them, and as precise. + sum_sq = torch.linalg.vector_norm(x, 2, dim=0, dtype=torch.float32).square() entry[0] = sum_sq if entry[0] is None else entry[0] + sum_sq entry[1] += x.shape[0] From 72512136e9983e88b42f6b085c9113de5fb1a215 Mon Sep 17 00:00:00 2001 From: Chenjie Luo Date: Fri, 30 Oct 2026 05:22:24 +0000 Subject: [PATCH 2/5] Keep imatrix's bounded float32 path off the GPU and for float64 Review: - On a CPU, vector_norm(dtype=float32) copies the whole input to float32 first (512 MiB for 32K x 4,096 bf16, against 16 MiB for the 1,024-row slices), and it rejects float64 inputs outright. The fused reduction now runs only for bf16, fp16 and fp32 inputs on a GPU; everything else squares float32 slices as before. - The dense statistic test also runs float64 inputs, and a GPU test checks the fused statistic of bf16 inputs against float64. - Drop the timing from the comment; it is in the PR description. Co-Authored-By: Claude Opus 5.5 Signed-off-by: Chenjie Luo --- modelopt/torch/quantization/model_calib.py | 15 ++++++++--- .../quantization/test_ggml_imatrix_gpu.py | 25 +++++++++++++++++++ .../torch/quantization/test_ggml_imatrix.py | 10 +++++--- 3 files changed, 43 insertions(+), 7 deletions(-) diff --git a/modelopt/torch/quantization/model_calib.py b/modelopt/torch/quantization/model_calib.py index 5e992026c21..127fbb5c0b9 100644 --- a/modelopt/torch/quantization/model_calib.py +++ b/modelopt/torch/quantization/model_calib.py @@ -2659,6 +2659,11 @@ def _expert_hook(_input_quantizer, _args, output): print_rank_0(f"GPTQ time: {time.time() - total_start:.2f}s") +# 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 + + @torch.no_grad() def imatrix_calibrate(model: nn.Module, forward_loop: ForwardLoop | None = None): """Encode weights once with their inputs' per-column second moments, llama.cpp's imatrix. @@ -2726,9 +2731,13 @@ def capture(quantizer, _weight, inputs): x = inputs.detach().reshape(-1, inputs.shape[-1]) if x.shape[0] == 0: return - # One reduction that squares and sums in float32 without copying the inputs to float32: - # 3-6x faster than squaring float32 slices of them, and as precise. - sum_sq = torch.linalg.vector_norm(x, 2, dim=0, dtype=torch.float32).square() + if x.is_cuda and x.dtype in (torch.bfloat16, torch.float16, torch.float32): + # On a GPU this squares and sums in float32 inside one kernel, without a float32 copy. + sum_sq = torch.linalg.vector_norm(x, 2, dim=0, dtype=torch.float32).square() + else: + # On a CPU that reduction copies the whole input to float32 first, and it rejects + # float64, so square float32 slices of the input instead. + sum_sq = sum(rows.float().square().sum(dim=0) for rows in x.split(_IMATRIX_ROW_CHUNK)) 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..d3421d08ea1 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,6 +23,7 @@ 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 IQ2_XXS_IMATRIX = { @@ -86,3 +88,26 @@ def timed_encode(self, weight, quantizer, second_moments): for gpu_b, b_start, b_end in spans if gpu_a != gpu_b ) + + +def test_imatrix_second_moments_of_bf16_inputs_on_gpu(monkeypatch): + """On a GPU, bf16 inputs take the fused float32 reduction, which 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(2500, 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) 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(): From bbe3b080f78436153768d969c7e3a34f6a422914 Mon Sep 17 00:00:00 2001 From: Chenjie Luo Date: Fri, 30 Oct 2026 16:37:08 +0000 Subject: [PATCH 3/5] Reduce imatrix GPU inputs in 512-row groups to bound scratch memory Review: one CUDA reduction over all of an input's rows stages partial sums in a global buffer once it covers about 1,000 rows, and that buffer grows with the input: 142 MiB for 32K x 4,096 bf16 rows and 512 MiB for 512K, against 48 MiB for the 1,024-row float32 slices. Reducing a (groups, 512, columns) view in one kernel, then adding the group sums and the rows left over, keeps every reduction below that size. Up to 512K rows the temporaries take no more memory than the slices did (1 MiB at 32K x 4,096), and about 1% of the input beyond, at the speed of the single reduction. A GPU test bounds the reduction's peak memory, and the bf16 statistic test covers inputs with only groups, only leftover rows, and both. Co-Authored-By: Claude Opus 5.5 Signed-off-by: Chenjie Luo --- modelopt/torch/quantization/model_calib.py | 30 ++++++++++++++----- .../quantization/test_ggml_imatrix_gpu.py | 20 +++++++++++-- 2 files changed, 40 insertions(+), 10 deletions(-) diff --git a/modelopt/torch/quantization/model_calib.py b/modelopt/torch/quantization/model_calib.py index 127fbb5c0b9..da67d249ec3 100644 --- a/modelopt/torch/quantization/model_calib.py +++ b/modelopt/torch/quantization/model_calib.py @@ -2662,6 +2662,28 @@ 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 +# Rows a GPU reduction sums at once. A CUDA reduction over about 1,000 rows or more stages partial +# sums in a global buffer that grows with the input, 512 MiB for 512K x 4,096 bf16 rows; summing +# 512-row groups first keeps the temporaries to about 1% of the input. +_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)) + # On a GPU vector_norm squares and sums in float32 inside the kernel, without a float32 copy. + full = x.shape[0] // _IMATRIX_GPU_GROUP * _IMATRIX_GPU_GROUP + parts = [] + if full: + groups = x[:full].reshape(-1, _IMATRIX_GPU_GROUP, x.shape[1]) + group_sums = torch.linalg.vector_norm(groups, 2, dim=1, dtype=torch.float32).square_() + parts.append(group_sums.sum(dim=0)) + if full < x.shape[0]: + parts.append(torch.linalg.vector_norm(x[full:], 2, dim=0, dtype=torch.float32).square_()) + return sum(parts[1:], parts[0]) @torch.no_grad() @@ -2731,13 +2753,7 @@ def capture(quantizer, _weight, inputs): x = inputs.detach().reshape(-1, inputs.shape[-1]) if x.shape[0] == 0: return - if x.is_cuda and x.dtype in (torch.bfloat16, torch.float16, torch.float32): - # On a GPU this squares and sums in float32 inside one kernel, without a float32 copy. - sum_sq = torch.linalg.vector_norm(x, 2, dim=0, dtype=torch.float32).square() - else: - # On a CPU that reduction copies the whole input to float32 first, and it rejects - # float64, so square float32 slices of the input instead. - 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 d3421d08ea1..d16be24d372 100644 --- a/tests/gpu/torch/quantization/test_ggml_imatrix_gpu.py +++ b/tests/gpu/torch/quantization/test_ggml_imatrix_gpu.py @@ -25,6 +25,7 @@ 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": [ @@ -90,8 +91,10 @@ def timed_encode(self, weight, quantizer, second_moments): ) -def test_imatrix_second_moments_of_bf16_inputs_on_gpu(monkeypatch): - """On a GPU, bf16 inputs take the fused float32 reduction, which gives the float64 statistic.""" +@pytest.mark.parametrize("rows", [300, 2048, 2500], ids=["remainder", "groups", "both"]) +def test_imatrix_second_moments_of_bf16_inputs_on_gpu(monkeypatch, rows): + """On a GPU, bf16 inputs take the fused float32 reduction, over 512-row groups and the rows + left over, which gives the float64 statistic.""" calls = [] iq2_xxs = GGML_FORMAT_REGISTRY["iq2_xxs"] @@ -104,10 +107,21 @@ def recording_quantize(weight, **kwargs): ) torch.manual_seed(0) scale = torch.randn(512).mul(1.5).exp() # columns of very different scale - inputs = (torch.randn(2500, 512) * scale).to("cuda", torch.bfloat16) + 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) + + +def test_imatrix_statistic_takes_no_input_sized_gpu_buffer(): + """The GPU reduction's temporaries stay well below a float32 copy of 1,024 rows, the slice the + statistic took before; one reduction over all rows stages a 142 MiB buffer for this input.""" + x = torch.randn(32768 + 300, 4096, device="cuda", dtype=torch.bfloat16) + torch.cuda.synchronize() + before = torch.cuda.memory_allocated() + torch.cuda.reset_peak_memory_stats() + _imatrix_sum_of_squares(x) + assert torch.cuda.max_memory_allocated() - before < 1024 * x.shape[1] * 4 From ee4941f462e7ea4d23850d3a838023db8ec2ed97 Mon Sep 17 00:00:00 2001 From: Chenjie Luo Date: Fri, 30 Oct 2026 17:28:56 +0000 Subject: [PATCH 4/5] Bound imatrix's GPU group sums to slices of 512 groups Review: summing every 512-row group at once still let the group sums, and the reduction over them, grow with the input past 512K rows, to 96 MiB at 1M x 4,096 bf16 rows against 48 MiB for the float32 slices. Inputs over 512 x 512 rows now go through the grouped reduction in slices of that many rows, so no reduction spans more than 512 rows or group sums: the temporaries stay at 8 MiB for 4,096 columns at any length, at the same speed. Smaller inputs take no extra step. The memory test compares the fused reduction's peak with the float32 slices' on the same input, at 32K rows and past 1M, and the bf16 statistic test covers an input split into slices. Drop the measured sizes from the code comment. Co-Authored-By: Claude Opus 5.5 Signed-off-by: Chenjie Luo --- modelopt/torch/quantization/model_calib.py | 8 +++-- .../quantization/test_ggml_imatrix_gpu.py | 35 +++++++++++++------ 2 files changed, 29 insertions(+), 14 deletions(-) diff --git a/modelopt/torch/quantization/model_calib.py b/modelopt/torch/quantization/model_calib.py index da67d249ec3..9142b336000 100644 --- a/modelopt/torch/quantization/model_calib.py +++ b/modelopt/torch/quantization/model_calib.py @@ -2662,9 +2662,8 @@ 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 -# Rows a GPU reduction sums at once. A CUDA reduction over about 1,000 rows or more stages partial -# sums in a global buffer that grows with the input, 512 MiB for 512K x 4,096 bf16 rows; summing -# 512-row groups first keeps the temporaries to about 1% of the input. +# Rows, or group sums, a GPU reduction sums at once: longer CUDA reductions stage partial sums in a +# global buffer that grows with the input. _IMATRIX_GPU_GROUP = 512 @@ -2674,6 +2673,9 @@ def _imatrix_sum_of_squares(x: torch.Tensor) -> torch.Tensor: # 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. full = x.shape[0] // _IMATRIX_GPU_GROUP * _IMATRIX_GPU_GROUP parts = [] diff --git a/tests/gpu/torch/quantization/test_ggml_imatrix_gpu.py b/tests/gpu/torch/quantization/test_ggml_imatrix_gpu.py index d16be24d372..fb1cb6c9e7a 100644 --- a/tests/gpu/torch/quantization/test_ggml_imatrix_gpu.py +++ b/tests/gpu/torch/quantization/test_ggml_imatrix_gpu.py @@ -91,10 +91,12 @@ def timed_encode(self, weight, quantizer, second_moments): ) -@pytest.mark.parametrize("rows", [300, 2048, 2500], ids=["remainder", "groups", "both"]) +@pytest.mark.parametrize( + "rows", [300, 2048, 2500, 512 * 512 + 700], ids=["remainder", "groups", "both", "slices"] +) def test_imatrix_second_moments_of_bf16_inputs_on_gpu(monkeypatch, rows): """On a GPU, bf16 inputs take the fused float32 reduction, over 512-row groups and the rows - left over, which gives the float64 statistic.""" + left over, in slices of 512 groups, which gives the float64 statistic.""" calls = [] iq2_xxs = GGML_FORMAT_REGISTRY["iq2_xxs"] @@ -116,12 +118,23 @@ def recording_quantize(weight, **kwargs): torch.testing.assert_close(calls[-1], expected, rtol=1e-5, atol=0) -def test_imatrix_statistic_takes_no_input_sized_gpu_buffer(): - """The GPU reduction's temporaries stay well below a float32 copy of 1,024 rows, the slice the - statistic took before; one reduction over all rows stages a 142 MiB buffer for this input.""" - x = torch.randn(32768 + 300, 4096, device="cuda", dtype=torch.bfloat16) - torch.cuda.synchronize() - before = torch.cuda.memory_allocated() - torch.cuda.reset_peak_memory_stats() - _imatrix_sum_of_squares(x) - assert torch.cuda.max_memory_allocated() - before < 1024 * x.shape[1] * 4 +@pytest.mark.parametrize( + ("rows", "columns"), [(32768 + 300, 4096), (4 * 512 * 512 + 300, 512)], ids=["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 From d5b4249ea47e7c923e931f4534d3d7b263d6fb22 Mon Sep 17 00:00:00 2001 From: Chenjie Luo Date: Fri, 30 Oct 2026 21:48:24 +0000 Subject: [PATCH 5/5] Reduce imatrix GPU inputs of up to 1,536 rows in one call At that size the hook's cost is its kernel launches. The grouped path's six launches made inputs of 513 to 1,023 rows slower than the float32 slices (0.061 vs 0.046 ms at 700 x 4,096), and cost Qwen3.8's typical 1,040-row expert input 0.079 ms on GB300, against 0.027 ms for one reduction. One reduction stages partial sums only from 1,021 rows, at 8 bytes per input element. The float32 slices peak at 12 bytes per element of a 1,024-row slice, so one reduction stays within them up to 1,536 rows. The bf16 statistic test runs a 1,500-row input through that call. The memory test adds 1,536 rows, where both peaks are equal, and 2,048, which a higher threshold would exceed. Co-Authored-By: Claude Opus 5.5 Signed-off-by: Chenjie Luo --- modelopt/torch/quantization/model_calib.py | 19 ++++++++++--------- .../quantization/test_ggml_imatrix_gpu.py | 11 +++++++---- 2 files changed, 17 insertions(+), 13 deletions(-) diff --git a/modelopt/torch/quantization/model_calib.py b/modelopt/torch/quantization/model_calib.py index 9142b336000..bdba6d82fb8 100644 --- a/modelopt/torch/quantization/model_calib.py +++ b/modelopt/torch/quantization/model_calib.py @@ -2662,8 +2662,8 @@ 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 -# Rows, or group sums, a GPU reduction sums at once: longer CUDA reductions stage partial sums in a -# global buffer that grows with the input. +# 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 @@ -2677,15 +2677,16 @@ def _imatrix_sum_of_squares(x: torch.Tensor) -> torch.Tensor: # 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 - parts = [] - if full: - groups = x[:full].reshape(-1, _IMATRIX_GPU_GROUP, x.shape[1]) - group_sums = torch.linalg.vector_norm(groups, 2, dim=1, dtype=torch.float32).square_() - parts.append(group_sums.sum(dim=0)) + 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]: - parts.append(torch.linalg.vector_norm(x[full:], 2, dim=0, dtype=torch.float32).square_()) - return sum(parts[1:], parts[0]) + sum_sq += torch.linalg.vector_norm(x[full:], 2, dim=0, dtype=torch.float32).square_() + return sum_sq @torch.no_grad() diff --git a/tests/gpu/torch/quantization/test_ggml_imatrix_gpu.py b/tests/gpu/torch/quantization/test_ggml_imatrix_gpu.py index fb1cb6c9e7a..4b93df1ee18 100644 --- a/tests/gpu/torch/quantization/test_ggml_imatrix_gpu.py +++ b/tests/gpu/torch/quantization/test_ggml_imatrix_gpu.py @@ -92,11 +92,12 @@ def timed_encode(self, weight, quantizer, second_moments): @pytest.mark.parametrize( - "rows", [300, 2048, 2500, 512 * 512 + 700], ids=["remainder", "groups", "both", "slices"] + "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, over 512-row groups and the rows - left over, in slices of 512 groups, which gives the float64 statistic.""" + """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"] @@ -119,7 +120,9 @@ def recording_quantize(weight, **kwargs): @pytest.mark.parametrize( - ("rows", "columns"), [(32768 + 300, 4096), (4 * 512 * 512 + 300, 512)], ids=["32K", "1M"] + ("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,