Código do erro: E0100

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

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

Mensagem de erro de exemplo :

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

Back-ends do XLA:TPU

Visão geral

Esse erro é gerado em:

  • Falhas de alocação de buffer iniciadas pelo usuário via jax.device_put

ou

  • Falhas de alocações de saída do programa agendado pelo usuário.

Essas falhas geralmente são causadas por alguns 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 enfileiradas, o ambiente de execução vai tentar novamente as alocações com falha.
  • Em OOMs causadas por uma fragmentação, o ambiente de execução pode acionar automaticamente uma desfragmentação e uma nova tentativa.
  • O ambiente de execução da TPU prioriza as alocações de buffer em vez de manter os programas carregados. Se uma alocação de buffer falhar devido à HBM insuficiente, o sistema vai remover os programas de TPU carregados até que haja memória suficiente para o buffer.

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

Depuração

  • Reduza o consumo de memória do modelo
    • Diminua 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 modelos ou fragmentação para distribuir parâmetros na HBM de vários núcleos ou hosts de TPU.
    • Encurte o tamanho 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 o consumo de memória.
    • Doação de buffer: use recursos de framework (como jax.jit(..., donate_argnums=...)) para sinalizar ao XLA que determinados buffers de entrada podem ser substituídos e reutilizados para saídas. Leia Doação de buffer para mais detalhes.
    • Otimize a estratégia de checkpoint:em vez de salvar o estado do modelo inteiro de uma só vez, salve apenas os pesos do modelo ou use uma estratégia de checkpoint fragmentada.
  • Aborde o layout e o preenchimento da memória
    • A memória da TPU é alocada em blocos, e o preenchimento pode aumentar o tamanho real dos tensores.
  • 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 após a conclusão da compilação do programa.

Consulte também Código de erro: E1000 para outras estratégias que podem ser usadas para reduzir a quantidade de memória que cada programa usa.

Ferramentas

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