Fehlercode: E1001

Kategorie:Kompilierzeit: Bereichsbezogener Vmem-Fehler (Out of Memory)

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

Beispiel 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 Vector Memory (VMEM), einen lokalen Scratchpad-Speicher, der ausschließlich vom TensorCore (TC) verwendet wird. Der Compiler verwaltet Vmem für verschiedene Arten von Zuweisungen:

  • Zuweisungen mit Anweisungsbereich:Temporärer Speicher in Vmem bei der Ausführung einer einzelnen HLO-Anweisung. Dazu gehören der Operandenspan-Puffer (z.B. für Double Buffering) und Register-Spills.
  • Zuweisungen mit Programmbereich:Zuweisungen, die über den Bereich einer einzelnen HLO-Anweisung hinausgehen. Dabei handelt es sich in der Regel um temporäre HLO-Variablen und Zwischenergebnisse, die Eingaben und/oder Ausgaben von HLO-Anweisungen sind.

Ein bereichsbezogener Vmem-Fehler (Out of Memory) zur Kompilierzeit tritt auf, wenn die Zuweisungen mit Anweisungsbereich das Zuweisungslimit für diese Anweisung überschreiten. Dieses Limit wird gesteuert

und

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 eine 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" ...
  • Bereichsbezogener Vmem-Fehler (Out of Memory) für benutzerdefinierten Kernel: Wenn der Fehler auf einen benutzerdefinierten Kernel verweist, fahren Sie mit Kernel neu abstimmen fort.
  • Vmem-Probleme ohne Kernel: Wenn der Vmem-Fehler (Out of Memory) aufgrund einer nicht benutzerdefinierten Kernel-Operation auftritt, ist dies wahrscheinlich ein interner Compilerfehler. Reichen Sie bitte einen Fehlerbericht zu XLA ein mit einem HLO-Dump.

Kernel neu abstimmen

Wenn der Fehler von einem benutzerdefinierten Kernel stammt, verwenden Sie die folgenden Methoden, um den Speicherbedarf des Kernels zu reduzieren:

  • Blockgrößen anpassen:Reduzieren Sie die Blockgrößen (Kachelgrößen) in der Kernelkonfiguration, um die bereichsbezogene Vmem-Nutzung zu verringern.
  • Bereichsbezogene Vmem-Limits pro Kernel festlegen: Fordern Sie mit dem Parameter vmem_limit_bytes explizit die erforderliche Menge an Arbeitsspeicher für diesen bestimmten Kernel an.
  • Speicherfarben ändern: Färben Sie die Ein- und Ausgaben des Kernels explizit in Vmem ein oder beschränken Sie sie darauf. Verwenden Sie dazu pallas.tpu.with_memory_space_constraint. Achten Sie darauf, nicht zu viele Ein- und Ausgaben in Vmem einzufärben, da dies zu einem allgemeinen Vmem-Fehler (Out of Memory) führen kann.
  • Vmem-Limit anpassen: Wenn die kernelspezifische Abstimmung schwierig ist oder das Problem viele Kernel betrifft, können Sie das globale Vmem-Limit mit dem Flag --xla_tpu_scoped_vmem_limit_kib anpassen.