Depuração de OOM em JAX a partir de um capture

O que Depuração de OOM em JAX a partir de um capture faz?

Visão geral

Um passo jitted do JAX obtém toda a sua memória temporária como um único bloco temp planejado. Portanto, a memória do passo é o plano do XLA, e o arquivo .xplane.pb contém tanto esse plano quanto os eventos do alocador em tempo de execução. O script scripts/oom_report.py lê ambos em 2 a 5 segundos, usando apenas a biblioteca padrão:

• plano: a atribuição de buffers de cada módulo, a partir de /host:metadata. Ele fornece os parâmetros (e quanto é doado), as saídas não aliadas, o bloco temp e o pico de heap. Para os buffers ativos nesse pico, ele informa tamanho, formato, função e linha Python. A função vem da pilha de nomes JAX: jvp( é forward (mantido para o backward), transpose( é backward, as saídas de custom-call s8 são workspace de biblioteca.
• tempo de execução: os eventos do alocador BFC. Eles fornecem a capacidade, bytes em uso no início de cada execução, o pico, esperas AllocatorRetry, execuções que nunca lançaram um kernel e os eventos do host entre as duas últimas execuções. Eles também fornecem o que estava ativo quando a última execução começou, agrupado pelo evento do host que o alocou: outputs of jit_step execution 3, BufferFromHostBuffer(s32[…]) (um device_put), ou "alocado antes do início do capture".

O relatório termina com linhas de diagnóstico, cada uma etiquetada como medido ou inferido. Execute-o primeiro:

python3 -P <diretório base desta skill>/scripts/oom_report.py <diretório do capture ou .xplane.pb> [--module 'jit_train_step(123)'] [--top 25] [--json out.json]
python3 -P <diretório base desta skill>/scripts/oom_report.py --self-test

Por padrão, ele detalha o módulo cuja última execução nunca lançou um kernel. Se não houver nenhum, ele pega o maior módulo que executou. Use --module para escolher outro, por exemplo, um passo de avaliação.

Da linha de diagnóstico à causa e correção

Linha de diagnóstico	Causa	Correção a propor; a economia vem dos bytes do relatório
o passo em si não cabe	o plano sozinho excede o pool	Siga a divisão de funções no pico (abaixo)
↳ principalmente forward, kept for the backward pass	ativações: crescem com batch × seq; atenção ingênua adiciona B·H·S·S por camada	jax.checkpoint / remat no bloco (economiza a maioria dos bytes forward-mantidos, custa cerca de um passo forward extra); um micro-batch menor com acúmulo de gradientes; uma sequência mais curta (buffers S² encolhem 4× quando S é reduzido à metade). Ofereça atenção fundida apenas como "verifique se os buffers B·H·S·S saem do plano"
↳ principalmente backward	gradientes e seus temporários	Fragmentar (FSDP) ou acumular; verifique cópias f32 de gradientes bf16
↳ workspace	workspace de biblioteca GEMM / conv	Geralmente pequeno; no ROCm são 64 MiB por GEMM hipBLASLt no gfx950
cresce X MiB por execução, = 1,00x as saídas não aliadas	o programa mantém as saídas de cada passo vivas (um vazamento entre passos): a lista de ativos mostra outputs of … execution 1, 2, …	Encontre o detentor nos eventos do host entre execuções (ex.: $<unknown> append = list.append). Em seguida, remova a referência, reduza dentro do passo (uma norma, não a árvore) ou faça jax.device_get. Um pool maior só atrasa o OOM
nenhum parâmetro é doado … pode economizar até X	parâmetros antigos e novos / estado do otimizador estão vivos juntos	donate_argnums nos argumentos de estado
outros arrays ativos ocupam espaço; o maior … veio de X	arrays fora do passo: um dataset ou cópia de avaliação colocados na GPU, pré-busca, outras saídas jit	Mantenha no host e faça streaming, libere ou reduza seu dtype. "Alocado antes do início do capture" significa perfilhar mais cedo para nomeá-lo
AllocatorRetry … mais longo N s, não lançou nenhum	o OOM ocorreu dentro do capture	Apenas evidência. A requisição com falha geralmente é o bloco temp desse módulo ou uma saída
XLA já rematerializou N instruções	o próprio remat do XLA já foi executado contra o limite	Espere que a correção venha do programa, não do XLA

Dê a cada correção sua economia estimada e seu custo, etiquetado como inferido. Verifique a estimativa contra a capacidade. "Aumente XLA_PYTHON_CLIENT_MEM_FRACTION" corrige um plano que não cabe apenas se a GPU tiver espaço. Capacidade bem abaixo do HBM sugere que a fração está definida; o capture não registra variáveis de ambiente.

Anatomia da memória do playbook em termos JAX

O Ultra-Scale Playbook ("Profiling the memory usage") divide a memória de treinamento em pesos, gradientes, estado do otimizador e ativações. No relatório:

• Pesos + estado do otimizador = os parâmetros do passo. Adam em fp32 são 12 B por parâmetro (p, m, v). Doado significa sem segunda cópia.
• Gradientes = a função backward no pico, mais saídas se o passo as retornar.
• Ativações = forward, kept for the backward pass. Modelos com Scan / nn.scan mantêm-nas como buffers empilhados [camadas, …] em um loop while.
• "Passo 1 cabe, passo N falha": em PyTorch isso é estado do otimizador criado no primeiro passo. Em JAX o estado existe antes do passo 1, então procure pela linha de crescimento. A primeira chamada também compila e faz autotune, portanto um pico antes da primeira execução é tempo de compilação.

Regras

• Use apenas o capture. Não reconstrua o modelo nem reexecute a carga de trabalho para testar uma correção: estime-a a partir dos bytes do plano e diga como confirmá-la (jit(...).lower(...).compile().memory_analysis() no lado do usuário).
• Não escreva dentro do diretório do capture. Os conversores do XProf (xspace_to_tool_names, memory_viewer, CLI xprof) escrevem arquivos de cache *.hlo_proto.pb ao lado do .xplane.pb. Se precisar de visualizações do XProf, copie o .xplane.pb para um diretório temporário primeiro.
• Etiquete cada número como medido, lido ou inferido.

Referência medida (um único MI355X, pools pequenos, 2026-11-06)

Capture	O que o relatório mostra
plano de atenção muito grande	precisa de 3,72 GiB > 2,79 GiB; pico de heap 81,6% forward-mantido; f32[8,8,1024,1024] ×5 na linha softmax; bloco temp 3.495 MiB = o "3,41 GiB" do erro; AllocatorRetry 10 s; 89 instruções de remat
gradientes mantidos a cada passo	em uso no início do passo 316 → 1.252 MiB, +104,0 MiB por passo = 1,00× as saídas não aliadas; plano 788 MiB cabe; $<unknown> append entre passos
sem doação	precisa de 3,88 GiB > 2,79 GiB; 1.902 MiB de parâmetros + 1.902 MiB de saídas, 0 doados; temp apenas 172 MiB
MaxText llama2-7b, 8 GPUs, FSDP (cabe)	temp 61,9 GiB; pico de heap 57% forward-mantido, os stacks do scan bf16[32,4,4096,11008] ×2 via while.50

Limitações

• Um host por arquivo: execuções multi-host têm um .xplane.pb por host; portanto, execute em cada um.
• Apenas HBM: apenas o espaço de memória 0 é detalhado. Offload de host aparece como memory space 5.
• Contiguidade: "cabe" compara bytes. BFC também precisa de um bloco contíguo, portanto uma situação marginal ainda pode falhar por fragmentação.
• Função "other": buffers definidos por operações criadas pelo XLA sem pilha de nomes JAX e sem aliasing que tenha uma.
• CUDA: não testado. Usa os mesmos eventos BFC e XLA, mas os nomes não foram verificados.
• TPU / eager: captures de TPU fornecem apenas totais do plano, sem pico de heap e sem execuções. Programas eager (op-por-op) aparecem como muitos módulos pequenos; portanto, leia a seção de tempo de execução nesses casos.