Código do erro: 0101

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

Esse erro indica que o ambiente de execução do XLA em um dispositivo de TPU não conseguiu carregar um executável de programa XLA compilado na HBM da TPU.

Exemplo de mensagem de erro:

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).

Backends do XLA:TPU

Visão geral

Esse erro geralmente é causado por um dos seguintes motivos:

• O tamanho do programa excede a HBM disponível: o programa XLA compilado, incluindo instruções, dados estáticos e constantes incorporadas, é maior do que a quantidade total de HBM livre disponível nos núcleos específicos da TPU em que o programa está sendo carregado.
• Fragmentação de HBM: embora o total de HBM livre no dispositivo possa ser suficiente no agregado, ele não está disponível em um único bloco contíguo grande o suficiente para caber em todo o programa.

É importante entender como o ambiente de execução da TPU prioriza a memória. As alocações de buffer têm privilégio sobre programas carregados. Se uma alocação de buffer falhar, o tempo de execução vai remover programas já carregados da HBM para liberar espaço. Isso pode levar a uma situação em que um programa que foi carregado com sucesso antes agora falha com um erro de falta de memória, porque a HBM está ocupada com mais buffers de dados.

Depuração

• Reduza o espaço ocupado pela memória do buffer: liberar a memória usada pelos buffers de dados deixa mais espaço para o próprio programa:
  • Diminuir o tamanho do lote: essa é uma das maneiras mais eficazes de reduzir a quantidade de memória usada para ativações.
  • Fragmentação de parâmetros: para modelos muito grandes, use paralelismo de modelo ou técnicas de fragmentação (como FSDP ou Megascale) para distribuir os parâmetros e a computação do modelo em vários núcleos ou hosts de TPU.
  • Reduzir o comprimento da sequência/contexto: para modelos que processam dados sequenciais (por exemplo, NLP), reduzir o comprimento da sequência pode diminuir significativamente o uso da memória.
  • Doação de buffer: use recursos do framework (por exemplo, jax.jit(..., donate_argnums=...)) para permitir que a XLA reutilize a memória dos buffers de entrada para armazenar a saída, reduzindo o uso máximo de memória.
• Reduza os requisitos de memória do programa para temporários:
  • Reduza o uso de memória dos programas para temporários usando a flag tpu_shared_memory_percent. Isso pode afetar negativamente a performance.
• Otimizar a estratégia de execução/reduzir a carga de veiculação:
  • Gerenciar o carregamento de programas: se você estiver compilando várias funções JIT, saiba que cada função pode resultar no carregamento de um programa. Tente estruturar sua carga de trabalho para minimizar o número de programas carregados simultaneamente.
• 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 despeje um relatório detalhado de todas as alocações atuais quando ocorrer um OOM, o que pode ser muito útil para depuração.
• Crie um perfil do seu programa: use o criador de perfil de memória do JAX ou do TensorFlow para ter uma visão detalhada do uso de memória do seu programa ao longo do tempo. Isso pode ajudar a identificar picos inesperados no consumo de memória.