CUDA out of memory dans PyTorch : diagnostic et solutions

Une erreur CUDA out of memory signifie qu'une allocation de mémoire sur le GPU n'a pas pu être satisfaite. Dans PyTorch, commencez par identifier l'étape en échec, le périphérique utilisé et la différence entre mémoire occupée par les tenseurs et mémoire réservée par l'allocateur. Réduisez ensuite la charge ou supprimez les références inutiles. Un appel à empty_cache() ne libère pas les tenseurs dont le programme a encore besoin.

Ce guide s'adresse aux développeurs qui diagnostiquent un problème de mémoire pendant l'entraînement ou l'inférence avec PyTorch sur un GPU NVIDIA. Il propose une démarche courte, des exemples d'inspection et une méthode pour vérifier chaque modification. Il ne suppose ni que toute erreur provient de la fragmentation, ni qu'un GPU plus grand soit la première réponse utile.

Consigner l'erreur avant de modifier les réglages

Conservez l'exception complète, la taille de l'allocation refusée, le modèle de GPU, la version de PyTorch, les dimensions des entrées, la taille du lot et l'étape interrompue. L'erreur survient-elle au chargement des poids, pendant la passe avant, le calcul des gradients, la mise à jour de l'optimiseur ou une requête plus volumineuse ? Gardez une entrée qui reproduit l'échec pour comparer les essais.

Procédez dans cet ordre :

1. Vérifiez quel GPU le processus voit et si une autre tâche l'occupe.
2. Mesurez la mémoire depuis le processus en échec.
3. Réduisez une dimension de la charge et répétez le même cas.
4. Recherchez les tenseurs conservés entre les itérations.
5. Examinez l'allocateur lorsque les mesures le justifient.

Un GPU indisponible, un noyau manquant ou une incompatibilité du pilote demandent un autre diagnostic. Pour la génération d'images, notre guide des versions CUDA pour Stable Diffusion aide à distinguer compatibilité logicielle et manque de mémoire.

Vérifier le GPU et les autres processus

Ces commandes consultent la mémoire des GPU et affichent le récapitulatif habituel des périphériques et processus. Elles n'arrêtent aucune tâche et ne modifient pas le pilote :

nvidia-smi --query-gpu=index,uuid,name,memory.total,memory.used,memory.free --format=csv
nvidia-smi

Comparez l'UUID du GPU au périphérique attribué à votre tâche. Les conteneurs et les réglages de visibilité peuvent modifier la numérotation logique vue par l'application. Sur un GPU partagé, identifiez le propriétaire d'un processus concurrent avant de lui demander de libérer des ressources. N'arrêtez pas un processus inconnu pour faire passer un essai.

La documentation NVIDIA de nvidia-smi décrit les champs et les limites selon l'environnement. Une valeur non prise en charge peut apparaître sous la forme N/A. Ces résultats décrivent le périphérique, sans expliquer à eux seuls les allocations internes de PyTorch.

Mesurer la mémoire allouée, réservée et libre

Chaque mesure répond à une question différente :

• memory_allocated() mesure la mémoire GPU occupée par les tenseurs PyTorch du processus.
• memory_reserved() mesure la mémoire gérée par son allocateur avec cache. Elle inclut les tenseurs alloués : n'additionnez pas les deux valeurs.
• mem_get_info() renvoie la mémoire libre et totale du périphérique via CUDA. Les autres processus peuvent influencer la mémoire libre.
• max_memory_allocated() indique le pic des allocations de tenseurs depuis la dernière remise à zéro des statistiques, ou depuis le démarrage du programme.

Insérez cette fonction dans le processus à examiner. Un nouveau processus Python ne permet pas d'inspecter l'allocateur PyTorch d'un autre processus :

import torch

def report_cuda_memory(label, device=None):
    if not torch.cuda.is_available():
        print(f"{label}: CUDA is unavailable")
        return None
    if device is None:
        device = torch.cuda.current_device()
    gib = 1024 ** 3
    free, total = torch.cuda.mem_get_info(device)
    values = {
        "allocated_GiB": torch.cuda.memory_allocated(device) / gib,
        "reserved_GiB": torch.cuda.memory_reserved(device) / gib,
        "peak_allocated_GiB": torch.cuda.max_memory_allocated(device) / gib,
        "device_free_GiB": free / gib,
        "device_total_GiB": total / gib,
    }
    print(label, values)
    return values

# Call inside your existing training or inference process.
report_cuda_memory("before workload")

Appelez-la à des étapes constantes : avant le lot, après la passe avant, après la rétropropagation et après la mise à jour de l'optimiseur. Comparez plusieurs itérations, dont la plus grande entrée prévue. Les appels constituent des observations successives, pas un instantané atomique ; un autre processus peut allouer de la mémoire entre deux mesures. Conservez les libellés et les résultats avec la configuration défaillante.

Choisir le premier changement selon le symptôme

• Le premier lot représentatif échoue : essayez un microlot ou une entrée plus petite et notez si l'étape en échec change.
• La mémoire allouée augmente entre des itérations comparables : examinez les listes, fonctions de rappel, métriques et références aux résultats précédents.
• Seules les grandes entrées échouent : relevez leurs dimensions et testez une limite explicite ou un traitement séparé.
• Le périphérique est très occupé alors que les tenseurs du processus utilisent peu de mémoire : recherchez les autres processus et les allocations extérieures à PyTorch.

Ce sont des hypothèses de départ. Un petit lot qui passe une fois fournit un indice, mais ne prouve pas que l'entraînement complet ou la période d'inférence la plus chargée tiendra. Gardez le même ordre des entrées afin qu'un lot plus facile ne donne pas l'impression d'avoir corrigé le problème.

