File size: 3,898 Bytes
e5235fb | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 | #!/usr/bin/env python3
"""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 # noqa: E402
from lowbit.quant import codebook as cb # noqa: E402
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)) # gguf-py lists ne0 first
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()
|