Criação de perfil gerenciada com o XProf

É possível usar a criação de perfil gerenciada do XProf ao usar o monitoramento de carga de trabalho do ML Diagnostics, a CLI do ML Diagnostics ou o SDK do ML Diagnostics. O XProf gerenciado é uma instância do XProf com um back-end escalonável que a plataforma de diagnóstico de ML provisiona e gerencia no seu projeto doGoogle Cloud , enquanto o XProf auto-hospedado exige que você implante e mantenha manualmente suas próprias instâncias do XProf para visualizar seus perfis.

O XProf é uma ferramenta de código aberto para criação de perfil e análise de desempenho de cargas de trabalho de machine learning (ML) e faz parte do ecossistema OpenXLA. Em comparação com o XProf auto-hospedado, o XProf gerenciado oferece os seguintes recursos:

  • Configuração zero: não requer instalação ou configuração manual de servidores de visualização do XProf ou das dependências deles.
  • Acesso centralizado a sessões de perfil de várias execuções: carrega sessões de perfil em diferentes execuções de ML com uma única instância gerenciada do XProf por projeto.
  • Escalonamento automático e carregamento rápido: acelera o carregamento de perfis grandes e escalona automaticamente os recursos de back-end quando vários usuários visualizam perfis ao mesmo tempo.
  • Colaboração e segurança integradas: gera links compartilháveis para sessões de perfil e protege as instâncias de criação de perfil contra vulnerabilidades de segurança.
  • Vários modos de captura de perfil: oferece suporte à captura de vários hosts, à criação de perfil sob demanda, à criação de perfil contínua no modo de snapshot e à captura automática de perfil do monitoramento de carga de trabalho.
  • Sem custo extra de infraestrutura: fornece recursos de computação gerenciados do XProf sem custo financeiro adicional, ao contrário da auto-hospedagem de uma instância do XProf.

Caminhos de implantação

É possível implantar e ativar o XProf gerenciado usando o SDK do ML Diagnostics, o monitoramento de cargas de trabalho ou a CLI do ML Diagnostics.

A plataforma ML Diagnostics implanta uma instância gerenciada do XProf com um back-end escalonável por projeto. Use essa instância para ver todas as sessões de perfil capturadas em todas as execuções de ML no projeto. A captura de perfis também requer um bucket do Cloud Storage associado à execução de ML para armazenar os dados de perfil capturados.

É possível ativar o XProf gerenciado e associar um bucket do Cloud Storage à sua execução de ML de uma das seguintes maneiras:

  • SDK do ML Diagnostics: quando você implanta sua primeira execução de ML usando o SDK, a plataforma do ML Diagnostics ativa e implanta automaticamente uma instância gerenciada do XProf no seu projeto. Para especificar o bucket do Cloud Storage que vai armazenar os perfis, defina o parâmetro gcs_path ao definir a execução de ML no SDK.
  • Monitoramento de carga de trabalho: o monitoramento de carga de trabalho não ativa o XProf gerenciado por padrão. Em vez disso, na primeira vez que a captura automática de perfil acionar uma sessão de perfil no seu projeto, a plataforma de diagnóstico de ML vai provisionar e ativar automaticamente uma instância gerenciada do XProf. Por padrão, o ML Diagnostics também provisiona um bucket regional do Cloud Storage no seu projeto.
  • CLI de diagnóstico de ML: é possível atualizar qualquer execução de ML para ativar a ferramenta XProf, que implanta a instância gerenciada do XProf no seu projeto, e associar um bucket do Cloud Storage à execução de ML.

Capturar perfis

O ML Diagnostics oferece quatro métodos para capturar perfis do XProf: captura programática, geração de perfis on demand, geração de perfis contínua no modo de snapshot e captura automática de perfis pelo monitoramento de cargas de trabalho. Esses métodos diferem na forma como uma captura é acionada e na janela de tempo que registram.

A tabela a seguir compara os métodos de captura de perfil:

Método de captura Descrição Frameworks compatíveis
Captura programática

Incorpora comandos de criação de perfil diretamente no código da carga de trabalho para iniciar e interromper a gravação de dados de perfil em etapas ou funções de treinamento especificadas.

