speedup: batched eval (46x, sieve-stub verified), run_sweep orchestrator (2-way parallel, idempotent, summaries), E1 jobs; 36 tests

1 files changed, 25 insertions, 17 deletions

diff --git a/src/train.py b/src/train.py
index 46bd8e7..aa0fb4b 100644
--- a/src/train.py
+++ b/src/train.py
@@ -12,7 +12,7 @@ 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, greedy_decode
+from src.model_api import build_model
 
 
 def set_seed(s: int) -> None:
@@ -23,28 +23,36 @@ def set_seed(s: int) -> None:
 
 @torch.no_grad()
 def evaluate(model, examples, cfg: Config):
-    """Token accuracy (teacher-forced) + exact-match accuracy (greedy) + per-example detail."""
+    """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."""
     model.eval()
-    token_correct = 0
-    token_total = 0
+    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)
     em_correct = 0
     per_example = []
-    for x, target in examples:
-        batch = make_batch([(x, target)], cfg)
-        out = model(batch["x"], batch["y_in"])
-        logits = out["logits"][0]
-        y = batch["y"][0]
-        mask = batch["y_mask"][0]
-        pred_tok = logits.argmax(-1)
-        token_correct += int((pred_tok[mask] == y[mask]).sum())
-        token_total += int(mask.sum())
-        xi = torch.tensor(x, dtype=torch.long).unsqueeze(0)
-        gen = greedy_decode(model, xi, cfg)[0].tolist()
-        pred = decode_tokens(gen, cfg)
+    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(x), target_n, pred, 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