Entraînez Votre Propre LLM de A à Z : Un Guide Complet de Bout en Bout
Un tutoriel pratique pour construire, pré-entraîner et aligner un LLM Transformer de zéro en utilisant PyTorch, couvrant SFT, DPO, PPO et GRPO.
Avez-vous déjà voulu entraîner votre propre grand modèle de langage (LLM) de zéro, pas seulement affiner un modèle existant ? Le projet open-source train-llm-from-scratch de FareedKhan-dev fournit un pipeline complet et transparent qui vous emmène du texte brut à un modèle aligné et capable de raisonnement. Ce n'est pas un wrapper boîte noire autour de bibliothèques existantes ; c'est un parcours éducatif pratique où chaque algorithme — de l'attention multi-têtes à GRPO — est implémenté en PyTorch pur.
Ce guide vous accompagnera tout au long du processus, en expliquant le pourquoi et le comment de chaque étape, afin que vous puissiez comprendre les fondamentaux et exécuter le code vous-même.
Pourquoi Construire un LLM de Zéro ?
La plupart des développeurs interagissent avec les LLM via des API de haut niveau comme transformers ou trl de Hugging Face. Bien que puissants, ils abstraient les mécanismes fondamentaux. Construire de zéro vous donne :
- Compréhension Approfondie : Vous apprenez exactement comment fonctionnent la tokenisation, l'attention et les boucles d'entraînement.
- Contrôle Total : Vous pouvez modifier n'importe quel composant — fonctions de perte, architectures ou stratégies d'entraînement — sans contraintes de bibliothèque.
- Valeur Éducative : C'est la meilleure façon de saisir des concepts comme le masquage causal, la modélisation de récompense et les méthodes de gradient de politique.
Ce projet est conçu pour les étudiants, développeurs et chercheurs qui veulent voir l'ensemble du pipeline en un seul endroit, du téléchargement des données à la génération de texte.
Le Pipeline en un Coup d'Œil
Le parcours suit un chemin clair et séquentiel :
- Préparation des Données : Texte brut → Identifiants de tokens → Tableaux stockés.
- Construction du Modèle : Construction d'un Transformer à partir de petits composants réutilisables (MLP, Attention, Blocs).
- Pré-entraînement : Prédiction du prochain token sur un grand corpus (The Pile).
- Génération de Texte : Échantillonnage à partir du modèle entraîné.
- Post-entraînement (Alignement) : Transformer un modèle de base en un assistant utile via SFT, Modélisation de Récompense, DPO, PPO et GRPO.
- Évaluation : Mesure des performances sur un benchmark comme GSM8K.
- Inférence et Chat : Interaction avec le modèle final.
Étape 1 : Préparation des Données
Un modèle ne comprend que des entiers. La première tâche est de convertir le texte en identifiants de tokens et de les stocker efficacement.
Tokenisation
Le projet utilise le tokenizer r50k_base d'OpenAI (de tiktoken), le même que celui utilisé par GPT-3. Le texte devient une liste d'entiers, et un token spécial <|endoftext|> (ID 50256) est ajouté à chaque document pour que le modèle apprenne les limites.
# Télécharger et tokeniser les données de The Pile
python scripts/data_download.py
python scripts/data_preprocess.py
Pour le chemin plus rapide et plus récent :
python scripts/prepare_pretrain_data.py --split val --out data/pile_dev.h5
python scripts/prepare_pretrain_data.py --split train --num_shards 1 --out data/pile_train.h5
Format de Chat et Masquage de Perte
Pour le post-entraînement, le modèle doit comprendre les rôles (utilisateur vs. assistant). Le projet utilise des marqueurs de texte brut :
<|user|>
Qu'est-ce que 13 + 29 ?<|endoftext|><|assistant|>
<think>13 + 29 = 42</think><answer>42</answer><|endoftext|>
Un masque de perte garantit que le modèle n'apprend que des tokens de l'assistant, pas de l'invite de l'utilisateur. Ceci est crucial pour SFT et RL.
def encode_chat(messages, add_generation_prompt=False):
ids, mask = [], []
for m in messages:
role = m["role"]
header_ids = _encode_ordinary(_header_for(role))
ids.extend(header_ids)
mask.extend([0] * len(header_ids))
content_ids = _encode_ordinary(m["content"])
is_completion = role == "assistant"
ids.extend(content_ids)
mask.extend([1 if is_completion else 0] * len(content_ids))
ids.append(EOT_ID)
mask.append(1 if is_completion else 0)
return ids, mask
Étape 2 : Construction du Transformer
Le Transformer est construit à partir de quatre petits modules PyTorch réutilisables. Cette conception modulaire rend le code facile à comprendre et à modifier.
Perceptron Multicouche (MLP)
La partie "réflexion" de chaque bloc. Il étend le vecteur de token, applique une ReLU, et le projette en retour.
class MLP(nn.Module):
def __init__(self, n_embed):
super().__init__()
self.hidden = nn.Linear(n_embed, 4 * n_embed)
self.relu = nn.ReLU()
self.proj = nn.Linear(4 * n_embed, n_embed)
def forward(self, x):
x = self.relu(self.hidden(x))
x = self.proj(x)
return x
Attention à Tête Unique
Permet à un token de regarder d'autres tokens. Il calcule les vecteurs de requête, clé et valeur, applique un masque causal (afin qu'un token ne puisse pas voir les tokens futurs), et prend une somme pondérée des valeurs.
class Head(nn.Module):
def __init__(self, head_size, n_embed, context_length):
super().__init__()
self.key = nn.Linear(n_embed, head_size, bias=False)
self.query = nn.Linear(n_embed, head_size, bias=False)
self.value = nn.Linear(n_embed, head_size, bias=False)
self.register_buffer('tril', torch.tril(torch.ones(context_length, context_length)))
def forward(self, x):
B, T, C = x.shape
k = self.key(x)
q = self.query(x)
scale_factor = 1 / math.sqrt(C)
attn_weights = q @ k.transpose(-2, -1) * scale_factor
attn_weights = attn_weights.masked_fill(self.tril[:T, :T] == 0, float('-inf'))
attn_weights = F.softmax(attn_weights, dim=-1)
v = self.value(x)
out = attn_weights @ v
return out
Attention Multi-Têtes
Exécute plusieurs têtes d'attention en parallèle, chacune apprenant différentes relations, puis concatène et projette les résultats.
Le Bloc Transformer
Combine l'attention et le MLP avec des connexions résiduelles pré-norm. C'est l'unité répétitive centrale.
class Block(nn.Module):
def __init__(self, n_head, n_embed, context_length):
super().__init__()
self.ln1 = nn.LayerNorm(n_embed)
self.attn = MultiHeadAttention(n_head, n_embed, context_length)
self.ln2 = nn.LayerNorm(n_embed)
self.mlp = MLP(n_embed)
def forward(self, x):
x = x + self.attn(self.ln1(x))
x = x + self.mlp(self.ln2(x))
return x
Le Transformer Complet
Enveloppe tout : embeddings de tokens, embeddings de position, une pile de blocs, une couche de normalisation finale, et une projection linéaire vers des logits de taille de vocabulaire.
class Transformer(nn.Module):
def __init__(self, n_head, n_embed, context_length, vocab_size, N_BLOCKS):
super().__init__()
self.token_embed = nn.Embedding(vocab_size, n_embed)
self.position_embed = nn.Embedding(context_length, n_embed)
self.attn_blocks = nn.ModuleList([Block(n_head, n_embed, context_length) for _ in range(N_BLOCKS)])
self.layer_norm = nn.LayerNorm(n_embed)
self.lm_head = nn.Linear(n_embed, vocab_size)
self.register_buffer('pos_idxs', torch.arange(context_length))
def forward(self, idx, targets=None):
x = self.forward_hidden(idx)
logits = self.lm_head(x)
loss = None
if targets is not None:
B, T, C = logits.shape
flat_logits = logits.reshape(B * T, C)
targets = targets.reshape(B * T).long()
loss = F.cross_entropy(flat_logits, targets)
return logits, loss
Étape 3 : Pré-entraînement du Modèle de Base
Le pré-entraînement est la prédiction du prochain token sur un grand corpus. Le modèle lit des fenêtres aléatoires de tokens, prédit le token suivant à chaque position, et met à jour ses poids via une perte d'entropie croisée.
# Entraîner un modèle de 13M paramètres
python scripts/train_transformer.py
# Pour des modèles plus grands avec optimisations mémoire
python scripts/train_transformer.py --amp --grad-checkpointing --grad-accum 8
La perte commence près de ln(vocab_size) (~10.8) et diminue à mesure que le modèle apprend les statistiques du langage. Un modèle de 77M paramètres entraîné sur 2 GPU L40 a atteint une perte d'environ 3,73 après 2000 étapes.
Étape 4 : Génération de Texte
Une fois entraîné, vous pouvez générer du texte en échantillonnant à partir de la distribution de sortie du modèle.
def generate(self, idx, max_new_tokens):
for _ in range(max_new_tokens):
idx_cond = idx[:, -self.context_length:]
logits, _ = self(idx_cond)
logits = logits[:, -1, :]
probs = F.softmax(logits, dim=-1)
idx_next = torch.multinomial(probs, num_samples=1)
idx = torch.cat((idx, idx_next), dim=1)
return idx
python scripts/generate_text.py --model_path models/transformer_B.pt --input_text "Le" --max_new_tokens 100
Étape 5 : Post-entraînement (Alignement)
Un modèle de base peut continuer du texte, mais il ne peut pas suivre des instructions. Le post-entraînement l'aligne en utilisant le fine-tuning supervisé (SFT) et l'apprentissage par renforcement (RL).
SFT (Fine-Tuning Supervisé)
Apprend au modèle à répondre dans un format de chat. La perte est calculée uniquement sur les tokens de l'assistant.
python scripts/prepare_sft_data.py --context_length 1024
torchrun --standalone --nproc_per_node=2 scripts/train_sft.py
Le Modèle de Récompense
Entraîne une petite tête linéaire au-dessus du modèle SFT pour noter les réponses. Utilise la perte de Bradley-Terry pour préférer les réponses choisies aux rejetées.
def bradley_terry_loss(chosen_rewards, rejected_rewards):
return -F.logsigmoid(chosen_rewards - rejected_rewards).mean()
DPO, ORPO et KTO
L'Optimisation Directe des Préférences (DPO) saute le modèle de récompense et la boucle RL, travaillant directement sur les paires de préférences. ORPO et KTO sont des variantes.
torchrun --standalone --nproc_per_node=2 scripts/train_dpo.py --loss_type dpo
PPO (Optimisation de Politique Proximale)
La boucle RLHF classique. Le modèle génère des réponses, les note avec un modèle de récompense, et met à jour via une perte de politique écrêtée.
def ppo_policy_loss(new_logp, old_logp, advantages, mask, clip=0.2):
ratio = torch.exp(new_logp - old_logp)
surr1 = ratio * advantages
surr2 = torch.clamp(ratio, 1.0 - clip, 1.0 + clip) * advantages
loss = -masked_mean(torch.min(surr1, surr2), mask)
return loss
GRPO (Optimisation de Politique Relative par Groupe)
La méthode de style DeepSeek-R1. Elle échantillonne un groupe de réponses pour chaque invite, les note, et utilise la moyenne et l'écart type du groupe comme ligne de base.
def group_advantages(rewards, group_size, eps=1e-4):
r = rewards.view(-1, group_size)
mean = r.mean(dim=1, keepdim=True)
std = r.std(dim=1, keepdim=True)
adv = (r - mean) / (std + eps)
return adv.reshape(-1)
Étape 6 : Évaluation
La métrique clé est la précision greedy GSM8K : le modèle doit produire la réponse correcte à l'intérieur des balises <answer>.
for s in base_pretrained sft dpo ppo grpo; do
python scripts/eval_post_training.py --ckpt models/$s.pt --label $s --limit 200 --append logs/table.jsonl
done
Étape 7 : Parler au Modèle
Utilisez le script de chat pour interagir avec n'importe quel point de contrôle.
python scripts/chat.py --ckpt models/sft.pt --prompt "Qu'est-ce que 13 + 29 ?"
Pour Commencer
Clonez le dépôt :
git clone https://github.com/FareedKhan-dev/train-llm-from-scratch.git cd train-llm-from-scratch pip install -e .Commencez petit : Entraînez d'abord le modèle de 13M paramètres. C'est rapide et vous donnera une idée du pipeline.
Passez à l'échelle : Augmentez
n_embedetn_blocksjusqu'à atteindre la limite mémoire de votre GPU.Suivez la chaîne de post-entraînement : Exécutez SFT, puis DPO, puis PPO ou GRPO, et observez l'amélioration de la précision GSM8K.
Ce projet est un joyau rare : c'est à la fois un outil pratique et un tutoriel complet. Que vous soyez un étudiant souhaitant comprendre en profondeur les Transformers ou un développeur cherchant à expérimenter avec des techniques d'alignement, train-llm-from-scratch est le point de départ parfait.