管理 TPU 资源

本页介绍如何使用以下各项管理 Cloud TPU 资源:

前提条件

在运行这些过程之前,您必须安装 Google Cloud CLI、创建 Google Cloud 项目并启用 Cloud TPU API。有关说明,请参阅设置 Cloud TPU 环境

如果您使用的是 Google Cloud CLI,则可以使用 Cloud Shell、Compute Engine 虚拟机或本地机器运行命令。借助 Cloud Shell,您无需安装任何软件即可与 Cloud TPU 交互。Cloud Shell 会在一段时间不活动后断开连接。如果您运行的是长时间运行的命令,建议在本地机器上安装 Google Cloud CLI。如需详细了解 Google Cloud CLI,请参阅 gcloud 参考文档

预配 Cloud TPU

您可以使用 gcloud、Google Cloud 控制台或 Cloud TPU API 来预配 Cloud TPU。

您可以通过以下两种方法使用 gcloud 预配 TPU:

  1. 使用已加入队列的资源:gcloud alpha compute tpus queued-resources create
  2. 使用 Create Node API:gcloud compute tpus tpu-vm create

最佳实践是使用已加入队列的资源来预配 TPU。当您对已加入队列的资源发出请求时,请求会被添加到由 Cloud TPU 服务维护的队列中。当请求的资源可用时,系统会将其分配给您的 Google Cloud 项目,供您立即独占使用。如需了解详情,请参阅代管式已加入队列的资源

使用多切片时,您必须使用已加入队列的资源并指定以下其他参数:

export NODE_COUNT=node_count
export NODE_PREFIX=your_tpu_prefix # Optional

其中:

  • ${NODE_COUNT} 是要创建的切片数量
  • ${NODE_PREFIX} 是您指定用于为每个切片生成名称的前缀。每个切片的前缀后面会附加一个数字。例如,如果您将 ${NODE_PREFIX} 设置为 mySlice,则切片将命名为:mySlice-0、mySlice-1,每个切片按数字编号继续。

如需详细了解多切片,请参阅多切片简介

使用 Create Node API 创建 Cloud TPU

创建 Cloud TPU 时,您必须指定 TPU 虚拟机映像(也称为 TPU 软件版本)。如需确定应使用哪个虚拟机映像,请参阅 TPU 虚拟机映像

您还需要根据 TensorCore 或 TPU 芯片指定 TPU 配置。如需了解详情,请参阅系统架构中有关您使用的 TPU 版本的部分。

gcloud

如需使用 Create Node API 创建 TPU,请使用 gcloud compute tpus tpu-vm create 命令。

以下命令使用基于 TensorCore 的配置:

$ gcloud compute tpus tpu-vm create tpu-name \
  --zone=us-central2-b \
  --accelerator-type=v4-8 \
  --version=tpu-software-version

命令标志说明

zone
您计划在其中创建 Cloud TPU 的地区
accelerator-type
加速器类型指定要创建的 Cloud TPU 的版本和大小。如需详细了解每个 TPU 版本支持的加速器类型,请参阅 TPU 版本
version
TPU 软件版本。
shielded-secure-boot(可选)
指定在启用安全启动的情况下创建 TPU 实例。这会隐式创建安全强化型虚拟机实例。请参阅什么是安全强化型虚拟机?

以下命令将创建具有特定拓扑的 TPU:

$ gcloud compute tpus tpu-vm create tpu-name \
  --zone=us-central2-b \
  --type=v4 \
  --topology=2x2x1 \
  --version=tpu-software-version

必需标志

tpu-name
您正在创建的 TPU 虚拟机的名称。
zone
您要在其中创建 Cloud TPU 的地区
type
您要使用的 TPU 版本。如需了解详情,请参阅 TPU 版本
topology
TPU 芯片的物理排列,指定每个维度中的芯片数量。如需详细了解每个 TPU 版本支持的拓扑,请参阅 TPU 版本
version
您要使用的 TPU 软件版本。如需了解详情,请参阅 TPU 软件版本

控制台

  1. 在 Google Cloud 控制台中,前往 TPU 页面:

    前往 TPU

  2. 点击创建 TPU

  3. 名称字段中,输入 TPU 的名称。

  4. 在“可用区”框中,选择要在其中创建 TPU 的可用区。

  5. TPU 类型框中,选择加速器类型。 加速器类型指定要创建的 Cloud TPU 的版本和大小。如需详细了解每个 TPU 版本支持的加速器类型,请参阅 TPU 版本

  6. TPU 软件版本框中,选择一个软件版本。创建 Cloud TPU 虚拟机时,TPU 软件版本指定了要安装的 TPU 运行时的版本。如需了解详情,请参阅 TPU 虚拟机映像

  7. 点击创建以创建资源。

curl

以下命令使用 curl 创建 TPU。

