| """Supervised / distillation training loop. |
| |
| Resumable (last.pt holds model, optimizer, scheduler, scaler, data position and |
| RNG), stops cleanly on SIGINT/SIGTERM or a time budget, and ends with |
| temperature calibration on the validation split and a test-set report. |
| """ |
| from __future__ import annotations |
|
|
| import math |
| import os |
| import random |
| import signal |
| import sys |
| import time |
| from pathlib import Path |
|
|
| import numpy as np |
| import torch |
|
|
| from minijev.calibration.metrics import (IsotonicConfidence, calibration_report, fit_temperature, |
| reliability_svg, softmax) |
| from minijev.datasets.batching import BackboneTable, BatchEncoder, EpochBatches, model_inputs, to_torch |
| from minijev.datasets.decision_data import DecisionData |
| from minijev.device import RuntimeProfile, autocast, resolve_runtime |
| from minijev.model import build_model, human_params |
| from minijev.runs import RunLogger, environment_info, new_run_dir |
| from minijev.tokenizers import make_tokenizer |
| from minijev.tasks import get_task |
|
|
| from .losses import decision_loss |
|
|
|
|
| def seed_everything(seed: int) -> None: |
| random.seed(seed) |
| np.random.seed(seed) |
| torch.manual_seed(seed) |
|
|
|
|
| def task_specs(cfg: dict) -> list[dict]: |
| """Single task (`task:` + `data:`) or multi-task (`tasks: [{key, task, data, weight}]`).""" |
| if cfg.get("tasks"): |
| sampling = cfg.get("sampling") or {} |
| out = [] |
| for t in cfg["tasks"]: |
| t = dict(t) |
| t.setdefault("key", t["task"]["name"]) |
| t["weight"] = float(sampling.get(t["key"], t.get("weight", 1.0))) |
| out.append(t) |
| return out |
| return [{"key": cfg["task"]["name"], "task": cfg["task"], "data": cfg["data"], "weight": 1.0}] |
|
|
|
|
| def make_task_for(cfg: dict, task_cfg: dict): |
| """Task plugin with the model's text tokenizer attached (byte or pretrained).""" |
| task = get_task(task_cfg["name"], task_cfg) |
| task.tokenizer = make_tokenizer(cfg["model"].get("text_encoder")) |
| return task |
|
|
|
|
| def load_splits(cfg: dict) -> dict[str, DecisionData]: |
| dcfg = cfg["data"] |
| d = Path(dcfg["dir"]) |
| if not (d / "train.npz").exists(): |
| gen = dcfg.get("generate") |
| if not gen: |
| raise FileNotFoundError(f"{d}/train.npz missing and no data.generate section") |
| import importlib |
| mod = importlib.import_module(dcfg.get("generator") or f"environments.{cfg['task']['name']}.dataset") |
| gen = dict(gen) |
| n = gen.pop("samples") |
| print(f"[data] generating {n:,} samples into {d} ...") |
| mod.generate_splits(d, n, **gen) |
| splits = {k: DecisionData.load(d / f"{k}.npz") for k in ("train", "val", "test")} |
| lim = dcfg.get("train_limit") |
| if lim and lim < len(splits["train"]): |
| splits["train"] = splits["train"].subset(np.arange(lim)) |
| return splits |
|
|
|
|
| def make_model(cfg: dict, task=None): |
| """task: one task, a list of tasks (multi-task: one state encoder per modality), or None (from cfg).""" |
| tasks = task if isinstance(task, (list, tuple)) else [task] if task is not None else \ |
| [make_task_for(cfg, t["task"]) for t in task_specs(cfg)] |
| encs, classes = {}, {} |
| for t in tasks: |
| encs.setdefault(t.state_encoder, t.encoder_config()) |
| if t.encoder_class(): |
| classes[t.state_encoder] = t.encoder_class() |
| return build_model(cfg["model"], encs, classes or None) |
|
|
|
|
| def param_groups(model, wd: float, lr: float | None = None, backbone_lr: float | None = None): |
| """AdamW groups: no decay on norms/biases/embeddings; a pretrained backbone can use its own (lower) lr.""" |
| groups = {} |
| for n, p in model.named_parameters(): |
| if not p.requires_grad: |
| continue |
| nd = p.ndim < 2 or "embed" in n or n.endswith((".pos", ".row", ".col", "cls", "type_emb")) |
| bb = backbone_lr is not None and n.startswith("backbone.") |
| groups.setdefault((nd, bb), []).append(p) |
| out = [] |
| for (nd, bb), ps in groups.items(): |
| g = {"params": ps, "weight_decay": 0.0 if nd else wd} |
| if bb: |
| g["lr"] = backbone_lr |
| elif lr is not None: |
| g["lr"] = lr |
| out.append(g) |
| return out |
|
|
|
|
| def lr_lambda(warmup: int, total: int, min_ratio: float): |
| def f(step): |
| if step < warmup: |
| return (step + 1) / max(1, warmup) |
| prog = min(1.0, (step - warmup) / max(1, total - warmup)) |
| return min_ratio + (1 - min_ratio) * 0.5 * (1 + math.cos(math.pi * prog)) |
| return f |
|
|
|
|
| @torch.no_grad() |
| def predict_logits(model, encoder: BatchEncoder, device, batch_size: int = 1024, max_rows: int | None = None, |
| profile: RuntimeProfile | None = None) -> tuple[np.ndarray, np.ndarray]: |
| model.eval() |
| n = len(encoder.data) if max_rows is None else min(max_rows, len(encoder.data)) |
| outs, tg = [], [] |
| for i in range(0, n, batch_size): |
| b = to_torch(encoder(np.arange(i, min(n, i + batch_size))), device) |
| with autocast(profile) if profile else torch.autocast("cpu", enabled=False): |
| out = model(**model_inputs(b)) |
| outs.append(out["logits"].float().cpu().numpy()) |
| tg.append(b["target"].cpu().numpy()) |
| return np.concatenate(outs), np.concatenate(tg) |
|
|
|
|
| def _acceptable(encoder: BatchEncoder, n: int): |
| """Rows' acceptable choices (soft label > 0), if the dataset has soft labels.""" |
| soft = encoder.data.soft |
| return None if soft is None else soft[:n] > 0 |
|
|
|
|
| class _Stop: |
| def __init__(self): |
| self.flag = False |
| self._prev = {} |
|
|
| def __enter__(self): |
| for s in (signal.SIGINT, signal.SIGTERM): |
| try: |
| self._prev[s] = signal.signal(s, self._handle) |
| except ValueError: |
| pass |
| return self |
|
|
| def _handle(self, *_): |
| if self.flag: |
| raise KeyboardInterrupt |
| print("\n[train] stop requested - saving checkpoint (Ctrl-C again to abort)") |
| self.flag = True |
|
|
| def __exit__(self, *a): |
| for s, h in self._prev.items(): |
| signal.signal(s, h) |
|
|
|
|
| class MultiTaskBatches: |
| """Batches for multi-task training: each batch comes from ONE task (its own state |
| encoder / choice set), picked with the configured sampling weights. Deterministic |
| in (seed, epoch, i), so resume skips exactly the consumed batches.""" |
|
|
| def __init__(self, encoders: dict, weights: dict, batch_size: int, n_batches: int, seed: int, epoch: int, |
| start_batch: int = 0): |
| self.keys = list(encoders) |
| w = np.array([weights[k] for k in self.keys], float) |
| self.p = w / w.sum() |
| self.encoders, self.bs, self.seed, self.epoch = encoders, batch_size, seed, epoch |
| self.range = range(start_batch, n_batches) |
|
|
| def __len__(self): |
| return len(self.range) |
|
|
| def __getitem__(self, i): |
| rng = np.random.default_rng((self.seed, self.epoch, self.range[i], 11)) |
| enc = self.encoders[self.keys[rng.choice(len(self.keys), p=self.p)]] |
| n = len(enc.data) |
| return enc(rng.choice(n, self.bs, replace=n < self.bs), rng) |
|
|
|
|
| def make_encoders(cfg: dict, model, task, split: DecisionData, choice_dropout: float = 0.0, seed: int = 0): |
| """BatchEncoder for one split; with a frozen pretrained backbone, questions and |
| choice labels are encoded once per vocabulary entry (BackboneTable).""" |
| mcfg = cfg["model"] |
| kw = dict(max_question_len=mcfg.get("max_question_len", 32), max_choice_len=mcfg.get("max_choice_len", 16)) |
| tables = {} |
| bb = getattr(model, "backbone", None) |
| if bb is not None and bb.frozen and cfg["training"].get("cache_text", True): |
| tok = task.tokenizer |
| tables = {"question_table": BackboneTable(split.questions, model, tok, kw["max_question_len"]), |
| "choice_table": BackboneTable([task.model_label(c) for c in split.choices], model, tok, |
| kw["max_choice_len"])} |
| return BatchEncoder(split, task, choice_dropout=choice_dropout, seed=seed, tokenizer=task.tokenizer, |
| **kw, **tables) |
|
|
|
|
| def train(cfg: dict, run_dir: str | Path | None = None, resume: bool = False, max_steps: int | None = None, |
| max_minutes: float | None = None, quiet: bool = False, final_eval: bool = True) -> dict: |
| tcfg = cfg["training"] |
| seed = int(cfg.get("seed", 0)) |
| seed_everything(seed) |
| specs = task_specs(cfg) |
| tasks = {s["key"]: make_task_for(cfg, s["task"]) for s in specs} |
| splits = {s["key"]: load_splits({**cfg, "task": s["task"], "data": s["data"]}) for s in specs} |
| weights = {s["key"]: s["weight"] for s in specs} |
| multi = len(specs) > 1 |
| model = make_model(cfg, list(tasks.values())) |
| if tcfg.get("init_checkpoint") and not resume: |
| ck = torch.load(tcfg["init_checkpoint"], map_location="cpu", weights_only=False) |
| model.load_state_dict(ck["model"]) |
| model.log_temperature.zero_() |
| print(f"[train] initialised from {tcfg['init_checkpoint']}") |
| n_params = model.num_parameters() |
| prof = resolve_runtime(tcfg, n_params) |
| torch.set_num_threads(prof.threads) |
| device = torch.device(prof.device) |
| model.to(device) |
|
|
| run_dir = Path(run_dir) if run_dir else new_run_dir(cfg.get("name", "run")) |
| ckpt_dir = run_dir / "checkpoints" |
| log = RunLogger(run_dir, cfg, resume=resume) |
| encs = {} |
| for k, sp in splits.items(): |
| encs[k] = {"train": make_encoders(cfg, model, tasks[k], sp["train"], tcfg.get("choice_dropout", 0.0), seed), |
| "val": make_encoders(cfg, model, tasks[k], sp["val"]), |
| "test": make_encoders(cfg, model, tasks[k], sp["test"])} |
| train_enc = {k: e["train"] for k, e in encs.items()} |
|
|
| eff_bs = int(tcfg.get("batch_size", 256)) |
| epochs = float(tcfg.get("epochs", 1)) |
| n_train = sum(len(sp["train"]) for sp in splits.values()) |
| batches_per_epoch = max(1, n_train // prof.micro_batch) |
| steps_per_epoch = max(1, batches_per_epoch // prof.grad_accum) |
| total_steps = max(1, int(steps_per_epoch * epochs)) |
| if max_steps: |
| total_steps = min(total_steps, max_steps) |
| lr = tcfg.get("lr", 1e-3) |
| opt = torch.optim.AdamW(param_groups(model, tcfg.get("weight_decay", 0.05), lr, tcfg.get("backbone_lr")), lr=lr, |
| betas=(0.9, 0.98), fused=prof.device == "cuda") |
| warmup = int(tcfg.get("warmup_steps", min(500, total_steps // 20))) |
| sched = torch.optim.lr_scheduler.LambdaLR(opt, lr_lambda(warmup, total_steps, tcfg.get("min_lr_ratio", 0.05))) |
| scaler = torch.amp.GradScaler(enabled=prof.amp_dtype == "fp16") |
| fwd = torch.compile(model) if prof.compile else model |
|
|
| step, epoch, batch_pos, best, elapsed_prev = 0, 0, 0, -1.0, 0.0 |
| samples_seen = 0 |
| session_samples = 0 |
| last = ckpt_dir / "last.pt" |
| if resume and last.exists(): |
| ck = torch.load(last, map_location=device, weights_only=False) |
| model.load_state_dict(ck["model"]) |
| opt.load_state_dict(ck["optimizer"]) |
| sched.load_state_dict(ck["scheduler"]) |
| scaler.load_state_dict(ck["scaler"]) |
| step, epoch, batch_pos, best = ck["step"], ck["epoch"], ck["batch_pos"], ck["best"] |
| elapsed_prev, samples_seen = ck.get("elapsed", 0.0), ck.get("samples_seen", 0) |
| torch.set_rng_state(ck["torch_rng"]) |
| print(f"[train] resumed from {last} at step {step} (epoch {epoch}, batch {batch_pos})") |
|
|
| breakdown = model.param_breakdown() |
| info = {"task": "+".join(tasks) if multi else next(iter(tasks.values())).name, "tasks": list(tasks), |
| "params": n_params, "params_human": human_params(n_params), "params_breakdown": breakdown, |
| "train_rows": n_train, "val_rows": sum(len(sp["val"]) for sp in splits.values()), |
| "test_rows": sum(len(sp["test"]) for sp in splits.values()), |
| "effective_batch": eff_bs, "total_steps": total_steps, "runtime": prof.to_dict(), |
| "env": environment_info()} |
| log.write_json("env.json", info) |
| if not quiet: |
| print(f"[train] {cfg.get('name')} | {human_params(n_params)} params (core+heads " |
| f"{human_params(breakdown['core_total'])}, text {human_params(breakdown['text'])}, trainable " |
| f"{human_params(breakdown['trainable'])}) | {prof.device} ({prof.device_name}) " |
| f"amp={prof.amp_dtype} micro={prof.micro_batch}x{prof.grad_accum} | " |
| f"{n_train:,} train rows ({', '.join(tasks)}) | {total_steps:,} steps") |
|
|
| def save(path, extra=None): |
| tmp = Path(str(path) + ".tmp") |
| torch.save({"model": model.state_dict(), "optimizer": opt.state_dict(), "scheduler": sched.state_dict(), |
| "scaler": scaler.state_dict(), "step": step, "epoch": epoch, "batch_pos": batch_pos, |
| "best": best, "elapsed": elapsed_prev + (time.time() - t_start), "samples_seen": samples_seen, |
| "torch_rng": torch.get_rng_state(), "config": cfg, **(extra or {})}, tmp) |
| os.replace(tmp, path) |
|
|
| def validate(): |
| reps = {} |
| for k, e in encs.items(): |
| lg, tg = predict_logits(model, e["val"], device, max_rows=eval_rows, profile=prof) |
| reps[k] = calibration_report(lg, tg, acceptable=_acceptable(e["val"], len(tg))) |
| w = np.array([weights[k] for k in reps]) |
| acc = float(np.dot(w / w.sum(), [r["accuracy"] for r in reps.values()])) |
| return acc, reps |
|
|
| loss_cfg = cfg.get("loss", {}) |
| eval_every = int(tcfg.get("eval_every", 1000)) |
| ckpt_every = int(tcfg.get("checkpoint_every", 1000)) |
| log_every = int(tcfg.get("log_every", 50)) |
| eval_rows = tcfg.get("eval_max_rows", 20000) |
| grad_clip = tcfg.get("grad_clip", 1.0) |
| t_start = time.time() |
| t_train = 0.0 |
| run_loss, run_n, window_samples, window_t = 0.0, 0, 0, time.time() |
| stopped_early = False |
|
|
| with _Stop() as stop: |
| while step < total_steps and not stop.flag: |
| if multi: |
| ds = MultiTaskBatches(train_enc, weights, prof.micro_batch, batches_per_epoch, seed, epoch, batch_pos) |
| else: |
| ds = EpochBatches(next(iter(train_enc.values())), prof.micro_batch, shuffle=True, seed=seed, |
| epoch=epoch, start_batch=batch_pos, drop_last=True) |
| loader = torch.utils.data.DataLoader(ds, batch_size=None, shuffle=False, num_workers=prof.num_workers, |
| pin_memory=prof.device == "cuda", |
| prefetch_factor=4 if prof.num_workers else None, |
| |
| multiprocessing_context="fork" if prof.num_workers and sys.platform == "linux" else None) |
| model.train() |
| micro = 0 |
| for batch in loader: |
| t0 = time.time() |
| b = to_torch(batch, device, non_blocking=True) |
| with autocast(prof): |
| out = fwd(**model_inputs(b)) |
| loss, terms = decision_loss(out["logits"], out["choice_mask"], b, loss_cfg, out.get("score")) |
| scaler.scale(loss / prof.grad_accum).backward() |
| micro += 1 |
| batch_pos += 1 |
| n_b = len(batch["target"]) |
| samples_seen += n_b |
| session_samples += n_b |
| window_samples += n_b |
| run_loss += float(loss.detach()) |
| run_n += 1 |
| if micro % prof.grad_accum: |
| t_train += time.time() - t0 |
| continue |
| if grad_clip: |
| scaler.unscale_(opt) |
| torch.nn.utils.clip_grad_norm_(model.parameters(), grad_clip) |
| scaler.step(opt) |
| scaler.update() |
| opt.zero_grad(set_to_none=True) |
| sched.step() |
| step += 1 |
| t_train += time.time() - t0 |
| if step % log_every == 0: |
| now = time.time() |
| sps = window_samples / max(1e-9, now - window_t) |
| log.log(step, train_loss=run_loss / run_n, lr=sched.get_last_lr()[0], |
| samples_per_sec=round(sps, 1), epoch=epoch, **{f"loss_{k}": v for k, v in terms.items()}) |
| if not quiet: |
| print(f" step {step:6d}/{total_steps} loss {run_loss / run_n:.4f} " |
| f"lr {sched.get_last_lr()[0]:.2e} {sps:,.0f} samples/s") |
| run_loss, run_n, window_samples, window_t = 0.0, 0, 0, now |
| if step % eval_every == 0 or step == total_steps: |
| acc, reps = validate() |
| first = next(iter(reps.values())) |
| extra = {"val_acceptable_acc": first["acceptable_accuracy"]} \ |
| if not multi and "acceptable_accuracy" in first else {} |
| per_task = {f"val_acc_{k}": r["accuracy"] for k, r in reps.items()} if multi else {} |
| log.log(step, val_acc=acc, val_nll=first["nll"], val_ece=first["ece"], |
| val_brier=first["brier"], **extra, **per_task) |
| if not quiet: |
| detail = " ".join(f"{k} {r['accuracy']:.4f}" for k, r in reps.items()) if multi else \ |
| f"nll {first['nll']:.4f} ece {first['ece']:.4f}" |
| print(f" [val] step {step} acc {acc:.4f} {detail}" |
| + (f" acceptable {first['acceptable_accuracy']:.4f}" if extra else "")) |
| if acc > best: |
| best = acc |
| torch.save({"model": model.state_dict(), "config": cfg, "step": step, |
| "val_accuracy": best}, ckpt_dir / "best.pt") |
| model.train() |
| if step % ckpt_every == 0: |
| save(last) |
| if max_minutes and (time.time() - t_start) / 60 > max_minutes: |
| print(f"[train] time budget of {max_minutes} min reached") |
| stopped_early = True |
| stop.flag = True |
| if step >= total_steps or stop.flag: |
| break |
| else: |
| epoch += 1 |
| batch_pos = 0 |
| continue |
| if batch_pos >= batches_per_epoch: |
| epoch, batch_pos = epoch + 1, 0 |
| save(last) |
| interrupted = stop.flag and not stopped_early and step < total_steps |
|
|
| train_seconds = elapsed_prev + (time.time() - t_start) |
| summary = {**{k: v for k, v in info.items() if k != "env"}, "steps_done": step, "samples_seen": samples_seen, |
| "train_seconds": round(train_seconds, 1), |
| "throughput_samples_per_sec": round(session_samples / max(t_train, 1e-9), 1), |
| "best_val_accuracy": best, "complete": step >= total_steps, "run_dir": str(run_dir)} |
| hourly = cfg.get("cost", {}).get("gpu_hourly_cost") or os.environ.get("MINIJEV_GPU_HOURLY_COST") |
| if hourly: |
| summary["cost"] = {"gpu_hourly_cost": float(hourly), |
| "estimated_cost": round(float(hourly) * train_seconds / 3600, 4)} |
| if final_eval and not interrupted: |
| summary.update(finalize(cfg, model, run_dir, {k: e["val"] for k, e in encs.items()}, |
| {k: e["test"] for k, e in encs.items()}, device, prof)) |
| log.write_json("summary.json", summary) |
| log.close() |
| return summary |
|
|
|
|
| def _pad_logits(parts: list[np.ndarray]) -> np.ndarray: |
| n = max(p.shape[1] for p in parts) |
| return np.concatenate([np.pad(p, ((0, 0), (0, n - p.shape[1])), constant_values=-1e9) for p in parts]) |
|
|
|
|
| def finalize(cfg, model, run_dir: Path, val_encs, test_encs, device, prof) -> dict: |
| """Best checkpoint -> temperature fitted on val (global + per task) -> calibrated test report -> final.pt.""" |
| from minijev.calibration.metrics import ece as ece_fn |
|
|
| if not isinstance(val_encs, dict): |
| val_encs, test_encs = {"task": val_encs}, {"task": test_encs} |
| ckpt_dir = run_dir / "checkpoints" |
| best = ckpt_dir / "best.pt" |
| if best.exists(): |
| model.load_state_dict(torch.load(best, map_location=device, weights_only=False)["model"]) |
| cal = cfg.get("calibration", {}) |
| per_task, V, VT, TL, TT = {}, [], [], [], [] |
| for k in val_encs: |
| vl, vt = predict_logits(model, val_encs[k], device, profile=prof) |
| tl, tt = predict_logits(model, test_encs[k], device, profile=prof) |
| Tk = fit_temperature(vl, vt) if cal.get("method", "temperature") == "temperature" else 1.0 |
| acc_ok = _acceptable(test_encs[k], len(tt)) |
| per_task[k] = {"temperature": Tk, "test_raw": calibration_report(tl, tt, 1.0, acceptable=acc_ok), |
| "test_calibrated": calibration_report(tl, tt, Tk, acceptable=acc_ok)} |
| V.append(vl), VT.append(vt), TL.append(tl), TT.append(tt) |
| val_lg, val_tg, test_lg, test_tg = _pad_logits(V), np.concatenate(VT), _pad_logits(TL), np.concatenate(TT) |
| T = fit_temperature(val_lg, val_tg) if cal.get("method", "temperature") == "temperature" else 1.0 |
| model.log_temperature.fill_(math.log(T)) |
| single = len(val_encs) == 1 |
| acc_ok = _acceptable(next(iter(test_encs.values())), len(test_tg)) if single else None |
| raw = calibration_report(test_lg, test_tg, 1.0, acceptable=acc_ok) |
| calibrated = calibration_report(test_lg, test_tg, T, acceptable=acc_ok) |
| |
| pv = softmax(val_lg, T) |
| iso = IsotonicConfidence().fit(pv.max(1), (pv.argmax(1) == val_tg).astype(float)) |
| pt = softmax(test_lg, T) |
| iso_ece = ece_fn(iso(pt.max(1)), (pt.argmax(1) == test_tg).astype(float)) |
| (run_dir / "reliability_raw.svg").write_text(reliability_svg(raw["reliability"], "Test - raw")) |
| (run_dir / "reliability_calibrated.svg").write_text( |
| reliability_svg(calibrated["reliability"], f"Test - temperature {T:.3f}")) |
| task_T = {k: v["temperature"] for k, v in per_task.items()} |
| torch.save({"model": model.state_dict(), "config": cfg, "temperature": T, "task_temperatures": task_T}, |
| ckpt_dir / "final.pt") |
| print(f"[calibration] T={T:.3f} | test acc {calibrated['accuracy']:.4f} | ECE {raw['ece']:.4f} -> " |
| f"{calibrated['ece']:.4f} | NLL {raw['nll']:.4f} -> {calibrated['nll']:.4f} | " |
| f"Brier {calibrated['brier']:.4f} | isotonic ECE {iso_ece:.4f}" |
| + (f" | acceptable acc {calibrated['acceptable_accuracy']:.4f}" if acc_ok is not None else "")) |
| if not single: |
| for k, v in per_task.items(): |
| c = v["test_calibrated"] |
| print(f" [{k}] T={v['temperature']:.3f} acc {c['accuracy']:.4f} ECE {v['test_raw']['ece']:.4f} -> " |
| f"{c['ece']:.4f} NLL {c['nll']:.4f}") |
| return {"temperature": T, "task_temperatures": task_T, "test_raw": raw, "test_calibrated": calibrated, |
| "test_isotonic_ece": iso_ece, "per_task": per_task, "final_checkpoint": str(ckpt_dir / "final.pt")} |
|
|
