使用 XProf 进行托管性能分析

使用 ML Diagnostics 工作负载监控、ML Diagnostics CLI 或 ML Diagnostics SDK 时，您可以使用代管式 XProf 性能分析。托管式 XProf 是一个具有可伸缩后端的 XProf 实例，由机器学习诊断平台在您的Google Cloud 项目中预配和管理，而自托管式 XProf 则需要您手动部署和维护自己的 XProf 实例，才能查看性能分析文件。

XProf 是一款开源性能分析和性能分析工具，适用于机器学习 (ML) 工作负载，是 OpenXLA 生态系统的一部分。与自托管 XProf 相比，代管式 XProf 具有以下功能：

• 零设置：无需手动安装或配置 XProf 可视化服务器或其依赖项。
• 集中式多运行性能剖析会话访问：通过每个项目一个代管式 XProf 实例，加载不同机器学习运行中的性能剖析会话。
• 自动伸缩和快速加载：加快大型个人资料的加载速度，并在多用户同时查看个人资料时自动伸缩后端资源。
• 内置协作和安全功能：为性能分析会话生成共享链接，同时保护性能分析实例免受安全漏洞的影响。
• 多种性能分析捕获模式：支持多主机捕获、按需性能分析、快照模式下的持续性能分析，以及通过工作负载监控自动捕获性能分析。
• 无额外基础设施费用：与自行托管 XProf 实例不同，该服务免费提供托管式 XProf 计算资源。

部署路径

您可以使用 ML Diagnostics SDK、工作负载监控或 ML Diagnostics CLI 部署和启用代管式 XProf。

ML Diagnostics 平台会为每个项目部署一个具有可伸缩后端的代管式 XProf 实例。您可以使用此实例查看项目中所有机器学习运行作业中捕获的所有性能分析会话。捕获性能剖析文件还需要一个与机器学习运行关联的 Cloud Storage 存储桶，用于存储捕获的性能剖析文件数据。

您可以通过以下任一方式启用代管式 XProf 并将 Cloud Storage 存储桶与机器学习运行相关联：

• ML Diagnostics SDK：当您使用该 SDK 部署第一个机器学习运行作业时，ML Diagnostics 平台会自动在您的项目中启用并部署一个代管式 XProf 实例。在 SDK 中定义机器学习运行时，您可以通过设置 gcs_path 参数来指定用于存储配置文件的 Cloud Storage 存储桶。
• 工作负载监控：工作负载监控默认不会启用代管式 XProf。相反，当自动性能分析文件捕获首次在项目中触发性能分析会话时，机器学习诊断平台会自动预配并启用代管式 XProf 实例。默认情况下，机器学习诊断还会在您的项目中预配一个区域级 Cloud Storage 存储桶。
• 机器学习诊断 CLI：您可以更新任何机器学习运行作业，以启用 XProf 工具，该工具会在您的项目中部署代管式 XProf 实例，并将 Cloud Storage 存储桶与机器学习运行作业相关联。

捕获性能剖析文件

机器学习诊断支持四种捕获 XProf 配置文件的方法：程序化捕获、按需性能分析、快照模式下的持续性能分析，以及通过工作负载监控自动捕获配置文件。这些方法的不同之处在于触发捕获的方式以及它们记录的时间范围。

下表比较了各种配置文件捕获方法：

捕获方法	说明	支持的框架
程序化捕获	直接在工作负载代码中嵌入性能分析命令，以便在指定的训练步骤或函数中开始和停止记录性能分析数据。 如果您希望获得已知训练步骤或特定代码部分的重现性分析，请使用此方法。	JAX、vLLM、Tunix、TorchTPU
按需性能分析	在工作负载运行期间，通过 Google Cloud 控制台或 gcloud CLI 在您指定的时长内触发实时配置文件捕获。 如果您想实时检查正在运行的工作负载，而无需修改步长范围或重启工作负载，请使用此方法。	JAX、vLLM、Tunix、TorchTPU
快照模式下的持续分析（预览版）	在 TPU 主机上的循环缓冲区中持续记录 2 GB 的设备跟踪记录，并在您触发快照时将最新的 2 GB 缓冲区保存到 Cloud Storage 存储桶。 当您需要检查导致间歇性速度减慢或异常的工作负载活动时，请使用此方法。	JAX、vLLM、Tunix
通过工作负载监控自动捕获性能剖析文件（预览版）	当工作负载监控性能下降分析器报告 "bad" 状态时，自动触发 XProf 配置文件捕获。您可以配置自动性能分析文件捕获，以触发按需捕获，或者在启用持续分析功能时触发持续分析快照。 如果您希望机器学习诊断在性能下降时立即捕获跟踪数据，而无需等待您手动触发捕获，请使用此方法。	JAX、vLLM、Tunix、TorchTPU

