Gradient Checkpointing

Définition

Le Gradient Checkpointing (aussi appelé Activation Checkpointing) est une technique d'optimisation mémoire pour l'entraînement de réseaux de neurones profonds qui réduit la consommation mémoire GPU en acceptant un surcoût de calcul. Sans cette technique, l'entraînement d'un LLM de 7B paramètres avec des séquences longues peut nécessiter plus de mémoire que disponible sur un seul GPU A100 80GB. Lors du backward pass (rétropropagation), le calcul des gradients nécessite les activations intermédiaires de toutes les couches calculées pendant le forward pass. Pour un modèle avec N couches et des séquences longues, stocker toutes ces activations peut requérir autant de mémoire que les paramètres du modèle eux-mêmes. Pour LLaMA 3 70B avec des séquences de 8K tokens, les activations peuvent dépasser 100 GB. Le Gradient Checkpointing résout ce problème en ne sauvegardant qu'un sous-ensemble des activations (les "checkpoints", par exemple une couche sur deux ou les entrées de chaque bloc Transformer). Lors du backward pass, les activations intermédiaires non sauvegardées sont recalculées à la volée depuis le dernier checkpoint. Ce re-calcul consomme ~33% de compute supplémentaire (un forward pass partiel par couche), mais réduit la mémoire des activations de O(N) à O(√N) dans la configuration optimale. En pratique, l'activation checkpointing est activée par défaut dans la plupart des scripts de fine-tuning (TRL, Axolotl, LLaMA-Factory) car elle permet de doubler la batch size ou d'entraîner des modèles plus grands sur des GPUs de capacité limitée. Combiné avec la quantification QLoRA (4-bit), le gradient checkpointing permet d'entraîner des modèles de 65B+ paramètres sur un seul GPU A100 80GB.

Trade-off mémoire vs compute

Sans checkpointing : mémoire activations = O(L × B × S × d) où L=couches, B=batch, S=seq, d=dim
Avec checkpointing granulaire (1/L checkpoints) : mémoire ≈ O(√L × B × S × d)
Surcoût compute : ~33% de forward pass supplémentaire
Gain mémoire pratique : 4-10× selon la configuration

Activation dans HuggingFace

from transformers import TrainingArguments

# Méthode 1 : TrainingArguments
args = TrainingArguments(
    gradient_checkpointing=True,
    gradient_checkpointing_kwargs={"use_reentrant": False}  # Recommandé pour éviter bugs
)

# Méthode 2 : directement sur le modèle
model.gradient_checkpointing_enable(
    gradient_checkpointing_kwargs={"use_reentrant": False}
)

# Note : désactiver model.use_cache=False quand gradient_checkpointing=True
model.config.use_cache = False  # Incompatible avec checkpointing

Activation offloading (ZeRO-Offload)

DeepSpeed ZeRO-Offload étend le concept en déchargeant les activations checkpointées vers la RAM CPU (généralement 10-50× plus grande que la VRAM). Avec NVLink, la bande passante CPU-GPU permet de maintenir des performances raisonnables tout en utilisant des centaines de GB de RAM pour les activations.

Articles liés

Expert en cybersécurité offensive et intelligence artificielle. Pentest, audit et développement IA sur-mesure.

Services

Ressources

Projets & Outils

Contact

Nous contacter

Mentions légales Confidentialité Politique éditoriale Plan du site ⭐ Source préférée

Un projet cybersécurité ?

Expert dispo · Réponse 24h

Devis