| """D-58 · llama.cpp's codebook formats as trainable students, started from llama.cpp's own files. |
| |
| A student starts from llama.cpp's imatrix quantisation of the teacher (`scripts/ptq_baseline.py`): its expert |
| tensors are read from that GGUF as raw blocks and kept. Every format here stores each block as |
| |
| weight = eff(d, s) · v |
| |
| where v is a small signed integer per weight (a codebook entry times its sign: the grid value of IQ2/IQ3, the |
| IQ4_NL table value), d is the block's fp16 scale and s its sub-blocks' integer scales: |
| |
| format block sub eff s range bytes (d at bytes 0-1 in every format) |
| IQ2_XXS 256 32 d·(0.5 + s)·0.25 0..15 66 s: top 4 bits of each sub-block's 2nd uint32 |
| IQ2_XS 256 16 d·(0.5 + s)·0.25 0..15 74 s: 16 nibbles at bytes 66-73 |
| IQ2_S 256 16 d·(0.5 + s)·0.25 0..15 82 s: 16 nibbles at bytes 74-81 |
| IQ3_XXS 256 32 d·(0.5 + s)·0.5 0..15 98 s: top 4 bits of the uint32s at bytes 66-97 |
| IQ3_S 256 32 d·(1 + 2s) 0..15 110 s: 8 nibbles at bytes 106-109 |
| IQ4_XS 256 32 d·(s − 32) 0..63 136 s: 4 low bits at bytes 4-7, 2 high at bytes 2-3 |
| IQ4_NL 32 32 d — 18 (no sub-scale) |
| |
| (ggml-quants.c's dequantize_row_*; gguf-py's `dequantize_blocks` are the reference the tests compare against.) |
| The codes — which codebook entry and which signs — are the bytes other than d and s. They are kept unchanged |
| ("frozen codes"): training moves d and s only, on their own grids, plus the bf16 attention and norms. The export |
| writes the new d and s into the original blocks, so an untrained student exports byte-identical to llama.cpp's |
| file, and every exported block is one llama.cpp wrote except for its scales. |
| |
| Trainable parameters per tensor, both fp32: `dlog` per block (d = d0·exp(dlog), rounded to fp16 in the forward) |
| and `slat` per sub-block (s = round(slat) clamped to the range). The gradient reaches both through the plain |
| straight-through estimator, like the latents of `quant.ternary`. v·eff is exact in fp32 (d: 11 significant bits, |
| the s factor ≤ 7, v ≤ 7), as ggml computes it; the forward then rounds to bf16 like every other student. |
| """ |
| from __future__ import annotations |
|
|
| from dataclasses import dataclass |
|
|
| import numpy as np |
| import torch |
|
|
|
|
| @dataclass(frozen=True) |
| class Format: |
| name: str |
| block: int |
| sub: int |
| nbytes: int |
| a: float |
| b: float |
| s_max: int |
| neutral_s: int |
| unit: float |
|
|
|
|
| FORMATS = { |
| "IQ2_XXS": Format("IQ2_XXS", 256, 32, 66, 0.25, 0.5, 15, 0, 0.125), |
| "IQ2_XS": Format("IQ2_XS", 256, 16, 74, 0.25, 0.5, 15, 0, 0.125), |
| "IQ2_S": Format("IQ2_S", 256, 16, 82, 0.25, 0.5, 15, 0, 0.125), |
| "IQ3_XXS": Format("IQ3_XXS", 256, 32, 98, 0.5, 0.5, 15, 0, 0.25), |
| "IQ3_S": Format("IQ3_S", 256, 32, 110, 2.0, 0.5, 15, 0, 1.0), |
| "IQ4_XS": Format("IQ4_XS", 256, 32, 136, 1.0, -32.0, 63, 33, 1.0), |
| "IQ4_NL": Format("IQ4_NL", 32, 32, 18, 1.0, 0.0, 0, 0, 1.0), |
| } |
|
|
|
|
| |
|
|
| def get_d(blocks: np.ndarray) -> np.ndarray: |
| return blocks[:, :2].copy().view(np.float16)[:, 0] |
|
|
|
|
| def set_d(blocks: np.ndarray, d: np.ndarray) -> None: |
| blocks[:, :2] = np.ascontiguousarray(d.astype(np.float16)).reshape(-1, 1).view(np.uint8) |
|
|
|
|
| def _nibbles(b: np.ndarray) -> np.ndarray: |
| return np.stack([b & 15, b >> 4], -1).reshape(b.shape[0], -1) |
|
|
|
|
| def _pack_nibbles(s: np.ndarray) -> np.ndarray: |
| s = s.astype(np.uint8).reshape(s.shape[0], -1, 2) |
| return s[..., 0] | (s[..., 1] << 4) |
|
|
|
|
| def _u32(blocks: np.ndarray, lo: int, hi: int) -> np.ndarray: |
| return blocks[:, lo:hi].copy().view("<u4") |
|
|
|
|
| def _set_u32_top4(blocks: np.ndarray, lo: int, hi: int, s: np.ndarray) -> None: |
| w = _u32(blocks, lo, hi) |
| w = (w & np.uint32(0x0FFFFFFF)) | (s.astype(np.uint32) << np.uint32(28)) |
| blocks[:, lo:hi] = w.view(np.uint8) |
|
|
|
|
| def get_s(fmt: Format, blocks: np.ndarray) -> np.ndarray: |
| """The sub-block scales, uint8 [n, block/sub] (IQ4_NL: an empty [n, 0]).""" |
| n = fmt.name |
| if n == "IQ2_XXS": |
| return (_u32(blocks, 2, 66)[:, 1::2] >> 28).astype(np.uint8) |
| if n == "IQ3_XXS": |
| return (_u32(blocks, 66, 98) >> 28).astype(np.uint8) |
| if n == "IQ2_XS": |
| return _nibbles(blocks[:, 66:74]) |
| if n == "IQ2_S": |
| return _nibbles(blocks[:, 74:82]) |
| if n == "IQ3_S": |
| return _nibbles(blocks[:, 106:110]) |
| if n == "IQ4_XS": |
| lo = _nibbles(blocks[:, 4:8]) |
| hi = blocks[:, 2:4].copy().view("<u2")[:, :1] >> (2 * np.arange(8, dtype=np.uint16)) |
| return (lo | ((hi & 3).astype(np.uint8) << 4)).astype(np.uint8) |
| if n == "IQ4_NL": |
| return np.zeros((blocks.shape[0], 0), np.uint8) |
| raise ValueError(n) |
|
|
|
|
| def set_s(fmt: Format, blocks: np.ndarray, s: np.ndarray) -> None: |
| n = fmt.name |
| if n == "IQ2_XXS": |
| w = _u32(blocks, 2, 66) |
| w[:, 1::2] = (w[:, 1::2] & np.uint32(0x0FFFFFFF)) | (s.astype(np.uint32) << np.uint32(28)) |
| blocks[:, 2:66] = w.view(np.uint8) |
| elif n == "IQ3_XXS": |
| _set_u32_top4(blocks, 66, 98, s) |
| elif n == "IQ2_XS": |
| blocks[:, 66:74] = _pack_nibbles(s) |
| elif n == "IQ2_S": |
| blocks[:, 74:82] = _pack_nibbles(s) |
| elif n == "IQ3_S": |
| blocks[:, 106:110] = _pack_nibbles(s) |
| elif n == "IQ4_XS": |
| s = s.astype(np.uint16) |
| blocks[:, 4:8] = _pack_nibbles(s & 15) |
| hi = ((s >> 4) << (2 * np.arange(8, dtype=np.uint16))).sum(-1).astype("<u2") |
| blocks[:, 2:4] = hi.reshape(-1, 1).view(np.uint8) |
| elif n != "IQ4_NL": |
| raise ValueError(n) |
|
|
|
|
| def eff_np(fmt: Format, d: np.ndarray, s: np.ndarray) -> np.ndarray: |
| """fp32 effective scale per sub-block [n, block/sub], as ggml computes it.""" |
| d = d.astype(np.float32).reshape(-1, 1) |
| if fmt.name == "IQ4_NL": |
| return d |
| if fmt.name == "IQ4_XS": |
| return d * (s.astype(np.float32) - np.float32(32)) |
| if fmt.name == "IQ3_S": |
| return d * (np.float32(1) + np.float32(2) * s.astype(np.float32)) |
| return d * (np.float32(0.5) + s.astype(np.float32)) * np.float32(fmt.a) |
|
|
|
|
| def values(fmt: Format, blocks: np.ndarray, gguf) -> np.ndarray: |
| """int8 [n, block]: each weight's v (codebook entry × sign), decoded by gguf-py with d = 1 and a neutral s.""" |
| b = blocks.copy() |
| set_d(b, np.ones(b.shape[0], np.float32)) |
| set_s(fmt, b, np.full((b.shape[0], fmt.block // fmt.sub if fmt.name != "IQ4_NL" else 0), fmt.neutral_s)) |
| v = gguf.quants.dequantize(b, getattr(gguf.GGMLQuantizationType, fmt.name)).reshape(b.shape[0], fmt.block) |
| v = v / np.float32(fmt.unit) |
| q = np.rint(v) |
| if not np.array_equal(q, v) or np.abs(q).max() > 127: |
| raise ValueError(f"{fmt.name}: decoded values are not small integers") |
| return q.astype(np.int8) |
|
|
|
|
| |
|
|
| def round_fp16(x: torch.Tensor) -> torch.Tensor: |
| """x rounded to fp16 (nearest even), as fp32, with the plain STE.""" |
| return x + (x.half().float() - x).detach() |
|
|
|
|
| def round_s(x: torch.Tensor, s_max: int) -> torch.Tensor: |
| return x + (x.round().clamp(0, s_max) - x).detach() |
|
|
|
|
| def eff_torch(fmt: Format, d: torch.Tensor, s: torch.Tensor) -> torch.Tensor: |
| """d [..., nblk], s [..., nblk·block/sub] (both already on their grids) → eff per sub-block, fp32.""" |
| k = fmt.block // fmt.sub |
| dd = d.unsqueeze(-1).expand(*d.shape, k).flatten(-2) |
| if fmt.name == "IQ4_NL": |
| return dd |
| if fmt.name == "IQ4_XS": |
| return dd * (s - 32.0) |
| if fmt.name == "IQ3_S": |
| return dd * (1.0 + 2.0 * s) |
| return dd * (0.5 + s) * fmt.a |
|
|
|
|
| def weight(fmt: Format, v: torch.Tensor, d0: torch.Tensor, dlog: torch.Tensor, slat: torch.Tensor | None |
| ) -> torch.Tensor: |
| """The student's fp32 weight [..., C] from int8 v [..., C], d0 / dlog [..., C/block], slat [..., C/sub].""" |
| d = round_fp16(d0 * torch.exp(dlog)) |
| s = round_s(slat, fmt.s_max) if slat is not None else None |
| e = eff_torch(fmt, d, s) |
| return (v.float().unflatten(-1, (-1, fmt.sub)) * e.unsqueeze(-1)).flatten(-2) |
|
|
|
|
| def trained_fields(fmt: Format, d0: torch.Tensor, dlog: torch.Tensor, slat: torch.Tensor | None |
| ) -> tuple[np.ndarray, np.ndarray]: |
| """(fp16 d [..., nblk], uint8 s [..., nsub]) exactly as `weight` uses them.""" |
| d = (d0 * torch.exp(dlog)).half().cpu().numpy() |
| s = (slat.round().clamp(0, fmt.s_max).to(torch.uint8).cpu().numpy() if slat is not None |
| else np.zeros((*d.shape[:-1], 0), np.uint8)) |
| return d, s |
|
|
|
|
| def dequant(fmt: Format, blocks: np.ndarray, v: np.ndarray) -> np.ndarray: |
| """fp32 weights [..., cols] from blocks [..., nblk, nbytes] and their v [..., cols] (no gguf-py needed).""" |
| flat = blocks.reshape(-1, fmt.nbytes) |
| e = eff_np(fmt, get_d(flat), get_s(fmt, flat)).reshape(*v.shape[:-1], -1) |
| return (v.astype(np.float32).reshape(*v.shape[:-1], -1, fmt.sub) * e[..., None]).reshape(v.shape) |
|
|
|
|
| def write_blocks(fmt: Format, blocks: np.ndarray, d: np.ndarray, s: np.ndarray) -> np.ndarray: |
| """A copy of `blocks` [..., nblk, nbytes] with d [..., nblk] and s [..., nsub] written in.""" |
| out = blocks.reshape(-1, fmt.nbytes).copy() |
| set_d(out, d.reshape(-1)) |
| if fmt.name != "IQ4_NL": |
| set_s(fmt, out, s.reshape(out.shape[0], -1)) |
| return out.reshape(blocks.shape) |
|
|
|
|
| |
| |
| |
| |
| |
| |
| |
| |
| |
| |
|
|
| GROUP = {"IQ3_S": 4, "IQ2_XXS": 8, "IQ4_NL": 1} |
| EVEN_SIGNS = ("IQ2_XXS",) |
| _CHUNK = 1 << 20 |
|
|
|
|
| def _nearest(a: torch.Tensor, grid: torch.Tensor) -> torch.Tensor: |
| """Index of the nearest grid row (L2) for each row of a [N, k]; grid [M, k] fp32. Ties: the lowest index.""" |
| half = 0.5 * (grid * grid).sum(-1) |
| out = torch.empty(a.shape[0], dtype=torch.int64, device=a.device) |
| for i in range(0, a.shape[0], _CHUNK): |
| out[i:i + _CHUNK] = (a[i:i + _CHUNK] @ grid.T - half).argmax(-1) |
| return out |
|
|
|
|
| _NEIGHBOURS = 16 |
| _TABLES: dict = {} |
|
|
|
|
| def _lattice(grid: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: |
| """(the grid's sorted coordinate values [L], for every point of values^k the indices of its 16 nearest grid |
| entries [L^k, 16], nearest first), as llama.cpp's quantisers search: round each magnitude to the nearest |
| coordinate value, then look only at that lattice point's neighbours. Cached per grid tensor.""" |
| key = (grid.data_ptr(), grid.device) |
| if key not in _TABLES: |
| vals = torch.unique(grid) |
| k = grid.shape[1] |
| pts = torch.cartesian_prod(*[vals] * k).reshape(-1, k).float() |
| d2 = ((pts[:, None, :] - grid[None]) ** 2).sum(-1) |
| _TABLES[key] = (vals, d2.topk(min(_NEIGHBOURS, grid.shape[0]), dim=-1, largest=False).indices) |
| return _TABLES[key] |
|
|
|
|
| def _nearest_lattice(a: torch.Tensor, grid: torch.Tensor) -> torch.Tensor: |
| """Nearest grid row for each row of a [N, k] among its rounded lattice point's neighbours (exact for a on the |
| grid: the point is its own nearest neighbour). Ties: the neighbour listed first.""" |
| vals, table = _lattice(grid) |
| L, k = vals.numel(), a.shape[1] |
| mid = (vals[1:] + vals[:-1]) / 2 |
| place = L ** torch.arange(k - 1, -1, -1, device=a.device) |
| out = torch.empty(a.shape[0], dtype=torch.int64, device=a.device) |
| for i in range(0, a.shape[0], _CHUNK): |
| c = a[i:i + _CHUNK] |
| p = (torch.bucketize(c, mid) * place).sum(-1) |
| cand = table[p] |
| d2 = ((c[:, None, :] - grid[cand]) ** 2).sum(-1) |
| out[i:i + _CHUNK] = cand.gather(1, d2.argmin(-1, keepdim=True))[:, 0] |
| return out |
|
|
|
|
| @torch.compiler.disable |
| def free_codes(fmt: Format, x: torch.Tensor, grid: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: |
| """x [..., C] (latent / eff, fp32) → (idx [..., C/k] int64, signed values [..., C] fp32) on the format's codebook. |
| IQ4_NL's grid is its 16 values [16, 1], searched per weight (exactly, by the midpoints) with signs included.""" |
| if fmt.name == "IQ4_NL": |
| g = grid[:, 0] |
| idx = torch.bucketize(x, (g[1:] + g[:-1]) / 2) |
| return idx, g[idx] |
| k = GROUP[fmt.name] |
| neg = x < 0 |
| if fmt.name in EVEN_SIGNS: |
| g8 = x.abs().unflatten(-1, (-1, 8)) |
| odd = (neg.unflatten(-1, (-1, 8)).sum(-1) % 2) == 1 |
| flip = torch.nn.functional.one_hot(g8.argmin(-1), 8).bool() & odd.unsqueeze(-1) |
| neg = neg ^ flip.flatten(-2) |
| a = x.abs().reshape(-1, k) |
| idx = _nearest_lattice(a, grid) |
| vals = grid[idx].reshape(x.shape) |
| return idx.reshape(*x.shape[:-1], -1), torch.where(neg, -vals, vals) |
|
|
|
|
| def weight_free(fmt: Format, latent: torch.Tensor, d0: torch.Tensor, dlog: torch.Tensor, |
| slat: torch.Tensor | None, grid: torch.Tensor) -> torch.Tensor: |
| """The free-code student's fp32 weight [..., C]: eff · (nearest entry × sign), STE to the latent.""" |
| d = round_fp16(d0 * torch.exp(dlog)) |
| s = round_s(slat, fmt.s_max) if slat is not None else None |
| e = eff_torch(fmt, d, s).unsqueeze(-1) |
| lat = latent.unflatten(-1, (-1, fmt.sub)) |
| with torch.no_grad(): |
| _, vals = free_codes(fmt, (lat / e.detach()).flatten(-2), grid) |
| w = e * vals.unflatten(-1, (-1, fmt.sub)) |
| return (w + (lat - lat.detach())).flatten(-2) |
|
|
|
|
| def encode(fmt: Format, idx: torch.Tensor, vals: torch.Tensor, d: np.ndarray, s: np.ndarray) -> np.ndarray: |
| """Blocks uint8 [..., nblk, nbytes] from codebook indices, signed values (for the signs), d and s.""" |
| lead = vals.shape[:-1] |
| n = int(np.prod(lead)) * (vals.shape[-1] // fmt.block) |
| neg = (vals < 0).reshape(n, fmt.block).cpu().numpy() |
| out = np.zeros((n, fmt.nbytes), np.uint8) |
| set_d(out, d.reshape(-1)) |
| if fmt.name == "IQ4_NL": |
| q = idx.reshape(n, 32).cpu().numpy().astype(np.uint8) |
| out[:, 2:18] = q[:, :16] | (q[:, 16:] << 4) |
| return out.reshape(*lead, -1, fmt.nbytes) |
| gi = idx.reshape(n, -1).cpu().numpy() |
| if fmt.name == "IQ3_S": |
| out[:, 2:66] = (gi & 0xFF).astype(np.uint8) |
| out[:, 66:74] = np.packbits(((gi >> 8) & 1).astype(np.uint8), axis=-1, bitorder="little") |
| out[:, 74:106] = np.packbits(neg.astype(np.uint8), axis=-1, bitorder="little") |
| elif fmt.name == "IQ2_XXS": |
| pat = np.packbits(neg.astype(np.uint8).reshape(n, 32, 8), axis=-1, bitorder="little")[..., 0] |
| if (np.unpackbits(pat[..., None], axis=-1).sum(-1) % 2).any(): |
| raise ValueError("IQ2_XXS sign groups must have an even count of negatives") |
| sidx = (pat & 0x7F).astype(np.uint32).reshape(n, 8, 4) |
| aux1 = (sidx << np.array([0, 7, 14, 21], np.uint32)).sum(-1).astype(np.uint32) |
| aux0 = gi.astype(np.uint8).reshape(n, 8, 4).copy().view("<u4")[..., 0] |
| out[:, 2:66] = np.stack([aux0, aux1], -1).astype("<u4").view(np.uint8).reshape(n, 64) |
| else: |
| raise NotImplementedError(f"free codes for {fmt.name}") |
| set_s(fmt, out, s.reshape(n, -1)) |
| return out.reshape(*lead, -1, fmt.nbytes) |
|
|
|
|
| def grid_from_gguf(fmt: Format, gguf) -> torch.Tensor: |
| """The format's codebook as fp32 [entries, k] (IQ4_NL and IQ4_XS: their 16 values as [16, 1]).""" |
| if fmt.name in ("IQ4_NL", "IQ4_XS"): |
| from lowbit.quant.ternary import IQ4NL_VALUES |
| return torch.tensor(IQ4NL_VALUES, dtype=torch.float32)[:, None] |
| cls = getattr(gguf.quants, fmt.name) |
| cls.init_grid() |
| return torch.from_numpy(np.ascontiguousarray(cls.grid.reshape(cls.grid_shape), dtype=np.float32)) |
|
|
