fix(gptq): use Cholesky factor of inverse Hessian for block updates

python/cactus/convert/quantization/cq.py

@@ -230,13 +230,26 @@ def _prepare_input_scale(input_scale: np.ndarray | None, k: int) -> np.ndarray:
     return np.clip(scale, 1e-8, 65504.0).astype(np.float32, copy=False)
 
 
-def _gptq_correct_group(work: np.ndarray, recon: np.ndarray, h_inv: np.ndarray | None, start: int, stop: int) -> None:
-    if h_inv is None or stop >= work.shape[1]:
+def _upper_cholesky_of_inverse(h: np.ndarray) -> np.ndarray | None:
+    h_inv = np.linalg.inv(h.astype(np.float64))
+    upper = np.linalg.cholesky(h_inv).T.astype(np.float32)
+    if not np.all(np.isfinite(upper)):
+        return None
+    return upper
+
+
+def _gptq_correct_group(work: np.ndarray, recon: np.ndarray, h_inv_upper: np.ndarray | None, start: int, stop: int) -> None:
+    if h_inv_upper is None or stop >= work.shape[1]:
         return
     try:
-        m_bb = h_inv[start:stop, start:stop].astype(np.float32, copy=False)
-        m_bs = h_inv[start:stop, stop:].astype(np.float32, copy=False)
-        update = np.linalg.solve(m_bb + np.eye(m_bb.shape[0], dtype=np.float32) * 1e-6, m_bs)
+        block = h_inv_upper[start:stop, start:stop]
+        block_to_rest = h_inv_upper[start:stop, stop:]
+        try:
+            from scipy.linalg import solve_triangular
+
+            update = solve_triangular(block, block_to_rest, lower=False)
+        except ImportError:
+            update = np.linalg.solve(block, block_to_rest)
         work[:, stop:] -= (work[:, start:stop] - recon) @ update
     except Exception:
         return
