Codice di errore: 1001

Categoria: Tempo di compilazione: errore di memoria virtuale esaurita con ambito

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 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 gli overflow dei registri.
  • Allocazioni con ambito del programma: allocazioni che vanno oltre l'ambito di una singola istruzione HLO. Si tratta in genere di risultati temporanei e intermedi HLO che sono input e/o output delle istruzioni HLO.

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

Questi errori sono in genere 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 dovrebbe 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" ...
  • Custom Kernel Scoped Vmem OOM: se l'errore indica un kernel personalizzato → vai a Ottimizzare il kernel.
  • Problemi di Vmem non del kernel: se l'errore Vmem OOM si verifica a causa di un'operazione non del kernel personalizzato, è probabile che si tratti di un bug interno del compilatore. Segnala un bug su XLA con un dump HLO.

Ottimizzare 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 limiti Vmem con ambito per kernel:richiedi esplicitamente la quantità di memoria richiesta per quel kernel specifico utilizzando il parametro vmem_limit_bytes.
  • Modifica della colorazione della memoria:colora/vincola esplicitamente gli input/output del kernel alla VMEM utilizzando pallas.tpu.with_memory_space_constraint. Fai attenzione a non colorare troppi input output in Vmem, in quanto ciò potrebbe causare un errore di memoria insufficiente complessivo.
  • Se la riorganizzazione specifica del kernel è difficile o il problema interessa molti kernel, puoi modificare il limite Vmem globale utilizzando il flag --xla_tpu_scoped_vmem_limit_kib.