Codice di errore: 1000

Categoria: Tempo di compilazione: HBM OOM

Questo errore indica che il programma richiede una memoria ad alta larghezza di banda (HBM) superiore a quella fisicamente disponibile sul dispositivo TPU.

Messaggi di errore di esempio:

RESOURCE_EXHAUSTED: XLA:TPU compile permanent error. Ran out of memory in memory space hbm. Used 49.34G of 32.00G hbm. Exceeded hbm capacity by 17.34G.

RESOURCE_EXHAUSTED: TPU TensorCore Hbm usage: 34.82G, SparseCore Hbm usage 174.10G, exceeding available bytes: 95.74G

Backend XLA: TPU

Panoramica

XLA esegue controlli per garantire che le dimensioni aggregate di tutte le allocazioni statiche necessarie rientrino nella HBM del dispositivo.

Il compilatore gestisce la capacità HBM fissa della TPU per diversi tipi di allocazioni:

• Input e output del programma: batch di addestramento, stati dell'ottimizzatore e così via.
• TensorCore + SparseCore Temporaries: memoria dinamica richiesta per calcoli intermedi (ad es. attivazioni, gradienti e così via).
• Binario compilato:il codice macchina per TensorCore (TC) e SparseCore (SC).
• Overhead di sistema:spazio riservato per XLA Runtime (ad es. buffer di inserimento nelle generazioni di TPU precedenti).
• Costanti:i valori costanti incorporati nel HLO IR vengono allocati su HBM.
• Elementi interni del compilatore: allocazioni a livello di programma e per HLO (ad es. informazioni di routing per i nodi nella mesh)

Questo errore si verifica quando il compilatore XLA non riesce a inserire tutte le allocazioni precedenti nella HBM del dispositivo.

Debug

Analizza attentamente il messaggio di errore e i log per determinare quale categoria di errore HBM OOM descrive meglio il tuo errore:

• TensorCore (TC) + SparseCore (SC) HBM Usage Exceeds Limit: Se l'errore suddivide esplicitamente l'utilizzo, ad esempio "Utilizzo Hbm TC: X, utilizzo Hbm SC Y". → Vai a Sezione 1. Bilancia l'utilizzo di HBM di TC e SC.
• Allocazioni inaspettatamente grandi: Se l'errore indica "Memoria esaurita nello spazio di memoria HBM", controlla i log per un'enumerazione delle allocazioni più grandi su HBM. Se sono presenti uno o più tensori inaspettatamente grandi (ad es. > 50% del limite HBM) → Vai alla Sezione 2. Allocazioni inaspettatamente grandi.
• Le allocazioni aggregate superano il limite HBM: Se l'errore indica "Memoria esaurita nello spazio di memoria HBM", ma nei log non sono presenti tensori di dimensioni insolitamente grandi → Vai alla Sezione 3. Aggregate Allocations Exceed HBM Limit.

Sezione 1. Bilanciare l'utilizzo di HBM di TC e SC

Se l'errore suddivide esplicitamente l'utilizzo, ad esempio: "TC Hbm usage: X, SC Hbm usage Y" confronta i due valori per identificare il collo di bottiglia

• Utilizzo elevato di SparseCore:
  • Ottimizza l'utilizzo dello stack HBM: il consumo di memoria dello stack HBM viene scalato con feature_width, max_unique_nz_per_row e logical_replica_count. Puoi ridurre l'utilizzo massimo dello stack modificando il flag --xla_sc_num_serialized_tables_to_optimize_hbm, che serializza l'elaborazione delle tabelle. Ciò comporta una riduzione del parallelismo.
  • Controlla l'overhead di padding:SparseCore allinea le tabelle di incorporamento a 32 byte (8 float). Tabelle con larghezze delle funzionalità ridotte (ad es. < 8 float) comportano un overhead di padding significativo, sprecando HBM.
  • Ridurre l'utilizzo dell'heap:valori elevati per maximum_parallel_iterations aumentano la quantità di dati di input precaricati nell'heap HBM. La riduzione di questo valore può liberare una quantità significativa di memoria.
  • Verifica lo sharding:assicurati che le tabelle di incorporamento siano suddivise correttamente in tutti i chip. Consulta la sezione Come si traducono i limiti nelle tabelle.
  • Per altre idee, consulta SC: Performance and memory bottlenecks.
• Utilizzo elevato di Tensor Core:
  • Vai alla sezione 2.
• Bilanciata
  • Se nessuno dei due è eccessivo individualmente, ma la somma è troppo alta, hai raggiunto la capacità del chip. Devi provare a ridurre l'utilizzo di entrambi i componenti. Segui i consigli riportati in tutte e tre le sezioni.

Sezione 2. Allocazioni inaspettatamente grandi

Se nei log sono presenti una o più allocazioni inaspettatamente grandi (> 50% del limite HBM), quasi mai si tratta di un problema di capacità hardware. In genere si tratta di un errore di configurazione. Controlla l'etichetta XLA (se presente) delle allocazioni di grandi dimensioni per suggerimenti sul codice sorgente JAX.

• Rimuovi artefatti di debug:
  • L'utilizzo di jax.debug.print() in esecuzioni su larga scala può forzare il compilatore a materializzare il tensore completo in HBM per trasferirlo alla CPU, interrompendo la fusione e aumentando l'utilizzo di memoria di picco. Rimuovi eventuali jax.debug.print() rimanenti.
