From 9a5912f7764dbc942b56f8114f2c878c1a6d44ff Mon Sep 17 00:00:00 2001 From: CaptainJack2491 Date: Sat, 29 Aug 2026 23:32:30 +0100 Subject: optimize E8: precompute causal mask buffer, add eval early exit, and enable compile_model in e8.csv --- src/train.py | 3 +++ 1 file changed, 3 insertions(+) (limited to 'src/train.py') 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): -- cgit v1.2.3