Merge #458: fmt=5 index codec tools — deployable bytes for the E8/IQ3 container
Tools-only step of the #452 plan: iq3_pack.py index codec + IQ3_XXS grid + quant_ablation hook, with round-trip tests (6/6 locally, CI 8/8). No engine changes. Co-authored-by: ZacharyZcR
This commit is contained in:
Vincenzo Fornaro
@@ -0,0 +1,96 @@
|
||||
"""fmt=5 codec oracle (#452 ladder step 2).
|
||||
|
||||
Checks the properties the container, the converter and the decode kernels all
|
||||
depend on: exact byte budget, deterministic encode, decode agreeing with a
|
||||
straight-from-the-spec reader, sign parity closure, and reconstruction quality
|
||||
matching the ablation that chose this scheme (#453).
|
||||
"""
|
||||
import os
|
||||
import sys
|
||||
import unittest
|
||||
|
||||
# The runtime path is dependency-free by design and CI keeps it that way, so the
|
||||
# offline-tooling tests skip rather than fail where numpy is absent.
|
||||
np = None
|
||||
P = None
|
||||
try:
|
||||
import numpy as np
|
||||
|
||||
sys.path.insert(0, os.path.join(os.path.dirname(__file__), "..", "tools"))
|
||||
import iq3_pack as P
|
||||
except ImportError: # pragma: no cover - exercised only on dependency-free CI
|
||||
pass
|
||||
|
||||
|
||||
def ref_decode(packed, K):
|
||||
"""Independent reader written straight from the layout comment — deliberately
|
||||
naive and loop-based, so a shared bug in the vectorized path shows up."""
|
||||
g = P.grid()
|
||||
nsb = K // P.QK
|
||||
rows = packed.reshape(-1, nsb * P.BLOCK_BYTES)
|
||||
out = np.zeros((len(rows), K), dtype=np.float32)
|
||||
for r in range(len(rows)):
|
||||
for sb in range(nsb):
|
||||
base = sb * P.BLOCK_BYTES
|
||||
d = float(rows[r, base + 96:base + 98].copy().view(np.float16)[0])
|
||||
for ib in range(P.QK // P.SUB):
|
||||
word = int(np.ascontiguousarray(
|
||||
rows[r, base + 64 + ib * 4:base + 64 + ib * 4 + 4]).view(np.uint32)[0])
|
||||
db = d * (0.5 + ((word >> 28) & 0xF)) * 0.5
|
||||
for l in range(4):
|
||||
seven = (word >> (7 * l)) & 0x7F
|
||||
bits = [(seven >> j) & 1 for j in range(7)]
|
||||
bits.append(sum(bits) & 1) # odd parity closes the 8th
|
||||
idx = int(rows[r, base + ib * 8 + l * 2 + 0])
|
||||
idx2 = int(rows[r, base + ib * 8 + l * 2 + 1])
|
||||
mags = list(g[idx]) + list(g[idx2])
|
||||
for j in range(8):
|
||||
pos = sb * P.QK + ib * P.SUB + l * 8 + j
|
||||
out[r, pos] = mags[j] * db * (-1.0 if bits[j] else 1.0)
|
||||
return out
|
||||
|
||||
|
||||
@unittest.skipIf(P is None, "numpy not available (offline-tooling test)")
|
||||
class TestIq3Pack(unittest.TestCase):
|
||||
def setUp(self):
|
||||
np.random.seed(4242)
|
||||
self.x = (np.random.randn(6, 1024) * 0.05).astype(np.float32)
|
||||
|
||||
def test_byte_budget(self):
|
||||
self.assertEqual(P.BLOCK_BYTES, 98)
|
||||
self.assertAlmostEqual(P.bpw(), 3.0625, places=6)
|
||||
packed = P.encode(self.x)
|
||||
self.assertEqual(packed.shape, (6, 1024 // P.QK * 98))
|
||||
self.assertEqual(packed.dtype, np.uint8)
|
||||
|
||||
def test_encode_is_deterministic(self):
|
||||
self.assertTrue(np.array_equal(P.encode(self.x), P.encode(self.x)))
|
||||
|
||||
def test_decode_matches_spec_reader(self):
|
||||
packed = P.encode(self.x)
|
||||
fast = P.decode(packed, 1024)
|
||||
slow = ref_decode(packed, 1024)
|
||||
self.assertTrue(np.allclose(fast, slow, rtol=1e-6, atol=1e-8),
|
||||
f"max |Δ| = {np.abs(fast - slow).max()}")
|
||||
|
||||
def test_sign_parity_closes(self):
|
||||
"""Every 8-weight block must have an even number of negatives — that is
|
||||
what lets the 8th sign be derived instead of stored."""
|
||||
y = P.decode(P.encode(self.x), 1024)
|
||||
neg = (y < 0).reshape(-1, 8).sum(-1)
|
||||
self.assertTrue(np.all(neg % 2 == 0), "a block stored odd negatives")
|
||||
|
||||
def test_reconstruction_quality(self):
|
||||
y = P.decode(P.encode(self.x), 1024)
|
||||
rel = np.sqrt(((y - self.x) ** 2).mean()) / np.sqrt((self.x ** 2).mean())
|
||||
# the torch model that won the #453 A/B measures ~0.195 on this input class
|
||||
self.assertLess(rel, 0.25, f"rel-RMSE {rel:.4f} — worse than the chosen scheme")
|
||||
self.assertGreater(rel, 0.05, f"rel-RMSE {rel:.4f} — implausibly good, check the test")
|
||||
|
||||
def test_shape_guard(self):
|
||||
with self.assertRaises(ValueError):
|
||||
P.encode(np.zeros((2, 300), dtype=np.float32))
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
@@ -0,0 +1,170 @@
|
||||
#!/usr/bin/env python3
|
||||
"""fmt=5 (E8/IQ3 grouped container) index codec — #452 ladder step 2.
|
||||
|
||||
The ablation (#453) proved the SCHEME: an IQ3_XXS-style codebook plus rotation
|
||||
matches our simulated E8 ball (51.5% vs 51.5% on OLMoE). That code quantizes to
|
||||
lattice points and keeps floats. This module produces the DEPLOYABLE bytes and
|
||||
reads them back, so the container, the converter and the decode kernels all
|
||||
agree on one layout.
|
||||
|
||||
Layout — one 256-weight super-block, 98 bytes, 3.0625 bpw:
|
||||
|
||||
[0 .. 63] uint8 grid index per 4-dim magnitude block (64 blocks)
|
||||
[64 .. 95] uint32 x8, one per 32-weight sub-block:
|
||||
bits 0..20 three 7-bit sign words (8 weights each,
|
||||
bit i set => weight i negative; the 8th sign
|
||||
is implied by odd parity)
|
||||
bits 21..27 the fourth 7-bit sign word
|
||||
bits 28..31 4-bit sub-scale code
|
||||
[96 .. 97] fp16 super-scale d
|
||||
|
||||
value(w) = d * (0.5 + code) * 0.5 * grid[idx][j] * 0.5 * sign
|
||||
|
||||
The last 0.5 is the half-unit convention of the published grid (magnitudes are
|
||||
stored doubled: 4,12,...,62 mean 2.0,6.0,...,31.0).
|
||||
|
||||
Odd-parity signs: llama.cpp stores 7 of every 8 signs and derives the 8th so the
|
||||
product of the eight is +1. The encoder therefore flips the smallest-magnitude
|
||||
weight of any block whose true signs violate that — the same cost the ablation
|
||||
priced in, now applied for real.
|
||||
"""
|
||||
import json
|
||||
import os
|
||||
import numpy as np
|
||||
|
||||
QK = 256 # weights per super-block
|
||||
SUB = 32 # weights per sub-block (one uint32 of signs+scale)
|
||||
BLOCK_BYTES = QK // 4 + (QK // SUB) * 4 + 2 # 64 + 32 + 2 = 98
|
||||
|
||||
_GRID = None
|
||||
|
||||
|
||||
def grid():
|
||||
"""[256,4] float32 magnitudes in weight units (published table is doubled)."""
|
||||
global _GRID
|
||||
if _GRID is None:
|
||||
path = os.path.join(os.path.dirname(__file__), "iq3xxs_grid.json")
|
||||
_GRID = np.asarray(json.load(open(path)), dtype=np.float32) * 0.5
|
||||
return _GRID
|
||||
|
||||
|
||||
def _nearest(mag4):
|
||||
"""[N,4] magnitudes -> [N] grid indices, argmin ||m-g||^2 without cdist."""
|
||||
g = grid()
|
||||
g2 = (g * g).sum(1)
|
||||
out = np.empty(len(mag4), dtype=np.uint8)
|
||||
for i in range(0, len(mag4), 1 << 16): # bounded working set
|
||||
c = mag4[i:i + (1 << 16)]
|
||||
out[i:i + len(c)] = np.argmin(g2 - 2.0 * (c @ g.T), axis=1).astype(np.uint8)
|
||||
return out
|
||||
|
||||
|
||||
def encode(x):
|
||||
"""float32 [..., K] (K % 256 == 0) -> packed uint8 [..., K//256 * 98]."""
|
||||
x = np.ascontiguousarray(x, dtype=np.float32)
|
||||
K = x.shape[-1]
|
||||
if K % QK:
|
||||
raise ValueError(f"fmt=5 needs K % {QK} == 0, got {K}")
|
||||
rows = x.reshape(-1, K)
|
||||
nsb = K // QK
|
||||
out = np.empty((len(rows), nsb * BLOCK_BYTES), dtype=np.uint8)
|
||||
|
||||
for sb in range(nsb):
|
||||
blk = rows[:, sb * QK:(sb + 1) * QK] # [R,256]
|
||||
sign = np.where(blk < 0, -1.0, 1.0).astype(np.float32)
|
||||
mag = np.abs(blk)
|
||||
|
||||
# parity fix: flip the smallest magnitude of every 8 whose product is -1
|
||||
s8 = sign.reshape(len(rows), QK // 8, 8)
|
||||
m8 = mag.reshape(len(rows), QK // 8, 8)
|
||||
viol = s8.prod(-1) < 0 # [R,32]
|
||||
amin = m8.argmin(-1)
|
||||
r, b = np.nonzero(viol)
|
||||
s8[r, b, amin[r, b]] *= -1.0
|
||||
sign = s8.reshape(len(rows), QK)
|
||||
|
||||
base = sb * BLOCK_BYTES
|
||||
# super-scale: RMS anchor, same statistic the ablation searches around
|
||||
d = np.sqrt((mag * mag).mean(-1, keepdims=True)) / 20.0 + 1e-12
|
||||
out[:, base + 96:base + 98] = d.astype(np.float16).view(np.uint8)
|
||||
d = d.astype(np.float16).astype(np.float32) # encode what we store
|
||||
|
||||
g = grid()
|
||||
for ib in range(QK // SUB):
|
||||
m = mag[:, ib * SUB:(ib + 1) * SUB] # [R,32]
|
||||
best_err = None
|
||||
best = None
|
||||
for code in range(16):
|
||||
db = d * (0.5 + code) * 0.5
|
||||
q = (m / np.maximum(db, 1e-20)).reshape(-1, 4)
|
||||
idx = _nearest(q)
|
||||
rec = g[idx].reshape(len(rows), SUB) * db
|
||||
err = ((rec - m) ** 2).sum(-1, keepdims=True)
|
||||
if best_err is None:
|
||||
best_err, best = err, (idx.reshape(len(rows), SUB // 4), code)
|
||||
else:
|
||||
take = (err < best_err)[:, 0]
|
||||
if take.any():
|
||||
keep_idx, keep_code = best
|
||||
ni = idx.reshape(len(rows), SUB // 4)
|
||||
keep_idx = np.where(take[:, None], ni, keep_idx)
|
||||
# per-row code: store alongside, resolved below
|
||||
keep_code = np.where(take, code, keep_code) if isinstance(
|
||||
keep_code, np.ndarray) else np.where(
|
||||
take, code, np.full(len(rows), keep_code))
|
||||
best = (keep_idx, keep_code)
|
||||
best_err = np.where(take[:, None], err, best_err)
|
||||
bidx, bcode = best
|
||||
if not isinstance(bcode, np.ndarray):
|
||||
bcode = np.full(len(rows), bcode)
|
||||
out[:, base + ib * 8:base + (ib + 1) * 8] = bidx.astype(np.uint8)
|
||||
|
||||
# signs: four 7-bit words for this sub-block + the 4-bit code
|
||||
s = sign[:, ib * SUB:(ib + 1) * SUB].reshape(len(rows), 4, 8)
|
||||
neg = (s < 0).astype(np.uint32)
|
||||
word = np.zeros(len(rows), dtype=np.uint32)
|
||||
for l in range(4):
|
||||
seven = np.zeros(len(rows), dtype=np.uint32)
|
||||
for j in range(7):
|
||||
seven |= neg[:, l, j] << j
|
||||
word |= seven << (7 * l)
|
||||
word |= (bcode.astype(np.uint32) & 0xF) << 28
|
||||
off = base + QK // 4 + ib * 4
|
||||
out[:, off:off + 4] = word.view(np.uint8).reshape(len(rows), 4) if False else \
|
||||
np.ascontiguousarray(word).view(np.uint8).reshape(len(rows), 4)
|
||||
return out.reshape(*x.shape[:-1], nsb * BLOCK_BYTES)
|
||||
|
||||
|
||||
def decode(packed, K):
|
||||
"""packed uint8 [..., K//256*98] -> float32 [..., K]. The kernels' reference."""
|
||||
packed = np.ascontiguousarray(packed, dtype=np.uint8)
|
||||
nsb = K // QK
|
||||
rows = packed.reshape(-1, nsb * BLOCK_BYTES)
|
||||
out = np.empty((len(rows), K), dtype=np.float32)
|
||||
g = grid()
|
||||
for sb in range(nsb):
|
||||
base = sb * BLOCK_BYTES
|
||||
d = rows[:, base + 96:base + 98].copy().view(np.float16).astype(np.float32)
|
||||
for ib in range(QK // SUB):
|
||||
idx = rows[:, base + ib * 8:base + (ib + 1) * 8] # [R,8]
|
||||
off = base + QK // 4 + ib * 4
|
||||
word = np.ascontiguousarray(rows[:, off:off + 4]).view(np.uint32).reshape(-1)
|
||||
code = (word >> 28) & 0xF
|
||||
db = d[:, 0] * (0.5 + code) * 0.5 # [R]
|
||||
mag = g[idx].reshape(len(rows), SUB) # [R,32]
|
||||
sgn = np.ones((len(rows), 4, 8), dtype=np.float32)
|
||||
for l in range(4):
|
||||
seven = (word >> (7 * l)) & 0x7F
|
||||
par = 0
|
||||
for j in range(7):
|
||||
bit = (seven >> j) & 1
|
||||
sgn[:, l, j] = np.where(bit == 1, -1.0, 1.0)
|
||||
par ^= bit
|
||||
sgn[:, l, 7] = np.where(par == 1, -1.0, 1.0) # odd parity closes the block
|
||||
out[:, sb * QK + ib * SUB:sb * QK + (ib + 1) * SUB] = \
|
||||
mag * sgn.reshape(len(rows), SUB) * db[:, None]
|
||||
return out.reshape(*packed.shape[:-1], K)
|
||||
|
||||
|
||||
def bpw():
|
||||
return BLOCK_BYTES * 8 / QK
|
||||
Reference in New Issue
Block a user