investigando-tracos-jax

Investigação de traços JAX

Visão geral

Uma investigação em sete etapas de uma captura: o veredito, reconciliação, faixas temporais (swimlanes), kernel → HLO → linha Python, roofline, o que mudar e a lista de confiabilidade.

Tudo vem da captura que o usuário forneceu: <host>.xplane.pb e/ou <host>.trace.json.gz. Não reexecute a carga de trabalho, não despeje XLA nem execute outro profiler. Onde o traço não puder responder, coloque na lista "não pode ser confiado".

Objetivo

Após o relatório, o leitor sabe:

• (a) o que mudar para tornar a etapa mais rápida, como uma estimativa com suas premissas;
• (b) onde este traço não pode suportar a atribuição, com contagens.

Ferramentas

Os caminhos dos scripts são relativos ao diretório base deste skill, aquele mostrado quando o skill carrega. Execute a partir de um diretório temporário: ambos os scripts escrevem suas saídas no diretório atual.

B=<diretório base deste skill>
python3 -P $B/../analyzing-jax-launch-overhead/scripts/skip_stats.py <captura>/<host>.trace.json.gz
python3 -P $B/scripts/attribute_kernels.py <diretório da captura> [--program ID] [--user-re REGEX]

• skip_stats.py (apenas biblioteca padrão): o veredito, o orçamento de tempo por fluxo, o evento do host durante o tempo ocioso e as métricas do artigo SKIP. Como lê-lo: o skill analyzing-jax-launch-overhead.
• attribute_kernels.py (precisa dos pacotes jax e xprof): todo kernel GPU → instrução HLO → op_name → linha Python. Ele lê apenas o .xplane.pb, copiando-o para ./attribution/xplane/ primeiro, porque o XProf escreve os protos HLO ao lado dele. Ele escreve kernels.csv, ops.csv, lines.csv, how.csv, summary.txt e o texto HLO do programa para ./attribution/.
  • exact: o próprio hlo_op do tracer, para kernels lançados um por chamada.
  • name: um kernel de grafo cujo nome é o único candidato em sua janela.
  • order: kernels compartilhados, GEMMs de biblioteca, cópias e coletivos, correspondidos em ordem de agendamento.
  • none: nenhum candidato restante.
  • Janelas são as instruções entre kernels exatos vizinhos, no cronograma do programa com loops while expandidos por sua contagem de iterações conhecida.
  • Validação: em uma captura do MiniGPT, ele escolheu a mesma instrução para todos os 147.840 kernels, como uma pipeline que concordava 100% com uma execução do mesmo programa com buffers de comando desligados.

Cinco técnicas, cada uma com sua ferramenta

Campo	Técnica	Ferramenta aqui
Engenharia de desempenho	Caminho crítico; roofline	O orçamento de fluxo de skip_stats.py; um roofline apenas com os picos corretos (veja armadilhas)
Rastreamento distribuído	Juntar eventos de host e GPU por ID de correlação; encontrar retardatários	skip_stats.py → rows.csv; a dispersão dos tempos de fim de etapa entre GPUs
Compiladores	Proveniência da fonte através de passes do compilador	attribute_kernels.py → lines.csv
Contabilidade forense	Totais devem reconciliar entre registros independentes	As verificações na etapa 2
Design de informação	Faixas temporais alinhadas, uma por fluxo de GPU	Uma etapa, cada GPU × fluxo, em um eixo de tempo

Etapas

O plano são estas sete etapas, todas executadas nos arquivos fornecidos. Uma pergunta que apenas uma nova captura poderia responder (o que uma chamada de host faz internamente, o efeito de uma correção, uma varredura de lotes) vai para a lista "precisa de uma nova captura" da etapa 7, e decidir se captura é por conta do usuário.

1. Veredito e causa do lado do host: execute skip_stats.py.
2. Reconciliar. Cada verificação ou é válida ou vai para a lista de confiabilidade:
   • Eventos GPU por dispositivo: json = xplane;
   • cada evento GPU se junta a uma chamada de lançamento do host;
   • kernels por etapa são constantes;
   • todas as GPUs fazem o mesmo trabalho (contagens de kernel iguais por etapa, ou a diferença explicada);
   • para cada fluxo, as parcelas de computação + coletivo + memcpy + ociosidade somam ao intervalo.
3. Faixas temporais. Para a etapa mediana, coloque cada GPU × fluxo em um eixo de tempo, mais os eventos do host da thread de lançamento. Isso mostra o caminho crítico e os retardatários: qual GPU termina por último, e o que essa GPU espera.
4. Kernel → HLO → linha Python: execute attribute_kernels.py.
   • Relate a parcela do tempo de GPU de cada classe de correspondência e os kernels deixados sem candidato.
   • Use lines.csv e ops.csv para as tabelas por linha e por operação.
   • --user-re seleciona quais quadros contam como código do usuário. Por padrão, é o primeiro quadro fora de site-packages e jax.
5. Roofline, apenas para operações cujos FLOPs e bytes você pode obter de suas formas HLO: o <program>(<id>).hlo.txt que attribute_kernels.py escreve em ./attribution/. Use o pico para o tipo de dado da operação, citado da especificação do fornecedor e rotulado como lido; o traço não pode medir picos.
6. O que mudar. Classifique as causas por sua parcela do caminho crítico. Para cada uma, estime o tempo da etapa se ela desaparecesse, por exemplo, etapa − ociosidade do host no fluxo crítico. Rotule como inferido, não executado e nomeie a alavanca como uma hipótese para o usuário testar (donate_argnums, buffers de comando desligados, menos coletivos ou maiores).
7. Onde este traço não pode ser confiado: liste cada item com contagens ou parcelas. Em seguida, adicione a lista "precisa de uma nova captura": cada pergunta em aberto e qual captura a responderia.

Armadilhas (medidas em capturas ROCm MI355X)

• As visualizações de operações do XProf omitem kernels de grafo. hlo_stats cobriu 37,6% do tempo de kernel GPU em uma captura de 4 GPUs. Os outros 62,4% (kernels de buffer de comando) aparecem em kernel_stats como operação "XlaModule". Verifique a cobertura antes de citar hlo_stats, op_profile ou roofline_model.
• Os picos de roofline do XProf para gfx950 são 2.517 TFLOP/s e 3.624 GiB/s, seja qual for o tipo de dado. GEMMs FP32 precisam do pico FP32.
• O tempo de kernel RCCL/NCCL inclui espera por pares. Não é tempo de transferência.
• Coletivos assíncronos: seu hlo_op (por exemplo, all-gather.34) não está no HLO armazenado na captura. attribute_kernels.py mantém o rótulo e pega a linha Python da instrução collective-start correspondente, por ordem.
• Linhas de memcpy ROCm podem mostrar kind:Unknown size:0.
• O tracer Python infla o tempo do host.
• Um programa pode aparecer duas vezes (dois arquivos .hlo_proto.pb com IDs diferentes). attribute_kernels.py usa o ID do programa com mais tempo de GPU, a partir do xplane; o json não tem program_id.

Relatório

• Markdown: achados primeiro.
• Números: cada um vem com seu intervalo entre etapas e GPUs, rotulado como medido, lido ou inferido.
• Seções: veredito e causa, reconciliação, tabela por GPU/fluxo, principais operações e linhas Python, o que mudar (estimativas) e a lista de confiabilidade.
• Verificação de fatos: antes de enviar, um subagente novo deve verificar cada afirmação contra as tabelas.
• Página HTML: se o usuário quiser, o skill writing-trace-report-pages a constrói a partir deste relatório e suas tabelas.