分析 PyTorch XLA 工作负载

性能分析是分析和提升模型性能的一种方式。尽管它的功能远不止于此，但有时将分析视为在设备 (TPU) 和主机 (CPU) 上运行的计时操作和代码部分会很有帮助。本指南简要介绍了如何分析代码以进行训练或推断。如需详细了解如何分析生成的配置文件，请参阅以下指南。

• 在 TPU 虚拟机上调试 PyTorch XLA 性能 - 第 1 部分
• 在 TPU 虚拟机上调试 PyTorch XLA 性能 - 第 2 部分
• 在 TPU 虚拟机上调试 PyTorch XLA 性能 - 第 3 部分

开始使用

创建 TPU

1. 导出环境变量：

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

   导出变量说明

   TPU name

   您要用于 Cloud TPU 的名称。

   zone

   您计划在其中创建 Cloud TPU 的区域。

   project ID

   您用于训练和剖析模型的项目 ID。

   accelerator-type

   加速器类型指定要创建的 Cloud TPU 的版本和大小。如需详细了解每个 TPU 版本支持的加速器类型，请参阅 TPU 版本。

   version

   Cloud TPU 运行时版本。导出的变量中会显示默认值，但您也可以使用受支持配置列表中的某个配置。

2. 启动 TPU 资源

   $ gcloud compute tpus tpu-vm create ${TPU_NAME} \
   --zone us-central2-b \
   --accelerator-type ${ACCELERATOR_TYPE} \
   --version ${RUNTIME_VERSION} \
   --project $PROJECT_ID \
   --subnetwork=tpusubnet
   

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

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

分析

注意 ：v4-8 包含 4 个 TPU 设备。每个工作器仅分析 4 个可用设备中的 1 个 TPU 设备。我们很快将支持为每个工作器对所有 4 个 TPU 设备进行性能分析。

您可以通过 capture_profile.py 手动捕获配置文件，也可以使用 torch_xla.debug.profiler API 在训练脚本中以编程方式捕获配置文件。

启动配置文件服务器

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

注意 ：您可以参考此示例脚本，其中包含以下步骤。

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

该服务器可在 main 函数的开头启动。

注意 ：如果您使用的是 Lightning 模块，请不要启动服务器，而是将 profiler="xla" 传递给 trainer，它会自动在端口 9012 启动服务器。

您现在可以捕获配置文件，如下一部分所述。该脚本会分析一台 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 = ...

重要提示： 请避免使用 xp.Trace 封装 xm.mark_step()。在这种情况下，代码将崩溃，并显示类似于 "RuntimeError: Expecting scope to be empty but it is train_loop." xp.StepTrace 的错误消息，这表示代码不会出现此问题，因为它在退出作用域时自动添加 xm.mark_step()。此外，使用 torchdynamo 时，您可能会遇到此错误，因为 torchdynamo 可能会在内部调用 xm.mark_step。当前的权宜解决方法不是使用 xp.Trace。

如果您使用的是光照，则可以跳过添加轨迹的步骤，因为代码的某些部分会自动完成添加。但是，如果您想添加其他跟踪记录，欢迎将它们插入到训练循环中。

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

手动捕获

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 自动触发配置文件，而不是通过触发脚本手动捕获配置文件。

注意 ：trace_detached API 仅适用于每夜 build 或 2.2 版本之后的版本。如需在早期的 torch_xla 版本中以程序化方式捕获数据，您可以复制方法的正文以直接创建后台捕获线程。

例如，如需在特定周期和步骤自动捕获配置文件，您可以将训练脚本配置为使用 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 连接到 TensorBoard 并启用端口转发。例如，

$ 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

注意 ：.xplane 文件必须置于捕获配置文件时生成的同一嵌套结构。例如：~/profiles/plugins/profile/2023_04_10_21_40_22/localhost_9012.xplane.pb

注意 ：如果您遇到“重复的插件”错误，可能是因为 TensorBoard 有多个软件包。卸载所有这些软件包，然后再次运行 pip install 命令。例如，您可能需要运行以下命令：

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

运行 TensorBoard

在本地浏览器中，转到 http://localhost:6006/，然后从下拉菜单中选择 PROFILE 以加载配置文件。

如需了解 TensorBoard 工具以及如何解读输出，请参阅 TPU 工具。