Fehlercode: 0101

Kategorie:Laufzeit: Fehler bei der Programmzuteilung

Dieser Fehler gibt an, dass die XLA-Laufzeit auf einem TPU-Gerät ein kompiliertes ausführbares XLA-Programm nicht in den HBM der TPU laden konnte.

Beispiel für eine Fehlermeldung:

XlaRuntimeError: RESOURCE_EXHAUSTED: Error loading program 'jit_embedding_pipeline_step_fn': Attempting to reserve 29.49G at the bottom of memory. That was not possible. There are 147.64M free, 0B reserved, and 147.64M reservable. Scope: unknown..: while running replica 0 and partition 34 of a replicated computation (other replicas may have failed as well).

XLA-Back-Ends:TPU

Übersicht

Dieser Fehler hat in der Regel eine der folgenden Ursachen:

  • Die Programmgröße überschreitet den verfügbaren HBM: Das kompilierte XLA-Programm, einschließlich seiner Anweisungen, statischen Daten und aller eingebetteten Konstanten, ist größer als die Gesamtmenge an kostenlosem HBM, die derzeit auf den spezifischen TPU-Kernen verfügbar ist, auf denen das Programm geladen wird.
  • HBM-Fragmentierung: Obwohl der gesamte kostenlose HBM auf dem Gerät insgesamt ausreichend sein könnte, ist er nicht in einem einzelnen, zusammenhängenden Block verfügbar, der groß genug ist, um das gesamte Programm aufzunehmen.

Es ist wichtig zu verstehen, wie die TPU-Laufzeit den Arbeitsspeicher priorisiert. Pufferzuweisungen haben Vorrang vor geladenen Programmen. Wenn eine Pufferzuweisung fehlschlägt, werden bereits geladene Programme von der Laufzeit aus dem HBM entfernt, um Speicherplatz freizugeben. Das kann dazu führen, dass ein Programm, das zuvor erfolgreich geladen wurde, jetzt mit einem OOM-Fehler fehlschlägt, weil der HBM jetzt mit mehr Datenpuffern belegt ist.

Debugging

  • Zwischenspeicher-Arbeitsspeicherbedarf reduzieren: Wenn Sie den von Datenpuffern verwendeten Arbeitsspeicher freigeben, bleibt mehr Platz für das Programm selbst:
    • Batchgröße verringern: Dies ist eine der effektivsten Möglichkeiten, die für Aktivierungen verwendete Speichermenge zu reduzieren.
    • Parameter-Sharding: Bei sehr großen Modellen sollten Sie Modellparallelität oder Sharding-Techniken (z. B. FSDP oder Megascale) verwenden, um die Parameter und Berechnungen des Modells auf mehrere TPU-Kerne oder Hosts zu verteilen.
    • Sequenz-/Kontextlänge verkürzen: Bei Modellen, die sequenzielle Daten verarbeiten (z.B. Bei Modellen für die Verarbeitung natürlicher Sprache (Natural Language Processing, NLP) kann die Arbeitsspeichernutzung durch Verringern der Sequenzlänge erheblich reduziert werden.
    • Puffer spenden: Verwenden Sie Framework-Funktionen (z.B. jax.jit(..., donate_argnums=...)), damit XLA den Arbeitsspeicher von Eingabepuffern zum Speichern der Ausgabe wiederverwenden kann. Dadurch wird die Spitzenauslastung des Arbeitsspeichers reduziert.
  • Reduzieren Sie den Speicherbedarf des Programms für temporäre Variablen:
    • Mit dem Flag tpu_shared_memory_percent können Sie die Arbeitsspeichernutzung von temporären Variablen in Programmen reduzieren. Das kann sich negativ auf die Leistung auswirken.
  • Ausführungsstrategie optimieren/Bereitstellungslast reduzieren:
    • Programm laden: Wenn Sie mehrere Funktionen JIT-kompilieren, kann jede Funktion dazu führen, dass ein Programm geladen wird. Versuchen Sie, Ihre Arbeitslast so zu strukturieren, dass die Anzahl der gleichzeitig geladenen Programme minimiert wird.
  • Speicherlecks vermeiden:
    • Achten Sie darauf, dass Verweise auf jax.Array-Objekte nicht länger als vorgesehen aufrechterhalten werden. Wenn Sie jax.Array-Objekte beibehalten, kann dies die automatische Freigabe auch nach Abschluss der Programmkompilierung verhindern.

Tools

  • Aktivieren Sie das Flag tpu_log_allocations_on_oom. Der Zuweisungsmanager gibt dann bei einem OOM-Fehler einen detaillierten Bericht aller aktuellen Zuweisungen aus, was beim Debuggen sehr hilfreich sein kann.
  • Programm profilieren: Verwenden Sie den JAX Memory Profiler oder den TensorFlow Profiler, um einen detaillierten Überblick über die Speicherauslastung Ihres Programms im Zeitverlauf zu erhalten. So lassen sich unerwartete Spitzen beim Arbeitsspeicherverbrauch erkennen.