Cloud TPU PyTorch/XLA 用户指南
使用 PyTorch/XLA 运行机器学习工作负载
本指南将介绍如何使用 PyTorch 在 v4 TPU 上执行简单的计算。
基本设置
创建一个运行 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
使用 SSH 连接到 TPU 虚拟机:
gcloud compute tpus tpu-vm ssh your-tpu-name \ --zone=us-central2-b \ --accelerator-type=v4-8
设置 PJRT 或 XRT TPU 设备配置。
印度卢比
(vm)$ export PJRT_DEVICE=TPU
XRT
(vm)$ export XRT_TPU_CONFIG="localservice;0;localhost:51011"
如果要使用 Cloud TPU v4 进行训练,请同时设置以下环境变量:
(vm)$ export TPU_NUM_DEVICES=4
执行简单的计算
在 TPU 虚拟机上启动 Python 解释器:
(vm)$ python3
导入以下 PyTorch 软件包:
import torch import torch_xla.core.xla_model as xm
输入以下脚本:
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 虚拟机上使用的默认 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
这种方法的优点是,编译缓存在训练运行过程中保持不变。在单独的进程中运行 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