管理 TPU 资源

本页面介绍了如何创建、列出、停止、启动、删除和连接 使用 Create Node API 的 Cloud TPU。当您创建节点时,系统会调用 使用 Google Cloud CLI 运行 gcloud compute tpus tpu-vm create 命令 以及使用 Google Cloud 控制台创建 TPU 时。使用“创建” Node API,系统会立即处理您的请求。如果容量不足 处理请求,该请求将失败。

最佳实践是使用 加入队列的资源,而不是使用 Create Node API。 当您对已加入队列的资源发出请求时,系统会将该请求添加到由 Cloud TPU 服务。当请求的资源可用时 分配给您的 Google Cloud 项目,供您立即独占使用。对于 如需了解详情,请参阅管理已加入队列的资源

使用多切片时,您必须使用排队的资源。有关 请参阅 多切片简介

如果您想使用 Google Kubernetes Engine (GKE) 管理 TPU 资源, 首先必须创建一个 GKE 集群然后将节点池添加到 添加到集群的 TPU 切片。如需了解详情,请参阅 GKE 中的 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 参考文档

使用 Create Node API 创建 Cloud TPU

您可以使用 Google Cloud 控制台 gcloud 创建 Cloud TPU, 或 Cloud TPU API。

创建 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 版本
version
您要使用的 TPU 软件版本。如需了解详情,请参阅 TPU 软件版本

控制台

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

    前往 TPU

  2. 点击创建 TPU

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

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

  5. TPU 类型框中,选择一个加速器类型。 加速器类型指定 Cloud 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.17.0-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 虚拟机的名称。

运行启动脚本

您可以通过指定 --metadata startup-script 标志。以下 命令来创建 TPU 虚拟机。

$ gcloud compute tpus tpu-vm create tpu-name \
    --zone=us-central2-b \
    --accelerator-type=tpu-type \
    --version=tpu-vm-tf-2.17.0-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

$ gcloud compute tpus tpu-vm ssh tpu-name

如需连接到特定的 TPU 虚拟机 使用 SSH 时,使用 --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 名称,以及 每个切片前缀和附加到它的编号。要运行 命令,请使用 --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. 点击 删除