Criação de perfis de desempenho do PyTorch XLA

Visão geral

Neste guia, você verá como usar as ferramentas de desempenho do Cloud TPU e o recurso de análise automática de métricas com o PyTorch. Essas ferramentas ajudam você a depurar e otimizar o desempenho da carga de trabalho de treinamento.

Conceitos

Se você for novo no PyTorch / XLA, consulte os documentos API_Guide e solução de problemas do PyTorch. Para o Cloud TPU, consulte o documento de conceitos.

Importante: use este guia com a VM de TPU ou a configuração de nó de TPU. As duas arquiteturas de VM são descritas em Arquitetura do sistema. As ferramentas de desempenho das duas arquiteturas são mostradas em seções separadas, abaixo. Na primeira seção, descrevemos as ferramentas de desempenho de criação de perfil do PyTorch usando a configuração do nó da TPU. A segunda seção descreve as ferramentas de desempenho do PyTorch para a configuração da VM do TPU.

Nó TPU + criação de perfil PyTorch/XLA

Criar recursos relacionados à Cloud TPU

Crie e inicialize os recursos necessários relacionados à Cloud TPU.

1. Crie variáveis para o ID do projeto, seu bucket do Cloud Storage e a zona a ser usada nos recursos da TPU.

   export PROJECT_ID=PROJECT_ID
   export BUCKET_NAME=BUCKET_NAME
   export ZONE=ZONE
   
   gcloud --project=$PROJECT_ID compute project-info add-metadata \
   --metadata BUCKET_NAME=$BUCKET_NAME
   

2. Criar uma instância de VM do Compute Engine. É onde todos os modelos e scripts do Python são armazenados.

   gcloud compute instances create profiler-tutorial-vm \
     --project=${PROJECT_ID} \
     --zone=${ZONE} \
     --machine-type=n1-standard-16 \
     --image-project=ml-images \
     --image-family=torch-xla \
     --boot-disk-size=300GB \
     --scopes=https://www.googleapis.com/auth/cloud-platform
   

3. Crie um recurso de TPU.

   gcloud compute tpus create profiler-tutorial-tpu \
    --project=${PROJECT_ID} \
    --zone=${ZONE} \
    --network=default \
    --version=pytorch-1.8 \
    --accelerator-type=v3-8
   

4. Criar um bucket do Cloud Storage

   Primeiro, instale a CLI gsutil se ainda não a tiver instalado: instruções de instalação.

5. Use o comando gsutil mb para criar um bucket do Cloud Storage em que todos os artefatos de criação de perfil são armazenados. Substitua as variáveis de região e nome de bucket pelos valores que você usará no treinamento.

   gsutil mb -p ${PROJECT_ID} -c standard -l REGION gs://${BUCKET_NAME}
   

   onde:

   • REGION é a região em que você criou o Cloud TPU, por exemplo, europe-west4;

6. Crie uma conta de serviço para o projeto da Cloud TPU.

   gcloud beta services identity create --service tpu.googleapis.com --project $PROJECT_ID
   

   O comando retorna uma conta de serviço do Cloud TPU com o formato a seguir:

   [e-mail]
   

   Por exemplo: [e-mail]

7. Exporte a conta de serviço e conceda permissões a ela no bucket de armazenamento. Substitua account-number pelo PROJECT_NUMBER retornado na saída da criação da conta de serviço.

   export SERVICE_ACCOUNT=[e-mail]
   gsutil acl ch -u $SERVICE_ACCOUNT:READER gs://${BUCKET_NAME}
   gsutil acl ch -u $SERVICE_ACCOUNT:WRITER gs://${BUCKET_NAME}
   

Configurar o TensorBoard

1. ssh para a porta de encaminhamento de VMs 9001 na VM para a porta 9001 na máquina local; Essa porta é usada para abrir a IU do TensorBoard no navegador local.

     gcloud compute ssh profiler-tutorial-vm \
      --project ${PROJECT_ID} \
      --zone ${ZONE} \
      --ssh-flag="-4 -L 9001:localhost:9001"
   

2. Crie um ambiente conda dedicado à instalação do TensorBoard:

   conda create -y -n tensorboard python=3.6
   conda activate tensorboard
   pip install tf-nightly==2.6.0.dev20210511 tb-nightly==2.6.0a20210512 tbp-nightly==2.5.0a20210511
   

