Código do erro: 1000

Categoria:tempo de compilação: HBM OOM

Esse erro indica que o programa requer mais memória de alta largura de banda (HBM) do que está disponível fisicamente no dispositivo TPU.

Exemplos de mensagens de erro:

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.

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

Back-ends do XLA:TPU

Visão geral

O XLA realiza verificações para garantir que o tamanho agregado de todas as alocações estáticas necessárias caiba na HBM do dispositivo.

O compilador gerencia a capacidade fixa de HBM da TPU para vários tipos de alocações:

• Entradas e saídas do programa:lotes de treinamento, estados do otimizador etc.
• Temporários do TensorCore + SparseCore:memória dinâmica necessária para cálculos intermediários (por exemplo, ativações, gradientes etc.).
• Binário compilado:o código de máquina para TensorCore (TC) e SparseCore (SC).
• Sobrecarga do sistema:espaço reservado para o tempo de execução do XLA (por exemplo, buffers de inclusão em gerações mais antigas de TPU).
• Constantes:valores constantes incorporados na representação intermediária de alto nível são alocados na HBM.
• Internos do compilador:alocações no nível do programa e por HLO (por exemplo, informações de roteamento para nós na malha)

Esse erro ocorre quando o compilador XLA não consegue ajustar todas as alocações acima na HBM do dispositivo.

Depuração

Analise cuidadosamente a mensagem de erro e os registros para determinar qual categoria de HBM OOM abaixo descreve melhor seu erro:

• O uso de HBM do TensorCore (TC) + SparseCore (SC) excedeu o limite: Se o erro detalhar explicitamente o uso, por exemplo, "Uso de Hbm do TC: X, uso de Hbm do SC Y". → Ir para Seção 1. Balancear o uso de HBM de TC e SC.
• Alocações inesperadamente grandes: se o erro for "Ran out of memory in memory space HBM", verifique os registros para uma enumeração das maiores alocações em HBM. Se um ou mais tensores inesperadamente grandes (por exemplo, > 50% do limite de HBM) estiverem presentes, vá para a Seção 2. Alocações inesperadamente grandes.
• As alocações agregadas excedem o limite de HBM: se o erro for "Ficou sem memória no espaço de memória HBM", mas não houver tensores inesperadamente grandes nos registros, vá para Seção 3. As alocações agregadas excedem o limite de HBM.

Seção 1. Equilibrar o uso de HBM de TC e SC

Se o erro detalhar explicitamente o uso, por exemplo, "Uso de Hbm de TC: X, uso de Hbm de SC Y" compare os dois valores para identificar o gargalo

• Uso elevado do SparseCore:
  • Otimizar o uso da pilha de HBM:o consumo de memória da pilha de HBM é dimensionado com feature_width, max_unique_nz_per_row e logical_replica_count. Para reduzir o uso máximo da pilha, ajuste a flag --xla_sc_num_serialized_tables_to_optimize_hbm, que serializa o processamento de tabelas. Isso reduz o paralelismo.
  • Verificar a sobrecarga de padding:o SparseCore alinha tabelas de incorporação a 32B (8 floats). Tabelas com larguras de recursos pequenas (por exemplo, < 8 floats) incur significant padding overhead, wasting HBM.
  • Reduza o uso de heap:valores altos para maximum_parallel_iterations aumentam a quantidade de dados de entrada pré-buscados no heap HBM. Diminuir esse valor pode liberar uma quantidade significativa de memória.
  • Verificar o sharding:garanta que as tabelas de incorporação sejam mod-sharded corretamente em todos os chips. Consulte Como os limites se traduzem em tabelas.
  • Confira SC: gargalos de desempenho e memória para mais ideias.
• Uso alto do TensorCore:
  • Acesse a Seção 2.
• Balanceado
  • Se nenhum deles for individualmente excessivo, mas a soma for muito alta, você estará na capacidade do chip. Tente reduzir o uso dos dois componentes. Siga as recomendações em todas as três seções.

Seção 2. Alocações inesperadamente grandes

Se uma ou mais alocações inesperadamente grandes estiverem presentes nos registros (> 50% do limite de HBM), quase nunca será um problema de capacidade de hardware. Normalmente, é um erro de configuração. Verifique o rótulo XLA (se presente) das alocações grandes para dicas sobre o código-fonte JAX.

• Remover artefatos de depuração:
  • Usar jax.debug.print() em execuções em grande escala pode forçar o compilador a materializar o tensor completo em HBM para transferi-lo para a CPU, interrompendo a fusão e aumentando o uso máximo de memória. Remova os jax.debug.print()s restantes.
