single-run trainer

train.py

@@ -0,0 +1,197 @@
+"""Train one pointer-chasing model (fixed-depth or recurrent).
+
+Example:
+  python train.py --arch rec --depth 1 --iters 4 --seed 0 \
+      --run-name rec1x4-seed0 --hub-repo eac123/ptr-recurrent-depth \
+      --space-id <trackio space id>
+
+Trains on fresh random mappings every batch (no fixed train set), k ~ U{1..8},
+and logs to trackio. Saves a checkpoint + training history JSON locally and
+uploads both to the Hub repo.
+"""
+
+import argparse
+import json
+import os
+import time
+
+import numpy as np
+import torch
+import torch.nn.functional as F
+
+from ptr.common import sample_train_batch, sample_eval_batch
+from ptr.flops import count_params, flops_budget_summary, step_flops
+from ptr.models import PtrTransformer
+
+PROBE_KS = (16, 32, 64)
+
+
+def make_parser():
+    p = argparse.ArgumentParser()
+    p.add_argument("--arch", choices=["fixed", "rec"], required=True)
+    p.add_argument("--depth", type=int, default=4)
+    p.add_argument("--iters", default="1",
+                   help="int, or 'dynamic' for R ~ U{2..8} per batch")
+    p.add_argument("--d-model", type=int, default=128)
+    p.add_argument("--S", type=int, default=64)
+    p.add_argument("--k-max", type=int, default=8)
+    p.add_argument("--steps", type=int, default=25000)
+    p.add_argument("--batch", type=int, default=256)
+    p.add_argument("--lr", type=float, default=3e-4)
+    p.add_argument("--warmup", type=int, default=1000)
+    p.add_argument("--seed", type=int, default=0)
+    p.add_argument("--run-name", required=True)
+    p.add_argument("--out-dir", default="out")
+    p.add_argument("--hub-repo", default=None)
+    p.add_argument("--space-id", default=None,
+                   help="trackio space id for the metrics dashboard")
+    p.add_argument("--project", default="ptr-recurrent-depth")
+    p.add_argument("--eval-every", type=int, default=2500)
+    p.add_argument("--eval-n", type=int, default=400)
+    return p
+
+
+def resolve_iters(args):
+    if args.iters == "dynamic":
+        return ("dynamic", 2, 8)
+    return int(args.iters)
+
+
+def sample_iters(spec, rng):
+    if isinstance(spec, tuple):
+        return int(rng.integers(spec[1], spec[2] + 1))
+    return spec
+
+
+def lr_at(step, args):
+    if step < args.warmup:
+        return args.lr * (step + 1) / args.warmup
+    t = (step - args.warmup) / max(1, args.steps - args.warmup)
+    return args.lr * 0.5 * (1 + np.cos(np.pi * t))
+
+
+@torch.no_grad()
+def probe(model, rng, device, S, ks, n, n_iters):
+    accs = {}
+    for k in ks:
+        b = sample_eval_batch(rng, n, S, k=k)
+        f = torch.as_tensor(b["f"], device=device)
+        x = torch.as_tensor(b["x"], device=device)
+        kt = torch.as_tensor(b["k"], device=device)
+        y = torch.as_tensor(b["answer"], device=device)
+        logits = model(f, x, kt, n_iters=n_iters)
+        accs[k] = float((logits.argmax(-1) == y).float().mean())
+    return accs
+
+
+def main():
+    args = make_parser().parse_args()
+    os.makedirs(args.out_dir, exist_ok=True)
+    torch.manual_seed(args.seed)
+    np.random.seed(args.seed)
+    rng = np.random.default_rng(args.seed)
+    device = "cuda" if torch.cuda.is_available() else "cpu"
+    iters_spec = resolve_iters(args)
+
+    model = PtrTransformer(args.S, args.d_model, 4, 4 * args.d_model,
+                           depth=args.depth, zero_init=True).to(device)
+    nparams = count_params(model)
+    budget = flops_budget_summary(model, args.S, args.batch, iters_spec,
+                                  args.steps)
+    print(f"[budget] {json.dumps(budget)}")
+    # cheap assertions up front
+    assert budget["params_total"] > 0 and budget["train_flops_total"] > 0
+
+    opt = torch.optim.AdamW(model.parameters(), lr=args.lr, weight_decay=0.01)
+
+    trackio = None
+    if args.space_id:
+        import trackio
+        trackio.init(project=args.project, space_id=args.space_id,
+                     config={**vars(args), **budget})
+
+    history = {"loss": [], "id_acc": [], "ood": [], "iters": [],
+               "flops_cum": 0.0}
+    flops_cum = 0.0
+    t0 = time.time()
+    bytes_per_batch = 0
+    for step in range(args.steps):
+        lr = lr_at(step, args)
+        for g in opt.param_groups:
+            g["lr"] = lr
+        n_iters = sample_iters(iters_spec, rng)
+        b = sample_train_batch(rng, args.batch, args.S, k_max=args.k_max)
+        f = torch.as_tensor(b["f"], device=device)
+        x = torch.as_tensor(b["x"], device=device)
+        k = torch.as_tensor(b["k"], device=device)
+        y = torch.as_tensor(b["answer"], device=device)
+        logits = model(f, x, k, n_iters=n_iters)
+        loss = F.cross_entropy(logits, y)
+        opt.zero_grad(set_to_none=True)
+        loss.backward()
+        torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
+        opt.step()
+        flops_cum += step_flops(model, args.S, args.batch, n_iters, True)
+
+        if step % 100 == 0 or step == args.steps - 1:
+            ips = 100 / max(1e-9, time.time() - t0) if step > 0 else 0.0
+            t0 = time.time()
+            rec = {"step": step, "loss": round(float(loss), 4), "lr": lr,
+                   "iters": n_iters, "steps_per_s": round(ips, 2),
+                   "flops_cum": flops_cum}
+            history["loss"].append(rec)
+            if device == "cuda" and step == 100:
+                bytes_per_batch = torch.cuda.max_memory_allocated()
+                print(f"[mem] peak GPU memory {bytes_per_batch/1e6:.1f} MB")
+            if trackio:
+                trackio.log({"train/loss": float(loss), "train/lr": lr,
+                             "train/steps_per_s": ips,
+                             "train/flops_cum": flops_cum,
+                             "train/iters": n_iters}, step=step)
+        if (step + 1) % args.eval_every == 0 or step == args.steps - 1:
+            n_eval = sample_iters(iters_spec, rng) if args.arch == "rec" \
+                else 1
+            id_accs = probe(model, rng, device, args.S,
+                            list(range(1, args.k_max + 1)), args.eval_n,
+                            n_eval)
+            ood = probe(model, rng, device, args.S, PROBE_KS, args.eval_n,
+                        n_eval)
+            id_mean = float(np.mean(list(id_accs.values())))
+            print(f"[eval] step {step+1} id_acc(k=1..8) {id_mean:.4f} "
+                  f"ood {ood}")
+            history["id_acc"].append({"step": step + 1, "mean": id_mean,
+                                      "per_k": id_accs})
+            history["ood"].append({"step": step + 1, "acc": ood})
+            if trackio:
+                trackio.log({"eval/id_acc_mean": id_mean,
+                             "eval/ood_k16": ood[16], "eval/ood_k32": ood[32],
+                             "eval/ood_k64": ood[64]}, step=step + 1)
+            history["flops_cum"] = flops_cum
+
+    ckpt_path = os.path.join(args.out_dir, f"{args.run_name}.pt")
+    torch.save({"state_dict": model.state_dict(),
+                "config": vars(args), "budget": budget,
+                "n_iters_spec": iters_spec}, ckpt_path)
+    hist_path = os.path.join(args.out_dir, f"{args.run_name}_history.json")
+    with open(hist_path, "w") as fh:
+        json.dump(history, fh, indent=1)
+
+    if args.hub_repo:
+        from huggingface_hub import HfApi
+        api = HfApi()
+        api.upload_file(path_or_fileobj=ckpt_path,
+                        path_in_repo=f"ckpts/{args.run_name}.pt",
+                        repo_id=args.hub_repo, repo_type="model")
+        api.upload_file(path_or_fileobj=hist_path,
+                        path_in_repo=f"results/{args.run_name}_history.json",
+                        repo_id=args.hub_repo, repo_type="model")
+        print(f"[hub] uploaded ckpt + history to {args.hub_repo}")
+    if trackio:
+        trackio.finish()
+    print(f"[done] {args.run_name}: final loss "
+          f"{history['loss'][-1]['loss']}, "
+          f"id_acc {history['id_acc'][-1]['mean']:.4f}")
+
+
+if __name__ == "__main__":
+    main()