diff options
Diffstat (limited to 'src')
| -rw-r--r-- | src/models/transformer.py | 8 | ||||
| -rw-r--r-- | src/train.py | 3 |
2 files changed, 9 insertions, 2 deletions
diff --git a/src/models/transformer.py b/src/models/transformer.py index 2b986ce..a128ad9 100644 --- a/src/models/transformer.py +++ b/src/models/transformer.py @@ -36,6 +36,11 @@ class TransformerBaseline(PrimeModel): self.blocks = nn.ModuleList([CausalBlock(d, cfg.n_heads) for _ in range(cfg.n_layers)]) self.ln_f = nn.LayerNorm(d) self.head = nn.Linear(d, cfg.vocab) + self.register_buffer( + "causal_triu", + torch.triu(torch.ones(256, 256, dtype=torch.bool), diagonal=1), + persistent=False, + ) def forward(self, x: torch.Tensor, y_in: torch.Tensor) -> dict: cfg = self.cfg @@ -61,8 +66,7 @@ class TransformerBaseline(PrimeModel): # causal mask: input positions (j < T_in) fully visible; output positions causal # (True = blocked, per torch.nn.MultiheadAttention bool convention) blocked = torch.zeros(T, T, dtype=torch.bool, device=device) - causal = torch.triu(torch.ones(T, T, dtype=torch.bool, device=device), diagonal=1) - blocked[:, T_in:] = causal[:, T_in:] + blocked[:, T_in:] = self.causal_triu[:T, T_in:T] key_pad = seq == cfg.pad_id # (B,T) True = ignore h = e for blk in self.blocks: 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): |