code/tests/test_ternary.py

"""T2.1 acceptance: fast path = reference, STE identity, Q-A zero fraction, packing round-trips,
idempotence. CPU only."""
import pytest
import torch
from hypothesis import given, settings
from hypothesis import strategies as st

from lowbit.quant.ternary import (
    QuantSpec,
    canonical_scale,
    dequant_packed,
    export_packed,
    fake_quant,
    init_scale,
    pack_2bit,
    pack_trits,
    reference_codes_and_scale,
    ternary_codes,
    unpack_2bit,
    unpack_trits,
    zero_fraction,
)

BF16_ULP = 2.0**-7


def gen(shape, seed=0, std=0.02):
    return torch.randn(*shape, generator=torch.Generator().manual_seed(seed)) * std


@pytest.mark.parametrize("method", ["absmean", "twn"])
@pytest.mark.parametrize("group", [64, 128])
def test_fast_matches_reference(method, group):
    spec = QuantSpec(method=method, group_size=group)
    w = gen((12, 384), seed=group)
    ref_q, ref_s = reference_codes_and_scale(w, spec)
    s = init_scale(w, spec)
    q = ternary_codes(w, s, spec)
    # scales: fp32 vs fp64 summation order may differ by <1 bf16 ulp; everything else is exact
    assert torch.allclose(s, ref_s, rtol=BF16_ULP, atol=0)
    same = (s == ref_s).repeat_interleave(group, -1)
    assert same.float().mean() > 0.99
    assert torch.equal(q[same], ref_q[same])


@pytest.mark.parametrize("method", ["absmean", "twn"])
def test_stacked_experts_match_per_expert(method):
    spec = QuantSpec(method=method)
    w = gen((5, 48, 384), seed=3)  # [E, out, in]
    stacked = fake_quant(w, spec)
    per = torch.stack([fake_quant(w[e], spec) for e in range(5)])
    assert torch.equal(stacked, per)


def test_forward_is_exact_and_bf16_safe():
    spec = QuantSpec()
    w = gen((32, 1024), seed=1)
    w_hat = fake_quant(w, spec)
    s = init_scale(w, spec)
    q = ternary_codes(w, s, spec).float()
    expect = (q.unflatten(-1, (-1, 128)) * s.unsqueeze(-1)).flatten(-2)
    assert torch.equal(w_hat, expect)
    # the bf16 cast used by the matmul loses nothing: values are bf16(s) * {-1,0,1}
    assert torch.equal(w_hat.to(torch.bfloat16).float(), w_hat)
    assert set(torch.unique(q).tolist()) <= {-1.0, 0.0, 1.0}


@pytest.mark.parametrize("method", ["absmean", "twn"])
def test_plain_ste_gradient_is_identity(method):
    w = gen((16, 256), seed=2).requires_grad_()
    g = gen((16, 256), seed=9, std=1.0)
    fake_quant(w, QuantSpec(method=method)).backward(g)
    assert torch.equal(w.grad, g)


def test_clipped_ste_masks_large_ratios():
    spec = QuantSpec(ste="clipped")
    w = gen((16, 256), seed=4)
    w[0, :5] = 1.0  # far outside 1.5 s
    w.requires_grad_()
    g = torch.ones_like(w)
    fake_quant(w, spec).backward(g)
    s = init_scale(w.detach(), spec).repeat_interleave(128, -1)
    keep = (w.detach().abs() / s) <= 1.5
    assert torch.equal(w.grad, keep.float())
    assert not keep[0, :5].any()


def test_lsq_scale_gets_gradient_and_init_matches_absmean():
    w = gen((8, 256), seed=5)
    s0 = init_scale(w, QuantSpec(method="lsq"))
    assert torch.equal(s0, init_scale(w, QuantSpec(method="absmean")))
    s = s0.clone().requires_grad_()
    wl = w.clone().requires_grad_()
    out = fake_quant(wl, QuantSpec(method="lsq"), scale=s)
    assert torch.equal(out.detach(), fake_quant(w, QuantSpec()))
    out.pow(2).sum().backward()
    assert s.grad is not None and s.grad.abs().sum() > 0
    assert torch.equal(wl.grad, 2 * out.detach())  # identity STE to the latent


def test_lsq_scale_gradient_matches_finite_difference_outside_rounding_edges():
    # away from rounding boundaries q is locally constant, so d(sum(s*q))/ds = sum(q) per group,
    # and LSQ's (q - w/s) term only differs inside; check the outside branch exactly
    w = torch.tensor([[3.0, -3.0, 2.0, 2.5] * 32])
    s = torch.ones(1, 1).requires_grad_()
    fake_quant(w, QuantSpec(method="lsq"), scale=s).sum().backward()
    q_sum = ternary_codes(w, s.detach(), QuantSpec(method="lsq")).float().sum()
    assert torch.allclose(s.grad, (q_sum / 128**0.5).reshape(1, 1))


def test_q_a_zero_fraction_on_gaussian():
    w = torch.randn(2048, 1024, generator=torch.Generator().manual_seed(0))
    s = init_scale(w, QuantSpec())
    zf = zero_fraction(ternary_codes(w, s, QuantSpec()))
    assert abs(zf - 0.31) <= 0.01, zf