3. Para testar a instalação, execute o servidor do TensorBoard na VM do Compute Engine e tente se conectar ao servidor em http://localhost:9001/#profile. sua máquina local:

    # Get bucket name
    BUCKET_NAME=$(curl "http://metadata.google.internal/computeMetadata/v1/project/attributes/BUCKET_NAME" -H "Metadata-Flavor: Google")
   

    tensorboard --logdir gs://${BUCKET_NAME} --port 9001
   

Ao acessar http://localhost:9001/#profile em sua máquina local, você verá algo assim:

Se você acessar http://localhost:9001, também poderá acessar a página de perfil acima selecionando a opção PROFILE no menu suspenso do canto superior direito, ao lado do botão UPLOAD .

Criar o perfil do modelo

Para manter o servidor do TensorBoard ativo, inicie uma nova janela de terminal e ssh na VM do GCE novamente, desta vez sem usar -L. opção de encaminhamento de portas).

1. Na nova janela do terminal, exporte novamente o ID do projeto, o ambiente do bucket de armazenamento e as variáveis de zona, porque ele está em um novo shell.

   export PROJECT_ID=PROJECT_ID
   export ZONE=ZONE
   

2. ssh na VM:

     gcloud compute ssh profiler-tutorial-vm \
      --project ${PROJECT_ID} \
      --zone ${ZONE}
   

   Importante: Para se preparar para configurar o ambiente de treinamento, armazene o valor de TPU_IP_ADDRESS retornado pelo seguinte comando, já que ele é usado posteriormente:

   conda activate torch-xla-1.8
   PROJECT_ID=$(curl "http://metadata.google.internal/computeMetadata/v1/project/project-id" -H "Metadata-Flavor: Google")
   export TPU_IP_ADDRESS=$(gcloud compute tpus describe profiler-tutorial-tpu --zone=${ZONE} --project=${PROJECT_ID} \
   --format="value(ipAddress)")
   echo TPU_IP_ADDRESS=${TPU_IP_ADDRESS}
   export XRT_TPU_CONFIG="tpu_worker;0;$TPU_IP_ADDRESS:8470"
   

3. Verifique se os testes de integração estão funcionando de ponta a ponta no ambiente:

   python /usr/share/torch-xla-1.8/pytorch/xla/test/test_profiler.py # takes <1 min
   >

4. Antes de começar a treinar, edite as seguintes linhas em /usr/share/torch-xla-1.8/pytorch/xla/test/test_profile_mp_mnist.py:

   Alteração:

   accuracy = train_mnist(flags, dynamic_graph=True, fetch_often=True)
   

   Para:

   accuracy = train_mnist(flags, dynamic_graph=False, fetch_often=False)
   

   Os dois argumentos acima em train_mnist causam artificialmente gráficos e buscas de tensores, que serão explorados posteriormente na seção "Análise automática". Por enquanto, como você está apenas analisando a TPU, o exemplo a seguir é executado com desempenho nominal.

5. Inicie uma execução de treinamento usada para criar o perfil do servidor:

   PT_XLA_DEBUG=1 XLA_HLO_DEBUG=1 python /usr/share/torch-xla-1.8/pytorch/xla/test/test_profile_mp_mnist.py --num_epochs 1000 --fake_data
   

Criação de perfil de TPU (servidor)

Quando o treinamento estiver em execução, acesse http://localhost:9001/#profile e capture um perfil usando as seguintes instruções:

A página a seguir será recarregada automaticamente:

As ferramentas compatíveis são mostradas na lista suspensa Ferramentas do painel esquerdo:

• Página de visão geral (isso não inclui a exibição do analisador de pipeline de entrada, que é fundamentalmente diferente entre o TensorFlow/TPU e a PyTorch / XLA TPU).
• Visualizador de memória
• Perfil das operações
• Visualizador de pod
• Estatísticas do TensorFlow (estatísticas no nível do framework, ou seja, estatísticas do PyTorch)
• Visualizador do Trace (exige o navegador Chrome)

Página de visão geral

Esta página mostra uma visão geral de um perfil capturado. Este exemplo mostra um tempo de inatividade muito alto porque está treinando um pequeno modelo no conjunto de dados MNIST.

Visualizador de memória

