Cloud TPU PyTorch/XLA 用户指南

重要提示 ：您可以使用两种不同架构的 TPU：TPU 节点和 TPU 虚拟机。本教程假定您使用的是 TPU 虚拟机。系统架构中描述了这两种虚拟机架构。除非您知道自己需要使用 TPU 节点，否则我们建议您使用 TPU 虚拟机。如需详细了解如何在使用 TPU 节点在 TPU Pod 上使用 PyTorch，请参阅在 Cloud TPU Pod 上训练 PyTorch 模型。

使用 PyTorch/XLA 运行机器学习工作负载

本指南将介绍如何使用 PyTorch 在 v4 TPU 上执行简单的计算。

基本设置

1. 创建一个运行 Pytorch 2.0 的 TPU 虚拟机运行时的 v4 TPU 的 TPU 虚拟机：

     gcloud compute tpus tpu-vm create your-tpu-name \
     --zone=us-central2-b \
     --accelerator-type=v4-8 \
     --version=tpu-vm-v4-pt-2.0

   注意 ：如果您想安装其他版本的 PyTorch，请将 tpu-vm-v4-pt-2.0 替换为您要使用的 TPU 运行时版本。如需查看使用 PyTorch 的 TPU 运行时的列表，请参阅支持 PyTorch 的版本。

2. 使用 SSH 连接到 TPU 虚拟机：

     gcloud compute tpus tpu-vm ssh your-tpu-name \
     --zone=us-central2-b \
     --accelerator-type=v4-8

   要点 ：在本指南中，(vm) $ 前缀表示您应该在 TPU 虚拟机实例上运行该命令。

3. 设置 PJRT 或 XRT TPU 设备配置。

   注意 ：有两种 PyTorch/XLA 运行时选项：PJRT 和 XRT。除非有理由使用 XRT，否则我们建议您使用 PJRT。如需详细了解不同的运行时配置，并且您有使用 XRT 的理由，请参阅 PJRT 运行时文档。

   印度卢比

       (vm)$ export PJRT_DEVICE=TPU
    

   XRT

       (vm)$ export XRT_TPU_CONFIG="localservice;0;localhost:51011"
    

4. 如果要使用 Cloud TPU v4 进行训练，请同时设置以下环境变量：

     (vm)$ export TPU_NUM_DEVICES=4

执行简单的计算

1. 在 TPU 虚拟机上启动 Python 解释器：

   (vm)$ python3

2. 导入以下 PyTorch 软件包：

   import torch
   import torch_xla.core.xla_model as xm

3. 输入以下脚本：

   dev = xm.xla_device()
   t1 = torch.randn(3,3,device=dev)
   t2 = torch.randn(3,3,device=dev)
   print(t1 + t2)

   注意：OpKernel ('op: "TPURoundRobin" device_type: "CPU"') for unknown op 错误消息的出现符合预期，将在未来版本中更正。

   输出如下所示：

   tensor([[-0.2121,  1.5589, -0.6951],
          [-0.7886, -0.2022,  0.9242],
          [ 0.8555, -1.8698,  1.4333]], device='xla:1')
   

在单设备 TPU 上运行 Resnet

此时，您可以运行您喜欢的任何 PyTorch / XLA 代码。例如，您可以使用虚构数据运行 ResNet 模型：

(vm)$ git clone --recursive https://github.com/pytorch/xla.git
(vm)$ python3 xla/test/test_train_mp_imagenet.py --fake_data --model=resnet50 --num_epochs=1

ResNet 示例训练 1 个周期，大约需要 7 分钟。它将返回如下所示的输出：

Epoch 1 test end 20:57:52, Accuracy=100.00 Max Accuracy: 100.00%

ResNet 训练结束后，删除 TPU 虚拟机。

(vm)$ exit

$ gcloud compute tpus tpu-vm delete tpu-name \
--zone=zone

删除操作可能需要几分钟时间才能完成。通过运行 gcloud compute tpus list --zone=${ZONE} 验证资源是否已删除。

高级主题

对于具有可扩缩、频繁分配的模型，与 C/C++ 运行时函数 malloc 相比，tcmalloc 可以提高性能。TPU 虚拟机上使用的默认 malloc 为 tcmalloc。您可以通过取消设置 LD_PRELOAD 环境变量，强制 TPU 虚拟机软件使用标准 malloc：

   (vm)$ unset LD_PRELOAD

在前面的示例中（简单计算和 ResNet50），PyTorch/XLA 程序在与 Python 解释器相同的过程中启动本地 XRT 服务器。您还可以选择在单独的进程中启动 XRT 本地服务：

(vm)$ python3 -m torch_xla.core.xrt_run_server --port 51011 --restart

注意：前面指定的端口与您在 export XRT_TPU_CONFIG="localservice;0;localhost:51011" 命令中指定的端口相同。

这种方法的优点是，编译缓存在训练运行过程中保持不变。在单独的进程中运行 XLA 服务器时，服务器端日志记录信息将写入 /tmp/xrt_server_log。

(vm)$ ls /tmp/xrt_server_log/

server_20210401-031010.log

TPU 虚拟机性能分析

如需详细了解如何在 TPU 虚拟机上剖析您的模型，请参阅 PyTorch XLA 性能剖析。

PyTorch/XLA TPU Pod 示例

如需了解在 TPU 虚拟机 Pod 上运行 PyTorch/XLA 的设置信息和示例，请参阅 PyTorch TPU 虚拟机 Pod。

TPU 虚拟机上的 Docker

本部分介绍如何在预安装了 PyTorch/XLA 的 TPU 虚拟机上运行 Docker。

注意 ：目前，只有单个设备 TPU（v2-8、v3-8）支持在 TPU 虚拟机上运行 Docker。

可用的 Docker 映像

您可以参考 GitHub README 文件来查找所有可用的 TPU 虚拟机 Docker 映像。

在 TPU 虚拟机上运行 Docker 映像

注意 ：最后一个命令输出 root:/# 是您在执行 Docker 映像后看到的提示。

(tpuvm): sudo docker pull gcr.io/tpu-pytorch/xla:nightly_3.8_tpuvm
(tpuvm): sudo docker run --privileged  --shm-size 16G --name tpuvm_docker -it -d  gcr.io/tpu-pytorch/xla:nightly_3.8_tpuvm
(tpuvm): sudo docker exec --privileged -it tpuvm_docker /bin/bash
(pytorch) root:/#

注意 ：您必须使用 --privileged 标志运行容器，以确保 TPU 设备在容器中可见。此容器中预安装了 pytorch conda 环境。以 (pytorch) 开头的每个命令都应在容器中运行。

验证 libtpu

注意：libtpu 应该预安装在较新的 TPUVM Docker 映像上，但手动验证始终是更好的方法。

如需验证是否已安装 libtpu，请运行以下命令：

(pytorch) root:/# ls /root/anaconda3/envs/pytorch/lib/python3.8/site-packages/ | grep libtpu

这应该会生成类似于以下内容的输出：

libtpu
libtpu_nightly-0.1.dev20220518.dist-info

如果未显示任何结果，您可以使用以下命令手动安装相应的 libtpu：

(pytorch) root:/# pip install torch_xla[tpuvm]

验证tcmalloc

tcmalloc 是我们在 TPU 虚拟机上使用的默认 malloc。如需了解详情，请参阅此部分。此库应预安装在较新的 TPU 虚拟机 Docker 映像上，但手动验证始终是更好的做法。您可以运行以下命令，验证是否已安装该库。

(pytorch) root:/# echo $LD_PRELOAD

这应该会生成类似于以下内容的输出：

/usr/lib/x86_64-linux-gnu/libtcmalloc.so.4

如果未设置 LD_PRELOAD，您可以手动运行：

(pytorch) root:/# sudo apt-get install -y google-perftools
(pytorch) root:/# export LD_PRELOAD="/usr/lib/x86_64-linux-gnu/libtcmalloc.so.4"

验证设备

您可以通过运行以下命令来验证 TPU 虚拟机是否可用：

(pytorch) root:/# ls /dev | grep accel

这应该会生成以下结果。

accel0
accel1
accel2
accel3

如果未显示任何结果，则很可能是因为您没有使用 --privileged 标志启动容器。

运行模型

您可以通过运行以下命令来验证 TPU 虚拟机是否可用：

(pytorch) root:/# export XRT_TPU_CONFIG="localservice;0;localhost:51011"
(pytorch) root:/# python3 pytorch/xla/test/test_train_mp_imagenet.py --fake_data --num_epochs 1