diff --git a/python/cactus/convert/quantization/cq.py b/python/cactus/convert/quantization/cq.py index 2d3942389..c3ffa23e0 100644 --- a/python/cactus/convert/quantization/cq.py +++ b/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: diff --git a/python/cactus/convert/tests/test_cq.py b/python/cactus/convert/tests/test_cq.py index 2cd816e25..9e0337449 100644 --- a/python/cactus/convert/tests/test_cq.py +++ b/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