创建 TPU 实例

本文档将介绍如何创建 TPU 虚拟机 (VM) 实例、连接到该实例，以及在该实例上运行计算。您在本快速入门中创建的 TPU 实例使用按需付费（默认消费选项）。

准备工作

  1. 登录您的 Google Cloud 账号。如果您是 Google Cloud新手，请 创建一个账号来评估我们的产品在实际场景中的表现。新客户还可获享 $300 赠金，用于运行、测试和部署工作负载。

  2. 安装 Google Cloud CLI。

  3. 如果您使用的是外部身份提供方 (IdP)，则必须先使用联合身份登录 gcloud CLI。

  4. 如需初始化 gcloud CLI，请运行以下命令： 

    gcloud init

  5. 创建或选择 Google Cloud 项目。

    选择或创建项目所需的角色

    • 选择项目：选择项目不需要特定的 IAM 角色，您可以选择已获授角色的任何项目。
    • 创建项目：如需创建项目，您需要拥有 Project Creator 角色 (roles/resourcemanager.projectCreator)，该角色包含 resourcemanager.projects.create 权限。了解如何授予角色。

    • 创建 Google Cloud 项目：

      gcloud projects create PROJECT_ID

      将 PROJECT_ID 替换为您要创建的 Google Cloud 项目的名称。

    • 选择您创建的 Google Cloud 项目：

      gcloud config set project PROJECT_ID

      将 PROJECT_ID 替换为您的 Google Cloud 项目名称。

  6. 如果您要使用现有项目来完成本指南，请验证您是否拥有完成本指南所需的权限。如果您创建了新项目，则您已拥有所需的权限。

  7. 验证是否已为您的 Google Cloud 项目启用结算功能。

  8. 启用 Compute Engine API：

    启用 API 所需的角色

    如需启用 API，您需要拥有 Service Usage Admin IAM 角色 (roles/serviceusage.serviceUsageAdmin)，该角色包含 serviceusage.services.enable 权限。了解如何授予角色。 

    gcloud services enable compute.googleapis.com

  9. 安装 Google Cloud CLI。

  10. 如果您使用的是外部身份提供方 (IdP)，则必须先使用联合身份登录 gcloud CLI。

  11. 如需初始化 gcloud CLI，请运行以下命令： 

    gcloud init

  12. 创建或选择 Google Cloud 项目。

    选择或创建项目所需的角色

    • 选择项目：选择项目不需要特定的 IAM 角色，您可以选择已获授角色的任何项目。
    • 创建项目：如需创建项目，您需要拥有 Project Creator 角色 (roles/resourcemanager.projectCreator)，该角色包含 resourcemanager.projects.create 权限。了解如何授予角色。

    • 创建 Google Cloud 项目：

      gcloud projects create PROJECT_ID

      将 PROJECT_ID 替换为您要创建的 Google Cloud 项目的名称。

    • 选择您创建的 Google Cloud 项目：

      gcloud config set project PROJECT_ID

      将 PROJECT_ID 替换为您的 Google Cloud 项目名称。

  13. 如果您要使用现有项目来完成本指南，请验证您是否拥有完成本指南所需的权限。如果您创建了新项目，则您已拥有所需的权限。

  14. 验证是否已为您的 Google Cloud 项目启用结算功能。

  15. 启用 Compute Engine API：

    启用 API 所需的角色

    如需启用 API，您需要拥有 Service Usage Admin IAM 角色 (roles/serviceusage.serviceUsageAdmin)，该角色包含 serviceusage.services.enable 权限。了解如何授予角色。 

    gcloud services enable compute.googleapis.com
    16.

所需的角色

如需获得创建 TPU 实例、使用 SSH 连接到该实例并运行命令所需的权限，请让您的管理员为您授予项目的以下 IAM 角色：

如需详细了解如何授予角色，请参阅管理对项目、文件夹和组织的访问权限。

您也可以通过自定义角色或其他预定义角色来获取所需的权限。

创建 TPU 实例

在本地终端中或使用 Cloud Shell 运行以下命令。

  1. 如需创建 TPU 实例，请运行以下 gcloud compute instances create 命令：

    gcloud compute instances create quickstart-tpu \
    --machine-type=ct6e-standard-4t \
    --zone=us-east5-b \
    --image-project=ubuntu-os-accelerator-images \
    --image-family=ubuntu-accelerator-2204-amd64-with-tpu-v6e

  2. （可选）使用 gcloud compute instances describe 命令验证实例是否正在运行：

    gcloud compute instances describe quickstart-tpu --zone=us-east5-b

    在输出中，检查 status 字段以验证实例是否正在运行或仍在等待。

连接到 TPU 实例

您可以使用 Google Cloud 控制台或 gcloud CLI 连接到 TPU 实例。

如需使用 SSH 连接到虚拟机，请运行以下 gcloud compute ssh 命令：

gcloud compute ssh quickstart-tpu --project=PROJECT_ID --zone=us-east5-b

安装 JAX

在 TPU 实例上运行以下命令，以安装 JAX 框架：

pip install jax[tpu] -f https://storage.googleapis.com/jax-releases/libtpu_releases.html

运行系统检查

在 TPU 实例上运行以下命令，验证 JAX 是否可以访问 TPU 且可以运行基本操作：

  1. 启动 Python 3 解释器：

    python3

  2. 导入 JAX 模块：

    >>> import jax

  3. 显示可用的 TPU 核心数：

    >>> jax.device_count()

    输出显示 TPU 核心数。核心数取决于您使用的 TPU 版本。如需了解详情，请参阅 TPU 版本。

执行计算

  1. 运行以下命令，将两个数字相加：

    >>> jax.numpy.add(1, 1)

    输出应类似如下所示：

    >>> Array(2, dtype=int32, weak_type=True)

  2. 退出 Python 解释器：

    >>> exit()

清理

为避免因本页面中使用的资源导致您的 Google Cloud 账号产生费用，请删除包含这些资源的 Google Cloud 项目。

为避免因本页中使用的资源导致您的 Google Cloud 账号产生费用，请使用 gcloud compute instances delete 命令删除 TPU 实例：

gcloud compute instances delete quickstart-tpu --zone=us-east5-b

后续步骤