Use esse método quando quiser perfis reproduzíveis de etapas de treinamento conhecidas ou seções de código específicas.

JAX, vLLM, Tunix, TorchTPU
Criação de perfil sob demanda

Aciona uma captura de perfil em tempo real no Google Cloud console ou na CLI gcloud durante um período especificado enquanto a carga de trabalho está em execução.

Use esse método quando quiser inspecionar uma carga de trabalho em execução em tempo real sem modificar os intervalos de etapas ou reiniciar a carga de trabalho.

JAX, vLLM, Tunix, TorchTPU
Criação contínua de perfis no modo de snapshot (pré-lançamento)

Registra continuamente um perfil de 2 GB de traces de dispositivos em um buffer circular nos hosts da TPU e salva o buffer de 2 GB mais recente no bucket do Cloud Storage quando você aciona um snapshot.

Use esse método quando precisar inspecionar a atividade da carga de trabalho que leva a uma lentidão ou anomalia intermitente.

JAX, vLLM, Tunix
Captura automática de perfil do monitoramento de cargas de trabalho (pré-lançamento)

Aciona automaticamente uma captura de perfil do XProf quando um analisador de degradação de desempenho de monitoramento de carga de trabalho relata um estado "bad". É possível configurar a captura automática de perfil para acionar uma captura on demand ou, se a criação de perfil contínua estiver ativada, um snapshot de criação de perfil contínua.

Use esse método quando quiser que o ML Diagnostics capture dados de trace no momento em que o desempenho diminuir, sem esperar que você acione manualmente uma captura.

JAX, vLLM, Tunix, TorchTPU

Captura programática

Com a captura programática, você incorpora comandos de criação de perfil diretamente no código da carga de trabalho para especificar explicitamente quando iniciar e interromper a gravação de dados de perfil. Todos os perfis gravados no bucket do Cloud Storage especificado pelo parâmetro gcs_path para sua execução de ML aparecem automaticamente como sessões de perfil na instância gerenciada do XProf do projeto, que pode ser visualizada no console doGoogle Cloud ou usando a Google Cloud CLI.

É possível configurar a captura programática de uma das seguintes maneiras:

Criação de perfil sob demanda

A criação de perfil on demand permite acionar uma captura de perfil em tempo real enquanto a carga de trabalho está em execução. Acionar uma captura sob demanda registra um perfil durante um período especificado. É possível ativar a criação de perfil sob demanda diretamente no código da carga de trabalho ou usando o SDK de diagnóstico de ML.

Pré-requisitos

Antes de capturar um perfil on demand, conclua as seguintes etapas de configuração:

  1. Inicie um servidor XProf na porta 9999 na carga de trabalho para que ele responda aos gatilhos de captura on demand:

    • Diretamente no JAX: chame o método profiler.start_server na porta 9999:

      profiler.start_server(port=9999)
      
    • SDK do ML Diagnostics: defina on_demand_xprof=True em machinelearning_run ao definir sua execução de ML.

  2. Configure o orquestrador de cluster ou o ambiente de computação:

Capturar um perfil sob demanda

É possível acionar uma captura de perfil on demand no console doGoogle Cloud ou na CLI gcloud. Ao acionar uma captura, especifique os seguintes parâmetros:

  • Duração: a duração da captura, em milissegundos ou segundos. A duração máxima que você pode capturar depende das características da carga de trabalho e da quantidade de dados que cabe no limite de 2 GB do perfil.
  • Destinos: os nós de destino em que você quer capturar o perfil:
    • Console doGoogle Cloud : todos os nós que executam a carga de trabalho são preenchidos automaticamente, e é possível selecionar todos ou um subconjunto de nós.
    • CLI gcloud: liste os destinos do profiler da carga de trabalho e acione uma captura para todos ou um subconjunto desses destinos.
  • Níveis de rastreamento de dispositivo, host e Python: o nível de detalhe a ser capturado para os rastreamentos de dispositivo, host e Python.

Criação contínua de perfis no modo de snapshot

