Le modèle Transformer.
Fichier : packages/minilab/model.py
Votre tokenizer transforme le roman en identifiants. Il faut maintenant le modèle qui apprendra à prédire le suivant. Cette étape rassemble six leçons : les embeddings (03), la prédiction du token suivant (04), les positions (05), l’attention et son masque causal (06), le bloc Transformer (07) et la tête de langage (08). Vous allez les écrire vous-même, sans aucune couche toute faite : PyTorch ne fournit que les tenseurs et les opérations élémentaires.
Comprendre l’architecture du fichier
Le modèle reçoit un lot de séquences d’identifiants, de forme (lot, tokens), et rend un logit par token du vocabulaire pour chaque position, de forme (lot, tokens, 1024). Si on lui donne aussi les cibles, il calcule la loss d’entropie croisée de la leçon 10.
identifiants → embedding du token + embedding de la position
→ 4 blocs : h = x + Attention(LayerNorm(x))
sortie = h + FeedForward(LayerNorm(h))
→ LayerNorm finale → logits = h · Eᵀ (partage des poids)Les dimensions sont choisies pour le processeur : d_model = 128, découpé en 4 têtes de 32 dimensions (d_k), 4 blocs, un feed-forward de 512 neurones (4 × d_model, comme GPT-2) et une fenêtre de contexte de 128 tokens. Cela donne environ 940 000 paramètres : assez pour apprendre la forme du français de Jules Verne, assez peu pour s’entraîner en une minute.
Le fichier contient aussi deux pièces qui ne servent pas tout de suite. Le dropout (leçon 11) vaut 0 par défaut ; l’étape 3 l’activera pendant l’entraînement. Le cache KV (leçon 15) n’est utilisé que si on le passe au modèle ; l’étape 4 s’en servira. Les écrire dès maintenant vous évite de revenir sur ce fichier.
Construire pas à pas : le cache et la configuration
KVCache garde, pour chaque couche, les clés et les valeurs déjà calculées, et compte combien de tokens sont passés par les projections. Config regroupe tous les hyperparamètres ; deux propriétés en déduisent d_k et d_ff, pour qu’ils restent toujours cohérents avec d_model.
"""Étape 2 — un petit Transformer causal, écrit pièce par pièce (leçons 03 à 08).
tokens → embeddings + positions → n_layer blocs → LayerNorm finale → logits
Chaque bloc est pre-norm, comme dans GPT-2 (leçon 07) :
h = x + Attention(LayerNorm(x)) sortie = h + FeedForward(LayerNorm(h))
PyTorch ne fournit que les tenseurs, les gradients automatiques et les couchesélémentaires (Embedding, Linear, LayerNorm). L'attention elle-même, avec son masquecausal, est écrite à la main : ni nn.MultiheadAttention, ni nn.TransformerEncoder.
uv run python -m packages.minilab.model"""
import mathfrom dataclasses import asdict, dataclass
import torchfrom torch import nnfrom torch.nn import functional as F
class KVCache: """Leçon 15 : les clés et les valeurs déjà calculées, gardées couche par couche. Avec le masque causal, elles ne changent jamais : chaque nouveau token ne calcule plus que ses propres q, k et v. Ce cache servira à l'étape 4."""
def __init__(self, n_layer: int) -> None: self.keys: list[torch.Tensor | None] = [None] * n_layer self.values: list[torch.Tensor | None] = [None] * n_layer self.length = 0 # nombre de tokens déjà dans le cache self.projected = 0 # tokens passés par W_Q, W_K, W_V, toutes couches comprises
@dataclass(frozen=True)class Config: vocab_size: int = 1024 context_length: int = 128 # la fenêtre de contexte (leçon 05) d_model: int = 128 n_head: int = 4 n_layer: int = 4 dropout: float = 0.0 # part des neurones éteints pendant l'entraînement (leçon 11)
@property def d_k(self) -> int: # dimensions d'une tête : d_model = n_head × d_k return self.d_model // self.n_head
@property def d_ff(self) -> int: # la couche cachée du feed-forward (leçon 07) return 4 * self.d_modelConstruire pas à pas : l’attention causale multi-têtes
C’est le cœur du modèle, la leçon 06 en code. Les quatre matrices W_Q, W_K, W_V et W_O sont des couches nn.Linear sans biais. split_heads découpe chaque vecteur de 128 nombres en 4 tranches de 32 : une par tête. Puis on calcule les scores q @ kᵀ / √d_k, on masque l’avenir avec un triangle de -inf, on applique softmax ligne par ligne, et l’on fait la somme pondérée des valeurs.
Le masque mérite un coup d’œil. Sans cache, c’est le triangle de la leçon 06 : triu(1) bloque tout ce qui est au-dessus de la diagonale. Avec un cache, les nouvelles requêtes arrivent après past tokens déjà connus : le triangle est simplement décalé de past colonnes.
class CausalSelfAttention(nn.Module): """Leçon 06 : requêtes, clés et valeurs, produit scalaire mis à l'échelle, masque causal, softmax."""
def __init__(self, config: Config) -> None: super().__init__() if config.d_model % config.n_head: raise ValueError("d_model doit être un multiple de n_head.") self.n_head, self.d_k = config.n_head, config.d_k # Quatre matrices apprises de d_model × d_model (sans biais, comme dans les leçons). self.W_Q = nn.Linear(config.d_model, config.d_model, bias=False) self.W_K = nn.Linear(config.d_model, config.d_model, bias=False) self.W_V = nn.Linear(config.d_model, config.d_model, bias=False) self.W_O = nn.Linear(config.d_model, config.d_model, bias=False) self.dropout = nn.Dropout(config.dropout)
def split_heads(self, x: torch.Tensor) -> torch.Tensor: """(lot, tokens, d_model) → (lot, n_head, tokens, d_k) : une tranche par tête.""" batch, tokens, _ = x.shape return x.view(batch, tokens, self.n_head, self.d_k).transpose(1, 2)
def forward( self, x: torch.Tensor, cache: KVCache | None = None, layer: int = 0 ) -> torch.Tensor: batch, tokens, d_model = x.shape q, k, v = (self.split_heads(W(x)) for W in (self.W_Q, self.W_K, self.W_V)) past = 0 if cache is not None: cache.projected += tokens if cache.keys[layer] is not None: # les nouvelles clés et valeurs après les anciennes past = cache.keys[layer].shape[2] k = torch.cat([cache.keys[layer], k], dim=2) v = torch.cat([cache.values[layer], v], dim=2) cache.keys[layer], cache.values[layer] = k, v # Chaque requête contre chaque clé, divisé par √d_k. scores = q @ k.transpose(-2, -1) / math.sqrt(self.d_k) # Masque causal : un token ne regarde jamais les tokens qui le suivent. Avec un # cache, les nouvelles requêtes sont aux positions past, past + 1… et voient # toutes les clés jusqu'à la leur. future = torch.ones(tokens, past + tokens, dtype=torch.bool, device=x.device).triu(past + 1) scores = scores.masked_fill(future, float("-inf")) weights = self.dropout(F.softmax(scores, dim=-1)) # la matrice d'attention, ligne par ligne heads = weights @ v # la somme pondérée des valeurs # On met les têtes bout à bout, puis W_O les fusionne. return self.dropout(self.W_O(heads.transpose(1, 2).reshape(batch, tokens, d_model)))Construire pas à pas : le bloc et le modèle complet
Le feed-forward élargit à 512 neurones, applique GELU, puis revient à 128 (leçon 07). Le bloc assemble les deux sous-couches en pre-norm, chacune entourée de sa connexion résiduelle. Le modèle complet additionne l’embedding du token et celui de sa position, traverse les blocs, puis la LayerNorm finale.
Deux détails viennent des leçons. Les poids sont initialisés petits, avec un écart type de 0,02 comme GPT-2 : au départ, toutes les prédictions sont presque égales. Et la tête de langage réutilise la table d’embeddings (leçon 08) : le logit d’un token est le produit scalaire entre le vecteur final et l’embedding de ce token. Elle n’ajoute donc aucun paramètre.
class FeedForward(nn.Module): """Leçon 07 : élargir (d_model → 4 × d_model), GELU, puis rétrécir."""
def __init__(self, config: Config) -> None: super().__init__() self.W_1 = nn.Linear(config.d_model, config.d_ff) self.W_2 = nn.Linear(config.d_ff, config.d_model) self.dropout = nn.Dropout(config.dropout)
def forward(self, x: torch.Tensor) -> torch.Tensor: return self.dropout(self.W_2(F.gelu(self.W_1(x), approximate="tanh")))
class Block(nn.Module): """Leçon 07 : un bloc pre-norm avec deux connexions résiduelles."""
def __init__(self, config: Config) -> None: super().__init__() self.ln_1 = nn.LayerNorm(config.d_model) self.attention = CausalSelfAttention(config) self.ln_2 = nn.LayerNorm(config.d_model) self.feed_forward = FeedForward(config)
def forward( self, x: torch.Tensor, cache: KVCache | None = None, layer: int = 0 ) -> torch.Tensor: h = x + self.attention(self.ln_1(x), cache, layer) return h + self.feed_forward(self.ln_2(h))
class MiniLabModel(nn.Module): def __init__(self, config: Config) -> None: super().__init__() self.config = config self.token_embedding = nn.Embedding(config.vocab_size, config.d_model) # leçon 03 self.position_embedding = nn.Embedding(config.context_length, config.d_model) # leçon 05 self.blocks = nn.ModuleList(Block(config) for _ in range(config.n_layer)) self.ln_f = nn.LayerNorm(config.d_model) # LayerNorm finale (leçon 08) self.dropout = nn.Dropout(config.dropout) self.apply(self._init_weights)
@staticmethod def _init_weights(module: nn.Module) -> None: # De petits poids aléatoires, comme GPT-2 : les premières prédictions sont # presque uniformes. if isinstance(module, nn.Linear | nn.Embedding): nn.init.normal_(module.weight, mean=0.0, std=0.02) if isinstance(module, nn.Linear) and module.bias is not None: nn.init.zeros_(module.bias)
def forward( self, ids: torch.Tensor, targets: torch.Tensor | None = None, cache: KVCache | None = None ) -> tuple[torch.Tensor, torch.Tensor | None]: """Les logits de chaque position ; avec des cibles, la loss d'entropie croisée (leçon 10). Les cibles égales à -100 sont ignorées : c'est le masquage de la loss de la leçon 12. Avec un cache, ids ne contient que les nouveaux tokens, placés après ceux du cache.""" tokens = ids.shape[1] past = cache.length if cache is not None else 0 if past + tokens > self.config.context_length: raise ValueError("La séquence dépasse la fenêtre de contexte.") positions = torch.arange(past, past + tokens, device=ids.device) x = self.dropout(self.token_embedding(ids) + self.position_embedding(positions)) for layer, block in enumerate(self.blocks): x = block(x, cache, layer) if cache is not None: cache.length += tokens # Tête de langage avec partage des poids (leçon 08) : logit = h · E[token]. logits = self.ln_f(x) @ self.token_embedding.weight.T loss = None if targets is not None: loss = F.cross_entropy(logits.reshape(-1, logits.shape[-1]), targets.reshape(-1)) return logits, lossConstruire pas à pas : sauvegarde et démonstration
save et load enregistrent la configuration avec les poids : un checkpoint se recharge sans avoir à se souvenir des dimensions. main compte les paramètres couche par couche et mesure la loss d’un modèle qui n’a encore rien appris.
def count_parameters(module: nn.Module) -> int: return sum(p.numel() for p in module.parameters())
def save(model: MiniLabModel, path) -> None: path.parent.mkdir(parents=True, exist_ok=True) torch.save({"config": asdict(model.config), "state": model.state_dict()}, path)
def load(path) -> MiniLabModel: checkpoint = torch.load(path, map_location="cpu", weights_only=True) model = MiniLabModel(Config(**checkpoint["config"])) model.load_state_dict(checkpoint["state"]) return model.eval()
def main() -> None: from packages.minilab.tokenizer import get_tokenizer, read_corpus
torch.manual_seed(0) tokenizer = get_tokenizer() config = Config(vocab_size=tokenizer.vocab_size) model = MiniLabModel(config) print(f"Configuration : {asdict(config)} (d_k = {config.d_k}, d_ff = {config.d_ff})\n") rows = [ ("Embeddings des tokens", model.token_embedding), ("Embeddings de position", model.position_embedding), ] rows += [(f"Bloc {i + 1}", block) for i, block in enumerate(model.blocks)] rows += [("LayerNorm finale", model.ln_f)] block = model.blocks[0] for name, module in rows: print(f" {name:<24} {count_parameters(module):>9,}".replace(",", " ")) if module is block: attention, feed_forward = ( count_parameters(block.attention), count_parameters(block.feed_forward), ) print(f" dont attention {attention:>9,} (4 × d_model²)".replace(",", " ")) print( f" dont feed-forward {feed_forward:>9,} (≈ 8 × d_model²)".replace(",", " ") ) print(f" {'Tête de langage':<24} {0:>9} (partage les poids des embeddings)") print(f" {'Total':<24} {count_parameters(model):>9,}".replace(",", " "))
ids = torch.tensor(tokenizer.encode(read_corpus()[:20_000])) batch = torch.stack([ids[i : i + config.context_length + 1] for i in range(0, 8 * 500, 500)]) with torch.no_grad(): _, loss = model(batch[:, :-1], batch[:, 1:]) print(f"\nLoss initiale, poids aléatoires : {loss.item():.3f}") print(f"Attendu, au hasard : ln({config.vocab_size}) = {math.log(config.vocab_size):.3f}") print( f"Perplexité initiale : {math.exp(loss.item()):.0f} (sur {config.vocab_size} tokens possibles)" )
if __name__ == "__main__": main()Tester et observer dans le terminal
- Lancer l’étape
uv run python -m packages.minilab.model - Lancer les tests de l’étape
uv run pytest tests/minilab/test_model.py
Configuration : {'vocab_size': 1024, 'context_length': 128, 'd_model': 128, 'n_head': 4, 'n_layer': 4, 'dropout': 0.0} (d_k = 32, d_ff = 512)
Embeddings des tokens 131 072
Embeddings de position 16 384
Bloc 1 197 760
dont attention 65 536 (4 × d_model²)
dont feed-forward 131 712 (≈ 8 × d_model²)
Bloc 2 197 760
Bloc 3 197 760
Bloc 4 197 760
LayerNorm finale 256
Tête de langage 0 (partage les poids des embeddings)
Total 938 752
Loss initiale, poids aléatoires : 6.952
Attendu, au hasard : ln(1024) = 6.931
Perplexité initiale : 1045 (sur 1024 tokens possibles)Vérifiez les ordres de grandeur de la leçon 07 : l’attention d’un bloc compte exactement 4 × 128² = 65 536 paramètres, le feed-forward environ le double, 131 712, soit 8 × 128² plus ses biais. Au total, 938 752 paramètres. Et la loss initiale, 6,95, est presque exactement ln(1 024) = 6,93 : le modèle devine au hasard parmi 1 024 tokens, comme le prédisait la leçon 10. Un écart important ici signalerait un bug.
Les tests vérifient aussi que le masque causal fonctionne vraiment : modifier les derniers tokens d’une séquence ne doit pas changer les logits des positions précédentes.
Dans main, passez à Config(vocab_size=tokenizer.vocab_size, n_layer=8) : le nombre de paramètres augmente de 4 × 197 760, et la loss initiale reste proche de 6,93. Essayez aussi n_head=3 : le modèle refuse, car 128 ne se divise pas en 3 têtes égales.
Le modèle additionne les embeddings des tokens et des positions, traverse 4 blocs pre-norm (attention causale multi-têtes, puis feed-forward GELU, chacun avec une connexion résiduelle), puis la LayerNorm finale ; la tête de langage réutilise la table d’embeddings. Avec d_model = 128, 4 têtes de 32 et d_ff = 512 : 938 752 paramètres. Avant l’entraînement, la loss vaut ln(1 024) ≈ 6,93.