Fehlercode: 0100

Kategorie:Laufzeit: Fehler bei der Pufferzuweisung

Dieser Fehler weist darauf hin, dass der Speicherzuweiser 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:

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

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

  • Wenn Deallocations in der Warteschlange stehen, versucht die Laufzeit, fehlgeschlagene Zuweisungen noch einmal auszuführen.
  • 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 Arbeitsspeicher 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 des Modells reduzieren:
    • Batchgröße verringern: Wenn Sie die Batchgröße verringern, wird der Speicherverbrauch direkt gesenkt.
    • Parameter-Sharding: Bei sehr großen Modellen sollten 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 Eingabesequenzlänge den Arbeitsspeicherbedarf erheblich verringern.
    • Pufferzuweisung: 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.
    • Checkpoint-Strategie optimieren: Anstatt den gesamten Modellstatus auf einmal zu speichern, sollten Sie nur die Modellgewichte speichern oder eine Sharded-Checkpointing-Strategie verwenden.
  • Adressspeicherlayout und Auffüllung:
    • 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.

Tools

Aktivieren Sie das Flag tpu_log_allocations_on_oom. Der Zuweisungsmodul gibt dann einen detaillierten Bericht aller aktuellen Zuweisungen aus, wenn ein OOM-Fehler auftritt. Das kann beim Debuggen sehr hilfreich sein.