本教程介绍如何使用 TPU Node 配置将模型从单个 Cloud TPU(v2-8 或 v3-8)扩容到 Cloud TPU Pod。TPU Pod 中的 Cloud TPU 加速器通过高带宽互连来连接,因而能够高效地扩容训练作业。
如需详细了解 Cloud TPU Pod 产品,请参阅 Cloud TPU 产品页面或此 Cloud TPU 演示。
下图简要介绍分布式集群设置。虚拟机的实例组已连接到 TPU Pod。每组(含 8 个 TPU 核心)都需要一个虚拟机。虚拟机为 TPU 核心馈送数据,所有训练都在 TPU Pod 上进行。
目标
- 设置 Compute Engine 实例组和 Cloud TPU Pod,以便使用 PyTorch/XLA 进行训练
- 在 Cloud TPU Pod 上运行 PyTorch/XLA 训练
准备工作
在 Cloud TPU Pod 上开始分布式训练之前,请验证您的模型可在单个 v2-8 或 v3-8 Cloud TPU 设备上正常训练。如果您的模型在单个设备上出现明显的性能问题,请参阅最佳做法和问题排查指南。
单个 TPU 设备成功训练后,请执行以下步骤,在 Cloud TPU Pod 上设置和训练:
配置 gcloud
命令
使用 gcloud
配置 Google Cloud 项目:
为项目 ID 创建一个变量。
export PROJECT_ID=project-id
将项目 ID 设置为 gcloud
中的默认项目
gcloud config set project ${PROJECT_ID}
当您第一次在新的 Cloud Shell 虚拟机中运行此命令时,系统会显示 Authorize Cloud Shell
页面。点击页面底部的 Authorize
,以允许 gcloud
使用您的凭据进行 API 调用。
使用 gcloud
配置默认区域:
gcloud config set compute/zone europe-west4-a
[可选] 捕获虚拟机磁盘映像
您可以使用您用于训练单个 TPU 的虚拟机中的磁盘映像(已安装数据集、软件包等)。在创建映像之前,请使用 gcloud
命令停止虚拟机:
gcloud compute instances stop vm-name
接下来,使用 gcloud
命令创建虚拟机映像:
gcloud compute images create image-name \ --source-disk instance-name \ --source-disk-zone europe-west4-a \ --family=torch-xla \ --storage-location europe-west4
通过虚拟机映像创建实例模板
创建默认实例模板。创建实例模板时,您可以使用您在上述步骤中创建的虚拟机映像,或者您可以使用 Google 提供的公开 PyTorch/XLA 映像。如需创建实例模板,请使用 gcloud
命令:
gcloud compute instance-templates create instance-template-name \ --machine-type n1-standard-16 \ --image-project=${PROJECT_ID} \ --image=image-name \ --scopes=https://www.googleapis.com/auth/cloud-platform
通过实例模板创建实例组
gcloud compute instance-groups managed create instance-group-name \ --size 4 \ --template template-name \ --zone europe-west4-a
通过 SSH 连接到您的 Compute Engine 虚拟机
创建实例组后,请通过 SSH 连接到实例组中的某个实例(虚拟机)。使用 gcloud
命令列出实例组中的所有实例:
gcloud compute instance-groups list-instances instance-group-name
通过 SSH 连接到使用 list-instances
命令列出的某个实例。
gcloud compute ssh instance-name --zone=europe-west4-a
验证实例组中的虚拟机是否可以相互通信
使用 nmap
命令验证实例组中的虚拟机是否可以相互通信。从您连接的虚拟机运行 nmap
命令,将 instance-name 替换为实例组中另一个虚拟机的实例名称。
(vm)$ nmap -Pn -p 8477 instance-name
Starting Nmap 7.40 ( https://nmap.org ) at 2019-10-02 21:35 UTC Nmap scan report for pytorch-20190923-n4tx.c.jysohntpu.internal (10.164.0.3) Host is up (0.00034s latency). PORT STATE SERVICE 8477/tcp closed unknown
只要 STATE 字段未显示 filtered,就表示防火墙规则设置正确。
创建 Cloud TPU Pod
gcloud compute tpus create tpu-name \ --zone=europe-west4-a \ --network=default \ --accelerator-type=v2-32 \ --version=pytorch-1.13
在 Pod 上运行分布式训练
从虚拟机会话窗口中,导出 Cloud TPU 名称并激活 conda 环境。
(vm)$ export TPU_NAME=tpu-name
(vm)$ conda activate torch-xla-1.13
运行训练脚本:
(torch-xla-1.13)$ python -m torch_xla.distributed.xla_dist \ --tpu=$TPU_NAME \ --conda-env=torch-xla-1.13 \ --env XLA_USE_BF16=1 \ --env ANY_OTHER=ENV_VAR \ -- python /usr/share/torch-xla-1.13/pytorch/xla/test/test_train_mp_imagenet.py \ --fake_data
运行完上述命令后,您应该会看到如下所示的输出(注意,这使用的是 --fake_data
)。在 v3-32 TPU Pod 上,训练需要大约 1/2 小时。
2020-08-06 02:38:29 [] Command to distribute: "python" "/usr/share/torch-xla-nightly/pytorch/xla/test/test_train_mp_imagenet.py" "--fake_data" 2020-08-06 02:38:29 [] Cluster configuration: {client_workers: [{10.164.0.43, n1-standard-96, europe-west4-a, my-instance-group-hm88}, {10.164.0.109, n1-standard-96, europe-west4-a, my-instance-group-n3q2}, {10.164.0.46, n1-standard-96, europe-west4-a, my-instance-group-s0xl}, {10.164.0.49, n1-standard-96, europe-west4-a, my-instance-group-zp14}], service_workers: [{10.131.144.61, 8470, v3-32, europe-west4-a, pytorch-nightly, my-tpu-slice}, {10.131.144.59, 8470, v3-32, europe-west4-a, pytorch-nightly, my-tpu-slice}, {10.131.144.58, 8470, v3-32, europe-west4-a, pytorch-nightly, my-tpu-slice}, {10.131.144.60, 8470, v3-32, europe-west4-a, pytorch-nightly, my-tpu-slice}]} 2020-08-06 02:38:31 10.164.0.43 [0] % Total % Received % Xferd Average Speed Time Time Time Current 2020-08-06 02:38:31 10.164.0.43 [0] Dload Upload Total Spent Left Speed 100 19 100 19 0 0 2757 0 --:--:-- --:--:-- --:--:-- 3166 2020-08-06 02:38:34 10.164.0.43 [0] % Total % Received % Xferd Average Speed Time Time Time Current 2020-08-06 02:38:34 10.164.0.43 [0] Dload Upload Total Spent Left Speed 100 19 100 19 0 0 2623 0 --:--:-- --:--:-- --:--:-- 2714 2020-08-06 02:38:37 10.164.0.46 [2] % Total % Received % Xferd Average Speed Time Time Time Current 2020-08-06 02:38:37 10.164.0.46 [2] Dload Upload Total Spent Left Speed 100 19 100 19 0 0 2583 0 --:--:-- --:--:-- --:--:-- 2714 2020-08-06 02:38:37 10.164.0.49 [3] % Total % Received % Xferd Average Speed Time Time Time Current 2020-08-06 02:38:37 10.164.0.49 [3] Dload Upload Total Spent Left Speed 100 19 100 19 0 0 2530 0 --:--:-- --:--:-- --:--:-- 2714 2020-08-06 02:38:37 10.164.0.109 [1] % Total % Received % Xferd Average Speed Time Time Time Current 2020-08-06 02:38:37 10.164.0.109 [1] Dload Upload Total Spent Left Speed 100 19 100 19 0 0 2317 0 --:--:-- --:--:-- --:--:-- 2375 2020-08-06 02:38:40 10.164.0.46 [2] % Total % Received % Xferd Average Speed Time Time Time Current 2020-08-06 02:38:40 10.164.0.49 [3] % Total % Received % Xferd Average Speed Time Time Time Current 2020-08-06 02:38:40 10.164.0.46 [2] Dload Upload Total Spent Left Speed 2020-08-06 02:38:40 10.164.0.49 [3] Dload Upload Total Spent Left Speed 100 19 100 19 0 0 2748 0 --:--:-- --:--:-- --:--:-- 3166 100 19 100 19 0 0 2584 0 --:--:-- --:--:-- --:--:-- 2714 2020-08-06 02:38:40 10.164.0.109 [1] % Total % Received % Xferd Average Speed Time Time Time Current 2020-08-06 02:38:40 10.164.0.109 [1] Dload Upload Total Spent Left Speed 100 19 100 19 0 0 2495 0 --:--:-- --:--:-- --:--:-- 2714 2020-08-06 02:38:43 10.164.0.49 [3] % Total % Received % Xferd Average Speed Time Time Time Current 2020-08-06 02:38:43 10.164.0.49 [3] Dload Upload Total Spent Left Speed 100 19 100 19 0 0 2654 0 --:--:-- --:--:-- --:--:-- 2714 2020-08-06 02:38:43 10.164.0.43 [0] % Total % Received % Xferd Average Speed Time Time Time Current 2020-08-06 02:38:43 10.164.0.43 [0] Dload Upload Total Spent Left Speed 100 19 100 19 0 0 2784 0 --:--:-- --:--:-- --:--:-- 3166 2020-08-06 02:38:43 10.164.0.46 [2] % Total % Received % Xferd Average Speed Time Time Time Current 2020-08-06 02:38:43 10.164.0.46 [2] Dload Upload Total Spent Left Speed 100 19 100 19 0 0 2691 0 --:--:-- --:--:-- --:--:-- 3166 2020-08-06 02:38:43 10.164.0.109 [1] % Total % Received % Xferd Average Speed Time Time Time Current 2020-08-06 02:38:43 10.164.0.109 [1] Dload Upload Total Spent Left Speed 100 19 100 19 0 0 2589 0 --:--:-- --:--:-- --:--:-- 2714 2020-08-06 02:38:57 10.164.0.109 [1] ==> Preparing data.. 2020-08-06 02:38:57 10.164.0.109 [1] | Training Device=xla:0/14 Epoch=1 Step=0 Loss=6.87500 Rate=258.47 GlobalRate=258.47 Time=02:38:57 2020-08-06 02:38:57 10.164.0.109 [1] ==> Preparing data.. 2020-08-06 02:38:57 10.164.0.109 [1] | Training Device=xla:0/15 Epoch=1 Step=0 Loss=6.87500 Rate=149.45 GlobalRate=149.45 Time=02:38:57 2020-08-06 02:38:57 10.164.0.43 [0] ==> Preparing data.. 2020-08-06 02:38:57 10.164.0.43 [0] Epoch 1 train begin 02:38:52 2020-08-06 02:38:57 10.164.0.43 [0] | Training Device=xla:1/0 Epoch=1 Step=0 Loss=6.87500 Rate=25.72 GlobalRate=25.72 Time=02:38:57 2020-08-06 02:38:57 10.164.0.43 [0] ==> Preparing data.. 2020-08-06 02:38:57 10.164.0.43 [0] | Training Device=xla:0/6 Epoch=1 Step=0 Loss=6.87500 Rate=89.01 GlobalRate=89.01 Time=02:38:57 2020-08-06 02:38:57 10.164.0.43 [0] ==> Preparing data.. 2020-08-06 02:38:57 10.164.0.43 [0] | Training Device=xla:0/1 Epoch=1 Step=0 Loss=6.87500 Rate=64.15 GlobalRate=64.15 Time=02:38:57 2020-08-06 02:38:57 10.164.0.43 [0] ==> Preparing data.. 2020-08-06 02:38:57 10.164.0.43 [0] | Training Device=xla:0/2 Epoch=1 Step=0 Loss=6.87500 Rate=93.19 GlobalRate=93.19 Time=02:38:57 2020-08-06 02:38:57 10.164.0.43 [0] ==> Preparing data.. 2020-08-06 02:38:57 10.164.0.43 [0] | Training Device=xla:0/7 Epoch=1 Step=0 Loss=6.87500 Rate=58.78 GlobalRate=58.78 Time=02:38:57 2020-08-06 02:38:57 10.164.0.109 [1] ==> Preparing data.. 2020-08-06 02:38:57 10.164.0.109 [1] Epoch 1 train begin 02:38:56 2020-08-06 02:38:57 10.164.0.109 [1] | Training Device=xla:1/8 Epoch=1 Step=0 Loss=6.87500 Rate=100.43 GlobalRate=100.43 Time=02:38:57 2020-08-06 02:38:57 10.164.0.109 [1] ==> Preparing data.. 2020-08-06 02:38:57 10.164.0.109 [1] | Training Device=xla:0/13 Epoch=1 Step=0 Loss=6.87500 Rate=66.83 GlobalRate=66.83 Time=02:38:57 2020-08-06 02:38:57 10.164.0.109 [1] ==> Preparing data.. 2020-08-06 02:38:57 10.164.0.109 [1] | Training Device=xla:0/11 Epoch=1 Step=0 Loss=6.87500 Rate=64.28 GlobalRate=64.28 Time=02:38:57 2020-08-06 02:38:57 10.164.0.109 [1] ==> Preparing data.. 2020-08-06 02:38:57 10.164.0.109 [1] | Training Device=xla:0/10 Epoch=1 Step=0 Loss=6.87500 Rate=73.17 GlobalRate=73.17 Time=02:38:57 2020-08-06 02:38:57 10.164.0.109 [1] ==> Preparing data.. 2020-08-06 02:38:57 10.164.0.109 [1] | Training Device=xla:0/9 Epoch=1 Step=0 Loss=6.87500 Rate=27.29 GlobalRate=27.29 Time=02:38:57 2020-08-06 02:38:57 10.164.0.109 [1] ==> Preparing data.. 2020-08-06 02:38:57 10.164.0.109 [1] | Training Device=xla:0/12 Epoch=1 Step=0 Loss=6.87500 Rate=110.29 GlobalRate=110.29 Time=02:38:57 2020-08-06 02:38:57 10.164.0.46 [2] ==> Preparing data.. 2020-08-06 02:38:57 10.164.0.46 [2] | Training Device=xla:0/20 Epoch=1 Step=0 Loss=6.87500 Rate=100.85 GlobalRate=100.85 Time=02:38:57 2020-08-06 02:38:57 10.164.0.46 [2] ==> Preparing data.. 2020-08-06 02:38:57 10.164.0.46 [2] | Training Device=xla:0/22 Epoch=1 Step=0 Loss=6.87500 Rate=93.52 GlobalRate=93.52 Time=02:38:57 2020-08-06 02:38:57 10.164.0.46 [2] ==> Preparing data.. 2020-08-06 02:38:57 10.164.0.46 [2] | Training Device=xla:0/23 Epoch=1 Step=0 Loss=6.87500 Rate=165.86 GlobalRate=165.86 Time=02:38:57
清理
为避免因本教程中使用的资源导致您的 Google Cloud 帐号产生费用,请删除包含这些资源的项目,或者保留项目但删除各个资源。
与 Compute Engine 虚拟机断开连接:
(vm)$ exit
删除实例组:
gcloud compute instance-groups managed delete instance-group-name
删除 TPU Pod:
gcloud compute tpus delete ${TPU_NAME} --zone=europe-west4-a
删除实例组模板:
gcloud compute instance-templates delete instance-template-name
[可选] 删除您的虚拟机磁盘映像:
gcloud compute images delete image-name
后续步骤
试用 PyTorch Colab:
- 在 Cloud TPU 上开始使用 PyTorch
- 在 TPU 上训练 MNIST
- 使用 Cifar10 数据集在 TPU 上训练 ResNet18
- 使用预训练的 ResNet50 模型进行推理
- 快速神经风格转移
- 在 Fashion MNIST 上使用多核心训练 AlexNet
- 在 Fashion MNIST 上使用单核心训练 AlexNet