Auteur : Riley Debug – spécialiste en débogage d’IA et ingénieur ML ops
L’angoissant “CUDA out of memory” est un obstacle courant pour quiconque travaille avec des modèles d’apprentissage profond dans PyTorch. Vous avez méticuleusement conçu votre modèle, préparé vos données et commencé l’entraînement, pour finalement être confronté à ce message frustrant. C’est un signe clair que votre GPU n’a pas suffisamment de mémoire pour contenir tous les tenseurs et calculs nécessaires pour l’opération actuelle. Cela n’est pas simplement une gêne ; cela ralentit vos progrès, fait perdre un temps précieux et peut constituer un goulot d’étranglement important dans le développement de solutions IA puissantes.
Ce guide est conçu pour vous fournir une compréhension approfondie des raisons pour lesquelles ces erreurs se produisent dans PyTorch et, plus important encore, vous fournir une boîte à outils pratique de stratégies pour les surmonter. Nous explorerons diverses techniques, des ajustements simples aux considérations architecturales plus avancées, afin de vous assurer que vous pouvez gérer efficacement vos ressources GPU et maintenir vos pipelines d’entraînement en bon état de fonctionnement. Explorons comment diagnostiquer, prévenir et corriger les erreurs CUDA out of memory dans PyTorch, vous permettant de construire et d’entraîner des modèles plus grands et plus complexes.
Comprendre l’utilisation de la mémoire GPU dans PyTorch
Avant de pouvoir corriger les erreurs “CUDA out of memory”, il est crucial de comprendre ce qui consomme de la mémoire GPU lors d’une exécution d’entraînement avec PyTorch. Plusieurs composants contribuent à l’empreinte mémoire totale, et identifier les principaux coupables est la première étape vers une optimisation efficace.
Tenseurs et paramètres du modèle
Chaque tenseur dans votre modèle, y compris les données d’entrée, les activations intermédiaires et les paramètres apprenables du modèle (poids et biais), se trouvent sur le GPU si vous les y avez déplacés. La taille de ces tenseurs est directement liée à l’utilisation de la mémoire. Les modèles plus grands avec plus de couches et de paramètres nécessitent naturellement plus de mémoire. De même, des images d’entrée à plus haute résolution ou des longueurs de séquence plus longues entraîneront des tenseurs d’entrée plus grands.
Activations intermédiaires (passage avant)
Lors du passage avant, PyTorch doit stocker les activations de chaque couche. Ces valeurs intermédiaires sont essentielles pour le calcul des gradients lors du passage arrière (rétropropagation). Pour les réseaux profonds, l’accumulation de ces activations peut être substantielle. Par exemple, un ResNet avec de nombreux blocs générera de nombreuses cartes de caractéristiques qui doivent être conservées en mémoire.
Gradients (passage arrière)
Lorsque le passage arrière commence, les gradients sont calculés pour chaque paramètre. Ces gradients occupent également de la mémoire GPU. Le moteur de différentiation automatique de PyTorch (Autograd) gère ce processus, mais la mémoire allouée pour les gradients peut être significative, en particulier pour les modèles avec un grand nombre de paramètres.
État des optimizers
Des optimizers comme Adam, RMSprop ou Adagrad maintiennent des états internes pour chaque paramètre (c’est-à-dire, des tampons de moment, des estimations de variance). Ces états sont souvent aussi grands que les paramètres eux-mêmes, doublant ou triplant en fait la mémoire requise uniquement pour les paramètres.
Taille de lot
Peut-être le facteur le plus simple à considérer est la taille de lot. Une taille de lot plus grande signifie que plus d’échantillons d’entrée et leurs activations intermédiaires correspondantes sont traités simultanément. Bien que de plus grands lots puissent parfois conduire à des estimations de gradients plus stables et à une convergence d’entraînement plus rapide, ils sont un facteur principal de consommation de mémoire GPU.
Frais généraux internes de PyTorch
Au-delà des données spécifiques de votre modèle, PyTorch lui-même a des frais généraux internes pour gérer les contextes CUDA, les allocateurs de mémoire et d’autres composants opérationnels. Bien que généralement plus petits que la mémoire des tenseurs, cela fait partie de l’utilisation totale.
Diagnostics initiaux et corrections rapides
Lorsque l’erreur “CUDA out of memory” apparaît, ne paniquez pas. Commencez par ces étapes immédiates pour diagnostiquer et potentiellement résoudre rapidement le problème.
Vider le cache CUDA de PyTorch
Parfois, l’allocateur de mémoire de PyTorch peut conserver de la mémoire mise en cache même après que les tenseurs ne sont plus utilisés, entraînant de la fragmentation ou une vue inexacte de la mémoire disponible. Vider explicitement le cache peut libérer de l’espace.
import torch
torch.cuda.empty_cache()
Il est conseillé de l’appeler périodiquement, surtout après avoir supprimé de grands tenseurs ou avant d’en allouer de nouveaux. Notez que cela ne vide que le cache interne de PyTorch, pas la mémoire activement utilisée par les tenseurs.
Réduire la taille de lot
C’est souvent le premier pas le plus efficace et le plus simple. Une taille de lot plus petite réduit directement le nombre d’échantillons traités simultanément, diminuant ainsi la mémoire nécessaire pour les entrées, les activations intermédiaires et les gradients.
# Taille de lot d'origine
batch_size = 64
# Si OOM, essayez
batch_size = 32
# Ou même
batch_size = 16
Réduisez votre taille de lot par moitié de manière itérative jusqu’à ce que l’erreur disparaisse. Soyez conscient qu’une taille de lot très petite pourrait affecter la stabilité de l’entraînement ou la vitesse de convergence, donc c’est un compromis.
Supprimer les tenseurs et variables inutiles
Assurez-vous de ne pas conserver de grands tenseurs ou de variables qui ne sont plus nécessaires. Le ramasse-miettes de Python finira par les libérer, mais les supprimer explicitement peut libérer de la mémoire plus tôt. N’oubliez pas de les déplacer vers le CPU ou de les détacher s’ils font partie du graphe de calcul et que vous souhaitez conserver leurs données mais pas leur historique de gradients.
# Exemple : Si vous avez un grand tenseur 'temp_data' qui n'est plus nécessaire
del temp_data
# Appelez également explicitement le ramasse-miettes
import gc
gc.collect()
torch.cuda.empty_cache() # Appelez à nouveau après la suppression
Surveiller l’utilisation de la mémoire GPU
Des outils tels que nvidia-smi (dans votre terminal) ou les fonctions de rapport de mémoire intégrées de PyTorch peuvent vous donner des aperçus sur la consommation de mémoire de votre GPU. Cela aide à identifier si un autre processus consomme de la mémoire ou si votre script PyTorch est le seul coupable.
nvidia-smi
Au sein de PyTorch, vous pouvez obtenir des statistiques de mémoire détaillées :
print(torch.cuda.memory_summary(device=None, abbreviated=False))
Cela fournit une répartition de la mémoire allouée par rapport à la mémoire réservée, et peut parfois indiquer une fragmentation.
Techniques avancées d’optimisation de la mémoire
Lorsque les corrections rapides ne suffisent pas, ou que vous devez entraîner de vraiment grands modèles, des techniques plus sophistiquées sont nécessaires. Ces méthodes impliquent souvent des compromis entre mémoire, temps de calcul et complexité du code.
Accumulation de gradients
L’accumulation de gradients vous permet de simuler une taille de lot effective plus grande sans augmenter l’empreinte mémoire d’un seul passage avant/arrière. Au lieu de mettre à jour les poids après chaque lot, vous accumulez les gradients sur plusieurs petits lots, puis effectuez une seule mise à jour.
model = MyModel().to(device)
optimizer = torch.optim.Adam(model.parameters(), lr=1e-3)
accumulation_steps = 4 # Accumuler les gradients sur 4 mini-batchs
for epoch in range(num_epochs):
for i, (inputs, labels) in enumerate(dataloader):
inputs, labels = inputs.to(device), labels.to(device)
outputs = model(inputs)
loss = criterion(outputs, labels)
loss = loss / accumulation_steps # Normaliser la perte pour l'accumulation
loss.backward() # Accumuler les gradients
if (i + 1) % accumulation_steps == 0:
optimizer.step() # Effectuer la mise à jour de l'optimiseur
optimizer.zero_grad() # Effacer les gradients
# S'assurer que les gradients accumulés restants sont appliqués à la fin de l'époque
if (i + 1) % accumulation_steps != 0:
optimizer.step()
optimizer.zero_grad()
Cette technique est puissante pour l’entraînement avec de grandes tailles de lot effectives sur des GPUs avec une mémoire limitée.
Point de contrôle des gradients (point de contrôle des activations)
Comme discuté, les activations intermédiaires occupent une mémoire significative. Le point de contrôle des gradients aborde ce problème en ne stockant pas toutes les activations intermédiaires durant le passage avant. Au lieu de cela, il les recompute durant le passage arrière pour les segments qui nécessitent des gradients. Cela réduit considérablement la mémoire mais augmente le temps de calcul, car certaines parties du passage avant sont exécutées deux fois.
import torch.utils.checkpoint as checkpoint
class CheckpointBlock(torch.nn.Module):
def __init__(self, layer):
super().__init__()
self.layer = layer
def forward(self, x):
# Utiliser le point de contrôle pour le passage avant de cette couche
return checkpoint.checkpoint(self.layer, x)
# Exemple d'utilisation : envelopper un grand bloc séquentiel
model = MyLargeModel()
# Remplacer une grande partie séquentielle par une version checkpointée
# Par exemple, si votre modèle a `self.encoder = nn.Sequential(...)`
# Vous pourriez envelopper l'encodeur :
# self.encoder = CheckpointBlock(nn.Sequential(*encoder_layers))
C’est particulièrement utile pour des réseaux très profonds où stocker toutes les activations est impossible.
Entraînement en précision mixte (FP16/BF16)
L’entraînement en précision mixte consiste à effectuer certaines opérations en précision inférieure (FP16 ou BF16) tout en conservant d’autres en FP32. Cela peut réduire de moitié l’empreinte mémoire pour les poids, les activations et les gradients, et accélérer souvent l’entraînement sur des GPUs modernes (comme les architectures NVIDIA Volta, Turing, Ampere, Ada Lovelace) qui ont des cœurs Tensor conçus pour les calculs FP16.
Le module torch.cuda.amp de PyTorch facilite la mise en œuvre de cela :
from torch.cuda.amp import autocast, GradScaler
model = MyModel().to(device)
optimizer = torch.optim.Adam(model.parameters(), lr=1e-3)
scaler = GradScaler() # Pour la stabilité FP16
for epoch in range(num_epochs):
for inputs, labels in dataloader:
inputs, labels = inputs.to(device), labels.to(device)
optimizer.zero_grad()
with autocast(): # Les opérations à l'intérieur de ce contexte seront exécutées en FP16 lorsque c'est possible
outputs = model(inputs)
loss = criterion(outputs, labels)
scaler.scale(loss).backward() # Échelle de la perte pour éviter le débordement des gradients FP16
scaler.step(optimizer) # Déséchelle des gradients et mise à jour des poids
scaler.update() # Met à jour le scaler pour l'itération suivante
La précision mixte est une technique puissante qui offre souvent à la fois des économies de mémoire et des gains de performance.
Délestage vers le CPU (Délestage CPU)
Pour des modèles extrêmement grands ou des tenseurs intermédiaires, vous pourriez envisager de déplacer certaines parties de votre modèle ou des tenseurs spécifiques vers le CPU lorsqu’ils ne sont pas activement utilisés, puis de les ramener sur le GPU lorsque nécessaire. Cela est plus complexe à gérer et introduit une surcharge significative en raison du transfert de données entre le CPU et le GPU, mais cela peut être un dernier recours pour des modèles qui, autrement, ne tiendraient pas.
# Exemple : Déplacer un grand tenseur vers le CPU après son utilisation
large_tensor_on_gpu = torch.randn(10000, 10000).to(device)
# ... calculs utilisant large_tensor_on_gpu ...
# Lorsqu'il n'est plus nécessaire sur le GPU
large_tensor_on_gpu = large_tensor_on_gpu.cpu()
# Ou simplement supprimer s'il n'est pas nécessaire du tout
del large_tensor_on_gpu
torch.cuda.empty_cache()
Pour les couches du modèle, cela implique souvent de diviser le modèle et de déplacer des blocs séquentiels entre les appareils.
Considérations Architecturales et de Conception de Code
Au-delà des techniques spécifiques d’optimisation de la mémoire, la manière dont vous concevez votre modèle et écrivez votre code PyTorch peut avoir un impact significatif sur l’utilisation de la mémoire GPU.
Architectures de Modèle Efficaces
Certaines architectures de modèles sont intrinsèquement plus gourmandes en mémoire que d’autres. Par exemple, les modèles avec des couches très larges ou ceux qui génèrent de nombreuses cartes de caractéristiques haute résolution (par exemple, dans les tâches de segmentation) consommeront plus de mémoire. Envisagez d’utiliser des alternatives plus économes en mémoire si possible :
- Convolutions Séparables en Profondeur : Souvent utilisées dans des architectures mobiles (par exemple, MobileNet), celles-ci peuvent réduire considérablement les paramètres et le calcul par rapport aux convolutions standard.
- Partage de Paramètres : Réutiliser les poids dans différentes parties du réseau peut économiser de la mémoire.
- Élagage et Quantification : Bien qu’elles soient généralement appliquées après l’entraînement, celles-ci peuvent être envisagées pour le déploiement et pourraient influencer les choix de conception pour des environnements à mémoire limitée.
Opérations In-place
Les opérations PyTorch créent souvent de nouveaux tenseurs pour leur sortie. Les opérations in-place (indiquées par un underscore final, par exemple, x.add_(y) au lieu de x = x + y) modifient directement le tenseur sans allouer de nouvelle mémoire pour le résultat. Bien qu’elles puissent économiser de la mémoire, utilisez-les avec prudence car elles peuvent casser le graphique de calcul si elles ne sont pas gérées correctement, en particulier lorsqu’elles sont utilisées sur des tenseurs qui nécessitent des gradients.
# Économies de mémoire (in-place)
x.relu_() # Modifie x directement
# Crée un nouveau tenseur
x = torch.relu(x)
Éviter les Clones/Copies de Tenseurs Inutiles
Faites attention aux opérations qui créent implicitement des copies de tenseurs. Par exemple, le découpage d’un tenseur pourrait parfois créer une vue, mais d’autres opérations pourraient créer une copie complète. Utilisez explicitement .clone() seulement lorsqu’une copie profonde est vraiment nécessaire, sinon, travaillez avec des vues lorsque cela est possible.
# Crée une vue (pas de nouvelle mémoire pour les données)
view_tensor = original_tensor[0:10]
# Crée un nouveau tenseur (nouvelle mémoire)
cloned_tensor = original_tensor.clone()
Utiliser torch.no_grad() pour l’Inference/L’Évaluation
Lors de l’évaluation ou de l’inférence, vous n’avez pas besoin de calculer ou de stocker des gradients. En enveloppant votre code d’inférence dans le gestionnaire de contexte torch.no_grad(), vous empêchez Autograd de construire le graphique de calcul, ce qui économise une quantité significative de mémoire en ne stockant pas les activations intermédiaires pour la rétropropagation.
model.eval() # Met le modèle en mode évaluation
with torch.no_grad():
for inputs, labels in val_dataloader:
inputs, labels = inputs.to(device), labels.to(device)
outputs = model(inputs)
# ... calculer les métriques ...
model.train() # Remet le modèle en mode entraînement
C’est une pratique fondamentale pour quiconque travaille avec PyTorch et peut souvent prévenir les erreurs OOM lors des étapes de validation.
Profilage de l’Utilisation de la Mémoire
Pour des cas complexes, PyTorch fournit de puissants outils de profilage qui peuvent déterminer exactement quelles opérations consomment le plus de mémoire. Le module torch.profiler (ou l’ancien torch.autograd.profiler) peut enregistrer les allocations de mémoire CUDA.
import torch
from torch.profiler import profile, record_function, ProfilerActivity
# Exemple de profilage d'un passage avant/arrière unique
model = MyModel().to(device)
inputs = torch.randn(4, 3, 224, 224).to(device)
labels = torch.randint(0, 10, (4,)).to(device)
optimizer = torch.optim.Adam(model.parameters(), lr=1e-3)
criterion = torch.nn.CrossEntropyLoss()
with profile(activities=[
ProfilerActivity.CPU, ProfilerActivity.CUDA], record_shapes=True) as prof:
with record_function("model_inference"):
outputs = model(inputs)
with record_function("loss_computation"):
loss = criterion(outputs, labels)
with record_function("backward_pass"):
loss.backward()
with record_function("optimizer_step"):
optimizer.step()
optimizer.zero_grad()
print(prof.key_averages().table(sort_by="cuda_memory_usage", row_limit=10))
La sortie du profileur montrera une mémoire détaillée
Articles Connexes
- Pratiques de l’équipe de test des systèmes d’IA
- Débogage des applications LLM : Un guide pratique pour le dépannage de l’IA
- Maîtriser le test des pipelines d’IA : Conseils, astuces et exemples pratiques
🕒 Published: