Fehlercode: E1000

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: TPU TensorCore Hbm usage: 34.82G, SparseCore Hbm usage 174.10G, exceeding available bytes: 95.74G

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.

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.
• Temporäre TPU-Speicher: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:

• „TC Hbm usage: X, SC Hbm usage Y“: Wenn der Fehler die HBM-Nutzung explizit aufschlüsselt, überschreitet die aggregierte TensorCore- (TC) + SparseCore-Nutzung (SC) das HBM-Limit. → Weiter zu Szenario 1. HBM-Nutzung von TC und SC ausgleichen:
• „Ran out of memory in memory space HBM“ (Speicherplatz im HBM-Speicher ist erschöpft): Prüfen Sie die Logs auf eine Auflistung der größten Zuweisungen im HBM.
  • Falls ein oder mehrere unerwartet große Tensoren vorhanden sind (z.B. > 50% des HBM-Limits) → Gehen Sie zu Szenario 2. Nicht genügend Arbeitsspeicher aufgrund unerwartet großer Zuweisungen.
  • Wenn in den Logs keine unerwartet großen Tensoren vorhanden sind, fahren Sie mit Szenario 3 fort. Nicht genügend Arbeitsspeicher aufgrund von aggregierten Zuweisungen.

Szenario 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“, bedeutet das, dass die aggregierte TensorCore- (TC) + SparseCore-Nutzung (SC) das HBM-Limit überschreitet. Vergleichen Sie die beiden Werte, um den Engpass zu ermitteln:

• Hohe SparseCore-Nutzung
  • HBM-Stack-Nutzung optimieren:Der Arbeitsspeicherverbrauch des HBM-Stacks skaliert mit feature_width, max_unique_nz_per_row und logical_replica_count. Sie können die maximale Stacknutzung 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. Bei Tabellen mit geringer Merkmalsbreite (z.B. weniger als 8 Gleitkommazahlen) ist der Padding-Overhead erheblich, was zu einer Verschwendung von HBM führt.
  • 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 korrekt ü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 Szenario 2 fort.
• Ausgeglichen
  • Wenn keiner der 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. Folgen Sie den Empfehlungen in allen drei Abschnitten.

Beispielszenario 2: Speicher ist aufgrund unerwartet großer Zuweisungen nicht ausreichend

Wenn die Fehlermeldung „Ran out of memory in memory space HBM“ angezeigt wird und in den Logs eine oder mehrere unerwartet große Zuweisungen vorhanden sind (> 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 umfangreichen Läufen verwenden, kann der Compiler gezwungen werden, den gesamten 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 beheben
  • Falsche Mesh-Formen oder fehlende Sharding-Anmerkungen können dazu führen, dass der Compiler standardmäßig Replikation 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.

Beispielszenario 3: Nicht genügend Arbeitsspeicher aufgrund von aggregierten Zuweisungen

Wenn die Fehlermeldung „Ran out of memory in memory space HBM“ angezeigt wird und keine unerwartet großen Tensoren in den Logs vorhanden sind, ist die Kapazität des Programms erschöpft, da die Gesamtsumme der Zuweisungen das HBM-Limit überschreitet. In diesem Fall ist es oft hilfreich, das Speicherprofil zu visualisieren, um die spezifischen Puffer zu identifizieren, die zur Spitzenlast 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.

Szenario 3.A: 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 Verringern der Batchgröße lässt sich die Arbeitsspeichernutzung oft reduzieren.
• 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. Diese Änderung kann sich auf das Verhalten des Modells auswirken und sollte daher sorgfältig abgewogen werden.

Szenario 3.B: 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 zu groß für die vorhandene Topologie 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:Mit Host-Offloading-Techniken können große Tensoren in den Host-CPU-Arbeitsspeicher ausgelagert werden, z.B. Aktivierungs-Offloading und Offloading des Optimiererstatus.

Szenario 3.C: Padding und Ausrichtung von Tensoren 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 Arbeitsspeichernutzung 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 über einen Puffer im Xprof Memory Viewer bewegen, wird die Karte mit den Pufferdetails angezeigt, die Pufferdetails einschließlich der Padding-Informationen enthält.
  • Beispiel: Eine Form von (129, 1024) könnte auf (256, 1024) aufgefüllt werden, 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. Diese Änderung kann sich auf das Verhalten des Modells auswirken und sollte daher sorgfältig geprüft werden.

Szenario 3.D: Wichtige XLA-Flags für den Arbeitsspeicher optimieren

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

Szenario 3.E: XLA-Rematerialisierungsdurchlauf optimieren/manuelle Prüfpunkte

Wenn das Modell fast in den Arbeitsspeicher passt, können Sie den Dekorator jax.checkpoint mit jax.grad verwenden, um manuell zu steuern, welche Zwischenwerte im Forward-Pass gespeichert und im Backward-Pass neu berechnet werden. So können Sie Rechenzyklen gegen HBM tauschen.

Alternativ können Sie den XLA::Rematerialization-Pass erzwingen, um die Speichereinsparungen zu priorisieren. Dies kann jedoch zu langsameren Kompilierungen führen:

Änderungen an XLA-Flags sollten nur als letztes Mittel in Betracht gezogen werden, da sie sich negativ auf die Leistung auswirken können.

Szenario 3.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.