• Corrigir formas de malha ou fragmentação ineficientes:
  • Formas de malha incorretas ou anotações de fragmentação ausentes podem fazer com que o compilador use Replicação por padrão, forçando-o a tentar ajustar tensores muito grandes em um único chip.
  • Confira os formatos das alocações grandes e verifique se o sharding foi especificado e propagado corretamente pelo XLA.

Seção 3. As alocações agregadas excedem o limite de HBM

Se o programa ficar sem capacidade devido ao total agregado de alocações exceder o limite de HBM, é útil visualizar o perfil de memória para identificar os buffers específicos que contribuem para o pico de uso. Consulte Depurar erros de OOM com o XProf para um guia detalhado sobre como identificar os principais contribuidores de memória.

Depois de identificar alguns dos principais contribuidores, siga estas etapas para otimizar a ocupação da memória.

A. Verificar o preenchimento e o alinhamento de tensores

Formas de tensores ineficientes são uma causa comum e silenciosa de erros de falta de memória nas TPUs. Para alcançar o desempenho máximo nas TPUs, o XLA faz o padding das dimensões do tensor, geralmente para múltiplos de 128 na dimensão menor e 8 na segunda menor. Esse padding afeta matrizes de entrada e tensores intermediários (temporários de HLO), aumentando potencialmente o uso de memória de forma significativa, especialmente com tamanhos de dimensão pequenos. Consulte Layouts de matriz.

• Auditar formatos de buffers grandes: (na TPU v5 com layouts padrão)
  • Passar o cursor sobre um buffer no Xprof Memory Viewer mostra o card de detalhes do buffer, que contém informações sobre o buffer, incluindo informações de padding.
  • Exemplo: um formato de (129, 1024) pode ser preenchido com (256, 1024), resultando em quase 50% de desperdício de memória.
  • Correção:uma forma de (128, 1024) não requer padding e causa 0% de desperdício de memória.
• Alinhe as dimensões:garanta que todas as dimensões grandes de tensores (tamanho do lote, dimensão de embedding, tamanho oculto) sejam múltiplos de 128.

B. Ajustar a configuração

É possível resolver erros de falta de memória com estes ajustes de configuração:

• Reduza o tamanho do lote:a memória necessária para ativações e gradientes intermediários é diretamente proporcional ao tamanho do lote. Reduzir o tamanho do lote pode ajudar a diminuir o uso de memória.
• Doar buffers de entrada:ao usar jax.jit, especifique donate_argnums para os parâmetros do modelo. Isso permite que o XLA substitua a memória de entrada pela saída.
• Ativar a precisão mista (bfloat16): use bfloat16 ou quantização (int8 etc.) para os maiores tensores no programa, se a arquitetura do modelo e os requisitos de qualidade permitirem.

C. Otimizar arquitetura e fragmentação

Se as mudanças na configuração forem insuficientes, a topologia do modelo poderá ser muito grande para a configuração de hardware atual.

• Use gerações mais recentes de TPUs:as TPUs mais recentes geralmente oferecem mais HBM por chip. Mude para gerações mais recentes de TPUs, se disponíveis.
• Executar em uma topologia de chip maior:se os pesos do modelo forem muito grandes para a topologia atual, tente fragmentá-los em mais chips.
• Implemente técnicas avançadas de fragmentação:
  • Conheça abordagens mais avançadas de paralelismo de dados, tensores ou pipelines.
  • Especifique dicas de fragmentação para valores e saídas intermediárias.
• Use o descarregamento do host JAX:descarregue tensores grandes para a memória da CPU do host, por exemplo, descarregamento de ativação e descarregamento do estado do otimizador.

D. Ajuste as flags do XLA que afetam a memória:

As flags de memória principais podem ser ajustadas para trocar o desempenho por um uso menor de memória. No entanto, elas devem ser usadas como último recurso, porque podem afetar negativamente o desempenho.

E. Ajustar a transmissão de rematerialização do XLA / Checkpointing manual

Se o modelo estiver quase cabendo na memória, force a transmissão XLA::Rematerialization para priorizar a economia de memória, possivelmente ao custo de compilações mais lentas:

Como alternativa, use o decorador jax.checkpoint com jax.grad para controlar manualmente quais intermediários são salvos na transmissão direta e quais são recalculados na transmissão indireta, trocando ciclos de computação por HBM.

F. Usar ferramentas avançadas de criação de perfil

Depurar erros de OOM com o XProf oferece um tutorial sobre como usar o Visualizador de memória do XProf para visualizar a visão do compilador sobre o uso de HBM.

Essa ferramenta permite ver a alocação de memória de pico e os tempos de vida do buffer, o que é crucial para entender exatamente o que consome HBM no ponto de utilização de pico. Para a configuração geral de criação de perfil, consulte Como começar a usar o Xprof e Criação de perfil do TensorBoard.