cj401-amd/debugging-jax-oom

Débogage d'un OOM JAX à partir d'une capture

Aperçu

Une étape JIT compilée de JAX alloue toute sa mémoire temporaire sous forme d'un seul bloc planifié. Ainsi, la mémoire de l'étape correspond au plan de XLA, et le fichier .xplane.pb contient à la fois ce plan et les événements d'exécution de l'allocateur. scripts/oom_report.py lit les deux en 2 à 5 secondes, en utilisant uniquement la bibliothèque standard :

• plan : l'affectation des tampons de chaque module, depuis /host:metadata. Il donne les paramètres (et combien sont donnés), les sorties non aliasées, le bloc temporaire et le pic du tas. Pour les tampons actifs à ce pic, il donne la taille, la forme, le rôle et la ligne Python. Le rôle provient de la pile de noms JAX : jvp( est direct (conservé pour la rétropropagation), transpose( est rétrograde, les sorties de call-custom s8 sont l'espace de travail de la bibliothèque.
• exécution : les événements de l'allocateur BFC. Ils donnent la capacité, les octets utilisés au début de chaque exécution, le pic, les attentes AllocatorRetry, les exécutions qui n'ont jamais lancé de noyau, et les événements hôte entre les deux dernières exécutions. Ils donnent également ce qui était vivant au début de la dernière exécution, regroupé par l'événement hôte qui l'a alloué : outputs of jit_step execution 3, BufferFromHostBuffer(s32[…]) (un device_put), ou « alloué avant le début de la capture ».

Le rapport se termine par des lignes de diagnostic, chacune étiquetée « mesurée » ou « déduite ». Exécutez-le d'abord :

python3 -P <chemin du répertoire de base de cette compétence>/scripts/oom_report.py <répertoire de capture ou .xplane.pb> [--module 'jit_train_step(123)'] [--top 25] [--json out.json]
python3 -P <chemin du répertoire de base de cette compétence>/scripts/oom_report.py --self-test

Par défaut, il analyse le module dont la dernière exécution n'a jamais lancé de noyau. S'il n'y en a pas, il prend le plus grand module qui a été exécuté. Utilisez --module pour en choisir un autre, par exemple une étape d'évaluation.

De la ligne de diagnostic à la cause et au correctif

