Codice di errore: E1001

Categoria: Compile Time: Scoped Vmem OOM

Questo errore indica che il programma richiede più memoria vettoriale con ambito (Vmem) di quella allocata.

Messaggi di errore di esempio:

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(...) ...

Backend XLA: TPU

Panoramica

Le TPU hanno una memoria vettoriale (VMEM) che è una memoria scratchpad locale utilizzata esclusivamente da TensorCore (TC). Il compilatore gestisce la Vmem per diversi tipi di allocazioni:

  • Allocazioni con ambito di istruzione: spazio di archiviazione temporaneo in Vmem durante l'esecuzione di una singola istruzione HLO. Sono inclusi il buffer di intervallo degli operandi (ad es. per il doppio buffering) e le fuoriuscite di registri.
  • Allocazioni con ambito di programma: allocazioni che vanno oltre l'ambito di una singola istruzione HLO. In genere si tratta di risultati temporanei e intermedi di HLO che sono input e/o output delle istruzioni HLO.

Si verifica un errore di memoria insufficiente (OOM) di Vmem con ambito di compilazione quando le allocazioni con ambito di istruzione superano il limite di allocazione per quell'istruzione. Questo limite è controllato

e

In genere questi errori sono causati da un bug del compilatore interno o da un kernel personalizzato che supera il limite di allocazione.

Debug

Analizza attentamente il messaggio di errore per identificare se l'errore deriva da un kernel personalizzato o da un HLO standard. Un errore dovuto a un kernel personalizzato deve avere la seguente firma:

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" ...
  • OOM di Vmem con ambito di kernel personalizzato: se l'errore rimanda a un kernel personalizzato, vai a Riassegnare il kernel.
  • Problemi di Vmem non kernel: se l'errore OOM di Vmem si verifica a causa di un'operazione non kernel personalizzata, è probabile che si tratti di un bug del compilatore interno. Invia una segnalazione di bug su XLA con un dump HLO.

Riassegnare il kernel

Se l'errore ha origine da un kernel personalizzato, utilizza le seguenti tecniche per ridurre il requisito di memoria del kernel:

  • Regola le dimensioni dei blocchi: riduci le dimensioni dei blocchi (dimensioni dei riquadri) nella configurazione del kernel per ridurre l'utilizzo di Vmem con ambito.
  • Imposta i limiti di Vmem con ambito per kernel: richiedi esplicitamente la quantità di memoria richiesta per quel kernel specifico utilizzando ilparametro vmem_limit_bytes
  • Modifica la colorazione della memoria: colora/vincola esplicitamente gli input/output del kernel a Vmem utilizzandopallas.tpu.with_memory_space_constraint. Fai attenzione a non colorare troppi input/output in Vmem, in quanto ciò potrebbe causare un errore OOM di Vmem complessivo.
  • Regola il limite di Vmem: se la riassegnazione specifica del kernel è difficile o il problema riguarda molti kernel, puoi regolare il limite di Vmem globale utilizzando il flag --xla_tpu_scoped_vmem_limit_kib.