| """GPTQ-style error-compensated ternary latents (T5.0). |
| |
| GPTQ quantises a weight's input columns one at a time and pushes each column's rounding error onto the columns not |
| yet quantised, weighted by the inverse Hessian of the layer's inputs (H = XᵀX). Here the quantiser is the student's |
| own (Q-A: ternary codes, per-row absmean scale per group of 128 input columns), and the result is a *latent*: the |
| trainer re-derives scale and codes from it on every forward (`fake_quant`). So the latent of each group is the |
| error-compensated values GPTQ actually rounded. The group's scale is their absmean, solved by re-running the group |
| until the scale it uses is the absmean of the values it produces, which makes `fake_quant(latent)` give GPTQ's codes. |
| """ |
| from __future__ import annotations |
|
|
| import torch |
|
|
| from lowbit.quant.ternary import QuantSpec, init_scale, ternary_codes |
|
|
|
|
| def _codes(w: torch.Tensor, s: torch.Tensor) -> torch.Tensor: |
| """Ternary codes of one column [E, R] against per-row scales s [E, R] (as `ternary_codes`, absmean).""" |
| return torch.sign(w) * (w.abs() > 0.5 * s) |
|
|
|
|
| def _inverse_cholesky_upper(H: torch.Tensor, damp: float, tries: int = 4) -> tuple[torch.Tensor, torch.Tensor, int]: |
| """Upper Cholesky factor of H⁻¹ per batch element, in fp64. An element that is not positive-definite retries |
| with 10× the dampening; after `tries` it gets the identity (round-to-nearest, no compensation). Returns |
| (U, dampening used per element, number of elements left on the identity).""" |
| E, C, _ = H.shape |
| eye = torch.eye(C, dtype=H.dtype, device=H.device) |
| mean_diag = torch.diagonal(H, dim1=-2, dim2=-1).mean(-1) |
| U = eye.expand(E, C, C).clone() |
| used = torch.full((E,), float("nan"), dtype=H.dtype, device=H.device) |
| todo = torch.ones(E, dtype=torch.bool, device=H.device) |
| d = damp |
| for _ in range(tries): |
| Hd = H[todo] + (d * mean_diag[todo])[:, None, None] * eye |
| L, info = torch.linalg.cholesky_ex(Hd) |
| Uc, info2 = torch.linalg.cholesky_ex(torch.cholesky_inverse(L), upper=True) |
| good = (info == 0) & (info2 == 0) & torch.isfinite(Uc).flatten(1).all(-1) |
| idx = todo.nonzero()[:, 0] |
| U[idx[good]] = Uc[good] |
| used[idx[good]] = d |
| todo[idx[good]] = False |
| if not todo.any(): |
| break |
| d *= 10 |
| return U, used, int(todo.sum()) |
|
|
|
|
| @torch.no_grad() |
| def gptq_ternary(w: torch.Tensor, h: torch.Tensor, spec: QuantSpec, damp: float = 0.01, |
| iters: int = 3) -> tuple[torch.Tensor, dict]: |
| """w [E, R, C] (rows = outputs, C = inputs), h [E, C, C] → error-compensated latent [E, R, C] in fp32. |
| |
| Batched over E (a layer's experts). Inputs with a zero Hessian diagonal (never active) are rounded like any |
| other column but pass no error on (their diagonal is set to 1). |
| """ |
| if spec.method != "absmean": |
| raise NotImplementedError("GPTQ init is built for Q-A (absmean) only") |
| g = spec.group_size |
| E, R, C = w.shape |
| W = w.detach().float().clone() |
| H = h.detach().double().clone() |
| diag = torch.diagonal(H, dim1=-2, dim2=-1) |
| diag[diag == 0] = 1.0 |
| U, damp_used, n_identity = _inverse_cholesky_upper(H, damp) |
| U = U.float() |
| latent = torch.empty_like(W) |
| used = torch.empty(E, R, C, dtype=torch.int8, device=W.device) |
| passes = [] |
| for g0 in range(0, C, g): |
| g1 = min(g0 + g, C) |
| start = W[:, :, g0:g1].clone() |
| s = init_scale(start, QuantSpec(method="absmean", group_size=g1 - g0)) |
| for it in range(iters): |
| Wg = start.clone() |
| err = torch.empty_like(Wg) |
| for j in range(g1 - g0): |
| i = g0 + j |
| col = Wg[:, :, j] |
| c = _codes(col, s[..., 0]) |
| used[:, :, i] = c.to(torch.int8) |
| q = c * s[..., 0] |
| e = (col - q) / U[:, i, i][:, None] |
| err[:, :, j] = e |
| if j + 1 < g1 - g0: |
| Wg[:, :, j + 1:] -= e[:, :, None] * U[:, i, i + 1:g1][:, None, :] |
| s_new = init_scale(Wg, QuantSpec(method="absmean", group_size=g1 - g0)) |
| same = torch.equal(s_new, s) |
| s = s_new |
| if same: |
| break |
| passes.append(it + 1) |
| latent[:, :, g0:g1] = Wg |
| if g1 < C: |
| W[:, :, g1:] -= torch.bmm(err, U[:, g0:g1, g1:]) |
| codes_latent = ternary_codes(latent, init_scale(latent, spec), spec) |
| return latent, {"group_passes_max": max(passes), "group_passes_mean": sum(passes) / len(passes), |
| "zero_fraction": float((codes_latent == 0).float().mean()), |
| "code_mismatch": float((codes_latent != used).float().mean()), |
| "damp_max": float(damp_used.nan_to_num(0).max()), "damp_raised": int((damp_used > damp).sum()), |
| "identity_fallback": n_identity} |
|
|
