Criar o perfil do modelo em VMs do Cloud TPU

A criação de perfil é uma das principais ferramentas para otimizar o desempenho dos modelos no Cloud TPU. A principal ferramenta de criação de perfil é o XProf, que está disponível no repositório do GitHub do OpenXLA/XProf. O XProf permite a criação de perfil de todos os frameworks baseados em XLA, como JAX, PyTorch/XLA e TensorFlow/Keras.

Como capturar perfis

Para usar o XProf, primeiro capture um perfil do modelo. Há duas maneiras de fazer isso:

  1. Captura programática
  2. Captura on demand (também chamada de captura manual)

Com a captura programática, é necessário anotar o código do modelo para especificar onde nele você quer capturar perfis. Normalmente, você captura um perfil para algumas etapas de treinamento ou cria um perfil de um bloco específico de código no modelo. Há diferentes maneiras de capturar traces nos frameworks JAX, PyTorch/XLA e TensorFlow: início ou interrupção de trace com base em API ou em gerenciador de contexto. Em frameworks de nível mais alto, como o MaxText, é possível ativar a captura de perfil ativando uma flag profiler=xplane ao executar o MaxText.

Use a captura de perfil on demand quando quiser capturar perfis on demand ou quando não tiver a captura de perfil programática ativada. Isso pode ser útil quando você encontra um problema com as métricas do modelo durante a execução e quer capturar perfis por algum período para diagnosticar o problema.

Para ativar a captura de perfil on demand, ainda é necessário iniciar o servidor XProf no código. Ao ativar profiler.start_server, um servidor XProf será iniciado na carga de trabalho de ML para detectar o gatilho da captura on demand e começar a capturar perfis. É possível acionar a captura de perfil on demand usando a interface do TensorBoard ou a CLI com a ferramenta XProfiler.

Saiba como ativar as capturas programática e on demand para diferentes frameworks em:

Para as capturas de perfil programática e on demand, é necessário especificar o local em que você quer armazenar os perfis. É possível armazenar os perfis em um diretório na VM de TPU com um caminho semelhante a /profiles/run-name ou usar o Cloud Storage com um caminho semelhante a gs://bucket-name/run-name/.

Com os caminhos mencionados acima, seus perfis serão capturados em:

/profiles/run-name/plugins/profile/session1/

ou

gs://bucket-name/run-name/plugins/profile/session1/.

Uma sessão representa uma captura de perfis durante uma execução de modelo. Por exemplo, se você capturar perfis em uma execução de treinamento das etapas 1 a 3 e, em seguida, das etapas 8 a 10, esses perfis serão da mesma execução, mas a primeira captura será session1 e a segunda será session2.

As diferentes sessões serão marcadas com carimbos de data diferentes em cada execução. É possível capturar perfis em diferentes sessões de maneira programática, on demand ou usando uma combinação das duas. Isso é útil quando você quer comparar perfis de diferentes partes do treinamento de modelo. Por exemplo, para comparar perfis do início e do fim do treinamento.

Como visualizar perfis com o TensorBoard

Para exibir os perfis do XProf no TensorBoard, instale o plug-in tensorboard-plugin-profile.

Para exibir os perfis capturados no Google Cloud, recomendamos usar o TensorBoard em uma VM de TPU ou a ferramenta XProfiler para hospedar o TensorBoard em uma VM ou um pod do GKE.

Se você tiver capturado os perfis na VM de TPU, confira esses perfis usando o TensorBoard na VM.

Para instalar o XProf e o TensorBoard em uma VM de TPU, execute estes comandos:

pip install tensorboard_plugin_profile tensorboard
pip install xprof

Para exibir os perfis capturados com o TensorBoard, execute o seguinte:

$ tensorboard --logdir=profiles/run-name

ou

$ tensorboard --logdir=profiles

Para carregar vários perfis de várias execuções, aponte o TensorBoard para o diretório raiz que contém todas as execuções e perfis, em vez de para o subdiretório com perfis de uma execução específica.

Biblioteca Cloud-Diagnostics-XProf para visualizar perfis do Tensorboard no Google Cloud

A cloud-diagnostics-xprof (também conhecida como a biblioteca do XProfiler) facilita a hospedagem do TensorBoard e a exibição de perfis no Google Cloud. Para mais informações, consulte o repositório do GitHub da cloud-diagnostics-xprof.