Mostra a memória do dispositivo (HBM) usada por tensor e op HLO PyTorch. O visualizador de memória captura uma visualização da memória por módulo HLO. Portanto, haverá alguns módulos como gráficos de alocação de dados do dispositivo (entrada e rótulos enviados em lote). Para ver o uso de memória de um módulo HLO específico, selecione o módulo na lista suspensa Hosts à esquerda:

Ao visualizar o módulo HLO selecionado, você tem uma visão holística da linha do tempo de ocupação HBM do módulo. Isso é ordenado por tamanho de alocação, tamanho de execução do programa e tamanho do preenchimento.

Cada uma das alocações de buffer pode ser mais inspecionada passando o cursor sobre elas. Por exemplo, para informar qual alocação está ocupando a maior parte da HBM do dispositivo:

No exemplo acima, (1) corresponde às anotações torch_xla.debug.profiler.Trace adicionadas pelo código do usuário. Ao inspecionar o código test_profile_mp_mnist.py, ele corresponde a esta linha:

   class MNIST(nn.Module):
   ...
     def forward(self, x):
       with xp.Trace('conv1'):
         x = F.relu(F.max_pool2d(self.conv1(x), 2))
         x = self.bn1(x)
   ...
   

Também no namespace op test_mnist, é possível ver que esse módulo HLO corresponde ao loop de avaliação porque tem o gerenciador de contexto xp.trace('test_mnist').

Perfil de operação XLA

O perfil de operações é uma ferramenta da Cloud TPU que exibe as estatísticas de desempenho das operações do XLA executadas durante o período de criação do perfil. O perfil de operações mostra:

• o uso do Cloud TPU como uma porcentagem do tempo gasto em operações por categoria e da utilização de TPUs FLOPS;
• as operações mais demoradas, Essas operações são metas em potencial para otimização.
• Detalhes de operações individuais, incluindo forma, padding e expressões que usam a operação

Use o perfil de operações para encontrar boas metas de otimização. Por exemplo, se o modelo alcança apenas 5% dos FLOPS de pico da TPU, é possível usar a ferramenta para identificar quais operações do XLA estão demorando mais para serem executadas e quantos FLOPS de TPU consomem.

Descrição de cada item:

1. Seção de visão geral. Mostra o uso da Cloud TPU e fornece sugestões para otimização.
2. Painel de controle. Contém controles que permitem definir o número de operações exibidas na tabela, quais operações são exibidas e como são classificadas.
3. Tabela de operações. Uma tabela que lista as principais categorias de operação do TensorFlow associadas às operações do XLA. Essas operações são classificadas por porcentagem do uso da Cloud TPU.
4. Cards de detalhes de operações. Detalhes sobre a operação que aparece quando você passa o cursor sobre uma operação na tabela. Isso inclui a utilização de FLOPS, a expressão em que a operação é usada e o layout da operação (ajuste).

Visualizador de pod

Consulte ferramentas de TPU para uma descrição completa da ferramenta de visualização de pods.

Estatísticas do framework (estatísticas do TensorFlow/PyTorch)

As estatísticas do framework fornecem uma análise detalhada das estatísticas de PyTorch e XRT em dispositivos TPU e hosts de TPU.

Visualizador de traces

O visualizador de traces é uma ferramenta de análise de desempenho do Cloud TPU. O visualizador de traces usa o visualizador de criação de perfis de eventos de rastreamento do Chrome para que requer o uso do navegador Chrome.

O visualizador de traces contém uma linha do tempo que mostra:

• Durações das operações executadas pelo modelo do TensorFlow.
• qual parte do sistema (TPU ou máquina host) executou uma operação. Para o PyTorch/XLA geralmente, a máquina host trabalha principalmente na compilação/alocação/desalocação do buffer, enquanto a TPU executa o treinamento do modelo real.
• O visualizador de traces permite identificar problemas de desempenho no modelo e, em seguida, seguir as etapas para resolvê-los. Detalhando quais operações PyTorch / XLA levam mais tempo para serem executadas.

É possível adicionar diretamente traces para medir quanto tempo determinadas partes do modelo levam para serem executadas. Para isso, adicione anotações xp.Trace(NAME). Por exemplo, o trace a seguir mostra:

1. Gerado por anotações de usuário explícitas no código de modelo de test_profile_mp_mnist.py.
2. Operações do PyTorch executadas (pré-redução).
3. Nome do módulo HLO gerado automaticamente pelo PyTorch / XLA.
4. Operações de XLA executadas no dispositivo (fusível)

