Crie perfis do seu modelo em VMs do Cloud TPU

A criação de perfis é uma das principais ferramentas para otimizar o desempenho dos seus modelos na Cloud TPU. A ferramenta de criação de perfis principal chama-se XProf e está disponível no repositório do GitHub OpenXLA/XProf. O XProf suporta a criação de perfis de todas as frameworks baseadas em XLA, incluindo JAX, Pytorch XLA e Tensorflow/Keras.

Capturar perfis

Para usar o XProf, primeiro tem de capturar um perfil do seu modelo. Existem duas formas de capturar perfis:

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

Com a captura programática, tem de anotar o código do modelo para especificar onde no código quer capturar perfis. Normalmente, captura um perfil para alguns passos de preparação ou cria um perfil de um bloco de código específico no seu modelo. Existem diferentes formas de capturar rastreios nos diferentes frameworks JAX, Pytorch XLA e TensorFlow, quer seja com base no início/paragem do rastreio baseado na API ou no gestor de contexto. Em frameworks de nível superior, como o MaxText, pode ativar a captura de perfis ativando apenas uma flag profiler=xplane quando executa o MaxText.

Pode usar a captura de perfis a pedido quando quiser capturar perfis de forma ad hoc ou quando não ativou a captura de perfis programática. Isto pode ser útil quando deteta um problema com as métricas do modelo durante a execução e quer capturar perfis nesse instante durante algum tempo para diagnosticar o problema.

Para ativar a captura de perfis a pedido, continua a ter de iniciar o servidor XProf no seu código. A ativação do profiler.start_server inicia um servidor XProf na sua carga de trabalho de ML que ouve o acionador de captura a pedido para iniciar a captura de perfis. Pode acionar a captura de perfis a pedido através da IU do TensorBoard ou da CLI com a ferramenta XProfiler.

Pode saber como ativar a captura programática e a pedido para diferentes frameworks em:

Para a captura de perfis programática e a pedido, tem de especificar a localização onde quer armazenar os perfis capturados. Pode optar por armazenar os perfis num 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, os seus perfis sã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 seu modelo. Por exemplo, se capturar perfis numa execução de formação dos passos 1 a 3 e, posteriormente, capturar perfis dos passos 8 a 10, estes são perfis da mesma execução, mas a primeira captura dos passos 1 a 3 é session1 e a segunda captura dos passos 8 a 10 é session2.

As diferentes sessões são indicadas com datas/horas diferentes em cada execução. Pode capturar perfis em diferentes sessões de forma programática ou a pedido, ou uma combinação de ambas. Isto é útil quando quer comparar perfis de diferentes partes da preparação do modelo. Por exemplo, comparar perfis desde o início da preparação com perfis do final da preparação.

Visualizar perfis com o TensorBoard

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

Para visualizar os perfis capturados no Google Cloud, recomendamos que use o TensorBoard numa VM de TPU ou a ferramenta XProfiler para alojar o TensorBoard numa VM ou num pod do GKE.

Se capturou os seus perfis na VM da TPU, pode vê-los com o TensorBoard nessa VM da TPU.

Para instalar o XProf e o TensorBoard numa VM de TPU, pode executar estes comandos:

pip install tensorboard_plugin_profile tensorboard
pip install xprof

Para ver os perfis capturados com o TensorBoard, execute o seguinte comando:

$ tensorboard --logdir=profiles/run-name

ou

$ tensorboard --logdir=profiles

