Codice di errore: 0101

Categoria: Runtime: Program Allocation Failure

Questo errore indica che il runtime XLA su un dispositivo TPU non è riuscito a caricare un eseguibile del programma XLA compilato nella HBM della TPU.

Esempio di messaggio di errore:

XlaRuntimeError: RESOURCE_EXHAUSTED: Error loading program 'jit_embedding_pipeline_step_fn': Attempting to reserve 29.49G at the bottom of memory. That was not possible. There are 147.64M free, 0B reserved, and 147.64M reservable. Scope: unknown..: while running replica 0 and partition 34 of a replicated computation (other replicas may have failed as well).

Backend XLA: TPU

Panoramica

Questo errore è in genere causato da uno dei seguenti motivi:

• Program Size Exceeds Available HBM: Il programma XLA compilato, incluse le istruzioni, i dati statici e le costanti incorporate, è più grande della quantità totale di HBM libera attualmente disponibile sui core TPU specifici in cui viene caricato il programma.
• Frammentazione HBM: sebbene l'HBM libero totale sul dispositivo possa essere sufficiente in totale, non è disponibile in un singolo blocco contiguo abbastanza grande da contenere l'intero programma.

È importante capire come il runtime TPU assegna la priorità alla memoria. Le allocazioni del buffer hanno la precedenza sui programmi caricati. Se l'allocazione di un buffer non va a buon fine, il runtime espellerà i programmi già caricati dalla HBM per liberare spazio. Ciò può portare a una situazione in cui un programma che prima veniva caricato correttamente ora non viene caricato e viene visualizzato un errore di memoria insufficiente, perché la memoria HBM è ora occupata da più buffer di dati.

Debug

• Riduzione dello spazio occupato dalla memoria buffer: liberare la memoria utilizzata dai buffer di dati lascerà più spazio per il programma stesso:
  • Diminuisci le dimensioni del batch: questo è uno dei modi più efficaci per ridurre la quantità di memoria utilizzata per le attivazioni.
  • Partizionamento dei parametri: per modelli molto grandi, utilizza il parallelismo del modello o tecniche di partizionamento (come FSDP o Megascale) per distribuire i parametri e il calcolo del modello su più host o core TPU.
  • Riduzione della lunghezza della sequenza/del contesto: per i modelli che elaborano dati sequenziali (ad es. NLP), la riduzione della lunghezza della sequenza può diminuire significativamente l'utilizzo della memoria.
  • Donazione del buffer: utilizza le funzionalità del framework (ad es. jax.jit(..., donate_argnums=...)) per consentire a XLA di riutilizzare la memoria dei buffer di input per memorizzare l'output, riducendo l'utilizzo di memoria di picco.
• Riduci i requisiti di memoria del programma per i file temporanei:
  • Riduci l'utilizzo della memoria dei programmi per i file temporanei utilizzando il flag tpu_shared_memory_percent. Tieni presente che questa operazione potrebbe influire negativamente sul rendimento.
• Ottimizza la strategia di esecuzione/riduci il carico di pubblicazione:
  • Gestisci il caricamento del programma: se stai compilando JIT più funzioni, tieni presente che ogni funzione può comportare il caricamento di un programma. Prova a strutturare il workload in modo da ridurre al minimo il numero di programmi caricati contemporaneamente.
• 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 genera un report dettagliato di tutte le allocazioni correnti quando si verifica un errore di memoria insufficiente, il che può essere molto utile per il debug.
• Profila il tuo programma: utilizza il profiler di memoria JAX o il profiler TensorFlow per visualizzare in dettaglio l'utilizzo della memoria del tuo programma nel tempo. Ciò può aiutare a identificare picchi imprevisti nel consumo di memoria.