Quando a criação contínua de perfis está ativada, a plataforma de diagnósticos de ML registra continuamente um perfil de 2 GB em um buffer circular nos hosts de TPU. Quando você aciona um snapshot, o ML Diagnostics salva o perfil mais recente de 2 GB do buffer como uma sessão de perfil no bucket do Cloud Storage. Esse snapshot permite inspecionar a atividade da carga de trabalho que levou a um problema. A criação de perfil contínua só oferece suporte ao trace de dispositivos. Ela não coleta traces de host ou do Python.

Pré-requisitos

Antes de iniciar a criação de perfil contínua, verifique se o ambiente atende aos seguintes requisitos:

  • Versões de biblioteca compatíveis:
    • libtpu versão 0.0.43 ou mais recente.
    • jax versão 0.11.0 ou mais recente.
  • Configuração do orquestrador de cluster ou do ambiente de computação:

Inicie o servidor XProf no código da carga de trabalho chamando profiler.start_server(port=9999) ou defina on_demand_xprof=True em machinelearning_run ao usar o SDK de diagnóstico de ML.

Iniciar o perfilamento contínuo

É possível iniciar a criação de perfil contínua diretamente no código da carga de trabalho usando a CLI gcloud ou no console do Google Cloud . Depois de iniciado, o perfilamento contínuo registra perfis de 2 GB em um buffer circular nos hosts de TPU especificados.

Ao iniciar a criação de perfil contínua usando a CLI gcloud ou o console doGoogle Cloud , especifique os seguintes parâmetros:

  • Destinos: os nós de destino para criar perfis continuamente:
    • Console doGoogle Cloud : todos os nós que executam a carga de trabalho são preenchidos automaticamente, e é possível selecionar todos ou um subconjunto de nós.
    • CLI gcloud: liste os destinos do profiler para a carga de trabalho e inicie a criação de perfil contínua para todos ou um subconjunto desses destinos.
  • Nível do rastreador de dispositivo: o nível de rastreamento do dispositivo a ser capturado continuamente. O rastreamento de host e Python é desativado durante a criação de perfil contínua.

É possível iniciar a criação de perfil contínua diretamente do código da carga de trabalho do Python, na CLI gcloud ou no console do Google Cloud .

Código de carga de trabalho em Python

Para iniciar a criação de perfil contínua diretamente do código da carga de trabalho do Python, chame a função start_continuous_profiling:

from xprof.api import continuous_profiling_snapshot
continuous_profiling_snapshot.start_continuous_profiling('localhost:9999', {})

Para mais informações, consulte Criação de perfil do XProf JAX.

CLI da gcloud

Para iniciar a criação de perfil contínua na CLI gcloud, liste os destinos ativos do profiler e inicie a criação de perfil contínua nesses nós usando os seguintes comandos:

PROJECT_ID="PROJECT_ID"
LOCATION="LOCATION"
RUN_ID="RUN_ID"

# 1. Discover active ProfilerTarget IDs on the ML run:
TARGETS=$(gcloud alpha mldiagnostics profiler-target list \
  --machine-learning-run="${RUN_ID}" \
  --location="${LOCATION}" \
  --project="${PROJECT_ID}" \
  --format="csv[no-heading](name.basename())" | paste -sd, -)

# 2. Start continuous profiling across the target nodes:
gcloud alpha mldiagnostics machine-learning-run start-continuous-profiling \
  --machine-learning-run="${RUN_ID}" \
  --location="${LOCATION}" \
  --project="${PROJECT_ID}" \
  --targets="${TARGETS}" \
  --device-tracer-level=enabled

Substitua:

  • PROJECT_ID: o ID do projeto Google Cloud .
  • LOCATION: a região em que a execução de ML está localizada.
  • RUN_ID: o ID da execução de ML.
Console doGoogle Cloud

Para iniciar a criação contínua de perfis no console do Google Cloud , siga estas etapas:

  1. No console do Google Cloud , acesse Cluster Director ou GKE > IA/ML > Executar diagnósticos e selecione sua execução de ML.
  2. Abra a guia Perfis.
  3. Na barra de ações de cima, clique em Iniciar criação de perfil contínua.
  4. Selecione os nós para criar o perfil e o nível do rastreador de dispositivos e clique em Iniciar.