• Correggere forme mesh o partizionamento inefficienti:
  • Forme mesh errate o annotazioni di sharding mancanti possono causare l'impostazione predefinita del compilatore su Replication, forzando il compilatore a tentare di adattare tensori molto grandi su un singolo chip
  • Controlla le forme delle allocazioni di grandi dimensioni e verifica che lo sharding sia specificato e propagato correttamente da XLA.

Sezione 3. Le allocazioni aggregate superano il limite HBM

Se il programma esaurisce la capacità perché la somma aggregata delle allocazioni supera il limite HBM, spesso è utile visualizzare il profilo di memoria per identificare i buffer specifici che contribuiscono al picco di utilizzo. Consulta Eseguire il debug degli errori OOM con XProf per una guida passo passo per identificare i principali contributori di memoria.

Una volta identificati alcuni dei principali collaboratori, segui questi passaggi per ottimizzare l'utilizzo della memoria.

A. Controlla il riempimento e l'allineamento del tensore

Le forme dei tensori inefficienti sono una causa comune e silenziosa di errori di esaurimento della memoria sulle TPU. Per ottenere le massime prestazioni sulle TPU, XLA esegue il padding delle dimensioni dei tensori, in genere a multipli di 128 per la dimensione meno significativa e 8 per la seconda meno significativa. Questo padding influisce sia sugli array di input sia sui tensori intermedi (HLO temporanei), gonfiando potenzialmente l'utilizzo della memoria in modo significativo, soprattutto con dimensioni ridotte. Consulta Layout degli array.

• Controlla le forme dei buffer di grandi dimensioni: (su TPU v5 con layout predefiniti)
  • Se passi il mouse sopra un buffer in Xprof Memory Viewer viene visualizzata la scheda dei dettagli del buffer, che contiene i dettagli del buffer, incluse le informazioni sul padding.
  • Esempio: una forma di (129, 1024) potrebbe essere riempita fino a (256, 1024), con conseguente spreco di memoria di quasi il 50%.
  • Correzione: una forma di (128, 1024) non richiede padding e comporta uno spreco di memoria dello 0%.
• Allinea dimensioni:assicurati che tutte le dimensioni dei tensori di grandi dimensioni (dimensione batch, dimensione incorporamento, dimensione nascosta) siano multipli di 128.

B. Modifica della configurazione

Spesso puoi risolvere gli errori di esaurimento della memoria con queste modifiche alla configurazione:

• Ridurre le dimensioni del batch:la memoria necessaria per le attivazioni intermedie e i gradienti è direttamente proporzionale alle dimensioni del batch. Ridurre la dimensione del batch spesso può contribuire a ridurre l'utilizzo della memoria.
• Donate Input Buffers:quando utilizzi jax.jit, specifica donate_argnums per i parametri del modello. In questo modo, XLA può sovrascrivere la memoria di input con l'output.
• Abilita la precisione mista (bfloat16): utilizza bfloat16 o la quantizzazione (int8 e così via) per i tensori più grandi del programma se l'architettura del modello e i requisiti di qualità lo consentono.

C. Ottimizza l'architettura e lo sharding

Se le modifiche alla configurazione non sono sufficienti, la topologia del modello potrebbe essere troppo grande per la configurazione hardware attuale.

• Utilizza generazioni di TPU più recenti:le TPU più recenti in genere offrono più HBM per chip; passa a generazioni di TPU più recenti, se disponibili.
• Esegui su una topologia di chip più grande:se i pesi del modello sono troppo grandi per la topologia esistente, puoi provare a suddividerli in più chip.
• Implementa tecniche di partizionamento avanzate:
  • Esplora approcci più avanzati per il parallelismo di dati, tensori o pipeline.
  • Specifica suggerimenti per lo sharding per valori e output intermedi.
• Utilizza il trasferimento host JAX:trasferisci i tensori di grandi dimensioni alla memoria della CPU host. Ad esempio, trasferimento dell'attivazione e trasferimento dello stato dell'ottimizzatore.

D. Ottimizza i flag XLA che influiscono sulla memoria delle chiavi:

I flag di memoria chiave possono essere ottimizzati per bilanciare le prestazioni con un utilizzo inferiore della memoria. Tuttavia, questi devono essere utilizzati come ultima risorsa perché possono influire negativamente sulle prestazioni.

E. Ottimizza il passaggio di rematerializzazione XLA / il checkpoint manuale

Se il modello è quasi pronto per essere inserito in memoria, puoi forzare il passaggio XLA::Rematerialization per dare la priorità al risparmio di memoria, potenzialmente a costo di compilazioni più lente:

In alternativa, utilizza il decoratore jax.checkpoint con jax.grad per controllare manualmente quali intermedi vengono salvati nel passaggio in avanti rispetto a quelli ricalcolati nel passaggio all'indietro, scambiando cicli di calcolo con HBM.

F. Utilizzare strumenti di profilazione avanzati

Debug degli errori OOM con XProf fornisce un tutorial sull'utilizzo del visualizzatore di memoria XProf per visualizzare la visualizzazione dell'utilizzo di HBM da parte del compilatore.

Questo strumento ti consente di visualizzare l'allocazione di memoria e la durata dei buffer di picco, il che è fondamentale per capire esattamente cosa consuma HBM nel punto di picco di utilizzo. Per la configurazione generale della profilazione, vedi Introduzione a Xprof e Profilazione di TensorBoard.