Criar perfis de cargas de trabalho do PyTorch XLA

O perfil é uma maneira de analisar e melhorar o desempenho dos modelos. Embora exista muito mais do que isso, às vezes é útil pensar no perfil como operações de tempo e partes do código executadas em dispositivos (TPUs) e hosts (CPUs). Este guia oferece uma visão geral rápida de como criar um perfil do código para treinamento ou inferência. Para mais informações sobre como analisar os perfis gerados, consulte os guias abaixo.

• Depuração de desempenho do PyTorch XLA em VMs de TPU: parte 1
• Depuração de desempenho do PyTorch XLA em VMs de TPU: parte 2
• Depuração de desempenho do PyTorch XLA em VMs de TPU - parte 3

Primeiros passos

Criar TPU

1. Exporte variáveis de ambiente:

   $ export TPU_NAME=your_tpu_name
   $ export ZONE=us-central2-b
   $ export PROJECT_ID=project-id
   $ export ACCELERATOR_TYPE=v4-8
   $ export RUNTIME_VERSION=tpu-vm-v4-pt-2.0

   Exportar descrições de variáveis

   TPU name

   O nome que você quer usar para a Cloud TPU.

   zone

   A zona em que você planeja criar a Cloud TPU.

   project ID

   O ID do projeto que você está usando para treinar e criar o perfil do modelo.

   accelerator-type

   O tipo de acelerador especifica a versão e o tamanho da Cloud TPU que você quer criar. Para mais informações sobre os tipos de aceleradores compatíveis com cada versão de TPU, consulte Versões de TPU.

   version

   A versão do ambiente de execução do Cloud TPU. Um valor padrão é mostrado na variável exportada, mas você também pode usar um da lista de configurações com suporte.

2. Iniciar os recursos da TPU

   $ gcloud compute tpus tpu-vm create ${TPU_NAME} \
   --zone us-central2-b \
   --accelerator-type ${ACCELERATOR_TYPE} \
   --version ${RUNTIME_VERSION} \
   --project $PROJECT_ID \
   --subnetwork=tpusubnet

3. Mova o código para o diretório principal na VM da TPU usando o comando gcloud scp. Exemplo:

   $ gcloud compute tpus tpu-vm scp my-code-file ${TPU_NAME}: --zone ${ZONE}

Criação de perfil

Observação: um v4-8 inclui quatro dispositivos de TPU. Cada trabalhador cria perfis de apenas um dispositivo TPU dos quatro disponíveis. O perfil de todos os quatro dispositivos de TPU para cada worker será suportado em breve.

Um perfil pode ser capturado manualmente pelo capture_profile.py ou de maneira programática no script de treinamento usando as APIs torch_xla.debug.profiler.

Como iniciar o servidor de perfil

Para capturar um perfil, um servidor de perfil precisa estar em execução no script de treinamento. Inicie um servidor com um número de porta de sua escolha, por exemplo, 9012, conforme mostrado no comando a seguir.

Observação: consulte este script de exemplo que inclui as etapas a seguir.

import torch_xla.debug.profiler as xp
server = xp.start_server(9012)

O servidor pode ser iniciado logo no início da função main.

Observação: se você estiver usando o módulo Lightning, em vez de iniciar um servidor, transmita profiler="xla" para o trainer. O servidor será iniciado automaticamente na porta 9012.

Agora é possível capturar perfis conforme descrito na seção a seguir. O script cria perfis de tudo o que acontece em um dispositivo TPU.

Adicionar rastros

Se você também quiser criar um perfil de operações na máquina host, adicione xp.StepTrace ou xp.Trace ao código. Essas funções rastreiam o código Python na máquina host. Isso pode ser considerado como o tempo necessário para executar o código Python no host (CPU) antes de transmitir o "gráfico" para o dispositivo TPU. Portanto, é mais útil para analisar a sobrecarga de rastreamento. É possível adicionar isso dentro do loop de treinamento em que o código processa lotes de dados, por exemplo,

for step, batch in enumerate(train_dataloader):
    with xp.StepTrace('Training_step', step_num=step): 
       ...

ou envolva partes individuais do código com

 with xp.Trace('loss'): 
    loss = ...

Importante: evite agrupar xm.mark_step() com xp.Trace. O código vai trava nesse caso com uma mensagem de erro semelhante a: "RuntimeError: Expecting scope to be empty but it is train_loop." xp.StepTrace não terá esse problema, já que ele adiciona automaticamente xm.mark_step() ao sair do escopo. Além disso, você pode receber esse erro ao usar torchdynamo, porque ele pode invocar xm.mark_step internamente. A solução alternativa atual é não usar xp.Trace.

