| |
| """Repack a student GGUF's expert gate/up tensors from Q2_0 to llama.cpp's TQ1_0 (1.69 bits per weight). |
| |
| python scripts/tq1_repack.py <in.gguf> <out.gguf> |
| |
| TQ1_0 holds ternary codes for blocks of 256 weights with ONE fp16 scale, and needs rows whose length is a |
| multiple of 256. Our groups are 128 wide with their own scales, so each TQ1_0 block covers two groups: |
| - the trained codes are kept exactly; |
| - the block's scale is the least-squares fit to both groups' values, mean |w| over its nonzero codes; |
| - ffn_down_exps (rows of 384) cannot be TQ1_0 and stays Q2_0, as does everything else. |
| Prints the size change and the scale approximation's relative weight error. |
| """ |
| from __future__ import annotations |
|
|
| import json |
| import sys |
| from pathlib import Path |
|
|
| import numpy as np |
|
|
| REPO = Path(__file__).resolve().parent.parent |
| sys.path.insert(0, str(REPO / "src")) |
| from lowbit.export.gguf import dequant_q2_0, gguf_module, write_gguf |
|
|
| LLAMA = Path("/path/to/llama.cpp") |
| TQ_BLOCK = 256 |
|
|
|
|
| def to_tq1(raw: np.ndarray, gguf) -> tuple[np.ndarray, float]: |
| """Q2_0 bytes [..., nb · 18] → TQ1_0 bytes [..., cols / 256 · 54], and the relative L2 weight error.""" |
| lead = raw.shape[:-1] |
| w = dequant_q2_0(raw) |
| cols = w.shape[-1] |
| if cols % TQ_BLOCK: |
| raise ValueError(f"rows of {cols} are not a multiple of {TQ_BLOCK}") |
| blk = w.reshape(-1, TQ_BLOCK) |
| nz = blk != 0 |
| d = np.abs(blk).sum(-1) / np.maximum(nz.sum(-1), 1) |
| codes = np.sign(blk) |
| x = (codes * d[:, None]).astype(np.float32) |
| err = float(np.linalg.norm(x - blk) / max(np.linalg.norm(blk), 1e-30)) |
| q = gguf.quants.quantize(x.reshape(*lead, cols), gguf.GGMLQuantizationType.TQ1_0) |
| back = gguf.quants.dequantize(q, gguf.GGMLQuantizationType.TQ1_0).reshape(-1, TQ_BLOCK) |
| if not np.array_equal(np.sign(back), codes): |
| raise RuntimeError("TQ1_0 changed a code") |
| return q, err |
|
|
|
|
| def main() -> None: |
| src, dst = sys.argv[1], sys.argv[2] |
| gguf = gguf_module(LLAMA) |
| r = gguf.GGUFReader(src) |
| T = gguf.GGMLQuantizationType |
| tensors, errs, moved = {}, [], 0 |
| for t in r.tensors: |
| data = np.asarray(t.data) |
| if t.tensor_type == T.Q2_0 and (t.name.endswith("ffn_gate_exps.weight") or t.name.endswith("ffn_up_exps.weight")): |
| q, e = to_tq1(data, gguf) |
| tensors[t.name] = (T.TQ1_0, q) |
| errs.append(e) |
| moved += data.nbytes - q.nbytes |
| else: |
| tensors[t.name] = (t.tensor_type, data) |
| info = write_gguf(src, dst, tensors, gguf) |
| out = {"in_bytes": Path(src).stat().st_size, "out_bytes": info["bytes"], |
| "saved_pct": round(100 * (1 - info["bytes"] / Path(src).stat().st_size), 1), |
| "tq1_tensors": len(errs), "scale_rel_err_mean": round(float(np.mean(errs)), 4), |
| "scale_rel_err_max": round(float(np.max(errs)), 4)} |
| print(json.dumps(out, indent=1)) |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