@pytest.mark.parametrize("method", ["absmean", "twn"])
@pytest.mark.parametrize("group", [64, 128])
def test_quantising_ternary_is_idempotent(method, group):
    """Codes are idempotent under every quantiser. Values are idempotent under Q-B (its scale is the mean
    of the nonzeros) but not Q-A: the absmean of a ternary group is s * (nonzero fraction), so the codes
    survive and the scale shrinks (DECISIONS.md D-09)."""
    spec = QuantSpec(method=method, group_size=group)
    w = gen((24, 384), seed=6)
    once = fake_quant(w, spec)
    s1, s2 = init_scale(w, spec), init_scale(once, spec)
    q1, q2 = ternary_codes(w, s1, spec), ternary_codes(once, s2, spec)
    assert torch.equal(q1, q2)
    if method == "twn":
        assert torch.equal(fake_quant(once, spec), once)
    else:
        nz = (q1 != 0).float().unflatten(-1, (-1, group)).mean(-1)
        assert torch.allclose(s2, canonical_scale(s1 * nz), rtol=BF16_ULP, atol=0)


def test_lsq_requantising_with_its_own_scale_is_idempotent():
    spec = QuantSpec(method="lsq")
    w = gen((24, 384), seed=10)
    s = init_scale(w, spec)
    once = fake_quant(w, spec, scale=s)
    assert torch.equal(fake_quant(once, spec, scale=s), once)


def _llama_quantize_row_q2_0_ref(x: list[float]) -> tuple[float, bytes]:
    """Port of mainline llama.cpp quantize_row_q2_0_ref for one 64-element block (ggml-quants.c):
    d = max|x|, q = clamp(roundf(x / d) + 1, 0, 3), element j at bits 2*(j%4) of byte j/4."""
    import math
    d = max(abs(v) for v in x)
    inv = 1.0 / d if d > 0 else 0.0
    qs = bytearray(16)
    for j, v in enumerate(x):
        r = v * inv
        q = int(math.floor(abs(r) + 0.5)) * (1 if r >= 0 else -1) + 1  # roundf: half away from zero
        qs[j // 4] |= max(0, min(3, q)) << ((j % 4) * 2)
    return d, bytes(qs)


def test_pack_2bit_matches_llama_q2_0_layout():
    # a g128 ternary group exports as two g64 Q2_0 blocks sharing its scale (T2.7)
    spec = QuantSpec()
    w = gen((4, 256), seed=11)
    p = export_packed(w, spec)
    deq = dequant_packed(p, torch.float32)
    for r in range(4):
        for blk in range(4):  # 256 / 64
            vals = deq[r, blk * 64:(blk + 1) * 64].tolist()
            d, qs = _llama_quantize_row_q2_0_ref(vals)
            ours = bytes(p["codes"][r, blk * 16:(blk + 1) * 16].tolist())
            assert qs == ours
            assert d == float(p["scales"][r, blk // 2]) or not any(vals)


def test_scales_are_exact_in_fp16():
    w = gen((64, 1024), seed=7, std=0.03)
    s = init_scale(w, QuantSpec())
    assert torch.equal(s.to(torch.float16).float(), s)
    assert torch.equal(canonical_scale(s), s)


@settings(max_examples=50, deadline=None)
@given(st.lists(st.integers(-1, 1), min_size=1, max_size=300))
def test_trit_pack_round_trip(vals):
    q = torch.tensor(vals, dtype=torch.int8)
    p = pack_trits(q)
    assert p.numel() == -(-len(vals) // 5)
    assert torch.equal(unpack_trits(p, len(vals)), q)


@settings(max_examples=50, deadline=None)
@given(st.integers(1, 64).flatmap(lambda n: st.lists(st.integers(-1, 1), min_size=4 * n, max_size=4 * n)))
def test_2bit_pack_round_trip(vals):
    q = torch.tensor(vals, dtype=torch.int8)
    p = pack_2bit(q)
    assert p.dtype == torch.uint8 and p.numel() == len(vals) // 4
    assert torch.equal(unpack_2bit(p), q)


def test_packed_export_reproduces_fake_quant_exactly():
    for spec in [QuantSpec(), QuantSpec(method="twn", group_size=64)]:
        w = gen((3, 64, 384), seed=8)
        expect = fake_quant(w, spec).to(torch.bfloat16)
        assert torch.equal(dequant_packed(export_packed(w, spec)), expect)


def test_marco_expert_shapes():
    # gate/up: [232, 384, 1024] -> 8 groups/row; down: [232, 1024, 384] -> 3 groups/row
    spec = QuantSpec()
    assert init_scale(torch.zeros(2, 384, 1024) + 0.01, spec).shape == (2, 384, 8)
    assert init_scale(torch.zeros(2, 1024, 384) + 0.01, spec).shape == (2, 1024, 3)
    with pytest.raises(ValueError):
        init_scale(torch.zeros(4, 100), spec)