Se você estiver usando o Lighting, pode pular a adição de rastros, já que isso é feito automaticamente em algumas partes do código. No entanto, se você quiser adicionar outros rastros, insira-os no ciclo de treinamento.

Você poderá capturar a atividade do dispositivo após a compilação inicial. Aguarde até que o modelo inicie as etapas de treinamento ou inferência.

Captura manual

O script capture_profile.py do repositório Pytorch XLA permite capturar um perfil rapidamente. Para fazer isso, copie o arquivo de perfil de captura diretamente para a VM da TPU. O comando a seguir copia para o diretório principal.

$ gcloud compute tpus tpu-vm ssh ${TPU_NAME} \
--zone us-central2-b \
--worker=all \
--command="wget https://raw.githubusercontent.com/pytorch/xla/master/scripts/capture_profile.py"

Enquanto o treinamento está em execução, siga estas etapas para capturar um perfil:

$ gcloud compute tpus tpu-vm ssh ${TPU_NAME} \
--zone us-central2-b \
--worker=all \
--command="python3 capture_profile.py --service_addr "localhost:9012" --logdir ~/profiles/ --duration_ms 2000"

Esse comando salva arquivos .xplane.pb no logdir. É possível mudar o diretório de registro ~/profiles/ para o local e o nome de sua preferência. Também é possível salvar diretamente no bucket do Cloud Storage. Para fazer isso, defina logdir como gs://your_bucket_name/.

Captura programática

Em vez de capturar o perfil manualmente, acionar um script. é possível configurar o script de treinamento para acionar automaticamente um perfil usando a API torch_xla.debug.profiler.trace_detached no script de treinamento.

Observação: a API trace_detached só está disponível em builds noturnos ou após a versão 2.2. Para capturar programaticamente em versões anteriores do torch_xla, copie o corpo do método para criar a linha de execução de captura em segundo plano diretamente.

Por exemplo, para capturar automaticamente um perfil em uma época e etapa específicas, configure o script de treinamento para consumir as variáveis de ambiente PROFILE_STEP, PROFILE_EPOCH e PROFILE_LOGDIR:

import os
import torch_xla.debug.profiler as xp

# Within the training script, read the step and epoch to profile from the
# environment.
profile_step = int(os.environ.get('PROFILE_STEP', -1))
profile_epoch = int(os.environ.get('PROFILE_EPOCH', -1))
...

for epoch in range(num_epoch):
   ...
   for step, data in enumerate(epoch_dataloader):
      if epoch == profile_epoch and step == profile_step:
         profile_logdir = os.environ['PROFILE_LOGDIR']
         # Use trace_detached to capture the profile from a background thread
         xp.trace_detached('localhost:9012', profile_logdir)
      ...

Isso vai salvar os arquivos .xplane.pb no diretório especificado pela variável de ambiente PROFILE_LOGDIR.

Análise no TensorBoard

Para analisar melhor os perfis, use o TensorBoard com o plug-in do TPU TensorBoard
na mesma máquina ou em outra (recomendado).

Para executar o TensorBoard em uma máquina remota, conecte-se a ela usando SSH e ative o encaminhamento de portas. Por exemplo,

$ ssh -L 6006:localhost:6006 remote server address

ou

$ gcloud compute tpus tpu-vm ssh $TPU_NAME --zone=$ZONE --ssh-flag="-4 -L 6006:localhost:6006"

Na máquina remota, instale os pacotes necessários e inicie o TensorBoard, assumindo que você tenha perfis nessa máquina em ~/profiles/. Se você armazenou os perfis em outro diretório ou bucket do Cloud Storage, especifique os caminhos corretamente, por exemplo, gs://your_bucket_name/profiles.

(vm)$ pip install tensorflow-cpu tensorboard-plugin-profile 

(vm)$ tensorboard --logdir ~/profiles/ --port 6006

Observação: é importante ter arquivos .xplane na mesma estrutura aninhada que é gerada quando você captura um perfil. Por exemplo: ~/profiles/plugins/profile/2023_04_10_21_40_22/localhost_9012.xplane.pb

Observação: se você receber o erro "plugins duplicados", isso pode ser porque há vários pacotes com o TensorBoard. Desinstale todos esses pacotes e execute os comandos pip install novamente. Por exemplo, talvez seja necessário executar:

(vm)$ pip uninstall tensorflow tf-nightly tensorboard tb-nightly tbp-nightly

Como executar o TensorBoard

No navegador local, acesse: http://localhost:6006/ e escolha PROFILE no menu suspenso para carregar seus perfis.

Consulte Ferramentas de TPU para informações sobre as ferramentas do TensorBoard e como interpretar a saída.