PyTorch XLA 性能剖析

概览

本指南将指导您使用 Cloud TPU 性能工具和 PyTorch 指标自动分析功能。这些工具可帮助您调试和优化训练工作负载的性能。

概念

如果您刚接触 PyTorch/XLA,请参阅 PyTorch API_GUIDE问题排查文档。对于 Cloud TPU,请参阅概念文档

TPU 节点 + PyTorch/XLA 性能剖析

创建和初始化所需的 Cloud TPU 相关资源。

  1. 为您的项目 ID、Cloud Storage 存储分区和用于 TPU 资源的可用区创建变量。

    export PROJECT_ID=PROJECT_ID
    export BUCKET_NAME=BUCKET_NAME
    export ZONE=ZONE
    
    gcloud --project=$PROJECT_ID compute project-info add-metadata \
    --metadata BUCKET_NAME=$BUCKET_NAME
    
  2. 创建一个 Compute Engine 虚拟机实例。 这是存储您的所有 Python 脚本和模型的地方。

    gcloud compute instances create profiler-tutorial-vm \
      --project=${PROJECT_ID} \
      --zone=${ZONE} \
      --machine-type=n1-standard-16 \
      --image-project=ml-images \
      --image-family=torch-xla \
      --boot-disk-size=300GB \
      --scopes=https://www.googleapis.com/auth/cloud-platform
    
  3. 创建 TPU 资源。

    gcloud compute tpus create profiler-tutorial-tpu \
     --project=${PROJECT_ID} \
     --zone=${ZONE} \
     --network=default \
     --version=pytorch-1.8 \
     --accelerator-type=v3-8
    
  4. 创建 Cloud Storage 存储分区。

    如果您尚未安装 gsutil CLI,请进行安装:安装说明

  5. 使用 gsutil mb 命令创建用于存储所有性能剖析工件的 Cloud Storage 存储分区。将 region 和 bucket-name 变量替换为您将用于训练的值。

    gsutil mb -p ${PROJECT_ID} -c standard -l REGION gs://${BUCKET_NAME}
    

    其中:

    • REGION 是您创建 Cloud TPU 的区域,例如 europe-west4
  6. 为 Cloud TPU 项目创建服务帐号。

    gcloud beta services identity create --service tpu.googleapis.com --project $PROJECT_ID
    

    该命令将返回以下格式的 Cloud TPU 服务帐号:

    service-PROJECT_NUMBER@cloud-tpu.iam.gserviceaccount.com
    

    例如:service-164006649440@cloud-tpu.iam.gserviceaccount.com

  7. 导出服务帐号并向服务帐号授予存储分区权限。将 account-number 替换为在服务帐号创建输出中返回的 PROJECT_NUMBER。

    export SERVICE_ACCOUNT=service-ACCOUNT_NUMBER@cloud-tpu.iam.gserviceaccount.com
    gsutil acl ch -u $SERVICE_ACCOUNT:READER gs://${BUCKET_NAME}
    gsutil acl ch -u $SERVICE_ACCOUNT:WRITER gs://${BUCKET_NAME}
    

设置 TensorBoard

  1. 通过 ssh 连接到您的虚拟机,将虚拟机的端口 9001 转发到本地机器上的端口 9001。此端口用于在本地浏览器中打开 TensorBoard 界面。

      gcloud compute ssh profiler-tutorial-vm \
       --project ${PROJECT_ID} \
       --zone ${ZONE} \
       --ssh-flag="-4 -L 9001:localhost:9001"
    
  2. 创建专用于 TensorBoard 安装的 conda 环境:

    conda create -y -n tensorboard python=3.6
    conda activate tensorboard
    pip install tf-nightly==2.6.0.dev20210511 tb-nightly==2.6.0a20210512 tbp-nightly==2.5.0a20210511
    
  3. 在 Compute Engine 虚拟机上运行 TensorBoard 服务器以测试安装,然后在本地机器上访问 http://localhost:9001/#profile 以尝试连接到服务器:

     # Get bucket name
     BUCKET_NAME=$(curl "http://metadata.google.internal/computeMetadata/v1/project/attributes/BUCKET_NAME" -H "Metadata-Flavor: Google")
    
     tensorboard --logdir gs://${BUCKET_NAME} --port 9001
    

