fix(train): enforce epoch permutation without replacement, add max cache OOM guard, and handle empty split tensors

1 files changed, 132 insertions, 56 deletions

diff --git a/src/train.py b/src/train.py
index d8c03de..a0874c4 100644
--- a/src/train.py
+++ b/src/train.py
@@ -27,10 +27,23 @@ def resolve_device(device_str: str) -> torch.device:
     return torch.device(device_str)
 
 
-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."""
+def make_gpu_dataset(examples: list, cfg: Config, device: torch.device, max_cache_size: int = 10_000) -> dict[str, torch.Tensor | int | bool]:
+    """Create static tensors in VRAM for small/medium datasets. For large datasets (>max_cache_size),
+    signals that chunked CPU streaming should be used to prevent VRAM OOM."""
+    from src.data import _global_lengths
+    in_max, out_max = _global_lengths(cfg)
     if not examples:
-        return {"size": 0}
+        return {
+            "x": torch.empty((0, in_max), dtype=torch.long, device=device),
+            "y": torch.empty((0, out_max), dtype=torch.long, device=device),
+            "y_in": torch.empty((0, out_max), dtype=torch.long, device=device),
+            "y_mask": torch.empty((0, out_max), dtype=torch.bool, device=device),
+            "size": 0,
+            "is_cached": True,
+        }
+    if len(examples) > max_cache_size:
+        return {"size": len(examples), "is_cached": False}
+
     batch = make_batch(examples, cfg)
     return {
         "x": batch["x"].to(device, non_blocking=True),
@@ -38,12 +51,13 @@ def make_gpu_dataset(examples: list, cfg: Config, device: torch.device) -> dict[
         "y_in": batch["y_in"].to(device, non_blocking=True),
         "y_mask": batch["y_mask"].to(device, non_blocking=True),
         "size": len(examples),
+        "is_cached": True,
     }
 
 
 @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."""
+def evaluate_gpu(model, gpu_ds: dict[str, torch.Tensor | int | bool], examples: list, cfg: Config, device: torch.device):
+    """Batched evaluation: operates directly on VRAM tensors if cached, or streams in eval_batch_size chunks."""
     model.eval()
     N = gpu_ds["size"]
     if N == 0:
@@ -55,13 +69,21 @@ def evaluate_gpu(model, gpu_ds: dict[str, torch.Tensor | int], examples: list, c
     token_total = 0
     em_correct = 0
     per_example = []
+    is_cached = gpu_ds.get("is_cached", False)
 
     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]
+        if is_cached:
+            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]
+        else:
+            batch = make_batch(chunk_examples, 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)
@@ -183,11 +205,17 @@ def main() -> None:
     best_path = os.path.join(out_dir, "best.pt")
 
     N_train = train_gpu["size"]
-    B = min(cfg.batch_size, N_train) if N_train > 0 else 1
+    n_batches = math.ceil(N_train / cfg.batch_size) if N_train > 0 else 0
+    rng = random.Random(cfg.seed)
     step = 0
     done = False
     first_row = True
 
+    if N_train == 0:
+        print("Warning: empty training split; skipping training.")
+        torch.save(model.state_dict(), os.path.join(out_dir, "last.pt"))
+        return
+
     def should_eval(cur_step: int) -> bool:
         if cur_step >= cfg.max_train_steps:
             return True
@@ -233,53 +261,101 @@ 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:
-        if B >= N_train:
-            x = train_gpu["x"]
-            y = train_gpu["y"]
-            y_in = train_gpu["y_in"]
-            mask = train_gpu["y_mask"]
+        if train_gpu.get("is_cached", False):
+            perm = torch.randperm(N_train, device=device)
+            for bi in range(n_batches):
+                idx = perm[bi * cfg.batch_size: (bi + 1) * cfg.batch_size]
+                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:
+                            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
+                if step >= cfg.max_train_steps:
+                    break
         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:
-                    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
+            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 should_eval(step):
+                    _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}")