对 PyTorch XLA 工作负载进行性能分析

性能分析是一种用于分析和提高模型性能的方法。虽然还有更多内容，但有时将性能分析视为用时操作以及在设备（TPU）和主机（CPU）上运行的代码的一部分是有帮助的。本指南简要介绍了如何对代码进行性能分析，以用于训练或推理。如需详细了解如何分析生成的性能分析文件，请参阅以下指南。

创建 TPU

导出环境变量：

export PROJECT_ID=your-project-id
export TPU_NAME=your-tpu-name
export ZONE=us-central2-b
export ACCELERATOR_TYPE=v4-8
export RUNTIME_VERSION=tpu-vm-v4-pt-2.0

环境变量说明

变量	说明
`PROJECT_ID`	您的 Google Cloud 项目 ID。使用现有项目或创建新项目。
`TPU_NAME`	TPU 的名称。
`ZONE`	要在其中创建 TPU 虚拟机的可用区。如需详细了解支持的可用区，请参阅 TPU 区域和可用区。
`ACCELERATOR_TYPE`	加速器类型用于指定您要创建的 Cloud TPU 的版本和大小。如需详细了解每个 TPU 版本支持的加速器类型，请参阅 TPU 版本。
`RUNTIME_VERSION`	Cloud TPU 软件版本。

启动 TPU 资源

$ gcloud compute tpus tpu-vm create ${TPU_NAME} \
    --zone ${ZONE} \
    --accelerator-type ${ACCELERATOR_TYPE} \
    --version ${RUNTIME_VERSION} \
    --project ${PROJECT_ID}

使用以下命令，在 TPU 切片中的所有 TPU 虚拟机上安装 torch_xla。您还需要安装训练脚本所需的任何其他依赖项。

gcloud compute tpus tpu-vm ssh ${TPU_NAME} \
    --zone=${ZONE} \
    --project=${PROJECT_ID} \
    --worker=all \
    --command="pip install torch==2.6.0 torch_xla[tpu]==2.6.0 torchvision -f https://storage.googleapis.com/libtpu-releases/index.html -f https://storage.googleapis.com/libtpu-wheels/index.html"

使用 gcloud scp 命令将代码移至 TPU 虚拟机上的主目录。例如：

$ gcloud compute tpus tpu-vm scp my-code-file ${TPU_NAME}:directory/target-file --zone ${ZONE}

性能分析文件

您可以通过 capture_profile.py 手动捕获性能分析文件，也可以在训练脚本中使用 torch_xla.debug.profiler API 以程序化方式捕获性能分析文件。

启动性能分析服务器

为了捕获性能分析文件，必须在训练脚本中运行性能分析服务器。使用您选择的端口号启动服务器，例如 9012，如以下命令所示。

import torch_xla.debug.profiler as xp
server = xp.start_server(9012)

服务器可以直接在 main 函数的开头启动。

现在，您可以按照下一部分中所述捕获性能分析文件。该脚本会对一个 TPU 设备上发生的所有操作进行性能分析。

添加跟踪记录

如果您还想对宿主机上的操作进行性能分析，可以在代码中添加 xp.StepTrace 或 xp.Trace。这些函数会跟踪宿主机上的 Python 代码。（您可以将其视为测量在将“图”传递给 TPU 设备之前，在主机 [CPU] 上执行 Python 代码所需的时间。因此，它主要用于分析跟踪记录开销）。您可以在代码处理批量数据的训练循环中添加此代码，例如：

for step, batch in enumerate(train_dataloader):
   with xp.StepTrace('Training_step', step_num=step):
   ...

或使用以下命令封装代码的各个部分

with xp.Trace('loss'):
   loss = ...

如果您使用的是 Lighting，则可以跳过添加跟踪记录，因为代码的某些部分会自动执行此操作。不过，如果您想添加其他跟踪记录，建议您将其插入训练循环中。

您将能够在初始编译后捕获设备活动；等待模型开始其训练或推理步骤。

手动捕获

来自 Pytorch XLA 仓库的 capture_profile.py 脚本可用于快速捕获性能分析文件。为此，您可以直接将捕获性能分析文件复制到 TPU 虚拟机。以下命令会将其复制到主目录。

$ gcloud compute tpus tpu-vm ssh ${TPU_NAME} \
    --zone us-central2-b \
    --worker=all \
    --command="wget https://raw.githubusercontent.com/pytorch/xla/master/scripts/capture_profile.py"

在训练运行期间，执行以下命令以捕获性能分析文件：

$ gcloud compute tpus tpu-vm ssh ${TPU_NAME} \
    --zone us-central2-b \
    --worker=all \
    --command="python3 capture_profile.py --service_addr "localhost:9012" --logdir ~/profiles/ --duration_ms 2000"

此命令会将 .xplane.pb 文件保存在 logdir 中。您可以将日志记录目录 ~/profiles/ 更改为您偏好的位置和名称。您也可以直接将该文件保存在 Cloud Storage 存储桶中。为此，请将 logdir 设置为 gs://your_bucket_name/。

程序化捕获

您可以将训练脚本配置为使用训练脚本中的 torch_xla.debug.profiler.trace_detached API 自动触发性能分析，而不是通过触发脚本手动捕获性能分析文件。

例如，如需在特定周期和步骤自动捕获性能分析文件，您可以将训练脚本配置为使用 PROFILE_STEP、PROFILE_EPOCH 和 PROFILE_LOGDIR 环境变量：

import os
import torch_xla.debug.profiler as xp

# Within the training script, read the step and epoch to profile from the
# environment.
profile_step = int(os.environ.get('PROFILE_STEP', -1))
profile_epoch = int(os.environ.get('PROFILE_EPOCH', -1))
...

for epoch in range(num_epoch):
   ...
   for step, data in enumerate(epoch_dataloader):
      if epoch == profile_epoch and step == profile_step:
         profile_logdir = os.environ['PROFILE_LOGDIR']
         # Use trace_detached to capture the profile from a background thread
         xp.trace_detached('localhost:9012', profile_logdir)
      ...

这会将 .xplane.pb 文件保存在由 PROFILE_LOGDIR 环境变量指定的目录中。

TensorBoard 中的分析

如需进一步分析性能分析文件，您可以在同一机器或其他机器（推荐）上将 TensorBoard 与 TPU TensorBoard 插件结合使用。

如需在远程机器上运行 TensorBoard，请使用 SSH 连接到该机器并启用端口转发。例如，

$ ssh -L 6006:localhost:6006 remote server address

或

$ gcloud compute tpus tpu-vm ssh ${TPU_NAME} --zone=${ZONE} --ssh-flag="-4 -L 6006:localhost:6006"

在远程机器上，安装所需的软件包并启动 TensorBoard（假设您在该机器上的 ~/profiles/ 下有性能分析文件）。如果您将性能分析文件存储在其他目录或 Cloud Storage 存储桶中，请务必正确指定路径，例如 gs://your_bucket_name/profiles。

(vm)$ pip install tensorflow-cpu tensorboard-plugin-profile

(vm)$ tensorboard --logdir ~/profiles/ --port 6006

(vm)$ pip uninstall tensorflow tf-nightly tensorboard tb-nightly tbp-nightly

运行 TensorBoard

在本地浏览器中，访问 http://localhost:6006/，然后从下拉菜单中选择 PROFILE 以加载性能分析文件。

如需了解 TensorBoard 工具以及如何解读输出，请参阅在 Cloud TPU 虚拟机上对模型进行性能分析。

TensorBoard 性能分析文件捕获页面