code/train/train_head.py · nandukmelath/nirnay at main

Nirnay v2 (bge-base, default) + v1 (fast/): weights, package, code, benchmarks

00ee773 verified 5 days ago

17.6 kB

"""Train the Nirnay fusion head on precomputed frozen-encoder states, calibrate, export.
Objective: log score (cross-entropy) on every question + ranked probability score on ordinal
(`score`) questions. Both are strictly proper scoring rules, the same family Laya's RLCD
rewards with REINFORCE on noisy logits; here the gradient of the rule is taken exactly.
  python train\\train_head.py                        (NIRNAY_TAG=e5s by default)
"""
import json
import math
import os
import random
import shutil
import sys
import time
from collections import defaultdict

import numpy as np
import torch
import torch.nn.functional as F

HERE = os.path.dirname(os.path.abspath(__file__))
ROOT = os.path.dirname(HERE)
sys.path.insert(0, ROOT)

from nirnay.encoder import load_hf_encoder  # noqa: E402
from nirnay.engine import temp_bucket  # noqa: E402
from nirnay.model import QTYPES, DecisionHead, DeepFusionHead, StateAdapter  # noqa: E402

TAG = os.environ.get("NIRNAY_TAG", "e5s")
CACHE = os.path.join(HERE, "cache", TAG)
RUN = os.environ.get("NIRNAY_RUN", TAG + "-b%s" % os.environ.get("NIRNAY_BLOCKS", "2"))
RUND = os.path.join(HERE, "runs", RUN)
BLOCKS = int(os.environ.get("NIRNAY_BLOCKS", "2"))
BATCH = int(os.environ.get("NIRNAY_BATCH", "64"))
EPOCHS = float(os.environ.get("NIRNAY_EPOCHS", "2"))
LR = float(os.environ.get("NIRNAY_LR", "3e-4"))
WARMUP = int(os.environ.get("NIRNAY_WARMUP", "300"))
MAX_OPTS = int(os.environ.get("NIRNAY_MAX_OPTS", "64"))
MAX_STATE = int(os.environ.get("NIRNAY_MAX_STATE", "320"))   # training-time cap on state tokens
EVAL_EVERY = int(os.environ.get("NIRNAY_EVAL_EVERY", "600"))
SPLIT = int(os.environ.get("NIRNAY_SPLIT", "0"))          # >0: train encoder layers SPLIT+1.. on the state
ADAPTER_LR = float(os.environ.get("NIRNAY_ADAPTER_LR", "5e-5"))
FUSION = os.environ.get("NIRNAY_FUSION", "stack")          # "stack": frozen encoder + fusion blocks; "deep": top layers fused
WORLD = int(os.environ.get("WORLD_SIZE", "1"))             # >1: one process per CPU socket, gloo all-reduce
RANK = int(os.environ.get("RANK", "0"))
SEED = 1234
QNAMES = {v: k for k, v in QTYPES.items()}


class Model(torch.nn.Module):
    """Decision head, optionally preceded by a trainable state adapter (split encoder)."""

    def __init__(self, head, adapter=None, pool="cls"):
        super().__init__()
        self.head, self.adapter, self.pool_kind = head, adapter, pool

    def forward(self, S, sm, Q, qm, O, om, ok, t, sv, ov):
        if self.adapter is not None:
            S = self.adapter(S, sm)
            if self.pool_kind == "cls":
                v = S[:, 0]
            else:
                m = sm.unsqueeze(-1).to(S.dtype)
                v = (S * m).sum(1) / m.sum(1).clamp_min(1.0)
            sv = F.normalize(v, dim=-1)
        return self.head(S, sm, Q, qm, O, om, ok, t, sv, ov)


class Store:
    def __init__(self, d):
        self.meta = json.load(open(os.path.join(d, "meta.json")))
        self.arr, self.off, self.vec = {}, {}, {}
        for n in ("states", "questions", "options"):
            mid = (FUSION == "deep") or (SPLIT and n == "states")
            f = n + ("_mid" if mid else "") + ".f16.npy"
            self.arr[n] = np.load(os.path.join(d, f), mmap_mode="r")
            self.off[n] = np.load(os.path.join(d, n + "_off.npy"))
            self.vec[n] = np.load(os.path.join(d, n + "_vec.npy"))
        # small tables fully in RAM
        for n in ("questions", "options"):
            self.arr[n] = np.asarray(self.arr[n])
        self.index = [json.loads(l) for l in open(os.path.join(d, "index.jsonl"), encoding="utf-8")]

    def rows(self, n, i):
        return self.arr[n][self.off[n][i]:self.off[n][i + 1]]