Para informações mais detalhadas, consulte a documentação genérica da TPU para o visualizador de traces, mas ignore as seções sobre pipeline de entrada e outras partes específicas do TensorFlow, já que elas não são relevantes na contexto deste documento.

Criação de perfil do cliente PyTorch / XLA

Assim como na criação de perfil da TPU enquanto a execução do modelo estava em andamento, agora você criará o perfil do lado do cliente PyTorch / XLA durante o treinamento. A principal ferramenta de monitoramento usada no lado do cliente é o visualizador de traces.

É preciso iniciar o servidor de criação de perfil no script de treinamento. Por exemplo, consulte, que pode ser consultado no TensorBoard para capturar um trace.

Para capturar traces de vários processos, cada processo pode iniciar servidores de perfil em portas diferentes (por exemplo, adicionando xm.get_ordinal() a um número de porta base) e, em seguida, fornecendo uma lista de localhost:port concatenado por ','. O TensorBoard não é compatível com a visualização de traces de vários processos ao mesmo tempo. Portanto, você verá diferentes menus suspensos de Host para cada processo.

O diagrama a seguir mostra um trace de amostra:

Da mesma forma que diferentes traces de namespace podem ser adicionados a traces da TPU, é possível usar a mesma API para adicioná-los aos traces do cliente (xp.Trace(NAME)). Como esse modelo é pequeno e está sendo usado com pequenas imagens MNIST, os tempos das etapas serão curtos e não necessariamente uniformes. Como exercício, tente adicionar os traces e inicializar um servidor do criador de perfil semelhante ao do nosso exemplo ao test_train_mp_imagenet.py --fake_data para receber traces do ResNet50.

Os traces têm metadados adicionais que podem ser inspecionados. Por exemplo, os traces TransferToServer e TransferFromServer, mostram o número exato de tensores que estão sendo enviados e recebidos e o tamanho total deles:

Para compilações de gráficos XLA, é possível ver o hash do gráfico que pode ser útil para diagnosticar problemas:

Além disso, em vez de criar perfis da IU do TensorBoard, também fornecemos a API para criar a criação de perfil programática da TPU e do cliente usando o PyTorch / XLA: xp.trace().

Análise automática de métricas

Nesta seção, você verá como usar o modo de depuração para detectar problemas de desempenho, como:

• Gráficos dinâmicos / compilações contínuas
• Compilação de gráficos muito lenta
• Execução de grafos muito lentas
• Transferências frequentes de CPU XLA→
• HBM do dispositivo repetido para hospedar RAM
• Desfragmentação de HBM repetida
• Operações aten:: desaceleradas

Antes de iniciar o treinamento, reverta as seguintes linhas em /usr/share/torch-xla-1.8/pytorch/xla/test/test_profile_mp_mnist.py:

Alteração:

   accuracy = train_mnist(flags, dynamic_graph=False, fetch_often=False)
   

Para:

   accuracy = train_mnist(flags, dynamic_graph=True, fetch_often=True)
   

Essas alterações causam artificialmente compilações e buscas de tensor. dynamic_graph=True altera artificialmente o tamanho do lote para cada etapa, fazendo com que os gráficos XLA diminuídos sejam diferentes em cada etapa e recompilação. fetch_often=True insere chamadas loss.item() em todas as etapas, resultando na busca de valores de tensor do dispositivo em cada etapa, diminuindo o desempenho.

Como executar um exemplo de script de treinamento:

   PT_XLA_DEBUG=1 python /usr/share/torch-xla-1.8/pytorch/xla/test/test_profile_mp_mnist.py --fake_data --num_cores=1
   

Durante a depuração, a prática recomendada é executá-la com --num_cores=1, porque ela simplifica o processo de depuração. Alguns dos resultados de amostra são assim:

