Millaguie
b300ba053d
test-dt3 checks the layout against an independent packer written from the format spec: single-trit position mapping for all 256 (plane, pos) pairs, structural byte-position checks at the region boundaries (79/80, 119/120), exact round-trips with negative scales, byte parity of the in-tree quantizer on already-ternary inputs, and the vec_dot against a hand-made sum over the known trits (catches any path that reads the padding 5th trit of the qh bytes). test-dt3-rust-parity.py packs known trits with ternaria's pack_dt3 and verifies that dequantize_row_dt3 (via test-dt3 --dequant) reproduces d1*t1 + d2*t2 bit-exactly. Also wires DT3 into the test-quantize-fns thresholds (ternary class).
80 lines
2.8 KiB
Python
80 lines
2.8 KiB
Python
#!/usr/bin/env python3
|
|
# Parity check between ternaria's Rust DT3 packer and llama.cpp's C
|
|
# dequantization: blocks packed by pack_dt3 must dequantize (through
|
|
# dequantize_row_dt3, exposed by `test-dt3 --dequant`) to exactly
|
|
# d1*t1 + d2*t2 computed independently in numpy from the same trits and
|
|
# fp16-rounded scales.
|
|
#
|
|
# Must run inside the ternaria environment, e.g.:
|
|
# cd /path/to/ternaria && uv run python /path/to/llama.cpp/tests/test-dt3-rust-parity.py \
|
|
# /path/to/llama.cpp/build/bin/test-dt3
|
|
|
|
import subprocess
|
|
import sys
|
|
import tempfile
|
|
from pathlib import Path
|
|
|
|
import numpy as np
|
|
|
|
from ternaria._core import pack_dt3
|
|
|
|
QK_DT3 = 128
|
|
BLOCK_BYTES = 56
|
|
|
|
|
|
def main() -> int:
|
|
if len(sys.argv) != 2:
|
|
print(f"usage: {sys.argv[0]} /path/to/test-dt3", file=sys.stderr)
|
|
return 1
|
|
test_bin = Path(sys.argv[1])
|
|
|
|
rng = np.random.default_rng(20260810)
|
|
|
|
rows, cols = 8, 1024
|
|
t1 = rng.integers(-1, 2, size=(rows, cols)).astype(np.int8)
|
|
t2 = rng.integers(-1, 2, size=(rows, cols)).astype(np.int8)
|
|
s1 = rng.normal(size=(rows, cols // QK_DT3)).astype(np.float32)
|
|
s2 = (0.25 * rng.normal(size=(rows, cols // QK_DT3))).astype(np.float32)
|
|
|
|
# edge cases: an all-zero block, a block with negative scales, and a
|
|
# block that is non-trivial only in the qh region (elements 120..127)
|
|
t1[0, :QK_DT3] = 0
|
|
t2[0, :QK_DT3] = 0
|
|
s1[0, 0] = 0.0
|
|
s2[0, 0] = 0.0
|
|
s1[0, 1] = -abs(s1[0, 1])
|
|
s2[0, 1] = -abs(s2[0, 1])
|
|
t1[1, :QK_DT3] = 0
|
|
t2[1, :QK_DT3] = 0
|
|
t1[1, 120:128] = [-1, 1, 0, -1, 1, -1, 0, 1]
|
|
t2[1, 120:128] = [1, -1, 1, 0, 0, 1, -1, -1]
|
|
|
|
raw = np.asarray(pack_dt3(t1, s1, t2, s2), dtype=np.uint8)
|
|
assert raw.size == rows * (cols // QK_DT3) * BLOCK_BYTES, raw.size
|
|
|
|
with tempfile.TemporaryDirectory() as tmp:
|
|
raw_path = Path(tmp) / "dt3.bin"
|
|
out_path = Path(tmp) / "out.f32"
|
|
raw_path.write_bytes(raw.tobytes())
|
|
subprocess.run([str(test_bin), "--dequant", str(raw_path), str(out_path)], check=True)
|
|
got = np.fromfile(out_path, dtype=np.float32).reshape(rows, cols)
|
|
|
|
# what the packed bytes mean: fp16-rounded scales times the trits
|
|
d1 = s1.astype(np.float16).astype(np.float32).repeat(QK_DT3, axis=1)
|
|
d2 = s2.astype(np.float16).astype(np.float32).repeat(QK_DT3, axis=1)
|
|
expected = d1 * t1.astype(np.float32) + d2 * t2.astype(np.float32)
|
|
|
|
if not np.array_equal(got, expected):
|
|
bad = np.nonzero(got != expected)
|
|
print(f"FAILED: {len(bad[0])} of {got.size} elements differ", file=sys.stderr)
|
|
r, c = bad[0][0], bad[1][0]
|
|
print(f"first mismatch at ({r}, {c}): got {got[r, c]}, expected {expected[r, c]}", file=sys.stderr)
|
|
return 1
|
|
|
|
print(f"ok: {got.size} weights bit-exact between Rust pack_dt3 and C dequantize_row_dt3")
|
|
return 0
|
|
|
|
|
|
if __name__ == "__main__":
|
|
sys.exit(main())
|