Ligne de diagnostic	Cause	Correctif à proposer ; l'économie provient des octets du rapport
l'étape elle-même ne tient pas	le plan dépasse à lui seul la capacité du pool	Suivez la répartition des rôles au pic (ci-dessous)
↳ principalement forward, kept for the backward pass	activations : elles croissent avec batch × séquence ; l'attention naïve ajoute B·H·S·S par couche	jax.checkpoint / remat sur le bloc (économise la plupart des octets conservés en avant, coûte environ un passage avant supplémentaire) ; un micro-batch plus petit avec accumulation de gradient ; une séquence plus courte (les tampons S² rétrécissent d'un facteur 4 lorsque S est divisé par 2). Proposez l'attention fusionnée uniquement comme « vérifiez que les tampons B·H·S·S quittent le plan »
↳ principalement backward	gradients et leurs temporaires	Shard (FSDP) ou accumulez ; vérifiez les copies f32 de gradients bf16
↳ workspace	espace de travail GEMM/conv de la librairie	Généralement petit ; sur ROCm, il fait 64 MiB par GEMM hipBLASLt sur gfx950
grossit de X MiB par exécution, = 1.00x les sorties non-aliasées	le programme maintient en vie les sorties de chaque étape (une fuite entre les étapes) : la liste des vivants montre outputs of … execution 1, 2, …	Trouvez le détenteur dans les événements hôte entre les exécutions (par exemple $<unknown> append = list.append). Ensuite, supprimez la référence, réduisez à l'intérieur de l'étape (une norme, pas l'arbre), ou utilisez jax.device_get. Un pool plus grand ne fait que retarder l'OOM
aucun paramètre n'est donné … peut économiser jusqu'à X	les paramètres anciens et nouveaux / l'état de l'optimiseur sont vivants ensemble	donate_argnums sur les arguments d'état
d'autres tableaux vivants prennent la place ; le plus grand … vient de X	tableaux en dehors de l'étape : un jeu de données ou une copie d'évaluation placée sur le GPU, une prélecture, d'autres sorties JIT	Gardez-le sur l'hôte et diffusez-le, libérez-le, ou réduisez son type de données. « Alloué avant le début de la capture » signifie profiler plus tôt pour le nommer
AllocatorRetry … le plus long N s, n'a rien lancé	l'OOM s'est produite pendant la capture	Preuve uniquement. La requête ayant échoué est généralement le bloc temporaire de ce module, ou une sortie
XLA a déjà rematérialisé N instructions	le propre remat de XLA a déjà été exécuté jusqu'à la limite	Attendez-vous à ce que la correction vienne du programme, pas de XLA

Donnez à chaque correctif son économie estimée et son coût, étiquetés « déduit ». Vérifiez l'estimation par rapport à la capacité. « Augmenter XLA_PYTHON_CLIENT_MEM_FRACTION » corrige un plan qui ne tient pas seulement si le GPU a de la place. Une capacité bien inférieure à la HBM suggère que la fraction est définie ; la capture n'enregistre pas les variables d'environnement.

Anatomie mémoire du playbook en termes JAX

Le Ultra-Scale Playbook (« Profiling the memory usage ») divise la mémoire d'apprentissage en poids, gradients, état de l'optimiseur et activations. Dans le rapport :

• Poids + état de l'optimiseur = les paramètres de l'étape. Adam en fp32 fait 12 B par paramètre (p, m, v). Donné signifie pas de seconde copie.
• Gradients = le rôle backward au pic, plus les sorties si l'étape les retourne.
• Activations = forward, kept for the backward pass. Les modèles Scan / nn.scan les conservent sous forme de tampons empilés [layers, …] sur une boucle while.
• « L'étape 1 tient, l'étape N échoue » : dans PyTorch, c'est l'état de l'optimiseur créé à la première étape. Dans JAX, l'état existe avant l'étape 1, donc cherchez plutôt la ligne de croissance. Le premier appel compile et auto-optimise également, donc un pic avant la première exécution est lié à la compilation.

Règles

• Utilisez uniquement la capture. Ne reconstruisez pas le modèle et ne réexécutez pas la charge de travail pour tester un correctif : estimez-le à partir des octets du plan, et dites comment le confirmer (jit(...).lower(...).compile().memory_analysis() côté utilisateur).
• N'écrivez pas dans le répertoire de capture. Les convertisseurs XProf (xspace_to_tool_names, memory_viewer, CLI xprof) écrivent des fichiers de cache *.hlo_proto.pb à côté du .xplane.pb. Si vous avez besoin de vues XProf, copiez d'abord le .xplane.pb dans un répertoire temporaire.
• Étiquetez chaque nombre comme mesuré, lu ou déduit.

Référence mesurée (MI355X unique, petits pools, 2026-11-06)

Capture	Ce que montre le rapport
plan d'attention trop grand	besoin 3.72 GiB > 2.79 GiB ; pic de tas 81.6 % conservé en avant ; f32[8,8,1024,1024] ×5 à la ligne softmax ; bloc temporaire 3,495 MiB = l'erreur « 3.41GiB » ; AllocatorRetry 10 s ; 89 instructions de remat
gradients conservés à chaque étape	en utilisation au début de l'étape 316 → 1,252 MiB, +104.0 MiB par étape = 1.00× les sorties non-aliasées ; plan 788 MiB tient ; $<unknown> append entre les étapes
pas de donation	besoin 3.88 GiB > 2.79 GiB ; 1,902 MiB de paramètres + 1,902 MiB de sorties, 0 donnés ; temp seulement 172 MiB
MaxText llama2-7b, 8 GPU, FSDP (tient)	temp 61.9 GiB ; pic de tas 57 % conservé en avant, les piles scan bf16[32,4,4096,11008] ×2 via while.50

Limites

• Un hôte par fichier : les exécutions multi-hôtes ont un .xplane.pb par hôte, exécutez donc chaque fichier.
• HBM uniquement : seul l'espace mémoire 0 est décomposé. Le déchargement hôte apparaît comme memory space 5.
• Contiguïté : « tient » compare des octets. Le BFC nécessite également un bloc contigu, donc un quasi-ajustement peut encore échouer à cause de la fragmentation.
• Rôle « other » : tampons définis par des opérations créées par XLA sans pile de noms JAX, et sans alias en ayant une.
• CUDA : non testé. Il utilise les mêmes événements BFC et XLA, mais les noms ne sont pas vérifiés.
• TPU / eager : les captures TPU donnent uniquement les totaux du plan, sans pic de tas ni exécutions. Les programmes eager (opération par opération) apparaissent comme de nombreux petits modules, donc lisez la section d'exécution dans ce cas.