Epoch 1 train begin 01:18:05
| Training Device=xla:1/0 Step=0 Loss=0.00000 Rate=1905.00 GlobalRate=1904.72 Time=01:18:05
pt-xla-profiler: TransferFromServerTime too frequent: 3 counts during 3 steps
pt-xla-profiler: TransferFromServerTime too frequent: 4 counts during 4 steps
pt-xla-profiler: TransferFromServerTime too frequent: 5 counts during 5 steps
pt-xla-profiler: TransferFromServerTime too frequent: 6 counts during 6 steps
pt-xla-profiler: TransferFromServerTime too frequent: 7 counts during 7 steps
pt-xla-profiler: TransferFromServerTime too frequent: 8 counts during 8 steps
pt-xla-profiler: TransferFromServerTime too frequent: 9 counts during 9 steps
pt-xla-profiler: TransferFromServerTime too frequent: 10 counts during 10 steps
pt-xla-profiler: CompileTime too frequent: 21 counts during 11 steps
pt-xla-profiler: TransferFromServerTime too frequent: 11 counts during 11 steps
pt-xla-profiler: CompileTime too frequent: 23 counts during 12 steps

As linhas com o prefixo pt-xla-profiler correspondem à saída da análise de métricas automáticas. Neste exemplo, é possível ver que TransferFromServerTime foi visto com muita frequência, uma por etapa. Isso se deve ao loop de treinamento que recupera o valor de loss.item() em todas as etapas. Além disso, você poderá ver um aviso CompileTime demais, quando quando o modelo precisar ser recompilado repetidamente devido a formas dinâmicas no gráfico. O snippet de código a seguir causa esse tipo de problema: test_profile_mp_mnist.py

    for step, (data, target) in enumerate(loader):
      if dynamic_graph:
        # The batch dimension is different every step.
        index = max(-step, -flags.batch_size + 1)  # non-empty
        data, target = data[:-index, :, :, :], target[:-index]
      ...
      if fetch_often:
        # Fetch tensor value from XLA:TPU to CPU every step.
        loss_i = loss.item()

Depois, pressione Ctrl^C para sair do script de treinamento. Quando terminar, você verá um resumo das operações não reduzidas. Observe que aten::_local_scalar_dense é uma operação especial que corresponde à recuperação de tensores XLA de volta ao contexto da CPU.

Neste relatório, você vê que há dois lugares principais em que a operação aten::_local_scalar_dense está sendo chamada, ambos correspondem ao código-fonte de loss.item():

• test/test_profile_mp_mnist.py:158
• test/test_profile_mp_mnist.py:61

pt-xla-profiler: ================================================================================
pt-xla-profiler: Unlowered Op usage summary (more of these ops, lower performance)
pt-xla-profiler: Note: _local_scalar_dense typically indicates CPU context access
pt-xla-profiler: --------------------------------------------------------------------------------
pt-xla-profiler: FRAME (count=27):
pt-xla-profiler: Unlowered Op: "_local_scalar_dense"
pt-xla-profiler: Python Frames:
pt-xla-profiler:   train_loop_fn (test/test_profile_mp_mnist.py:158)
pt-xla-profiler:   train_mnist (test/test_profile_mp_mnist.py:184)
pt-xla-profiler:   _mp_fn (test/test_profile_mp_mnist.py:206)
pt-xla-profiler:   _start_fn (/home/jysohn/git/jysohn23/pytorch/xla/torch_xla/distributed/xla_multiprocessing.py:323)
pt-xla-profiler:   spawn (/home/jysohn/git/jysohn23/pytorch/xla/torch_xla/distributed/xla_multiprocessing.py:386)
pt-xla-profiler:    (test/test_profile_mp_mnist.py:216)
pt-xla-profiler:
pt-xla-profiler:
pt-xla-profiler: FRAME (count=2):
pt-xla-profiler: Unlowered Op: "_local_scalar_dense"
pt-xla-profiler: Python Frames:
pt-xla-profiler:   _train_update (test/test_profile_mp_mnist.py:61)
pt-xla-profiler:    (/home/jysohn/git/jysohn23/pytorch/xla/torch_xla/core/xla_model.py:700)
pt-xla-profiler:   _run_step_closures (/home/jysohn/git/jysohn23/pytorch/xla/torch_xla/core/xla_model.py:709)
pt-xla-profiler:   mark_step (/home/jysohn/git/jysohn23/pytorch/xla/torch_xla/core/xla_model.py:723)
pt-xla-profiler:   __exit__ (/home/jysohn/git/jysohn23/pytorch/xla/torch_xla/debug/profiler.py:153)
pt-xla-profiler:   train_loop_fn (test/test_profile_mp_mnist.py:162)
pt-xla-profiler:   train_mnist (test/test_profile_mp_mnist.py:184)
pt-xla-profiler:   _mp_fn (test/test_profile_mp_mnist.py:206)
pt-xla-profiler:   _start_fn (/home/jysohn/git/jysohn23/pytorch/xla/torch_xla/distributed/xla_multiprocessing.py:323)
pt-xla-profiler:   spawn (/home/jysohn/git/jysohn23/pytorch/xla/torch_xla/distributed/xla_multiprocessing.py:386)
pt-xla-profiler:    (test/test_profile_mp_mnist.py:216)
pt-xla-profiler:
pt-xla-profiler:
pt-xla-profiler: ================================================================================