当您在本地机器上访问 http://localhost:9001/#profile 时,应该会看到如下内容:

图片

如果您访问 http://localhost:9001,也可以访问上面的性能剖析页面,方法是选择右上角“上传”按钮旁边的“性能剖析” (PROFILE) 选项:

图片

剖析模型性能

要让 TensorBoard 服务器保持活动状态,请在本地机器上启动新的终端窗口并再次通过 ssh 连接到 GCE 虚拟机(这次不使用用于端口转发的 -L 选项)。

  1. 因为这是在新的 Shell 中,因此请在新的终端窗口中再次导出项目 ID、存储分区环境和可用区变量。

    export PROJECT_ID=PROJECT_ID
    export ZONE=ZONE
    
  2. 通过 ssh 连接到虚拟机:

      gcloud compute ssh profiler-tutorial-vm \
       --project ${PROJECT_ID} \
       --zone ${ZONE}
    
    conda activate torch-xla-1.8
    PROJECT_ID=$(curl "http://metadata.google.internal/computeMetadata/v1/project/project-id" -H "Metadata-Flavor: Google")
    export TPU_IP_ADDRESS=$(gcloud compute tpus describe profiler-tutorial-tpu --zone=${ZONE} --project=${PROJECT_ID} \
    --format="value(ipAddress)")
    echo TPU_IP_ADDRESS=${TPU_IP_ADDRESS}
    export XRT_TPU_CONFIG="tpu_worker;0;$TPU_IP_ADDRESS:8470"
    
  3. 验证集成测试能否在您的环境中端到端运行:

    python /usr/share/torch-xla-1.8/pytorch/xla/test/test_profiler.py # takes <1 min
    >
  4. 在开始训练之前,请在 /usr/share/torch-xla-1.8/pytorch/xla/test/test_profile_mp_mnist.py 中修改以下行:

    更改:

    accuracy = train_mnist(flags, dynamic_graph=True, fetch_often=True)
    
    至:
    accuracy = train_mnist(flags, dynamic_graph=False, fetch_often=False)
    

    train_mnist 中的上述两个参数会刻意导致动态图和张量提取,下文的“自动指标分析”部分将对此进行进一步介绍。目前,由于您仅剖析 TPU 的性能,以下示例以名义性能运行。

  5. 启动用于服务器性能剖析的训练运行:

    PT_XLA_DEBUG=1 XLA_HLO_DEBUG=1 python /usr/share/torch-xla-1.8/pytorch/xla/test/test_profile_mp_mnist.py --num_epochs 1000 --fake_data
    

TPU(服务器)性能剖析

训练运行后,请访问 http://localhost:9001/#profile,并按照以下说明捕获性能剖析文件:

图片

以下页面会自动重新加载:

图片

支持的工具显示在左侧窗格的工具下拉菜单中:

  • 概览页面(不包括显示 Input Pipeline Analyzer,它在 TensorFlow/TPU 和 PyTorch / XLA TPU 中完全不同)。
  • Memory Viewer
  • Op Profile
  • Pod Viewer
  • Tensorflow Stats(框架级统计信息;即 PyTorch 统计信息)
  • Trace Viewer(需要 Chrome 浏览器)

概览页面

此页面显示捕获的性能剖析文件的概览。此示例显示了很高的空闲时间,因为它是在 MNIST 数据集上训练一个很小的模型。

Memory Viewer

显示每个张量和 HLO PyTorch 操作使用的设备内存 (HBM)。Memory Viewer 会捕获每个 HLO 模块的内存视图,因此会有某些模块图,例如设备数据分配(输入和批量发送的标签)。如需查看特定 HLO 模块的内存使用量,请从左侧的主机下拉列表中选择该模块:

图片

选择 HLO 模块后,您可以查看模块执行 HBM 占用空间时间轴的全面视图,按照分配大小、程序执行大小和填充大小排序。