O repositório da cloud-diagnostics-xprof oferece as seguintes melhorias em relação à execução local do XProf e do TensorBoard:

  • Configuração e empacotamento das dependências do XProf e do TensorBoard.
  • Armazenamento dos perfis no Cloud Storage, o que pode ser útil para retenção de longo prazo e análise pós-execução. Os perfis locais capturados serão excluídos depois que o pesquisador concluir a execução.
  • Carregamento rápido de vários perfis grandes ao provisionar o TensorBoard em uma VM do Compute Engine ou em um pod do GKE, com a opção de mudar o tipo de máquina com base nas necessidades do usuário para velocidade de carregamento e custo.
  • Criação de um link para compartilhar perfis e colaborar com membros da equipe e engenheiros do Google.
  • Criação de perfil on demand facilitada para cargas de trabalho no GKE e no Compute Engine para escolher qualquer host que execute a carga de trabalho para capturar perfis.

Antes de usar a biblioteca do XProfiler, capture perfis para o código da carga de trabalho de maneira programática ou iniciando o servidor de perfil para fazer a captura on demand mais tarde.

Para configurar a biblioteca do XProfiler, configure a CLI do Google Cloud e um ambiente virtual do Python. Depois, execute:

pip install cloud-diagnostics-xprof

Isso vai instalar todas as dependências necessárias do XProf e do TensorBoard.

Em seguida, para criar uma VM ou um pod do GKE para hospedar o TensorBoard, execute este comando:

xprofiler create -z $ZONE -l $GCS_PATH

ou

xprofiler create --GKE -z $ZONE -l $GCS_PATH

Substitua $ZONE por qualquer zona e $GCS_PATH pelo caminho dos traces de perfil. É possível especificar o diretório raiz que contém todos os traces de perfil de várias execuções ou um conjunto específico de traces de perfil de uma execução.

Por exemplo, se você capturar perfis nos seguintes diretórios:

gs://bucket-name/run1/plugins/profile/session1/<profile.xplane.pb>
gs://bucket-name/run1/plugins/profile/session2/<profile.xplane.pb>
gs://bucket-name/run2/plugins/profile/session1/<profile.xplane.pb>

Será possível definir o caminho do GCS para o diretório raiz (gs://bucket-name):

xprofiler create -z $ZONE -l gs://bucket-name/

Na interface do TensorBoard, você vai encontrar todos os perfis de todas as execuções e sessões como run1/session1, run1/session2 e run2/session1.

Por padrão, o xprofiler create cria uma VM do Compute Engine, especificamente uma c4-highmem-8. É possível mudar o tipo de máquina com a flag -m. Se você quiser criar uma instância do TensorBoard em um pod do GKE em vez de em uma VM do Compute Engine, transmita a flag –GKE para xprofiler create. Hospedar uma instância do TensorBoard em um pod do GKE pode facilitar o gerenciamento do TensorBoard com o restante da carga de trabalho implantada no GKE.

A VM ou o pod do GKE que hospeda o TensorBoard torna o carregamento de vários perfis grandes muito mais rápido no Google Cloud do que no TensorBoard hospedado localmente. Com base no comparativo, serão gerados perfis com cerca de 1 GB de carga em poucos minutos para a primeira carga usando a VM c4-highmem-8 padrão. Além disso, você pode escolher o tipo de máquina ideal com base nas suas necessidades de desempenho e custo.

Depois de executar xprofiler create, você vai receber uma saída como esta:

Instance for gs://<bucket> has been created.
You can access it via the following:
1. https://<id>-dot-us-<region>.notebooks.googleusercontent.com.
2. xprofiler connect -z <zone> -l gs://bucket-name -m ssh
Instance is hosted at xprof-97db0ee6-93f6-46d4-b4c4-6d024b34a99f VM.

A primeira opção é um link em que você pode clicar para conferir os perfis do XProf em um TensorBoard. Esse link permite compartilhar perfis com a equipe e também com engenheiros do Google que ajudam na otimização do desempenho no Google Cloud.

Você controla quem tem acesso ao link com base nas permissões definidas para o bucket do Cloud Storage que armazena os dados do perfil. Se o link não funcionar por algum motivo, será possível usar o comando xprofiler connect para se conectar por SSH à instância do TensorBoard e conferir os perfis.

Se você ativou o servidor do Cloud Profiler no código da carga de trabalho, um processo que é necessário para as capturas programáticas e on demand, e quer capturar perfis on demand, há duas maneiras de fazer isso:

a. Clique no botão Capturar perfil no TensorBoard. Agora escolha qualquer host de dispositivo em que a carga de trabalho é executada para criar o perfil. Disponibilizamos a captura on demand para cargas de trabalho executadas no GKE ou no Compute Engine.

b. Use o comando de captura do XProfiler e forneça informações como zona, bucket do Cloud Storage, framework, nomes de pod/VM host e duração da captura em milissegundos. Essas são as mesmas informações que você precisa inserir na interface do TensorBoard.

Para saber como usar a biblioteca cloud-diagnostics-xprof, confira esta página do GitHub.