Agora, execute a análise de métricas automáticas no script abaixo, que contém uma operação não reduzida (a partir da nossa versão 1.8): _ctc_loss:

PT_XLA_DEBUG=1 python <<EOF
import torch
import torch_xla.core.xla_model as xm

dev = xm.xla_device()
t = torch.randn(50, 16, 20).log_softmax(2).to(dev)
target = torch.randint(low=1, high=20, size=(16, 30), dtype=torch.long).to(dev)
input_lengths = torch.full(size=(16,), fill_value=50, dtype=torch.long).to(dev)
target_lengths = torch.randint(low=10, high=30, size=(16,), dtype=torch.long).to(dev)

for _ in range(10):
  loss = torch.nn.CTCLoss()(t, target, input_lengths, target_lengths)
  xm.mark_step()
EOF

A execução do script acima com PT_XLA_DEBUG=1 precisa ter uma aparência semelhante à seguinte:

…
pt-xla-profiler: TransferFromServerTime too frequent: 30 counts during 10 steps
pt-xla-profiler: Op(s) not lowered: aten::_ctc_loss,  Please open a GitHub issue with the above op lowering requests.
pt-xla-profiler: ================================================================================
pt-xla-profiler: Unlowered Op usage summary (more of these ops, lower performance)
pt-xla-profiler: Note: _local_scalar_dense typically indicates CPU context access
pt-xla-profiler: --------------------------------------------------------------------------------
pt-xla-profiler: FRAME (count=10):
pt-xla-profiler: Unlowered Op: "_ctc_loss"
pt-xla-profiler: Python Frames:
pt-xla-profiler:   ctc_loss (/anaconda3/envs/torch-xla-1.8/lib/python3.6/site-packages/torch/nn/functional.py:2305)
pt-xla-profiler:   forward (/anaconda3/envs/torch-xla-1.8/lib/python3.6/site-packages/torch/nn/modules/loss.py:1593)
pt-xla-profiler:   _call_impl (/anaconda3/envs/torch-xla-1.8/lib/python3.6/site-packages/torch/nn/modules/module.py:889)
pt-xla-profiler:    (:11)
pt-xla-profiler:
pt-xla-profiler:
pt-xla-profiler: ================================================================================

O analisador de métricas automáticas mostra que a 11a linha do STDIN está causando essa operação não reduzida (ou seja, a linha com torch.nn.CTCLoss()). Atualmente, a operação ctc_loss não foi reduzida. , por isso você vê o relatório acima. Também é possível ver alguns avisos para TransferFromServerTime, já que os tensores estão inicialmente em XLA:TPU antes da execução. Porém, como a operação não é reduzida, primeiro você precisa transferir os tensores XLA de volta para CPU, execute a operação aten:: na CPU e a transfira de volta.

Caso você queira gravar a saída pt-xla-profiler em um arquivo, defina PT_XLA_DEBUG=1 e PT_XLA_DEBUG_FILE=$PATH_TO_FILE.

Limpeza

Saia da VM e exclua o bucket da TPU, da VM e do Cloud Storage executando os comandos a seguir:

(vm)$ exit

gcloud compute instances delete profiler-tutorial-vm \
  --zone=${ZONE} \
  --project=${PROJECT_ID}

gcloud compute tpus delete profiler-tutorial-tpu \
  --zone=${ZONE} \
  --project=${PROJECT_ID} \
  --async

gsutil rm -fr gs://${BUCKET_NAME}

VM de TPU + criação de perfil do PyTorch/XLA

Use esta seção para criar o perfil do PyTorch/XLA usando a arquitetura de VM da TPU.

Variáveis de ambiente de exportação

1. Crie variáveis para o ID do projeto e a zona a ser usada nos recursos da TPU.

   export PROJECT_ID=PROJECT_ID
   export ZONE=ZONE
   

Criar uma Cloud TPU

