Codice di errore: E1000

Categoria: Tempo di compilazione: HBM OOM

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

Esempi di messaggi di errore:

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

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.

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 ecc.
• TPU temporaries:memoria dinamica richiesta per i 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 infeed nelle generazioni di TPU precedenti).
• Costanti:i valori costanti incorporati nell'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 di esaurimento della memoria HBM di seguito descrive meglio il tuo errore:

• "Utilizzo HBM TC: X, utilizzo HBM SC Y": se l'errore suddivide esplicitamente l'utilizzo di HBM, l'utilizzo aggregato di TensorCore (TC) + SparseCore (SC) supera il limite HBM. → Vai allo Scenario 1. Bilancia l'utilizzo di HBM di TC e SC.
• "Ran out of memory in memory space HBM": controlla i log per un'enumerazione delle allocazioni più grandi su HBM.
  • Nel caso in cui siano presenti uno o più tensori inaspettatamente grandi (ad es. > 50% del limite HBM) → Vai allo Scenario 2. Memoria esaurita a causa di allocazioni inaspettatamente grandi.
  • Se nei log non sono presenti tensori di dimensioni inaspettatamente grandi → Vai allo Scenario 3. Memoria insufficiente a causa delle allocazioni aggregate.

Scenario 1 Bilanciare l'utilizzo di HBM di TC e SC

Se l'errore suddivide esplicitamente l'utilizzo, ad esempio "Utilizzo HBM TC: X, utilizzo HBM SC Y", significa che l'utilizzo aggregato di TensorCore (TC) + SparseCore (SC) supera il limite HBM. 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 aumenta in base a 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). Le 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 ulteriori idee, consulta SC: Performance and memory bottlenecks.
• Utilizzo elevato di Tensor Core
  • Vai allo Scenario 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.

Scenario 2. Memoria esaurita a causa di allocazioni inaspettatamente grandi

Se visualizzi il messaggio di errore "Memoria esaurita nello spazio di memoria HBM" e 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.

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

Scenario 3. Memoria insufficiente a causa delle allocazioni aggregate

Se visualizzi il messaggio di errore "Ran out of memory in memory space HBM" e nei log non sono presenti tensori di dimensioni inaspettatamente grandi, il programma esaurisce la capacità perché la somma aggregata delle allocazioni supera il limite HBM. In questo caso, 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 il footprint della memoria.

Scenario 3.A Regola la 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 le dimensioni del batch può spesso contribuire a ridurre la memoria utilizzata.
• Buffer di input di donazione: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.
• Attiva la precisione mista (bfloat16): utilizza bfloat16 o la quantizzazione (int8 ecc.) per i tensori più grandi del programma se l'architettura del modello e i requisiti di qualità lo consentono. Tieni presente che questa modifica può influire sul comportamento del modello e deve essere valutata attentamente.

Scenario 3.B: 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 intermedi e output.
• Utilizza l'offload dell'host JAX:le tecniche di offload dell'host consentono all'utente di eseguire l'offload di tensori di grandi dimensioni nella memoria della CPU host (ad es. offload dell'attivazione e offload dello stato dell'ottimizzatore).

Scenario 3.C Controlla il padding e l'allineamento dei tensori

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 la memoria utilizzata 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 informazioni tra cui quelle 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 del batch, dimensione incorporamento, dimensione nascosta) siano multipli di 128. Tieni presente che questa modifica può influire sul comportamento del modello e deve essere valutata attentamente.

Scenario 3.D: ottimizza i flag XLA che influiscono sulla memoria delle chiavi

I flag di memoria chiave possono essere ottimizzati per bilanciare le prestazioni con una minore memoria utilizzata. Tuttavia, questa strategia deve essere utilizzata come ultima risorsa, in quanto può influire negativamente sul rendimento.

Scenario 3.E Tune XLA rematerialization pass/manual checkpointing

Se il modello è quasi pronto per essere caricato in memoria, puoi utilizzare 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.

In alternativa, puoi forzare il passaggio XLA::Rematerialization per dare la priorità al risparmio di memoria, potenzialmente a scapito di compilazioni più lente:

Tieni presente che le modifiche ai flag XLA devono essere utilizzate come ultima risorsa, in quanto possono influire negativamente sul rendimento.

Scenario 3.F Utilizza 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 della memoria di picco e la durata dei buffer, 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.