"""Train Nirnay on a CUDA GPU (Kaggle, Colab, any box) without the precompute cache. The CPU pipeline (precompute.py + train_head.py) stores the frozen encoder's layer-SPLIT token states for every training text on disk (~40 GB), because recomputing them on a CPU is the expensive part. On a GPU that is cheap: * questions and options (few, reused by many records) are encoded once at start-up exactly as precompute.py does: the full frozen encoder, keeping the layer-SPLIT states and the pooled final-layer vector; * states are encoded per batch up to layer SPLIT (their pooled vectors are zeros, as in the mid-only cache). Model, loss, batching, dev selection, temperature calibration and export are train_head.py's own code, so the result is the model the CPU pipeline trains, only faster. Mixed precision (fp16 autocast + loss scaling) on the GPU. NIRNAY_RECORDS=records.jsonl NIRNAY_BACKBONE=BAAI/bge-base-en-v1.5 NIRNAY_POOL=cls NIRNAY_PREFIX=none NIRNAY_TRAIN_STATE_TOKENS=256 NIRNAY_SPLIT=8 NIRNAY_MAX_STATE=256 NIRNAY_EPOCHS=1.6 NIRNAY_RUN=v3 NIRNAY_EXPORT=out/nirnay-v3 python train/train_gpu.py NIRNAY_DEVICE=cpu runs it on a CPU (tests); NIRNAY_LIMIT_RECORDS=n reads only the first n records. """ import json import math import os import random import sys import time import numpy as np import torch HERE = os.path.dirname(os.path.abspath(__file__)) ROOT = os.path.dirname(HERE) sys.path.insert(0, ROOT) sys.path.insert(0, HERE) os.environ.setdefault("NIRNAY_FUSION", "deep") os.environ.setdefault("NIRNAY_SPLIT", "8") import train_head as th # noqa: E402 (reads the NIRNAY_* hyper-parameters from the environment) from nirnay.encoder import load_hf_encoder # noqa: E402 from nirnay.model import QTYPES, DeepFusionHead # noqa: E402 from nirnay.render import render_question, render_state # noqa: E402 BACKBONE = os.environ.get("NIRNAY_BACKBONE", "intfloat/multilingual-e5-small") POOL = os.environ.get("NIRNAY_POOL", "mean") PREFIX = os.environ.get("NIRNAY_PREFIX", "query: ") PREFIX = "" if PREFIX.strip().lower() == "none" else PREFIX REC = os.environ.get("NIRNAY_RECORDS", os.path.join(HERE, "data", "records.jsonl")) LIMIT = int(os.environ.get("NIRNAY_LIMIT_RECORDS", "0")) LIM = {"state_tokens": int(os.environ.get("NIRNAY_TRAIN_STATE_TOKENS", "320")), "question_tokens": 64, "option_tokens": 32} DEV = torch.device(os.environ.get("NIRNAY_DEVICE", "cuda" if torch.cuda.is_available() else "cpu")) AMP = DEV.type == "cuda" TAB = torch.float16 if AMP else torch.float32 # dtype of the question / option state tables CHUNK_TOKENS = 16384 # padded tokens per encoder call at start-up # Attention-memory budget per forward pass (sum over samples of the fused layers' score elements per head). A batch # over budget is split into micro-batches whose gradients accumulate into one optimizer step: the same update as # the whole batch (padding is masked), with bounded GPU memory. ~8M fits a 16 GB T4 with room to spare. MICRO_BUDGET = float(os.environ.get("NIRNAY_MICRO_BUDGET", "8e6")) assert th.FUSION == "deep" and th.SPLIT > 0, "train_gpu.py trains deep fusion (NIRNAY_FUSION=deep, NIRNAY_SPLIT>0)" def autocast(): return torch.autocast("cuda", dtype=torch.float16, enabled=AMP) class LiveStore: """train_head.Store's interface (meta, index, off, vec) with token ids instead of cached states.""" def __init__(self, enc, tok): probe = tok.encode("a", add_special_tokens=True).ids self.cls, self.sep, self.pad = probe[0], probe[-1], enc.cfg.pad_id self.prefix_ids = tok.encode(PREFIX, add_special_tokens=False).ids if PREFIX else [] t0 = time.time() states, questions, options, index = {}, {}, {}, [] with open(REC, encoding="utf-8") as f: # exactly precompute.py's numbering and canonical option order for line in f: r = json.loads(line) try: nq, qtext, otexts = render_question("q", r["q"]) except Exception: continue perm = sorted(range(len(otexts)), key=lambda i: (otexts[i], i)) st = render_state(r["state"]) si = states.setdefault(st, len(states)) qi = questions.setdefault(qtext, len(questions)) oi = [options.setdefault(otexts[i], len(options)) for i in perm] index.append({"s": si, "q": qi, "o": oi, "y": perm.index(int(r["y"])), "t": QTYPES[nq["type"]], "src": r["src"], "split": r["split"]}) if LIMIT and len(index) >= LIMIT: break self.index = index print("records %d | unique states %d questions %d options %d | %.0fs" % (len(index), len(states), len(questions), len(options), time.time() - t0), flush=True) self.ids, self.off = {}, {} for name, texts, budget in (("states", states, LIM["state_tokens"]), ("questions", questions, LIM["question_tokens"]), ("options", options, LIM["option_tokens"])): room = budget - 2 - len(self.prefix_ids) seqs = [[self.cls] + self.prefix_ids + e.ids[:max(1, room)] + [self.sep] for e in tok.encode_batch(list(texts), add_special_tokens=False)] lens = np.array([len(s) for s in seqs], dtype=np.int64) off = np.zeros(len(seqs) + 1, dtype=np.int64) off[1:] = np.cumsum(lens) self.ids[name] = np.concatenate([np.asarray(s, dtype=np.int64) for s in seqs]) self.off[name] = off print(" tokenized %-9s %7d texts %10d tokens" % (name, len(seqs), int(off[-1])), flush=True) d = enc.cfg.hidden self.vec = {"states": np.zeros((len(states), d), np.float32)} # mid-only states: no pooled vector self.mid, self.vecd = {}, {} for name in ("questions", "options"): self._encode_table(enc, name) self.meta = {"backbone": BACKBONE, "pool": POOL, "prefix": PREFIX, "limits": LIM, "d": d, "mid_layer": th.SPLIT, "state_mid_only": True, "n_records": len(index), "n_states": len(states), "n_questions": len(questions), "n_options": len(options), "special_ids": {"cls": self.cls, "sep": self.sep, "pad": self.pad}, "encoder": enc.cfg.to_dict(), "source_dir": enc.source_dir} def batch_ids(self, name, items): off = self.off[name] lens = [int(off[i + 1] - off[i]) for i in items] L = max(lens) ids = torch.full((len(items), L), self.pad, dtype=torch.long) mask = torch.zeros((len(items), L), dtype=torch.bool) for r, (i, n) in enumerate(zip(items, lens)): ids[r, :n] = torch.from_numpy(self.ids[name][off[i]:off[i] + n]) mask[r, :n] = True return ids.to(DEV, non_blocking=True), mask.to(DEV, non_blocking=True) @torch.no_grad() def _encode_table(self, enc, name): """Full frozen encoder over every question / option: layer-SPLIT states + pooled final vector.""" t0 = time.time() off = self.off[name] n = len(off) - 1 lens = off[1:] - off[:-1] mid = torch.zeros((int(off[-1]), enc.cfg.hidden), dtype=TAB, device=DEV) vec = torch.zeros((n, enc.cfg.hidden), dtype=torch.float32, device=DEV) order = np.argsort(lens, kind="stable") i = 0 while i < n: j = i while j < n and lens[order[j]] * (j - i + 1) <= max(CHUNK_TOKENS, lens[order[i]]): j += 1 j = max(j, i + 1) idx = order[i:j].tolist() ids, mask = self.batch_ids(name, idx) with autocast(): H, M = enc(ids, mask, keep_layer=th.SPLIT) vec[idx] = enc.pool(H.float(), mask) for r, t in enumerate(idx): mid[off[t]:off[t + 1]] = M[r, :lens[t]].to(TAB) i = j self.mid[name], self.vecd[name] = mid, vec self.vec[name] = vec.cpu().numpy() print(" encoded %-9s %7d texts in %.0fs" % (name, n, time.time() - t0), flush=True) def rows(self, n, i): # train_head.Store compatibility (questions / options only) return self.mid[n][self.off[n][i]:self.off[n][i + 1]] def make_batch(st, enc, recs, rng): """train_head.make_batch, with the states encoded live and every tensor on the device.""" d = st.meta["d"] B = len(recs) opts, golds = [], [] for r in recs: o, y = list(r["o"]), r["y"] if len(o) > th.MAX_OPTS: keep = [y] + rng.sample([i for i in range(len(o)) if i != y], th.MAX_OPTS - 1) keep.sort() y = keep.index(y) o = [o[i] for i in keep] opts.append(o) golds.append(y) ids, mask = st.batch_ids("states", [r["s"] for r in recs]) with torch.no_grad(), autocast(): S = enc(ids, mask, stop_layer=th.SPLIT) n = min(th.MAX_STATE, S.shape[1]) sm = mask[:, :n] S = (S[:, :n] * sm.unsqueeze(-1)).to(TAB) # zeros at padding, exactly as the cached batches qo, oo = st.off["questions"], st.off["options"] Lq = max(int(qo[r["q"] + 1] - qo[r["q"]]) for r in recs) K = max(len(o) for o in opts) Lo = max(int(oo[i + 1] - oo[i]) for o in opts for i in o) Q = torch.zeros((B, Lq, d), dtype=TAB, device=DEV) qm = torch.zeros((B, Lq), dtype=torch.bool, device=DEV) O = torch.zeros((B, K, Lo, d), dtype=TAB, device=DEV) om = torch.zeros((B, K, Lo), dtype=torch.bool, device=DEV) ok = torch.zeros((B, K), dtype=torch.bool, device=DEV) ov = torch.zeros((B, K, d), dtype=torch.float32, device=DEV) for b, (r, o) in enumerate(zip(recs, opts)): x = st.rows("questions", r["q"]) Q[b, :len(x)] = x qm[b, :len(x)] = True for j, i in enumerate(o): x = st.rows("options", i) O[b, j, :len(x)] = x om[b, j, :len(x)] = True ok[b, :len(o)] = True ov[b, :len(o)] = st.vecd["options"][o] sv = torch.zeros((B, d), dtype=torch.float32, device=DEV) t = torch.tensor([r["t"] for r in recs], device=DEV) return (S, sm, Q, qm, O, om, ok, t, sv, ov), torch.tensor(golds, device=DEV) def micro_batches(st, recs, budget): """Split a batch into near-equal parts whose attention cost stays under `budget`.""" so, qo, oo = st.off["states"], st.off["questions"], st.off["options"] n = min(th.MAX_STATE, max(int(so[r["s"] + 1] - so[r["s"]]) for r in recs)) lq = max(int(qo[r["q"] + 1] - qo[r["q"]]) for r in recs) k = max(min(len(r["o"]), th.MAX_OPTS) for r in recs) lo = max(int(oo[i + 1] - oo[i]) for r in recs for i in r["o"]) per = k * lo * (n + lq + lo) + lq * (n + lq) + n * n parts = max(1, math.ceil(len(recs) * per / budget)) size = math.ceil(len(recs) / parts) return [recs[i:i + size] for i in range(0, len(recs), size)] @torch.no_grad() def collect_logits(model, st, enc, recs, bs=128): model.eval() rng = random.Random(0) res = [] recs = sorted(recs, key=lambda r: th.shape_key(st, r)) for i in range(0, len(recs), bs): part = [r for r in recs[i:i + bs] if len(r["o"]) <= th.MAX_OPTS] # as train_head: longer ones are skipped if not part: continue for mb in micro_batches(st, part, 3 * MICRO_BUDGET): # no activations kept: a larger budget is safe xb, yb = make_batch(st, enc, mb, rng) with autocast(): lg = model(*xb) for r, z, y in zip(mb, lg.float().cpu().numpy(), yb.tolist()): res.append((r, z, y)) model.train() return res def main(): from tokenizers import Tokenizer torch.manual_seed(th.SEED) rng = random.Random(th.SEED) torch.set_num_threads(int(os.environ.get("NIRNAY_THREADS", "4"))) os.makedirs(th.RUND, exist_ok=True) print("device", DEV, torch.cuda.get_device_name(DEV) if DEV.type == "cuda" else "", "| amp", AMP, flush=True) enc = load_hf_encoder(BACKBONE, POOL, PREFIX).to(DEV).eval() for p in enc.parameters(): p.requires_grad_(False) tok = Tokenizer.from_file(os.path.join(enc.source_dir, "tokenizer.json")) tok.no_truncation() tok.no_padding() st = LiveStore(enc, tok) tr = [r for r in st.index if r["split"] == "train"] dev = [r for r in st.index if r["split"] == "dev"] cal = [r for r in st.index if r["split"] == "cal"] if os.environ.get("NIRNAY_DEV_LIMIT"): dev = dev[:int(os.environ["NIRNAY_DEV_LIMIT"])] print("train %d dev %d cal %d" % (len(tr), len(dev), len(cal)), flush=True) d = enc.cfg.hidden head = DeepFusionHead(list(enc.layers)[th.SPLIT:], d, pool=POOL) pre = [p for nm, p in head.named_parameters() if nm.startswith("layers.")] new = [p for nm, p in head.named_parameters() if not nm.startswith("layers.")] for p in pre + new: p.requires_grad_(True) model = th.Model(head, None, POOL).to(DEV) if os.environ.get("NIRNAY_INIT"): model.load_state_dict(torch.load(os.environ["NIRNAY_INIT"], map_location="cpu", weights_only=True)) print("warm start from", os.environ["NIRNAY_INIT"], flush=True) opt = torch.optim.AdamW([{"params": new, "lr": th.LR}, {"params": pre, "lr": th.ADAPTER_LR}], weight_decay=0.01, betas=(0.9, 0.98)) scaler = (torch.amp.GradScaler("cuda", enabled=AMP) if hasattr(torch, "amp") and hasattr(torch.amp, "GradScaler") else torch.cuda.amp.GradScaler(enabled=AMP)) total = int(math.ceil(len(tr) / th.BATCH) * th.EPOCHS) sched = torch.optim.lr_scheduler.LambdaLR( opt, lambda s: min(1.0, (s + 1) / th.WARMUP) * 0.5 * (1 + math.cos(math.pi * min(1.0, s / max(1, total))))) log = {"run": th.RUN, "backbone": BACKBONE, "fusion": th.FUSION, "split": th.SPLIT, "batch": th.BATCH, "lr": th.LR, "adapter_lr": th.ADAPTER_LR, "epochs": th.EPOCHS, "device": str(DEV), "amp": AMP, "evals": []} best = (-1.0, None) step, t0, run_loss = 0, time.time(), 0.0 model.train() def evaluate(): m = th.metrics(collect_logits(model, st, enc, dev)) m.update({"step": step, "loss": round(run_loss, 4), "elapsed_min": round((time.time() - t0) / 60, 1)}) log["evals"].append({k: v for k, v in m.items() if k != "by_src"}) print("\n dev acc %.4f macro %.4f nll %.4f (step %d, %.0f min)" % (m["acc"], m["macro_acc"], m["nll"], step, (time.time() - t0) / 60), flush=True) json.dump(log, open(os.path.join(th.RUND, "log.json"), "w"), indent=2) return m while step < total: for batch in th.batches_by_length(st, tr, th.BATCH, rng): opt.zero_grad(set_to_none=True) loss = 0.0 for mb in micro_batches(st, batch, MICRO_BUDGET): # gradient accumulation: one step per batch xb, yb = make_batch(st, enc, mb, rng) with autocast(): logits = model(*xb) part = th.loss_fn(logits.float(), yb, xb[7]) * (len(mb) / len(batch)) scaler.scale(part).backward() loss += part.detach() scaler.unscale_(opt) torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0) scaler.step(opt) scaler.update() sched.step() step += 1 run_loss = 0.98 * run_loss + 0.02 * float(loss) if step > 1 else float(loss) if step % 100 == 0: el = time.time() - t0 print("step %d/%d loss %.4f %.3f s/step ETA %.0f min" % (step, total, run_loss, el / step, el / step * (total - step) / 60), flush=True) if step % th.EVAL_EVERY == 0 or step >= total: m = evaluate() if m["macro_acc"] > best[0]: best = (m["macro_acc"], m) torch.save(model.state_dict(), os.path.join(th.RUND, "model_best.pt")) if step >= total: break print("\nbest dev macro acc %.4f" % best[0], flush=True) model.load_state_dict(torch.load(os.path.join(th.RUND, "model_best.pt"), map_location=DEV, weights_only=True)) cal_rows = collect_logits(model, st, enc, cal) temps = th.fit_temperatures(cal_rows) print("temperatures", temps, flush=True) m_raw, m_cal = th.metrics(cal_rows), th.metrics(cal_rows, temps) log["best_dev"] = best[1] log["cal_raw"] = {k: v for k, v in m_raw.items() if k != "by_src"} log["cal_calibrated"] = {k: v for k, v in m_cal.items() if k != "by_src"} log["temperatures"] = temps log["train_minutes"] = round((time.time() - t0) / 60, 1) json.dump(log, open(os.path.join(th.RUND, "log.json"), "w"), indent=2) model.cpu() enc.cpu() th.export(model, enc, st.meta, temps, log) if __name__ == "__main__": main()