Criar perfil de cargas de trabalho PyTorch XLA
O perfil é uma maneira de analisar e melhorar o desempenho dos modelos. Embora há muito mais a ser feito. Às vezes, pode ser útil pensar na criação de perfil como operações e partes do código que são executadas em dispositivos (TPUs) e hosts (CPUs). Este guia fornece uma visão geral rápida de como criar o perfil do código para treinamento ou inferência. Para mais informações sobre como analisar os perfis gerados, consulte nos guias a seguir.
- 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
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 o 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 do 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 padrão é mostrado no arquivo mas você também pode usar uma da lista de configurações compatíveis.
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
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
É possível capturar um perfil manualmente usando capture_profile.py
ou
de forma programática, a partir do script de treinamento,
APIs do torch_xla.debug.profiler
.
Como iniciar o servidor de perfil
Para capturar um perfil, um servidor de perfil deve estar em execução no
script de treinamento. Inicie um servidor com um número de porta de sua escolha, por exemplo
9012
, como mostrado no comando 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
.
Agora é possível capturar perfis conforme descrito na próxima seção. O script cria perfis de tudo o que acontece em um dispositivo TPU.
Adicionar rastros
Se você também quiser
para criar o perfil de operações na máquina host, adicione xp.StepTrace
ou
xp.Trace
no código. Essas funções rastreiam o código Python no
máquina host.
Pense nisso como a medição de quanto tempo leva para executar
no host (CPU) antes de passar o "gráfico" ao dispositivo TPU. Portanto, é muito útil para analisar a sobrecarga de rastreamento. Você pode
adicionar isso ao 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 unir partes individuais do código com
with xp.Trace('loss'):
loss = ...
Se estiver usando a iluminação, você pode pular a adição de rastros, já que isso é feito automaticamente. em algumas partes do código. No entanto, se quiser adicionar mais traces, bem-vindos a inseri-los dentro do loop de treinamento.
Você poderá capturar a atividade do dispositivo após a compilação inicial. wait 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 a captura rápida de um perfil. Você pode fazer isso copiando o
capturar arquivo de perfil
diretamente na VM da TPU. O comando a seguir o 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 geração de registros
~/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 ao acionar um script, você pode configure seu script de treinamento para acionar automaticamente um perfil usando o método torch_xla.debug.profiler.trace_detached. no script de treinamento.
Por exemplo, para capturar automaticamente um perfil em um período e uma etapa específicos,
é possível configurar o script de treinamento para consumir PROFILE_STEP
,
PROFILE_EPOCH
e PROFILE_LOGDIR
ambiente
variáveis:
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 salvará os arquivos .xplane.pb
no diretório especificado pelo
PROFILE_LOGDIR
.
Análise no TensorBoard
Para analisar ainda mais os perfis, use o TensorBoard.
com o plug-in do TensorBoard da TPU
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 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 sua máquina remota, instale os pacotes necessários e inicie o TensorBoard
supondo que você tenha perfis nessa máquina em ~/profiles/
. Se você armazenou
os perfis em outro diretório ou bucket do Cloud Storage,
especificar caminhos corretamente, por exemplo, gs://your_bucket_name/profiles
;
(vm)$ pip install tensorflow-cpu tensorboard-plugin-profile
(vm)$ tensorboard --logdir ~/profiles/ --port 6006
(vm)$ pip uninstall tensorflow tf-nightly tensorboard tb-nightly tbp-nightly
Como executar o TensorBoard
No seu navegador local, acesse:
http://localhost:6006/
e escolha PROFILE
no menu suspenso para carregar seus perfis.
Consulte Ferramentas de TPU para informações sobre a as ferramentas do TensorBoard e como interpretar a saída.