程序化捕获

使用程序化捕获时，您可以直接在工作负载代码中嵌入性能分析命令，以明确指定何时开始和停止记录性能分析数据。写入到机器学习运行作业的 gcs_path 参数所指定的 Cloud Storage 存储桶中的任何配置文件都会自动显示为项目的代管式 XProf 实例中的配置文件会话，您可以在Google Cloud 控制台中或使用 Google Cloud CLI 查看这些会话。

您可以通过以下任一方式配置程序化捕获：

• ML Diagnostics SDK：使用 SDK 中的 prof 方法。这些方法封装了底层框架级性能分析 API，因此您的性能分析代码在 JAX、vLLM、Tunix 和 TorchTPU 中保持一致。
• 框架原生性能分析 API：直接调用机器学习框架的内置性能分析方法，并将输出保存到机器学习运行的 Cloud Storage 存储桶：
  • JAX：分析计算
  • vLLM：在 Cloud TPU 上对 vLLM 进行性能分析
  • TorchTPU：使用 XProf 进行性能剖析

按需性能分析

注意：TorchTPU 不支持直接通过框架启用按需性能分析。如需将按需性能分析与 TorchTPU 搭配使用，请使用 ML Diagnostics SDK。

借助按需性能分析，您可以在工作负载运行时实时触发性能分析文件捕获。触发按需捕获后，系统会在您指定的时长内记录性能分析数据。您可以直接在工作负载代码中启用按需性能分析，也可以使用 ML Diagnostics SDK 启用。

前提条件

在捕获按需性能剖析文件之前，请完成以下设置步骤：

1. 在工作负载中的端口 9999 上启动 XProf 服务器，使其响应按需捕获触发器：

   • 直接在 JAX 中：在端口 9999 上调用 profiler.start_server 方法：

     profiler.start_server(port=9999)
     

   • ML Diagnostics SDK：在定义机器学习运行作业时，在 machinelearning_run 中设置 on_demand_xprof=True。

2. 配置集群编排器或计算环境：

   • Google Kubernetes Engine (GKE) 或自行管理的 Kubernetes：部署算子和网络钩子或 ML 诊断 DaemonSet。
   • Cluster Director 管理的 Slurm 或自行管理的 Slurm：使用 Slurm prolog 脚本部署机器学习诊断守护程序。
   • Compute Engine：使用实例启动脚本部署机器学习诊断守护程序。

按需捕获性能分析

您可以通过 Google Cloud 控制台或 gcloud CLI 触发按需性能分析文件捕获。触发捕获时，请指定以下参数：

• 时长：捕获时长，以毫秒或秒为单位。您可以捕获的最长时长取决于您的工作负载特征以及 2 GB 性能分析限制内可容纳的数据量。
• 目标：您要捕获其配置文件的目标节点：
  • Google Cloud 控制台：运行工作负载的所有节点都会自动填充，您可以选择所有节点或部分节点。
  • gcloud CLI：列出工作负载的性能分析器目标，然后触发对所有或部分目标的捕获。
• 设备、主机和 Python 跟踪器级别：要为设备、主机和 Python 跟踪记录捕获的详细程度。

快照模式下的持续分析

预览

此功能受服务专用条款的“通用服务条款”部分中的“非正式版产品条款”的约束。 非正式版功能“按原样”提供，可能只能获得有限的支持。 如需了解详情，请参阅发布阶段说明。

启用持续分析后，机器学习诊断平台会在 TPU 主机的循环缓冲区中持续记录 2 GB 的性能分析数据。触发快照时，机器学习诊断会将缓冲区中最新的 2 GB 性能分析数据保存为 Cloud Storage 存储桶中的性能分析会话。通过此快照，您可以检查导致问题的工作负载活动。持续性能分析仅支持设备跟踪；它不会收集主机或 Python 跟踪记录。