图片

您可以将鼠标悬停在缓冲区分配上方以查看详细信息。例如,如需了解什么分配占用了大多数设备 HBM:

图片

在上例中,(1) 与用户代码添加的 torch_xla.debug.profiler.Trace 注释对应。查看 test_profile_mp_mnist.py 代码,它对应于此行

   class MNIST(nn.Module):
   ...
     def forward(self, x):
       with xp.Trace('conv1'):
         x = F.relu(F.max_pool2d(self.conv1(x), 2))
         x = self.bn1(x)
   ...
   

此外,从 test_mnist 操作命名空间可以看出,此 HLO 模块对应于评估循环,因为它具有 xp.trace('test_mnist') 上下文管理器。

XLA Op Profile

Op Profile 是一款 Cloud TPU 工具,它可展示在性能剖析期间执行的 XLA 操作的性能统计信息。Op Profile 会显示以下内容:

  • 您的应用使用 Cloud TPU 的程度,按类别和 TPU FLOPS 利用率显示各操作所用时间的百分比。
  • 最耗时的操作。这些操作可能有待优化。
  • 各个操作的详细信息,包括形状、填充和使用该操作的表达式。

您可以使用 Op Profile 来查找适合优化的目标。例如,如果您的模型的 TPU 峰值 FLOPS 仅达到 5%,您可以使用此工具找出执行时间最长的 XLA 操作并查看它们使用了多少 TPU FLOPS。

图片

每个部分的说明:

  1. 概览部分。显示 Cloud TPU 利用率并提供优化建议。
  2. 控制面板。通过包含的控件,您可设置表格中显示的操作数量、显示哪些操作,以及这些操作的排序方式。
  3. 操作表。列出与 XLA 操作相关联的主要 TensorFlow 操作类别的表。这些操作按 Cloud TPU 使用率百分比排序。
  4. 操作详细信息卡片。将鼠标悬停在表中的某个操作上时,会出现一张卡片,上面显示有关此操作的详细信息,其中包括 FLOPS 利用率、使用该操作的表达式以及操作布局(匹配)。

Pod Viewer

如需查看 Pod Viewer 工具的完整说明,请参阅 TPU 工具

Framework Stats(Tensorflow/PyTorch 统计信息)

Framework Stats 提供在 TPU 设备和 TPU 主机上运行的 PyTorch 和 XRT 操作的统计信息明细。

Trace Viewer

Trace Viewer 是一款 Cloud TPU 性能分析工具。Trace Viewer 使用 Chrome 跟踪记录事件分析查看器,因此它需要使用 Chrome 浏览器。

Trace Viewer 包含一个时间轴,其中显示了以下内容:

  • 由 TensorFlow 模型执行的操作的时长。
  • 系统的哪个部分(TPU 还是主机)执行了操作。对于 PyTorch/XLA,通常主机主要处理编译和缓冲区分配/取消分配,而 TPU 执行实际的模型训练。
  • 您可以通过 Trace Viewer 找出模型中的性能问题,然后采取措施来加以解决。展开细目后,您可以确定哪些 PyTorch/XLA 操作执行时间最长。

请注意,您可以通过添加 xp.Trace(NAME) 直接添加跟踪记录来衡量模型某些部分的执行时长。例如,以下跟踪记录显示:

图片

  1. 由模型代码 test_profile_mp_mnist.py 中的显式用户注释生成。
  2. 执行的 PyTorch 操作(降低之前)。
  3. PyTorch/XLA 自动生成的 HLO 模块名称。
  4. 在设备上执行的 XLA 操作(融合)。

如需了解详情,请参阅 Trace Viewer 的通用 TPU 文档,但请忽略关于输入流水线和其他 特定于 TensorFlow 的部分,因为这些部分与本文内容无关。

PyTorch/XLA 客户端性能剖析

现在您将在训练期间对 PyTorch/XLA 客户端进行性能剖析,这与在模型执行期间分析 TPU 类似。在客户端使用的主要监控工具是 Trace Viewer

