OOM-Fehler mit XProf beheben

OOM-Fehler (Out-of-Memory) treten auf, wenn die HBM-Kapazität (High-Bandwidth Memory) des Beschleunigers (GPU oder TPU) erschöpft ist. Einige häufige Ursachen für OOM-Probleme und Debugging-Techniken werden in der Dokumentation zu E1000 – HBM-OOM zur Kompilierzeit und in der JAX-Dokumentation zur GPU-Speicherzuweisung beschrieben.

Auf dieser Seite wird beschrieben, wie Sie mit dem Memory Viewer-Tool von XProf die Speicherauslastung Ihres JAX-Programms visualisieren, Spitzenlasten identifizieren und OOM-Fehler beheben können. Dazu sind folgende Schritte erforderlich:

  1. Führen Sie Ihr Programm mit jax.profiler.trace aus, um das Profil zu erfassen.
  2. Starten Sie XProf im Hintergrund und verwenden Sie das Tool zur Speicheransicht, um Details zur Speicherauslastung aufzurufen.

Beispielprogramm

Das folgende JAX-Programm führt zu einem OOM-Fehler:

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()

Auf einer TPU-Maschine schlägt dieses Programm mit folgender Fehlermeldung fehl:

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>

Auf einer GPU-Maschine sieht der Fehler so aus: XlaRuntimeError: RESOURCE_EXHAUSTED: Out of memory while trying to allocate 214748364800 bytes.

XProf ausführen

Installieren Sie xprof (pip install xprof) und starten Sie eine XProf-Instanz, indem Sie das Verzeichnis angeben, in dem das Profil gespeichert ist:

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

Rufen Sie die Instanz auf (auf einem lokalen Computer unter http://localhost:6006). Wählen Sie im Drop-down-Menü Tools die Option Memory Viewer aus und wählen Sie im Toolfenster „Memory Viewer“ im Drop-down-Menü Memory Types die Option HBM aus (normalerweise standardmäßig ausgewählt).

XProf Memory Viewer-Seite für das oben genannte Beispielprogramm

In der Dokumentation zum Tool „XProf: Memory Viewer“ werden die Komponenten des Tools und die präsentierten Informationen beschrieben.

Konzentrieren Sie sich auf den Abschnitt HLO-Vorgänge bei maximaler Speicherzuweisung, in dem drei Pufferdiagramme für den Punkt der maximalen Speichernutzung angezeigt werden. Der Puffer umfasst: * Programmeingaben und ‑ausgaben: Trainingsbatches, Optimiererstatus usw. * TensorCore- und SparseCore-Temporaries: Dynamischer Speicher, der für Zwischenberechnungen (z. B. Aktivierungen, Gradienten usw.) erforderlich ist.

Wenn Sie den Mauszeiger auf die Pufferdiagramme bewegen, erhalten Sie weitere Details zum Vorgang, z. B. Größe, Form und Zuweisungstyp. So können Sie Vorgänge mit möglicherweise hohen oder langlebigen temporären Daten, große Eingabe-/Zwischen-/Ausgabetensoren mit ineffizientem Padding usw. identifizieren und bewerten, die zum maximalen Speicherbedarf beitragen und angepasst oder optimiert werden müssen.

Spezifische Debugging-Techniken finden Sie unter E1000: Debugging.