使用 PyTorch 在 Cloud TPU 虚拟机上运行计算

快速入门：使用 PyTorch 在 Cloud TPU 虚拟机上运行计算

使用 PyTorch 在 Cloud TPU 虚拟机上运行计算

本快速入门介绍如何创建 Cloud TPU、如何安装 PyTorch，以及如何在 Cloud TPU 上运行简单的计算。如需更深入地了解如何在 Cloud TPU 上训练模型，请参阅 Cloud TPU PyTorch 教程之一。

准备工作

重要提示：您可以将本指南与 TPU 虚拟机或 TPU 节点配置搭配使用。系统架构中介绍了这两种虚拟机架构。您使用的 gcloud 命令取决于您使用的 TPU 配置。 在此快速入门中，每个 gcloud 命令都显示在标签式部分中。选择您要使用的 TPU 配置的标签页，网页会显示相应的 gcloud 命令。除非您知道需要使用 TPU 节点，否则我们建议您使用 TPU 虚拟机。

按照本快速入门操作之前，您必须先创建一个 Google Cloud Platform 帐号，安装 Google Cloud CLI 并配置 gcloud 命令。如需了解详情，请参阅设置帐号和 Cloud TPU 项目

使用 gcloud 创建 Cloud TPU

使用 gcloud 命令启动 Compute Engine 虚拟机和 Cloud TPU。 使用的命令取决于您使用的是 TPU 虚拟机还是 TPU 节点。如需详细了解这两种虚拟机架构，请参阅系统架构。如需详细了解 gcloud 命令，请参阅 gcloud 参考文档。

TPU 虚拟机

注意：如需确定启动 TPU 虚拟机时要指定的 TPU 软件版本，请参阅 Cloud TPU 软件版本。

如需在默认用户项目中创建 TPU 虚拟机，请运行网络和计算/区域运行：

$ gcloud compute tpus tpu-vm create tpu-name \
--zone=us-central1-b \
--accelerator-type=v3-8 \
--version=tpu-vm-pt-1.11

命令标志说明

zone

您计划在其中创建 Cloud TPU 的可用区。

accelerator-type

要创建的 Cloud TPU 的类型。

version

Cloud TPU 软件版本。

创建 TPU 时，如果要指定默认网络和子网，可以传递额外的 --network 和 --subnetwork 标志。如果您不想使用默认网络，则必须传递 --network 标志。--subnetwork 标志是可选的，可用于为您正在使用的任何网络（默认或用户指定）指定默认子网。如需详细了解这些标志，请参阅 gcloud API 参考页面。

TPU 节点

1. 为 PyTorch 创建 TPU 节点时，首先要创建一个 Compute Engine 虚拟机实例。

   gcloud compute instances create tpu-name \
   --zone=us-central1-b  \
   --machine-type=n1-standard-16  \
   --image-family=torch-xla \
   --image-project=ml-images  \
   --boot-disk-size=200GB \
   --scopes=https://www.googleapis.com/auth/cloud-platform

   命令标志说明

   project

   您的 Cloud 项目 ID

   name

   要创建的 Cloud TPU 的名称。

   zone

   您计划在其中创建 Cloud TPU 的可用区。

   disk-size

   gcloud 命令所创建虚拟机的硬盘大小（以 GB 为单位）。

   machine-type

   要创建的 Compute Engine 虚拟机的机器类型。

   tf-version

   在虚拟机上安装的 gcloud 版 TensorFlow。

   accelerator-type

   要创建的 Cloud TPU 的类型。

2. 接下来，创建 TPU 实例。

   gcloud compute tpus create tpu-name \
   --zone=us-central1-b \
   --network=default \
   --version=pytorch-1.11 \
   --accelerator-type=v3-8

   命令标志说明

   name

   要创建的 Cloud TPU 的名称。

   zone

   拟在其中创建 Cloud TPU 的可用区。

   network

   TPU 连接到的网络。

   version

   虚拟机上安装的 PyTorch gcloud 版本。

   accelerator-type

   要创建的 Cloud TPU 的类型。

注意：如果您有多个项目，则必须使用 --project 标志指定项目 ID。

连接到 Cloud TPU 虚拟机

TPU 虚拟机

$ gcloud compute tpus tpu-vm ssh tpu-name \
  --zone us-central1-b

TPU 节点

gcloud compute ssh tpu-name --zone=us-central1-b

设置 XRT TPU 设备配置

TPU 虚拟机

配置 Torch-XLA 环境。

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

TPU 节点

1. 查找 TPU 节点的 IP 地址。

   (vm)$ gcloud compute tpus describe tpu-name --zone=us-central1-b

2. 配置 Torch-XLA 环境。 请务必将 your-tpu-ip-address 替换为 TPU 的 IP 地址。

   (vm)$ conda activate torch-xla-1.11
   (vm)$ export TPU_IP_ADDRESS=your-tpu-ip-address
   (vm)$ export XRT_TPU_CONFIG="tpu_worker;0;$TPU_IP_ADDRESS:8470"

执行简单的计算：

1. 在当前目录中创建一个名为 tpu-test.py 的文件，并将以下脚本复制粘贴到其中。

   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)
   

2. 运行脚本：

     (vm)$ python3 tpu-test.py

   脚本的输出显示计算结果：

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

   注意：在 TPU 虚拟机上运行时，系统会显示以下错误消息。这是一个已知问题，可以忽略。此问题将在即将推出的 TPU 软件版本中得到解决。

   OpKernel ('op: "TPURoundRobin" device_type: "CPU"') for unknown op: TPURoundRobin
   OpKernel ('op: "TpuHandleToProtoKey" device_type: "CPU"') for unknown op: TpuHandleToProtoKey
   

清理

为避免因本页中使用的资源导致您的 Google Cloud 帐号产生费用，请按照以下步骤操作。

1. 断开与 Compute Engine 实例的连接（如果您尚未这样做）：

   (vm)$ exit
   

   您的提示符现在应为 username@projectname，表明您位于 Cloud Shell 中。

2. 删除您的 Cloud TPU。

   TPU 虚拟机

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

   TPU 节点

   $ gcloud compute tpus execution-groups delete tpu-name \
   --zone=us-central1-b

   此命令的输出应确认 TPU 已被删除。

后续步骤

详细了解 Cloud TPU 虚拟机：

• 在 TPU Pod 切片上运行 PyTorch 代码
• Cloud TPU 用户指南
• Cloud TPU 系统架构
• PyTorch/XLA 文档