Optimisation et performance PyTorch : Entraîner plus vite et plus efficacement

10 min de lecture 1994 mots

tl;dr: Guide complet pour accélérer PyTorch : mixed precision FP16 (2× speedup), DDP multi-GPU (scaling linéaire), gradient accumulation, torch.compile (jusqu'à 2× speedup), profiling pour identifier les bottlenecks. Inclut best practices et optimisations mémoire.

L’entraînement de modèles deep learning peut prendre des heures, jours, voire semaines. Maîtriser les techniques d’optimisation PyTorch vous permettra de réduire drastiquement les temps d’entraînement (2-10×) tout en utilisant efficacement les ressources GPU/CPU disponibles.

Dans cet article, vous allez :

⚡ Utiliser Mixed Precision (AMP) pour entraîner 2× plus vite
🖥️ Maîtriser le Distributed Training (DDP) multi-GPU
💾 Optimiser la mémoire avec gradient checkpointing
📊 Profiler votre code pour identifier les bottlenecks
🚀 Accélérer avec torch.compile() (PyTorch 2.0+)
🧠 Appliquer les best practices pour maximiser les performances

Temps de lecture : 40 minutes Niveau : Intermédiaire à Avancé

💡 Prérequis : Maîtrise de PyTorch (articles 1-6), expérience en entraînement de modèles. Accès à un GPU recommandé pour tester les optimisations.

Code et visualisation PyTorch pour l’optimisation et les algorithmes d’entraînement dans le développement de modèles d’IA

Benchmark initial : Connaître votre baseline

Avant d’optimiser, mesurez vos performances actuelles :

import torch
import time

def benchmark_training_step(model, dataloader, criterion, optimizer, device):
    """Mesurer la vitesse d'un epoch"""
    model.train()

    start_time = time.time()
    total_samples = 0

    for batch_idx, (data, targets) in enumerate(dataloader):
        data, targets = data.to(device), targets.to(device)

        optimizer.zero_grad()
        outputs = model(data)
        loss = criterion(outputs, targets)
        loss.backward()
        optimizer.step()

        total_samples += data.size(0)

        if batch_idx >= 100:  # Benchmark sur 100 batches
            break

    elapsed = time.time() - start_time
    samples_per_sec = total_samples / elapsed

    print(f"⏱️  Samples/sec: {samples_per_sec:.2f}")
    print(f"⏱️  Time per batch: {elapsed / (batch_idx + 1) * 1000:.2f} ms")

    return samples_per_sec

# Baseline
baseline_speed = benchmark_training_step(model, train_loader, criterion, optimizer, device)

🔎 Tip
Règle d’or : Toujours mesurer avant et après chaque optimisation. Certaines optimisations peuvent ne pas fonctionner sur votre architecture/dataset spécifique.

Mixed Precision Training (AMP)

Mixed Precision utilise FP16 (float16) pour les calculs tout en conservant FP32 pour la stabilité.

Avantages

✅ 2-3× plus rapide sur GPUs modernes (Tensor Cores)
✅ Réduit la mémoire de ~40%
✅ Même précision finale (avec loss scaling)
✅ Facile à implémenter (3 lignes de code)

Implémentation avec autocast et GradScaler

import torch
from torch.cuda.amp import autocast, GradScaler

# Setup
model = Model().cuda()
optimizer = torch.optim.Adam(model.parameters(), lr=1e-3)
criterion = nn.CrossEntropyLoss()
scaler = GradScaler()  # Pour éviter les underflows FP16

# Training loop
model.train()
for epoch in range(num_epochs):
    for data, targets in train_loader:
        data, targets = data.cuda(), targets.cuda()

        optimizer.zero_grad()

        # Forward pass en mixed precision
        with autocast():
            outputs = model(data)
            loss = criterion(outputs, targets)

        # Backward avec scaling
        scaler.scale(loss).backward()

        # Optimizer step avec unscaling
        scaler.step(optimizer)
        scaler.update()

    print(f'Epoch {epoch}: Loss = {loss.item():.4f}')

Que se passe-t-il ?

autocast() : Choisit automatiquement FP16 ou FP32 pour chaque opération
- Matmuls, convolutions → FP16 (rapide)
- Softmax, layer norm → FP32 (stabilité)
GradScaler : Multiplie la loss par un facteur (ex: 65536) avant backward
- Évite que les gradients FP16 deviennent trop petits (underflow)
- Divise les gradients avant optimizer step

Comparaison FP32 vs AMP

def compare_fp32_vs_amp(model, dataloader, device):
    """Comparer les performances FP32 vs AMP"""

    # FP32
    model_fp32 = model.to(device)
    optimizer = torch.optim.Adam(model_fp32.parameters())

    start = time.time()
    for data, targets in dataloader:
        data, targets = data.to(device), targets.to(device)
        optimizer.zero_grad()
        loss = nn.functional.cross_entropy(model_fp32(data), targets)
        loss.backward()
        optimizer.step()
    fp32_time = time.time() - start

    # AMP
    model_amp = model.to(device)
    optimizer = torch.optim.Adam(model_amp.parameters())
    scaler = GradScaler()

    start = time.time()
    for data, targets in dataloader:
        data, targets = data.to(device), targets.to(device)
        optimizer.zero_grad()
        with autocast():
            loss = nn.functional.cross_entropy(model_amp(data), targets)
        scaler.scale(loss).backward()
        scaler.step(optimizer)
        scaler.update()
    amp_time = time.time() - start

    print(f"FP32 time: {fp32_time:.2f}s")
    print(f"AMP time: {amp_time:.2f}s")
    print(f"Speedup: {fp32_time / amp_time:.2f}×")

Résultats typiques :

GPU moderne (A100, V100) : 2-3× speedup
GPU ancien (K80, GTX 1080) : Pas de speedup (pas de Tensor Cores)

Distributed Training : Multi-GPU

DataParallel vs DistributedDataParallel

Caractéristique	DataParallel	DistributedDataParallel
Performance	Correct	Excellent
Scaling	Non linéaire	Linéaire
Multi-node	❌ Non	✅ Oui
Recommandé	❌ Legacy	✅ Production
Complexité	Facile	Moyenne

Toujours utiliser DDP en 2025 (sauf si vous avez vraiment besoin de DP pour compatibilité legacy).

DistributedDataParallel (DDP)

Principe : Chaque GPU exécute une copie du modèle, les gradients sont moyennés.

GPU 0: Batch 0-31   → Forward → Backward → Gradients ┐
GPU 1: Batch 32-63  → Forward → Backward → Gradients ├→ All-Reduce → Update weights
GPU 2: Batch 64-95  → Forward → Backward → Gradients ┘

Implémentation DDP

train_ddp.py :

import torch
import torch.distributed as dist
import torch.multiprocessing as mp
from torch.nn.parallel import DistributedDataParallel as DDP
from torch.utils.data.distributed import DistributedSampler

def setup(rank, world_size):
    """Initialiser le process group"""
    dist.init_process_group(
        backend='nccl',  # NCCL pour GPU, 'gloo' pour CPU
        init_method='env://',
        world_size=world_size,
        rank=rank
    )

def cleanup():
    """Nettoyer"""
    dist.destroy_process_group()

def train_ddp(rank, world_size, epochs, batch_size):
    """Fonction de training pour chaque GPU"""

    # Setup
    setup(rank, world_size)
    torch.cuda.set_device(rank)

    # Model
    model = Model().to(rank)
    model = DDP(model, device_ids=[rank])

    # Dataset avec DistributedSampler
    dataset = YourDataset()
    sampler = DistributedSampler(
        dataset,
        num_replicas=world_size,
        rank=rank,
        shuffle=True
    )
    dataloader = DataLoader(
        dataset,
        batch_size=batch_size,
        sampler=sampler,
        num_workers=4,
        pin_memory=True
    )

    # Optimizer
    optimizer = torch.optim.Adam(model.parameters(), lr=1e-3)
    criterion = nn.CrossEntropyLoss()

    # Training loop
    for epoch in range(epochs):
        sampler.set_epoch(epoch)  # Important pour le shuffling

        model.train()
        for data, targets in dataloader:
            data, targets = data.to(rank), targets.to(rank)

            optimizer.zero_grad()
            outputs = model(data)
            loss = criterion(outputs, targets)
            loss.backward()
            optimizer.step()

        # Logging uniquement sur rank 0
        if rank == 0:
            print(f'Epoch {epoch}: Loss = {loss.item():.4f}')

    # Sauvegarder uniquement depuis rank 0
    if rank == 0:
        torch.save(model.module.state_dict(), 'model_ddp.pth')

    cleanup()

def main():
    world_size = torch.cuda.device_count()  # Nombre de GPUs
    epochs = 10
    batch_size = 32  # Per GPU

    # Lancer un process par GPU
    mp.spawn(
        train_ddp,
        args=(world_size, epochs, batch_size),
        nprocs=world_size,
        join=True
    )

if __name__ == '__main__':
    main()

Lancement :

# Single node, 4 GPUs
python train_ddp.py

# Multi-node avec torchrun (recommandé)
torchrun --nproc_per_node=4 --nnodes=2 --node_rank=0 \
  --master_addr="192.168.1.1" --master_port=29500 \
  train_ddp.py

DDP + Mixed Precision

Combiner pour maximum speedup :

from torch.cuda.amp import autocast, GradScaler

def train_ddp_amp(rank, world_size):
    setup(rank, world_size)

    model = DDP(Model().to(rank), device_ids=[rank])
    optimizer = torch.optim.Adam(model.parameters())
    scaler = GradScaler()  # Un scaler par process

    for epoch in range(epochs):
        for data, targets in dataloader:
            data, targets = data.to(rank), targets.to(rank)

            optimizer.zero_grad()

            with autocast():
                outputs = model(data)
                loss = criterion(outputs, targets)

            scaler.scale(loss).backward()
            scaler.step(optimizer)
            scaler.update()

    cleanup()

Speedup attendu :

4 GPUs : ~3.7-3.9× speedup (efficacité ~95%)
8 GPUs : ~7.2-7.6× speedup (efficacité ~90-95%)

Gradient Accumulation

Problème : Vous voulez batch_size=256 mais votre GPU n’a que 8GB (max batch_size=64).

Solution : Accumuler les gradients sur plusieurs mini-batches avant de faire optimizer.step().

accumulation_steps = 4  # Effective batch_size = 64 * 4 = 256

model.train()
optimizer.zero_grad()

for batch_idx, (data, targets) in enumerate(dataloader):
    data, targets = data.cuda(), targets.cuda()

    # Forward
    outputs = model(data)
    loss = criterion(outputs, targets)

    # Normaliser par accumulation_steps
    loss = loss / accumulation_steps

    # Backward (accumule les gradients)
    loss.backward()

    # Update uniquement tous les N batches
    if (batch_idx + 1) % accumulation_steps == 0:
        optimizer.step()
        optimizer.zero_grad()

Avec AMP :

scaler = GradScaler()
accumulation_steps = 4

for batch_idx, (data, targets) in enumerate(dataloader):
    data, targets = data.cuda(), targets.cuda()

    with autocast():
        outputs = model(data)
        loss = criterion(outputs, targets) / accumulation_steps

    scaler.scale(loss).backward()

    if (batch_idx + 1) % accumulation_steps == 0:
        scaler.step(optimizer)
        scaler.update()
        optimizer.zero_grad()

torch.compile() : PyTorch 2.0+ Speedup

torch.compile() (PyTorch 2.0+) compile le modèle pour des optimisations automatiques.

import torch

model = Model().cuda()

# Compiler le modèle
model = torch.compile(model, mode='default')
# Modes: 'default', 'reduce-overhead', 'max-autotune'

# Entraînement normal
optimizer = torch.optim.Adam(model.parameters())

for epoch in range(epochs):
    for data, targets in train_loader:
        data, targets = data.cuda(), targets.cuda()

        optimizer.zero_grad()
        outputs = model(data)  # Compilé automatiquement
        loss = criterion(outputs, targets)
        loss.backward()
        optimizer.step()

Modes de compilation :

'default' : Équilibre vitesse/compilation
'reduce-overhead' : Réduit l’overhead Python (bon pour petits modèles)
'max-autotune' : Optimisations agressives (long temps de compilation)

Speedups typiques :

Transformers : 1.3-1.7× speedup
CNNs : 1.2-1.5× speedup
RNNs : 1.1-1.3× speedup

Important : La première itération compile (lent), les suivantes sont rapides.

Profiling : Identifier les bottlenecks

PyTorch Profiler

from torch.profiler import profile, ProfilerActivity, schedule

with profile(
    activities=[ProfilerActivity.CPU, ProfilerActivity.CUDA],
    schedule=schedule(wait=1, warmup=1, active=3, repeat=2),
    on_trace_ready=torch.profiler.tensorboard_trace_handler('./log/profiler'),
    record_shapes=True,
    profile_memory=True,
    with_stack=True
) as prof:
    for step, (data, targets) in enumerate(train_loader):
        data, targets = data.cuda(), targets.cuda()

        optimizer.zero_grad()
        outputs = model(data)
        loss = criterion(outputs, targets)
        loss.backward()
        optimizer.step()

        prof.step()  # Signal au profiler

        if step >= 10:
            break

# Voir les résultats
print(prof.key_averages().table(sort_by="cuda_time_total", row_limit=10))

Analyser avec TensorBoard :

tensorboard --logdir=./log/profiler

Profiler Simplifié

import torch.autograd.profiler as profiler

with profiler.profile(record_shapes=True, use_cuda=True) as prof:
    outputs = model(data)
    loss = criterion(outputs, targets)
    loss.backward()

print(prof.key_averages().table(sort_by="cuda_time_total", row_limit=10))

Exemple de sortie :

--------------------------  ------------  ------------  ------------
Name                        CPU time      CUDA time     Calls
--------------------------  ------------  ------------  ------------
aten::conv2d                5.23ms        45.12ms       64
aten::batch_norm            2.11ms        18.34ms       64
aten::relu                  1.02ms        8.91ms        64
aten::linear                3.45ms        32.78ms       2
...

Analyse :

Si aten::conv2d domine : Modèle bien optimisé
Si aten::copy_ élevé : Problème de transferts CPU↔GPU
Si temps CPU » temps CUDA : Bottleneck dans le DataLoader

Optimisations mémoire

Gradient Checkpointing

Trade-off : Temps CPU contre mémoire GPU.

import torch.utils.checkpoint as checkpoint

class CheckpointedModel(nn.Module):
    def __init__(self):
        super().__init__()
        self.layer1 = HeavyLayer()
        self.layer2 = HeavyLayer()
        self.layer3 = HeavyLayer()

    def forward(self, x):
        # Checkpointer les layers lourds
        x = checkpoint.checkpoint(self.layer1, x)
        x = checkpoint.checkpoint(self.layer2, x)
        x = self.layer3(x)
        return x

Résultat :

Mémoire : -50% (permet batch_size 2× plus grand)
Temps : +20-30% (recalcul du forward pendant backward)

Flash attention (pour transformers)

# Standard attention : O(N²) mémoire
attention = nn.MultiheadAttention(embed_dim=512, num_heads=8)

# Flash Attention : O(N) mémoire, 2-4× plus rapide
from flash_attn import flash_attn_func

def flash_attention(q, k, v):
    """
    q, k, v: (batch, seq_len, head_dim)
    """
    return flash_attn_func(q, k, v, causal=False)

Best practices de performance

DataLoader optimisations

dataloader = DataLoader(
    dataset,
    batch_size=128,
    num_workers=4,          # 4-8× num GPUs
    pin_memory=True,        # Transferts CPU→GPU plus rapides
    prefetch_factor=2,      # Prefetch 2 batches par worker
    persistent_workers=True # Garder workers entre epochs
)

Transferts asynchrones

for data, targets in dataloader:
    data = data.cuda(non_blocking=True)  # Async transfer
    targets = targets.cuda(non_blocking=True)

    # GPU peut commencer pendant que le transfer se termine
    outputs = model(data)

Fused optimizers

# Optimizers fusionnés (plus rapides)
from apex.optimizers import FusedAdam

optimizer = FusedAdam(model.parameters(), lr=1e-3)
# Au lieu de torch.optim.Adam

Éviter les synchronisations inutiles

# ❌ Lent : force synchronisation CPU↔GPU
for data, targets in dataloader:
    loss = model(data, targets)
    total_loss += loss.item()  # .item() synchronise !

# ✅ Rapide : accumuler puis synchroniser
losses = []
for data, targets in dataloader:
    loss = model(data, targets)
    losses.append(loss)

total_loss = torch.stack(losses).mean().item()  # Une seule sync

torch.backends.cudnn

import torch.backends.cudnn as cudnn

# Activer autotuner (trouve les meilleurs algos)
cudnn.benchmark = True  # Si input size constant

# Rendre déterministe (reproductibilité)
cudnn.deterministic = True
cudnn.benchmark = False

Checklist d’optimisation

Appliquez dans cet ordre pour gains cumulatifs :

Profiler → Identifier les bottlenecks
DataLoader → num_workers, pin_memory, prefetch
Mixed Precision (AMP) → +2× speedup
torch.compile() → +1.3× speedup (PyTorch 2.0+)
DDP Multi-GPU → +N× speedup (N GPUs)
Gradient Accumulation → Larger effective batch_size
Gradient Checkpointing → Si OOM (Out Of Memory)
cudnn.benchmark → +5-10% speedup

Speedup cumulatif attendu : 5-10× sur setup multi-GPU optimisé !

Exercices Pratiques

Optimiser un entraînement

Prenez un modèle existant et optimisez-le progressivement.

Objectifs :

Mesurer le baseline (samples/sec, mémoire GPU)
Appliquer AMP : mesurer le speedup
Profiler avec PyTorch Profiler : identifier le top-3 des opérations lentes
Optimiser le DataLoader (num_workers, pin_memory)
Appliquer torch.compile() si PyTorch 2.0+
Créer un rapport avant/après avec métriques

Métriques :

Samples/sec
Temps par epoch
Utilisation mémoire GPU
Accuracy finale (doit rester identique !)

Scaling multi-GPU

Testez le scaling DDP sur plusieurs GPUs.

Objectifs :

Entraîner avec 1, 2, 4, 8 GPUs (si disponible)
Mesurer le speedup pour chaque configuration
Calculer l’efficacité de scaling : speedup / num_gpus
Combiner DDP + AMP pour maximum speedup
Identifier à quel point le scaling devient sous-optimal

Memory vs Speed Trade-offs

Explorez gradient checkpointing et accumulation.

Objectifs :

Trouver le batch_size maximum sans checkpointing
Appliquer gradient checkpointing : quel batch_size maintenant ?
Mesurer le trade-off temps/mémoire
Utiliser gradient accumulation pour simuler un très grand batch (1024+)
Comparer la convergence avec différents batch_size effectifs

Conclusion

Vous maîtrisez maintenant les optimisations PyTorch avancées :

✅ Mixed Precision : Entraîner 2-3× plus vite avec FP16

✅ DDP : Scaling linéaire sur plusieurs GPUs

✅ torch.compile() : Optimisations automatiques (PyTorch 2.0+)

✅ Profiling : Identifier et résoudre les bottlenecks

✅ Optimisations mémoire : Checkpointing, flash attention

✅ Best Practices : DataLoader, transferts async, fused optimizers

Prochaines Étapes

➡️ Prochain article : PyTorch Lightning et Best Practices - Structurer votre code professionnel, logging, callbacks, et workflows modernes.

Pour aller plus loin :

PyTorch Performance Tuning Guide : Guide officiel
NVIDIA Apex : Outils d’optimisation avancés
PyTorch Distributed Tutorials : DDP en profondeur
Flash Attention : Attention optimisée

Vos entraînements PyTorch sont maintenant 10× plus rapides !