path: root/src/train.py

"""Training loop. Usage: python -m src.train [model] [seed] [--flag ...]"""
import csv
import json
import math
import os
import random
import sys

import numpy as np
import torch
import torch.nn as nn

from src.config import Config, parse_args
from src.data import build_examples, decode_tokens, get_splits, make_batch
from src.model_api import build_model


def set_seed(s: int) -> None:
    random.seed(s)
    np.random.seed(s)
    torch.manual_seed(s)


def resolve_device(device_str: str) -> torch.device:
    if device_str == "auto":
        return torch.device("cuda" if torch.cuda.is_available() else "cpu")
    return torch.device(device_str)


@torch.no_grad()
def evaluate(model, examples, cfg: Config, device: torch.device | None = None):
    """Token accuracy (teacher-forced) + exact-match accuracy (batched greedy) + per-example detail.

    BATCHED & CHUNKED: examples evaluated in chunks of eval_batch_size. Valid because
    models are batch-invariant by construction (global fixed layout, pad no-ops).
    Chunking prevents GPU VRAM OOM on large datasets (e.g. range_end = 100k)."""
    if device is None:
        try:
            device = next(model.parameters()).device
        except StopIteration:
            device = torch.device("cpu")
    model.eval()
    if not examples:
        return 0.0, 0.0, []

    eval_bs = getattr(cfg, "eval_batch_size", 512)
    use_amp = getattr(cfg, "use_amp", True) and device.type == "cuda"
    token_correct = 0
    token_total = 0
    em_correct = 0
    per_example = []

    for bi in range(0, len(examples), eval_bs):
        chunk = examples[bi: bi + eval_bs]
        batch = make_batch(chunk, cfg)
        x = batch["x"].to(device)
        y = batch["y"].to(device)
        y_in_eval = batch["y_in"].to(device)
        mask = batch["y_mask"].to(device)

        with torch.amp.autocast(device_type="cuda", dtype=torch.float16, enabled=use_amp):
            out = model(x, y_in_eval)
            logits = out["logits"]                                  # (B,T_out,vocab)
        pred_tok = logits.argmax(-1)
        token_correct += int((pred_tok[mask] == y[mask]).sum().item())
        token_total += int(mask.sum().item())

        # batched greedy decode on device
        B = x.shape[0]
        max_out_len = getattr(cfg, "max_out_len", 6)
        from src.data import _global_lengths
        _, out_max = _global_lengths(cfg)
        max_len = max(max_out_len, out_max)

        y_in = torch.full((B, 1), cfg.eos_id, dtype=torch.long, device=device)
        with torch.amp.autocast(device_type="cuda", dtype=torch.float16, enabled=use_amp):
            for _ in range(max_len):
                o = model(x, y_in)
                nxt = o["logits"][:, -1].argmax(-1)
                y_in = torch.cat([y_in, nxt[:, None]], dim=1)
        gen = y_in[:, 1:].cpu()

        for i, (xrow, target) in enumerate(chunk):
            pred = decode_tokens(gen[i].tolist(), cfg)
            target_n = decode_tokens(target, cfg)
            ok = pred == target_n
            em_correct += int(ok)
            per_example.append((tuple(xrow), target_n, pred, ok))

    model.train()
    return token_correct / max(1, token_total), em_correct / len(examples), per_example


def log_example_rows(log_examples, val_per_example):
    """Format the fixed 10 logged val inputs as '42:ok' strings, in fixed order."""
    by_x = {x: (target, pred, ok) for x, target, pred, ok in val_per_example}
    rows = []
    for x, _target in log_examples:
        t, p, ok = by_x.get(tuple(x), (-1, -1, False))
        label = "".join(map(str, x))  # "42" in both vocab modes
        rows.append(f"{label}:{'ok' if ok else f'{p}~{t}'}")
    return ";".join(rows)


