diff options
Diffstat (limited to 'src/models/transformer.py')
| -rw-r--r-- | src/models/transformer.py | 8 |
1 files changed, 6 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: |