"""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 # weights per block (one fp16 d) sub: int # weights per sub-scale (= block when there is none) nbytes: int # bytes per block a: float # eff = d·a·(s + b); IQ4_NL has no s (eff = d) b: float s_max: int # s in 0..s_max neutral_s: int # the s that makes eff = d·unit when decoding v (exact) unit: float # a·(neutral_s + b) 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), } # ---------------------------------------------------------------- block fields (numpy, uint8 [n, nbytes]) 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(" 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("> (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(" 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) # ---------------------------------------------------------------- the trainable form (torch) 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) # ---------------------------------------------------------------- free codes (D-58's second mode) # # The codes become trainable too: an fp32 latent per weight, as in `quant.ternary`, starting at llama.cpp's # dequantised weights. The forward divides the latent by its sub-block's eff, finds the nearest codebook entry for # each group of k magnitudes (Euclidean, exact fp32), takes the latent's signs, and multiplies back by eff; the # backward is the plain STE to the latent and the exact gradient to eff (so d and s still train). At the start the # latent is exactly eff·entry·sign, so every group finds its own entry at distance 0 and the student is llama.cpp's # file. Formats whose 8 signs are stored in 7 bits (IQ2_XXS, IQ2_XS, IQ3_XXS: an even count of negatives) flip # the sign of the smallest magnitude in an 8-weight group with an odd count, as llama.cpp's quantiser does. # Grids come from gguf-py at import (`scripts/cb_import.py` writes grids.safetensors), so training needs no gguf-py. GROUP = {"IQ3_S": 4, "IQ2_XXS": 8, "IQ4_NL": 1} EVEN_SIGNS = ("IQ2_XXS",) _CHUNK = 1 << 20 # groups per nearest-entry matmul ([chunk, entries] fp32 scores) 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() # point p = sum r_i·L^(k-1-i) 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] # [n, 16] 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) # ties go to the lower value, as `_nearest`'s argmax return idx, g[idx] k = GROUP[fmt.name] neg = x < 0 if fmt.name in EVEN_SIGNS: # an odd count of negatives in an 8-group: flip the smallest magnitude 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) # [..., nsub, 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)) # gradient to eff (d, s) through the chosen values return (w + (lat - lat.detach())).flatten(-2) # and the plain STE to the latent 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] # [n, 32] 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) # [n, 8] aux0 = gi.astype(np.uint8).reshape(n, 8, 4).copy().view(" 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))