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 é chamada de XProf, que está disponível no repositório OpenXLA/XProf do GitHub. O XProf oferece suporte à criação de perfil de todos os frameworks baseados em XLA, incluindo JAX, Pytorch XLA e Tensorflow/Keras.

Como capturar perfis

Para usar o XProf, primeiro capture um perfil do seu modelo. Há duas maneiras de capturar perfis:

  1. Captura programática
  2. Captura sob demanda (também conhecida como captura manual)

Com a captura programática, você precisa anotar o código do modelo para especificar onde no código 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 rastreamentos nos diferentes frameworks JAX, Pytorch XLA e TensorFlow, seja com base em API ou parar rastreamento ou com base em gerenciador de contexto. Em frameworks de nível mais alto, como o MaxText, é possível ativar a captura de perfil apenas ativando uma flag profiler=xplane ao executar o MaxText.

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

Para ativar a captura de perfil sob demanda, ainda é necessário iniciar o servidor XProf no seu código. Ao ativar o profiler.start_server, um servidor XProf será iniciado na sua carga de trabalho de ML, que detecta o gatilho de captura sob demanda para começar a capturar perfis. É possível acionar a captura de perfil sob demanda pela interface do TensorBoard ou pela CLI com a ferramenta XProfiler.

Saiba como ativar a captura programática e sob demanda para diferentes frameworks em:

Para a captura de perfil programática e sob demanda, é necessário especificar o local onde você quer armazenar os perfis capturados. É possível armazenar os perfis em um diretório na VM da 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 do modelo. Por exemplo, se você capturar perfis em uma execução de treinamento das etapas 1 a 3 e depois das etapas 8 a 10, esses serão perfis da mesma execução, mas a primeira captura das etapas 1 a 3 será session1 e a segunda das etapas 8 a 10 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 forma programática, sob demanda ou uma combinação das duas. Isso é útil quando você quer comparar perfis de diferentes partes do treinamento de modelo. Por exemplo, comparar perfis do início e do fim do treinamento.

Como visualizar perfis com o TensorBoard

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

Para visualizar 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ê capturou seus perfis na VM de TPU, é possível conferir os perfis capturados usando o TensorBoard nessa VM de TPU.

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 ver os perfis capturados com o TensorBoard, execute:

$ 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 do subdiretório com perfis de uma execução específica.

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

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

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

  • Configuração e empacotamento das dependências do XProf e do TensorBoard
  • Armazene seus 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 perfis grandes e múltiplos ao provisionar o TensorBoard em uma VM do Compute Engine ou 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.
  • Criar um link para compartilhar perfis e colaborar com membros da equipe e engenheiros do Google
  • Perfilamento on demand mais fácil de cargas de trabalho no GKE e no Compute Engine para escolher qualquer host que execute sua carga de trabalho e capturar perfis

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

Para configurar a biblioteca XProfiler, você precisa ter configurado a Google Cloud CLI e um ambiente virtual do Python. Depois, basta executar:

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 o comando a seguir:

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 rastreamentos de perfil. É possível especificar o diretório raiz que contém todos os traces de perfil para várias execuções ou um conjunto específico de traces de perfil para 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>

Você pode 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, run2/session1.

Por padrão, o xprofiler create cria uma VM do Compute Engine, especificamente uma VM 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 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 perfis grandes e múltiplos muito mais rápido no Google Cloud do que no TensorBoard hospedado localmente. Com base nos nossos comparativos de mercado, você vai ver perfis da ordem de 1 GB de carga em poucos minutos para o primeiro carregamento usando a VM c4-highmem-8 padrão. Além disso, você pode escolher o tipo de máquina certo com base nas suas necessidades de desempenho e custo.

Depois de executar xprofiler create, você vai ver uma saída semelhante a 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 ver seus perfis do XProf em um TensorBoard. Esse link permite compartilhar perfis com sua equipe e também com engenheiros do Google que ajudam na otimização de performance 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 seu perfil. Se o link não funcionar por algum motivo, também oferecemos uma maneira de se conectar à instância do TensorBoard usando SSH para ver seus perfis com o comando xprofiler connect.

Se você ativou o servidor do Cloud Profiler no código da carga de trabalho (necessário para capturas programáticas e sob demanda) e quer fazer o perfil sob demanda, há duas maneiras de fazer isso:

a. Clique no botão Capturar perfil no TensorBoard. Agora você pode escolher qualquer host de dispositivo em que sua carga de trabalho está sendo executada para ser analisado. Oferecemos suporte à captura sob demanda para cargas de trabalho executadas no GKE ou no Compute Engine.

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

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