feat(gpu): add CUDA support, chunked eval, bisect prime search, and scaling to 100k (48 tests green)

1 files changed, 90 insertions, 42 deletions

diff --git a/src/train.py b/src/train.py
index 406defd..e8870df 100644
--- a/src/train.py
+++ b/src/train.py
@@ -21,38 +21,69 @@ def set_seed(s: int) -> None:
     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):
+def evaluate(model, examples, cfg: Config, device: torch.device | None = None):
     """Token accuracy (teacher-forced) + exact-match accuracy (batched greedy) + per-example detail.
 
-    BATCHED: all examples in ONE forward pass (greedy decode batch-wide). Valid because
-    models are batch-invariant by construction (global fixed layout, pad no-ops — see
-    tests/test_codex_fixes.py). ~50x fewer forwards than the per-example version."""
+    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()
-    batch = make_batch(examples, cfg)
-    out = model(batch["x"], batch["y_in"])
-    logits = out["logits"]                                  # (B,T_out,vocab)
-    y = batch["y"]
-    mask = batch["y_mask"]
-    pred_tok = logits.argmax(-1)
-    token_correct = int((pred_tok[mask] == y[mask]).sum())
-    token_total = int(mask.sum())
-    # batched greedy decode
-    B = batch["x"].shape[0]
-    y_in = torch.full((B, 1), cfg.eos_id, dtype=torch.long)
-    for _ in range(cfg.max_out_len):
-        o = model(batch["x"], y_in)
-        nxt = o["logits"][:, -1].argmax(-1)
-        y_in = torch.cat([y_in, nxt[:, None]], dim=1)
-    gen = y_in[:, 1:]                                        # (B, max_out_len)
+    if not examples:
+        return 0.0, 0.0, []
+
+    eval_bs = getattr(cfg, "eval_batch_size", 512)
+    token_correct = 0
+    token_total = 0
     em_correct = 0
     per_example = []
-    for i, (xrow, target) in enumerate(examples):
-        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))
+
+    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)
+
+        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)
+        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
 
@@ -62,7 +93,7 @@ def log_example_rows(log_examples, val_per_example):
     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[tuple(x)]
+        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)
@@ -71,6 +102,7 @@ def log_example_rows(log_examples, val_per_example):
 def main() -> None:
     cfg = parse_args()
     set_seed(cfg.seed)
+    device = resolve_device(cfg.device)
     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):
@@ -78,24 +110,36 @@ def main() -> None:
                          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),
+        "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({
-            "python": sys.version.split()[0],
-            "torch": torch.__version__,
-            "numpy": np.__version__,
-            "device": "cpu",
-            "torch_threads": torch.get_num_threads(),
-            "cmd": sys.argv,
-        }, fh, indent=2)
+        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)
-    print(f"model={cfg.model} params={model.param_count()} train={len(train_ex)} val={len(val_ex)} "
+    model = build_model(cfg).to(device)
+    print(f"model={cfg.model} params={model.param_count()} device={device} "
+          f"train={len(train_ex)} val={len(val_ex)} "
           f"vocab={cfg.vocab} eos={cfg.eos_id} pad={cfg.pad_id}")
 
-    opt = torch.optim.AdamW(model.parameters(), lr=cfg.lr, weight_decay=cfg.weight_decay)
+    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)
     ce = nn.CrossEntropyLoss(reduction="none")
 
     log_examples = sorted(val_ex, key=lambda ex: (len(ex[0]), ex[0]))[: cfg.log_n_examples]
@@ -149,12 +193,16 @@ def main() -> None:
             if not sl:
                 continue
             batch = make_batch(sl, cfg)
-            out = model(batch["x"], batch["y_in"])
+            x = batch["x"].to(device)
+            y = batch["y"].to(device)
+            y_in = batch["y_in"].to(device)
+            mask = batch["y_mask"].to(device)
+            out = model(x, y_in)
             logits = out["logits"]
-            y_safe = batch["y"].clamp(max=cfg.vocab - 1)   # CE index guard for pad positions
+            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) * batch["y_mask"].float()
-            loss_tokens = loss_tokens.sum() / batch["y_mask"].sum().clamp(min=1)
+                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