注意：TorchTPU 框架不支持持续性能分析。

前提条件

在开始持续性能分析之前，请确保您的环境满足以下要求：

• 支持的库版本：
  • libtpu 0.0.43 版或更高版本。
  • jax 0.11.0 版或更高版本。
• 集群编排器或计算环境设置：
  • GKE 或自行管理的 Kubernetes：部署 ML Diagnostics DaemonSet。
  • Cluster Director 管理的 Slurm 或自行管理的 Slurm：使用 Slurm prolog 脚本部署机器学习诊断守护程序。
  • Compute Engine：使用实例启动脚本部署机器学习诊断守护程序。

通过调用 profiler.start_server(port=9999) 在工作负载代码中启动 XProf 服务器，或在使用 ML Diagnostics SDK 时在 machinelearning_run 中设置 on_demand_xprof=True。

开始持续分析

您可以使用 gcloud CLI 或在 Google Cloud 控制台中直接在工作负载代码中启动持续性能分析。启动后，持续分析会在指定 TPU 主机上的循环缓冲区中记录 2 GB 的性能分析数据。

使用 gcloud CLI 或Google Cloud 控制台启动持续性能分析时，请指定以下参数：

• 目标：要持续分析的目标节点：
  • Google Cloud 控制台：运行工作负载的所有节点都会自动填充，您可以选择所有节点或部分节点。
  • gcloud CLI：列出工作负载的性能分析器目标，然后开始对所有或部分目标进行持续性能分析。
• 设备跟踪器级别：要持续捕获的设备跟踪级别。在持续性能分析期间，主机和 Python 跟踪功能处于停用状态。

您可以直接从 Python 工作负载代码、gcloud CLI 或 Google Cloud 控制台中启动持续性能分析。

Python 工作负载代码

如需直接从 Python 工作负载代码启动持续性能分析，请调用 start_continuous_profiling 函数：

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

如需了解详情，请参阅 XProf JAX 性能分析。

gcloud CLI

如需从 gcloud CLI 启动持续性能分析，请使用以下命令列出活跃的性能分析器目标，并在这些节点上启动持续性能分析：

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

替换以下内容：

• PROJECT_ID：您的 Google Cloud 项目 ID。
• LOCATION：机器学习运行所在的区域。
• RUN_ID：机器学习运行的 ID。

Google Cloud 控制台

如需在 Google Cloud 控制台中启动持续分析，请完成以下步骤：

1. 在 Google Cloud 控制台中，依次前往 Cluster Director 或 GKE > AI/机器学习 > 运行诊断，然后选择您的机器学习运行。
2. 打开分析结果标签页。
3. 在顶部操作栏中，点击开始持续性能分析。
4. 选择要分析的节点和设备跟踪器级别，然后点击开始。

持续分析开始后，正在捕获持续分析横幅会将该会话的开始时间显示为 Since START_TIME，操作栏会显示拍摄快照和停止持续分析。

捕获快照

在持续分析运行期间，您可以捕获快照，将循环缓冲区中最新的 2 GB 性能分析数据保存为 Cloud Storage 存储桶中的性能分析会话。您可以使用 gcloud CLI 或在 Google Cloud 控制台中捕获快照。

捕获持续分析快照时，请指定以下参数：

• 目标：运行持续性能分析的所有或部分目标节点：
  • Google Cloud 控制台：系统会自动填充为持续性能分析配置的所有节点，您可以选择所有节点或部分节点。
  • gcloud CLI：指定运行持续性能分析的所有或部分性能分析器目标。
• 性能分析器会话名称：捕获的快照会话的名称。

您可以使用 gcloud CLI 或在Google Cloud 控制台中捕获快照。

gcloud CLI

如需使用 gcloud CLI 捕获持续性能分析快照，请运行以下命令：

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

替换以下内容：

• PROJECT_ID：您的 Google Cloud 项目 ID。
• LOCATION：机器学习运行所在的区域。
• RUN_ID：机器学习运行的 ID。
• SNAPSHOT_NAME：捕获的配置文件会话的名称。
• TARGET_NODES：要捕获的目标节点的英文逗号分隔列表。

Google Cloud 控制台

如需在 Google Cloud 控制台中捕获持续性能分析快照，请完成以下步骤：

