File size: 1,487 Bytes
37d8122 6857ec0 37d8122 6857ec0 37d8122 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 | """Bloc Transformer pre-norm avec résidus — écrit à la main.
Forme pre-norm (justifiée dans docs/architecture.md §2.2) :
x = x + Attention(LayerNorm(x))
x = x + FeedForward(LayerNorm(x))
La normalisation est appliquée *avant* chaque sous-couche. Le gradient traverse
la connexion résiduelle sans être multiplié par une norme, ce qui stabilise
l'entraînement des modèles profonds.
"""
from __future__ import annotations
import torch.nn as nn
from src.model.attention import MultiHeadSelfAttention
from src.model.ffn import FeedForward
class TransformerBlock(nn.Module):
"""Un bloc = attention causale + FFN, chacun précédé d'une LayerNorm."""
def __init__(
self,
n_embd: int,
n_head: int,
dropout: float = 0.0,
bias: bool = False,
ffn_mult: int = 4,
rope_theta: float = 10000.0,
use_rope: bool = True,
causal: bool = True,
):
super().__init__()
self.ln_1 = nn.LayerNorm(n_embd, bias=bias)
self.attn = MultiHeadSelfAttention(
n_embd, n_head, dropout=dropout, bias=bias, rope_theta=rope_theta,
use_rope=use_rope, causal=causal,
)
self.ln_2 = nn.LayerNorm(n_embd, bias=bias)
self.ffn = FeedForward(n_embd, mult=ffn_mult, dropout=dropout, bias=bias)
def forward(self, x, cos=None, sin=None):
x = x + self.attn(self.ln_1(x), cos, sin)
x = x + self.ffn(self.ln_2(x))
return x
|