@@ -266,7 +279,7 @@ def quantize_hadamard(
     groups = k // GROUP_SIZE
     indices = np.zeros((n, k), dtype=np.uint8)
     norms = np.zeros((n, groups), dtype=np.float16)
-    h_inv = None
+    h_inv_upper = None
     if use_gptq and hessian is not None:
         try:
             h = np.asarray(hessian, dtype=np.float32)
@@ -277,9 +290,9 @@ def quantize_hadamard(
                 s = np.clip(input_scale.astype(np.float32), 1e-6, None)
                 h = h / (s[:, None] * s[None, :])
                 h = h + np.eye(h.shape[0], dtype=np.float32) * (0.01 * np.mean(np.diag(h)) + 1e-6)
-                h_inv = np.linalg.inv(h)
+                h_inv_upper = _upper_cholesky_of_inverse(h)
         except Exception:
-            h_inv = None
+            h_inv_upper = None
     for g in range(groups):
         start, stop = g * GROUP_SIZE, (g + 1) * GROUP_SIZE
         group = work[:, start:stop]
@@ -289,8 +302,8 @@ def quantize_hadamard(
         recon = (codebook[idx] @ rot.T) * row_norms[:, None]
         indices[:, start:stop] = idx
         norms[:, g] = row_norms.astype(np.float16)
-        _gptq_correct_group(work, recon, h_inv, start, stop)
-    return CQTensor(indices=indices, norms=norms, input_scale=input_scale.astype(np.float16), bits=bits, gptq_used=bool(use_gptq and h_inv is not None))
+        _gptq_correct_group(work, recon, h_inv_upper, start, stop)
+    return CQTensor(indices=indices, norms=norms, input_scale=input_scale.astype(np.float16), bits=bits, gptq_used=bool(use_gptq and h_inv_upper is not None))
 
 
 def quantize_orthogonal(weight, bits: int = 4, seed: int = 1234, input_scale: np.ndarray | None = None) -> CQTensor:

python/cactus/convert/tests/test_cq.py

@@ -155,3 +155,89 @@ def test_pointwise_conv1d_int8_preserves_rank3_shape(tmp_path):
     assert scales_bytes == 3 * 2 * 2
     assert group_size == GROUP_SIZE
     assert num_groups == 2
+
+
+def _correlated_inputs(rng: np.random.Generator, samples: int, k: int, rank: int = 32) -> np.ndarray:
+    latent = rng.standard_normal((samples, rank), dtype=np.float32)
+    mixing = rng.standard_normal((rank, k), dtype=np.float32)
+    noise = 0.1 * rng.standard_normal((samples, k), dtype=np.float32)
+    return (latent @ mixing + noise).astype(np.float32)
+
+
+def _damp_like_cq(h: np.ndarray) -> np.ndarray:
+    return h + np.eye(h.shape[0], dtype=np.float32) * (0.01 * np.mean(np.diag(h)) + 1e-6)
+
+
+def test_gptq_cholesky_block_update_matches_explicit_obs():
+    from cactus.convert.quantization import cq as cq_mod
+
+    rng = np.random.default_rng(10)
+    k = 3 * cq_mod.GROUP_SIZE
+    x = _correlated_inputs(rng, 2000, k)
+    h = _damp_like_cq((x.T @ x / x.shape[0]).astype(np.float32))
+    h64 = h.astype(np.float64)
+
+    build_factor = getattr(cq_mod, "_gptq_cholesky_factor", None)
+    factor = build_factor(h) if build_factor is not None else np.linalg.cholesky(np.linalg.inv(h64)).T.astype(np.float32)
+
+    work = rng.standard_normal((8, k), dtype=np.float32)
+    work_ref = work.astype(np.float64)
+    for g in range(k // cq_mod.GROUP_SIZE):
+        start, stop = g * cq_mod.GROUP_SIZE, (g + 1) * cq_mod.GROUP_SIZE
+        recon = np.round(work[:, start:stop], 1).astype(np.float32)
+        cq_mod._gptq_correct_group(work, recon, factor, start, stop)
+        if stop < k:
+            # Reference: the inverse Hessian restricted to the still-unquantized columns,
+            # recomputed from scratch each step (equal to the Schur-updated OBS inverse).
+            hinv_cur = np.linalg.inv(h64[start:, start:])
+            gsz = stop - start
+            update_ref = np.linalg.solve(hinv_cur[:gsz, :gsz], hinv_cur[:gsz, gsz:])
+            err = work_ref[:, start:stop] - recon.astype(np.float64)
+            work_ref[:, stop:] -= err @ update_ref
+        np.testing.assert_allclose(work[:, stop:], work_ref[:, stop:], rtol=1e-4, atol=1e-4)
+        work_ref[:, start:stop] = work[:, start:stop]
+
+
+def test_gptq_reduces_calibrated_error(tmp_path):
+    rng = np.random.default_rng(11)
+    n, k = 16, 384
+    w = rng.standard_normal((n, k), dtype=np.float32)
+    x = _correlated_inputs(rng, 4096, k)
+    h = (x.T @ x / x.shape[0]).astype(np.float32)
+
+    plain = quantize_hadamard(w, bits=4)
+    gptq = quantize_hadamard(w, bits=4, hessian=h, use_gptq=True)
+    assert not plain.gptq_used
+    assert gptq.gptq_used
+
+    def reconstruct(cq, name):
+        path = tmp_path / name
+        write_cq_tensor(path, cq)
+        return dequantize_cq_file(path, read_header(path), torch.float32, 4).numpy()
+
+    w_plain = reconstruct(plain, "plain.weights")
+    w_gptq = reconstruct(gptq, "gptq.weights")
+    err_plain = np.linalg.norm(x @ (w - w_plain).T)
+    err_gptq = np.linalg.norm(x @ (w - w_gptq).T)
+    assert np.isfinite(err_gptq)
+    assert err_gptq < err_plain, (err_gptq, err_plain)
+
+
+def test_gptq_disabled_on_bad_hessian():
+    rng = np.random.default_rng(12)
+    w = rng.standard_normal((4, 256), dtype=np.float32)
+    baseline = quantize_hadamard(w, bits=4)
+
+    negative = quantize_hadamard(w, bits=4, hessian=-np.eye(256, dtype=np.float32) * 1e6, use_gptq=True)
+    assert not negative.gptq_used
+    np.testing.assert_array_equal(negative.indices, baseline.indices)
+
+
+def test_gptq_disabled_on_nan_hessian():
+    rng = np.random.default_rng(12)
+    w = rng.standard_normal((4, 256), dtype=np.float32)
+    baseline = quantize_hadamard(w, bits=4)
+    nan_h = np.full((256, 256), np.nan, dtype=np.float32)
+    broken = quantize_hadamard(w, bits=4, hessian=nan_h, use_gptq=True)
+    np.testing.assert_array_equal(broken.indices, baseline.indices)
+    assert not broken.gptq_used