1. 在 Google Cloud 控制台中，前往 Cluster Director 或 GKE > AI/机器学习 > 运行诊断，选择您的机器学习运行作业，然后打开配置文件标签页。
2. 在顶部操作栏中，点击创建快照。
3. 在捕获持续分析快照窗格中，输入分析名称，选择目标节点，然后点击提交。

快照会显示在配置文件表中。如果链接未立即显示，请点击刷新。

通过工作负载监控自动捕获性能剖析文件

预览

此功能受服务专用条款的“通用服务条款”部分中的“非正式版产品条款”的约束。 非正式版功能“按原样”提供，可能只能获得有限的支持。 如需了解详情，请参阅发布阶段说明。

借助工作负载监控中的自动性能分析捕获功能，当性能下降分析器报告 "bad" 状态时，机器学习诊断会自动捕获 XProf 性能分析，而无需等待您在问题发生后手动触发按需捕获。您可以将自动性能分析文件捕获配置为触发按需捕获，或者在启用持续性能分析时触发持续性能分析快照。通过此过程，您可以使用在事件发生时捕获的性能分析文件来诊断问题的根本原因。

如需了解设置说明和前提条件，请参阅工作负载监控文档中的自动配置文件捕获。

智能体 AI 个人资料分析

预览

此功能受服务专用条款的“通用服务条款”部分中的“非正式版产品条款”的约束。 非正式版功能“按原样”提供，可能只能获得有限的支持。 如需了解详情，请参阅发布阶段说明。

智能体 AI 性能分析功能会自动分析为 TPU 工作负载捕获的 XProf 性能分析，并根据原始 XProf 跟踪数据生成 HTML 分析报告。分析代理使用 Vertex AI 上的 Gemini 模型运行 XProf 性能分析工具，总结性能瓶颈，并推荐代码或配置更改。启用所需的 API 后，机器学习诊断会为平台中捕获的每个配置文件生成智能体 AI 配置文件分析报告。

智能体 AI 配置文件分析使用与开源 XProf 技能代码库相同的底层工具和功能。ML Diagnostics 不会直接使用技能文件或安装开源技能所需的 XProf CLI，而是导入底层 XProf 工具实现，以便智能体可以直接调用它们。因此，机器学习诊断中的分析输出相当于在本地 xplane 文件上使用 XProf CLI 运行开源 XProf 技能。

前提条件

在使用智能体 AI 配置文件分析之前，请在 Google Cloud 项目中完成以下设置步骤。

启用必需的 API

启用 Vertex AI API (aiplatform.googleapis.com)，以便分析代理可以在您的项目中运行 Gemini 模型：

• Google Cloud 控制台：在 API 库中，启用 Vertex AI API。

• gcloud CLI：运行以下命令：

  gcloud services enable aiplatform.googleapis.com
  

授予 IAM 角色

向您的用户账号授予以下 Identity and Access Management (IAM) 角色之一：

• Cluster Director Viewer (roles/hypercomputecluster.viewer)：授予对 Cluster Director 和机器学习诊断资源的只读权限，包括 Google Cloud 控制台中的 Run Diagnostics > Profiles（运行诊断 > 配置文件）下生成的分析报告。
• 自定义角色 (hypercomputecluster.locations.analysis)：添加到自定义角色后，授予在 Google Cloud 控制台中的运行诊断 > 配置文件下查看生成的分析报告的权限。

工作原理

当以编程方式、按需、持续或自动捕获性能分析完成后，一个隔离的后台作业会对 XProf 性能分析运行以下分析工作流：

1. 广泛的数据收集：并行运行 15 个 XProf 工具，以收集健康状况快照，其中包括步长时间、硬件利用率、内存用量、Roofline 边界、最慢的操作和编译器效率低下问题。此步骤会按训练步骤对时间进行归一化，并将正常操作与瓶颈分开。
2. 重点 AI 调查：查看收集的健康状况快照。如果工作负载的计算利用率较高、空闲和通信开销较低，且内存不足风险较低，则智能体将跳过其他工具调用，并生成 OPTIMAL 报告。如果检测到瓶颈，智能体就会形成假设，并使用其他 XProf 工具运行最多四项有针对性的后续检查，以检查计算图中的周围操作、跨主机的延迟时间以及自定义内核。
3. 报告生成：将 XProf 中的指标与 AI 智能体的根本原因分析和建议的修复措施相结合。生成的 summary.html 报告会保存到您的 Cloud Storage 存储桶中，并显示在Google Cloud 控制台中的运行诊断 > 配置文件下。

