Cloud TPU PyTorch/XLA 用户指南

使用 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
    
  2. 使用 SSH 连接到 TPU 虚拟机:

      gcloud compute tpus tpu-vm ssh your-tpu-name \
      --zone=us-central2-b \
      --accelerator-type=v4-8
    
  3. 设置 PJRT 或 XRT TPU 设备配置。

    印度卢比

        (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)
    

    输出如下所示:

    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 虚拟机上使用的默认 malloctcmalloc。您可以通过取消设置 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

这种方法的优点是,编译缓存在训练运行过程中保持不变。在单独的进程中运行 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。

可用的 Docker 映像

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

在 TPU 虚拟机上运行 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:/#

验证 libtpu

如需验证是否已安装 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