| |
| """D-58: a codebook student's start, read from a llama.cpp GGUF of the teacher (`scripts/ptq_baseline.py`). |
| |
| python scripts/cb_import.py --gguf $MT_DATA/ptq/gguf/teacher-iq3_s-iq4nldown-imat.gguf --out $MT_DATA/cbinit/iq3s |
| |
| Writes <out>/experts.safetensors with, per layer and projection (gate, up, down): `blocks` (the GGUF's raw blocks, |
| uint8 [E, rows, nblk, nbytes]), `v` (int8 [E, rows, cols]: codebook entry × sign per weight), `d` (fp16 |
| [E, rows, nblk]) and `s` (uint8 [E, rows, nsub]); and <out>/manifest.json with each tensor's format and the |
| source file's sha256. Everything outside the routed experts comes from the teacher, as in every PTQ baseline |
| (attention, embeddings bf16; norms, router f32), so this file and the GGUF describe the same model. |
| """ |
| from __future__ import annotations |
|
|
| import argparse |
| import hashlib |
| import json |
| import re |
| import sys |
| from pathlib import Path |
|
|
| import numpy as np |
| import torch |
|
|
| REPO = Path(__file__).resolve().parent.parent |
| sys.path.insert(0, str(REPO / "src")) |
| from lowbit.export.gguf import gguf_module |
| from lowbit.quant import codebook as cb |
|
|
| LLAMA = "/path/to/llama.cpp" |
|
|
|
|
| def sha256(p: Path) -> str: |
| h = hashlib.sha256() |
| with open(p, "rb") as f: |
| for chunk in iter(lambda: f.read(1 << 24), b""): |
| h.update(chunk) |
| return h.hexdigest() |
|
|
|
|
| def write_grids(out: Path, names: set, gguf) -> None: |
| """grids.safetensors: each format's codebook (`quant.codebook.grid_from_gguf`), for free-code training.""" |
| from safetensors.torch import save_file as save_pt |
| save_pt({n: cb.grid_from_gguf(cb.FORMATS[n], gguf).contiguous() for n in sorted(names)}, str(out / "grids.safetensors")) |
|
|
|
|
| def main() -> None: |
| p = argparse.ArgumentParser() |
| p.add_argument("--gguf", required=True) |
| p.add_argument("--out", required=True) |
| p.add_argument("--llama", default=LLAMA, help="a directory holding llama.cpp's gguf-py (spark: third_party)") |
| p.add_argument("--grids-only", action="store_true", help="only (re)write grids.safetensors for the manifest's formats") |
| a = p.parse_args() |
| from safetensors.numpy import save_file |
| gguf = gguf_module(a.llama) |
| out = Path(a.out) |
| out.mkdir(parents=True, exist_ok=True) |
| if a.grids_only: |
| write_grids(out, set(json.loads((out / "manifest.json").read_text())["formats"].values()), gguf) |
| return |
| tensors, formats = {}, {} |
| for t in gguf.GGUFReader(a.gguf).tensors: |
| m = re.fullmatch(r"blk\.(\d+)\.ffn_(gate|up|down)_exps\.weight", t.name) |
| if not m: |
| continue |
| name = gguf.GGMLQuantizationType(t.tensor_type).name |
| if name not in cb.FORMATS: |
| sys.exit(f"{t.name} is {name}, not a codebook format ({sorted(cb.FORMATS)})") |
| fmt = cb.FORMATS[name] |
| E, rows, cols = (int(x) for x in reversed(t.shape)) |
| nblk = cols // fmt.block |
| b = np.asarray(t.data).reshape(-1, fmt.nbytes) |
| key = f"layers.{m.group(1)}.{m.group(2)}" |
| formats[key] = name |
| tensors[f"{key}.blocks"] = b.reshape(E, rows, nblk, fmt.nbytes).copy() |
| tensors[f"{key}.v"] = cb.values(fmt, b, gguf).reshape(E, rows, cols) |
| tensors[f"{key}.d"] = cb.get_d(b).reshape(E, rows, nblk) |
| tensors[f"{key}.s"] = cb.get_s(fmt, b).reshape(E, rows, -1) |
| print(key, name, (E, rows, cols), flush=True) |
| save_file(tensors, str(out / "experts.safetensors")) |
| write_grids(out, set(formats.values()), gguf) |
| (out / "manifest.json").write_text(json.dumps({"source": str(a.gguf), "source_sha256": sha256(Path(a.gguf)), |
| "formats": formats}, indent=1)) |
| print(json.dumps({"out": str(out), "tensors": len(formats), |
| "bytes": (out / "experts.safetensors").stat().st_size})) |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
