Gradient Checkpointing

Le checkpointing des gradients est une technique d’entraînement qui économise de la mémoire en ne stockant que certaines activations intermédiaires de la passe avant et en recalculant les autres pendant la rétropropagation. Également appelée checkpointing des activations, cette technique échange du calcul supplémentaire contre une réduction de l’utilisation maximale de la mémoire. Malgré son nom, elle stocke des points de contrôle des activations, et non des gradients de paramètres ou des fichiers de modèle ; elle est donc particulièrement utile lorsque les tenseurs d’activation empêchent un réseau neuronal de tenir dans la mémoire GPU disponible.

Fonctionnement du checkpointing des gradients#

Lors d’une passe avant standard, un réseau neuronal calcule des tenseurs intermédiaires appelés activations. Le système de différentiation automatique conserve les activations nécessaires au calcul ultérieur des gradients, comme l’explique la mécanique de l’autograd de PyTorch. Les réseaux profonds, les lots volumineux, les images haute résolution et les longues séquences d’entrée peuvent rendre ces tenseurs sauvegardés très gourmands en mémoire.

Le checkpointing des gradients divise le réseau en segments :

1. La passe avant stocke les entrées ou les activations de frontière de certains segments.
2. Les autres activations intermédiaires de ces segments sont supprimées.
3. Lors de la passe arrière, chaque segment avec points de contrôle exécute de nouveau la passe avant pour reconstruire les valeurs manquantes.
4. Les activations reconstruites servent immédiatement au calcul des gradients.

L’API de checkpointing des activations de PyTorch expose ce comportement via torch.utils.checkpoint. Des notions équivalentes existent sous les noms de checkpointing des gradients et rematérialisation dans JAX et de checkpointing des blocs tape dans TensorFlow.

Le placement des points de contrôle détermine le compromis. En général, plus on en place dans le réseau, plus on économise de mémoire, mais plus on répète d’opérations. Un placement sélectif autour des blocs gourmands en activations peut offrir un meilleur équilibre que le recalcul de tout le réseau.

Compromis et techniques apparentées#

Le checkpointing des gradients laisse généralement inchangés les paramètres du modèle, les gradients et les états de l’optimiseur. Il cible principalement la mémoire occupée par les activations, et ralentit l’entraînement, car certains calculs de la passe avant sont répétés. 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.

Cette technique diffère de plusieurs techniques apparentées :

• L’accumulation des gradients traite plusieurs micro-lots avant de mettre à jour les poids du modèle, ce qui crée un lot effectif plus grand sans charger tous les échantillons simultanément. Le checkpointing des gradients réduit plutôt les activations conservées pour chaque micro-lot.
• La précision mixte utilise des types de données moins précis pour certaines opérations. Le flux de travail de précision mixte automatique de PyTorch peut réduire l’utilisation de 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 de lot diminue l’utilisation de la mémoire en traitant moins d’échantillons à la fois. Le checkpointing peut permettre d’utiliser un lot plus grand ou une résolution d’entrée supérieure.
• Les points de contrôle d’entraînement sauvegardent les poids et l’état de l’optimiseur pour permettre une reprise ou une inférence ultérieure. Le guide PyTorch sur les points de contrôle de modèle décrit ce mécanisme de persistance, qui n’a rien à voir avec le recalcul des activations.

Le code avec points de contrôle doit produire des résultats fonctionnellement cohérents lors de son exécution initiale et de son recalcul. Un état mutable, des transferts entre appareils ou un caractère aléatoire non contrôlé dans une région avec points de contrôle peuvent provoquer des erreurs ou des gradients incorrects. Dans les Transformers de type décodeur, il peut également falloir désactiver un cache clé-valeur pendant l’entraînement avec checkpointing, car l’état mis en cache pour l’inférence peut entrer en conflit avec la reconstruction du graphe de calcul de la passe avant.

Exemple PyTorch#

L’exemple suivant applique le checkpointing à un bloc gourmand en mémoire pendant 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()

Seules les entrées du bloc et les informations de frontière requises sont conservées ; ses activations internes sont reconstruites pendant loss.backward(). Dans les projets réels, compare les exécutions avec et sans checkpointing à l’aide d’une mesure telle que le pic de mémoire GPU allouée de PyTorch.

Applications concrètes#

• Vision par ordinateur haute résolution : La segmentation médicale, la détection aérienne et l’inspection industrielle peuvent nécessiter l’entraînement sur de grandes images dont les cartes de caractéristiques occupent plus de mémoire que les poids du modèle. Le checkpointing de certaines étapes du réseau de base peut préserver la résolution des images ou permettre d’ajouter des échantillons par lot, plutôt que de réduire fortement les dimensions des entrées.

• Entraînement de Transformers sur de longues séquences : Le stockage des activations augmente rapidement avec la longueur des séquences et le nombre de couches. Recalculer les blocs Transformer peut permettre de traiter des contextes plus longs ou d’utiliser des micro-lots plus grands sur le même accélérateur. Le guide NVIDIA sur le recalcul des activations illustre le recalcul complet et sélectif pour les couches Transformer.

Conseils pratiques#

Utilise le checkpointing des gradients lorsque le profilage montre que les activations, plutôt que les paramètres ou l’état de l’optimiseur, dominent l’utilisation de la mémoire. Commence par les grands blocs répétés et évalue le pic mémoire, le temps d’itération et le comportement en validation avant d’étendre son utilisation.

Pour les configurations d’entraînement Ultralytics YOLO documentées, les premiers réglages de gestion de la mémoire sont 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 de la taille des lots en fonction de la mémoire GPU disponible. Si le matériel local reste insuffisant, l’entraînement dans le cloud d’Ultralytics Platform fournit des GPU cloud configurables pour les sessions d’entraînement gérées. Le checkpointing des gradients peut compléter ces réglages lorsqu’une architecture PyTorch personnalisée exige une gestion plus fine de la mémoire occupée par les activations.