code/train/train_head.py

"""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()