Ferramenta Trace Viewer

Você pode usar o Trace Viewer para visualizar a linha do tempo dos eventos que ocorreram durante a sessão de criação de perfil. Ele mostra as durações das operações executadas pelo modelo em diferentes partes do sistema, como o host (CPU) e os aceleradores (GPUs ou TPUs). Isso permite entender como o modelo utiliza recursos de hardware, identificar gargalos de desempenho e otimizar o modelo para uma execução mais rápida. A interface do Visualizador de trace é baseada na usada em chrome://tracing e, portanto, exige que você use o navegador Chrome.

Plataformas compatíveis

TPU: compatível

GPU: compatível

O Trace Viewer para GPUs agrupa eventos em linhas do tempo por chip de GPU e as streams neles. Já para TPUs, ele organiza eventos em linhas do tempo por chip e núcleo de TPU. Em ambos os casos, o Trace Viewer também apresenta eventos em linhas do tempo por linha de execução para a CPU host à qual os aceleradores estão conectados.

Como interagir com a linha do tempo

O Trace Viewer oferece várias ferramentas e métodos para navegar e examinar a linha do tempo:

  • Navegação: você pode usar os seguintes atalhos de teclado:
    • W:aumenta o zoom.
    • S:diminui o zoom.
    • A Mover para a esquerda.
    • D:Mover para a direita.
  • Seletor de ferramentas: clique em uma ferramenta no seletor ou use o atalho de teclado correspondente:
    • Ferramenta de seleção (1 ou !): clique em um evento para selecioná-lo e conferir os detalhes dele no painel "Detalhes". Selecione vários eventos usando Ctrl + clique para ver um resumo.
    • Ferramenta de panorama (2 ou @): arraste para mover a linha do tempo na horizontal ou na vertical.
    • Ferramenta de zoom (3 ou #): arraste para aumentar o zoom em uma área específica da linha do tempo.
    • Ferramenta de tempo (4 ou $): arraste para marcar um intervalo de tempo. A duração do intervalo marcado será exibida. Você também pode usar a chave "m" para marcar uma seleção e determinar a duração total.
  • Ampliar eventos selecionados (f): selecione um ou mais eventos e pressione a tecla "f" para ampliar rapidamente essa parte da linha do tempo. Isso é útil para se concentrar em uma etapa específica do treinamento.

Ao processar um grande número de eventos de trace, o Visualizador de traces funciona no modo de streaming. Isso significa que ele carrega dados sob demanda à medida que você desliza e faz zoom na linha do tempo, de forma semelhante ao funcionamento dos aplicativos de mapa. Se você aumentar o zoom mais rápido do que os dados podem ser carregados, talvez você veja uma representação de baixa resolução dos dados até que o carregamento seja concluído.

Componentes da interface do Trace Viewer

Estes são os principais componentes da interface do Visualizador de traces:

  • O eixo do tempo é horizontal na parte de cima, mostrando o tempo em relação ao início do trace.
  • As linhas do tempo são organizadas por seções e faixas, com rótulos localizados no eixo vertical à esquerda. Cada seção representa um elemento de processamento (por exemplo, um nó de dispositivo ou linhas de execução do host) e pode ser expandida ou recolhida. Em cada seção, há faixas, que são cronogramas para atividades específicas.
  • Os eventos são blocos retangulares coloridos nas faixas da linha do tempo, que representam a duração de uma operação ou um metaevento, como uma etapa de treinamento. A cor dos eventos não tem um significado específico.
  • O painel de detalhes mostra mais informações sobre os eventos selecionados no painel da linha do tempo, como nome, horário de início e duração.

Seções e faixas típicas

O Trace Viewer fornece as seções e faixas a seguir.

  • Uma seção para cada nó da TPU, com as seguintes faixas:
    • Módulo XLA: o programa XLA que está sendo executado.
    • Operações do XLA: mostra as operações de HLO do XLA que foram executadas no núcleo da TPU. Cada operação de framework de nível mais alto (JAX, TensorFlow ou PyTorch, por exemplo) é traduzida em uma ou várias operações do XLA, que são em seguida compiladas para serem executadas na TPU. Ao clicar em uma operação XLA, você pode conferir links para a operação no Visualizador de gráficos, além de outras informações sobre a operação, como os tempos de início/parada, a duração e o trace de pilha de origem (se disponibilizado pelo framework e pelo compilador).
    • XLA TraceMe: anotações especificadas pelo usuário no código que descrevem unidades lógicas de trabalho que ele pretende rastrear. Ainda é possível ver dados aqui, mesmo que você não tenha adicionado nenhuma anotação. Elas geralmente são adicionadas pelo XLA (por exemplo, cores de barreira) ou pelo próprio XProf (por exemplo, entradas de rastreamento descartadas).
    • Etapa: mostra a duração das etapas de treinamento executadas nesse núcleo da TPU, se devidamente anotadas no programa ou framework do usuário.
    • Framework Ops: mostra as operações do framework (JAX, Tensorflow ou PyTorch, por exemplo) executadas no núcleo da TPU, se devidamente anexadas ao programa ou framework do usuário.
    • Escopo do nome do framework: para cada operação de framework, uma visualização do stack trace. Para encurtar, essa faixa aparece apenas para um dispositivo.
    • Código-fonte: caminho para o código-fonte que está sendo executado, se disponível no perfil.
    • Unidade escalar: para TPUs, eventos executados na unidade escalar, mostrados se presentes no perfil.
    • TensorCore Sync Flags: mecanismo de sincronização em TPUs, mostrado se presente no perfil.
    • Desativação do host: operações que movem dados de forma assíncrona entre a memória do host e a memória do acelerador. Normalmente, há operações de início e parada correspondentes que aparecem na linha de operações XLA, indicando que o acelerador está se preparando para a transferência de dados (por exemplo, marcando regiões de memória de origem/destino como "em uso" durante a transferência). Várias linhas de transferência de host podem estar presentes se houver várias operações de transferência executadas em paralelo, exigindo que o visualizador de rastros exiba vários eventos simultaneamente.
  • Uma seção para cada nó Sparsecore: algumas gerações de TPU (por exemplo, A TPU v5p e a TPU v6e) são equipadas com uma ou mais unidades SparseCore, além das unidades MXU de computação densa. Os módulos, operações e TraceMes associados a esses núcleos vão aparecer nesta seção.
  • Uma seção para cada nó de GPU, com as seguintes faixas:
    • Módulos XLA, operações do framework, escopo do nome do framework, etapas, código-fonte. Todas são semelhantes às seções de TPU.
    • Uma faixa por stream, com o nome do stream também incluindo informações sobre os tipos de operações executadas no stream (Memcpy, Compute etc.).
    • O XLA TraceMe não tem suporte para GPUs.
    • As operações XLA aparecem em seções de GPU, mas atualmente não são sempre precisas, já que são derivadas dos dados de stream. Portanto, elas não podem explicar totalmente o modelo de execução da GPU, em que pode haver um mapeamento N:M de operações XLA para os kernels reais que são executados nos diferentes streams e a programação dinâmica de vários streams em diferentes SMs no hardware.
  • Uma seção para cada componente (por exemplo, um threadpool) em execução na CPU da máquina host, com uma faixa por thread, no caso de threadpools. É aqui que você vai encontrar os rastros do Python, se eles foram ativados durante a coleta de perfil.

Observe que apenas as operações XLA para TPUs e dados de stream para GPUs são diretamente fundamentadas no perfil coletado. Todas as outras linhas são "linhas derivadas", envolvendo informações opcionais de "sideband" fornecidas pelo compilador, anotações opcionais do usuário e/ou heurísticas aplicadas pelo XProf. Portanto, essas linhas derivadas podem ou não aparecer em determinados perfis.

Outros recursos

  • É possível pesquisar nomes de eventos específicos usando a barra de pesquisa "Encontrar eventos". No momento, a pesquisa é feita apenas na janela de tempo visível na tela, e não no rastro completo.
  • Eventos de fluxo: ativar essa opção clicando no botão "Eventos de fluxo" na barra superior adiciona visualizações, vinculando eventos em uma linha ou linha de conversa a eventos em outra linha de conversa. Por exemplo, pode haver uma seta desenhada da operação no host que enfileira ou inicia o trabalho para um acelerador para a operação no acelerador que executa esse trabalho. O XProf determina esses links por meio de uma combinação de anotações do usuário, heurísticas integradas e informações de pós-processamento que recebe de diferentes componentes (por exemplo, drivers do CUPTI, IDs de inicialização do kernel, informações de execução da TPU etc.
  • Clicar em uma operação XLA mostra mais informações no painel de detalhes. Por exemplo, ele vincula a operação na ferramenta Visualizador de gráficos. Ele também pode fornecer ponteiros para o código-fonte e/ou o stack trace do Python, a operação do framework que causou a geração dessa operação XLA etc. (se presente no perfil). Ele também pode mostrar FLOPS (número de operações de ponto flutuante executadas pela operação) e bytes acessados pela operação. Essas informações são adquiridas estaticamente do XLA durante a compilação, em vez de informações de execução do perfil.