Para carregar vários perfis de várias execuções, direcione o TensorBoard para o diretório raiz que contém todas as suas execuções e perfis, em vez de para o subdiretório que tem 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 biblioteca XProfiler) facilita o alojamento do TensorBoard e a visualização de perfis no Google Cloud. Para mais informações, consulte o repositório do GitHub cloud-diagnostics-xprof.

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

  • Configuração e embalagem das dependências do XProf e do TensorBoard
  • Armazene os seus perfis no Cloud Storage, o que pode ser útil para a retenção a longo prazo e a análise pós-execução (os perfis locais capturados são eliminados depois de o investigador terminar a execução)
  • Carregamento rápido de perfis grandes e vários perfis através do aprovisionamento do TensorBoard na VM do Compute Engine ou no pod do GKE, com a opção de alterar o tipo de máquina com base nas necessidades do utilizador para a velocidade de carregamento e o custo
  • Crie um link para partilhar perfis e colaborar com membros da equipa e engenheiros da Google
  • Criação de perfis a pedido mais fácil de cargas de trabalho no GKE e no Compute Engine para escolher qualquer anfitrião que execute a sua carga de trabalho para capturar perfis

Antes de usar a biblioteca XProfiler, tem de capturar perfis para o código da carga de trabalho de forma programática ou iniciando o servidor de perfis, para poder fazer a captura de perfis a pedido mais tarde.

Para configurar a biblioteca XProfiler, tem de ter configurado a Google Cloud CLI e um ambiente virtual Python. Em seguida, basta executar:

pip install cloud-diagnostics-xprof

Esta ação instala todas as dependências necessárias do XProf e do TensorBoard.

Em seguida, para criar uma VM ou um pod do GKE para alojar o TensorBoard, execute o seguinte 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 para os rastreios do seu perfil. Pode especificar o diretório raiz que contém todos os rastreios de perfis para várias execuções ou um conjunto específico de rastreios de perfis para uma execução.

Por exemplo, se 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>

Pode definir o caminho do GCS para o diretório raiz (gs://bucket-name):

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

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

Por predefinição, o xprofiler create cria uma VM do Compute Engine, especificamente uma VM c4-highmem-8. Pode alterar o tipo de máquina com a flag -m. Se quiser criar uma instância do TensorBoard num pod do GKE em vez de numa VM do Compute Engine, pode transmitir a flag –GKE para xprofiler create. Alojar uma instância do TensorBoard num pod do GKE pode facilitar a gestão do TensorBoard juntamente com o resto da sua carga de trabalho implementada no GKE.

A VM ou o pod do GKE que aloja o TensorBoard torna o carregamento de perfis grandes e de vários perfis muito mais rápido do que o TensorBoard alojado localmente. Google Cloud Com base nos nossos testes de referência, verá perfis da ordem de 1 GB carregados em poucos minutos para o primeiro carregamento com a VM c4-highmem-8 predefinida. Além disso, pode escolher o tipo de máquina certo com base nas suas necessidades de desempenho e custo.

Depois de executar xprofiler create, vai ver um resultado semelhante ao seguinte:

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 no qual pode clicar para ver os seus perfis do XProf num TensorBoard. Esta associação permite a partilha de perfis com a sua equipa e também com os engenheiros da Google que ajudam na otimização do desempenho no Google Cloud.

Controla quem tem acesso ao link com base nas autorizações definidas para o contentor do Cloud Storage que armazena os dados do seu perfil. Caso a associação não funcione por algum motivo, também oferecemos uma forma de estabelecer ligação à instância do TensorBoard através de SSH para ver os seus perfis com o comando xprofiler connect.

Se ativou o servidor do Cloud Profiler no código da carga de trabalho (necessário para capturas programáticas e a pedido) e quer realizar a criação de perfis a pedido, pode fazê-lo de duas formas:

a. Clique no botão Capturar perfil no TensorBoard. Agora, pode escolher qualquer dispositivo anfitrião no qual a sua carga de trabalho esteja a ser executada para criar um perfil. Temos compatibilidade com a captura a pedido para cargas de trabalho em execução no GKE ou no Compute Engine.

b. Use o comando de captura XProfiler, fornecendo informações como a zona, o contentor do Cloud Storage, a framework, os nomes de VM/pod do anfitrião e a duração da captura em milissegundos. Estas são as mesmas informações que tem de introduzir na IU do TensorBoard.

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