Gradient Checkpointing

Le gradient checkpointing est une technique d'entraînement permettant d'économiser de la mémoire, qui stocke uniquement les activations intermédiaires sélectionnées de la passe avant et recalcule les autres lors de la rétropropagation. Également appelé activation checkpointing, il échange des calculs supplémentaires contre une utilisation de la mémoire de pointe inférieure. Malgré son nom, la technique enregistre des activations plutôt que des gradients de paramètres ou des fichiers de modèles, ce qui la rend particulièrement précieuse lorsque les tenseurs d'activation empêchent un réseau de neurones de tenir dans la mémoire GPU disponible.

Comment fonctionne le gradient checkpointing#

Lors d'une passe avant standard, un réseau de neurones calcule des tenseurs intermédiaires appelés activations. Le système de différenciation automatique conserve les activations nécessaires pour calculer les gradients ultérieurement, comme décrit dans la section mécanique de l'autograd PyTorch. Les réseaux profonds, les lots volumineux, les images haute résolution et les longues séquences d'entrée peuvent amener ces tenseurs enregistrés à consommer une mémoire considérable.

Le gradient checkpointing divise le réseau en segments :

1. La passe avant stocke les entrées ou les activations de limites pour les segments sélectionnés.
2. D'autres activations intermédiaires à l'intérieur de ces segments sont supprimées.
3. Pendant la passe arrière, chaque segment avec point de contrôle s'exécute à nouveau vers l'avant pour reconstruire les valeurs manquantes.
4. Les activations reconstruites sont immédiatement utilisées pour calculer les gradients.

L'API d'activation checkpointing PyTorch expose ce comportement via torch.utils.checkpoint. Des concepts équivalents apparaissent sous le nom de gradient checkpointing et rematérialisation JAX et de checkpointing de bande TensorFlow.

Le placement des points de contrôle détermine le compromis. Le placement de points de contrôle sur un plus grand nombre de régions permet généralement d'économiser davantage de mémoire mais répète un plus grand nombre d'opérations. Un placement sélectif autour des blocs lourds en activations peut offrir un meilleur équilibre que le recalcul de l'ensemble du réseau.

Compromis et techniques associées#

Le gradient checkpointing laisse généralement les paramètres du modèle, les gradients et les états de l'optimiseur inchangés. Sa cible principale est la mémoire des activations, et l'entraînement devient plus lent car certains calculs avant s'exécutent deux fois. Le résultat exact dépend de l'architecture, des limites des points de contrôle, de la forme des lots et du matériel.

Il diffère de plusieurs techniques associées :

• L'accumulation de gradients traite plusieurs micro-lots avant de mettre à jour les poids du modèle, créant ainsi un lot effectif plus grand sans charger chaque échantillon simultanément. Le gradient checkpointing réduit quant à lui les activations conservées pour chaque micro-lot.
• La précision mixte utilise des types de données de plus basse précision pour certaines opérations sélectionnées. Le flux de travail de précision mixte automatique PyTorch peut réduire la mémoire et accélérer les opérations compatibles, tandis que le checkpointing ajoute délibérément des calculs.
• La réduction de la taille des lots diminue la mémoire en traitant moins d'échantillons ensemble. Le checkpointing peut permettre à un lot plus important ou à une résolution d'entrée supérieure de rester réalisable.
• Les points de contrôle d'entraînement enregistrent les poids et l'état de l'optimiseur pour la récupération ou une inférence ultérieure. Le guide des points de contrôle de modèles PyTorch décrit ce mécanisme de persistance, qui n'est pas lié au recalcul des activations.

Le code soumis à des points de contrôle doit être fonctionnellement cohérent entre ses exécutions avant d'origine et recalculées. Un état mutable, des transferts d'appareils ou un caractère aléatoire non contrôlé à l'intérieur d'une région dotée de points de contrôle peuvent provoquer des erreurs ou des gradients incorrects. Dans les transformeurs de type décodeur, un cache clé-valeur peut également devoir être désactivé pendant l'entraînement avec points de contrôle, car l'état d'inférence mis en cache peut entrer en conflit avec la reconstruction du graphe avant.

Exemple PyTorch#

L'exemple suivant applique un point de contrôle à un bloc gourmand en mémoire lors d'une étape d'entraînement :

import torch
from torch import nn
from torch.utils.checkpoint import checkpoint

torch.manual_seed(0)

block = nn.Sequential(
    nn.Linear(1024, 4096),
    nn.ReLU(),
    nn.Linear(4096, 1024),
)
optimizer = torch.optim.AdamW(block.parameters())
inputs = torch.randn(8, 1024, requires_grad=True)
targets = torch.zeros_like(inputs)

optimizer.zero_grad(set_to_none=True)
outputs = checkpoint(block, inputs, use_reentrant=False)
loss = nn.functional.mse_loss(outputs, targets)
loss.backward()
optimizer.step()

Seuls les entrées du bloc et les informations de limites requises sont conservées ; ses activations internes sont reconstruites pendant loss.backward(). Les projets réels doivent comparer les exécutions avec et sans points de contrôle à l'aide d'une mesure telle que la mémoire GPU allouée de pointe PyTorch.

Applications concrètes#

• Vision par ordinateur haute résolution : La segmentation médicale, la détection aérienne et l'inspection industrielle peuvent s'entraîner sur de grandes images dont les cartes de caractéristiques consomment plus de mémoire que les poids du modèle. Le placement de points de contrôle sur des étages de squelette sélectionnés peut préserver la résolution des images ou autoriser des échantillons supplémentaires par lot au lieu de réduire agressivement la taille des entrées.

• Entraînement de transformeurs à longues séquences : Le stockage des activations augmente rapidement à mesure que la longueur des séquences et le nombre de couches augmentent. Le recalcul des blocs de transformeurs peut permettre à des contextes plus longs ou à des micro-lots plus grands de tenir sur le même accélérateur. Le guide de recalcul des activations NVIDIA illustre le recalcul complet et sélectif pour les couches de transformeurs.

Conseils pratiques#

Utilisez le gradient checkpointing lorsque le profilage montre que les activations, plutôt que les paramètres ou l'état de l'optimiseur, dominent la mémoire. Commencez par de grands blocs répétés et effectuez des benchmarks sur la mémoire de pointe, le temps d'itération et le comportement de validation avant d'étendre la couverture.

Pour les configurations d'entraînement Ultralytics YOLO documentées, les contrôles de mémoire de première ligne incluent la précision mixte automatique, la taille des images et la taille physique des lots. La référence Ultralytics AutoBatch explique la sélection automatique des lots en fonction de la mémoire GPU disponible. Lorsque le matériel local reste insuffisant, l'entraînement dans le cloud de la plateforme Ultralytics fournit des GPU cloud configurables pour des exécutions d'entraînement gérées. Le gradient checkpointing peut compléter ces contrôles lorsqu'une architecture PyTorch personnalisée nécessite une gestion plus fine de la mémoire des activations.