| """Training for Groot 1: proper-scoring objective with calibration built in. |
| |
| Objective per question (soft target t over k options, model probabilities p): |
| |
| L = w_soft * CE(t, p) + w_hard * CE(onehot(argmax t), p) + w_brier * ||p - t||^2 |
| + w_ord * E_p|i - argmax t| (score questions only: ordinal distance) |
| + w_conf * BCE(conf_logit, 1[argmax p == argmax t]) (correctness head) |
| |
| CE with a soft target and the Brier term are proper scoring rules, so the |
| model is rewarded for honest probabilities; this is the supervised core of |
| "reinforcement learning for calibrated decisions". Augmentation permutes |
| choice options (the model must not learn list positions) and, optionally, |
| replaces option names by neutral ids. |
| |
| Usage (single GPU): |
| groot-train --preset groot-1-s --train data/train.jsonl --val data/val.jsonl --out runs/groot-1-s |
| Multi-GPU: |
| torchrun --nproc_per_node 8 -m groot1.train --preset groot-1-s ... |
| """ |
| from __future__ import annotations |
|
|
| import argparse |
| import datetime |
| import json |
| import math |
| import mmap |
| import os |
| import random |
| import shutil |
| import sys |
| import time |
| from pathlib import Path |
| from typing import Any, Iterator, Sequence |
|
|
| import numpy as np |
| import torch |
| import torch.distributed as dist |
| import torch.nn.functional as F |
| from torch import nn |
|
|
| from .calibration import Record, evaluate_records, fit_temperatures |
| from .config import QTYPE_IDS, GrootConfig, preset |
| from .data import read_jsonl |
| from .model import Batch, GrootDecisionModel, collate |
| from .serialize import Serializer |
|
|
|
|
| |
| class RowStore: |
| """Memory-light random access to a JSONL file of decision rows. |
| |
| Only an int64 offset per line is kept in RAM; rows are parsed on demand from |
| an mmap of the file, so 4 DDP ranks x N loader workers share one page cache |
| instead of each holding a full Python copy of the dataset. Picklable (the |
| mmap is reopened lazily), so it works with fork and forkserver workers. |
| """ |
|
|
| def __init__(self, path: str | Path, limit: int | None = None): |
| self.path = str(path) |
| offsets = [0] |
| with open(self.path, "rb") as f: |
| for line in f: |
| offsets.append(offsets[-1] + len(line)) |
| if limit and len(offsets) > limit: |
| break |
| starts = np.array(offsets[:-1], dtype=np.int64) |
| ends = np.array(offsets[1:], dtype=np.int64) |
| keep = (ends - starts) > 1 |
| self.starts, self.ends = starts[keep], ends[keep] |
| if len(self.starts) == 0: |
| raise ValueError(f"{path}: no rows") |
| self._mm = None |
|
|
| def __len__(self) -> int: |
| return len(self.starts) |
|
|
| def estimate_lengths(self, early_state_tokens: int, chars_per_token: float = 3.2) -> tuple[np.ndarray, np.ndarray]: |
| """Approximate (question-sequence tokens, state tokens) per row, from characters. Cached on the store.""" |
| if getattr(self, "_qlen", None) is not None: |
| return self._qlen, self._slen |
| q = np.empty(len(self), dtype=np.int32) |
| st = np.empty(len(self), dtype=np.int32) |
| for i in range(len(self)): |
| r = self[i] |
| opts = r.get("options") or [] |
| qchars = len(r.get("question") or "") + sum(len(o) for o in opts) |
| schars = max(0, self.length_key(i) - qchars - 200) |
| stok = int(schars / chars_per_token) |
| st[i] = stok |
| q[i] = int(qchars / chars_per_token) + 2 * len(opts) + 8 + min(stok, early_state_tokens) |
| self._qlen, self._slen = q, st |
| return q, st |
|
|
| def length_key(self, i: int) -> int: |
| return int(self.ends[i] - self.starts[i]) |
|
|
| def __getitem__(self, i: int) -> dict: |
| if self._mm is None: |
| with open(self.path, "rb") as f: |
| self._mm = mmap.mmap(f.fileno(), 0, access=mmap.ACCESS_READ) |
| return json.loads(self._mm[int(self.starts[i]) : int(self.ends[i])]) |
|
|
| def take(self, n: int | None = None) -> list[dict]: |
| n = len(self) if n is None else min(n, len(self)) |
| return [self[i] for i in range(n)] |
|
|
| def __getstate__(self): |
| d = dict(self.__dict__) |
| d["_mm"] = None |
| return d |
|
|
|
|
| def permute_options(row: dict, rng: random.Random) -> dict: |
| if row["type"] != "choice" or len(row["options"]) < 2: |
| return row |
| perm = list(range(len(row["options"]))) |
| rng.shuffle(perm) |
| return dict(row, options=[row["options"][j] for j in perm], target=[row["target"][j] for j in perm]) |
|
|
|
|
| def build_batch(rows: Sequence[dict], ser: Serializer, rng: random.Random | None, max_k: int | None = None) -> tuple[Batch, torch.Tensor]: |
| """Encode rows; states shared by several rows are encoded once (q_state map).""" |
| states, q_state, questions, targets = [], [], [], [] |
| key_to_idx: dict[str, int] = {} |
| for row in rows: |
| if rng is not None: |
| row = permute_options(row, rng) |
| key = row.get("state_key") or str(row["state"]) |
| if key not in key_to_idx: |
| key_to_idx[key] = len(states) |
| states.append(ser.encode_state(row["state"])) |
| q_state.append(key_to_idx[key]) |
| _, q = ser.encode_row(row, states[key_to_idx[key]]) |
| questions.append(q) |
| targets.append([float(t) for t in row["target"]][: q.n_options]) |
| batch = collate(states, questions, ser.pad, q_state) |
| K = int(batch.q_count.max()) |
| target = torch.tensor([t + [0.0] * (K - len(t)) for t in targets], dtype=torch.float32) |
| return batch, target |
|
|
|
|
| def epoch_plan(store: RowStore, batch_size: int, world: int, seed: int, epoch: int, |
| qlen: np.ndarray | None = None, slen: np.ndarray | None = None, |
| max_tokens: int | None = None, max_attn: float | None = None, chunk_tokens: int = 512, max_chunks: int = 8) -> list[list[int]]: |
| """Batches of row indices for one epoch, identical on every rank. |
| |
| Global shuffle -> sort by length inside mega-chunks -> greedy batches capped by |
| `batch_size` rows and, when given, by padded token budgets: |
| rows * Lmax <= max_tokens (question sequences, activations) |
| rows * Lmax^2 <= max_attn (block-masked attention matrices) |
| rows * state_tokens <= 4*max_tokens (state chunks) |
| -> shuffle batches -> truncate to a multiple of `world`, so every rank gets |
| exactly the same number of steps (no DDP hang at the end of an epoch). |
| """ |
| rng = np.random.default_rng([seed, epoch]) |
| idx = rng.permutation(len(store)) |
| key = qlen if qlen is not None else (store.ends - store.starts) |
| mega = batch_size * 64 |
| batches: list[list[int]] = [] |
| for start in range(0, len(idx), mega): |
| chunk = idx[start : start + mega] |
| chunk = chunk[np.argsort(key[chunk], kind="stable")] |
| cur: list[int] = [] |
| lmax = 0 |
| smax = 0 |
| for j in chunk.tolist(): |
| lj = int(qlen[j]) if qlen is not None else 0 |
| sj = min(int(slen[j]), chunk_tokens * max_chunks) if slen is not None else 0 |
| nl, ns, n = max(lmax, lj), max(smax, sj), len(cur) + 1 |
| over = n > batch_size |
| if max_tokens and qlen is not None: |
| over = over or (n * nl > max_tokens) or (n * ns > 4 * max_tokens) |
| if max_attn and qlen is not None: |
| over = over or (n * nl * nl > max_attn) |
| if over and cur: |
| batches.append(cur) |
| cur, lmax, smax = [], 0, 0 |
| nl, ns = lj, sj |
| cur.append(j) |
| lmax, smax = nl, ns |
| if cur: |
| batches.append(cur) |
| order = rng.permutation(len(batches)) |
| batches = [batches[i] for i in order] |
| usable = (len(batches) // world) * world |
| if usable == 0: |
| raise ValueError(f"not enough rows for {world} ranks") |
| return batches[:usable] |
|
|
|
|
| class EpochBatches(torch.utils.data.Dataset): |
| """One item = one fully collated training batch (tokenized in loader workers).""" |
|
|
| def __init__(self, store: RowStore, batches: list[list[int]], ser: Serializer, seed: int, epoch: int, permute: bool, start: int = 0): |
| self.store, self.batches, self.ser = store, batches, ser |
| self.seed, self.epoch, self.permute, self.start = seed, epoch, permute, start |
|
|
| def __len__(self) -> int: |
| return len(self.batches) - self.start |
|
|
| def __getitem__(self, i: int): |
| j = i + self.start |
| rows = [self.store[k] for k in self.batches[j]] |
| rng = random.Random(f"{self.seed}:{self.epoch}:{j}") if self.permute else None |
| return build_batch(rows, self.ser, rng) |
|
|
|
|
| def _first(items): |
| return items[0] |
|
|
|
|
| def make_loader(ds: EpochBatches, workers: int) -> torch.utils.data.DataLoader: |
| kw: dict[str, Any] = dict(batch_size=1, shuffle=False, collate_fn=_first, num_workers=workers) |
| if workers > 0: |
| kw.update(prefetch_factor=4, persistent_workers=False) |
| if sys.platform.startswith("linux"): |
| kw["multiprocessing_context"] = "fork" |
| return torch.utils.data.DataLoader(ds, **kw) |
|
|
|
|
| |
| def decision_loss( |
| logits: torch.Tensor, |
| target: torch.Tensor, |
| marker_mask: torch.Tensor, |
| qtype: torch.Tensor, |
| conf_logit: torch.Tensor, |
| weights: dict[str, float], |
| ) -> tuple[torch.Tensor, dict[str, float]]: |
| logp = logits.masked_fill(~marker_mask, -1e4).log_softmax(-1) |
| p = logp.exp() |
| soft_ce = -(target * logp).sum(-1) |
| hard = target.argmax(-1) |
| hard_ce = -logp.gather(1, hard[:, None]).squeeze(1) |
| brier = ((p - target) ** 2 * marker_mask).sum(-1) |
| idx = torch.arange(logits.shape[1], device=logits.device)[None, :].float() |
| ordinal = (p * (idx - hard[:, None].float()).abs()).sum(-1) * (qtype == QTYPE_IDS["score"]).float() |
| correct = (logits.argmax(-1) == hard).float() |
| conf_bce = F.binary_cross_entropy_with_logits(conf_logit, correct) |
| loss = ( |
| weights["soft"] * soft_ce.mean() |
| + weights["hard"] * hard_ce.mean() |
| + weights["brier"] * brier.mean() |
| + weights["ord"] * ordinal.mean() |
| + weights["conf"] * conf_bce |
| ) |
| parts = { |
| "loss": float(loss.detach()), |
| "soft_ce": float(soft_ce.detach().mean()), |
| "brier": float(brier.detach().mean()), |
| "acc": float(correct.mean()), |
| "conf_bce": float(conf_bce.detach()), |
| } |
| return loss, parts |
|
|
|
|
| |
| @torch.no_grad() |
| def evaluate(model: GrootDecisionModel, rows: list[dict], ser: Serializer, device, batch_size: int = 16, max_rows: int | None = None) -> tuple[list[Record], dict]: |
| model.eval() |
| records: list[Record] = [] |
| subset = rows[:max_rows] if max_rows else rows |
| for start in range(0, len(subset), batch_size): |
| chunk = subset[start : start + batch_size] |
| batch, target = build_batch(chunk, ser, None) |
| batch = batch.to(device) |
| with torch.autocast(device_type=device.type, dtype=torch.bfloat16, enabled=device.type == "cuda"): |
| out = model(batch) |
| logits = out["logits"].float().cpu() |
| conf = torch.sigmoid(out["confidence_logit"].float()).cpu() |
| for i, row in enumerate(chunk): |
| k = len(row["options"]) |
| records.append(Record(logits=logits[i, :k].numpy(), target=np.array(row["target"][:k], dtype=np.float64), qtype=QTYPE_IDS[row["type"]], correct_conf=float(conf[i]))) |
| report = evaluate_records(records).to_dict() |
| model.train() |
| return records, report |
|
|
|
|
| |
| def setup_distributed() -> tuple[int, int, torch.device]: |
| if "WORLD_SIZE" in os.environ and int(os.environ["WORLD_SIZE"]) > 1: |
| cuda = torch.cuda.is_available() |
| dist.init_process_group("nccl" if cuda else "gloo", timeout=datetime.timedelta(minutes=90)) |
| rank, world = dist.get_rank(), dist.get_world_size() |
| if cuda: |
| local = int(os.environ.get("LOCAL_RANK", rank)) % torch.cuda.device_count() |
| torch.cuda.set_device(local) |
| return rank, world, torch.device("cuda", local) |
| return rank, world, torch.device("cpu") |
| if torch.cuda.is_available(): |
| return 0, 1, torch.device("cuda") |
| if torch.backends.mps.is_available(): |
| return 0, 1, torch.device("mps") |
| return 0, 1, torch.device("cpu") |
|
|
|
|
| def barrier(world: int) -> None: |
| if world > 1: |
| dist.barrier() |
|
|
|
|
| |
| def save_checkpoint(out: Path, model: GrootDecisionModel, opt: torch.optim.Optimizer, tok, state: dict) -> None: |
| """Atomic resumable checkpoint at out/ckpt (weights + optimizer + trainer state).""" |
| tmp, final = out / "ckpt.tmp", out / "ckpt" |
| if tmp.exists(): |
| shutil.rmtree(tmp) |
| model.save_pretrained(tmp, tokenizer=tok) |
| torch.save(opt.state_dict(), tmp / "optimizer.pt") |
| (tmp / "trainer_state.json").write_text(json.dumps(state, indent=2) + "\n") |
| if final.exists(): |
| shutil.rmtree(final) |
| tmp.rename(final) |
|
|
|
|
| def load_checkpoint(ckpt: Path, model: GrootDecisionModel, opt: torch.optim.Optimizer, device) -> dict: |
| from safetensors.torch import load_file |
|
|
| model.load_state_dict(load_file(str(ckpt / "model.safetensors")), strict=True) |
| opt.load_state_dict(torch.load(ckpt / "optimizer.pt", map_location=device, weights_only=True)) |
| return json.loads((ckpt / "trainer_state.json").read_text()) |
|
|
|
|
| def lr_at(step: int, total: int, warmup: int, peak: float, floor: float = 0.1) -> float: |
| if step < warmup: |
| return peak * step / max(1, warmup) |
| progress = (step - warmup) / max(1, total - warmup) |
| return peak * (floor + (1 - floor) * 0.5 * (1 + math.cos(math.pi * min(1.0, progress)))) |
|
|
|
|
| def main(argv: list[str] | None = None) -> None: |
| ap = argparse.ArgumentParser(prog="groot-train") |
| ap.add_argument("--preset", default="groot-1-s") |
| ap.add_argument("--backbone", default=None, help="override backbone HF id or local path") |
| ap.add_argument("--init", default=None, help="resume/initialise from a Groot checkpoint directory") |
| ap.add_argument("--train", required=True) |
| ap.add_argument("--val", default=None) |
| ap.add_argument("--out", required=True) |
| ap.add_argument("--epochs", type=float, default=2.0) |
| ap.add_argument("--batch-size", type=int, default=16) |
| ap.add_argument("--grad-accum", type=int, default=1) |
| ap.add_argument("--lr", type=float, default=5e-5) |
| ap.add_argument("--head-lr", type=float, default=2e-4) |
| ap.add_argument("--weight-decay", type=float, default=0.01) |
| ap.add_argument("--warmup", type=int, default=500) |
| ap.add_argument("--clip", type=float, default=1.0) |
| ap.add_argument("--freeze-encoder", action="store_true") |
| ap.add_argument("--grad-checkpoint", action="store_true") |
| ap.add_argument("--w-soft", type=float, default=1.0) |
| ap.add_argument("--w-hard", type=float, default=0.3) |
| ap.add_argument("--w-brier", type=float, default=0.3) |
| ap.add_argument("--w-ord", type=float, default=0.1) |
| ap.add_argument("--w-conf", type=float, default=0.2) |
| ap.add_argument("--no-permute", action="store_true") |
| ap.add_argument("--eval-every", type=int, default=1000) |
| ap.add_argument("--eval-rows", type=int, default=2000) |
| ap.add_argument("--save-every", type=int, default=2000) |
| ap.add_argument("--log-every", type=int, default=50) |
| ap.add_argument("--limit", type=int, default=None, help="use only the first N training rows (debug)") |
| ap.add_argument("--tiny", action="store_true", help="random tiny encoder (tests)") |
| ap.add_argument("--seed", type=int, default=17) |
| ap.add_argument("--workers", type=int, default=0, help="loader workers per rank (tokenization off the training thread)") |
| ap.add_argument("--resume", action="store_true", help="continue from <out>/ckpt if it exists") |
| ap.add_argument("--ckpt-minutes", type=float, default=15.0, help="resumable checkpoint interval") |
| ap.add_argument("--final-eval-rows", type=int, default=20000) |
| ap.add_argument("--max-steps", type=int, default=None, help="stop after N optimizer steps (smoke tests)") |
| ap.add_argument("--find-unused", action="store_true", help="DDP find_unused_parameters") |
| ap.add_argument("--max-tokens", type=int, default=None, help="padded question tokens per batch (rows * Lmax)") |
| ap.add_argument("--max-attn", type=float, default=None, help="padded attention budget per batch (rows * Lmax^2)") |
| ap.add_argument("--max-state-chunks", type=int, default=8, help="truncate states to N chunks during training") |
| ap.add_argument("--save-every-eval", action="store_true", help="also save out/step_N at every evaluation") |
| ap.add_argument("--spike-factor", type=float, default=1.8, help="rollback when fast loss EMA > factor * best slow EMA") |
| ap.add_argument("--max-rollbacks", type=int, default=3) |
| ap.add_argument("--chunk-pos", choices=["learned", "sinusoidal"], default=None, help="override the chunk-position scheme") |
| ap.add_argument("--set", action="append", default=[], help="config override key=value (e.g. state_retrieval=true)") |
| args = ap.parse_args(argv) |
|
|
| rank, world, device = setup_distributed() |
| os.environ.setdefault("TOKENIZERS_PARALLELISM", "false") |
| torch.manual_seed(args.seed) |
| out = Path(args.out) |
| if rank == 0: |
| out.mkdir(parents=True, exist_ok=True) |
| barrier(world) |
|
|
| from transformers import AutoTokenizer |
|
|
| cfg = preset(args.preset, **({"backbone": args.backbone} if args.backbone else {})) |
| if args.init: |
| model = GrootDecisionModel.from_pretrained(args.init, device="cpu") |
| tok = AutoTokenizer.from_pretrained(Path(args.init) / "tokenizer") |
| cfg = model.cfg |
| elif args.tiny: |
| from huggingface_hub import snapshot_download |
|
|
| tok = AutoTokenizer.from_pretrained(snapshot_download(cfg.backbone, allow_patterns=["tokenizer*", "special_tokens_map.json"])) |
| torch.manual_seed(args.seed) |
| model = GrootDecisionModel.tiny(cfg, vocab_size=len(tok), layers=2) |
| else: |
| tok = AutoTokenizer.from_pretrained(cfg.backbone) |
| model = GrootDecisionModel.from_backbone(cfg) |
| if args.chunk_pos: |
| model.cfg = GrootConfig.from_dict({**model.cfg.to_dict(), "chunk_pos": args.chunk_pos}) |
| cfg = model.cfg |
| if args.set: |
| from .config import parse_overrides |
|
|
| model.cfg = GrootConfig.from_dict({**model.cfg.to_dict(), **parse_overrides(args.set)}) |
| cfg = model.cfg |
| model.sync_chunk_pos_trainable() |
| model.to(device) |
| if args.grad_checkpoint and hasattr(model.encoder, "gradient_checkpointing_enable"): |
| model.encoder.gradient_checkpointing_enable() |
| if args.freeze_encoder: |
| for p in model.encoder.parameters(): |
| p.requires_grad_(False) |
| train_cfg = GrootConfig.from_dict({**cfg.to_dict(), "max_chunks": min(cfg.max_chunks, args.max_state_chunks)}) |
| ser = Serializer(tok, train_cfg) |
|
|
| store = RowStore(args.train, args.limit) |
| qlen = slen = None |
| if args.max_tokens or args.max_attn: |
| t_est = time.time() |
| qlen, slen = store.estimate_lengths(cfg.early_state_tokens) |
| if rank == 0: |
| print(json.dumps({"length_estimate_s": round(time.time() - t_est, 1), "qlen_p50": int(np.median(qlen)), |
| "qlen_p99": int(np.percentile(qlen, 99)), "qlen_max": int(qlen.max())}), flush=True) |
| plan_kw = dict(qlen=qlen, slen=slen, max_tokens=args.max_tokens, max_attn=args.max_attn, |
| chunk_tokens=cfg.chunk_tokens, max_chunks=args.max_state_chunks) |
| val_rows = RowStore(args.val).take(max(args.eval_rows, args.final_eval_rows)) if args.val else [] |
| if rank == 0: |
| print(json.dumps({"params": model.count_parameters(), "train_rows": len(store), "val_rows": len(val_rows), |
| "device": str(device), "world": world, "workers": args.workers}), flush=True) |
|
|
| enc_params = [p for n, p in model.named_parameters() if n.startswith("encoder.") and p.requires_grad] |
| head_params = [p for n, p in model.named_parameters() if not n.startswith("encoder.") and p.requires_grad] |
| groups = [{"params": enc_params, "lr": args.lr}, {"params": head_params, "lr": args.head_lr}] |
| opt = torch.optim.AdamW(groups, weight_decay=args.weight_decay, betas=(0.9, 0.98)) |
|
|
| state = {"step": 0, "epoch": 0, "batch_pos": 0, "best": -1.0, "lr_scale": 1.0, "rollbacks": 0} |
| ckpt = out / "ckpt" |
| if args.resume and ckpt.exists(): |
| state = load_checkpoint(ckpt, model, opt, device) |
| if rank == 0: |
| print(json.dumps({"resumed": state}), flush=True) |
|
|
| ddp = model |
| if world > 1: |
| ddp = nn.parallel.DistributedDataParallel(model, device_ids=[device.index] if device.type == "cuda" else None, |
| find_unused_parameters=args.find_unused) |
|
|
| batches_per_rank = len(epoch_plan(store, args.batch_size, world, args.seed, 0, **plan_kw)) // world |
| steps_per_epoch = max(1, batches_per_rank // args.grad_accum) |
| total_steps = max(1, int(steps_per_epoch * args.epochs)) |
| if args.max_steps: |
| total_steps = min(total_steps, args.max_steps) |
| weights = {"soft": args.w_soft, "hard": args.w_hard, "brier": args.w_brier, "ord": args.w_ord, "conf": args.w_conf} |
| log_path = out / "train_log.jsonl" |
|
|
| def log(rec: dict) -> None: |
| print(json.dumps(rec), flush=True) |
| with log_path.open("a") as f: |
| f.write(json.dumps(rec) + "\n") |
|
|
| step, epoch, best = state["step"], state["epoch"], state["best"] |
| lr_scale, rollbacks = state.get("lr_scale", 1.0), state.get("rollbacks", 0) |
| ema_fast = ema_slow = None |
| best_slow = float("inf") |
| last_gn = 0.0 |
| start_pos = state["batch_pos"] |
| t0 = time.time() |
| last_ckpt = time.time() |
| win_tokens, win_rows, win_t = 0, 0, time.time() |
| model.train() |
| done = step >= total_steps |
| while not done: |
| plan = epoch_plan(store, args.batch_size, world, args.seed, epoch, **plan_kw)[rank::world] |
| loader = make_loader(EpochBatches(store, plan, ser, args.seed + rank, epoch, not args.no_permute, start_pos), args.workers) |
| pos = start_pos |
| micro = 0 |
| for batch, target in loader: |
| pos += 1 |
| batch, target = batch.to(device), target.to(device) |
| win_tokens += int(batch.chunk_mask.sum()) + int(batch.q_mask.sum()) |
| win_rows += int(target.shape[0]) |
| with torch.autocast(device_type=device.type, dtype=torch.bfloat16, enabled=device.type == "cuda"): |
| outp = ddp(batch) |
| K = int(batch.q_count.max()) |
| marker_mask = torch.arange(K, device=device)[None, :] < batch.q_count[:, None] |
| loss, parts = decision_loss(outp["logits"].float(), target, marker_mask, batch.q_type, outp["confidence_logit"].float(), weights) |
| (loss / args.grad_accum).backward() |
| micro += 1 |
| if micro % args.grad_accum: |
| continue |
| step += 1 |
| for g, base in zip(opt.param_groups, (args.lr, args.head_lr)): |
| g["lr"] = lr_at(step, total_steps, args.warmup, base) * lr_scale |
| last_gn = float(torch.nn.utils.clip_grad_norm_(model.parameters(), args.clip)) |
| opt.step() |
| opt.zero_grad(set_to_none=True) |
| |
| lv = torch.tensor([parts["loss"]], device=device) |
| if world > 1: |
| dist.all_reduce(lv) |
| lv /= world |
| lval = float(lv) |
| ema_fast = lval if ema_fast is None else 0.9 * ema_fast + 0.1 * lval |
| ema_slow = lval if ema_slow is None else 0.99 * ema_slow + 0.01 * lval |
| if step > args.warmup + 200: |
| best_slow = min(best_slow, ema_slow) |
| diverged = not math.isfinite(lval) or (step > args.warmup + 300 and ema_fast > args.spike_factor * best_slow) |
| if diverged and rollbacks < args.max_rollbacks and ckpt.exists(): |
| barrier(world) |
| saved = load_checkpoint(ckpt, model, opt, device) |
| rollbacks += 1 |
| lr_scale *= 0.5 |
| ema_fast = ema_slow = None |
| if rank == 0: |
| log({"step": step, "event": "rollback", "to_step": saved["step"], "lr_scale": lr_scale, "rollbacks": rollbacks, |
| "loss_fast": None if not math.isfinite(lval) else round(lval, 4), "best_slow": round(best_slow, 4)}) |
| barrier(world) |
| continue |
| if rank == 0 and step % args.log_every == 0: |
| dt = max(1e-6, time.time() - win_t) |
| rec = {"step": step, "epoch": epoch, "lr": opt.param_groups[0]["lr"], "elapsed": round(time.time() - t0, 1), "grad_norm": round(last_gn, 3), |
| "tok_s": round(win_tokens * world / dt), "rows_s": round(win_rows * world / dt, 1), **parts} |
| if device.type == "cuda": |
| rec["mem_gb"] = round(torch.cuda.max_memory_allocated(device) / 2**30, 2) |
| log(rec) |
| win_tokens, win_rows, win_t = 0, 0, time.time() |
| if val_rows and step % args.eval_every == 0: |
| if rank == 0: |
| _, report = evaluate(model, val_rows, ser, device, args.batch_size, args.eval_rows) |
| log({"step": step, "val": report}) |
| if report["accuracy"].get("all", -1) > best: |
| best = report["accuracy"]["all"] |
| model.save_pretrained(out / "best", tokenizer=tok) |
| if args.save_every_eval: |
| model.save_pretrained(out / f"step_{step}", tokenizer=tok) |
| barrier(world) |
| want_ckpt = step % args.save_every == 0 or (time.time() - last_ckpt) > args.ckpt_minutes * 60 |
| if world > 1: |
| flag = torch.tensor([1 if want_ckpt else 0], device=device) |
| dist.broadcast(flag, 0) |
| want_ckpt = bool(flag.item()) |
| if want_ckpt: |
| if rank == 0: |
| save_checkpoint(out, model, opt, tok, {"step": step, "epoch": epoch, "batch_pos": pos, "best": best, |
| "lr_scale": lr_scale, "rollbacks": rollbacks}) |
| last_ckpt = time.time() |
| barrier(world) |
| if step >= total_steps: |
| done = True |
| break |
| start_pos = 0 |
| if not done: |
| epoch += 1 |
| barrier(world) |
| if rank == 0: |
| if val_rows: |
| records, report = evaluate(model, val_rows, ser, device, args.batch_size, args.final_eval_rows) |
| temps = fit_temperatures(records, n_types=len(cfg.qtypes)) |
| model.temperature.copy_(torch.tensor(temps, dtype=model.temperature.dtype, device=model.temperature.device)) |
| calibrated = evaluate_records(records, temps).to_dict() |
| log({"final_val": report, "calibrated": calibrated, "temperatures": temps}) |
| (out / "eval.json").write_text(json.dumps({"final_val": report, "calibrated": calibrated}, indent=2)) |
| model.save_pretrained(out / "final", tokenizer=tok) |
| print(f"saved {out / 'final'}", flush=True) |
| if world > 1: |
| dist.barrier() |
| dist.destroy_process_group() |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