Consulte a Guia do usuário de VMs da TPU e, após a configuração, crie uma VM v3-8 de TPU, que vem comtorch ,torch_xla, torchvision etensorboard pré-instalado.

1. Crie um recurso de TPU.

   gcloud alpha compute tpus tpu-vm create profiler-tutorial-tpu-vm \
    --project=${PROJECT_ID} \
    --zone=${ZONE} \
    --version=v2-alpha \
    --accelerator-type=v3-8
   

Inicialização do servidor do TensorBoard

1. Conecte-se por SSH à VM, instale o tensorboard-plugin-profile e inicie um servidor do tensorboard.

     gcloud alpha compute tpus tpu-vm ssh profiler-tutorial-tpu-vm \
      --project ${PROJECT_ID} \
      --zone ${ZONE} \
      --ssh-flag="-4 -L 9001:localhost:9001"
   

     pip3 install tf-nightly==2.6.0.dev20210511 tb-nightly==2.6.0a20210512 tbp-nightly==2.5.0a20210511
   

     tensorboard --logdir ./tensorboard --port 9001
   

Ao visualizar a saída do TensorBoard em http://localhost:9001 na máquina local, você verá algo assim:

Se você visualizar a saída do TensorBoard em http://localhost:9001, também poderá acessar a página de perfil acima. Basta selecionar a opção PROFILE no menu suspenso do canto superior direito, ao lado do botão UPLOAD:

Criar o perfil do modelo

Em uma nova janela de terminal no ambiente de desenvolvimento, exporte as mesmas variáveis de ambiente mostradas acima e ssh para a VM de TPU:

1. Na nova janela do terminal, exporte as variáveis de zona e ID do projeto novamente, já que isso está em um novo shell.

   export PROJECT_ID=PROJECT_ID
   export ZONE=ZONE
   

2. ssh na VM:

     gcloud alpha compute tpus tpu-vm ssh profiler-tutorial-tpu-vm \
      --project ${PROJECT_ID} \
      --zone ${ZONE}
   

3. Clone o repositório PyTorch/XLA e execute o teste e2e:

     git clone -b r1.8 https://github.com/pytorch/xla
     export XRT_TPU_CONFIG="localservice;0;localhost:51011"
     python3 xla/test/test_profiler.py  # takes <1 min
   >

4. Antes de iniciar o treinamento, edite as seguintes linhas em xla/test/test_profile_mp_mnist.py:

   Alteração:

       accuracy = train_mnist(flags, dynamic_graph=True, fetch_often=True)
   

   Para:

       accuracy = train_mnist(flags, dynamic_graph=False, fetch_often=False)
   

   Os dois argumentos acima em train_mnist causam artificialmente gráficos e buscas de tensores, que serão explorados posteriormente na seção "Análise automática". Por enquanto, como você está apenas analisando a TPU, o exemplo a seguir é executado com desempenho nominal.

5. Inicie uma execução de treinamento:

    XLA_HLO_DEBUG=1 python3 xla/test/test_profile_mp_mnist.py --num_epochs 1000 --fake_data
   

Criação de perfis de TPU + cliente

Quando o treinamento estiver em execução, veja a saída do TensorBoard em http://localhost:9001 e capture um perfil usando as seguintes instruções:

Você verá a seguinte página recarregada:

Atualmente, na configuração da VM da TPU, somente a ferramenta visualizador de traces é selecionada. Portanto, na lista suspensa Ferramentas, selecione trace_viewer e inspecione os traces. Na configuração da VM da TPU, você vê os traces do cliente e do dispositivo da TPU em uma visualização completa:

Limpeza

1. Saia da VM e exclua o bucket da TPU, da VM e do Cloud Storage executando os seguintes comandos:

   (vm)$ exit
   

Exclua a VM da TPU que você criou:

1. Exclua os recursos do Cloud TPU e do Compute Engine.

   $ gcloud alpha compute tpus tpu-vm delete profiler-tutorial-tpu-vm \
     --project ${PROJECT_ID} --zone=${ZONE}
   

2. Verifique se os recursos foram excluídos executando o seguinte comando. A exclusão pode levar vários minutos. Uma resposta como esta indica que suas instâncias foram excluídas com sucesso.

   $ gcloud alpha compute tpus tpu-vm list --project ${PROJECT_ID} --zone=${ZONE}
   

   Listed 0 items.