$ curl -X POST -H "Authorization: Bearer $(gcloud auth print-access-token)" -H "Content-Type: application/json" -d "{accelerator_type: 'v4-8', \
runtime_version:'tpu-vm-tf-2.16.1-pjrt', \
network_config: {enable_external_ips: true}, \
shielded_instance_config: { enable_secure_boot: true }}" \
https://tpu.googleapis.com/v2/projects/project-id/locations/us-central2-b/nodes?node_id=node_name

必填字段

runtime_version
您要使用的 Cloud TPU 运行时版本。
project
已注册的 Google Cloud 项目的名称。
zone
您要在其中创建 Cloud TPU 的地区
node_name
您要创建的 TPU 虚拟机的名称。

运行启动脚本

通过在创建 TPU 虚拟机时指定 --metadata startup-script 标志,您可以在每个 TPU 虚拟机上运行启动脚本。以下命令使用启动脚本创建一个 TPU 虚拟机。

$ gcloud compute tpus tpu-vm create tpu-name \
    --zone=us-central2-b \
    --accelerator-type=tpu-type \
    --version=tpu-vm-tf-2.16.1-pjrt \
    --metadata startup-script='#! /bin/bash
      pip3 install numpy
      EOF'

连接到 Cloud TPU

gcloud

使用 SSH 连接到您的 Cloud TPU:

$ gcloud compute tpus tpu-vm ssh tpu-name --zone=zone

当您请求大于单个主机的切片时,Cloud TPU 会为每个主机创建一个 TPU 虚拟机。每个主机的 TPU 芯片的数量取决于 TPU 版本

如需安装二进制文件或运行代码,请使用 tpu-vm ssh command 连接到每个 TPU 虚拟机。

$ gcloud compute tpus tpu-vm ssh tpu-name

如需使用 SSH 连接到特定的 TPU 虚拟机,请使用 --worker 标志,该标志遵循从 0 开始的索引:

$ gcloud compute tpus tpu-vm ssh tpu-name --worker=1

如需通过单个命令在所有 TPU 虚拟机上运行某个命令,请使用 --worker=all--command 标志:

$ gcloud compute tpus tpu-vm ssh tpu-name \
  --project=your_project_ID \
  --zone=zone \
  --worker=all \
  --command='pip install "jax[tpu]==0.4.20" -f https://storage.googleapis.com/jax-releases/libtpu_releases.html'

对于多切片,您可以使用枚举的 TPU 名称(每个切片前缀和附加编号)在单个虚拟机上运行命令。如需在所有切片中的所有 TPU 虚拟机上运行命令,请使用 --node=all--worker=all--command 标志,并附带 --batch-size 标志(可选)。

$ gcloud compute tpus queued-resources ssh ${QUEUED_RESOURCE_ID} \
  --project=project_ID \
  --zone=zone \
  --node=all \
  --worker=all \
  --command='pip install "jax[tpu]==0.4.20" -f https://storage.googleapis.com/jax-releases/libtpu_releases.html' \
  --batch-size=4

控制台

如需在 Google Cloud 控制台中连接到您的 TPU,请使用浏览器中的 SSH:

  1. 在 Google Cloud 控制台中,前往 TPU 页面:

    前往 TPU

  2. 在 TPU 虚拟机列表中,点击要连接的 TPU 虚拟机所在行中的 SSH

列出您的 Cloud TPU 资源

您可以列出指定可用区中的所有 Cloud TPU。

gcloud

$ gcloud compute tpus tpu-vm list --zone=zone

控制台

在 Google Cloud 控制台中,前往 TPU 页面:

前往 TPU

检索有关 Cloud TPU 的信息

您可以检索关于指定 Cloud TPU 的信息。

gcloud

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

控制台

  1. 在 Google Cloud 控制台中,前往 TPU 页面:

    前往 TPU

  2. 点击您的 Cloud TPU 的名称。控制台会显示 Cloud TPU 详情页面。

停止 Cloud TPU 资源

您可以停止单个 Cloud TPU 以停止产生费用,而不会丢失虚拟机的配置和软件。

gcloud

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

控制台

  1. 在 Google Cloud 控制台中,前往 TPU 页面:

    前往 TPU

  2. 选中您的 Cloud TPU 旁边的复选框。

  3. 点击 停止

启动 Cloud TPU 资源

您可以在 Cloud TPU 停止后启动它。

gcloud

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

控制台

  1. 在 Google Cloud 控制台中,前往 TPU 页面:

    前往 TPU

  2. 选中您的 Cloud TPU 旁边的复选框。

  3. 点击 开始

删除 Cloud TPU

在会话结束时删除 TPU 虚拟机切片。

gcloud

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

命令标志说明

zone
您打算删除 Cloud TPU 的区域

控制台

  1. 在 Google Cloud 控制台中,前往 TPU 页面:

    前往 TPU

  2. 选中您的 Cloud TPU 旁边的复选框。

  3. 点击 删除

