管理 TPU 资源

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

Cloud TPU 有两种虚拟机架构,即 TPU 节点和 TPU 虚拟机。系统架构中介绍了这两种虚拟机架构。您可以将本文档中介绍的 gcloud 命令用于这两种 TPU 配置。您使用的 gcloud 命令取决于您使用的 TPU 配置。每个 gcloud 命令都显示在标签式部分中。选择您要使用的 TPU 配置的标签页,并且网页会显示相应的 gcloud 命令。除非您知道需要使用 TPU 节点,否则我们建议您使用 TPU 虚拟机。对于 Cloud TPU v4 及更高版本,仅支持 TPU 虚拟机架构。

前提条件

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

如果您使用的是 Google Cloud CLI,则可以使用 Google Cloud Shell、Compute Engine 虚拟机或在本地安装 Google Cloud CLI。借助 Google Cloud Shell,您无需安装任何软件即可与 Cloud TPU 交互。Google 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 alpha compute tpus tpu-vm create

最佳实践是使用已加入队列的资源进行预配。 当您请求已加入队列的资源时,请求会被添加到由 Cloud TPU 服务维护的队列中。当请求的资源可用时,系统就会将它分配给您的 Google Cloud 项目,供您立即专用使用。

如需使用排入队列的资源创建 TPU,请参阅排入队列的资源

如果您使用多切片,请参阅多切片简介了解详情。

使用 Multislice 时,请在请求已加入队列的资源时指定以下额外参数:

export NODE_COUNT=node_count
export NODE_PREFIX=your_tpu_prefix # Optional

使用 Create Node API 创建 Cloud TPU

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

TPU 虚拟机

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

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

$ gcloud alpha 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 alpha 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
您要创建的 Cloud TPU 的版本。如需详细了解 TPU 版本,请参阅 TPU 版本
topology
如需了解支持的拓扑,请参阅拓扑部分。
version
您要使用的 TPU 软件版本。如需了解详情,请参阅 TPU 软件版本

如需详细了解支持的 TPU 类型和拓扑,请参阅 TPU 版本

TPU 节点

$ gcloud compute tpus execution-groups create --name=tpu-name \
  --zone=us-central1-a \
  --tf-version=2.12.0 \
  --machine-type=n1-standard-1 \
  --accelerator-type=v3-8

命令标志说明

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

tf-version
gcloud 命令在您的虚拟机上安装的 Tensorflow 版本。

machine-type
要创建的 Compute Engine 虚拟机的 机器类型

accelerator-type
加速器类型指定要创建的 Cloud TPU 的版本和大小。如需详细了解每个 TPU 版本支持的加速器类型,请参阅 TPU 版本

Google Cloud 控制台中创建 Cloud TPU

  1. 前往 Google Cloud 控制台
  2. 在导航菜单中,选择 Compute Engine > TPU
  3. 点击“创建 TPU 节点”
  4. 在“名称”框中,输入 TPU 实例名称。
  5. 在“可用区”框中,选择要在其中创建 TPU 的可用区。
  6. TPU 设置下,选择 TPU 虚拟机架构TPU 节点架构。TPU 配置确定您是将 TPU 创建为 TPU 虚拟机还是 TPU 节点。如需了解详情,请参阅系统架构
  7. 对于“TPU 类型”,选择要创建的 TPU 类型
  8. 对于 TPU 软件版本,请选择软件版本。创建 Cloud TPU 虚拟机时,TPU 软件版本指定了要安装的 TPU 运行时的版本。创建 Cloud TPU 节点时,通过 TPU 软件版本您可选择安装在节点虚拟机上的机器学习框架。无需进行其他设置。如需了解详情,请参阅支持的模型
  9. 点击“创建”以创建资源。

使用 curl 创建 Cloud TPU 虚拟机

以下命令使用 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

您可以使用 SSH 连接到 TPU。

TPU 虚拟机

使用 TPU 虚拟机时,您必须使用 SSH 明确连接到您的 TPU 虚拟机。

  • 使用 gcloud compute tpus tpu-vm ssh 命令通过 SSH 连接到 TPU 虚拟机。

    当您请求包含 4 个以上芯片的切片时,Cloud TPU 会为每组 4 个芯片创建一个 TPU 虚拟机。

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

    gcloud compute tpus tpu-vm ssh ${TPU_NAME}
    

    如需通过 SSH 访问特定的 TPU 虚拟机或在每个 TPU 虚拟机上安装二进制文件,请使用 --worker 标志,该标志后跟索引(从 0 开始):

    gcloud compute tpus tpu-vm ssh ${TPU_NAME} --worker=1
    

    如果您有多个 TPU 虚拟机,请使用 --worker=all--command 标志同时在所有 TPU 虚拟机上运行命令。例如:

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

    对于 Multislice,您可以在单个虚拟机上运行命令(使用枚举的 tpu-name),也可以使用 --node=all--worker=all--command 标志在 Multislice 的所有切片的所有 TPU 虚拟机上运行该命令,并带有可选的 --batch-size 字段。

    gcloud alpha 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
    
    $ gcloud compute tpus tpu-vm ssh tpu-name --zone=zone
  • 执行以下操作,使用“在浏览器中使用 SSH”:

    1. 在 Google Cloud 控制台中,转到 TPU 页面:

      前往 TPU

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

TPU 节点