您必须在训练脚本中启动性能剖析服务器。例如,参阅此处,您可以从 TensorBoard 查询以捕获跟踪记录。

图片

为了捕获多个进程的跟踪记录,每个进程可以在不同的端口上启动性能剖析服务器(例如,通过向基本端口号添加 xm.get_ordinal()),然后提供由“,”连接的 localhost:port 列表。Tensorboard 不支持同时查看多个进程的跟踪记录,因此您会看到每个进程有不同的主机下拉菜单。

下图显示了一个示例跟踪记录:

图片

与为 TPU 跟踪记录添加不同的命名空间跟踪记录的方式类似,您可以使用相同的 API 将它们添加到客户端跟踪记录 (xp.Trace(NAME))。请注意,由于此模型很小,并且与较小的 MNIST 图片一起使用,因此单步用时很短并且不一定统一。作为练习,您可以尝试添加跟踪记录,启动与我们的 test_train_mp_imagenet.py --fake_data 示例类似的性能剖析器服务器,以获取 ResNet50 的跟踪记录。

跟踪记录具有其他可供查看的元数据。例如,TransferToServer 和 TransferFromServer 跟踪记录,显示发送和接收的张量的确切数量及其总大小:

图片

对于 XLA 图编译,您可以看到可帮助诊断问题的图哈希:

图片

此外,除了通过 TensorBoard 界面进行分析,我们还提供用于从 PyTorch/XLA 以编程方式分析 TPU 和客户端的 API:xp.trace()

自动指标分析

在本部分中,您将了解如何使用调试模式检测性能问题,例如:

  • 动态图/持续编译
  • 非常缓慢的图编译
  • 非常缓慢的图表执行
  • 频繁的 XLA→CPU 传输
  • 重复的设备 HBM 以托管 RAM 交换
  • 重复的 HBM 碎片整理
  • 未降低的 aten:: 操作

在开始训练之前,请还原 /usr/share/torch-xla-1.8/pytorch/xla/test/test_profile_mp_mnist.py 中的以下行:

更改:

   accuracy = train_mnist(flags, dynamic_graph=False, fetch_often=False)
   
至:
   accuracy = train_mnist(flags, dynamic_graph=True, fetch_often=True)
   

这些更改会刻意导致编译和张量提取。dynamic_graph=True 会刻意更改每一步的批次大小,使每一步的 XLA 降低图都不同,并重新编译。fetch_often=True 会在每一步插入 loss.item() 调用,导致每一步都从设备提取张量值,从而降低性能。

运行示例训练脚本:

   PT_XLA_DEBUG=1 python /usr/share/torch-xla-1.8/pytorch/xla/test/test_profile_mp_mnist.py --fake_data --num_cores=1
   

调试时,最佳做法是使用 --num_cores=1 运行,因为它可以简化调试过程。部分示例输出如下所示:

Epoch 1 train begin 01:18:05
| Training Device=xla:1/0 Step=0 Loss=0.00000 Rate=1905.00 GlobalRate=1904.72 Time=01:18:05
pt-xla-profiler: TransferFromServerTime too frequent: 3 counts during 3 steps
pt-xla-profiler: TransferFromServerTime too frequent: 4 counts during 4 steps
pt-xla-profiler: TransferFromServerTime too frequent: 5 counts during 5 steps
pt-xla-profiler: TransferFromServerTime too frequent: 6 counts during 6 steps
pt-xla-profiler: TransferFromServerTime too frequent: 7 counts during 7 steps
pt-xla-profiler: TransferFromServerTime too frequent: 8 counts during 8 steps
pt-xla-profiler: TransferFromServerTime too frequent: 9 counts during 9 steps
pt-xla-profiler: TransferFromServerTime too frequent: 10 counts during 10 steps
pt-xla-profiler: CompileTime too frequent: 21 counts during 11 steps
pt-xla-profiler: TransferFromServerTime too frequent: 11 counts during 11 steps
pt-xla-profiler: CompileTime too frequent: 23 counts during 12 steps

