| """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) |
| |
| 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) |
| 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) |
| |
| 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 |
| 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()) |
|
|
|
|
| def test_lsq_scale_gradient_matches_finite_difference_outside_rounding_edges(): |
| |
| |
| 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 |
| qs[j // 4] |= max(0, min(3, q)) << ((j % 4) * 2) |
| return d, bytes(qs) |
|
|
|
|
| def test_pack_2bit_matches_llama_q2_0_layout(): |
| |
| 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): |
| 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(): |
| |
| 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) |
|
|