def make_batch(st: Store, recs, rng, train: bool):
    d = st.meta["d"]
    B = len(recs)
    opts, golds = [], []
    for r in recs:
        o, y = list(r["o"]), r["y"]
        if len(o) > MAX_OPTS:
            keep = [y] + rng.sample([i for i in range(len(o)) if i != y], MAX_OPTS - 1)
            keep.sort()
            y = keep.index(y)
            o = [o[i] for i in keep]
        opts.append(o)
        golds.append(y)
    n = min(MAX_STATE, max(st.off["states"][r["s"] + 1] - st.off["states"][r["s"]] for r in recs))
    Lq = max(st.off["questions"][r["q"] + 1] - st.off["questions"][r["q"]] for r in recs)
    K = max(len(o) for o in opts)
    Lo = max(st.off["options"][i + 1] - st.off["options"][i] for o in opts for i in o)
    S = np.zeros((B, n, d), np.float32)
    sm = np.zeros((B, n), bool)
    Q = np.zeros((B, Lq, d), np.float32)
    qm = np.zeros((B, Lq), bool)
    O = np.zeros((B, K, Lo, d), np.float32)
    om = np.zeros((B, K, Lo), bool)
    ok = np.zeros((B, K), bool)
    sv = np.zeros((B, d), np.float32)
    ov = np.zeros((B, K, d), np.float32)
    for b, (r, o) in enumerate(zip(recs, opts)):
        x = st.rows("states", r["s"])[:MAX_STATE]
        S[b, :len(x)] = x
        sm[b, :len(x)] = True
        sv[b] = st.vec["states"][r["s"]]
        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, j] = True
            ov[b, j] = st.vec["options"][i]
    t = torch.tensor([r["t"] for r in recs])
    return (torch.from_numpy(S), torch.from_numpy(sm), torch.from_numpy(Q), torch.from_numpy(qm), torch.from_numpy(O),
            torch.from_numpy(om), torch.from_numpy(ok), t, torch.from_numpy(sv), torch.from_numpy(ov)), torch.tensor(golds)


def loss_fn(logits, y, qtype):
    ce = F.cross_entropy(logits, y, reduction="none")
    is_score = qtype == QTYPES["score"]
    if is_score.any():
        p = torch.softmax(logits[is_score], -1)
        tgt = F.one_hot(y[is_score], logits.shape[1]).float()
        k = torch.isfinite(logits[is_score]).sum(-1).clamp(min=2).float()
        rps = ((p.cumsum(-1) - tgt.cumsum(-1)) ** 2).sum(-1) / (k - 1)
        ce = ce.clone()
        ce[is_score] = ce[is_score] + rps
    return ce.mean()


