Fehlercode: 1001

Kategorie:Kompilierungszeit: Scoped Vmem OOM

Dieser Fehler gibt an, dass das Programm mehr Scoped Vector Memory (Vmem) benötigt, als zugewiesen wurde.

Beispiele für Fehlermeldungen:

RESOURCE_EXHAUSTED: Ran out of memory in memory space vmem while allocating on stack for %my-custom-kernel = bf16[2048,4096]{1,0:T(8,128)(2,1)} custom-call(...) ...

XLA-Back-Ends:TPU

Übersicht

TPUs haben einen Vektorspeicher (Vector Memory, VMEM), der ein lokaler Scratchpad-Speicher ist, der ausschließlich vom TensorCore (TC) verwendet wird. Der Compiler verwaltet Vmem für verschiedene Arten von Zuweisungen:

  • Zuweisungen auf Anweisungsebene:Temporärer Speicher in Vmem während der Ausführung einer einzelnen HLO-Anweisung. Dazu gehören der Operanden-Span-Puffer (z.B. für Double-Buffering) und Register-Spills.
  • Programmbezogene Zuweisungen:Zuweisungen, die über den Umfang einer einzelnen HLO-Anweisung hinausgehen. Dabei handelt es sich in der Regel um temporäre HLO-Variablen und Zwischenergebnisse, die Ein- und/oder Ausgaben von HLO-Anweisungen sind.

Ein Vmem-OOM mit Kompilierzeitbereich tritt auf, wenn die anweisungsbezogenen Zuweisungen das Zuweisungslimit für diese Anweisung überschreiten. Dieses Limit wird kontrolliert

Diese Fehler werden in der Regel durch einen internen Compilerfehler oder durch einen benutzerdefinierten Kernel verursacht, der sein Zuweisungslimit überschreitet.

Debugging

Analysieren Sie die Fehlermeldung sorgfältig, um festzustellen, ob der Fehler auf einen benutzerdefinierten Kernel oder einen Standard-HLO zurückzuführen ist. Ein Fehler aufgrund eines benutzerdefinierten Kernels sollte die folgende Signatur haben:

Ran out of memory in memory space vmem while allocating on stack for %my-custom-call = <output-shape> custom-call(<params>), custom_call_target="tpu_custom_call" ...
  • Custom Kernel Scoped Vmem OOM: Wenn der Fehler auf einen benutzerdefinierten Kernel verweist → Kernel neu abstimmen
  • Vmem-Probleme ohne Kernel: Wenn der Vmem-OOM-Fehler aufgrund eines Vorgangs ohne benutzerdefinierten Kernel auftritt, handelt es sich wahrscheinlich um einen internen Compilerfehler. Bitte melden Sie einen Fehler zu XLA mit einem HLO-Dump.

Kernel neu abstimmen

Wenn der Fehler auf einen benutzerdefinierten Kernel zurückzuführen ist, können Sie mit den folgenden Methoden den Speicherbedarf des Kernels reduzieren:

  • Blockgrößen anpassen:Reduzieren Sie die Blockgrößen (Kachelgrößen) in Ihrer Kernelkonfiguration, um die Verwendung von Scoped Vmem zu verringern.
  • Per-Kernel Scoped Vmem Limits festlegen:Fordern Sie mit dem vmem_limit_bytes-Parameter explizit die erforderliche Menge an Arbeitsspeicher für diesen bestimmten Kernel an.
  • Speicherfarben ändern:Die Ein-/Ausgaben des Kernels für VMEM können mit pallas.tpu.with_memory_space_constraint explizit eingefärbt/eingeschränkt werden. Achten Sie jedoch darauf, nicht zu viele Ein-/Ausgaben für Vmem zu kennzeichnen, da dies zu einem allgemeinen VMEM-OOM führen kann.
  • Wenn die kernelspezifische Anpassung schwierig ist oder das Problem viele Kerne betrifft, können Sie das globale Vmem-Limit mit dem Flag --xla_tpu_scoped_vmem_limit_kib anpassen.