Rédigé par : Riley Debug – spécialiste du débogage AI et ingénieur ML ops
L’angoissant problème “CUDA out of memory” est un obstacle courant pour quiconque travaille avec des modèles d’apprentissage profond dans PyTorch. Vous avez minutieusement conçu votre modèle, préparé vos données et commencé l’entraînement, pour être finalement confronté à ce message frustrant. C’est un signal clair que votre GPU ne dispose pas d’une mémoire suffisante pour contenir tous les tenseurs et calculs nécessaires à l’opération actuelle. Ce n’est pas seulement une gêne ; cela freine vos progrès, fait perdre un temps précieux et peut constituer un goulet d’étranglement significatif dans le développement de solutions AI puissantes.
Ce guide a pour but de vous fournir une compréhension approfondie des raisons pour lesquelles ces erreurs se produisent dans PyTorch et, plus important encore, de vous fournir une boîte à outils pratique de stratégies pour les surmonter. Nous explorerons différentes techniques, des ajustements simples aux considérations architecturales plus avancées, en veillant à ce que vous puissiez gérer efficacement vos ressources GPU et garder vos pipelines d’entraînement en bon état de marche. Explorons comment diagnostiquer, prévenir et corriger les erreurs CUDA out of memory dans PyTorch, vous permettant ainsi 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 essentiel de comprendre ce qui consomme de la mémoire GPU durant un entraînement 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 de 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), résident sur le GPU si vous les y avez déplacés. La taille de ces tenseurs est directement corrélé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 volumineux.
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 considérable. 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, des 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, surtout pour les modèles avec un grand nombre de paramètres.
États de l’optimiseur
Les optimiseurs comme Adam, RMSprop ou Adagrad maintiennent des états internes pour chaque paramètre (par exemple, des buffers de momentum, des estimations de variance). Ces états sont souvent aussi grands que les paramètres eux-mêmes, doublant ou triplant efficacement la mémoire requise pour les paramètres seuls.
Taille de lot
Peut-être le facteur le plus simple 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 des lots plus importants puissent parfois conduire à des estimations de gradient plus stables et à une convergence plus rapide de l’entraînement, ils sont un moteur principal de la consommation de mémoire GPU.
Surcharge interne de PyTorch
Au-delà des données spécifiques de votre modèle, PyTorch a également une certaine surcharge interne pour gérer les contextes CUDA, les allocateurs de mémoire et d’autres composants opérationnels. Bien que généralement plus petite que la mémoire des tenseurs, elle fait partie de l’utilisation totale.
Diagnostics initiaux et solutions rapides
Lorsque l’erreur “CUDA out of memory” se produit, ne paniquez pas. Commencez par ces étapes immédiates pour diagnostiquer et potentiellement résoudre rapidement le problème.
Effacer 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 soient plus utilisés, entraînant une fragmentation ou une vue inexacte de la mémoire disponible. Effacer explicitement le cache peut libérer de l’espace.
import torch
torch.cuda.empty_cache()
Il est bon de l’appeler périodiquement, surtout après avoir supprimé de grands tenseurs ou avant d’en allouer de nouveaux. Notez que cela ne fait que vider le cache interne de PyTorch, pas la mémoire activement utilisée par les tenseurs.
Réduire la taille de lot
C’est souvent la première étape la plus efficace et la 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 originale
batch_size = 64
# Si OOM, essayez
batch_size = 32
# Ou même
batch_size = 16
Divisez votre taille de lot par deux de manière itérative jusqu’à ce que l’erreur disparaisse. Soyez conscient qu’une très petite taille de lot 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 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 voulez conserver leurs données mais pas leur historique de gradient.
# Exemple : Si vous avez un grand tenseur 'temp_data' qui n'est plus nécessaire
del temp_data
# De plus, appelez explicitement le ramasse-miettes
import gc
gc.collect()
torch.cuda.empty_cache() # Appelez à nouveau après suppression
Surveiller l’utilisation de la mémoire GPU
Des outils comme 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
Dans 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 mémoire
Lorsque les solutions rapides ne suffisent pas, ou si vous devez entraîner des modèles véritablement grands, 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 et effectuez ensuite 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-lots
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 accumulation
loss.backward() # Accumuler les gradients
if (i + 1) % accumulation_steps == 0:
optimizer.step() # Effectuer la mise à jour d'optimisation
optimizer.zero_grad() # Réinitialiser 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 GPU avec une mémoire limitée.
Checkpointing de gradient (Checkpointing d’activation)
Comme discuté, les activations intermédiaires occupent une mémoire significative. Le checkpointing de gradient aborde ce problème en ne stockant pas toutes les activations intermédiaires lors du passage avant. Au lieu de cela, il les recompute lors du 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 des 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 checkpointing pour le passage avant de cette couche
return checkpoint.checkpoint(self.layer, x)
# Exemple d'utilisation : enveloppez un grand bloc séquentiel
model = MyLargeModel()
# Remplacez 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))
Ceci est particulièrement utile pour les réseaux très profonds où stocker toutes les activations est impossible.
Entraînement à précision mixte (FP16/BF16)
L’entraînement à précision mixte consiste à effectuer certaines opérations en précision inférieure (FP16 ou BF16) tout en maintenant 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 les GPU 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 cette mise en œuvre :
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 dans ce contexte seront exécutées en FP16 lorsque cela est possible
outputs = model(inputs)
loss = criterion(outputs, labels)
scaler.scale(loss).backward() # Mise à l'échelle de la perte pour prévenir le sous-flux dans les gradients FP16
scaler.step(optimizer) # Déséchelle les gradients et met à jour les poids
scaler.update() # Met à jour le scaler pour la prochaine itération
La précision mixte est une technique puissante qui offre souvent à la fois des économies de mémoire et des gains de performance.
Déchargement vers le CPU (CPU Offloading)
Pour les modèles ou tenseurs intermédiaires extrêmement volumineux, vous pourriez envisager de déplacer des parties de votre modèle ou des tenseurs spécifiques vers le CPU lorsqu’ils ne sont pas utilisés activement, puis de les ramener sur le GPU au besoin. C’est plus complexe à gérer et introduce un surcoût important en raison du transfert de données entre le CPU et le GPU, mais cela peut être un recours ultime pour les modèles qui, autrement, ne tiendraient pas.
# Exemple : Déplacez 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 du tout nécessaire
del large_tensor_on_gpu
torch.cuda.empty_cache()
Pour les couches de modèles, 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 façon dont vous concevez votre modèle et rédigez 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) consomment plus de mémoire. Envisagez d’utiliser des alternatives plus économe en mémoire si possible :
- Convolutions séparables par profondeur : Souvent utilisées dans les architectures mobiles (par exemple, MobileNet), elles peuvent réduire de manière significative le nombre de paramètres et le calcul par rapport aux convolutions standard.
- Partage de paramètres : La réutilisation des poids dans différentes parties du réseau peut économiser de la mémoire.
- Élagage et quantification : Bien que généralement appliqués après l’entraînement, ils peuvent être envisagés pour le déploiement et peuvent influencer les choix de conception pour des environnements à mémoire contrainte.
Opérations en place
Les opérations PyTorch créent souvent de nouveaux tenseurs pour leur sortie. Les opérations en place (indiquées par un underscore final, par exemple, x.add_(y) au lieu de x = x + y) modifient le tenseur directement 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 perturber le graphe 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.
# Économie de mémoire (en place)
x.relu_() # Modifie x directement
# Crée un nouveau tenseur
x = torch.relu(x)
Éviter les clones/copies de tenseurs inutiles
Soyez attentif aux opérations qui créent implicitement des copies de tenseurs. Par exemple, le découpage d’un tenseur peut 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()
Utilisation de torch.no_grad() pour l’inférence/évaluation
Lors de l’évaluation ou de l’inférence, vous n’avez pas besoin de calculer ou de stocker les gradients. Envelopper votre code d’inférence dans le gestionnaire de contexte torch.no_grad() empêche Autograd de construire le graphe de calcul, ce qui permet d’économiser 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() # Réinitialise le modèle en mode entraînement
C’est une pratique essentielle pour quiconque travaillant avec PyTorch et cela peut souvent prévenir des erreurs OOM lors des étapes de validation.
Profilage de l’utilisation de la mémoire
Pour les cas complexes, PyTorch fournit de puissants outils de profilage qui peuvent identifier 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 affichera des détails sur l’utilisation de la mémoire.
Articles connexes
- Pratiques de test de l’équipe système IA
- Débogage des applications LLM : Un guide pratique pour le dépannage IA
- Maîtriser le test de pipeline IA : Conseils, astuces et exemples pratiques
🕒 Published: