File size: 2,533 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 | #!/usr/bin/env python3
"""D-58 gate: a codebook student exactly as imported, before any training, exported packed. Its GGUF must be the
llama.cpp source file's expert tensors byte for byte (`t27_gguf.py build` then `cb_gate.py compare` on gen).
python scripts/cb_gate.py export --config configs/p6/CB3a.yaml --out $MT_DATA/runs/p6-gate-CB3a/packed/tok_0
python scripts/cb_gate.py compare --ours <gguf> --source <llama.cpp gguf>
"""
from __future__ import annotations
import argparse
import json
import os
import sys
from pathlib import Path
import numpy as np
import torch
import yaml
REPO = Path(__file__).resolve().parent.parent
sys.path.insert(0, str(REPO / "src"))
from lowbit.model import student as st # noqa: E402
from lowbit.quant.ternary import QuantSpec # noqa: E402
def export(a) -> None:
cfg = yaml.safe_load(open(a.config))
prec = st.Precision.load(REPO / cfg["model"]["precision_map"])
prec.experts = QuantSpec(method=cfg["quant"]["method"], group_size=cfg["quant"]["group_size"])
model = st.load_student(cfg["model"]["teacher"], prec, quant=True, device=a.device)
print(json.dumps(st.export_packed_student(model, a.out))[:300])
def compare(a) -> None:
sys.path.insert(0, os.environ.get("MT_LLAMA", "/path/to/llama.cpp") + "/gguf-py")
import gguf
ours = {t.name: t for t in gguf.GGUFReader(a.ours).tensors}
src = {t.name: t for t in gguf.GGUFReader(a.source).tensors}
diff = {"experts": [], "other": []}
for name, t in src.items():
same = (name in ours and ours[name].tensor_type == t.tensor_type
and np.asarray(ours[name].data).tobytes() == np.asarray(t.data).tobytes())
if not same:
diff["experts" if "_exps" in name else "other"].append(name)
print(json.dumps({"tensors": len(src), "missing": sorted(set(src) - set(ours)),
"expert_mismatches": diff["experts"], "other_mismatches": diff["other"][:20],
"n_other_mismatches": len(diff["other"])}))
def main() -> None:
p = argparse.ArgumentParser()
sub = p.add_subparsers(dest="cmd", required=True)
e = sub.add_parser("export")
e.add_argument("--config", required=True)
e.add_argument("--out", required=True)
e.add_argument("--device", default="cuda")
c = sub.add_parser("compare")
c.add_argument("--ours", required=True)
c.add_argument("--source", required=True)
a = p.parse_args()
export(a) if a.cmd == "export" else compare(a)
if __name__ == "__main__":
main()
|