| """D-58: codebook formats against llama.cpp's own files (gguf-py's dequantisers are the reference).""" |
| from pathlib import Path |
|
|
| import numpy as np |
| import pytest |
| import torch |
|
|
| from lowbit.export.gguf import gguf_module |
| from lowbit.quant import codebook as cb |
|
|
| PTQ = Path("/path/to/data/ptq/gguf") |
| LLAMA = "/path/to/llama.cpp" |
| SOURCES = { |
| "IQ2_XXS": ("teacher-iq2_xxs-imat.gguf", "blk.5.ffn_up_exps.weight"), |
| "IQ2_XS": ("teacher-iq2_xs-imat.gguf", "blk.5.ffn_up_exps.weight"), |
| "IQ2_S": ("teacher-iq2_s-iq4nldown-imat.gguf", "blk.5.ffn_gate_exps.weight"), |
| "IQ3_XXS": ("teacher-iq3_xxs-imat.gguf", "blk.5.ffn_gate_exps.weight"), |
| "IQ3_S": ("teacher-iq3_s-iq4nldown-imat.gguf", "blk.5.ffn_gate_exps.weight"), |
| "IQ4_XS": ("teacher-iq4_xs-imat.gguf", "blk.5.ffn_up_exps.weight"), |
| "IQ4_NL": ("teacher-iq3_s-iq4nldown-imat.gguf", "blk.5.ffn_down_exps.weight"), |
| } |
| N = 4096 |
|
|
|
|
| @pytest.fixture(scope="module") |
| def gguf(): |
| return gguf_module(LLAMA) |
|
|
|
|
| def _blocks(gguf, name, n=N): |
| path, tensor = SOURCES[name] |
| if not (PTQ / path).exists(): |
| pytest.skip(f"{path} not on this machine") |
| t = next(t for t in gguf.GGUFReader(str(PTQ / path)).tensors if t.name == tensor) |
| fmt = cb.FORMATS[name] |
| assert t.tensor_type == getattr(gguf.GGMLQuantizationType, name) |
| return fmt, np.asarray(t.data).reshape(-1, fmt.nbytes)[:n].copy() |
|
|
|
|
| @pytest.mark.parametrize("name", list(SOURCES)) |
| def test_fields_round_trip(gguf, name): |
| fmt, b = _blocks(gguf, name) |
| out = cb.write_blocks(fmt, b, cb.get_d(b), cb.get_s(fmt, b)) |
| assert np.array_equal(out, b) |
|
|
|
|
| @pytest.mark.parametrize("name", list(SOURCES)) |
| def test_decode_matches_ggml(gguf, name): |
| fmt, b = _blocks(gguf, name) |
| ref = gguf.quants.dequantize(b, getattr(gguf.GGMLQuantizationType, name)).reshape(len(b), fmt.block) |
| v = cb.values(fmt, b, gguf) |
| e = cb.eff_np(fmt, cb.get_d(b), cb.get_s(fmt, b)) |
| got = (v.astype(np.float32).reshape(len(b), -1, fmt.sub) * e[..., None]).reshape(len(b), -1) |
| assert np.array_equal(got, ref) |
| |
| d0 = torch.from_numpy(cb.get_d(b).astype(np.float32)) |
| s0 = torch.from_numpy(cb.get_s(fmt, b).astype(np.float32)) |
| w = cb.weight(fmt, torch.from_numpy(v.reshape(-1)), d0, torch.zeros_like(d0), |
| s0.reshape(-1) if fmt.name != "IQ4_NL" else None) |
| assert torch.equal(w, torch.from_numpy(ref.reshape(-1))) |
|
|
|
|
| @pytest.mark.parametrize("name", list(SOURCES)) |
| def test_new_scales_dequantise_as_trained(gguf, name): |
| fmt, b = _blocks(gguf, name) |
| v = torch.from_numpy(cb.values(fmt, b, gguf).reshape(-1)) |
| d0 = torch.from_numpy(cb.get_d(b).astype(np.float32)) |
| g = torch.Generator().manual_seed(0) |
| dlog = 0.05 * torch.randn(d0.shape, generator=g) |
| slat = None |
| if fmt.name != "IQ4_NL": |
| s0 = torch.from_numpy(cb.get_s(fmt, b).astype(np.float32)).reshape(-1) |
| slat = s0 + 1.5 * torch.randn(s0.shape, generator=g) |
| w = cb.weight(fmt, v, d0, dlog, slat) |
| d, s = cb.trained_fields(fmt, d0, dlog, slat) |
| nb = cb.write_blocks(fmt, b, d, s.reshape(len(b), -1)) |
| ref = gguf.quants.dequantize(nb, getattr(gguf.GGMLQuantizationType, name)).reshape(-1) |
| assert torch.equal(w, torch.from_numpy(ref)) |
|
|
|
|
| def test_ste_reaches_scales(): |
| fmt = cb.FORMATS["IQ3_S"] |
| v = torch.randint(-15, 16, (2, 256), dtype=torch.int8) |
| d0 = torch.full((2, 1), 0.01) |
| dlog = torch.zeros(2, 1, requires_grad=True) |
| slat = torch.full((2, 8), 5.0, requires_grad=True) |
| cb.weight(fmt, v, d0, dlog, slat).sum().backward() |
| assert dlog.grad.abs().sum() > 0 and slat.grad.abs().sum() > 0 |
|
|
|
|
| def _fake_init(gguf, tmp_path, cfg): |
| """A cb_import-style directory for a tiny model, from real llama.cpp blocks (IQ3_S gate/up, IQ4_NL down).""" |
| from safetensors.numpy import save_file |
| import json |
| E, I, H = cfg.num_experts, cfg.moe_intermediate_size, cfg.hidden_size |
| f3, b3 = _blocks(gguf, "IQ3_S") |
| fn, bn = _blocks(gguf, "IQ4_NL", 4 * N) |
| need3, needn = E * I * (H // 256), E * H * (I // 32) |
| assert len(b3) >= 2 * need3 and len(bn) >= needn |
| t, formats = {}, {} |
| for i in range(cfg.num_hidden_layers): |
| for k, (fmt, src, rows, cols) in {"gate": (f3, b3[:need3], I, H), "up": (f3, b3[need3:2 * need3], I, H), |
| "down": (fn, bn[:needn], H, I)}.items(): |
| key = f"layers.{i}.{k}" |
| formats[key] = fmt.name |
| t[f"{key}.blocks"] = src.reshape(E, rows, cols // fmt.block, fmt.nbytes).copy() |
| t[f"{key}.v"] = cb.values(fmt, src, gguf).reshape(E, rows, cols) |
| t[f"{key}.d"] = cb.get_d(src).reshape(E, rows, -1) |
| t[f"{key}.s"] = cb.get_s(fmt, src).reshape(E, rows, -1) |
| save_file(t, str(tmp_path / "experts.safetensors")) |
| (tmp_path / "manifest.json").write_text(json.dumps({"formats": formats})) |
| return t |
|
|
|
|
| def test_student_round_trip(gguf, tmp_path): |
| """The codebook student's weights are the source blocks' dequantisation; after a training-like change the packed |
| export dequantises to the student's weights; gate/up keep llama.cpp's order (gate first).""" |
| from test_student import tiny |
| from lowbit.model import student as st |
| from lowbit.quant.ternary import QuantSpec |
| from lowbit.export.gguf import packed_blocks, student_tensors |
| base = tiny(hidden=256) |
| (tmp_path / "init").mkdir() |
| t = _fake_init(gguf, tmp_path / "init", base.config) |
| prec = st.Precision.load(Path(__file__).resolve().parent.parent / "configs/precision/cb_iq3s.yaml") |
| prec.codebook = str(tmp_path / "init") |
| prec.experts = QuantSpec(method="codebook", group_size=32) |
| m = st.build_student(base, prec, quant=True) |
| T = gguf.GGMLQuantizationType |
| ex = m.model.layers[1].mlp.experts |
| deq = lambda b, q: gguf.quants.dequantize(b, q).reshape(*b.shape[:2], -1) |
| ref_gu = np.concatenate([deq(t[f"layers.1.{k}.blocks"], T.IQ3_S) for k in ("gate", "up")], 1) |
| assert np.array_equal(st._expert_weight(ex, "gate_up").detach().numpy(), ref_gu) |
| ref_d = deq(t["layers.1.down.blocks"], T.IQ4_NL) |
| assert np.array_equal(st._expert_weight(ex, "down").detach().numpy(), ref_d) |
| |
| m.train() |
| x = torch.randint(0, 256, (2, 32), generator=torch.Generator().manual_seed(1)) |
| st.logits(m, st.final_hidden(m, x)).float().pow(2).mean().backward() |
| assert ex.gate_up_dlog.grad.abs().sum() > 0 and ex.gate_up_slat.grad.abs().sum() > 0 |
| assert ex.down_dlog.grad.abs().sum() > 0 and not hasattr(ex, "down_slat") |
| |
| st.export_packed_student(m, tmp_path / "p0") |
| sd, packed, man = st.read_packed(tmp_path / "p0") |
| tens = student_tensors({}, packed, gguf) |
| for k in ("gate", "up", "down"): |
| assert tens[f"blk.1.ffn_{k}_exps.weight"][1].tobytes() == t[f"layers.1.{k}.blocks"].tobytes() |
| |
| g = torch.Generator().manual_seed(2) |
| with torch.no_grad(): |
| for layer in m.model.layers: |
| e = layer.mlp.experts |
| e.gate_up_dlog.add_(0.05 * torch.randn(e.gate_up_dlog.shape, generator=g)) |
| e.gate_up_slat.add_(1.5 * torch.randn(e.gate_up_slat.shape, generator=g)) |
| e.down_dlog.add_(0.05 * torch.randn(e.down_dlog.shape, generator=g)) |
| st.export_packed_student(m, tmp_path / "p1") |
| sd, packed, man = st.read_packed(tmp_path / "p1") |
| for name in ("gate_up", "down"): |
| key = f"model.layers.1.mlp.experts.{name}_proj" |
| qt, raw = packed_blocks(packed[key], gguf) |
| got = gguf.quants.dequantize(raw, qt).reshape(packed[key]["shape"]) |
| assert np.array_equal(got, st._expert_weight(ex, name).detach().numpy()) |
| assert torch.equal(st.dequant_codebook(key, packed[key], torch.float32), st._expert_weight(ex, name).detach()) |
| assert man["codebook_formats"]["model.layers.0.mlp.experts.gate_up_proj"] == "IQ3_S" |
|
|
|
|
| FREE = ("IQ3_S", "IQ2_XXS", "IQ4_NL") |
|
|
|
|
| @pytest.mark.parametrize("name", FREE) |
| def test_free_codes_start_is_the_file(gguf, name): |
| """Latents at llama.cpp's dequantised weights find llama.cpp's own codes: the re-encoded blocks are identical.""" |
| fmt, b = _blocks(gguf, name) |
| ref = gguf.quants.dequantize(b, getattr(gguf.GGMLQuantizationType, name)).reshape(len(b), fmt.block) |
| grid = cb.grid_from_gguf(fmt, gguf) |
| d0 = torch.from_numpy(cb.get_d(b).astype(np.float32))[:, None] |
| s0 = torch.from_numpy(cb.get_s(fmt, b).astype(np.float32)) |
| lat = torch.from_numpy(ref.copy()) |
| w = cb.weight_free(fmt, lat, d0, torch.zeros_like(d0), None if name == "IQ4_NL" else s0, grid) |
| assert torch.equal(w, lat) |
| e = cb.eff_torch(fmt, d0, s0).unsqueeze(-1) |
| idx, vals = cb.free_codes(fmt, (lat.unflatten(-1, (-1, fmt.sub)) / e).flatten(-2), grid) |
| out = cb.encode(fmt, idx, vals, cb.get_d(b), cb.get_s(fmt, b)) |
| assert np.array_equal(out.reshape(b.shape), b) |
|
|
|
|
| @pytest.mark.parametrize("name", FREE) |
| def test_free_codes_moved_dequantise_as_trained(gguf, name): |
| fmt, b = _blocks(gguf, name) |
| ref = gguf.quants.dequantize(b, getattr(gguf.GGMLQuantizationType, name)).reshape(len(b), fmt.block) |
| grid = cb.grid_from_gguf(fmt, gguf) |
| g = torch.Generator().manual_seed(3) |
| lat = torch.from_numpy(ref.copy()) * (1 + 0.3 * torch.randn(ref.shape, generator=g)) |
| d0 = torch.from_numpy(cb.get_d(b).astype(np.float32))[:, None] |
| dlog = 0.05 * torch.randn(d0.shape, generator=g) |
| slat = None if name == "IQ4_NL" else torch.from_numpy(cb.get_s(fmt, b).astype(np.float32)) + torch.randn(len(b), fmt.block // fmt.sub, generator=g) |
| w = cb.weight_free(fmt, lat, d0, dlog, slat, grid) |
| d, s = cb.trained_fields(fmt, d0, dlog, slat) |
| e = cb.eff_torch(fmt, torch.from_numpy(d.astype(np.float32)), torch.from_numpy(s.astype(np.float32))).unsqueeze(-1) |
| idx, vals = cb.free_codes(fmt, (lat.unflatten(-1, (-1, fmt.sub)) / e).flatten(-2), grid) |
| out = cb.encode(fmt, idx, vals, d, s) |
| got = gguf.quants.dequantize(out.reshape(b.shape), getattr(gguf.GGMLQuantizationType, name)).reshape(w.shape) |
| assert torch.equal(w, torch.from_numpy(got)) |
| assert (out.reshape(b.shape) != b).any() |
|
|
|
|
| def test_student_free_round_trip(gguf, tmp_path): |
| """Free codes: the start is the file; latents get gradients; a moved student exports to its own weights.""" |
| from test_student import tiny |
| from lowbit.model import student as st |
| from lowbit.quant.ternary import QuantSpec |
| from lowbit.export.gguf import packed_blocks |
| from safetensors.torch import save_file |
| base = tiny(hidden=256) |
| (tmp_path / "init").mkdir() |
| t = _fake_init(gguf, tmp_path / "init", base.config) |
| save_file({n: cb.grid_from_gguf(cb.FORMATS[n], gguf) for n in ("IQ3_S", "IQ4_NL")}, str(tmp_path / "init/grids.safetensors")) |
| prec = st.Precision.load(Path(__file__).resolve().parent.parent / "configs/precision/cb_iq3s.yaml") |
| prec.codebook, prec.codebook_mode = str(tmp_path / "init"), "free" |
| prec.experts = QuantSpec(method="codebook", group_size=32) |
| m = st.build_student(base, prec, quant=True) |
| ex = m.model.layers[1].mlp.experts |
| T = gguf.GGMLQuantizationType |
| deq = lambda b, q: gguf.quants.dequantize(b, q).reshape(*b.shape[:2], -1) |
| ref = np.concatenate([deq(t[f"layers.1.{k}.blocks"], T.IQ3_S) for k in ("gate", "up")], 1) |
| assert np.array_equal(st._expert_weight(ex, "gate_up").detach().numpy(), ref) |
| m.train() |
| x = torch.randint(0, 256, (2, 32), generator=torch.Generator().manual_seed(1)) |
| st.logits(m, st.final_hidden(m, x)).float().pow(2).mean().backward() |
| assert ex.gate_up_latent.grad.abs().sum() > 0 and ex.down_latent.grad.abs().sum() > 0 |
| assert ex.gate_up_slat.grad.abs().sum() > 0 |
| g = torch.Generator().manual_seed(4) |
| with torch.no_grad(): |
| for layer in m.model.layers: |
| e = layer.mlp.experts |
| for n in ("gate_up", "down"): |
| lat = getattr(e, f"{n}_latent") |
| lat.mul_(1 + 0.3 * torch.randn(lat.shape, generator=g)) |
| st.export_packed_student(m, tmp_path / "p") |
| sd, packed, man = st.read_packed(tmp_path / "p") |
| for name in ("gate_up", "down"): |
| key = f"model.layers.1.mlp.experts.{name}_proj" |
| qt, raw = packed_blocks(packed[key], gguf) |
| got = gguf.quants.dequantize(raw, qt).reshape(packed[key]["shape"]) |
| assert np.array_equal(got, st._expert_weight(ex, name).detach().numpy()) |
|
|
|
|
| def test_codebook_export_resume(gguf, tmp_path): |
| """A student loaded from its own export exports the same file again (blocks and dense tensors).""" |
| from test_student import tiny |
| from lowbit.model import student as st |
| from lowbit.quant.ternary import QuantSpec |
| from safetensors.torch import load_file |
| base = tiny(hidden=256) |
| (tmp_path / "init").mkdir() |
| _fake_init(gguf, tmp_path / "init", base.config) |
| prec = st.Precision.load(Path(__file__).resolve().parent.parent / "configs/precision/cb_iq3s.yaml") |
| prec.codebook = str(tmp_path / "init") |
| prec.experts = QuantSpec(method="codebook", group_size=32) |
| m = st.build_student(base, prec, quant=True) |
| g = torch.Generator().manual_seed(5) |
| with torch.no_grad(): |
| for layer in m.model.layers: |
| e = layer.mlp.experts |
| e.gate_up_dlog.add_(0.05 * torch.randn(e.gate_up_dlog.shape, generator=g)) |
| e.gate_up_slat.add_(1.5 * torch.randn(e.gate_up_slat.shape, generator=g)) |
| e.down_dlog.add_(0.05 * torch.randn(e.down_dlog.shape, generator=g)) |
| layer.self_attn.q_proj.weight.add_(0.01) |
| st.export_packed_student(m, tmp_path / "a") |
| m2 = st.build_student(tiny(hidden=256), prec, quant=True) |
| st.load_codebook_export(m2, tmp_path / "a") |
| st.export_packed_student(m2, tmp_path / "b") |
| a, b = load_file(str(tmp_path / "a/model.safetensors")), load_file(str(tmp_path / "b/model.safetensors")) |
| assert a.keys() == b.keys() and all(torch.equal(a[k], b[k]) for k in a) |
|
|
