Código do erro: 0100

Categoria:ambiente de execução: falha na alocação de buffer

Esse erro indica que o alocador de memória do tempo de execução do XLA:TPU não encontrou um bloco de memória adequado na HBM do acelerador para a alocação solicitada.

Exemplo de mensagem de erro:

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

Backends do XLA:TPU

Visão geral

Esse erro é gerado em

• falhas de alocação de buffer iniciada pelo usuário via jax.device_put ou
• falhas nas alocações de saída do programa programado pelo usuário.

Essas falhas geralmente ocorrem por dois motivos:

• Falta de memória (OOM): o usuário está tentando alocar um bloco de memória maior que a quantidade total de memória livre disponível na HBM da TPU.
• Fragmentação de memória:a alocação falha porque nenhum bloco livre contíguo no espaço de memória é grande o suficiente para atender ao tamanho solicitado. A quantidade total de memória livre é suficiente para a alocação, mas está espalhada pelo espaço de memória em blocos pequenos e não contíguos.

O ambiente de execução da TPU tem vários mecanismos para tentar novamente falhas de alocação, incluindo:

• Se houver desalocações na fila, o ambiente de execução vai tentar novamente as alocações com falha.
• Em OOMs causados por uma fragmentação, o ambiente de execução pode acionar automaticamente uma desfragmentação e uma nova tentativa.
• O tempo de execução da TPU prioriza alocações de buffer em vez de manter programas carregados. Se uma alocação de buffer falhar devido à HBM insuficiente, o sistema vai desalojar programas de TPU carregados até que haja memória suficiente para o buffer.

Portanto, um erro encontrado após as mitigações acima geralmente exige uma ação do usuário.

Depuração

• Reduza o consumo de memória do modelo:
  • Diminuir o tamanho do lote: reduzir o tamanho do lote diminui diretamente o uso de memória.
  • Fragmentação de parâmetros: para modelos muito grandes, use técnicas como paralelismo de modelo ou fragmentação para distribuir parâmetros na HBM de vários núcleos ou hosts de TPU.
  • Diminua o comprimento da sequência/contexto: para modelos que operam em sequências (como modelos de linguagem), reduzir o comprimento da sequência de entrada pode diminuir significativamente a área de ocupação da memória.
  • Doação de buffer: use recursos do framework (como jax.jit(..., donate_argnums=...)) para sinalizar ao XLA que determinados buffers de entrada podem ser substituídos e reutilizados para saídas.
  • Otimize a estratégia de checkpoint: em vez de salvar todo o estado do modelo de uma vez, salve apenas os pesos do modelo ou use uma estratégia de checkpoint fragmentada.
• Layout e padding da memória de endereço:
  • A memória da TPU é alocada em partes, e o preenchimento pode aumentar o tamanho real dos tensores.
• Verifique se não há vazamentos de memória:
  • Verifique se as referências a objetos jax.Array não estão sendo mantidas por mais tempo do que o pretendido. Manter objetos jax.Array pode impedir a desalocação automática, mesmo depois que a compilação do programa for concluída.

Ferramentas

Ative a flag tpu_log_allocations_on_oom para que o alocador gere um relatório detalhado de todas as alocações atuais quando ocorrer um OOM, o que pode ser muito útil para depuração.