Como criar perfil do JAX em GPUs com XProf e Diagnóstico de ML

Para otimizar modelos JAX em grande escala em GPUs, é necessário ter visibilidade total dos gargalos de desempenho. Este guia fornece um fluxo de trabalho completo de ponta a ponta (E2E) para executar e criar perfis de cargas de trabalho JAX em GPUs (como NVIDIA L4) usando o Google Cloud ML Diagnostics e o XProf. Ao usar essas ferramentas, é possível identificar operações ineficientes, otimizar o uso de recursos de computação e acelerar as execuções de treinamento.

Ao seguir este guia, você vai aprender a:

  1. Instrumente um loop de treinamento simples do JAX para criação de perfis.
  2. Contêinere a carga de trabalho com o suporte adequado do CUDA.
  3. Implante a carga de trabalho no Google Kubernetes Engine (GKE) usando o JobSet.
  4. Capture e visualize perfis de performance de forma dinâmica.

Pré-requisitos

Antes de começar, verifique se você atende a estes requisitos:

  • Ter um projeto do Google Cloud com o faturamento ativado.
  • Um cluster do GKE com suporte a GPU (por exemplo, NVIDIA L4).
  • Um bucket do Google Cloud Storage (GCS) para armazenar perfis.
  • As CLIs gcloud e kubectl instaladas e configuradas.
  • A Identidade da carga de trabalho configurada para o cluster do GKE acessar o GCS.

Etapa 1: instrumentar a carga de trabalho do JAX

Primeiro, crie um script de treinamento do JAX (por exemplo, train.py). Usamos o SDK google-cloud-mldiagnostics para interagir com a infraestrutura de criação de perfil gerenciada.

[!WARNING] O script abaixo inclui um loop infinito para manter a GPU ocupada para demonstrações de criação de perfil sob demanda. Não se esqueça de interromper manualmente o job ou excluir os recursos do GKE depois de concluir para evitar custos desnecessários de faturamento.

import logging
import os
import time
from google_cloud_mldiagnostics import machinelearning_run
from google_cloud_mldiagnostics import xprof
import jax
import jax.numpy as jnp
import numpy as np

logging.basicConfig(level=logging.INFO)

def main():
    logging.info("Starting JAX training job...")

    # Coordinates multihost collective operations and healthchecks
    jax.distributed.initialize()

    logging.info(
        f"JAX initialized: process_index={jax.process_index()}, "
        f"process_count={jax.process_count()}"
    )

    # Syncs metadata with the mldiag hook & launches reverse proxy daemons
    machinelearning_run(
        name=f"jax-gpu-run-{int(time.time())}",
        configs={"learning_rate": 1e-5, "batch_size": 8192},
        project=os.environ.get("PROJECT_ID", "<your-project-id>"),
        region=os.environ.get("REGION", "us-central1"),
        gcs_path=os.environ.get("GCS_BUCKET", "gs://<your-gcs-bucket>"),
        on_demand_xprof=True,
    )

    key = jax.random.PRNGKey(0)
    size = 4096
    matrix = jax.random.normal(key, (size, size), dtype=jnp.float32)

    def train_step(x):
        return jnp.dot(x, x)

    train_step = jax.jit(train_step)

    # Triggers XLA compilation ahead of tracing steps so compilation overhead isn't profiled
    matrix = train_step(matrix)
    matrix.block_until_ready() # Wait for compilation to complete.

    prof = xprof()
    prof.start(session_id="warmup_phase")

    for _ in range(5):
        matrix = train_step(matrix)
        matrix.block_until_ready()

    prof.stop()
    logging.info("Programmatic profile capture complete.")

    logging.info("Entering training loop. Ready for on-demand profiling...")
    try:
        while True:
            # Continuously pump steps keeping GPUs occupied for on-demand capture triggers
            matrix = train_step(matrix)
            matrix.block_until_ready() # Ensure GPU work completes before next step.
            time.sleep(0.5)
    except KeyboardInterrupt:
        logging.info("Training loop interrupted.")

if __name__ == "__main__":
    main()

Etapa 2: contêinerização (Dockerfile)

Crie um Dockerfile para empacotar seu script JAX com as dependências necessárias do CUDA e o SDK do ML Diagnostics.

