Collect imatrix statistics with one fused float32 reduction

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]
 

tests/gpu/torch/quantization/test_ggml_imatrix_gpu.py

@@ -15,14 +15,17 @@
 
 """GPU tests for imatrix calibration of the GGML formats."""
 
+import dataclasses
 import time
 
 import pytest
 import torch
 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

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():