ZacharyZcR
e48346162a
#453 settled the scheme; this produces the actual bytes. iq3_pack.encode/decode implement the container layout so the converter, the engine and the decode kernels can all be written against one spec: 98 bytes per 256 weights = 3.0625 bpw [ 0..63] uint8 grid index per 4-dim magnitude block [64..95] uint32 x8 — four 7-bit sign words + 4-bit sub-scale per 32 [96..97] fp16 super-scale Signs use the published odd-parity trick: 7 of every 8 are stored and the 8th is derived, so the encoder flips the smallest-magnitude weight of any block whose true signs would violate parity — the cost the #453 ablation priced in, now actually paid. Sub-scale is searched over all 16 codes against the stored (fp16-rounded) super-scale, so encode-time and decode-time arithmetic agree. GLM-5.2 routed experts under this container: 372.7 -> 281.2 GB (-24.6%), and a 176 GB VRAM tier holds ~12,180 experts instead of ~9,190 (+33%). tests/test_iq3_pack.py: byte budget, deterministic encode, decode checked value-by-value against an independent loop-based reader written straight from the layout, sign-parity closure, and reconstruction quality in the band the chosen scheme measured (rel-RMSE 0.195 vs the torch model's 0.195).
89 lines
3.6 KiB
Python
89 lines
3.6 KiB
Python
"""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
|
|
|
|
import numpy as np
|
|
|
|
sys.path.insert(0, os.path.join(os.path.dirname(__file__), "..", "tools"))
|
|
import iq3_pack as P # noqa: E402
|
|
|
|
|
|
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
|
|
|
|
|
|
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()
|