此工作流可确保您的性能剖析数据和分析结果在项目中保持私密。 原始性能分析器跟踪记录和生成的 HTML 报告存储在您的 Cloud Storage 存储桶中，所有 AI 模型调用都通过 Google Cloud 项目中的 Vertex AI 运行。

XProf 工具

下表介绍了代理在分析期间运行的 XProf 工具：

类别	XProf 工具	说明
工作负载概览和步骤时间	get_xprof_overview、get_step_trace、get_kpi_metrics、get_device_information、get_hosts、get_utilization_viewer、get_smart_suggestions	这些工具会在广泛数据收集期间运行，而 get_hosts 工具也会在重点 AI 调查期间运行，以检查以下内容： 硬件类型、主机和芯片数量以及硬件峰值限制 平均步长时间以及计算、网络通信和主机或数据加载空闲时间的细分 矩阵乘法单元 (MXU)、TensorCore 和内存带宽利用率，以及内置 XProf 建议
Roofline 和最慢的操作	get_roofline_model，get_hlo_op_profile	这些工具会在广泛的数据收集和重点 AI 调查期间运行，以检查以下内容： 工作负载或特定操作是否受计算或内存带宽限制 按操作类型（例如矩阵乘法、注意力机制和集合操作）划分的用时，以及指向 Python 源代码行的链接
编译器低效检测器	detect_layout_mismatch_copies、detect_unfused_reshapes、detect_unnecessary_convert_reduce、list_hlo_modules	这些工具会在广泛的数据收集期间运行，以检查以下内容： 操作之间的张量布局不匹配导致额外的数据副本 未融合的 reshape、transpose 或 copy 操作，会浪费内存带宽 归约内部不必要的精度转换，例如 bf16 到 f32
内存和 OOM 风险	get_memory_profile，get_peak_allocations	这些工具会在广泛的数据收集期间运行，以检查以下内容： 设备上的高带宽内存 (HBM) 用量峰值与总容量和内存碎片化的对比 内存不足 (OOM) 风险是LOW、MEDIUM还是HIGH，以及内存中最大的张量
计算图探索器	get_hlo_neighborhood	此工具在 AI 重点调查期间运行，用于检查 XLA 图中慢速操作之前和之后的操作，并确定哪些因素会阻止编译器融合或强制数据移动。
时间轴和内核计时	aggregate_xplane_events，get_kernel_stats	这些工具会在 AI 专项调查期间运行，以检查以下内容： 跨主机的时间轴统计信息，用于区分数据加载器停滞与网络通信延迟 各个设备内核的执行时间和调用次数
自定义内核性能分析器	llo_load，llo_query	这些工具在 AI 调查期间运行，用于检查自定义内核（例如 Pallas 或 Mosaic）的低级执行轨迹，这些内核在标准高级性能分析中不透明。

AI 受众群体特征分析报告

当捕获的分析数据的分析完成后，机器学习诊断会将生成的 summary.html 报告保存到您的 Cloud Storage 存储桶，并将其显示在 Google Cloud 控制台中的运行诊断 > 分析数据下。您可以使用此报告查看工作负载的整体运行状况、检查主要性能瓶颈，并确定建议的代码或配置更改。

每份生成的报告都包含以下部分：

• 分析摘要：显示 OPTIMAL 或 NEEDS_ATTENTION 的总体状态、3 到 5 句话的摘要，以及关键效果指标，例如步长时间、内存用量和 Roofline 效率。
• 按优先级排列的瓶颈和可采取行动的修复措施：列出按优先级排列的发现结果，从 P0（严重问题）到 P2（次要问题）。每项发现结果都包含：
  • What's happening：对根本原因的解释，包括张量形状和布局（如果相关）。
  • 证据：直接引用测量瓶颈的 XProf 指标和操作。
  • 预期影响和修复：受影响的步长时间或内存量，以及建议的框架、分片或编译器更改。
• 详细的 XProf 工具分析：提供以下分析的支持数据和详细细分：
  • Roofline 分析
  • 内存分析
  • 高级优化器 (HLO) 操作性能分析
  • 内核执行统计信息
  • 自定义内核和低级优化器分析