Fehlercode: 1000

Kategorie:Kompilierungszeit: HBM-OOM

Dieser Fehler weist darauf hin, dass das Programm mehr HBM (High Bandwidth Memory, Speicher mit hoher Bandbreite) benötigt, als auf dem TPU-Gerät physisch verfügbar ist.

Beispiele für Fehlermeldungen:

RESOURCE_EXHAUSTED: XLA:TPU compile permanent error. Ran out of memory in memory space hbm. Used 49.34G of 32.00G hbm. Exceeded hbm capacity by 17.34G.

RESOURCE_EXHAUSTED: TPU TensorCore Hbm usage: 34.82G, SparseCore Hbm usage 174.10G, exceeding available bytes: 95.74G

XLA-Back-Ends:TPU

Übersicht

XLA führt Prüfungen durch, um sicherzustellen, dass die Gesamtgröße aller erforderlichen statischen Zuweisungen in den HBM des Geräts passt.

Der Compiler verwaltet die feste HBM-Kapazität der TPU für verschiedene Arten von Zuweisungen:

• Programmeingaben und ‑ausgaben:Trainingsbatches, Optimierungsstatus usw.
• TensorCore + SparseCore Temporaries:Dynamischer Speicher, der für Zwischenberechnungen (z. B. Aktivierungen, Gradienten usw.) erforderlich ist.
• Kompilierte Binärdatei:Der Maschinencode für TensorCore (TC) und SparseCore (SC).
• System-Overhead:Reservierter Speicherplatz für die XLA-Laufzeit (z.B. Infeed-Puffer bei älteren TPU-Generationen).
• Konstanten:Konstante Werte, die in die HLO-IR eingebettet sind, werden auf HBM zugewiesen.
• Compiler-Interna:Zuweisungen auf Programmebene und pro HLO (z.B. Routinginformationen für Knoten im Mesh)

Dieser Fehler tritt auf, wenn der XLA-Compiler nicht alle oben genannten Zuweisungen in den HBM des Geräts einpassen kann.

Debugging

Analysieren Sie die Fehlermeldung und die Logs sorgfältig, um festzustellen, welche der folgenden Kategorien von HBM-OOM am besten auf Ihren Fehler zutrifft:

• TensorCore (TC) + SparseCore (SC) HBM Usage Exceeds Limit (HBM-Nutzung von TensorCore (TC) + SparseCore (SC) überschreitet das Limit): Wenn der Fehler die Nutzung explizit aufschlüsselt, z.B. „TC Hbm usage: X, SC Hbm usage Y“. → Direkt zu Abschnitt 1. HBM-Nutzung von TC und SC ausgleichen:
• Unerwartet große Zuweisungen: Wenn der Fehler „Ran out of memory in memory space HBM“ (Speicherplatz HBM ist nicht mehr verfügbar) lautet, sehen Sie in den Logs nach, welche Zuweisungen auf HBM am größten sind. Falls ein oder mehrere unerwartet große Tensoren (z.B. > 50% des HBM-Limits) vorhanden sind → Gehen Sie zu Abschnitt 2. Unerwartet große Zuweisungen:
• Zusammengefasste Zuweisungen überschreiten HBM-Grenzwert: Wenn im Fehler „Ran out of memory in memory space HBM“ steht, aber keine unerwartet großen Tensoren in den Logs vorhanden sind, fahren Sie mit Abschnitt 3 fort. Die aggregierten Zuweisungen überschreiten das HBM-Limit.

Bereich 1. HBM-Nutzung von TC und SC ausgleichen

Wenn der Fehler die Nutzung explizit aufschlüsselt, z.B. „TC Hbm usage: X, SC Hbm usage Y“: Vergleichen Sie die beiden Werte, um den Engpass zu ermitteln.

