optimize E8: precompute causal mask buffer, add eval early exit, and enable compile_model in e8.csv

1 files changed, 3 insertions, 0 deletions

diff --git a/src/train.py b/src/train.py
index 2523325..8ce5394 100644
--- a/src/train.py
+++ b/src/train.py
@@ -105,6 +105,9 @@ def evaluate_gpu(model, gpu_ds: dict[str, torch.Tensor | int | bool], examples:
                 o = model(x, y_in)
                 nxt = o["logits"][:, -1].argmax(-1)
                 y_in = torch.cat([y_in, nxt[:, None]], dim=1)
+                # If all sequences in the batch have emitted at least one EOS, we can exit safely
+                if ((y_in[:, 1:] == cfg.eos_id).any(dim=1)).all():
+                    break
         gen = y_in[:, 1:].cpu()
 
         for i, ex in enumerate(chunk_examples):