Depois que o recurso é iniciado, o banner Capturando perfis contínuos mostra o horário de início da sessão como Since START_TIME, e a barra de ações mostra Criar snapshot e Parar o perfilamento contínuo.

Capturar um snapshot

Enquanto a criação de perfil contínua estiver em execução, você poderá capturar um snapshot para salvar o perfil mais recente de 2 GB do buffer circular como uma sessão de perfil no bucket do Cloud Storage. É possível capturar um snapshot usando a CLI gcloud ou no console do Google Cloud .

Ao capturar um snapshot de criação contínua de perfis, especifique os seguintes parâmetros:

  • Destinos: todos ou um subconjunto dos nós de destino em que a criação de perfil contínua está em execução:
    • Console doGoogle Cloud : todos os nós configurados para criação de perfil contínua são preenchidos automaticamente, e é possível selecionar todos ou um subconjunto de nós.
    • CLI gcloud: especifique todos ou um subconjunto dos destinos do criador de perfil em que a criação de perfil contínua está em execução.
  • Nome da sessão do criador de perfil: um nome para a sessão de captura de snapshot.

É possível capturar um snapshot usando a CLI gcloud ou no console doGoogle Cloud .

CLI da gcloud

Para capturar um snapshot de criação de perfil contínua usando a CLI gcloud, execute o seguinte comando:

PROJECT_ID="PROJECT_ID"
LOCATION="LOCATION"
RUN_ID="RUN_ID"

gcloud alpha mldiagnostics profiler-session capture "SNAPSHOT_NAME" \
  --machine-learning-run="${RUN_ID}" \
  --location="${LOCATION}" \
  --project="${PROJECT_ID}" \
  --targets="TARGET_NODES" \
  --kind=snapshot

Substitua:

  • PROJECT_ID: o ID do projeto Google Cloud .
  • LOCATION: a região em que a execução de ML está localizada.
  • RUN_ID: o ID da execução de ML.
  • SNAPSHOT_NAME: um nome para a sessão de perfil capturada.
  • TARGET_NODES: uma lista separada por vírgulas de nós de destino a serem capturados.
Console doGoogle Cloud

Para capturar um snapshot de criação contínua de perfis no console do Google Cloud , siga estas etapas:

  1. No console do Google Cloud , acesse Cluster Director ou GKE > IA/ML > Executar diagnósticos, selecione sua execução de ML e abra a guia Perfis.
  2. Na barra de ações superior, clique em Criar snapshot.
  3. No painel Capturar snapshot de criação contínua de perfis, insira um nome de perfil, selecione os nós de destino e clique em Enviar.

A captura de tela aparece na tabela Perfis. Se o link não aparecer imediatamente, clique em Atualizar.

Captura automática de perfis do monitoramento de cargas de trabalho

Com a captura automática de perfil do monitoramento de cargas de trabalho, o ML Diagnostics captura automaticamente um perfil do XProf quando um analisador de degradação de desempenho informa um estado "bad", em vez de esperar que você acione manualmente uma captura sob demanda após um problema. É possível configurar a captura automática de perfil para acionar uma captura on demand ou, se a criação de perfil contínua estiver ativada, um snapshot de criação de perfil contínua. Esse processo permite diagnosticar a causa raiz de um problema usando um perfil capturado no momento do evento.

Para instruções de configuração e pré-requisitos, consulte Captura automática de perfil na documentação de monitoramento de cargas de trabalho.

Análise de perfil de IA agêntica

A análise de perfil de IA agêntica analisa automaticamente os perfis de desempenho do XProf capturados para cargas de trabalho de TPU e gera um relatório de análise em HTML com base nos dados de trace brutos do XProf. Usando modelos do Gemini na Vertex AI, o agente de análise executa ferramentas de criação de perfil do XProf, resume gargalos de desempenho e recomenda mudanças de código ou configuração. Quando as APIs necessárias são ativadas, o ML Diagnostics gera um relatório de análise de perfil de IA agêntica para cada perfil capturado na plataforma.

