diff --git a/python/cactus/convert/quantization/cq.py b/python/cactus/convert/quantization/cq.py index 2d3942389..0502ec31b 100644 --- a/python/cactus/convert/quantization/cq.py +++ b/python/cactus/convert/quantization/cq.py @@ -230,13 +230,15 @@ 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 _gptq_correct_group(work: np.ndarray, recon: np.ndarray, u_factor: np.ndarray | None, start: int, stop: int) -> None: + if u_factor 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) + from scipy.linalg import solve_triangular + + u_bb = u_factor[start:stop, start:stop].astype(np.float32, copy=False) + u_bs = u_factor[start:stop, stop:].astype(np.float32, copy=False) + update = solve_triangular(u_bb, u_bs, lower=False) work[:, stop:] -= (work[:, start:stop] - recon) @ update except Exception: return @@ -266,7 +268,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 + u_factor = None if use_gptq and hessian is not None: try: h = np.asarray(hessian, dtype=np.float32) @@ -277,9 +279,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) + u_factor = np.linalg.cholesky(np.linalg.inv(h)).T except Exception: - h_inv = None + u_factor = None for g in range(groups): start, stop = g * GROUP_SIZE, (g + 1) * GROUP_SIZE group = work[:, start:stop] @@ -289,8 +291,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, u_factor, start, stop) + return CQTensor(indices=indices, norms=norms, input_scale=input_scale.astype(np.float16), bits=bits, gptq_used=bool(use_gptq and u_factor is not None)) def quantize_orthogonal(weight, bits: int = 4, seed: int = 1234, input_scale: np.ndarray | None = None) -> CQTensor: diff --git a/python/cactus/convert/tests/test_cq.py b/python/cactus/convert/tests/test_cq.py index 2cd816e25..3c05ef2ac 100644 --- a/python/cactus/convert/tests/test_cq.py +++ b/python/cactus/convert/tests/test_cq.py @@ -15,7 +15,15 @@ ) from cactus.convert.export.qdq import FLAG_INTERLEAVED_4ROW, FLAG_ORTHOGONAL_ROTATION, dequantize_cq_file, read_header from cactus.convert.interleave_orthogonal_cq4 import interleave_orthogonal_cq4_file -from cactus.convert.quantization.cq import PRECISION_CQ, pack_indices_lsb, quantize_hadamard, quantize_orthogonal, write_cq_tensor +from cactus.convert.quantization.cq import ( + GROUP_SIZE as CQ_GROUP_SIZE, + PRECISION_CQ, + _gptq_correct_group, + pack_indices_lsb, + quantize_hadamard, + quantize_orthogonal, + write_cq_tensor, +) def test_pack_indices_lsb_bits(): @@ -155,3 +163,76 @@ 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 test_gptq_multiblock_compensation_matches_progressive_obs(): + rng = np.random.default_rng(12345) + k = 384 + b = CQ_GROUP_SIZE + assert k == 3 * b + + # Construct deterministic SPD correlated Hessian + x = rng.standard_normal((k * 2, k)).astype(np.float32) + h = (x.T @ x) / float(k * 2) + h = h + np.eye(k, dtype=np.float32) * (0.01 * np.mean(np.diag(h)) + 1e-6) + + h_inv_global = np.linalg.inv(h) + u_factor = np.linalg.cholesky(h_inv_global).T + + work_orig = rng.standard_normal((8, k)).astype(np.float32) + recon_orig = rng.standard_normal((8, k)).astype(np.float32) + tol = 1e-4 + + for g in range(k // b - 1): + start, stop = g * b, (g + 1) * b + work = work_orig.copy() + work_static = work_orig.copy() + + # 1. Independent ground-truth OBS reference from remaining sub-Hessian + h_rem = h[start:, start:] + h_rem_inv = np.linalg.inv(h_rem) + ref_bb = h_rem_inv[0:b, 0:b] + ref_bs = h_rem_inv[0:b, b:] + ref_update = np.linalg.solve(ref_bb, ref_bs) + + residual = work_orig[:, start:stop] - recon_orig[:, start:stop] + expected_work_rem = work_orig[:, stop:] - residual @ ref_update + + # 2. Production Cholesky implementation passes OBS invariant + _gptq_correct_group(work, recon_orig[:, start:stop], u_factor, start, stop) + actual_work_rem = work[:, stop:] + actual_diff = float(np.max(np.abs(actual_work_rem - expected_work_rem))) + assert actual_diff <= tol, f"Block {g} failed progressive OBS invariant: {actual_diff} > {tol}" + + # 3. Old static-inverse implementation violates the invariant on correlated blocks g >= 1 + old_m_bb = h_inv_global[start:stop, start:stop] + old_m_bs = h_inv_global[start:stop, stop:] + old_update = np.linalg.solve(old_m_bb + np.eye(b, dtype=np.float32) * 1e-6, old_m_bs) + old_work_rem = work_static[:, stop:] - residual @ old_update + old_diff = float(np.max(np.abs(old_work_rem - expected_work_rem))) + + if g == 0: + assert old_diff <= tol, f"Block 0 unexpectedly diverged on static inverse: {old_diff} > {tol}" + else: + assert old_diff > tol, f"Old static inverse unexpectedly passed tolerance on block {g}: {old_diff} <= {tol}" + + +def test_quantize_hadamard_with_gptq_multiblock(): + rng = np.random.default_rng(2026) + k = 384 + n = 16 + w = rng.standard_normal((n, k)).astype(np.float32) + x = rng.standard_normal((k * 2, k)).astype(np.float32) + h = (x.T @ x) / float(k * 2) + + cq_gptq = quantize_hadamard(w, bits=4, hessian=h, use_gptq=True) + cq_rtn = quantize_hadamard(w, bits=4, hessian=h, use_gptq=False) + + assert cq_gptq.gptq_used is True + assert cq_rtn.gptq_used is False + assert cq_gptq.indices.shape == (n, k) + assert cq_gptq.norms.shape == (n, k // CQ_GROUP_SIZE) + # Block 0 has no prior compensation, so index decisions match RTN + assert np.array_equal(cq_gptq.indices[:, :CQ_GROUP_SIZE], cq_rtn.indices[:, :CQ_GROUP_SIZE]) + # Multi-block compensation actively modifies subsequent quantization decisions + assert not np.array_equal(cq_gptq.indices[:, CQ_GROUP_SIZE:], cq_rtn.indices[:, CQ_GROUP_SIZE:])