def shape_key(st, r):
    """(option count, longest option, state length): batching on this keeps padding small."""
    oo = st.off["options"]
    k = min(len(r["o"]), MAX_OPTS)
    lo = max(oo[i + 1] - oo[i] for i in r["o"])
    so = st.off["states"]
    return (k, (lo + 3) // 4, so[r["s"] + 1] - so[r["s"]])


def batch_cost(st, batch):
    """Padded tokens the fused layers process for this batch (what step time scales with)."""
    so, qo, oo = st.off["states"], st.off["questions"], st.off["options"]
    n = min(MAX_STATE, max(so[r["s"] + 1] - so[r["s"]] for r in batch))
    lq = max(qo[r["q"] + 1] - qo[r["q"]] for r in batch)
    k = max(min(len(r["o"]), MAX_OPTS) for r in batch)
    lo = max(oo[i + 1] - oo[i] for r in batch for i in r["o"])
    return len(batch) * (n + lq + k * lo)


def batches_by_length(st, recs, bs, rng):
    recs = list(recs)
    rng.shuffle(recs)
    out = []
    chunk = bs * 64
    for s in range(0, len(recs), chunk):
        part = sorted(recs[s:s + chunk], key=lambda r: shape_key(st, r))
        out += [part[i:i + bs] for i in range(0, len(part), bs)]
    if WORLD > 1:
        # pair batches of similar cost so the ranks finish each step together (DDP waits for the slowest)
        out.sort(key=lambda b: batch_cost(st, b))
        groups = [out[i:i + WORLD] for i in range(0, len(out) - WORLD + 1, WORLD)]
        rng.shuffle(groups)
        return [b for g in groups for b in g]
    rng.shuffle(out)
    return out


@torch.no_grad()
def collect_logits(head, st, recs, bs=128):
    head.eval()
    rng = random.Random(0)
    res = []
    recs = sorted(recs, key=lambda r: shape_key(st, r))
    for i in range(0, len(recs), bs):
        part = recs[i:i + bs]
        if any(len(r["o"]) > MAX_OPTS for r in part):
            for r in part:
                (xb, yb) = make_batch(st, [r], rng, False) if len(r["o"]) <= MAX_OPTS else (None, None)
                if xb is None:
                    continue
                lg = head(*xb)
                res.append((r, lg[0].numpy(), int(yb[0])))
            continue
        xb, yb = make_batch(st, part, rng, False)
        lg = head(*xb).numpy()
        for r, z, y in zip(part, lg, yb.tolist()):
            res.append((r, z, y))
    head.train()
    return res


def metrics(rows, temps=None):
    by = defaultdict(lambda: [0, 0, 0.0])
    tot = [0, 0, 0.0]
    for r, z, y in rows:
        z = z[np.isfinite(z)]
        T = (temps or {}).get(temp_bucket(QNAMES[r["t"]], len(z)), 1.0)
        zz = z / T
        p = np.exp(zz - zz.max())
        p /= p.sum()
        c = int(np.argmax(p) == y)
        nll = -math.log(max(p[y], 1e-12))
        for acc in (by[r["src"]], tot):
            acc[0] += 1
            acc[1] += c
            acc[2] += nll
    out = {"acc": tot[1] / max(1, tot[0]), "nll": tot[2] / max(1, tot[0]), "n": tot[0]}
    out["by_src"] = {k: round(v[1] / v[0], 4) for k, v in sorted(by.items())}
    out["macro_acc"] = float(np.mean(list(out["by_src"].values()))) if by else 0.0
    return out


def fit_temperatures(rows):
    buckets = defaultdict(list)
    for r, z, y in rows:
        z = z[np.isfinite(z)]
        buckets[temp_bucket(QNAMES[r["t"]], len(z))].append((z, y))
    temps = {}
    grid = np.exp(np.linspace(math.log(0.2), math.log(5.0), 161))
    for b, items in sorted(buckets.items()):
        if len(items) < 30:
            continue
        best = (1e9, 1.0)
        for T in grid:
            nll = 0.0
            for z, y in items:
                zz = z / T
                m = zz.max()
                nll += -(zz[y] - m - math.log(np.exp(zz - m).sum()))
            best = min(best, (nll / len(items), float(T)))
        temps[b] = round(best[1], 4)
    return temps


def main():
    torch.manual_seed(SEED)
    rng = random.Random(SEED)
    torch.set_num_threads(int(os.environ.get("NIRNAY_THREADS", "16")))
    os.makedirs(RUND, exist_ok=True)
    st = Store(CACHE)
    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)
    meta = st.meta
    enc = load_hf_encoder(meta["backbone"], meta["pool"], meta["prefix"])
    d, heads, ff = enc.cfg.hidden, enc.cfg.heads, enc.cfg.ff
    adapter = None
    if FUSION == "deep":
        assert SPLIT and meta.get("mid_layer") == SPLIT, "deep fusion needs NIRNAY_SPLIT = cached mid layer"
        head = DeepFusionHead(list(enc.layers)[SPLIT:], d, pool=meta["pool"])
        pre = [p for n, p in head.named_parameters() if n.startswith("layers.")]
        new = [p for n, p in head.named_parameters() if not n.startswith("layers.")]
        groups = [{"params": new, "lr": LR}, {"params": pre, "lr": ADAPTER_LR}]
    else:
        head = DecisionHead(d, heads, ff, BLOCKS)
        head.init_from_encoder(enc.layers)
        if SPLIT:
            assert meta.get("mid_layer") == SPLIT, "cache holds layer %s states, not %d" % (meta.get("mid_layer"), SPLIT)
            adapter = StateAdapter(list(enc.layers)[SPLIT:], enc.cfg.ln_eps)
        groups = [{"params": list(head.parameters()), "lr": LR}]
        if adapter is not None:
            groups.append({"params": list(adapter.parameters()), "lr": ADAPTER_LR})
    model = Model(head, adapter, meta["pool"])
    if os.environ.get("NIRNAY_INIT"):   # warm start (e.g. a second epoch) from a saved model_best.pt
        model.load_state_dict(torch.load(os.environ["NIRNAY_INIT"], weights_only=True))
        print("warm start from", os.environ["NIRNAY_INIT"], flush=True)
    n_params = sum(p.numel() for p in model.parameters())
    print("trainable params %.2fM, fusion %s, blocks %d, split %d, world %d" % (n_params / 1e6, FUSION, BLOCKS, SPLIT, WORLD), flush=True)
    opt = torch.optim.AdamW(groups, weight_decay=0.01, betas=(0.9, 0.98))
    ddp = model
    if WORLD > 1:
        import torch.distributed as dist
        dist.init_process_group("gloo", init_method=os.environ["NIRNAY_DDP_INIT"], rank=RANK, world_size=WORLD)
        ddp = torch.nn.parallel.DistributedDataParallel(model)
    steps_per_epoch = math.ceil(len(tr) / BATCH / WORLD)
    total = int(steps_per_epoch * EPOCHS)
    sched = torch.optim.lr_scheduler.LambdaLR(
        opt, lambda s: min(1.0, (s + 1) / WARMUP) * 0.5 * (1 + math.cos(math.pi * min(1.0, s / max(1, total)))))
    log = {"run": RUN, "tag": TAG, "fusion": FUSION, "blocks": BLOCKS, "batch": BATCH, "world": WORLD, "lr": LR,
           "adapter_lr": ADAPTER_LR, "epochs": EPOCHS, "evals": []}
    best = (-1.0, None)
    step, t0, run_loss = 0, time.time(), 0.0
    model.train()
    ep = 0

    def evaluate():
        m = metrics(collect_logits(model, st, 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" % (m["acc"], m["macro_acc"], m["nll"]), flush=True)
        json.dump(log, open(os.path.join(RUND, "log.json"), "w"), indent=2)
        return m

    if os.environ.get("NIRNAY_EXPORT_ONLY") == "1":   # calibrate + export a saved checkpoint (NIRNAY_INIT), no training
        assert WORLD == 1 and os.environ.get("NIRNAY_INIT"), "export-only needs NIRNAY_INIT and a single process"
        best = (-1.0, None)
        m = evaluate()
        best = (m["macro_acc"], m)
        torch.save(model.state_dict(), os.path.join(RUND, "model_best.pt"))
        total = 0
    while step < total:
        # every rank builds the identical batch list (same seed) and takes its own slice; equal counts keep
        # the all-reduce in lockstep
        bl = batches_by_length(st, tr, BATCH, rng)
        bl = bl[:len(bl) // WORLD * WORLD]
        for batch in bl[RANK::WORLD]:
            xb, yb = make_batch(st, batch, rng, True)
            logits = ddp(*xb)
            loss = loss_fn(logits, yb, xb[7])
            opt.zero_grad(set_to_none=True)
            loss.backward()
            torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
            opt.step()
            sched.step()
            step += 1
            run_loss = 0.98 * run_loss + 0.02 * float(loss) if step > 1 else float(loss)
            if step % 50 == 0 and RANK == 0:
                el = time.time() - t0
                print("\r step %d/%d  loss %.4f  %.2f s/step  ETA %.0f min   " %
                      (step, total, run_loss, el / step, el / step * (total - step) / 60), end="", flush=True)
            if step % EVAL_EVERY == 0 or step >= total:
                if RANK == 0:
                    m = evaluate()
                    if m["macro_acc"] > best[0]:
                        best = (m["macro_acc"], m)
                        torch.save(model.state_dict(), os.path.join(RUND, "model_best.pt"))
                if WORLD > 1:
                    torch.distributed.barrier()
            if step >= total:
                break
        ep += 1
    if RANK != 0:
        torch.distributed.barrier()
        return
    print("\nbest dev macro acc %.4f" % best[0], flush=True)
    model.load_state_dict(torch.load(os.path.join(RUND, "model_best.pt"), weights_only=True))
    cal_rows = collect_logits(model, st, cal)
    temps = fit_temperatures(cal_rows)
    print("temperatures", temps, flush=True)
    m_raw, m_cal = metrics(cal_rows), 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
    json.dump(log, open(os.path.join(RUND, "log.json"), "w"), indent=2)
    export(model, enc, meta, temps, log)
    if WORLD > 1:
        torch.distributed.barrier()


def export(model, enc, meta, temps, log):
    from safetensors.torch import save_file
    out = os.environ.get("NIRNAY_EXPORT", os.path.join(ROOT, "models", "nirnay-" + RUN))
    os.makedirs(out, exist_ok=True)
    save_file({k: v.contiguous() for k, v in enc.state_dict().items()}, os.path.join(out, "encoder.safetensors"))
    save_file({k: v.contiguous() for k, v in model.head.state_dict().items()}, os.path.join(out, "head.safetensors"))
    if model.adapter is not None:
        save_file({k: v.contiguous() for k, v in model.adapter.state_dict().items()}, os.path.join(out, "adapter.safetensors"))
    shutil.copy(os.path.join(meta["source_dir"], "tokenizer.json"), os.path.join(out, "tokenizer.json"))
    cfg = {"name": "nirnay-" + RUN, "version": "1.0.0", "backbone": meta["backbone"], "encoder": enc.cfg.to_dict(),
           "head": {"blocks": BLOCKS, "heads": enc.cfg.heads, "ff": enc.cfg.ff}, "split_layer": SPLIT, "fusion": FUSION,
           "special_ids": meta["special_ids"],
           "limits": {"state_tokens": 512, "question_tokens": meta["limits"]["question_tokens"],
                      "option_tokens": meta["limits"]["option_tokens"], "max_fused_options": 128, "window_overlap": 64},
           "threads": 8, "temperature_by_bucket": temps,
           "training": {"records": meta["n_records"], "dev_macro_acc": round(log["best_dev"]["macro_acc"], 4),
                        "dev_acc": round(log["best_dev"]["acc"], 4), "step": log["best_dev"]["step"]}}
    json.dump(cfg, open(os.path.join(out, "nirnay_config.json"), "w", encoding="utf-8"), indent=2)
    print("exported ->", out, flush=True)


if __name__ == "__main__":
    main()

	"""Train the Nirnay fusion head on precomputed frozen-encoder states, calibrate, export.

	Objective: log score (cross-entropy) on every question + ranked probability score on ordinal
	(`score`) questions. Both are strictly proper scoring rules, the same family Laya's RLCD
	rewards with REINFORCE on noisy logits; here the gradient of the rule is taken exactly.

	python train\\train_head.py (NIRNAY_TAG=e5s by default)
	"""
	import json
	import math
	import os
	import random
	import shutil
	import sys
	import time
	from collections import defaultdict

	import numpy as np
	import torch
	import torch.nn.functional as F

	HERE = os.path.dirname(os.path.abspath(__file__))
	ROOT = os.path.dirname(HERE)
	sys.path.insert(0, ROOT)

	from nirnay.encoder import load_hf_encoder # noqa: E402
	from nirnay.engine import temp_bucket # noqa: E402
	from nirnay.model import QTYPES, DecisionHead, DeepFusionHead, StateAdapter # noqa: E402

	TAG = os.environ.get("NIRNAY_TAG", "e5s")
	CACHE = os.path.join(HERE, "cache", TAG)
	RUN = os.environ.get("NIRNAY_RUN", TAG + "-b%s" % os.environ.get("NIRNAY_BLOCKS", "2"))
	RUND = os.path.join(HERE, "runs", RUN)
	BLOCKS = int(os.environ.get("NIRNAY_BLOCKS", "2"))
	BATCH = int(os.environ.get("NIRNAY_BATCH", "64"))
	EPOCHS = float(os.environ.get("NIRNAY_EPOCHS", "2"))
	LR = float(os.environ.get("NIRNAY_LR", "3e-4"))
	WARMUP = int(os.environ.get("NIRNAY_WARMUP", "300"))
	MAX_OPTS = int(os.environ.get("NIRNAY_MAX_OPTS", "64"))
	MAX_STATE = int(os.environ.get("NIRNAY_MAX_STATE", "320")) # training-time cap on state tokens
	EVAL_EVERY = int(os.environ.get("NIRNAY_EVAL_EVERY", "600"))
	SPLIT = int(os.environ.get("NIRNAY_SPLIT", "0")) # >0: train encoder layers SPLIT+1.. on the state
	ADAPTER_LR = float(os.environ.get("NIRNAY_ADAPTER_LR", "5e-5"))
	FUSION = os.environ.get("NIRNAY_FUSION", "stack") # "stack": frozen encoder + fusion blocks; "deep": top layers fused
	WORLD = int(os.environ.get("WORLD_SIZE", "1")) # >1: one process per CPU socket, gloo all-reduce
	RANK = int(os.environ.get("RANK", "0"))
	SEED = 1234
	QNAMES = {v: k for k, v in QTYPES.items()}


	class Model(torch.nn.Module):
	"""Decision head, optionally preceded by a trainable state adapter (split encoder)."""

	def __init__(self, head, adapter=None, pool="cls"):
	super().__init__()
	self.head, self.adapter, self.pool_kind = head, adapter, pool

	def forward(self, S, sm, Q, qm, O, om, ok, t, sv, ov):
	if self.adapter is not None:
	S = self.adapter(S, sm)
	if self.pool_kind == "cls":
	v = S[:, 0]
	else:
	m = sm.unsqueeze(-1).to(S.dtype)
	v = (S * m).sum(1) / m.sum(1).clamp_min(1.0)
	sv = F.normalize(v, dim=-1)
	return self.head(S, sm, Q, qm, O, om, ok, t, sv, ov)


	class Store:
	def __init__(self, d):
	self.meta = json.load(open(os.path.join(d, "meta.json")))
	self.arr, self.off, self.vec = {}, {}, {}
	for n in ("states", "questions", "options"):
	mid = (FUSION == "deep") or (SPLIT and n == "states")
	f = n + ("_mid" if mid else "") + ".f16.npy"
	self.arr[n] = np.load(os.path.join(d, f), mmap_mode="r")
	self.off[n] = np.load(os.path.join(d, n + "_off.npy"))
	self.vec[n] = np.load(os.path.join(d, n + "_vec.npy"))
	# small tables fully in RAM
	for n in ("questions", "options"):
	self.arr[n] = np.asarray(self.arr[n])
	self.index = [json.loads(l) for l in open(os.path.join(d, "index.jsonl"), encoding="utf-8")]

	def rows(self, n, i):
	return self.arr[n][self.off[n][i]:self.off[n][i + 1]]


	def make_batch(st: Store, recs, rng, train: bool):
	d = st.meta["d"]
	B = len(recs)
	opts, golds = [], []
	for r in recs:
	o, y = list(r["o"]), r["y"]
	if len(o) > MAX_OPTS:
	keep = [y] + rng.sample([i for i in range(len(o)) if i != y], MAX_OPTS - 1)
	keep.sort()
	y = keep.index(y)
	o = [o[i] for i in keep]
	opts.append(o)
	golds.append(y)
	n = min(MAX_STATE, max(st.off["states"][r["s"] + 1] - st.off["states"][r["s"]] for r in recs))
	Lq = max(st.off["questions"][r["q"] + 1] - st.off["questions"][r["q"]] for r in recs)
	K = max(len(o) for o in opts)
	Lo = max(st.off["options"][i + 1] - st.off["options"][i] for o in opts for i in o)
	S = np.zeros((B, n, d), np.float32)
	sm = np.zeros((B, n), bool)
	Q = np.zeros((B, Lq, d), np.float32)
	qm = np.zeros((B, Lq), bool)
	O = np.zeros((B, K, Lo, d), np.float32)
	om = np.zeros((B, K, Lo), bool)
	ok = np.zeros((B, K), bool)
	sv = np.zeros((B, d), np.float32)
	ov = np.zeros((B, K, d), np.float32)
	for b, (r, o) in enumerate(zip(recs, opts)):
	x = st.rows("states", r["s"])[:MAX_STATE]
	S[b, :len(x)] = x
	sm[b, :len(x)] = True
	sv[b] = st.vec["states"][r["s"]]
	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, j] = True
	ov[b, j] = st.vec["options"][i]
	t = torch.tensor([r["t"] for r in recs])
	return (torch.from_numpy(S), torch.from_numpy(sm), torch.from_numpy(Q), torch.from_numpy(qm), torch.from_numpy(O),
	torch.from_numpy(om), torch.from_numpy(ok), t, torch.from_numpy(sv), torch.from_numpy(ov)), torch.tensor(golds)


	def loss_fn(logits, y, qtype):
	ce = F.cross_entropy(logits, y, reduction="none")
	is_score = qtype == QTYPES["score"]
	if is_score.any():
	p = torch.softmax(logits[is_score], -1)
	tgt = F.one_hot(y[is_score], logits.shape[1]).float()
	k = torch.isfinite(logits[is_score]).sum(-1).clamp(min=2).float()
	rps = ((p.cumsum(-1) - tgt.cumsum(-1)) ** 2).sum(-1) / (k - 1)
	ce = ce.clone()
	ce[is_score] = ce[is_score] + rps
	return ce.mean()


	def shape_key(st, r):
	"""(option count, longest option, state length): batching on this keeps padding small."""
	oo = st.off["options"]
	k = min(len(r["o"]), MAX_OPTS)
	lo = max(oo[i + 1] - oo[i] for i in r["o"])
	so = st.off["states"]
	return (k, (lo + 3) // 4, so[r["s"] + 1] - so[r["s"]])


	def batch_cost(st, batch):
	"""Padded tokens the fused layers process for this batch (what step time scales with)."""
	so, qo, oo = st.off["states"], st.off["questions"], st.off["options"]
	n = min(MAX_STATE, max(so[r["s"] + 1] - so[r["s"]] for r in batch))
	lq = max(qo[r["q"] + 1] - qo[r["q"]] for r in batch)
	k = max(min(len(r["o"]), MAX_OPTS) for r in batch)
	lo = max(oo[i + 1] - oo[i] for r in batch for i in r["o"])
	return len(batch) * (n + lq + k * lo)


	def batches_by_length(st, recs, bs, rng):
	recs = list(recs)
	rng.shuffle(recs)
	out = []
	chunk = bs * 64
	for s in range(0, len(recs), chunk):
	part = sorted(recs[s:s + chunk], key=lambda r: shape_key(st, r))
	out += [part[i:i + bs] for i in range(0, len(part), bs)]
	if WORLD > 1:
	# pair batches of similar cost so the ranks finish each step together (DDP waits for the slowest)
	out.sort(key=lambda b: batch_cost(st, b))
	groups = [out[i:i + WORLD] for i in range(0, len(out) - WORLD + 1, WORLD)]
	rng.shuffle(groups)
	return [b for g in groups for b in g]
	rng.shuffle(out)
	return out


	@torch.no_grad()
	def collect_logits(head, st, recs, bs=128):
	head.eval()
	rng = random.Random(0)
	res = []
	recs = sorted(recs, key=lambda r: shape_key(st, r))
	for i in range(0, len(recs), bs):
	part = recs[i:i + bs]
	if any(len(r["o"]) > MAX_OPTS for r in part):
	for r in part:
	(xb, yb) = make_batch(st, [r], rng, False) if len(r["o"]) <= MAX_OPTS else (None, None)
	if xb is None:
	continue
	lg = head(*xb)
	res.append((r, lg[0].numpy(), int(yb[0])))
	continue
	xb, yb = make_batch(st, part, rng, False)
	lg = head(*xb).numpy()
	for r, z, y in zip(part, lg, yb.tolist()):
	res.append((r, z, y))
	head.train()
	return res


	def metrics(rows, temps=None):
	by = defaultdict(lambda: [0, 0, 0.0])
	tot = [0, 0, 0.0]
	for r, z, y in rows:
	z = z[np.isfinite(z)]
	T = (temps or {}).get(temp_bucket(QNAMES[r["t"]], len(z)), 1.0)
	zz = z / T
	p = np.exp(zz - zz.max())
	p /= p.sum()
	c = int(np.argmax(p) == y)
	nll = -math.log(max(p[y], 1e-12))
	for acc in (by[r["src"]], tot):
	acc[0] += 1
	acc[1] += c
	acc[2] += nll
	out = {"acc": tot[1] / max(1, tot[0]), "nll": tot[2] / max(1, tot[0]), "n": tot[0]}
	out["by_src"] = {k: round(v[1] / v[0], 4) for k, v in sorted(by.items())}
	out["macro_acc"] = float(np.mean(list(out["by_src"].values()))) if by else 0.0
	return out


	def fit_temperatures(rows):
	buckets = defaultdict(list)
	for r, z, y in rows:
	z = z[np.isfinite(z)]
	buckets[temp_bucket(QNAMES[r["t"]], len(z))].append((z, y))
	temps = {}
	grid = np.exp(np.linspace(math.log(0.2), math.log(5.0), 161))
	for b, items in sorted(buckets.items()):
	if len(items) < 30:
	continue
	best = (1e9, 1.0)
	for T in grid:
	nll = 0.0
	for z, y in items:
	zz = z / T
	m = zz.max()
	nll += -(zz[y] - m - math.log(np.exp(zz - m).sum()))
	best = min(best, (nll / len(items), float(T)))
	temps[b] = round(best[1], 4)
	return temps


	def main():
	torch.manual_seed(SEED)
	rng = random.Random(SEED)
	torch.set_num_threads(int(os.environ.get("NIRNAY_THREADS", "16")))
	os.makedirs(RUND, exist_ok=True)
	st = Store(CACHE)
	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)
	meta = st.meta
	enc = load_hf_encoder(meta["backbone"], meta["pool"], meta["prefix"])
	d, heads, ff = enc.cfg.hidden, enc.cfg.heads, enc.cfg.ff
	adapter = None
	if FUSION == "deep":
	assert SPLIT and meta.get("mid_layer") == SPLIT, "deep fusion needs NIRNAY_SPLIT = cached mid layer"
	head = DeepFusionHead(list(enc.layers)[SPLIT:], d, pool=meta["pool"])
	pre = [p for n, p in head.named_parameters() if n.startswith("layers.")]
	new = [p for n, p in head.named_parameters() if not n.startswith("layers.")]
	groups = [{"params": new, "lr": LR}, {"params": pre, "lr": ADAPTER_LR}]
	else:
	head = DecisionHead(d, heads, ff, BLOCKS)
	head.init_from_encoder(enc.layers)
	if SPLIT:
	assert meta.get("mid_layer") == SPLIT, "cache holds layer %s states, not %d" % (meta.get("mid_layer"), SPLIT)
	adapter = StateAdapter(list(enc.layers)[SPLIT:], enc.cfg.ln_eps)
	groups = [{"params": list(head.parameters()), "lr": LR}]
	if adapter is not None:
	groups.append({"params": list(adapter.parameters()), "lr": ADAPTER_LR})
	model = Model(head, adapter, meta["pool"])
	if os.environ.get("NIRNAY_INIT"): # warm start (e.g. a second epoch) from a saved model_best.pt
	model.load_state_dict(torch.load(os.environ["NIRNAY_INIT"], weights_only=True))
	print("warm start from", os.environ["NIRNAY_INIT"], flush=True)
	n_params = sum(p.numel() for p in model.parameters())
	print("trainable params %.2fM, fusion %s, blocks %d, split %d, world %d" % (n_params / 1e6, FUSION, BLOCKS, SPLIT, WORLD), flush=True)
	opt = torch.optim.AdamW(groups, weight_decay=0.01, betas=(0.9, 0.98))
	ddp = model
	if WORLD > 1:
	import torch.distributed as dist
	dist.init_process_group("gloo", init_method=os.environ["NIRNAY_DDP_INIT"], rank=RANK, world_size=WORLD)
	ddp = torch.nn.parallel.DistributedDataParallel(model)
	steps_per_epoch = math.ceil(len(tr) / BATCH / WORLD)
	total = int(steps_per_epoch * EPOCHS)
	sched = torch.optim.lr_scheduler.LambdaLR(
	opt, lambda s: min(1.0, (s + 1) / WARMUP) * 0.5 * (1 + math.cos(math.pi * min(1.0, s / max(1, total)))))
	log = {"run": RUN, "tag": TAG, "fusion": FUSION, "blocks": BLOCKS, "batch": BATCH, "world": WORLD, "lr": LR,
	"adapter_lr": ADAPTER_LR, "epochs": EPOCHS, "evals": []}
	best = (-1.0, None)
	step, t0, run_loss = 0, time.time(), 0.0
	model.train()
	ep = 0

	def evaluate():
	m = metrics(collect_logits(model, st, 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" % (m["acc"], m["macro_acc"], m["nll"]), flush=True)
	json.dump(log, open(os.path.join(RUND, "log.json"), "w"), indent=2)
	return m

	if os.environ.get("NIRNAY_EXPORT_ONLY") == "1": # calibrate + export a saved checkpoint (NIRNAY_INIT), no training
	assert WORLD == 1 and os.environ.get("NIRNAY_INIT"), "export-only needs NIRNAY_INIT and a single process"
	best = (-1.0, None)
	m = evaluate()
	best = (m["macro_acc"], m)
	torch.save(model.state_dict(), os.path.join(RUND, "model_best.pt"))
	total = 0
	while step < total:
	# every rank builds the identical batch list (same seed) and takes its own slice; equal counts keep
	# the all-reduce in lockstep
	bl = batches_by_length(st, tr, BATCH, rng)
	bl = bl[:len(bl) // WORLD * WORLD]
	for batch in bl[RANK::WORLD]:
	xb, yb = make_batch(st, batch, rng, True)
	logits = ddp(*xb)
	loss = loss_fn(logits, yb, xb[7])
	opt.zero_grad(set_to_none=True)
	loss.backward()
	torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
	opt.step()
	sched.step()
	step += 1
	run_loss = 0.98 * run_loss + 0.02 * float(loss) if step > 1 else float(loss)
	if step % 50 == 0 and RANK == 0:
	el = time.time() - t0
	print("\r step %d/%d loss %.4f %.2f s/step ETA %.0f min " %
	(step, total, run_loss, el / step, el / step * (total - step) / 60), end="", flush=True)
	if step % EVAL_EVERY == 0 or step >= total:
	if RANK == 0:
	m = evaluate()
	if m["macro_acc"] > best[0]:
	best = (m["macro_acc"], m)
	torch.save(model.state_dict(), os.path.join(RUND, "model_best.pt"))
	if WORLD > 1:
	torch.distributed.barrier()
	if step >= total:
	break
	ep += 1
	if RANK != 0:
	torch.distributed.barrier()
	return
	print("\nbest dev macro acc %.4f" % best[0], flush=True)
	model.load_state_dict(torch.load(os.path.join(RUND, "model_best.pt"), weights_only=True))
	cal_rows = collect_logits(model, st, cal)
	temps = fit_temperatures(cal_rows)
	print("temperatures", temps, flush=True)
	m_raw, m_cal = metrics(cal_rows), 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
	json.dump(log, open(os.path.join(RUND, "log.json"), "w"), indent=2)
	export(model, enc, meta, temps, log)
	if WORLD > 1:
	torch.distributed.barrier()


	def export(model, enc, meta, temps, log):
	from safetensors.torch import save_file
	out = os.environ.get("NIRNAY_EXPORT", os.path.join(ROOT, "models", "nirnay-" + RUN))
	os.makedirs(out, exist_ok=True)
	save_file({k: v.contiguous() for k, v in enc.state_dict().items()}, os.path.join(out, "encoder.safetensors"))
	save_file({k: v.contiguous() for k, v in model.head.state_dict().items()}, os.path.join(out, "head.safetensors"))
	if model.adapter is not None:
	save_file({k: v.contiguous() for k, v in model.adapter.state_dict().items()}, os.path.join(out, "adapter.safetensors"))
	shutil.copy(os.path.join(meta["source_dir"], "tokenizer.json"), os.path.join(out, "tokenizer.json"))
	cfg = {"name": "nirnay-" + RUN, "version": "1.0.0", "backbone": meta["backbone"], "encoder": enc.cfg.to_dict(),
	"head": {"blocks": BLOCKS, "heads": enc.cfg.heads, "ff": enc.cfg.ff}, "split_layer": SPLIT, "fusion": FUSION,
	"special_ids": meta["special_ids"],
	"limits": {"state_tokens": 512, "question_tokens": meta["limits"]["question_tokens"],
	"option_tokens": meta["limits"]["option_tokens"], "max_fused_options": 128, "window_overlap": 64},
	"threads": 8, "temperature_by_bucket": temps,
	"training": {"records": meta["n_records"], "dev_macro_acc": round(log["best_dev"]["macro_acc"], 4),
	"dev_acc": round(log["best_dev"]["acc"], 4), "step": log["best_dev"]["step"]}}
	json.dump(cfg, open(os.path.join(out, "nirnay_config.json"), "w", encoding="utf-8"), indent=2)
	print("exported ->", out, flush=True)


	if __name__ == "__main__":
	main()