• Hohe SparseCore-Nutzung:
  • HBM-Stack-Nutzung optimieren:Der HBM-Stack-Speicherverbrauch skaliert mit feature_width, max_unique_nz_per_row und logical_replica_count. Sie können die maximale Stapelnutzung reduzieren, indem Sie das Flag --xla_sc_num_serialized_tables_to_optimize_hbm anpassen, das die Verarbeitung von Tabellen serialisiert. Dies geht jedoch auf Kosten der Parallelität.
  • Padding-Overhead prüfen:SparseCore richtet Einbettungstabellen an 32 Byte (8 Gleitkommazahlen) aus. Tabellen mit geringen Spaltenbreiten (z.B. < 8 Floats) verursachen einen erheblichen Padding-Overhead, wodurch HBM verschwendet wird.
  • Heap-Nutzung reduzieren:Hohe Werte für maximum_parallel_iterations erhöhen die Menge der Eingabedaten, die in den HBM-Heap vorab abgerufen werden. Wenn Sie diesen Wert verringern, kann das zu einer erheblichen Speicherfreigabe führen.
  • Sharding überprüfen:Prüfen Sie, ob die Einbettungstabellen ordnungsgemäß über alle Chips hinweg mod-sharded sind. Weitere Informationen finden Sie unter So werden Limits in Tabellen umgesetzt.
  • Weitere Ideen finden Sie unter SC: Performance and memory bottlenecks.
• Hohe TensorCore-Nutzung:
  • Fahren Sie mit Abschnitt 2 fort.
• Ausgeglichen
  • Wenn keiner der beiden Werte einzeln zu hoch ist, aber die Summe zu hoch ist, ist die Kapazität des Chips erreicht. Sie müssen versuchen, die Nutzung beider Komponenten zu reduzieren. Befolgen Sie die Empfehlungen in allen drei Abschnitten.

Bereich 2. Unerwartet große Zuweisungen

Wenn in den Logs eine oder mehrere unerwartet große Zuweisungen vorhanden sind (über 50% des HBM-Limits), handelt es sich fast nie um ein Problem mit der Hardwarekapazität. In der Regel handelt es sich um einen Konfigurationsfehler. Sehen Sie sich das XLA-Label (falls vorhanden) der großen Zuweisungen an, um Hinweise auf den JAX-Quellcode zu erhalten.

• Debugging-Artefakte entfernen:
  • Wenn Sie jax.debug.print() in groß angelegten Läufen verwenden, kann der Compiler gezwungen werden, den vollständigen Tensor im HBM zu materialisieren, um ihn an die CPU zu übertragen. Dadurch wird die Fusion unterbrochen und die maximale Arbeitsspeichernutzung erhöht. Entfernen Sie alle verbleibenden jax.debug.print().
• Ineffiziente Mesh-Formen oder Sharding-Probleme beheben:
  • Falsche Mesh-Formen oder fehlende Sharding-Anmerkungen können dazu führen, dass der Compiler standardmäßig Replication verwendet und versucht, sehr große Tensoren auf einem einzelnen Chip unterzubringen.
  • Prüfen Sie die Formen der großen Zuweisungen und vergewissern Sie sich, dass das Sharding von XLA korrekt angegeben und weitergegeben wird.

Bereich 3. Aggregierte Zuweisungen überschreiten das HBM-Limit

Wenn das Programm aufgrund der Gesamtsumme der Zuweisungen, die das HBM-Limit überschreitet, nicht mehr genügend Kapazität hat, ist es oft hilfreich, das Speicherprofil zu visualisieren, um die spezifischen Puffer zu identifizieren, die zur Spitzennutzung beitragen. Eine Schritt-für-Schritt-Anleitung zum Ermitteln der Hauptursachen für den Spitzenverbrauch von Arbeitsspeicher finden Sie unter OOM-Fehler mit XProf debuggen.

Nachdem Sie einige der wichtigsten Faktoren ermittelt haben, können Sie den Speicherbedarf mit den folgenden Schritten optimieren.

A. Tensor-Padding und ‑Ausrichtung prüfen

Ineffiziente Tensorformen sind eine häufige, stille Ursache für OOM-Fehler auf TPUs. Um die Spitzenleistung auf TPUs zu erzielen, füllt XLA Tensordimensionen auf – in der Regel auf Vielfache von 128 für die kleinste Dimension und 8 für die zweitkleinste. Diese Auffüllung wirkt sich sowohl auf Eingabearrays als auch auf temporäre Zwischen-Tensoren (HLO-Temporaries) aus und kann die Speichernutzung erheblich erhöhen, insbesondere bei kleinen Dimensionsgrößen. Weitere Informationen finden Sie unter Array-Layouts.