# Use an official NVIDIA CUDA base image compatible with JAX
FROM nvidia/cuda:13.2.1-cudnn-devel-ubuntu24.04

# Install Python, venv, and other OS dependencies
RUN apt-get update && apt-get install -y \
    python3-pip \
    python3-venv \
    python3-dev \
    git \
    curl \
    && rm -rf /var/lib/apt/lists/*

# Set up a virtual environment and update PATH to use it implicitly
ENV VIRTUAL_ENV=/opt/venv
RUN python3 -m venv $VIRTUAL_ENV
ENV PATH="$VIRTUAL_ENV/bin:$PATH"

# At this point, pip and python implicitly map to the virtual env!
# No need for --break-system-packages.

# Upgrade pip inside the venv
RUN pip install --upgrade pip

# Install JAX with CUDA support
RUN pip install --upgrade "jax[cuda13]" -f https://storage.googleapis.com/jax-releases/jax_cuda_releases.html

# Install ML Diagnostics SDK and XProf tools
RUN pip install --no-cache-dir \
    google-cloud-mldiagnostics \
    xprof-nightly

WORKDIR /app
COPY train.py .

CMD ["python3", "train.py"]

Crie e envie a imagem para o Artifact Registry:

docker build -t us-central1-docker.pkg.dev/<project-id>/<repo>/jax-gpu-workload:latest .
docker push us-central1-docker.pkg.dev/<project-id>/<repo>/jax-gpu-workload:latest

Etapa 3: implantação (manifesto do Kubernetes)

Implante a carga de trabalho usando um JobSet ou um job padrão do GKE. Para permitir que a plataforma de diagnóstico de ML injete metadados e encaminhe solicitações de perfil, aplique o rótulo managed-mldiagnostics-gke: "true". Para mais detalhes sobre como configurar o GKE para o ML Diagnostics, consulte o guia de configuração do GKE oficial.

apiVersion: jobset.x-k8s.io/v1alpha2
kind: JobSet
metadata:
  name: jax-gpu-job
  namespace: ai-workloads
  labels:
    managed-mldiagnostics-gke: "true"
spec:
  replicatedJobs:
  - name: gpu-nodes
    replicas: 1
    template:
      spec:
        parallelism: 1
        completions: 1
        backoffLimit: 0
        template:
          metadata:
            labels:
              managed-mldiagnostics-gke: "true"
          spec:
            # Must match the GKE Service Account with Workload Identity permissions
            serviceAccountName: <your-service-account>
            hostNetwork: true
            dnsPolicy: ClusterFirstWithHostNet
            nodeSelector:
              cloud.google.com/gke-accelerator: nvidia-l4 # Or other GPU
            containers:
            - name: workload
              image: us-central1-docker.pkg.dev/<project-id>/<repo>/jax-gpu-workload:latest
              imagePullPolicy: Always
              # Expose ports required for profile daemons
              ports:
              - containerPort: 8471 # JAX distributed coordinator port
              - containerPort: 8080 # ML Diagnostics agent/proxy port
              - containerPort: 9999 # XProf server port for on-demand profiling
              resources:
                limits:
                  nvidia.com/gpu: 1

Aplique o manifesto:

kubectl apply -f deploy.yaml

Etapa 4: captura e visualização

Captura programática

Se você incluiu prof.start() / prof.stop() no script, esses perfis são enviados automaticamente para o bucket do GCS no caminho: gs://<your-gcs-bucket>/<run-name>/plugins/profile/<session-id>/

Captura sob demanda

Como on_demand_xprof=True está definido em machinelearning_run, é possível capturar perfis dinamicamente enquanto o job está em execução.

Para instruções detalhadas sobre como usar a interface do TensorBoard para acionar perfis on demand, selecionar pods específicos e visualizar os rastreamentos capturados, consulte a documentação pública oficial: Diagnóstico de ML do Google Cloud: captura de perfil on demand.

Também é possível capturar perfis usando a CLI gcloud, conforme descrito no guia da CLI do ML Diagnostics.

Esta documentação pública se aplica a cargas de trabalho de TPU e GPU gerenciadas pelo ML Diagnostics.