分析 PyTorch XLA 工作负载

性能分析是一种分析和改进模型性能的方法。虽然还有很多其他内容,但有时将性能分析视为对设备(TPU)和主机(CPU)上运行的操作和代码部分进行时间测算会很有帮助。本指南简要介绍了如何对代码进行性能分析,以便进行训练或推理。如需详细了解如何分析生成的配置文件,请参阅以下指南。

开始使用

创建 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}

分析

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

启动配置文件服务器

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

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

服务器可以在 main 函数开头立即启动。

现在,您可以按照以下部分中的说明捕获配置文件。该脚本会对一台 TPU 设备上发生的所有活动进行性能分析。

添加轨迹

如果您还想对宿主机上的操作进行性能分析,可以在代码中添加 xp.StepTracexp.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_STEPPROFILE_EPOCHPROFILE_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 中的分析

如需进一步分析配置文件,您可以在同一台机器或其他机器(推荐)上将 TensorBoardTPU 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 工具以及如何解读输出,请参阅 TPU 工具

图片

图片