A análise de perfil da IA agêntica usa as mesmas ferramentas e recursos que o repositório de habilidades de código aberto do XProf. Em vez de usar o arquivo de habilidade diretamente ou instalar a CLI do XProf exigida pela habilidade de código aberto, o ML Diagnostics importa as implementações da ferramenta XProf subjacente para que o agente possa chamá-las diretamente. Como resultado, a saída da análise no ML Diagnostics é equivalente à execução das habilidades de código aberto do XProf com a CLI do XProf em arquivos xplane locais.

Pré-requisitos

Antes de usar a análise de perfil de IA agêntica, conclua as seguintes etapas de configuração no seu projeto do Google Cloud .

Ativar APIs obrigatórias

Ative a API Vertex AI (aiplatform.googleapis.com) para que o agente de análise possa executar modelos do Gemini no seu projeto:

  • Console doGoogle Cloud : na Biblioteca de APIs, ative a API Vertex AI.
  • CLI gcloud: execute o seguinte comando:

    gcloud services enable aiplatform.googleapis.com
    

Conceder papéis do IAM

Conceda à sua conta de usuário um dos seguintes papéis do Identity and Access Management (IAM):

  • Leitor do Cluster Director (roles/hypercomputecluster.viewer): concede acesso somente leitura aos recursos do Cluster Director e do ML Diagnostics, incluindo relatórios de análise gerados no console do Google Cloud em Executar diagnóstico > Perfis.
  • Função personalizada (hypercomputecluster.locations.analysis): concede permissão para visualizar relatórios de análise gerados no Google Cloud console em Executar diagnósticos > Perfis quando adicionado a uma função personalizada.

Como funciona

Quando uma captura de perfil programática, sob demanda, contínua ou automática é concluída, um job em segundo plano isolado executa o seguinte fluxo de trabalho de análise no perfil do XProf:

  1. Coleta de dados ampla: executa 15 ferramentas do XProf em paralelo para coletar um instantâneo de integridade que inclui tempo de etapa, utilização de hardware, uso da memória, limites de roofline, operações mais lentas e ineficiências do compilador. Essa etapa normaliza os tempos por etapa de treinamento e separa as operações íntegras dos gargalos.
  2. Investigação de IA focada: analisa o resumo de integridade coletado. Se a carga de trabalho tiver alta utilização de computação, baixa sobrecarga de inatividade e comunicação e baixo risco de falta de memória, o agente vai pular as chamadas de ferramentas adicionais e gerar um relatório OPTIMAL. Se gargalos forem detectados, o agente formará hipóteses e executará até quatro verificações de acompanhamento direcionadas usando outras ferramentas do XProf para inspecionar as operações ao redor no gráfico de computação, atrasos na linha do tempo entre hosts e kernels personalizados.
  3. Geração de relatórios: combina as métricas do XProf com a análise de causa raiz do agente de IA e as correções recomendadas. O relatório summary.html gerado é salvo no bucket do Cloud Storage e exibido no console doGoogle Cloud em Executar diagnósticos > Perfis.

Esse fluxo de trabalho mantém os dados e as análises do perfil particulares no projeto. Os traces brutos do criador de perfil e os relatórios HTML gerados são armazenados no bucket do Cloud Storage, e todas as chamadas de modelo de IA são executadas pela Vertex AI no projeto do Google Cloud .

Ferramentas do XProf

A tabela a seguir descreve as ferramentas do XProf que o agente executa durante a análise:

Categoria Ferramentas do XProf Descrição
Visão geral da carga de trabalho e tempo das etapas get_xprof_overview, get_step_trace, get_kpi_metrics, get_device_information, get_hosts, get_utilization_viewer, get_smart_suggestions

Essas ferramentas são executadas durante a coleta de dados ampla, e a ferramenta get_hosts também é executada durante a investigação de IA focada, para verificar o seguinte:

  • Tipo de hardware, número de hosts e chips e limites máximos de hardware
  • Tempo médio de etapa e detalhamento de computação, comunicação de rede e tempo de inatividade do host ou do carregamento de dados
  • Unidade de multiplicação de matriz (MXU), TensorCore e utilização da largura de banda de memória, além de recomendações integradas do XProf