Supprimer les références accidentellement conservées

Des métriques ou des résultats stockés peuvent maintenir inutilement l'historique de calcul. Pour journaliser une perte scalaire, conservez un nombre Python :

# loss is a scalar tensor from the existing training step.
# Keep the list outside the training loop.
loss_values.append(loss.detach().item())

La FAQ de PyTorch explique ce problème de rétention. Examinez les attributs des objets et les fonctions de rappel, ainsi que les variables locales. Supprimer une référence ne suffit pas si une autre retient encore le tenseur. Limitez également les historiques de diagnostic, même lorsqu'ils sont stockés en mémoire CPU.

Utiliser les réglages d'inférence pour l'inférence

Le mode d'évaluation et le suivi des gradients sont deux contrôles distincts. Pour un modèle qui renvoie un seul tenseur, cette fonction active le comportement d'évaluation et exécute le calcul sans suivi autograd :

import torch

def predict_to_cpu(model, batch):
    # model and batch must already be on the same device.
    # This example expects a single Tensor output.
    model.eval()
    with torch.inference_mode():
        return model(batch).cpu()

inference_mode() n'appelle pas model.eval() à votre place. Ses tenseurs ont aussi des restrictions lors d'un calcul ultérieur avec gradients. Utilisez no_grad() lorsque ces restrictions ne conviennent pas. Adaptez la sortie pour les dictionnaires ou tuples et réactivez le mode entraînement si vous reprenez l'entraînement.

Pour servir un modèle de langage, définissez aussi des limites de longueur d'entrée, de sortie et de requêtes simultanées. Notre guide du cache KV explique pourquoi loger les poids ne suffit pas à dimensionner l'inférence.

Réduire la mémoire d'entraînement en choisissant le compromis

Microlots et accumulation des gradients

De plus petits microlots peuvent réduire les données traitées simultanément. L'accumulation permet à plusieurs microlots de contribuer à une mise à jour de l'optimiseur. Respectez la normalisation de la perte et le rythme des mises à jour décrits dans les exemples d'accumulation PyTorch, surtout avec la mise à l'échelle des gradients. Traitez explicitement le dernier groupe s'il est incomplet. Vérifiez les effets sur l'entraînement, au-delà de la seule consommation mémoire.

Précision mixte

La précision mixte automatique utilise une précision réduite pour les opérations admissibles tout en conservant une précision supérieure où elle est nécessaire. Vérifiez le matériel, les opérateurs et le comportement numérique. Ne supposez pas que la mémoire totale sera divisée par deux. Comparez la perte et la qualité des résultats en plus du pic mémoire.

Recalcul des activations

L'activation checkpointing recalcule certaines opérations de la passe avant pendant la rétropropagation pour conserver moins d'intermédiaires. Il échange du calcul contre de la mémoire. PyTorch recommande de passer explicitement use_reentrant=False pour la variante non réentrante. Vérifiez les gradients et le temps d'exécution avec votre modèle. Ce mécanisme diffère d'une sauvegarde sur disque permettant de reprendre un entraînement.

Comprendre ce que libère empty_cache()

empty_cache() rend les blocs inutilisés du cache disponibles aux autres applications GPU. Il ne libère pas les tenseurs encore utilisés et n'agrandit pas le GPU. Il peut aider dans certains cas de fragmentation, mais ne remplace ni la réduction de la charge ni la suppression des références inutiles.

Avant de redémarrer le noyau de votre propre notebook, sauvegardez le travail nécessaire puis vérifiez si la même entrée provoque encore l'échec. Le redémarrage offre un point de comparaison propre ; il n'identifie pas la cause à lui seul.

Examiner la fragmentation à partir des mesures

Un grand écart entre mémoire réservée et allouée ne prouve pas une fragmentation. Consultez memory_summary() ou memory_stats() avant de régler l'allocateur. Les notes actuelles sur la mémoire CUDA utilisent PYTORCH_ALLOC_CONF ; PYTORCH_CUDA_ALLOC_CONF reste un alias de compatibilité. Les options dépendent du moteur d'allocation. Ainsi, max_split_size_mb est un réglage de dernier recours pour l'allocateur natif, ignoré par cudaMallocAsync.

Si le rapport court ne suffit pas, suivez la méthode des instantanés mémoire sur une reproduction limitée. Elle révèle l'historique des allocations, mais pas toutes celles effectuées hors de PyTorch, notamment par certaines bibliothèques de communication. Examinez les chemins et traces capturés avant de partager un instantané.

Vérifier la correction avant de changer de matériel

Répétez le cas défaillant, puis testez la plus grande entrée autorisée et une exécution prolongée. Relevez le pic mémoire, l'achèvement, le débit et les contrôles de qualité utiles. Modifiez une seule variable à la fois. Conservez la configuration réussie avec les versions des paquets pour éviter qu'un déploiement suivant annule discrètement la correction.

Si la charge exige toujours davantage de mémoire après des réductions raisonnables, comparez un GPU plus capacitaire avec une méthode de partitionnement ou de déchargement prise en charge. Le DistributedDataParallel classique utilise des répliques du modèle : ajouter un GPU ne met pas automatiquement sa mémoire à disposition d'une réplique. Notre guide PyTorch sur A100 couvre la configuration et les mesures propres à ce GPU.

Avant de louer une instance plus grande, définissez précisément la charge à terminer et les mesures à comparer. Un petit exemple réussi est un point de départ. Votre test d'acceptation doit inclure les tailles d'entrée, la simultanéité et la durée qui ont provoqué l'échec initial.