• Formen großer Puffer prüfen: (auf TPU v5 mit Standardlayouts)
  • Wenn Sie den Mauszeiger im Xprof Memory Viewer über einen Puffer bewegen, wird die Karte mit den Pufferdetails angezeigt, die auch Informationen zum Padding enthält.
  • Beispiel: Eine Form von (129, 1024) wird möglicherweise auf (256, 1024) aufgefüllt, was zu einer Speicherverschwendung von fast 50% führt.
  • Korrektur:Für eine Form von (128, 1024) ist kein Padding erforderlich und es kommt zu keinem Speicherverlust.
• Dimensionen ausrichten:Alle großen Tensordimensionen (Batchgröße, Einbettungsdimension, verborgene Größe) müssen Vielfache von 128 sein.

B. Konfiguration anpassen

Mit den folgenden Konfigurationsanpassungen lassen sich OOM-Fehler häufig beheben:

• Batchgröße reduzieren:Der für Zwischenaktivierungen und ‑gradienten benötigte Arbeitsspeicher ist direkt proportional zur Batchgröße. Durch Reduzieren der Batchgröße lässt sich die Arbeitsspeichernutzung oft verringern.
• Eingabepuffer spenden:Wenn Sie jax.jit verwenden, geben Sie donate_argnums für Ihre Modellparameter an. Dadurch kann XLA den Eingabespeicher mit der Ausgabe überschreiben.
• Gemischte Präzision (bfloat16) aktivieren: Verwenden Sie bfloat16 oder Quantisierung (int8 usw.) für die größten Tensoren im Programm, sofern die Modellarchitektur und die Qualitätsanforderungen dies zulassen.

C. Architektur und Sharding optimieren

Wenn Konfigurationsänderungen nicht ausreichen, ist die Modelltopologie möglicherweise zu groß für die aktuelle Hardwarekonfiguration.

• Neuere TPU-Generationen verwenden:Neuere TPUs bieten in der Regel mehr HBM pro Chip. Wechseln Sie zu neueren TPU-Generationen, sofern verfügbar.
• Auf einer größeren Chip-Topologie ausführen:Wenn die Modellgewichte für die vorhandene Topologie zu groß sind, können Sie versuchen, sie auf mehr Chips aufzuteilen.
• Erweiterte Sharding-Techniken implementieren:
  • Erwägen Sie komplexere Ansätze für Daten-, Tensor- oder Pipeline-Parallelität.
  • Geben Sie Sharding-Hinweise für Zwischenwerte und Ausgaben an.
• JAX-Host-Offloading verwenden:Lagern Sie große Tensoren in den Host-CPU-Speicher aus, z. B. Aktivierungs-Offloading und Offloading des Optimierungsstatus.

D. Wichtige XLA-Flags für den Arbeitsspeicherverbrauch optimieren:

Wichtige Speicher-Flags können so angepasst werden, dass die Leistung zugunsten einer geringeren Arbeitsspeichernutzung reduziert wird. Sie sollten jedoch nur als letztes Mittel verwendet werden, da sie sich negativ auf die Leistung auswirken können.

E. XLA-Rematerialization-Pass / Manuelle Prüfpunktausführung optimieren

Wenn das Modell fast in den Arbeitsspeicher passt, können Sie den XLA::Rematerialization-Pass erzwingen, um Arbeitsspeicher zu sparen. Das kann jedoch zu langsameren Kompilierungen führen:

Alternativ können Sie den Decorator jax.checkpoint mit jax.grad verwenden, um manuell zu steuern, welche Zwischenwerte beim Forward-Pass gespeichert und beim Backward-Pass neu berechnet werden. So können Sie Rechenzyklen gegen HBM tauschen.

F. Erweiterte Profiling-Tools verwenden

OOM-Fehler mit XProf debuggen enthält eine Anleitung zur Verwendung des XProf Memory Viewer, um die Sicht des Compilers auf die HBM-Nutzung zu visualisieren.

Mit diesem Tool können Sie die maximale Speicherzuweisung und die Pufferlebensdauer sehen. Das ist wichtig, um genau zu verstehen, was HBM zum Zeitpunkt der maximalen Auslastung verbraucht. Allgemeine Informationen zur Einrichtung der Profilerstellung finden Sie unter Erste Schritte mit Xprof und TensorBoard Profiling.