import argparse import json import os import random import time import torch import torch.nn.functional as F from oev.evaluate import load_model, pack_question from oev.tokenizer_hf import HFTokenPacker class Teacher: # one frozen teacher: forward a packed case, return per-question probs def __init__(self, path, device): self.model = load_model(path, device) self.packer = HFTokenPacker(self.model.cfg["backbone"]) self.device = device @torch.no_grad() def question_probs(self, state, question, max_len, tau): # distill cases always carry answers; the default mirrors rotate()'s contract q = pack_question(dict(question, answer=question.get("answer", question["options"][0]))) ids, anchors, _ = self.packer.pack(state, q, min(max_len, self.model.cfg["max_len"])) tids = torch.tensor([ids], device=self.device) pmask = torch.zeros(1, len(ids), dtype=torch.bool, device=self.device) apos = torch.tensor([anchors], device=self.device) logits = self.model(tids, pmask, apos)[0] return F.softmax(logits.float() / tau, dim=-1) def rotate(q, rng, p): # order-invariance augmentation: rotate options + answer in lockstep n = len(q["options"]) if n < 3 or rng.random() > p: return q k = rng.randrange(1, n) ops = q["options"] rot = ops[k:] + ops[:k] out = dict(q) out["options"] = rot ans = q.get("answer") if ans in ops: out["answer"] = rot[(ops.index(ans) - k) % n] return out def load_cases(data_dir, domains, max_cases, seed, per_domain=0): # with per_domain>0, balance every domain to that count: big domains # sampled down, small ones repeated (rotation makes repeats non-identical) rng = random.Random(seed) pools = [] for d in domains: path = os.path.join(data_dir, d, "train.jsonl") if not os.path.exists(path): print(f"[distill] skipping {d} (no train split)", flush=True) continue with open(path, encoding="utf-8") as f: rows = [json.loads(line) for line in f] if per_domain: if len(rows) >= per_domain: rows = rng.sample(rows, per_domain) else: rows = (rows * (per_domain // len(rows) + 1))[:per_domain] print(f"[distill] {d}: {len(rows)} cases (balanced)", flush=True) pools.append(rows) cases = [c for rows in pools for c in rows] rng.shuffle(cases) if max_cases and len(cases) > max_cases: cases = cases[:max_cases] return cases def run_epoch(student, teachers, cases, packer, device, opt, scaler, args, epoch, f): student.train() rng = random.Random(epoch + 1) order = torch.randperm(len(cases)) running, t0, n_seen = 0.0, time.time(), 0 running_ce, running_kl, n_ce, n_kl = 0.0, 0.0, 0, 0 for bi, start in enumerate(range(0, len(order), args.batch_size)): batch = [cases[i] for i in order[start:start + args.batch_size].tolist()] opt.zero_grad(set_to_none=True) with torch.autocast("cuda", enabled=device == "cuda"): losses, ce_parts, kl_parts = [], [], [] for case in batch: q = rotate(case["questions"][0], rng, args.rotate) state = case["state"] sq = dict(q, answer=q.get("answer", q["options"][0])) ids, anchors, _ = packer.pack(state, sq, student.cfg["max_len"]) tids = torch.tensor([ids], device=device) pmask = torch.zeros(1, len(ids), dtype=torch.bool, device=device) apos = torch.tensor([anchors], device=device) slogits = student(tids, pmask, apos)[0] logq = F.log_softmax(slogits.float() / args.tau, dim=-1) tprobs, n_scored = None, 0 for t in teachers: tp = t.question_probs(state, q, student.cfg['max_len'], args.tau) if tp.shape[0] == len(q["options"]): tprobs = tp if tprobs is None else (tprobs + tp) n_scored += 1 if tprobs is None: continue tprobs = tprobs / n_scored # mean over teachers that actually scored this option set kl = -(tprobs * logq).sum() * (args.tau ** 2) kl_parts.append(kl) ans = q.get("answer") if args.alpha and ans in q["options"]: gold = q["options"].index(ans) ce = F.cross_entropy(slogits.float().unsqueeze(0), torch.tensor([gold], device=device)) ce_parts.append(ce) losses.append((1.0 - args.alpha) * kl + args.alpha * ce) else: losses.append(kl) if not losses: continue loss = torch.stack(losses).mean() scaler.scale(loss).backward() scaler.step(opt) scaler.update() running += loss.item() * len(losses) n_seen += len(losses) if kl_parts: running_kl += torch.stack(kl_parts).sum().item() n_kl += len(kl_parts) if ce_parts: running_ce += torch.stack(ce_parts).sum().item() n_ce += len(ce_parts) if bi % args.log_every == 0: msg = (f"epoch {epoch} step {bi} loss {running / max(n_seen, 1):.4f} " f"(ce {running_ce / max(n_ce, 1):.3f} kl {running_kl / max(n_kl, 1):.3f}) " f"({n_seen / max(time.time() - t0, 1):.2f} cases/s)") print(msg, flush=True) f.write(msg + "\n") f.flush() if args.save_every and bi and bi % args.save_every == 0: torch.save({"config": student.cfg, "state": student.state_dict()}, os.path.join(args.out, f"student-e{epoch}-s{bi}.pt")) return running / max(n_seen, 1) def main(): ap = argparse.ArgumentParser() ap.add_argument("--teachers", required=True, help="comma-separated teacher checkpoint paths") ap.add_argument("--student-init", default=None, help="optional warm-start checkpoint for the student") ap.add_argument("--data-dir", default="data") ap.add_argument("--domains", default="banking77,ag_news,emotion", help="comma-separated domain dirs under --data-dir, each with train.jsonl") ap.add_argument("--out", default="checkpoints_distill") ap.add_argument("--epochs", type=int, default=1) ap.add_argument("--batch-size", type=int, default=8) ap.add_argument("--lr", type=float, default=2e-5) ap.add_argument("--rotate", type=float, default=0.5, help="option-rotation augmentation probability") ap.add_argument("--tau", type=float, default=1.0, help="distillation temperature") ap.add_argument("--alpha", type=float, default=0.3, help="weight on gold cross-entropy; (1-alpha) goes to teacher KL. 0 = pure mimicry") ap.add_argument("--per-domain", type=int, default=0, help="if >0, balance each domain to this many train cases") ap.add_argument("--max-cases", type=int, default=27000) ap.add_argument("--seed", type=int, default=1234) ap.add_argument("--log-every", type=int, default=100) ap.add_argument("--save-every", type=int, default=2000) ap.add_argument("--device", default="cuda" if torch.cuda.is_available() else "cpu") args = ap.parse_args() os.makedirs(args.out, exist_ok=True) device = args.device log = open(os.path.join(args.out, "distill.log"), "a") # noqa: SIM115 - kept open for the whole run if args.student_init: student = load_model(args.student_init, device) print(f"student warm-started from {args.student_init}", flush=True) else: first = load_model(args.teachers.split(",")[0].strip(), device) from oev.model import HFBackboneOEV student = HFBackboneOEV(first.cfg["backbone"]) student.cfg = first.cfg student.eval().to(device) print("student initialized from scratch", flush=True) teacher_paths = [p.strip() for p in args.teachers.split(",") if p.strip()] teachers = [] for i, tpath in enumerate(teacher_paths, 1): # each Teacher is a full backbone load: the slowest silent stretch in # this script, so every load announces itself print(f"loading teacher {i}/{len(teacher_paths)}: {tpath} (model load, 1-2 min each)", flush=True) teachers.append(Teacher(tpath, device)) print(f"{len(teachers)} teachers loaded", flush=True) packer = HFTokenPacker(student.cfg["backbone"]) cases = load_cases(args.data_dir, [d.strip() for d in args.domains.split(",")], args.max_cases, args.seed, per_domain=args.per_domain) print(f"{len(cases)} training cases", flush=True) if not cases: raise SystemExit("[distill] 0 training cases - check --data-dir/--domains; " "each domain needs {data_dir}/{domain}/train.jsonl " "(build with: python -m oev.convert && python -m oev.convert_banking77 " "&& python -m oev.convert_typed)") opt = torch.optim.AdamW(student.parameters(), lr=args.lr) scaler = torch.amp.GradScaler(enabled=device == "cuda") for epoch in range(args.epochs): tl = run_epoch(student, teachers, cases, packer, device, opt, scaler, args, epoch, log) msg = f"epoch {epoch} done - train loss {tl:.4f}" print(msg, flush=True) log.write(msg + "\n") log.flush() torch.save({"config": student.cfg, "state": student.state_dict()}, os.path.join(args.out, "student.pt")) print(f"saved -> {args.out}/student.pt", flush=True) if __name__ == "__main__": main()