分析 PyTorch XLA 工作负载
性能分析是分析和改进模型性能的一种方式。虽然还有很多其他内容,但有时将性能分析视为对设备(TPU)和主机(CPU)上运行的操作和代码部分进行时间测算会很有帮助。本指南简要介绍了如何对代码进行性能分析,以便进行训练或推理。如需详细了解如何分析生成的配置文件,请参阅以下指南。
- TPU 虚拟机上的 PyTorch XLA 性能调试 - 第 1 部分
- TPU 虚拟机上的 PyTorch XLA 性能调试 - 第 2 部分
- TPU 虚拟机上的 PyTorch XLA 性能调试 - 第 3 部分
开始使用
创建 TPU
导出环境变量:
$ 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 资源
$ gcloud compute tpus tpu-vm create ${TPU_NAME} \ --zone us-central2-b \ --accelerator-type ${ACCELERATOR_TYPE} \ --version ${RUNTIME_VERSION} \ --project $PROJECT_ID \ --subnetwork=tpusubnet
使用
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.StepTrace
或
xp.Trace
。这些函数会跟踪
宿主机。
(可以将其视为衡量执行 Python 代码所需的时间,
然后向主机 (CPU) 上传递“graph”连接到 TPU 设备。因此,它最适合用于分析跟踪开销)。您可以在代码处理批量数据的训练循环中添加此操作,例如:
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
以加载配置文件。
请参阅 TPU 工具,了解 TensorBoard 工具以及如何解读输出。