Le pré-entraînement.
Fichier : packages/minilab/train.py
Votre modèle existe, mais il ne sait rien : sa loss vaut ln(1 024), le pur hasard. Il est temps de l’entraîner sur Jules Verne. Cette étape met en œuvre les deux leçons sur l’apprentissage : la loss, le gradient et la descente de gradient (10), puis les lots, AdamW, le planning du taux d’apprentissage, le jeu de validation et le surapprentissage (11). C’est l’étape la plus longue à exécuter, environ 70 secondes.
Comprendre l’architecture du fichier
L’entrée est le corpus tokenisé, environ 58 000 tokens ; la sortie, le checkpoint checkpoints/minilab/base.pt, notre modèle de base. Entre les deux, une boucle de 800 pas : tirer un lot de 32 fenêtres de 128 tokens, calculer la loss, rétropropager, faire un pas d’AdamW.
Plusieurs choix sont des réglages mesurés, pas des valeurs au hasard. Le premier essai, sans dropout et avec un taux de 0,003, surapprenait dès le pas 200 : la loss de validation remontait pendant que celle d’entraînement plongeait. D’où trois décisions. Un dropout de 0,1 (leçon 11) : la meilleure loss de validation passe de 4,28 à 4,12. Un taux maximal plus doux, 0,0015, avec un warmup de 100 pas puis une décroissance en cosinus. Et l’arrêt précoce : on évalue tous les 100 pas, et le checkpoint gardé est celui de la meilleure loss de validation, pas le dernier.
Le corpus ne fait que 52 000 tokens d’entraînement : en 800 pas de 4 096 tokens, le modèle le relit environ 60 fois. C’est exactement la situation où le jeu de validation devient indispensable.
Construire pas à pas : imports et réglages
Tous les hyperparamètres sont réunis dans TrainingSettings, un dataclass figé : on les lit d’un coup d’œil, et le mode de démonstration du surapprentissage peut en changer deux sans toucher au reste. History garde les points des courbes et les meilleurs poids rencontrés.
"""Étape 3 — le pré-entraînement sur Jules Verne, avec un jeu de validation (leçons 10 et 11).
Auto-supervisé : chaque fenêtre du texte est son propre exercice, entrée = tokens[:-1] etcible = tokens[1:]. La loss est l'entropie croisée, la perplexité vaut exp(loss),l'optimiseur est AdamW, avec un warmup puis une décroissance en cosinus du tauxd'apprentissage. Le checkpoint gardé est celui de la meilleure loss de validation :c'est l'arrêt précoce.
uv run python -m packages.minilab.train # pré-entraînement (environ une minute) uv run python -m packages.minilab.train --overfit # 500 caractères seulement : surapprentissage"""
import argparseimport mathimport timefrom collections.abc import Callablefrom dataclasses import dataclass, field
import torch
from packages.minilab.model import Config, MiniLabModel, savefrom packages.minilab.tokenizer import CHECKPOINTS, Tokenizer, get_tokenizer, read_corpus
BASE_CHECKPOINT = CHECKPOINTS / "base.pt"DROPOUT = 0.1 # choisi sur la loss de validation : 4,12 au lieu de 4,28 sans dropout
@dataclass(frozen=True)class TrainingSettings: steps: int = 800 batch_size: int = 32 learning_rate: float = 1.5e-3 warmup: int = 100 weight_decay: float = 0.1 eval_every: int = 100 eval_batches: int = 8 seed: int = 0
@dataclassclass History: points: list[tuple[int, float, float]] = field(default_factory=list) # pas, train, val best_step: int = 0 best_val: float = math.inf best_state: dict[str, torch.Tensor] | None = NoneConstruire pas à pas : découpe, lots et taux d’apprentissage
split coupe la suite de tokens à 90 % : la fin du roman sert de validation. get_batch tire des fenêtres au hasard et applique le décalage auto-supervisé de la leçon 10 : la cible est l’entrée décalée d’un token. learning_rate calcule le taux de chaque pas : il monte linéairement pendant le warmup, puis descend en cosinus jusqu’au dixième du maximum.
def split(ids: list[int], train_share: float = 0.9) -> tuple[torch.Tensor, torch.Tensor]: """90 % pour apprendre, 10 % mis sous clé pour mesurer la généralisation (leçon 11).""" cut = int(len(ids) * train_share) return torch.tensor(ids[:cut]), torch.tensor(ids[cut:])
def get_batch( data: torch.Tensor, batch_size: int, length: int, generator: torch.Generator) -> tuple[torch.Tensor, torch.Tensor]: """Des fenêtres de length + 1 tokens tirées au hasard : la cible est l'entrée décalée d'un token.""" starts = torch.randint(0, len(data) - length - 1, (batch_size,), generator=generator) windows = torch.stack([data[s : s + length + 1] for s in starts.tolist()]) return windows[:, :-1], windows[:, 1:]
def learning_rate(step: int, settings: TrainingSettings) -> float: """Le warmup, puis une décroissance en cosinus jusqu'au dixième du taux maximal.""" if step < settings.warmup: return settings.learning_rate * (step + 1) / settings.warmup progress = (step - settings.warmup) / max(1, settings.steps - settings.warmup) return settings.learning_rate * (0.1 + 0.45 * (1 + math.cos(math.pi * progress)))Construire pas à pas : évaluation et boucle d’entraînement
estimate_loss mesure la loss sur 8 lots, toujours les mêmes grâce à une graine fixe : les points des courbes restent comparables d’une évaluation à l’autre. Elle passe le modèle en mode évaluation, ce qui coupe le dropout. greedy_continuation donne un aperçu de ce que le modèle écrit.
La boucle train est la leçon 10 en six lignes : lot, loss, zero_grad, loss.backward(), écrêtage du gradient à 1, puis optimizer.step(). L’écrêtage évite qu’un gradient exceptionnellement grand fasse faire au modèle un pas démesuré. Tous les 100 pas, on évalue, et l’on copie les poids s’ils battent le record de validation.
@torch.no_grad()def estimate_loss( model: MiniLabModel, data: torch.Tensor, settings: TrainingSettings, length: int) -> float: model.eval() generator = torch.Generator().manual_seed(1234) # les mêmes fenêtres à chaque évaluation losses = [ model(*get_batch(data, settings.batch_size, length, generator))[1].item() for _ in range(settings.eval_batches) ] model.train() return sum(losses) / len(losses)
@torch.no_grad()def greedy_continuation(model: MiniLabModel, tokenizer: Tokenizer, prompt: str, n: int) -> str: """Un aperçu rapide de ce que le modèle écrit (la vraie génération est l'étape 4).""" model.eval() ids = tokenizer.encode(prompt) for _ in range(n): logits, _ = model(torch.tensor([ids[-model.config.context_length :]])) ids.append(int(logits[0, -1].argmax())) model.train() return tokenizer.decode(ids)
def train( model: MiniLabModel, train_data: torch.Tensor, val_data: torch.Tensor, settings: TrainingSettings, on_eval: Callable[[int, float, float], None] | None = None,) -> History: torch.manual_seed(settings.seed) generator = torch.Generator().manual_seed(settings.seed) length = min(model.config.context_length, len(train_data) - 2, len(val_data) - 2) optimizer = torch.optim.AdamW( model.parameters(), lr=settings.learning_rate, betas=(0.9, 0.95), weight_decay=settings.weight_decay, ) history = History() model.train() for step in range(settings.steps + 1): if step % settings.eval_every == 0 or step == settings.steps: train_loss = estimate_loss(model, train_data, settings, length) val_loss = estimate_loss(model, val_data, settings, length) history.points.append((step, train_loss, val_loss)) if val_loss < history.best_val: # arrêt précoce : on garde les meilleurs poids history.best_step, history.best_val = step, val_loss history.best_state = {k: v.clone() for k, v in model.state_dict().items()} if on_eval: on_eval(step, train_loss, val_loss) if step == settings.steps: break for group in optimizer.param_groups: group["lr"] = learning_rate(step, settings) inputs, targets = get_batch(train_data, settings.batch_size, length, generator) _, loss = model(inputs, targets) optimizer.zero_grad() loss.backward() # la rétropropagation (leçon 10) torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0) optimizer.step() # un pas d'AdamW return historyConstruire pas à pas : le programme principal
main prépare les données, crée le modèle avec le dropout, affiche une ligne par évaluation, puis recharge les meilleurs poids avant de les sauvegarder. L’option --overfit remplace le corpus d’entraînement par ses 500 premiers caractères, pour voir le surapprentissage de près.
def main() -> None: parser = argparse.ArgumentParser(description="Pré-entraînement du modèle MiniLab.") parser.add_argument("--overfit", action="store_true", help="entraîner sur 500 caractères") args = parser.parse_args() torch.set_num_threads(max(1, torch.get_num_threads())) tokenizer = get_tokenizer() text = read_corpus() train_data, val_data = split(tokenizer.encode(text)) settings = TrainingSettings() if args.overfit: cut = int(len(text) * 0.9) train_data = torch.tensor(tokenizer.encode(text[:500])) val_data = torch.tensor(tokenizer.encode(text[cut:])) settings = TrainingSettings(steps=300, eval_every=30) torch.manual_seed(settings.seed) model = MiniLabModel(Config(vocab_size=tokenizer.vocab_size, dropout=DROPOUT)) print( f"Entraînement : {len(train_data):,} tokens · validation : {len(val_data):,} tokens".replace( ",", " " ) ) print( f"{settings.steps} pas, lots de {settings.batch_size} × {model.config.context_length} tokens, AdamW (taux max {settings.learning_rate})\n" ) print(" pas loss train loss val perplexité val durée") start = time.perf_counter()
def report(step: int, train_loss: float, val_loss: float) -> None: print( f"{step:4d} {train_loss:9.3f} {val_loss:8.3f} {math.exp(val_loss):14.1f} {time.perf_counter() - start:5.1f} s" ) if not args.overfit and step % 400 == 0 and step: print(f" « {greedy_continuation(model, tokenizer, 'Phileas Fogg', 24)} »")
history = train(model, train_data, val_data, settings, report) print( f"\nMeilleure loss de validation : {history.best_val:.3f} au pas {history.best_step} (perplexité {math.exp(history.best_val):.1f})" ) if args.overfit: final = history.points[-1] print( f"Surapprentissage : loss train {final[1]:.3f}, loss val {final[2]:.3f} → écart {final[2] - final[1]:.2f}" ) print("Le modèle récite ses 500 caractères ; aucun checkpoint n'est sauvegardé.") return model.load_state_dict(history.best_state) save(model, BASE_CHECKPOINT) print(f"Checkpoint sauvegardé : {BASE_CHECKPOINT.relative_to(CHECKPOINTS.parents[1])}") print(f"« {greedy_continuation(model, tokenizer, 'Phileas Fogg', 30)} »")
if __name__ == "__main__": main()Tester et observer dans le terminal
- Lancer le pré-entraînement, environ 70 s
uv run python -m packages.minilab.train - Lancer la démonstration du surapprentissage, environ 30 s
uv run python -m packages.minilab.train --overfit - Lancer les tests de l’étape
uv run pytest tests/minilab/test_train.py
Entraînement : 52 178 tokens · validation : 5 798 tokens
800 pas, lots de 32 × 128 tokens, AdamW (taux max 0.0015)
pas loss train loss val perplexité val durée
0 6.953 6.958 1051.5 0.7 s
100 4.763 4.869 130.2 9.8 s
200 3.996 4.351 77.5 18.3 s
300 3.686 4.239 69.3 27.2 s
400 3.478 4.200 66.7 36.1 s
« Phileas Fogg, et il ne pouvait pas pas à la porte.
--Et, répondit le consul, répondit le consul, »
500 3.302 4.176 65.1 45.0 s
600 3.153 4.143 63.0 54.2 s
700 3.037 4.116 61.3 63.2 s
800 2.971 4.118 61.5 72.4 s
« Phileas Fogg et lui dit-il.
--Et vous, répondit le consul, répondit le consul, répondit le consul.
»
Meilleure loss de validation : 4.116 au pas 700 (perplexité 61.3)
Checkpoint sauvegardé : checkpoints/minilab/base.pt
« Phileas Fogg, il semait un c'était un cloignée.
Le Parsi, le guide, il se trouvertant à »Les deux loss partent de 6,95, le hasard. La loss d’entraînement descend jusqu’à 2,97 ; celle de validation, plus lentement, atteint son minimum au pas 700, 4,116, soit une perplexité d’environ 61, puis remonte très légèrement au pas 800. C’est le checkpoint du pas 700 qui est sauvegardé. Le texte généré ressemble à du Jules Verne de loin, avec ses tirets de dialogue et son consul ; il tourne aussi en rond, comme le décodage glouton de la leçon 09.
Entraînement : 195 tokens · validation : 5 746 tokens 300 pas, lots de 32 × 128 tokens, AdamW (taux max 0.0015) pas loss train loss val perplexité val durée 0 6.972 6.956 1049.0 0.8 s 30 4.870 6.679 795.8 3.6 s 60 2.001 6.713 823.0 6.4 s 90 0.206 7.499 1805.6 9.1 s 120 0.024 8.330 4144.7 12.0 s ... 300 0.004 9.924 20407.2 28.6 s Meilleure loss de validation : 6.679 au pas 30 (perplexité 795.8) Surapprentissage : loss train 0.004, loss val 9.924 → écart 9.92 Le modèle récite ses 500 caractères ; aucun checkpoint n'est sauvegardé.
Sur 500 caractères, soit 195 tokens, le surapprentissage de la leçon 11 est spectaculaire : la loss d’entraînement tombe à 0,004, le modèle récite son extrait, pendant que la loss de validation grimpe de 6,68 à 9,92. Le meilleur modèle, au sens de la validation, était celui du pas 30.
Dans train.py, mettez DROPOUT = 0.0 et relancez l’entraînement : la loss d’entraînement descend plus bas, jusqu’à 2,37, mais la meilleure loss de validation n’est plus que de 4,21, atteinte dès le pas 300, contre 4,12 au pas 700 avec le dropout. Sans lui, le modèle commence à réciter plus tôt. Remettez 0,1 et relancez avant de passer à l’étape 4, qui utilise ce checkpoint.
Le pré-entraînement découpe le corpus en 90 % d’entraînement et 10 % de validation, tire des lots de 32 × 128 tokens et entraîne avec AdamW (taux maximal 0,0015, warmup de 100 pas, décroissance en cosinus, weight decay 0,1, dropout 0,1). La loss passe de 6,95 à environ 4,12 en validation, soit une perplexité de 61, en 800 pas et environ 70 secondes. L’arrêt précoce garde les poids de la meilleure validation ; sur 500 caractères, le modèle surapprend.