fix(qqq): accumulate Hessian in-place per device to reduce memory (#2981)

Author: Qubitium

gptqmodel/quantization/qqq.py

@@ -451,15 +451,20 @@ def add_batch(self, inp, out):
             pad = inp.new_zeros((self._tp_pad_cols, inp.shape[1]))
             inp = torch.cat((inp, pad), dim=0)
 
-        xtx = inp.matmul(inp.t())
-
         with self.lock:
             self.fwd_counter += 1
             existing = self._device_hessian_partials.get(dev)
             if existing is None:
-                self._device_hessian_partials[dev] = xtx.to(dtype=torch.float32).detach()
-            else:
-                existing.add_(xtx.to(dtype=torch.float32))
+                existing = torch.zeros(
+                    (self.columns, self.columns),
+                    dtype=torch.float32,
+                    device=dev,
+                )
+                self._device_hessian_partials[dev] = existing
+
+            # Accumulate xtx directly into the per-device buffer to avoid
+            # allocating a separate output tensor for every batch.
+            existing.addmm_(inp, inp.t(), beta=1.0, alpha=1.0)
             self._device_sample_counts[dev] = self._device_sample_counts.get(dev, 0) + batch_token_size
             self.nsamples += batch_token_size
 
@@ -474,17 +479,21 @@ def materialize_hessian(self, target_device: Optional[torch.device] = None) -> t
             self.nsamples = 0
             return self.H
 
+        # Merge per-device partials into a single Hessian on the target device.
+        # Pop and delete each partial as it is added so the peak memory stays
+        # at result + one partial instead of result + all partials.
         result = torch.zeros((self.columns, self.columns), dtype=torch.float32, device=target_device)
-        for partial in self._device_hessian_partials.values():
-            if partial.device != target_device or partial.dtype != torch.float32:
-                result.add_(partial.to(device=target_device, dtype=torch.float32))
-            else:
+        while self._device_hessian_partials:
+            dev, partial = self._device_hessian_partials.popitem()
+            if partial.device == target_device and partial.dtype == torch.float32:
                 result.add_(partial)
+            else:
+                result.add_(partial.to(device=target_device, dtype=torch.float32))
+            del partial
 
         result.mul_(2.0 / float(total_samples))
         self.H = result
         self.nsamples = total_samples
-        self._device_hessian_partials.clear()
         self._device_sample_counts.clear()
         return self.H
 
@@ -556,7 +565,7 @@ def quantize(
         Q = torch.zeros_like(W)
 
         damp = percdamp * torch.mean(torch.diag(H))
-        diag = torch.arange(self.columns, device=self.dev)
+        diag = torch.arange(self.columns, device=H.device)
         H[diag, diag] += damp
         threshold_raw, is_percent = resolve_threshold(self.fallback, self.expected_nsamples)
         fallback_configured = threshold_raw is not None

tests/test_qqq.py

@@ -10,13 +10,16 @@
 import tempfile
 import unittest
 
+import torch
+
 from datasets import load_dataset
 from models.model_test import ModelTest
 from parameterized import parameterized
 from transformers import AutoTokenizer
 
 from gptqmodel.nn_modules.qlinear.qqq import QQQLinear
 from gptqmodel.quantization import FORMAT, METHOD, QUANT_CONFIG_FILENAME
+from gptqmodel.quantization.qqq import QQQ
 from gptqmodel.utils.torch import torch_empty_cache
 
 
@@ -107,3 +110,57 @@ def assert_qqq_linear(self, model):
                 has_qqq = True
                 break
         self.assertTrue(has_qqq)
+
+
+class TestQQQHessian(unittest.TestCase):
+    """Verify QQQ Hessian accumulation matches the closed-form reference."""
+
+    def _reference_hessian(self, *tensors):
+        """Compute 2/N * (X^T X) for one or more (B, S, C) activation tensors."""
+        target = tensors[0].device
+        x = torch.cat([t.to(target) for t in tensors], dim=0).reshape(-1, tensors[0].shape[-1]).float()
+        return (2.0 / x.shape[0]) * x.t().matmul(x)
+
+    def test_single_batch_matches_reference(self):
+        torch.manual_seed(42)
+        layer = torch.nn.Linear(16, 8, dtype=torch.float32)
+        qqq = QQQ(layer)
+        x = torch.randn(4, 8, 16, dtype=torch.float32)
+        qqq.add_batch(x, None)
+        H = qqq.materialize_hessian()
+        H_ref = self._reference_hessian(x)
+        self.assertTrue(torch.allclose(H, H_ref, rtol=1e-4, atol=1e-5))
+
+    def test_multiple_batches_matches_reference(self):
+        torch.manual_seed(42)
+        layer = torch.nn.Linear(16, 8, dtype=torch.float32)
+        qqq = QQQ(layer)
+        batches = [torch.randn(2, 8, 16, dtype=torch.float32) for _ in range(3)]
+        for x in batches:
+            qqq.add_batch(x, None)
+        H = qqq.materialize_hessian()
+        H_ref = self._reference_hessian(*batches)
+        self.assertTrue(torch.allclose(H, H_ref, rtol=1e-4, atol=1e-5))
+
+    def test_multi_device_matches_reference(self):
+        if not torch.cuda.is_available():
+            self.skipTest("CUDA not available")
+        torch.manual_seed(42)
+        layer = torch.nn.Linear(16, 8, dtype=torch.float32, device="cpu")
+        qqq = QQQ(layer)
+        x_cpu = torch.randn(2, 8, 16, dtype=torch.float32, device="cpu")
+        x_gpu = torch.randn(2, 8, 16, dtype=torch.float32, device="cuda:0")
+        qqq.add_batch(x_cpu, None)
+        qqq.add_batch(x_gpu, None)
+        H = qqq.materialize_hessian()
+        H_ref = self._reference_hessian(x_cpu, x_gpu)
+        self.assertTrue(torch.allclose(H, H_ref, rtol=1e-4, atol=1e-5))
+
+    def test_empty_batch_is_noop(self):
+        torch.manual_seed(42)
+        layer = torch.nn.Linear(16, 8, dtype=torch.float32)
+        qqq = QQQ(layer)
+        qqq.add_batch(torch.zeros(0, 8, 16, dtype=torch.float32), None)
+        H = qqq.materialize_hessian()
+        self.assertEqual(H.shape, (16, 16))
+        self.assertTrue(torch.allclose(H, torch.zeros_like(H)))