def main() -> None:
    cfg = parse_args()
    set_seed(cfg.seed)
    device = resolve_device(cfg.device)
    use_amp = cfg.use_amp and device.type == "cuda"
    out_dir = os.path.join(cfg.out_dir, cfg.model, f"seed{cfg.seed}")
    csv_path = os.path.join(out_dir, "metrics.csv")
    if os.path.exists(csv_path):
        raise SystemExit(f"REFUSING to rerun in place: {csv_path} exists. Use a fresh --out_dir "
                         f"(reruns would corrupt the CSV and checkpoint provenance).")
    os.makedirs(out_dir, exist_ok=True)
    cfg.save(os.path.join(out_dir, "config.json"))
    meta = {
        "python": sys.version.split()[0],
        "torch": torch.__version__,
        "numpy": np.__version__,
        "device": str(device),
        "use_amp": use_amp,
        "compile_model": cfg.compile_model,
        "torch_threads": torch.get_num_threads(),
        "cmd": sys.argv,
    }
    if device.type == "cuda":
        meta["gpu_name"] = torch.cuda.get_device_name(device)
        meta["gpu_capability"] = list(torch.cuda.get_device_capability(device))
    with open(os.path.join(out_dir, "run_meta.json"), "w") as fh:
        json.dump(meta, fh, indent=2)

    train_in, val_in = get_splits(cfg)
    train_ex = build_examples(train_in, cfg)
    val_ex = build_examples(val_in, cfg)
    model = build_model(cfg).to(device)
    if cfg.compile_model and hasattr(torch, "compile") and device.type == "cuda":
        try:
            model = torch.compile(model, mode="reduce-overhead")
        except Exception as e:
            print(f"Warning: torch.compile failed ({e}), falling back to uncompiled model")

    print(f"model={cfg.model} params={model.param_count()} device={device} amp={use_amp} "
          f"train={len(train_ex)} val={len(val_ex)} "
          f"vocab={cfg.vocab} eos={cfg.eos_id} pad={cfg.pad_id}")

    opt_kwargs = {"lr": cfg.lr, "weight_decay": cfg.weight_decay}
    if device.type == "cuda":
        try:
            opt = torch.optim.AdamW(model.parameters(), fused=True, **opt_kwargs)
        except Exception:
            opt = torch.optim.AdamW(model.parameters(), **opt_kwargs)
    else:
        opt = torch.optim.AdamW(model.parameters(), **opt_kwargs)
    scaler = torch.amp.GradScaler("cuda", enabled=use_amp)
    ce = nn.CrossEntropyLoss(reduction="none")

    log_examples = sorted(val_ex, key=lambda ex: (len(ex[0]), ex[0]))[: cfg.log_n_examples]

    fieldnames = ["step", "train_loss", "train_token_acc", "train_em", "val_token_acc",
                  "val_em", "mean_halt_steps", "log_examples", "param_count"]
    best_val_em = -1.0
    patience_left = cfg.early_stop_patience
    best_path = os.path.join(out_dir, "best.pt")

    n_batches = math.ceil(len(train_ex) / cfg.batch_size)
    rng = random.Random(cfg.seed)
    step = 0
    done = False
    first_row = True

    def _eval_pass(cur_loss, halt_steps):
        nonlocal best_val_em, patience_left, done, first_row
        train_tok, train_em, _ = evaluate(model, train_ex, cfg)
        val_tok, val_em, val_per = evaluate(model, val_ex, cfg)
        mean_halt = float(halt_steps.mean()) if halt_steps is not None else float("nan")
        row = {
            "step": step, "train_loss": float(cur_loss), "train_token_acc": train_tok,
            "train_em": train_em, "val_token_acc": val_tok, "val_em": val_em,
            "mean_halt_steps": mean_halt,
            "log_examples": log_example_rows(log_examples, val_per),
            "param_count": model.param_count(),
        }
        with open(csv_path, "a", newline="") as fh:
            w = csv.DictWriter(fh, fieldnames=fieldnames)
            if first_row:
                w.writeheader()
                first_row = False
            w.writerow(row)
        if val_em > best_val_em:
            best_val_em = val_em
            torch.save(model.state_dict(), best_path)
        if val_em >= cfg.early_stop_em:
            patience_left -= 1
        else:
            patience_left = cfg.early_stop_patience
        if patience_left <= 0:
            done = True
        print(f"step {step} loss {float(cur_loss):.4f} train_em {train_em:.3f} "
              f"val_em {val_em:.3f} val_tok {val_tok:.3f} halt {mean_halt:.2f}")

    while step < cfg.max_train_steps and not done:
        rng.shuffle(train_ex)
        for bi in range(n_batches):
            sl = train_ex[bi * cfg.batch_size: (bi + 1) * cfg.batch_size]
            if not sl:
                continue
            batch = make_batch(sl, cfg)
            x = batch["x"].to(device)
            y = batch["y"].to(device)
            y_in = batch["y_in"].to(device)
            mask = batch["y_mask"].to(device)

            with torch.amp.autocast(device_type="cuda", dtype=torch.float16, enabled=use_amp):
                out = model(x, y_in)
                logits = out["logits"]
                y_safe = y.clamp(max=cfg.vocab - 1)   # CE index guard for pad positions
                loss_tokens = ce(logits.reshape(-1, cfg.vocab), y_safe.reshape(-1)).reshape(
                    logits.shape[0], -1) * mask.float()
                loss_tokens = loss_tokens.sum() / mask.sum().clamp(min=1)
                if cfg.halting and step >= cfg.halt_warmup_steps:
                    if step >= cfg.halt_ramp_end_steps:
                        lam = cfg.halt_penalty
                    else:
                        frac = (step - cfg.halt_warmup_steps) / max(1, cfg.halt_ramp_end_steps - cfg.halt_warmup_steps)
                        lam = cfg.halt_penalty * frac
                else:
                    lam = 0.0
                halt = out["halt_steps"]
                penalty = halt.float().mean() * lam if halt is not None and lam > 0 else 0.0
                loss = loss_tokens + penalty

            if cfg.lr_decay:
                # cosine 1e-3 -> 1e-4 over the full run (E3, locked in Addendum 4)
                frac = min(1.0, step / cfg.max_train_steps)
                lr_t = cfg.lr * 0.1 + 0.5 * (cfg.lr - cfg.lr * 0.1) * (1 + math.cos(math.pi * frac))
                opt.param_groups[0]["lr"] = lr_t

            opt.zero_grad()
            scaler.scale(loss).backward()
            scaler.step(opt)
            scaler.update()

            step += 1
            if step % cfg.eval_every == 0 or step >= cfg.max_train_steps:
                _eval_pass(loss.detach(), halt.detach() if halt is not None else None)
                if done:
                    break
            if step >= cfg.max_train_steps:
                break

    torch.save(model.state_dict(), os.path.join(out_dir, "last.pt"))
    print(f"DONE steps={step} best_val_em={best_val_em:.4f}")


if __name__ == "__main__":
    main()