前缀为 pt-xla-profiler 的行对应于自动指标分析输出。在此示例中,您可以看到 TransferFromServerTime 每一步出现一次,这过于频繁。这是因为训练循环在每一步都检索 loss.item() 的值。此外,当图中有动态形状而导致模型需要重复地重新编译时,您可能会看到 CompileTime too frequent 警告。以下代码段会导致此类问题:test_profile_mp_mnist.py

    for step, (data, target) in enumerate(loader):
      if dynamic_graph:
        # The batch dimension is different every step.
        index = max(-step, -flags.batch_size + 1)  # non-empty
        data, target = data[:-index, :, :, :], target[:-index]
      ...
      if fetch_often:
        # Fetch tensor value from XLA:TPU to CPU every step.
        loss_i = loss.item()

然后,您可以按 Ctrl^C 退出训练脚本,在最后应该会看到未降低的操作的摘要。请注意,aten::_local_scalar_dense 是一种特殊的操作,对应于将 XLA 张量检索回 CPU 上下文。

在这份报告中,您可以看到调用 aten::_local_scalar_dense 操作的两个主要位置,两者都对应于 loss.item() 的源代码:

  • test/test_profile_mp_mnist.py:158
  • test/test_profile_mp_mnist.py:61
pt-xla-profiler: ================================================================================
pt-xla-profiler: Unlowered Op usage summary (more of these ops, lower performance)
pt-xla-profiler: Note: _local_scalar_dense typically indicates CPU context access
pt-xla-profiler: --------------------------------------------------------------------------------
pt-xla-profiler: FRAME (count=27):
pt-xla-profiler: Unlowered Op: "_local_scalar_dense"
pt-xla-profiler: Python Frames:
pt-xla-profiler:   train_loop_fn (test/test_profile_mp_mnist.py:158)
pt-xla-profiler:   train_mnist (test/test_profile_mp_mnist.py:184)
pt-xla-profiler:   _mp_fn (test/test_profile_mp_mnist.py:206)
pt-xla-profiler:   _start_fn (/home/jysohn/git/jysohn23/pytorch/xla/torch_xla/distributed/xla_multiprocessing.py:323)
pt-xla-profiler:   spawn (/home/jysohn/git/jysohn23/pytorch/xla/torch_xla/distributed/xla_multiprocessing.py:386)
pt-xla-profiler:    (test/test_profile_mp_mnist.py:216)
pt-xla-profiler:
pt-xla-profiler:
pt-xla-profiler: FRAME (count=2):
pt-xla-profiler: Unlowered Op: "_local_scalar_dense"
pt-xla-profiler: Python Frames:
pt-xla-profiler:   _train_update (test/test_profile_mp_mnist.py:61)
pt-xla-profiler:    (/home/jysohn/git/jysohn23/pytorch/xla/torch_xla/core/xla_model.py:700)
pt-xla-profiler:   _run_step_closures (/home/jysohn/git/jysohn23/pytorch/xla/torch_xla/core/xla_model.py:709)
pt-xla-profiler:   mark_step (/home/jysohn/git/jysohn23/pytorch/xla/torch_xla/core/xla_model.py:723)
pt-xla-profiler:   __exit__ (/home/jysohn/git/jysohn23/pytorch/xla/torch_xla/debug/profiler.py:153)
pt-xla-profiler:   train_loop_fn (test/test_profile_mp_mnist.py:162)
pt-xla-profiler:   train_mnist (test/test_profile_mp_mnist.py:184)
pt-xla-profiler:   _mp_fn (test/test_profile_mp_mnist.py:206)
pt-xla-profiler:   _start_fn (/home/jysohn/git/jysohn23/pytorch/xla/torch_xla/distributed/xla_multiprocessing.py:323)
pt-xla-profiler:   spawn (/home/jysohn/git/jysohn23/pytorch/xla/torch_xla/distributed/xla_multiprocessing.py:386)
pt-xla-profiler:    (test/test_profile_mp_mnist.py:216)
pt-xla-profiler:
pt-xla-profiler:
pt-xla-profiler: ================================================================================