默认情况下,用于创建 TPU 节点的 gcloud 命令会自动尝试通过 SSH 连接到您的 TPU 节点。如果您使用的是 TPU 节点,并且未通过 gcloud 命令连接到 Compute Engine 实例,则可以通过运行以下命令进行连接:

$ gcloud compute ssh tpu-name \
  --zone=zone

创建 TPU 虚拟机后,您可以使用 SSH 连接到 TPU 虚拟机并运行以下命令来查看启动脚本的日志:

$ cat /var/log/syslog | grep startup-script

列出 Cloud TPU 资源

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

使用 gcloud 列出 Cloud TPU 资源

您使用的命令取决于您使用的是 TPU 虚拟机还是 TPU 节点。如需了解详情,请参阅系统架构

TPU 虚拟机

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

TPU 节点

$ gcloud compute tpus execution-groups list --zone=zone

此命令会列出指定可用区中的 Cloud TPU 资源。如果当前没有设置资源,则输出内容中虚拟机和 TPU 位置将只显示破折号。如果一个资源处于活动状态而另一资源非活动,则您将看到一条说明运行状况不佳的消息。此时,您需要启动或重启当前未运行的资源。

在 Google Cloud 控制台中列出 Cloud TPU 资源

  1. 前往 Google Cloud 控制台

  2. 在导航菜单中,选择 Compute Engine > TPU。 控制台会显示 TPU 页面。

检索有关 Cloud TPU 的信息

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

使用 gcloud 检索有关 Cloud TPU 的信息

您使用的命令取决于您使用的是 TPU 虚拟机还是 TPU 节点。如需了解详情,请参阅系统架构

TPU 虚拟机

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

TPU 节点

$ gcloud compute tpus execution-groups describe tpu-name \
  --zone=zone

Google Cloud 控制台中检索 Cloud TPU 的相关信息

  1. 前往 Google Cloud 控制台
  2. 在导航菜单中,选择 Compute Engine > TPU。 控制台会显示 TPU 页面。
  3. 点击您的 Cloud TPU 的名称。随即将显示 Cloud TPU 详情页面。

停用 Cloud TPU 资源

您可以停止单个 Cloud TPU 以停止产生费用,而不会丢失虚拟机的配置和软件。不支持停止通过已加入队列的资源 API 分配的 TPU Pod 或 TPU。如需让通过排队的资源 API 分配的 TPU 停止产生费用,您必须删除 TPU。

使用 gcloud 停止 Cloud TPU

用于停止 Cloud TPU 的命令取决于您使用的是 TPU 虚拟机还是 TPU 节点。如需了解详情,请参阅系统架构

TPU 虚拟机

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

TPU 节点

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

在 Google Cloud 控制台中停止 Cloud TPU

  1. 前往 Google Cloud 控制台

  2. 在导航菜单中,选择 Compute Engine > TPU。 控制台会显示 TPU 页面。

  3. 选中 Cloud TPU 旁边的复选框,然后点击“停止”

启动 Cloud TPU 资源

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

使用 gcloud 启动 Cloud TPU

您可以启动已停止的 Cloud TPU 以继续使用它。

用于启动已停止的 Cloud TPU 的命令取决于您使用的是 TPU 虚拟机还是 TPU 节点。如需了解详情,请参阅系统架构

TPU 虚拟机

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

TPU 节点

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

在 Google Cloud 控制台中启动 Cloud TPU

  1. 前往 Google Cloud 控制台

  2. 在导航菜单中,选择 Compute Engine > TPU。 控制台会显示 TPU 页面。

  3. 选中 Cloud TPU 旁边的复选框,然后点击“启动”

删除您的 TPU 虚拟机

您使用的命令取决于您使用的是 TPU 虚拟机还是 TPU 节点。如需了解详情,请参阅系统架构

TPU 虚拟机

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

 gcloud compute tpus tpu-vm delete ${TPU_NAME} --project=${PROJECT_ID}
 --zone=${ZONE} --quiet

命令标志说明

zone
您计划删除 Cloud TPU 的区域

TPU 节点

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

命令标志说明

zone
您计划删除 Cloud TPU 的区域

在 Google Cloud 控制台中删除 Cloud TPU

  1. 前往 Google Cloud 控制台

  2. 在导航菜单中,选择 Compute Engine > TPU。 控制台会显示 TPU 页面。

  3. 选中 Cloud TPU 旁边的复选框,然后点击“删除”

高级配置

自定义网络资源

创建 TPU 时,您可以选择指定网络和/或子网。您可以通过提交 gcloud 命令或 curl 调用来实现此目的。

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

--network [NETWORK] --subnetwork [SUBNETWORK]

如需在 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 的地区位于同一区域

有效格式

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。如果您想选择不自动迁移,请指定以下命令:

--internal-ips

或者使用 curl 调用:

network_config: {enable_external_ips: true}

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

自定义服务账号

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

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

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

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

使用 gcloud CLI 创建 TPU 虚拟机

$ gcloud alpha 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 创建 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 }}" \
  service_account: {email: 'your-service-account'} \
  https://tpu.googleapis.com/v2/projects/project-id/locations/us-central2-b/nodes?node_id=node_name

如需使用自定义服务帐号,您需要为 Google Cloud Storage 存储分区的服务帐号授权。如需了解详情,请参阅连接到 Cloud Storage 存储分区

自定义虚拟机 SSH 方法

  1. 为 SSH 设置防火墙。

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

    $ gcloud CLI 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 节点的位置。目前仅支持 us-central2-b
    • project-id:您在上面创建的项目。

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