| |
| """Quick KL check of packed students against the teacher on one GPU (PyTorch), for deciding on a pod whether a run |
| paid off before downloading it. Mirrors the llama.cpp harness: the first N held-out eval sequences of 4,096 tokens, |
| plain causal attention, KL(teacher || student) over the full vocabulary, scored on the second half of each sequence. |
| |
| python scripts/pod_kl.py --packed runs/A/packed/tok_N --packed runs/B/packed/tok_N [--n 40] |
| |
| Absolute values differ slightly from llama.cpp's (tokenisation of the decoded text, kernels); compare runs measured |
| here with each other. |
| """ |
| from __future__ import annotations |
|
|
| import argparse |
| import json |
| 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.model.student import load_packed_student |
|
|
| TEACHER = "/path/to/data/models/Marco-Nano-Instruct" |
| EVAL = "/path/to/data/shards/eval/tokens.u32" |
| CTX = 4096 |
|
|
|
|
| @torch.no_grad() |
| def logp(model, ids: torch.Tensor) -> torch.Tensor: |
| return torch.log_softmax(model(ids).logits[0, CTX // 2:].float(), -1) |
|
|
|
|
| def main() -> None: |
| p = argparse.ArgumentParser() |
| p.add_argument("--packed", action="append", required=True) |
| p.add_argument("--n", type=int, default=40) |
| a = p.parse_args() |
| from transformers import AutoModelForCausalLM |
| toks = np.memmap(EVAL, dtype=np.uint32, mode="r") |
| seqs = [torch.from_numpy(toks[i * CTX:(i + 1) * CTX].astype(np.int64))[None].cuda() for i in range(a.n)] |
| teacher = AutoModelForCausalLM.from_pretrained(TEACHER, dtype=torch.bfloat16).cuda().eval() |
| ref = [logp(teacher, s).cpu() for s in seqs] |
| del teacher |
| torch.cuda.empty_cache() |
| out = {} |
| for d in a.packed: |
| m = load_packed_student(d, device="cuda") |
| kl = n = 0.0 |
| same = 0 |
| for s, lt in zip(seqs, ref): |
| ls = logp(m, s) |
| lt = lt.cuda() |
| kl += (lt.exp() * (lt - ls)).sum(-1).sum().item() |
| same += (lt.argmax(-1) == ls.argmax(-1)).sum().item() |
| n += lt.shape[0] |
| out[d] = {"kl": round(kl / n, 5), "same_top1_pct": round(100 * same / n, 2), "tokens": int(n)} |
| print(json.dumps({d: out[d]}), flush=True) |
| del m |
| torch.cuda.empty_cache() |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