高级配置

指定自定义网络资源

创建 TPU 时,您可以选择指定网络或子网。

gcloud

如需使用 gcloud CLI 指定网络或子网,请使用以下命令标志:

--network NETWORK --subnetwork SUBNETWORK

curl

如需在 curl 调用中指定网络或子网,请将以下参数添加到请求正文中:

network_config: {network: 'NETWORK', subnet: 'SUBNETWORK', enable_external_ips: true}

网络

您可以选择指定要用于 TPU 的网络。如果未指定,则使用 default 网络。

有效的网络格式

https://www.googleapis.com/compute/{version}/projects/{proj-id}/global/networks/{network}
compute/{version}/projects/{proj-id}/global/networks/{network}
compute/{version}/projects/{proj-##}/global/networks/{network}
projects/{proj-id}/global/networks/{network}
projects/{proj-##}/global/networks/{network}
global/networks/{network}
{network}

子网

您可以指定要用于 TPU 的特定子网。指定的子网需要与运行 TPU 的地区位于同一区域

有效格式

https://www.googleapis.com/compute/{version}/projects/{proj-id}/regions/{region}/subnetworks/{subnetwork}
compute/{version}/projects/{proj-id}/regions/{region}/subnetworks/{subnetwork}
compute/{version}/projects/{proj-##}/regions/{region}/subnetworks/{subnetwork}
projects/{proj-id}/regions/{region}/subnetworks/{subnetwork}
projects/{proj-##}/regions/{region}/subnetworks/{subnetwork}
regions/{region}/subnetworks/{subnetwork}
{subnetwork}

启用专用 Google 访问通道

如需使用 SSH 连接到 TPU 虚拟机,您需要为 TPU 虚拟机添加访问配置,或者为 TPU 虚拟机所连接的子网启用专用 Google 访问通道

如需添加访问权限配置,您必须设置 enable_external_ips。当您创建 TPU 时,系统会默认设置 enable_external_ips

如果您想选择停用,请启用内部 IP:

gcloud

创建 TPU 时使用 --internal-ips 标志:

--internal-ips

curl

将以下参数添加到请求正文:

network_config: {enable_external_ips: false}

配置专用 Google 访问通道后,请使用 SSH 连接到虚拟机

附加自定义服务帐号

每个 TPU 虚拟机都有一个关联的服务帐号,用于代表您发出 API 请求。TPU 虚拟机使用此服务帐号调用 Cloud TPU API 以及访问 Cloud Storage 和其他服务。默认情况下,TPU 虚拟机使用默认 Compute Engine 服务帐号

您必须在创建 TPU 虚拟机的同一 Google Cloud 项目中定义该服务帐号。用于 TPU 虚拟机的自定义服务帐号必须具有 TPU Viewer 角色才能调用 Cloud TPU API。如果 TPU 虚拟机中运行的代码调用其他 Google Cloud 服务,则该代码必须具有访问这些服务所需的角色。

创建 TPU 时,您可以选择使用 --service-account 标志指定自定义服务帐号。如需详细了解服务帐号,请参阅服务帐号

使用以下命令指定自定义服务帐号。

gcloud

$ gcloud compute tpus tpu-vm create tpu-name \
  --zone=us-central2-b \
  --accelerator-type=tpu-type \
  --version=tpu-vm-tf-2.16.1-pjrt \
  --service-account=your-service-account

curl

$ curl -X POST -H "Authorization: Bearer $(gcloud auth print-access-token)" -H "Content-Type: application/json" -d "{accelerator_type: 'v4-8', \
runtime_version:'tpu-vm-tf-2.16.1-pjrt', \
network_config: {enable_external_ips: true}, \
shielded_instance_config: { enable_secure_boot: true }}" \
service_account: {email: 'your-service-account'} \
https://tpu.googleapis.com/v2/projects/project-id/locations/us-central2-b/nodes?node_id=node_name

启用自定义 SSH 方法

  1. 为 SSH 设置防火墙。

    默认网络已预先配置为允许所有虚拟机进行 SSH 访问。如果您不使用默认网络,或者已更改默认网络设置,则可能需要通过添加防火墙规则来明确启用 SSH 访问:

    $ gcloud compute tpus tpu-vm compute firewall-rules create \
      --network=network allow-ssh \
      --allow=tcp:22
    
  2. 使用 SSH 连接到 TPU 虚拟机。

    $ gcloud compute tpus tpu-vm ssh tpu-name \
      --zone=us-central2-b \
      --project=project-id
    

    必填字段

    • tpu-name:TPU 虚拟机的名称。
    • zone:您在其中创建了 TPU 虚拟机的区域。
    • project-id:您的 Google Cloud 项目的名称。

    如需查看可选字段的列表,请参阅 gcloud API 文档