使用 XProf 和 ML 诊断在 GPU 上分析 JAX

在 GPU 上优化大规模 JAX 模型需要深入了解性能瓶颈。本指南提供了一个全面的端到端 (E2E) 工作流，用于使用 Google Cloud ML Diagnostics 和 XProf 在 GPU（例如 NVIDIA L4）上运行和分析 JAX 工作负载。借助这些工具，您可以识别低效的操作、优化计算资源使用情况并加快训练运行速度。

按照本指南操作，您将了解如何：

  1. 为性能分析对简单的 JAX 训练循环进行插桩。
  2. 将工作负载容器化，并提供适当的 CUDA 支持。
  3. 使用 JobSet 在 Google Kubernetes Engine (GKE) 上部署工作负载。
  4. 动态捕获和直观呈现性能分析数据。

前提条件

开始之前，请确保您满足以下方面的要求：

  • 启用了结算功能的 Google Cloud 项目。
  • 支持 GPU 的 GKE 集群（例如 NVIDIA L4）。
  • 用于存储配置的 Google Cloud Storage (GCS) 存储桶。
  • 已安装并配置 gcloud 和 kubectl CLI。
  • 为 GKE 集群配置了 Workload Identity 以访问 GCS。

第 1 步：检测 JAX 工作负载

首先，创建一个 JAX 训练脚本（例如 train.py）。我们使用 google-cloud-mldiagnostics SDK 与受管分析基础架构进行交互。

[!WARNING] 下面的脚本包含一个无限循环，用于在按需分析演示中保持 GPU 处于忙碌状态。请务必在完成后手动停止作业或删除 GKE 资源，以免产生不必要的结算费用。

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()

第 2 步：容器化 (Dockerfile)

创建一个 Dockerfile，用于将 JAX 脚本与所需的 CUDA 依赖项和 ML Diagnostics SDK 打包在一起。

# 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"]

构建映像并将其推送到 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

第 3 步：部署（Kubernetes 清单）

使用 GKE JobSet 或标准作业部署工作负载。如需使机器学习诊断平台能够注入元数据并路由配置文件请求，请应用标签 managed-mldiagnostics-gke: "true"。如需详细了解如何为 ML Diagnostics 配置 GKE，请参阅官方 GKE 设置指南。

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

应用清单：

kubectl apply -f deploy.yaml

第 4 步：捕获和可视化

程序化捕获

如果您在脚本中添加了 prof.start() / prof.stop()，这些配置文件会自动上传到 GCS 存储桶中的以下路径：gs://<your-gcs-bucket>/<run-name>/plugins/profile/<session-id>/

按需捕获

由于 on_demand_xprof=True 是在 machinelearning_run 中设置的，因此您可以在作业运行时动态捕获配置文件。

如需详细了解如何使用 TensorBoard 界面触发按需分析、选择特定 pod 和查看捕获的轨迹，请参阅官方公开文档：Google Cloud ML 诊断 - 按需捕获分析。

您还可以按照 ML Diagnostics CLI 指南中所述，使用 gcloud CLI 捕获配置文件。

此公开文档适用于由 ML 诊断工具管理的 TPU 和 GPU 工作负载。