Codice di errore: 0100

Categoria: Runtime: Buffer Allocation Failure

Questo errore indica che l'allocatore di memoria del runtime XLA:TPU non è riuscito a trovare un blocco di memoria adatto nella HBM dell'acceleratore per l'allocazione richiesta.

Esempio di messaggio di errore:

ValueError: RESOURCE_EXHAUSTED: Error allocating device buffer: Attempting to allocate 8.00M. That was not possible. There are 6.43M free.; (0x0x1_HBM0)

Backend XLA: TPU

Panoramica

Questo errore viene generato su

• errori di allocazione del buffer avviata dall'utente tramite jax.device_put o
• errori di allocazione dell'output del programma pianificato dall'utente.

Questi errori sono in genere causati da un paio di motivi:

• Esaurimento della memoria (OOM): l'utente sta tentando di allocare un blocco di memoria più grande della quantità totale di memoria libera disponibile nella HBM della TPU.
• Frammentazione della memoria:l'allocazione non riesce perché nessun singolo blocco libero contiguo nello spazio di memoria è abbastanza grande da soddisfare le dimensioni richieste. La quantità totale di memoria libera è sufficiente per l'allocazione, ma è distribuita nello spazio di memoria in blocchi piccoli e non contigui.

Il runtime TPU dispone di diversi meccanismi per riprovare ad allocare le risorse in caso di errori, tra cui:

• Se sono presenti deallocazioni in coda, il runtime riprova le allocazioni non riuscite,
• In caso di errori di memoria insufficiente causati da una frammentazione, il runtime può attivare automaticamente una deframmentazione e un nuovo tentativo.
• Il runtime della TPU dà la priorità alle allocazioni del buffer rispetto al mantenimento dei programmi caricati. Se l'allocazione di un buffer non riesce a causa di HBM insufficiente, il sistema espelle i programmi TPU caricati finché non è disponibile memoria sufficiente per il buffer.

Pertanto, un errore riscontrato dopo le mitigazioni sopra indicate in genere richiede l'intervento dell'utente.

Debug

• Riduci l'utilizzo di memoria del modello:
  • Diminuisci dimensione batch: la riduzione della dimensione del batch riduce direttamente l'utilizzo della memoria.
  • Suddivisione dei parametri: per modelli molto grandi, utilizza tecniche come il parallelismo del modello o la suddivisione per distribuire i parametri nella HBM di più core TPU o host.
  • Riduzione della lunghezza della sequenza/del contesto: per i modelli che operano su sequenze (come i modelli linguistici), la riduzione della lunghezza della sequenza di input può ridurre significativamente l'utilizzo di memoria.
  • Buffer Donation: utilizza le funzionalità del framework (ad esempio jax.jit(..., donate_argnums=...)) per segnalare a XLA che determinati buffer di input possono essere sovrascritti e riutilizzati per gli output.
  • Ottimizza la strategia di checkpoint: anziché salvare l'intero stato del modello contemporaneamente, valuta la possibilità di salvare solo i pesi del modello o di utilizzare una strategia di checkpoint con partizionamento.
• Layout e padding della memoria degli indirizzi:
  • La memoria TPU viene allocata in blocchi e il padding può aumentare le dimensioni effettive dei tensori.
• Assicurati che non ci siano perdite di memoria:
  • Assicurati che i riferimenti agli oggetti jax.Array non vengano conservati più a lungo del previsto. Il mantenimento degli oggetti jax.Array potrebbe impedire la deallocazione automatica anche dopo il completamento della compilazione del programma.

Strumenti

Attiva il flag tpu_log_allocations_on_oom per cui l'allocatore genererà un report dettagliato di tutte le allocazioni correnti quando si verifica un errore di memoria insufficiente, il che può essere fondamentale per il debug.