现在,对以下脚本运行自动指标分析,该脚本包含未降低的操作(从我们的 1.8 版本开始),_ctc_loss

PT_XLA_DEBUG=1 python <<EOF
import torch
import torch_xla.core.xla_model as xm

dev = xm.xla_device()
t = torch.randn(50, 16, 20).log_softmax(2).to(dev)
target = torch.randint(low=1, high=20, size=(16, 30), dtype=torch.long).to(dev)
input_lengths = torch.full(size=(16,), fill_value=50, dtype=torch.long).to(dev)
target_lengths = torch.randint(low=10, high=30, size=(16,), dtype=torch.long).to(dev)

for _ in range(10):
  loss = torch.nn.CTCLoss()(t, target, input_lengths, target_lengths)
  xm.mark_step()
EOF

使用 PT_XLA_DEBUG=1 运行上述脚本,输出应如下所示:

…
pt-xla-profiler: TransferFromServerTime too frequent: 30 counts during 10 steps
pt-xla-profiler: Op(s) not lowered: aten::_ctc_loss,  Please open a GitHub issue with the above op lowering requests.
pt-xla-profiler: ================================================================================
pt-xla-profiler: Unlowered Op usage summary (more of these ops, lower performance)
pt-xla-profiler: Note: _local_scalar_dense typically indicates CPU context access
pt-xla-profiler: --------------------------------------------------------------------------------
pt-xla-profiler: FRAME (count=10):
pt-xla-profiler: Unlowered Op: "_ctc_loss"
pt-xla-profiler: Python Frames:
pt-xla-profiler:   ctc_loss (/anaconda3/envs/torch-xla-1.8/lib/python3.6/site-packages/torch/nn/functional.py:2305)
pt-xla-profiler:   forward (/anaconda3/envs/torch-xla-1.8/lib/python3.6/site-packages/torch/nn/modules/loss.py:1593)
pt-xla-profiler:   _call_impl (/anaconda3/envs/torch-xla-1.8/lib/python3.6/site-packages/torch/nn/modules/module.py:889)
pt-xla-profiler:    (:11)
pt-xla-profiler:
pt-xla-profiler:
pt-xla-profiler: ================================================================================

自动指标分析器显示 STDIN 的第 11 行导致了此未降低的操作(即具有 torch.nn.CTCLoss() 的行)。目前 ctc_loss 操作尚未降低,因此您会看到上述报告。您还可以看到针对 TransferFromServerTime 的警告,因为在执行前,张量最初在 XLA:TPU 中,但由于该操作未降低,因此您首先需要将 XLA 张量传输回 CPU,在 CPU 上执行 aten:: 操作,然后将其传回。

如果您想将 pt-xla-profiler 输出写入文件,请设置 PT_XLA_DEBUG=1PT_XLA_DEBUG_FILE=$PATH_TO_FILE

清理

通过运行以下命令退出虚拟机,然后删除 TPU、虚拟机和 Cloud Storage 存储分区:

(vm)$ exit
gcloud compute instances delete profiler-tutorial-vm \
  --zone=${ZONE} \
  --project=${PROJECT_ID}
gcloud compute tpus delete profiler-tutorial-tpu \
  --zone=${ZONE} \
  --project=${PROJECT_ID} \
  --async
gsutil rm -fr gs://${BUCKET_NAME}

TPU 虚拟机 + PyTorch/XLA 性能剖析

本部分使用 TPU 虚拟机架构来分析 PyTorch/XLA 性能。

导出环境变量

  1. 为项目 ID 和用于 TPU 资源的可用区创建变量。

    export PROJECT_ID=PROJECT_ID
    export ZONE=ZONE
    

创建 Cloud TPU

请参阅 TPU 虚拟机用户指南,设置完成后,创建一个预装了 torchtorch_xlatorchvisiontensorboard 的 v3-8 TPU 虚拟机。

  1. 创建 TPU 资源。

    gcloud alpha compute tpus tpu-vm create profiler-tutorial-tpu-vm \
     --project=${PROJECT_ID} \
     --zone=${ZONE} \
     --version=v2-alpha \
     --accelerator-type=v3-8
    

