Fehlersuche bei JAX OOM aus einem Profiling-Capture

Fehlersuche bei JAX OOM aus einem Capture

Überblick

Ein gejitter JAX-Schritt erhält seinen gesamten temporären Speicher als einen einzigen geplanten Temp-Block. Daher ist der Speicher des Schritts der Plan von XLA, und die .xplane.pb enthält sowohl diesen Plan als auch die Laufzeit-Ereignisse des Allokators. scripts/oom_report.py liest beides in 2–5 Sekunden ein, ausschließlich mit der Standardbibliothek:

• Plan: die Pufferzuweisung jedes Moduls, aus /host:metadata. Er liefert die Parameter (und wie viel davon gespendet wird), die nicht-aliassierten Ausgaben, den Temp-Block und den Heap-Peak. Für die Puffer, die zu diesem Peak leben, gibt er Größe, Form, Rolle und Python-Zeile an. Die Rolle stammt aus dem JAX-Namenstack: jvp( ist Forward (für den Rückwärtsdurchlauf behalten), transpose( ist Backward, s8-Custom-Call-Ausgaben sind Bibliotheks-Workspace.
• Laufzeit: die Ereignisse des BFC-Allokators. Sie geben die Kapazität, die zu Beginn jeder Ausführung belegten Bytes, den Peak, AllocatorRetry-Wartezeiten, Ausführungen, die nie einen Kernel gestartet haben, und die Host-Ereignisse zwischen den letzten beiden Ausführungen. Sie geben auch an, was zu Beginn der letzten Ausführung lebendig war, gruppiert nach dem Host-Ereignis, das es alloziert hat: outputs of jit_step execution 3, BufferFromHostBuffer(s32[…]) (ein device_put) oder „allocated before the capture started”.

Der Bericht endet mit Diagnosezeilen, jede als gemessen oder abgeleitet gekennzeichnet. Führen Sie es zuerst aus:

python3 -P <Basisverzeichnis dieses Skills>/scripts/oom_report.py <Capture-Verzeichnis oder .xplane.pb> [--module 'jit_train_step(123)'] [--top 25] [--json out.json]
python3 -P <Basisverzeichnis dieses Skills>/scripts/oom_report.py --self-test

Standardmäßig wird das Modul analysiert, dessen letzte Ausführung nie gestartet wurde. Falls es keines gibt, wird das größte ausgeführte Modul genommen. Mit --module kann ein anderes gewählt werden, z. B. ein Evaluierungsschritt.

Von der Diagnosezeile zur Ursache und Behebung

Diagnosezeile	Ursache	Vorzuschlagende Behebung; die Einsparung ergibt sich aus den Bytes des Berichts
der Schritt selbst passt nicht	der Plan allein übersteigt den Pool	Folgen Sie der Rollenaufteilung am Peak (unten)
↳ hauptsächlich forward, kept for the backward pass	Aktivierungen: sie wachsen mit Batch × Seq; naive Aufmerksamkeit fügt B·H·S·S pro Layer hinzu	jax.checkpoint / Remat auf dem Block (spart die meisten Forward-kept-Bytes, kostet etwa einen zusätzlichen Forward-Durchlauf); eine kleinere Micro-Batch mit Gradientenakkumulation; eine kürzere Sequenz (S²-Puffer schrumpfen um den Faktor 4, wenn S halbiert wird). Bieten Sie gefaltete Aufmerksamkeit nur als „Überprüfen, ob die B·H·S·S-Puffer den Plan verlassen” an
↳ hauptsächlich backward	Gradienten und deren Temporäre	Shard (FSDP) oder akkumulieren; auf f32-Kopien von bf16-Gradienten prüfen
↳ workspace	GEMM-/Conv-Bibliotheks-Temp	Normalerweise klein; auf ROCm sind es 64 MiB pro hipBLASLt-GEMM auf gfx950
wächst um X MiB pro Ausführung, = 1,00× die nicht-aliassierten Ausgaben	das Programm hält die Ausgaben jedes Schritts am Leben (ein Leck über Schritte hinweg): die Live-Liste zeigt outputs of … execution 1, 2, …	Finden Sie den Halter in den Host-Ereignissen zwischen den Ausführungen (z. B. $<unknown> append = list.append). Dann entfernen Sie die Referenz, reduzieren Sie innerhalb des Schritts (eine Norm, nicht der Baum) oder verwenden Sie jax.device_get. Ein größerer Pool verzögert nur den OOM
kein Parameter wird gespendet … kann bis zu X sparen	alte und neue Parameter/Optimiererzustand sind gleichzeitig live	donate_argnums auf die Zustandsargumente
andere live Arrays nehmen den Platz; das größte … kam von X	Arrays außerhalb des Schritts: ein Dataset oder eine Eval-Kopie auf der GPU, Prefetch, andere jit-Ausgaben	Auf dem Host behalten und streamen, freigeben oder den Datentyp verkleinern. „Allocated before the capture started” bedeutet, früher zu profilen, um es zu benennen
AllocatorRetry … längste N s, keine gestartet	Der OOM geschah innerhalb des Captures	Nur ein Beleg. Die fehlgeschlagene Anforderung ist normalerweise der Temp-Block des Moduls oder eine Ausgabe
XLA hat bereits N Anweisungen rematerialisiert	XLAs eigenes Remat wurde bereits gegen das Limit ausgeführt	Erwarten Sie, dass die Behebung vom Programm kommt, nicht von XLA

Geben Sie für jede Behebung die geschätzte Einsparung und ihre Kosten an, gekennzeichnet als abgeleitet. Überprüfen Sie die Schätzung anhand der Kapazität. „XLA_PYTHON_CLIENT_MEM_FRACTION erhöhen” behebt einen Plan, der nicht passt, nur wenn die GPU den Platz hat. Eine Kapazität deutlich unter HBM deutet darauf hin, dass der Bruch gesetzt ist; das Capture zeichnet keine Umgebungsvariablen auf.

Speicheranatomie des Playbooks in JAX-Begriffen

Das Ultra-Scale Playbook („Profiling the memory usage”) unterteilt das Trainingsspeicher in Gewichte, Gradienten, Optimiererzustand und Aktivierungen. Im Bericht:

• Gewichte + Optimiererzustand = die Parameter des Schritts. Adam in fp32 sind 12 B pro Parameter (p, m, v). Gespendet bedeutet keine zweite Kopie.
• Gradienten = die Rolle backward am Peak, plus Ausgaben, falls der Schritt sie zurückgibt.
• Aktivierungen = forward, kept for the backward pass. Scan-/nn.scan-Modelle halten sie als gestapelte [Layers, …]-Puffer in einer while-Schleife.
• „Schritt 1 passt, Schritt N schlägt fehl”: In PyTorch wird der Optimiererzustand beim ersten Schritt erstellt. In JAX existiert der Zustand vor Schritt 1, suchen Sie daher nach der Wachstumszeile. Der erste Aufruf kompiliert auch und führt Autotuning durch, daher ist ein Peak vor der ersten Ausführung Kompilierzeit.

Regeln

• Nur das Capture verwenden. Bauen Sie das Modell nicht neu oder führen Sie den Workload nicht erneut aus, um eine Behebung zu testen: Schätzen Sie sie aus den Bytes des Plans und sagen Sie, wie sie bestätigt werden kann (jit(...).lower(...).compile().memory_analysis() auf der Benutzerseite).
• Nicht in das Capture-Verzeichnis schreiben. Die Konverter von XProf (xspace_to_tool_names, memory_viewer, xprof CLI) schreiben *.hlo_proto.pb-Cache-Dateien neben die .xplane.pb. Wenn Sie XProf-Ansichten benötigen, kopieren Sie die .xplane.pb zuerst in ein temporäres Verzeichnis.
• Jede Zahl als gemessen, gelesen oder abgeleitet kennzeichnen.

Gemessene Referenz (Single MI355X, kleine Pools, 2026-11-06)

Capture	Was der Bericht zeigt
Aufmerksamkeitsplan zu groß	benötigt 3,72 GiB > 2,79 GiB; Heap-Peak 81,6 % Forward-kept; f32[8,8,1024,1024] ×5 an der Softmax-Zeile; Temp-Block 3.495 MiB = die „3,41 GiB” der Fehlermeldung; AllocatorRetry 10 s; 89 Remat-Anweisungen
Gradienten bei jedem Schritt behalten	belegt bei Schrittstart 316 → 1.252 MiB, +104,0 MiB pro Schritt = 1,00× die nicht-aliassierten Ausgaben; Plan 788 MiB passt; $<unknown> append zwischen den Schritten
keine Spende	benötigt 3,88 GiB > 2,79 GiB; 1.902 MiB Parameter + 1.902 MiB Ausgaben, 0 gespendet; Temp nur 172 MiB
MaxText llama2-7b, 8 GPUs, FSDP (passt)	Temp 61,9 GiB; Heap-Peak 57 % Forward-kept, die Scan-Stapel bf16[32,4,4096,11008] ×2 via while.50

Grenzen

• Ein Host pro Datei: Multi-Host-Läufe haben eine .xplane.pb pro Host, daher jede einzeln ausführen.
• Nur HBM: Nur Speicherraum 0 wird aufgeschlüsselt. Host-Offload wird als memory space 5 angezeigt.
• Kontiguität: „Passt” vergleicht Bytes. BFC benötigt auch einen zusammenhängenden Block, daher kann ein knappes Passen durch Fragmentierung dennoch fehlschlagen.
• Rolle „other”: Puffer, die von XLA-erzeugten Operationen ohne JAX-Namenstack und ohne Alias, der einen hat, definiert werden.
• CUDA: Nicht getestet. Es verwendet dieselben BFC- und XLA-Ereignisse, aber die Namen sind unbestätigt.
• TPU / Eager: TPU-Captures geben nur Pläne als Gesamtsummen an, ohne Heap-Peak und ohne Ausführungen. Eager (Op-für-Op)-Programme erscheinen als viele kleine Module, daher dort den Laufzeitabschnitt lesen.