Fehlercode: E0100

Kategorie:Laufzeit: Fehler bei der Pufferzuweisung

Dieser Fehler gibt an, dass der Speicher-Allocator der XLA:TPU-Laufzeit keinen geeigneten Speicherblock im HBM des Beschleunigers für die angeforderte Zuweisung finden konnte.

Beispiel für eine Fehlermeldung:

ValueError: RESOURCE_EXHAUSTED: Error allocating device buffer: Attempting to allocate 8.00M. That was not possible. There are 6.43M free.; (0x0x1_HBM0)

XLA-Back-Ends:TPU

Übersicht

Dieser Fehler wird in folgenden Fällen ausgegeben:

  • Fehler bei der vom Nutzer initiierten Pufferzuweisung über jax.device_put

oder

  • Fehler bei der Zuweisung der Ausgaben von vom Nutzer geplanten Programmen.

Diese Fehler haben in der Regel folgende Ursachen:

  • Out of Memory (OOM): Der Nutzer versucht, einen Speicherbereich zuzuweisen, der größer ist als der insgesamt auf dem HBM der TPU verfügbare kostenlose Speicher.
  • Speicherfragmentierung:Die Zuweisung schlägt fehl, weil kein einzelner zusammenhängender kostenloser Block im Speicherplatz groß genug ist, um die angeforderte Größe zu erfüllen. Der insgesamt verfügbare Speicher ist für die Zuweisung ausreichend, aber er ist über den Speicherplatz in kleinen, nicht zusammenhängenden Blöcken verteilt.

Die TPU-Laufzeit verfügt über eine Reihe von Mechanismen, um Zuweisungsfehler zu wiederholen, darunter:

  • Wenn Deallocationen in der Warteschlange stehen, wiederholt die Laufzeit fehlgeschlagene Zuweisungen.
  • Bei OOMs, die durch eine Fragmentierung verursacht werden, kann die Laufzeit automatisch eine Defragmentierung und einen Wiederholungsversuch auslösen.
  • Bei der TPU-Laufzeit haben Pufferzuweisungen Vorrang vor dem Laden von Programmen. Wenn eine Pufferzuweisung aufgrund von unzureichendem HBM fehlschlägt, werden geladene TPU-Programme aus dem System entfernt, bis genügend Arbeitsspeicher für den Puffer verfügbar ist.

Ein Fehler, der nach den oben genannten Maßnahmen auftritt, erfordert in der Regel eine Nutzeraktion.

Debugging

  • Speicherbedarf Ihres Modells reduzieren
    • Batchgröße verringern:Wenn Sie die Batchgröße verringern, wird der Speicherverbrauch direkt gesenkt.
    • Parameter-Sharding:Bei sehr großen Modellen können Sie Techniken wie Modellparallelität oder Sharding verwenden, um Parameter auf den HBM mehrerer TPU-Kerne oder Hosts zu verteilen.
    • Sequenz-/Kontextlänge verkürzen:Bei Modellen, die auf Sequenzen basieren (z. B. Sprachmodelle), kann die Reduzierung der Länge der Eingabesequenz den Arbeitsspeicherbedarf erheblich verringern.
    • Pufferübertragung: Verwenden Sie Framework-Funktionen (z. B. jax.jit(..., donate_argnums=...)), um XLA zu signalisieren, dass bestimmte Eingabepuffer überschrieben und für Ausgaben wiederverwendet werden können. Weitere Informationen findest du unter Puffer spenden.
    • Checkpoint-Strategie optimieren:Anstatt den gesamten Modellstatus auf einmal zu speichern, sollten Sie nur die Modellgewichte speichern oder eine Strategie mit aufgeteilten Checkpoints verwenden.
  • Speicherlayout und Padding adressieren
    • Der TPU-Speicher wird in Blöcken zugewiesen und das Padding kann die tatsächliche Größe von Tensoren erhöhen.
  • 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.

Weitere Strategien zur Reduzierung des von den einzelnen Programmen verwendeten Arbeitsspeichers finden Sie unter Fehlercode: E1000.

Tools

Aktivieren Sie das Flag tpu_log_allocations_on_oom. Der Zuweisungsmodul gibt dann bei einem Out-of-Memory-Fehler einen detaillierten Bericht aller aktuellen Zuweisungen aus, was beim Debuggen sehr hilfreich sein kann.