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.
Para mais informações sobre o desempenho do Cloud TPU com o PyTorch, consulte as seguintes postagens do blog:
- Como escalonar cargas de trabalho de aprendizado profundo PyTorch
- 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
- [Desempenho lento do tensor com PyTorch XLA](https://pytorch.org/blog/understanding-lazytensor-system-performance-with-pytorch-xla-on-cloud-tpu/
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.
Nó da TPU + criação de perfil PyTorch/XLA
Criar recursos relacionados ao Cloud TPU
Criar e inicializar os recursos necessários relacionados ao Cloud TPU.
Criar 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
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
Criar 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
Criar um bucket do Cloud Storage
Primeiro, instale a CLI
gsutil
se ainda não a tiver instalado: instruções de instalação.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}
em que:
- REGION é a região em que você criou o
Cloud TPU, por exemplo,
europe-west4
.
- REGION é a região em que você criou o
Cloud TPU, por exemplo,
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:
service-PROJECT_NUMBER@cloud-tpu.iam.gserviceaccount.com
Por exemplo: service-164006649440@cloud-tpu.iam.gserviceaccount.com
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=service-ACCOUNT_NUMBER@cloud-tpu.iam.gserviceaccount.com gsutil acl ch -u $SERVICE_ACCOUNT:READER gs://${BUCKET_NAME} gsutil acl ch -u $SERVICE_ACCOUNT:WRITER gs://${BUCKET_NAME}
Configurar o TensorBoard
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"
Criar 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
Para testar sua instalação, execute o servidor do TensorBoard na VM do Compute Engine e tente se conectar ao servidor em http://localhost:9001/#profile na 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 na máquina local, você verá algo parecido com o seguinte:
Se você acessar http://localhost:9001, também poderá acessar a página de perfil acima. Basta selecionar a opção PERFIL na lista suspensa no canto superior direito ao lado do botão ENVIAR:
Criar o perfil do modelo
Para manter o servidor do TensorBoard ativo, na máquina local,
inicie uma nova janela de terminal e ssh
na VM do GCE novamente,
(desta vez sem usar a opção -L
de encaminhamento de portas).
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
ssh
na VM:gcloud compute ssh profiler-tutorial-vm \ --project ${PROJECT_ID} \ --zone ${ZONE}
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"
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 >
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
:Mudar:
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 dinâmicos e buscas de tensores, que serão explorados posteriormente na seção "Análise de métricas automáticas". Por enquanto, como você está apenas analisando a TPU, o exemplo a seguir é executado com desempenho nominal.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 da 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 no 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 identificadores 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 da HBM de execução do módulo. Isso é ordenado por tamanho de alocação, tamanho de execução do programa e tamanho do padding.
Cada uma das alocações de buffer pode ser inspecionada ainda 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 op do XLA
O perfil de op é uma ferramenta do Cloud TPU que exibe as estatísticas de desempenho das operações do XLA executadas durante um período de criação de perfil. O perfil de operações mostra:
- O uso que o aplicativo faz do Cloud TPU como uma porcentagem do tempo gasto em operações por categoria e de utilização de FLOPS de TPU;
- as operações mais demoradas, Essas operações são alvos potenciais de otimização.
- Detalhes de operações individuais, incluindo forma, padding e expressões que usam a operação.
É possível usar o perfil de op para identificar bons alvos de otimização. Por exemplo, se o modelo atingir apenas 5% de FLOPS de pico da TPU, será possível usar a ferramenta para identificar quais operações do XLA são mais demoradas na execução e quantos FLOPS de TPU elas consomem.
Descrição de cada item:
- Seção de visão geral. Mostra a utilização do Cloud TPU e fornece sugestões para otimização.
- 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 elas são classificadas.
- Tabela op. Uma tabela que lista as principais categorias de operações do TensorFlow associadas às operações do XLA. Essas operações são classificadas por porcentagem de uso do Cloud TPU.
- Cartões de detalhes de operações Detalhes sobre a operação exibidos quando você passa o cursor sobre uma operação da tabela. Inclui a utilização de FLOPS, a expressão em que a operação é usada e o layout da operação (fit).
Visualizador de pod
Consulte Ferramentas da 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 operações PyTorch e XRT em execução em dispositivos 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 trace do Chrome para que use o navegador Chrome.
O visualizador de traces contém uma linha do tempo que mostra:
- As durações das operações que foram executadas pelo seu 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 e alocação/desalocação de buffer, enquanto a TPU executa o treinamento do modelo real.
- Com o visualizador de traces, é possível identificar problemas de desempenho do modelo e tomar providências para solucioná-los. Detalhando, é possível identificar 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 do xp.Trace(NAME)
.
Por exemplo, o trace a seguir é exibido:
- Gerado por anotações de usuário explícitas no código de modelo de test_profile_mp_mnist.py.
- Operações do PyTorch executadas (pré-redução).
- Nome do módulo HLO gerado automaticamente pelo PyTorch/XLA.
- Operações de XLA executadas no dispositivo (fundido).
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 no 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 criação 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 ','. Como o TensorBoard não é compatível com a visualização de traces de
vários processos ao mesmo tempo, você verá menus suspensos de Host diferentes
para cada um deles.
O diagrama a seguir mostra um trace de amostra:
Da mesma forma que diferentes traces de namespace podem ser adicionados a
traces de TPU, é possível usar a mesma API para adicioná-los aos traces do lado
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, você pode tentar 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 de 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 a criação de perfil programática da TPU e
do cliente do 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 gráficos muito lenta
- Transferências frequentes XLA→CPU
- HBM do dispositivo repetido para hospedar a troca de RAM
- Desfragmentação de HBM repetida
- Operações
aten::
não reduzidas
Antes de iniciar o treinamento, reverta as seguintes linhas em
/usr/share/torch-xla-1.8/pytorch/xla/test/test_profile_mp_mnist.py
:
Mudar:
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
muda 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, reduzindo 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
,
pois ele 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 vez 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 de forma muito frequente, 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 (desde 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 a saída 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 11ª 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 a CPU, executar a
operação aten::
na CPU e transferi-la 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 de TPU.
Exportar variáveis de ambiente
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 o
Guia do usuário de VMs da TPU
e, após a configuração, crie uma VM de TPU v3-8, que vem com torch
,
torch_xla
,
torchvision
e tensorboard
pré-instalados.
Criar 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
Conecte-se por SSH à VM, instale
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 no 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:
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
ssh
na VM:gcloud alpha compute tpus tpu-vm ssh profiler-tutorial-tpu-vm \ --project ${PROJECT_ID} \ --zone ${ZONE}
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 >
Antes de iniciar o treinamento, edite as seguintes linhas em
xla/test/test_profile_mp_mnist.py
:Mudar:
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 dinâmicos e buscas de tensores, que serão explorados posteriormente na seção "Análise de métricas automáticas". Por enquanto, como você está apenas analisando a TPU, o exemplo a seguir é executado com desempenho nominal.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 de TPU, somente a ferramenta visualizador
de traces está selecionada. Portanto, na lista suspensa Ferramentas,
selecione trace_viewer
e inspecione os traces. Na configuração
da VM de TPU, você vê os traces do lado do "cliente" e do lado do dispositivo da TPU
em uma visualização completa:
Limpeza
Saia da VM e exclua o bucket da TPU, da VM e do Cloud Storage executando os comandos a seguir:
(vm)$ exit
Exclua a VM da TPU que você criou:
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}
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.