Aller au contenu principal

Chapitre 3 : Techniques d'entraînement distribué

Introduction​

L'entraînement des LLMs avec des milliards de paramètres nécessite de distribuer le travail sur de nombreux GPU. Ce chapitre couvre les principales stratégies de parallélisme.

Parallélisme de données​

L'approche la plus simple​

import torch.distributed as dist
from torch.nn.parallel import DistributedDataParallel

# Chaque GPU traite des données différentes, même modèle
model = DistributedDataParallel(model)

Limitation : Le modèle doit tenir sur un seul GPU

Parallélisme de modèle​

Parallélisme de tenseur​

Diviser les couches individuelles sur les GPU :

[Matrice de poids de la couche]
   ↓ Diviser les colonnes
[GPU 0: W[:, :d/2]] [GPU 1: W[:, d/2:]]

Parallélisme de pipeline​

Différentes couches sur différents GPU :

GPU 0: Couches 0-11
GPU 1: Couches 12-23
GPU 2: Couches 24-35
GPU 3: Couches 36-47

ZeRO : Optimiseur sans redondance​

Trois étapes​

ÉtapeCe qui est partitionnéÉconomies de mémoire
ZeRO-1États de l'optimiseur4x
ZeRO-2+ Gradients8x
ZeRO-3+ ParamètresLinéaire avec les GPU

Implémentation avec DeepSpeed​

import deepspeed

model_engine, optimizer, _, _ = deepspeed.initialize(
    model=model,
    config={
        "zero_optimization": {
            "stage": 3,
            "offload_optimizer": {"device": "cpu"},
            "offload_param": {"device": "cpu"}
        },
        "bf16": {"enabled": True}
    }
)

Calcul d'entraînement​

Mise à l'échelle de Chinchilla​

Entraînement optimal : environ 20 tokens par paramètre

# Taille de modèle optimale en fonction du calcul
def optimal_model_size(compute_budget):
    # C ≈ 6 * N * D (N=paramètres, D=tokens)
    # Optimal : D ≈ 20N
    return compute_budget / (6 * 20)

Points clés à retenir​

Résumé
  1. Parallélisme de données pour les petits modèles
  2. Parallélisme de modèle lorsque le modèle ne tient pas sur un seul GPU
  3. ZeRO réduit considérablement les besoins en mémoire
  4. Précision mixte (bf16) accélère l'entraînement