Roofline e operações mais lentas get_roofline_model, get_hlo_op_profile

Essas ferramentas são executadas durante a coleta de dados ampla e a investigação de IA focada para verificar o seguinte:

  • Se a carga de trabalho ou as operações específicas estão limitadas pela largura de banda de computação ou de memória.
  • Tempo gasto por tipo de operação, como multiplicação de matrizes, atenção e operações coletivas, e links para linhas de código-fonte do Python
Detectores de ineficiência do compilador detect_layout_mismatch_copies, detect_unfused_reshapes, detect_unnecessary_convert_reduce, list_hlo_modules

Essas ferramentas são executadas durante a coleta de dados ampla para verificar o seguinte:

  • Cópias de dados extras causadas por incompatibilidades de layout de tensor entre operações
  • Operações reshape, transpose ou copy não combinadas que desperdiçam largura de banda de memória
  • Conversões de precisão desnecessárias, como bf16 para f32, dentro de reduções
Memória e risco de OOM get_memory_profile, get_peak_allocations

Essas ferramentas são executadas durante a coleta de dados ampla para verificar o seguinte:

  • Uso máximo da memória de alta largura de banda (HBM) no dispositivo em comparação com a capacidade total e a fragmentação da memória
  • Se o risco de falta de memória (OOM) é LOW, MEDIUM ou HIGH, e os maiores tensores na memória
Explorador de gráficos de computação get_hlo_neighborhood

Essa ferramenta é executada durante a investigação de IA focada para verificar as operações imediatamente antes e depois de uma operação lenta no gráfico XLA e identificar o que bloqueia a fusão do compilador ou força o movimento de dados.

Linha do tempo e tempo do kernel aggregate_xplane_events, get_kernel_stats

Essas ferramentas são executadas durante a investigação de IA focada para verificar o seguinte:

  • Estatísticas da linha do tempo em hosts para distinguir as interrupções do carregador de dados dos atrasos na comunicação de rede
  • Tempo de execução e contagem de chamadas para kernels de dispositivos individuais
Criador de perfil de kernel personalizado llo_load, llo_query

Essas ferramentas são executadas durante a investigação de IA focada para verificar os traces de execução de baixo nível para kernels personalizados, como Pallas ou Mosaic, que parecem opacos para a criação de perfil padrão de alto nível.

Relatório de análise de perfil de IA

Quando a análise de um perfil capturado é concluída, o ML Diagnostics salva o relatório summary.html gerado no bucket do Cloud Storage e o exibe no console do Google Cloud em Executar diagnóstico > Perfis. Use esse relatório para analisar a integridade geral da carga de trabalho, inspecionar os principais gargalos de desempenho e identificar as mudanças recomendadas de código ou configuração.

Cada relatório gerado inclui as seguintes seções:

  • Resumo da análise: mostra um status geral de OPTIMAL ou NEEDS_ATTENTION, um resumo de 3 a 5 frases e as principais métricas de performance, como tempo de etapa, uso da memória e eficiência do Roofline.
  • Gargalos priorizados e correções práticas: lista as descobertas priorizadas de P0 para problemas críticos a P2 para problemas menores. Cada descoberta inclui:
    • O que está acontecendo: uma explicação da causa raiz, incluindo dimensões e layouts de tensor, quando relevante.
    • Evidência: citações diretas das métricas e operações do XProf em que o gargalo foi medido.
    • Impacto e correção esperados: a quantidade de tempo de etapa ou memória afetada, combinada com as mudanças recomendadas de framework, fragmentação ou compilador.
  • Análise detalhada da ferramenta XProf: fornece dados de suporte e detalhamentos das seguintes análises:
    • Análise de roofline
    • Análise de memória
    • Análise de perfil de operação do otimizador de alto nível (HLO, na sigla em inglês)
    • Estatísticas de execução do kernel
    • Análise personalizada de kernel e otimizador de baixo nível