diff --git a/c/tools/iq3xxs_grid.json b/c/tools/iq3xxs_grid.json new file mode 100644 index 0000000..3cde1a2 --- /dev/null +++ b/c/tools/iq3xxs_grid.json @@ -0,0 +1 @@ +[[4, 4, 4, 4], [20, 4, 4, 4], [36, 4, 4, 4], [12, 12, 4, 4], [28, 12, 4, 4], [62, 12, 4, 4], [4, 20, 4, 4], [20, 20, 4, 4], [12, 28, 4, 4], [20, 36, 4, 4], [28, 62, 4, 4], [44, 62, 4, 4], [12, 4, 12, 4], [28, 4, 12, 4], [4, 12, 12, 4], [20, 12, 12, 4], [12, 20, 12, 4], [44, 20, 12, 4], [4, 28, 12, 4], [20, 28, 12, 4], [12, 36, 12, 4], [36, 44, 12, 4], [4, 62, 12, 4], [4, 4, 20, 4], [20, 4, 20, 4], [36, 4, 20, 4], [12, 12, 20, 4], [4, 20, 20, 4], [20, 20, 20, 4], [12, 28, 20, 4], [28, 28, 20, 4], [62, 28, 20, 4], [12, 44, 20, 4], [62, 44, 20, 4], [44, 62, 20, 4], [12, 4, 28, 4], [62, 4, 28, 4], [4, 12, 28, 4], [20, 12, 28, 4], [44, 20, 28, 4], [4, 62, 28, 4], [28, 12, 36, 4], [62, 28, 36, 4], [36, 36, 36, 4], [62, 44, 36, 4], [28, 62, 36, 4], [44, 62, 36, 4], [12, 4, 44, 4], [62, 4, 44, 4], [20, 28, 44, 4], [20, 44, 44, 4], [44, 28, 52, 4], [36, 52, 52, 4], [4, 12, 62, 4], [36, 12, 62, 4], [52, 12, 62, 4], [28, 36, 62, 4], [12, 52, 62, 4], [12, 4, 4, 12], [28, 4, 4, 12], [4, 12, 4, 12], [20, 12, 4, 12], [12, 20, 4, 12], [28, 20, 4, 12], [4, 28, 4, 12], [20, 28, 4, 12], [36, 28, 4, 12], [62, 36, 4, 12], [4, 44, 4, 12], [4, 4, 12, 12], [20, 4, 12, 12], [12, 12, 12, 12], [4, 20, 12, 12], [20, 20, 12, 12], [12, 4, 20, 12], [28, 4, 20, 12], [4, 12, 20, 12], [20, 12, 20, 12], [12, 20, 20, 12], [4, 28, 20, 12], [20, 62, 20, 12], [4, 4, 28, 12], [20, 4, 28, 12], [4, 20, 28, 12], [12, 28, 28, 12], [52, 36, 28, 12], [52, 52, 28, 12], [12, 4, 36, 12], [44, 4, 36, 12], [4, 44, 36, 12], [4, 20, 44, 12], [36, 20, 44, 12], [52, 36, 44, 12], [12, 62, 44, 12], [44, 4, 52, 12], [20, 20, 62, 12], [4, 36, 62, 12], [4, 4, 4, 20], [20, 4, 4, 20], [12, 12, 4, 20], [28, 12, 4, 20], [4, 20, 4, 20], [20, 20, 4, 20], [52, 20, 4, 20], [12, 28, 4, 20], [20, 36, 4, 20], [12, 4, 12, 20], [28, 4, 12, 20], [44, 4, 12, 20], [4, 12, 12, 20], [20, 12, 12, 20], [12, 20, 12, 20], [4, 28, 12, 20], [28, 52, 12, 20], [62, 52, 12, 20], [4, 62, 12, 20], [4, 4, 20, 20], [20, 4, 20, 20], [12, 12, 20, 20], [62, 12, 20, 20], [4, 20, 20, 20], [20, 20, 20, 20], [62, 28, 20, 20], [4, 36, 20, 20], [44, 44, 20, 20], [12, 4, 28, 20], [4, 12, 28, 20], [36, 12, 28, 20], [4, 62, 28, 20], [36, 62, 28, 20], [44, 28, 36, 20], [28, 44, 36, 20], [28, 4, 44, 20], [62, 20, 44, 20], [12, 36, 44, 20], [36, 62, 44, 20], [12, 4, 62, 20], [28, 4, 62, 20], [52, 12, 62, 20], [44, 36, 62, 20], [12, 4, 4, 28], [4, 12, 4, 28], [20, 12, 4, 28], [12, 20, 4, 28], [28, 20, 4, 28], [4, 44, 4, 28], [44, 52, 4, 28], [20, 62, 4, 28], [4, 4, 12, 28], [20, 4, 12, 28], [4, 20, 12, 28], [12, 28, 12, 28], [36, 36, 12, 28], [52, 36, 12, 28], [12, 4, 20, 28], [28, 4, 20, 28], [4, 12, 20, 28], [44, 20, 20, 28], [20, 44, 20, 28], [20, 62, 20, 28], [12, 12, 28, 28], [28, 28, 28, 28], [4, 28, 36, 28], [62, 36, 36, 28], [20, 62, 36, 28], [4, 4, 44, 28], [52, 4, 44, 28], [20, 20, 44, 28], [44, 44, 44, 28], [36, 12, 52, 28], [52, 28, 52, 28], [28, 52, 52, 28], [28, 28, 62, 28], [4, 52, 62, 28], [36, 4, 4, 36], [62, 12, 4, 36], [44, 28, 4, 36], [62, 28, 4, 36], [28, 44, 4, 36], [62, 44, 4, 36], [36, 62, 12, 36], [4, 20, 20, 36], [62, 28, 20, 36], [4, 36, 20, 36], [4, 52, 20, 36], [52, 52, 20, 36], [62, 4, 28, 36], [44, 36, 28, 36], [36, 4, 36, 36], [12, 44, 36, 36], [36, 52, 36, 36], [44, 20, 44, 36], [28, 36, 44, 36], [4, 62, 44, 36], [44, 4, 62, 36], [4, 12, 62, 36], [20, 12, 62, 36], [4, 28, 62, 36], [20, 12, 4, 44], [12, 36, 4, 44], [4, 62, 4, 44], [4, 4, 12, 44], [52, 4, 12, 44], [52, 20, 12, 44], [44, 44, 12, 44], [36, 12, 20, 44], [20, 28, 20, 44], [20, 62, 20, 44], [20, 4, 28, 44], [28, 44, 28, 44], [4, 12, 36, 44], [28, 20, 36, 44], [62, 20, 36, 44], [20, 62, 36, 44], [20, 4, 44, 44], [12, 28, 44, 44], [4, 44, 52, 44], [36, 20, 62, 44], [20, 36, 62, 44], [36, 20, 4, 52], [36, 36, 4, 52], [52, 36, 4, 52], [36, 52, 4, 52], [12, 20, 12, 52], [12, 52, 12, 52], [62, 12, 20, 52], [36, 52, 20, 52], [4, 28, 28, 52], [52, 28, 28, 52], [36, 36, 36, 52], [44, 4, 44, 52], [20, 44, 44, 52], [28, 28, 52, 52], [28, 4, 62, 52], [12, 20, 62, 52], [28, 4, 4, 62], [44, 4, 4, 62], [62, 4, 4, 62], [4, 12, 4, 62], [20, 28, 4, 62], [20, 44, 4, 62], [52, 20, 12, 62], [4, 36, 12, 62], [20, 12, 20, 62], [44, 36, 20, 62], [20, 44, 20, 62], [4, 4, 28, 62], [44, 12, 28, 62], [28, 28, 28, 62], [4, 52, 28, 62], [12, 20, 36, 62], [12, 36, 36, 62], [4, 4, 44, 62], [20, 4, 44, 62], [36, 20, 44, 62], [4, 28, 52, 62]] \ No newline at end of file diff --git a/c/tools/quant_ablation.py b/c/tools/quant_ablation.py index dd08715..a602ae8 100644 --- a/c/tools/quant_ablation.py +++ b/c/tools/quant_ablation.py @@ -117,11 +117,79 @@ def quantize_param(w, bits, group, rot=False, e8=""): def _grid_or_e8(x, bits, group, e8): + if e8 == "-iq3": + return _quant_iq3(x.float()) if e8: return _quant_e8(x.float(), group, bits, ball=(e8 == "-e8")) return _quant_last_dim(x, bits, group) +# -------------------------------------------------------------------------------------- +# IQ3_XXS-style codebook (#452 candidate (a)): llama.cpp's deployed 3.06-bpw scheme. +# 4-dim magnitude blocks quantized to a 256-entry lattice-subset grid (magnitudes on the +# odd ladder 4,12,..,62 in half-units), signs factored out per 8 weights with an odd-parity +# constraint (7 stored + 1 derived: a block whose true signs violate parity gets its +# smallest-magnitude sign flipped — modelled here so the ablation pays the real cost). +# Scales: fp16 super-scale per 256 + 4-bit sub-scale per 32, db = d*(0.5+s)*0.5. +# Grid extracted from ggml-common.h (MIT). +# -------------------------------------------------------------------------------------- +_IQ3_GRID = None +def _iq3_grid(device): + global _IQ3_GRID + if _IQ3_GRID is None or _IQ3_GRID.device != device: + import json, os + path = os.path.join(os.path.dirname(__file__), "iq3xxs_grid.json") + _IQ3_GRID = torch.tensor(json.load(open(path)), dtype=torch.float32, device=device) + return _IQ3_GRID # [256,4], half-unit magnitudes (value/2 = weight units) + +def _quant_iq3(x): + orig = x.shape + K = orig[-1] + assert K % 256 == 0, "iq3 needs multiples of 256 along the input dim" + xb = x.reshape(-1, 256) # super-blocks + grid = _iq3_grid(x.device) * 0.5 # weight units + out = torch.empty_like(xb) + signs = torch.sign(xb); signs[signs == 0] = 1.0 + mags = xb.abs() + for sb in range(8): # 8 sub-blocks of 32 + m = mags[:, sb*32:(sb+1)*32] # [N,32] + s = signs[:, sb*32:(sb+1)*32] + # per-8 sign parity: flip the smallest-|w| sign where the product is negative + s8 = s.reshape(-1, 4, 8) + m8 = m.reshape(-1, 4, 8) + viol = (s8.prod(-1) < 0) # odd number of minus signs + idxmin = m8.argmin(-1) + flip = torch.zeros_like(s8) + flip.scatter_(-1, idxmin[..., None], 1.0) + s8 = torch.where(viol[..., None].expand_as(s8) & (flip > 0), -s8, s8) + s = s8.reshape(-1, 32) + # sub-scale search: db candidates from the 4-bit code, super d from block RMS + d = m.pow(2).mean(-1, keepdim=True).sqrt() / 20.0 + 1e-12 # rough anchor + best = None + for code in range(16): + db = d * (0.5 + code) * 0.5 + q = m / db # [N,32] target magnitudes + q4 = q.reshape(-1, 4) # 4-dim grid blocks + # chunked argmin ||q-g||^2 = argmin(|g|^2 - 2 q.g): a full cdist on a + # 100M-param tensor materializes tens of GB — this stays at ~256 MB. + g2 = grid.pow(2).sum(-1) + idx = torch.empty(q4.shape[0], dtype=torch.long, device=q4.device) + CH = 1 << 18 + for i0 in range(0, q4.shape[0], CH): + cc = q4[i0:i0+CH] + idx[i0:i0+CH] = (g2 - 2.0 * (cc @ grid.T)).argmin(-1) + hit = grid[idx].reshape(-1, 8, 4) + rec = (hit.reshape(-1, 32) * db) + err = (rec - m).pow(2).sum(-1, keepdim=True) + if best is None: + best = (err, rec) + else: + take = err < best[0] + best = (torch.where(take, err, best[0]), torch.where(take, rec, best[1])) + out[:, sb*32:(sb+1)*32] = best[1] * s + return out.reshape(orig) + + def _rot_quant(x, bits, group, e8=""): """W -> Qn(W@Q) @ Q^T along the last (input) dim — see rotation() above.""" q = rotation(x.shape[-1], x.device) @@ -206,7 +274,7 @@ def _quant_e8(x, group, bits, ball): return best_out.reshape(shp) -SCHEME_RE = re.compile(r"^int(2|3|4|8)(?:-g(\d+))?(-e8u?)?(-rot)?(-nohead)?$") +SCHEME_RE = re.compile(r"^int(2|3|4|8)(?:-g(\d+))?(-e8u?|-iq3)?(-rot)?(-nohead)?$") def parse_scheme(name):