"""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)