perf(train): VRAM-resident tensor caching, zero-copy GPU batch sampling, and adaptive eval schedule

2 files changed, 106 insertions, 73 deletions

diff --git a/src/config.py b/src/config.py
index 4db66b0..945c259 100644
--- a/src/config.py
+++ b/src/config.py
@@ -31,6 +31,7 @@ class Config:
     weight_decay: float = 1.0
     max_train_steps: int = 200_000
     eval_every: int = 200
+    eval_schedule: str = "fixed"     # "fixed" | "adaptive" (more frequent early, coarser later for 4M runs)
     early_stop_em: float = 1.0
     early_stop_patience: int = 5
     batch_size: int = 32
diff --git a/src/train.py b/src/train.py
index ec7bb4e..d8c03de 100644
--- a/src/train.py
+++ b/src/train.py
@@ -27,20 +27,26 @@ def resolve_device(device_str: str) -> torch.device:
     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.
+def make_gpu_dataset(examples: list, cfg: Config, device: torch.device) -> dict[str, torch.Tensor | int]:
+    """Create static tensors directly in VRAM to eliminate per-step host-to-device transfers."""
+    if not examples:
+        return {"size": 0}
+    batch = make_batch(examples, cfg)
+    return {
+        "x": batch["x"].to(device, non_blocking=True),
+        "y": batch["y"].to(device, non_blocking=True),
+        "y_in": batch["y_in"].to(device, non_blocking=True),
+        "y_mask": batch["y_mask"].to(device, non_blocking=True),
+        "size": len(examples),
+    }
 
-    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")
+
+@torch.no_grad()
+def evaluate_gpu(model, gpu_ds: dict[str, torch.Tensor | int], examples: list, cfg: Config, device: torch.device):
+    """Batched evaluation operating directly on pre-allocated VRAM tensors."""
     model.eval()
-    if not examples:
+    N = gpu_ds["size"]
+    if N == 0:
         return 0.0, 0.0, []
 
     eval_bs = getattr(cfg, "eval_batch_size", 512)
@@ -50,22 +56,20 @@ def evaluate(model, examples, cfg: Config, device: torch.device | None = None):
     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)
+    for bi in range(0, N, eval_bs):
+        x = gpu_ds["x"][bi: bi + eval_bs]
+        y = gpu_ds["y"][bi: bi + eval_bs]
+        y_in_eval = gpu_ds["y_in"][bi: bi + eval_bs]
+        mask = gpu_ds["y_mask"][bi: bi + eval_bs]
+        chunk_examples = examples[bi: bi + eval_bs]
 
         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)
+            logits = out["logits"]
         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
@@ -80,7 +84,7 @@ def evaluate(model, examples, cfg: Config, device: torch.device | None = None):
                 y_in = torch.cat([y_in, nxt[:, None]], dim=1)
         gen = y_in[:, 1:].cpu()
 
-        for i, (xrow, target) in enumerate(chunk):
+        for i, (xrow, target) in enumerate(chunk_examples):
             pred = decode_tokens(gen[i].tolist(), cfg)
             target_n = decode_tokens(target, cfg)
             ok = pred == target_n
@@ -88,7 +92,19 @@ def evaluate(model, examples, cfg: Config, device: torch.device | None = None):
             per_example.append((tuple(xrow), target_n, pred, ok))
 
     model.train()
-    return token_correct / max(1, token_total), em_correct / len(examples), per_example
+    return token_correct / max(1, token_total), em_correct / N, per_example
+
+
+@torch.no_grad()
+def evaluate(model, examples, cfg: Config, device: torch.device | None = None):
+    """Token accuracy + exact-match accuracy + per-example detail (backward compatible wrapper)."""
+    if device is None:
+        try:
+            device = next(model.parameters()).device
+        except StopIteration:
+            device = torch.device("cpu")
+    gpu_ds = make_gpu_dataset(examples, cfg, device)
+    return evaluate_gpu(model, gpu_ds, examples, cfg, device)
 
 
 def log_example_rows(log_examples, val_per_example):
@@ -133,6 +149,9 @@ def main() -> None:
     train_in, val_in = get_splits(cfg)
     train_ex = build_examples(train_in, cfg)
     val_ex = build_examples(val_in, cfg)
+    train_gpu = make_gpu_dataset(train_ex, cfg, device)
+    val_gpu = make_gpu_dataset(val_ex, cfg, device)
+
     model = build_model(cfg).to(device)
     if cfg.compile_model and hasattr(torch, "compile") and device.type == "cuda":
         try:
@@ -163,16 +182,30 @@ def main() -> None:
     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)
+    N_train = train_gpu["size"]
+    B = min(cfg.batch_size, N_train) if N_train > 0 else 1
     step = 0
     done = False
     first_row = True
 
+    def should_eval(cur_step: int) -> bool:
+        if cur_step >= cfg.max_train_steps:
+            return True
+        if getattr(cfg, "eval_schedule", "fixed") in ("adaptive", "logarithmic"):
+            if cur_step < 10_000:
+                return cur_step % 200 == 0
+            elif cur_step < 50_000:
+                return cur_step % 1_000 == 0
+            elif cur_step < 500_000:
+                return cur_step % 5_000 == 0
+            else:
+                return cur_step % 10_000 == 0
+        return cur_step % cfg.eval_every == 0
+
     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)
+        train_tok, train_em, _ = evaluate_gpu(model, train_gpu, train_ex, cfg, device)
+        val_tok, val_em, val_per = evaluate_gpu(model, val_gpu, val_ex, cfg, device)
         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,
@@ -200,53 +233,52 @@ def main() -> None:
               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
+        if B >= N_train:
+            x = train_gpu["x"]
+            y = train_gpu["y"]
+            y_in = train_gpu["y_in"]
+            mask = train_gpu["y_mask"]
+        else:
+            idx = torch.randint(0, N_train, (B,), device=device)
+            x = train_gpu["x"][idx]
+            y = train_gpu["y"][idx]
+            y_in = train_gpu["y_in"][idx]
+            mask = train_gpu["y_mask"][idx]
+
+        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:
-                    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:
+                    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 should_eval(step):
+            _eval_pass(loss.detach(), halt.detach() if halt is not None else None)
+            if done:
                 break
 
     torch.save(model.state_dict(), os.path.join(out_dir, "last.pt"))