Tensorboard 服务器启动

  1. 通过 SSH 连接到虚拟机,安装 tensorboard-plugin-profile,然后启动 TensorBoard 服务器。

      gcloud alpha compute tpus tpu-vm ssh profiler-tutorial-tpu-vm \
       --project ${PROJECT_ID} \
       --zone ${ZONE} \
       --ssh-flag="-4 -L 9001:localhost:9001"
    
      pip3 install tf-nightly==2.6.0.dev20210511 tb-nightly==2.6.0a20210512 tbp-nightly==2.5.0a20210511
    
      tensorboard --logdir ./tensorboard --port 9001
    

当您在本地机器上访问 http://localhost:9001 查看 TensorBoard 输出时,应该会看到如下内容:

图片

如果您在 http://localhost:9001 查看 TensorBoard 输出,则还可以通过选择右上角“上传”按钮旁边的下拉菜单中的“性能剖析” (PROFILE) 选项来访问上面的性能剖析页面:

图片

剖析模型性能

开发环境中的新终端窗口中,导出与上文相同的环境变量并同通过 ssh 连接到您的 TPU 虚拟机:

  1. 在新终端窗口中,再次导出项目 ID 和可用区变量,因为这是在新的 shell 中。

    export PROJECT_ID=PROJECT_ID
    export ZONE=ZONE
    
  2. 通过 ssh 连接到虚拟机:

      gcloud alpha compute tpus tpu-vm ssh profiler-tutorial-tpu-vm \
       --project ${PROJECT_ID} \
       --zone ${ZONE}
    
  3. 克隆 PyTorch/XLA 代码库并运行我们的 e2e 测试:

      git clone -b r1.8 https://github.com/pytorch/xla
      export XRT_TPU_CONFIG="localservice;0;localhost:51011"
      python3 xla/test/test_profiler.py  # takes <1 min
    >
  4. 在开始训练之前,请修改 xla/test/test_profile_mp_mnist.py 中的以下行:

    更改:

        accuracy = train_mnist(flags, dynamic_graph=True, fetch_often=True)
    
    至:
        accuracy = train_mnist(flags, dynamic_graph=False, fetch_often=False)
    

    train_mnist 中的上述两个参数会刻意导致动态图和张量提取,下文的“自动指标分析”部分将对此进行进一步介绍。目前,由于您仅剖析 TPU 的性能,以下示例以名义性能运行。

  5. 启动训练运行:

     XLA_HLO_DEBUG=1 python3 xla/test/test_profile_mp_mnist.py --num_epochs 1000 --fake_data
    

TPU + 客户端性能剖析

训练运行后,访问 http://localhost:9001 查看 TensorBoard 输出,并按照以下说明捕获性能剖析文件:

图片

您应该会看到重新加载的以下页面:

图片

目前在 TPU 虚拟机设置中,仅 Trace Viewer 工具被选中,因此在工具下拉列表下选择 trace_viewer 并检查检查跟踪记录。您可以看到,在 TPU 虚拟机设置中,您可以在一个完整视图中同时看到“客户端”侧和 TPU 设备侧的跟踪记录:

图片

清理

  1. 通过运行以下命令退出虚拟机,然后删除 TPU、虚拟机和 Cloud Storage 存储分区:

    (vm)$ exit
    

删除您创建的 TPU 虚拟机:

  1. 删除您的 Cloud TPU 和 Compute Engine 资源。

    $ gcloud alpha compute tpus tpu-vm delete profiler-tutorial-tpu-vm \
      --project ${PROJECT_ID} --zone=${ZONE}
    
  2. 通过运行以下命令来验证资源已删除。删除操作可能需要几分钟时间才能完成。如下所示的响应表明实例已成功删除。

    $ gcloud alpha compute tpus tpu-vm list --project ${PROJECT_ID} --zone=${ZONE}
    
    Listed 0 items.