Déboguer les erreurs OOM avec XProf

Les erreurs de mémoire insuffisante (OOM, Out Of Memory) se produisent lorsque la capacité de mémoire à haut débit (HBM, High Bandwidth Memory) de l'accélérateur (GPU ou TPU) est épuisée. Certaines causes courantes des problèmes d'OOM et des techniques de débogage sont détaillées dans la documentation E1000 : OOM HBM au moment de la compilation et dans la documentation JAX sur l'allocation de mémoire GPU.

Cette page explique comment utiliser l'outil Lecteur de mémoire de XProf pour visualiser l'utilisation de la mémoire de votre programme JAX, identifier les instances d'utilisation maximale et déboguer les erreurs OOM. Cela implique les étapes suivantes :

  1. Exécutez votre programme avec jax.profiler.trace pour capturer le profil.
  2. Démarrez XProf en arrière-plan et utilisez l'outil Memory Viewer pour afficher les détails de l'utilisation de la mémoire.

Exemple de programme

Le programme JAX suivant entraîne une erreur OOM :

import jax
from jax import random
import jax.numpy as jnp


@jax.profiler.trace("/tmp/xprof")
@jax.jit
def oom():
    a = random.normal(random.PRNGKey(1), (327680, 327680), dtype=jnp.bfloat16)
    return a @ a


if __name__ == "__main__":
    oom()

Sur une machine TPU, ce programme échoue avec le message suivant :

XlaRuntimeError: RESOURCE_EXHAUSTED: Allocation (size=107374182400) would exceed memory (size=17179869184) :: #allocation7 [shape = 'u8[327680,327680]{1,0:T(8,128)(4,1)}', space=hbm, size = 0xffffffffffffffff, tag = 'output of xor_convert_fusion@{}'] :: <no-hlo-instruction>

(Sur une machine avec GPU, l'erreur ressemble à ceci : XlaRuntimeError: RESOURCE_EXHAUSTED: Out of memory while trying to allocate 214748364800 bytes.)

Exécuter XProf

Installez xprof (pip install xprof), puis démarrez une instance XProf en spécifiant le répertoire dans lequel le profil est stocké :

xprof --logdir=/tmp/xprof/ --port=6006

Accédez à l'instance (sur une machine locale, à l'adresse http://localhost:6006). Dans le menu déroulant Outils, sélectionnez Memory Viewer (Visualiseur de mémoire), puis dans la fenêtre d'outil Memory Viewer, sélectionnez HBM dans le menu déroulant Memory Types (Types de mémoire) (généralement sélectionné par défaut).

Page du lecteur de mémoire XProf pour l&#39;exemple de programme ci-dessus

La documentation de l'outil XProf : Lecteur de mémoire décrit les composants de l'outil et les informations présentées.

Concentrez-vous sur la section Opérations HLO au moment de l'allocation maximale de mémoire, qui affiche trois graphiques de tampon au moment du pic d'utilisation de la mémoire. La mémoire tampon inclut : * Entrées et sorties du programme : lots d'entraînement, états de l'optimiseur, etc. * Temporaires TensorCore et SparseCore : mémoire dynamique requise pour les calculs intermédiaires (comme les activations, les gradients, etc.)

Vous pouvez pointer sur les graphiques de mémoire tampon pour obtenir plus de détails sur l'opération, comme sa taille, sa forme, son type d'allocation, etc. Cela peut vous aider à identifier et à évaluer les opérations qui peuvent avoir des temporaires élevés ou de longue durée, les grands tenseurs d'entrée/intermédiaires/de sortie qui ont un remplissage inefficace, etc., qui contribuent à la mémoire maximale et doivent être ajustés ou optimisés.

Découvrez des techniques de débogage spécifiques dans E1000 : Débogage.