Auteur : Riley Debug – Spécialiste du débogage AI et ingénieur en ML ops
L’angoissant message d’erreur “CUDA out of memory” est un obstacle courant pour quiconque travaille avec des modèles d’apprentissage profond dans PyTorch. Vous avez soigneusement conçu votre modèle, préparé vos données et commencé l’entraînement, pour être confronté à ce message frustrant. C’est un signal clair indiquant que votre GPU n’a pas assez de mémoire pour contenir tous les tenseurs nécessaires et les calculs pour l’opération en cours. Ce n’est pas seulement une gêne ; cela bloque votre progression, gaspille un temps précieux et peut constituer un goulot d’étranglement significatif dans le développement de solutions IA puissantes.
Ce guide est conçu pour vous donner une compréhension approfondie des raisons pour lesquelles ces erreurs se produisent dans PyTorch et, plus important encore, pour 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. Explorons comment diagnostiquer, prévenir et corriger les erreurs de mémoire CUDA 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 pendant un entraînement PyTorch. Plusieurs composants contribuent à l’empreinte mémoire totale, et identifier les principaux responsables 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), réside 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 de plus haute résolution ou des longueurs de séquence plus longues conduiront à 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 importante, surtout pour les modèles avec un grand nombre de paramètres.
États des optimiseurs
Les optimiseurs comme Adam, RMSprop ou Adagrad maintiennent des états internes pour chaque paramètre (par exemple, des tampons de momentum, des estimations de variance). Ces états sont souvent aussi volumineux que les paramètres eux-mêmes, doublant ou triplant l’espace mémoire requis pour les seuls paramètres.
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 les plus grands lots puissent parfois conduire à des estimations de gradients plus stables et à une convergence d’entraînement plus rapide, ils représentent un facteur principal de 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, cela fait partie de l’utilisation totale.
Diagnostics initiaux et solutions rapides
Lorsque l’erreur “CUDA out of memory” survient, ne paniquez pas. Commencez par ces étapes immédiates pour diagnostiquer et potentiellement résoudre le problème rapidement.
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 soient plus utilisés, ce qui conduit à une 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 bon de l’appeler périodiquement, surtout après avoir supprimé de gros tenseurs ou avant d’en allouer de nouveaux. Notez que cela ne vide que le cache interne de PyTorch, et non 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, essayer
batch_size = 32
# Ou même
batch_size = 16
Divisez itérativement votre taille de lot par deux 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 collecteur de déchets de Python les libérera finalement, 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 collecteur de déchets
import gc
gc.collect()
torch.cuda.empty_cache() # Appelez à nouveau après la 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 informations 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 responsable.
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 ventilation de la mémoire allouée par rapport à celle réservée, et peut parfois donner des indices sur la fragmentation.
Techniques avancées d’optimisation de la mémoire
Lorsque les solutions rapides ne suffisent pas, ou que vous devez entraîner de très 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 efficace 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 des 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 l'étape d'optimisation
optimizer.zero_grad() # Réinitialiser les gradients
# S'assurer que tous les gradients accumulés restants sont appliqués à la fin de l'epoch
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 efficaces sur des GPU avec une mémoire limitée.
Point de contrôle de gradient (Point de contrôle d’activation)
Comme discuté, les activations intermédiaires prennent beaucoup de mémoire. Le point de contrôle de gradient résout ce problème en ne stockant pas toutes les activations intermédiaires durant le passage avant. Au lieu de cela, il les recalcule pendant le passage arrière pour les segments qui nécessitent des gradients. Ceci 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 avec point de contrôle
# 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 les réseaux très profonds où il est impossible de stocker toutes les activations.
Formation en précision mixte (FP16/BF16)
La formation en précision mixte implique d’effectuer certaines opérations en plus basse précision (FP16 ou BF16) tout en gardant 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ère souvent l’entraînement sur des GPU modernes (comme les architectures NVIDIA Volta, Turing, Ampere, Ada Lovelace) qui disposent de noyaux Tensor conçus pour les calculs FP16.
Le module torch.cuda.amp de PyTorch facilite 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 sous-flux dans les gradients FP16
scaler.step(optimizer) # Dé-échelle des gradients et mise à jour des poids
scaler.update() # Mise à jour de l'échelle 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écharger vers le CPU (CPU Offloading)
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 lorsqu’ils sont nécessaires. Cela est plus complexe à gérer et introduit un surcoût significatif en raison du transfert de données entre le CPU et le GPU, mais cela peut être un dernier recours pour les modèles qui ne tiendraient autrement 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 du tout nécessaire
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 dispositifs.
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 é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èle 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 fonctionnalités 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 les architectures mobiles (par exemple, MobileNet), celles-ci peuvent réduire considérablement le nombre de paramètres et le calcul par rapport aux convolutions standards.
- Partage de Paramètres : Réutiliser les poids à travers différentes parties du réseau peut faire économiser de la mémoire.
- Élagage et Quantification : Bien que généralement appliqués après l’entraînement, ceux-ci peuvent être envisagés pour le déploiement et pourraient influencer les choix de conception pour des environnements à mémoire contrainte.
Opérations In-Place
Les opérations PyTorch créent souvent de nouveaux tenseurs pour leur sortie. Les opérations in-place (dénotées par un trait de soulignement 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 faire économiser de la mémoire, utilisez-les avec précaution car elles peuvent casser le graphe de computation si elles ne sont pas gérées correctement, surtout lorsqu’elles sont utilisées sur des tenseurs qui nécessitent des gradients.
# Économie de mémoire (in-place)
x.relu_() # Modifie directement x
# Crée un nouveau tenseur
x = torch.relu(x)
Éviter les Clones/Copies Inutiles de Tenseurs
Faites attention 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 peuvent créer une copie complète. Utilisez explicitement .clone() uniquement lorsqu’une copie profonde est réellement 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’Inference/L’Évaluation
Lors de l’évaluation ou de l’inférence, vous n’avez pas besoin de calculer ou de stocker des gradients. Envelopper votre code d’inférence dans le gestionnaire de contexte torch.no_grad() empêche Autograd de construire le graphe de computation, ce qui fait é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() # Remet le modèle en mode entraînement
C’est une pratique fondamentale pour quiconque travaillant avec PyTorch et peut souvent éviter des erreurs de type OOM lors des étapes de validation.
Profilage de l’Utilisation de la Mémoire
Pour des cas complexes, PyTorch fournit des outils de profilage puissants 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 profil peut montrer des détails sur l’utilisation de la mémoire
Articles Connexes
- Pratiques de l’équipe de test du système d’IA
- Débogage des Apps LLM : Un guide pratique pour le dépannage de l’IA
- Maîtrise du test des pipelines d’IA : Conseils, astuces et exemples pratiques
🕒 Published: