Chapitre 2 : Entraînement Distribué

Stratégies de Parallélisme

Exemple PyTorch DDP

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

def setup(rank, world_size):
    dist.init_process_group("nccl", rank=rank, world_size=world_size)

def train(rank, world_size):
    setup(rank, world_size)
    
    model = MyModel().to(rank)
    model = DDP(model, device_ids=[rank])
    
    optimizer = torch.optim.Adam(model.parameters())
    
    for epoch in range(num_epochs):
        for batch in dataloader:
            outputs = model(batch)
            loss = criterion(outputs, targets)
            loss.backward()
            optimizer.step()
            optimizer.zero_grad()
    
    dist.destroy_process_group()

# Launch with torchrun
# torchrun --nproc_per_node=4 train.py

Points Clés

Résumé

1. Le parallélisme de données fonctionne pour la plupart des cas d'utilisation
2. Le parallélisme de modèle est nécessaire pour les modèles très grands
3. DeepSpeed/FSDP permettent d'entraîner des modèles plus grands que la mémoire du GPU

---

Suivant : Chapitre 3 : Entraînement Infonuagique