fix(qqq): per-device Hessian partials for balanced multi-GPU calibration (#2979)

QQQ.add_batch previously accumulated a single H tensor on self.dev, which
failed when balanced calibration placed module replicas on different GPUs.

- Add per-device Hessian partials and a thread-safe accumulate step, mirroring
  the GPTQ path.
- Materialize the global Hessian at quantize time from the merged partials,
  scaling by 2 / total_token_count (matching GPTQ math).
- Add QQQ to ModelTest's post-quant validation backend dispatch.
- Introduce tests/models/test_llama3_2_qqq.py with fast-mode arc/gsm8k baselines.

Fixes #2115.

Author: Qubitium

gptqmodel/quantization/qqq.py

@@ -4,8 +4,9 @@
 # Contact: qubitium@modelcloud.ai, x.com/qubitium
 
 import math
+import threading
 import time
-from typing import Optional
+from typing import Dict, Optional
 
 import torch
 import transformers
@@ -16,6 +17,7 @@
 from ..quantization.config import FallbackStrategy, SmoothMSE
 from ..quantization.quantizer import HF_OPTIMUM
 from ..utils import setup_logger
+from ..utils.device import get_device
 from .fallback_smooth import mse_optimal_quant, smooth_block
 from .gptq import get_number_of_rows_and_cols
 from .npu_linalg import npu_inverse_cholesky_factor
@@ -242,9 +244,10 @@ def __init__(self, module: nn.Module, qcfg: Optional[QuantizeConfig] = None):
         self.fallback = self.qcfg.fallback
         self.expected_nsamples: Optional[float] = None
 
-        self.H = torch.zeros((self.columns, self.columns),
-                             dtype=torch.float32,
-                             device = self.dev)
+        self.lock = threading.Lock()
+        self.H: Optional[torch.Tensor] = None
+        self._device_hessian_partials: Dict[torch.device, torch.Tensor] = {}
+        self._device_sample_counts: Dict[torch.device, int] = {}
 
     @staticmethod
     def _validate_module(module):
@@ -424,7 +427,6 @@ def add_batch(self, inp, out):
             self.out1 = out
         if len(inp.shape) == 2:
             inp = inp.unsqueeze(0)
-        tmp = inp.shape[0]
         if isinstance(self.layer, nn.Linear) or isinstance(
                 self.layer, transformers.Conv1D
         ):
@@ -441,13 +443,50 @@ def add_batch(self, inp, out):
             inp = unfold(inp)
             inp = inp.permute([1, 0, 2])
             inp = inp.flatten(1)
-        self.H *= self.nsamples / (self.nsamples + tmp)
-        self.nsamples += tmp
-        inp = math.sqrt(2 / self.nsamples) * inp.float()
+
+        dev = torch.device(get_device(inp))
+        batch_token_size = inp.shape[1]
+        inp = inp.float()
         if self._tp_pad_cols:
             pad = inp.new_zeros((self._tp_pad_cols, inp.shape[1]))
             inp = torch.cat((inp, pad), dim=0)
-        self.H += inp.matmul(inp.t())
+
+        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))
+            self._device_sample_counts[dev] = self._device_sample_counts.get(dev, 0) + batch_token_size
+            self.nsamples += batch_token_size
+
+    def materialize_hessian(self, target_device: Optional[torch.device] = None) -> torch.Tensor:
+        if target_device is None:
+            target_device = self.dev
+        target_device = torch.device(target_device)
+
+        total_samples = sum(self._device_sample_counts.values())
+        if total_samples == 0:
+            self.H = torch.zeros((self.columns, self.columns), dtype=torch.float32, device=target_device)
+            self.nsamples = 0
+            return self.H
+
+        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:
+                result.add_(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
 
     @torch.inference_mode()
     def quantize(
@@ -459,6 +498,8 @@ def quantize(
 
         resolved_strategy = resolve_fallback_strategy(self.fallback)
 
+        self.materialize_hessian()
+
         percdamp = self.qcfg.damp_percent
         groupsize = self.qcfg.group_size
         actorder = self.qcfg.desc_act
@@ -696,6 +737,8 @@ def free(self):
             self.inp1 = None
             self.out1 = None
         self.H = None
+        self._device_hessian_partials.clear()
+        self._device_sample_counts.clear()
         self.Losses = None
         self.Trace = None
         if hasattr(self, "quantizer"):

tests/models/model_test.py

@@ -1007,6 +1007,8 @@ def perform_post_quant_validation(self, model_path, trust_remote_code=False):
                 compare_backends = (BACKEND.MARLIN,)
             else:
                 compare_backends = (self.LOAD_BACKEND,)
+        elif format_family == FORMAT.QQQ:
+            compare_backends = (self.LOAD_BACKEND,)
         else:
             compare_backends = (BACKEND.MARLIN, BACKEND.GEMM)
         fallback_backend = None

tests/models/test_llama3_2_qqq.py

@@ -0,0 +1,102 @@
+# SPDX-FileCopyrightText: 2024-2025 ModelCloud.ai
+# SPDX-FileCopyrightText: 2024-2025 qubitium@modelcloud.ai
+# SPDX-License-Identifier: Apache-2.0
+# Contact: qubitium@modelcloud.ai, x.com/qubitium
+
+import os
+
+from model_test import ModelTest
+
+from gptqmodel import BACKEND
+from gptqmodel.quantization import FORMAT, METHOD
+
+
+# QQQ equivalent of tests/models/test_llama3_2.py.
+#
+# Fast-mode post-quant baselines (eager attention, group_size=128, 4bit QQQ):
+#   gsm8k_platinum_cot :: acc,num          0.4574
+#   arc_challenge        :: acc            0.3225
+#   arc_challenge        :: acc_norm       0.3541
+class TestLlama3_2_QQQ(ModelTest):
+    NATIVE_MODEL_ID = "/monster/data/model/Llama-3.2-1B-Instruct"
+    SAVE_PATH = os.environ.get(
+        "GPTQMODEL_LLAMA3_2_QQQ_SAVE_PATH",
+        "/tmp/llama3_2_qqq_saved_ckpt",
+    )
+    DELETE_QUANTIZED_MODEL = False
+    EVAL_BATCH_SIZE = 64
+    DATASET_CONCAT_SIZE = 2048
+    USE_FLASH_ATTN = False
+
+    METHOD = METHOD.QQQ
+    FORMAT = FORMAT.QQQ
+    LOAD_BACKEND = BACKEND.QQQ
+    QUANT_BACKEND = BACKEND.AUTO
+    BITS = 4
+    GROUP_SIZE = 128
+    DESC_ACT = False
+    SYM = True
+    ACT_GROUP_AWARE = False
+
+    EVAL_TASKS_FAST = {
+        "gsm8k_platinum_cot": {
+            "chat_template": True,
+            "evalution_use_model_path": True,
+            "evalution_batch_size": "auto",
+            "evalution_model_args": {
+                "dtype": "bfloat16",
+                "attn_implementation": "eager",
+                "device": "cuda:0",
+            },
+            "evalution_suite_kwargs": {
+                "batch_size": 32,
+                "max_new_tokens": 256,
+                "stream": True,
+            },
+            "acc,num": {
+                "value": 0.4574028122415219,
+                "floor_pct": 0.04,
+                "ceil_pct": 1.0,
+            },
+        },
+        "arc_challenge": {
+            "chat_template": True,
+            "acc": {
+                "value": 0.3225255972696246,
+                "floor_pct": 0.04,
+                "ceil_pct": 1.0,
+            },
+            "acc_norm": {
+                "value": 0.35409556313993173,
+                "floor_pct": 0.04,
+                "ceil_pct": 1.0,
+            },
+        },
+    }
+
+    EVAL_TASKS_SLOW = {
+        "gsm8k_platinum_cot": {
+            "chat_template": True,
+            "acc,num": {
+                "value": 0.4574028122415219,
+                "floor_pct": 0.10,
+                "ceil_pct": 0.10,
+            },
+        },
+        "arc_challenge": {
+            "chat_template": True,
+            "acc": {
+                "value": 0.3225255972696246,
+                "floor_pct": 0.10,
+                "ceil_pct": 0.10,
+            },
+            "acc_norm": {
+                "value": 0.35409556313993173,
+                "floor_pct": 0.10,
+                "ceil_pct": 0.10,
+            },
+        },
+    }
+
+    def test_llama3_2_qqq(self):
+        self.quantize_and_evaluate()