| """Train the heads classifier on an export bundle. Run on the pod as ``python -m training.train``. |
| |
| Classifier-plan §3.2-3.5, §4.3-4.4. Steps of ``--posts-per-step`` posts (32) in length buckets, split into |
| token-budget micro-batches (16k padded tokens) with per-step loss normalisation; masked targets; LoRA lr 2e-4, |
| heads lr 1e-3, AdamW wd 0.01, clip 1.0, cosine with 3% warmup; val every ``--eval-every`` steps (200, or four times |
| per epoch if that is fewer); ``best/`` = lowest mean val log-loss over all fields. The arm of the §4.4 screen is |
| ``--checkpoint base|it`` (the base checkpoint dir is ``--model``) times ``--pool last|mean``. Per-option class |
| weights come from the bundle's taxonomy (nsb-v2: pnl breakeven 3.0) unless ``--class-weights`` overrides them; the |
| resolved weights are recorded in config.json. |
| |
| Writes OUT/config.json, OUT/taxonomy.json, OUT/train_log.jsonl, OUT/last/ (resume point every eval: adapter, |
| heads.safetensors, optimizer, scheduler, RNG, data cursor, cumulative hours), OUT/best/, OUT/predictions/val.jsonl |
| and OUT/metrics_val.json for the best checkpoint. ``--max-hours`` is cumulative across resumes (state.json). |
| Exit code 3: stopped by --max-hours, --stop-after or SIGINT (``timeout``) at a saved resume point; rerun with |
| --resume. Exit code 4: non-finite loss or gradient (nothing saved; last/ keeps the previous good point). |
| """ |
| from __future__ import annotations |
|
|
| import argparse |
| import hashlib |
| import json |
| import math |
| import platform |
| import random |
| import time |
| from pathlib import Path |
|
|
| from .calibrate import nll, outputs |
| from .data import inference_batches, load_split, micro_batches, step_counts, steps, subset, write_jsonl |
| from .formats import FORMAT, encode |
| from .taxonomy import Taxonomy |
|
|
| |
| CORE = ("data", "model", "checkpoint", "pool", "subset", "epochs", "max_steps", "posts_per_step", "token_budget", |
| "max_len", "seed", "lr_lora", "lr_heads", "warmup", "weight_decay", "lora_r", "lora_alpha", "lora_dropout", |
| "head_dropout", "dtype", "eval_every", "val_limit", "class_weights") |
|
|
|
|
| def parse(argv=None): |
| p = argparse.ArgumentParser(description=__doc__.splitlines()[0]) |
| p.add_argument("--data", required=True, help="Bundle directory (train/val.jsonl, taxonomy.json, manifest.json)") |
| p.add_argument("--model", required=True, help="Local base checkpoint directory") |
| p.add_argument("--out", required=True) |
| p.add_argument("--arm", default="", help="Arm name for the record (B-last, I-last, B-mean, I-mean)") |
| p.add_argument("--checkpoint", default="base", choices=["base", "it"], help="Which Gemma 4 checkpoint --model is") |
| p.add_argument("--pool", default="last", choices=["last", "mean"]) |
| p.add_argument("--subset", type=int, default=0, help="Train on this many rows in seeded order (0 = all)") |
| p.add_argument("--epochs", type=float, default=2.0) |
| p.add_argument("--max-steps", type=int, default=0, help="Cap the schedule at this many optimiser steps") |
| p.add_argument("--posts-per-step", type=int, default=32) |
| p.add_argument("--token-budget", type=int, default=16384, help="Padded tokens per training micro-batch") |
| p.add_argument("--eval-token-budget", type=int, default=65536) |
| p.add_argument("--max-len", type=int, default=2048) |
| p.add_argument("--lr-lora", type=float, default=2e-4) |
| p.add_argument("--lr-heads", type=float, default=1e-3) |
| p.add_argument("--warmup", type=float, default=0.03) |
| p.add_argument("--weight-decay", type=float, default=0.01) |
| p.add_argument("--lora-r", type=int, default=32) |
| p.add_argument("--lora-alpha", type=int, default=32) |
| p.add_argument("--lora-dropout", type=float, default=0.05) |
| p.add_argument("--head-dropout", type=float, default=0.1) |
| p.add_argument("--class-weights", default=None, |
| help='Per-option loss weights as JSON, e.g. \'{"pnl": {"breakeven": 3.0}}\'; "{}" = none ' |
| "(default: the bundle taxonomy's class_weights)") |
| p.add_argument("--eval-every", type=int, default=200, help="Optimiser steps between val evaluations and saves") |
| p.add_argument("--val-limit", type=int, default=0, help="Evaluate on the first N val rows in seeded order") |
| p.add_argument("--log-every", type=int, default=10) |
| p.add_argument("--max-hours", type=float, default=0, help="Cumulative training hours across resumes") |
| p.add_argument("--stop-after", type=int, default=0, help="Stop this process after N more steps (resume test)") |
| p.add_argument("--force-longest", type=int, default=0, help="First run one forward/backward on the N longest rows") |
| p.add_argument("--seed", type=int, default=20261017) |
| p.add_argument("--resume", action="store_true", help="Continue from last/ if it exists (settings must match)") |
| p.add_argument("--device", default="cuda") |
| p.add_argument("--dtype", default="bfloat16", choices=["bfloat16", "float32"]) |
| p.add_argument("--no-gradient-checkpointing", action="store_true") |
| return p.parse_args(argv) |
|
|
|
|
| def versions() -> dict: |
| out = dict(python=platform.python_version()) |
| for pkg in ("torch", "transformers", "peft", "accelerate", "safetensors", "tokenizers", "huggingface_hub"): |
| try: |
| out[pkg] = getattr(__import__(pkg), "__version__", "unknown") |
| except Exception: |
| out[pkg] = None |
| try: |
| import torch |
|
|
| out["cuda"] = torch.version.cuda |
| if torch.cuda.is_available(): |
| out["gpu"] = torch.cuda.get_device_name(0) |
| except Exception: |
| pass |
| commit = Path(__file__).with_name(".source_commit") |
| out["source_commit"] = commit.read_text().strip() if commit.exists() else None |
| return out |
|
|
|
|
| def file_sha256(path: Path) -> str | None: |
| if not path.exists(): |
| return None |
| return hashlib.sha256(path.read_bytes()).hexdigest() |
|
|
|
|
| class Log: |
| def __init__(self, path: Path): |
| self.fh = open(path, "a", encoding="utf-8") |
|
|
| def __call__(self, **row): |
| row = dict(time=round(time.time(), 1), **row) |
| line = json.dumps(row, ensure_ascii=False) |
| self.fh.write(line + "\n") |
| self.fh.flush() |
| print(line, flush=True) |
|
|
|
|
| def prediction_rows(rows, logits_by_index, infos, split=None) -> list[dict]: |
| """Prediction rows hold ids and numbers only (never text).""" |
| return [{"id": r["id"], "split": split or r.get("split"), "n_tokens": infos[i]["n_tokens"], |
| "truncated": infos[i]["cut"], "logits": logits_by_index[i]} for i, r in enumerate(rows)] |
|
|
|
|
| def val_report(preds: list[dict], rows: list[dict], tax: Taxonomy) -> dict: |
| """Mean val log-loss over fields (selection), per-field log-loss and gated accuracy (uncalibrated).""" |
| pairs = list(zip(preds, rows)) |
| per_field = {f: nll(pairs, tax, f, 1.0) for f in tax.names} |
| scored = [v for v in per_field.values() if v is not None] |
| acc = {} |
| labels = [outputs(p["logits"], tax)["labels"] for p in preds] |
| for f in tax.names: |
| gold = [(lab[f], r["targets"].get(f)) for lab, r in zip(labels, rows) |
| if f not in r.get("mask", ()) and r["targets"].get(f) is not None] |
| acc[f] = sum(1 for a, g in gold if a == g) / len(gold) if gold else None |
| eval_pairs = [(p, r) for p, r in pairs if r.get("eval")] |
| screen = [nll(eval_pairs, tax, f, 1.0) for f in ("relevance", "direction", "commercial")] |
| screen = [v for v in screen if v is not None] |
| return {"selection_loss": sum(scored) / len(scored) if scored else float("inf"), "log_loss": per_field, |
| "accuracy": acc, "rows": len(rows), "eval_rows": len(eval_pairs), |
| "screen_primary": sum(screen) / len(screen) if screen else None} |
|
|
|
|
| def resolve_class_weights(tax: Taxonomy, override: str | None) -> Taxonomy: |
| """The taxonomy with ``--class-weights`` applied (JSON object; None keeps the bundle's weights).""" |
| if override is None: |
| return tax |
| value = json.loads(override) |
| if not isinstance(value, dict): |
| raise SystemExit("--class-weights must be a JSON object {field: {option: weight}}") |
| return tax.replace(class_weights=value) |
|
|
|
|
| def schedule_eval_every(eval_every: int, steps_per_epoch: int) -> int: |
| """200 steps, or four times per epoch if that is fewer steps (classifier-plan §4.3).""" |
| return max(1, min(eval_every, math.ceil(steps_per_epoch / 4))) |
|
|
|
|
| def main(argv=None) -> int: |
| args = parse(argv) |
| import torch |
|
|
| from .model import HUB_IDS, bos_id, build, load_tokenizer, pad, predict, rng_state, save_checkpoint, \ |
| set_rng_state, step_loss, tokenizer_fn |
|
|
| out = Path(args.out) |
| out.mkdir(parents=True, exist_ok=True) |
| cfg_path = out / "config.json" |
| tax = resolve_class_weights(Taxonomy.load(args.data), args.class_weights) |
| cfg = {k: v for k, v in vars(args).items() if k not in ("resume", "device", "stop_after", "force_longest")} |
| cfg["gradient_checkpointing"] = not args.no_gradient_checkpointing |
| cfg["class_weights"] = tax.class_weights or None |
| if cfg_path.exists(): |
| if not args.resume: |
| raise SystemExit(f"refusing to start a new run over {out}; use --resume or a new --out") |
| old = json.loads(cfg_path.read_text()) |
| changed = sorted(k for k in CORE if old.get(k) != cfg.get(k)) |
| if changed: |
| raise SystemExit(f"--resume with different settings: {changed}") |
| if old.get("max_hours") != cfg.get("max_hours"): |
| raise SystemExit("--resume with a different --max-hours; the cap is cumulative") |
| manifest_path = Path(args.data) / "manifest.json" |
| manifest = json.loads(manifest_path.read_text()) if manifest_path.exists() else {} |
| cfg.update(versions=versions(), format=FORMAT, base_model=HUB_IDS[args.checkpoint], |
| data_version=manifest.get("version"), data_manifest_sha256=file_sha256(manifest_path), |
| taxonomy_version=manifest.get("taxonomy_version"), render_version=manifest.get("render_version")) |
| cfg_path.write_text(json.dumps(cfg, indent=1)) |
| log = Log(out / "train_log.jsonl") |
|
|
| random.seed(args.seed) |
| torch.manual_seed(args.seed) |
| device = torch.device(args.device) |
| (out / "taxonomy.json").write_text(json.dumps(tax.to_json(), ensure_ascii=False, indent=1)) |
| tok = load_tokenizer(args.model) |
| tokenize, bos = tokenizer_fn(tok), bos_id(tok) |
|
|
| train_rows = subset(load_split(args.data, "train"), args.subset, args.seed) |
| val_rows = subset(load_split(args.data, "val"), args.val_limit, args.seed) |
| t0 = time.time() |
| enc_train = [encode(r["segments"], tokenize, args.max_len, bos) for r in train_rows] |
| enc_val = [encode(r["segments"], tokenize, args.max_len, bos) for r in val_rows] |
| ids = [e[0] for e in enc_train] |
| lengths = [len(x) for x in ids] |
| labels = [tax.encode(r["targets"], r.get("mask", ())) for r in train_rows] |
| val_ids = [e[0] for e in enc_val] |
| val_infos = [e[1] for e in enc_val] |
| val_batches = inference_batches([len(x) for x in val_ids], args.eval_token_budget) |
| log(event="data", train=len(train_rows), val=len(val_rows), tokens=sum(lengths), max_len=max(lengths), |
| mean_len=round(sum(lengths) / len(lengths), 1), truncated_rows=sum(1 for e in enc_train if e[1]["cut"]), |
| val_tokens=sum(len(x) for x in val_ids), seconds=round(time.time() - t0, 1)) |
|
|
| last_dir, best_dir = out / "last", out / "best" |
| resume = args.resume and (last_dir / "state.json").exists() |
| model = build(cfg, tax.sizes(), device, train=True) |
| if resume: |
| from peft import set_peft_model_state_dict |
| from safetensors.torch import load_file |
|
|
| from .model import load_heads |
|
|
| set_peft_model_state_dict(model.backbone, load_file(str(last_dir / "adapter" / "adapter_model.safetensors"))) |
| load_heads(model, last_dir / "heads.safetensors") |
| lora = [p for n, p in model.named_parameters() if p.requires_grad and "lora_" in n] |
| heads = [p for n, p in model.named_parameters() if p.requires_grad and "lora_" not in n] |
| opt = torch.optim.AdamW([dict(params=lora, lr=args.lr_lora), dict(params=heads, lr=args.lr_heads)], |
| weight_decay=args.weight_decay) |
| steps_per_epoch = math.ceil(len(train_rows) / args.posts_per_step) |
| total = max(1, int(steps_per_epoch * args.epochs)) |
| if args.max_steps: |
| total = min(total, args.max_steps) |
| eval_every = schedule_eval_every(args.eval_every, steps_per_epoch) |
| from transformers import get_cosine_schedule_with_warmup |
|
|
| sched = get_cosine_schedule_with_warmup(opt, int(args.warmup * total), total) |
| state = dict(step=0, epoch=0, step_in_epoch=0, best=None, best_step=None, hours_used=0.0) |
| if resume: |
| state = json.loads((last_dir / "state.json").read_text()) |
| trainer = torch.load(last_dir / "trainer.pt", map_location="cpu", weights_only=False) |
| opt.load_state_dict(trainer["opt"]) |
| sched.load_state_dict(trainer["sched"]) |
| set_rng_state(trainer["rng"]) |
| log(event="resume", **state) |
| log(event="start", trainable=sum(p.numel() for p in lora + heads), steps=total, steps_per_epoch=steps_per_epoch, |
| eval_every=eval_every, arm=args.arm, checkpoint=args.checkpoint, pool=args.pool) |
|
|
| amp = dict(device_type=device.type, dtype=torch.bfloat16, enabled=device.type == "cuda") |
| process_start = time.time() |
| hours_before = state["hours_used"] |
| weights = tax.weights() |
| class_weights = [None if v is None else torch.tensor(v, dtype=torch.float32, device=device) |
| for v in tax.class_weight_vectors()] |
|
|
| def hours_now(): |
| return hours_before + (time.time() - process_start) / 3600 |
|
|
| def save_last(): |
| state["hours_used"] = round(hours_now(), 4) |
| save_checkpoint(model, last_dir, tax.names, dict(state), |
| dict(opt=opt.state_dict(), sched=sched.state_dict(), rng=rng_state())) |
|
|
| def evaluate(): |
| t = time.time() |
| rng = rng_state() |
| logits = predict(model, val_ids, tok.pad_token_id, val_batches, device, tax.names) |
| set_rng_state(rng) |
| preds = prediction_rows(val_rows, logits, val_infos, "val") |
| rep = val_report(preds, val_rows, tax) |
| seconds = time.time() - t |
| log(event="eval", step=state["step"], selection_loss=round(rep["selection_loss"], 5), |
| screen_primary=rep["screen_primary"], |
| log_loss={f: (round(v, 4) if v is not None else None) for f, v in rep["log_loss"].items()}, |
| accuracy={f: (round(v, 4) if v is not None else None) for f, v in rep["accuracy"].items()}, |
| seconds=round(seconds, 1), posts_per_second=round(len(val_rows) / max(seconds, 1e-9), 1)) |
| if state["best"] is None or rep["selection_loss"] < state["best"]: |
| state["best"], state["best_step"] = rep["selection_loss"], state["step"] |
| save_checkpoint(model, best_dir, tax.names, dict(state)) |
| write_jsonl(out / "predictions" / "val.jsonl", preds) |
| (out / "metrics_val.json").write_text(json.dumps(dict(step=state["step"], **rep), indent=1)) |
|
|
| if args.force_longest and not resume: |
| |
| pool = sorted(load_split(args.data, "train"), key=lambda r: -sum(len(x["text"]) for x in r["segments"])) |
| cand = [(encode(r["segments"], tokenize, args.max_len, bos)[0], tax.encode(r["targets"], r.get("mask", ()))) |
| for r in pool[:4 * args.force_longest]] |
| cand = sorted(cand, key=lambda c: -len(c[0]))[:args.force_longest] |
| f_ids, f_labels = [c[0] for c in cand], [c[1] for c in cand] |
| f_len = [len(x) for x in f_ids] |
| parts = micro_batches(list(range(len(cand))), f_len, args.token_budget) |
| for mb in parts: |
| x, m, last = pad([f_ids[i] for i in mb], tok.pad_token_id, device) |
| y = torch.tensor([f_labels[i] for i in mb], device=device) |
| with torch.autocast(**amp): |
| loss, _ = step_loss(model(x, m, last), y, step_counts(f_labels), weights, class_weights) |
| loss.backward() |
| opt.zero_grad(set_to_none=True) |
| mem = round(torch.cuda.max_memory_allocated() / 2 ** 30, 2) if device.type == "cuda" else None |
| log(event="forced_longest", rows=len(cand), lengths=f_len, micro_batches=len(parts), |
| padded_tokens=[max(f_len[i] for i in mb) * len(mb) for mb in parts], max_memory_gb=mem) |
| if device.type == "cuda": |
| torch.cuda.reset_peak_memory_stats() |
|
|
| model.train() |
| try: |
| stop = run_steps(args, state, total, eval_every, model, opt, sched, lora, heads, ids, lengths, labels, tax, |
| weights, device, amp, tok, log, hours_now, evaluate, save_last, class_weights) |
| except KeyboardInterrupt: |
| log(event="interrupted", step=state["step"]) |
| save_last() |
| return 3 |
| if stop == 4: |
| return 4 |
| save_last() |
| if state["best"] is None: |
| evaluate() |
| log(event="done" if stop is None else "stopped", **state) |
| return stop or 0 |
|
|
|
|
| def run_steps(args, state, total, eval_every, model, opt, sched, lora, heads, ids, lengths, labels, tax, weights, |
| device, amp, tok, log, hours_now, evaluate, save_last, class_weights=None): |
| """The optimiser loop. Returns None (finished), 3 (cap or planned stop; the caller saves the resume point) |
| or 4 (non-finite loss or gradient: nothing is saved, so last/ keeps the previous good resume point).""" |
| import torch |
|
|
| from .model import pad, step_loss |
|
|
| def fresh_window(): |
| return dict(loss=0.0, tokens=0, padded=0, t=time.time(), field_sums=[0.0] * len(tax.names), |
| field_counts=[0] * len(tax.names), n=0) |
|
|
| window = fresh_window() |
| stop, done_here = None, 0 |
| while state["step"] < total and stop is None: |
| epoch_steps = steps(lengths, args.posts_per_step, args.seed * 1000 + state["epoch"]) |
| while state["step_in_epoch"] < len(epoch_steps) and state["step"] < total: |
| if args.max_hours and hours_now() >= args.max_hours: |
| stop = 3 |
| break |
| step = epoch_steps[state["step_in_epoch"]] |
| step_labels = [labels[i] for i in step] |
| counts = step_counts(step_labels) |
| step_total = 0.0 |
| for mb in micro_batches(step, lengths, args.token_budget): |
| x, m, last = pad([ids[i] for i in mb], tok.pad_token_id, device) |
| y = torch.tensor([labels[i] for i in mb], device=device) |
| with torch.autocast(**amp): |
| loss, per_field = step_loss(model(x, m, last), y, counts, weights, class_weights) |
| loss.backward() |
| step_total += float(loss.detach()) |
| window["field_sums"] = [a + b for a, b in zip(window["field_sums"], per_field)] |
| window["tokens"] += int(m.sum()) |
| window["padded"] += int(m.numel()) |
| window["field_counts"] = [a + sum(1 for r in step_labels if r[f] >= 0) |
| for f, a in enumerate(window["field_counts"])] |
| norm = float(torch.nn.utils.clip_grad_norm_(lora + heads, 1.0)) |
| if not (math.isfinite(step_total) and math.isfinite(norm)): |
| log(event="non_finite", step=state["step"] + 1, loss=step_total, grad_norm=norm) |
| opt.zero_grad(set_to_none=True) |
| return 4 |
| opt.step() |
| sched.step() |
| opt.zero_grad(set_to_none=True) |
| state["step"] += 1 |
| state["step_in_epoch"] += 1 |
| done_here += 1 |
| window["loss"] += step_total |
| window["n"] += 1 |
| if state["step"] % args.log_every == 0 or state["step"] == total: |
| dt = max(time.time() - window["t"], 1e-9) |
| row = dict(event="train", step=state["step"], epoch=state["epoch"], |
| loss=round(window["loss"] / window["n"], 4), |
| fields={f: round(s / c, 4) if c else None for f, s, c in |
| zip(tax.names, window["field_sums"], window["field_counts"])}, |
| lr=sched.get_last_lr()[0], tokens_per_second=round(window["tokens"] / dt), |
| padded_tokens_per_second=round(window["padded"] / dt), hours=round(hours_now(), 3)) |
| if device.type == "cuda": |
| row["max_memory_gb"] = round(torch.cuda.max_memory_allocated() / 2 ** 30, 1) |
| log(**row) |
| window = fresh_window() |
| if state["step_in_epoch"] >= len(epoch_steps): |
| state["epoch"] += 1 |
| state["step_in_epoch"] = 0 |
| if state["step"] % eval_every == 0 or state["step"] == total: |
| evaluate() |
| save_last() |
| if args.stop_after and done_here >= args.stop_after: |
| stop = 3 |
| break |
| if state["step_in_epoch"] == 0: |
| break |
| return stop |
|
|
|
|
| if __name__ == "__main__": |
| raise SystemExit(main()) |
|
|
