在 TPU 切片上运行 JAX 代码

在运行本文档中的命令之前，请确保已按照设置账号和 Cloud TPU 项目中的说明操作。

在单个 TPU 板上运行 JAX 代码后，您可以通过在 TPU 切片上运行代码来扩容代码。TPU 切片是通过专用高速网络连接相互连接的多个 TPU 板。本文档介绍了如何在 TPU 切片上运行 JAX 代码；如需了解更深入的信息，请参阅在多主机和多进程环境中使用 JAX。

创建 Cloud TPU 切片

1. 创建一些环境变量：

   export PROJECT_ID=your-project-id
   export TPU_NAME=your-tpu-name
   export ZONE=europe-west4-b
   export ACCELERATOR_TYPE=v5litepod-32
   export RUNTIME_VERSION=v2-alpha-tpuv5-lite

   环境变量说明

   变量	说明
   PROJECT_ID	您的 Google Cloud 项目 ID。使用现有项目或创建新项目。
   TPU_NAME	TPU 的名称。
   ZONE	要在其中创建 TPU 虚拟机的可用区。如需详细了解支持的可用区，请参阅 TPU 区域和可用区。
   ACCELERATOR_TYPE	加速器类型用于指定您要创建的 Cloud TPU 的版本和大小。如需详细了解每个 TPU 版本支持的加速器类型，请参阅 TPU 版本。
   RUNTIME_VERSION	Cloud TPU 软件版本。

2. 使用 gcloud 命令创建 TPU 切片。例如，如需创建 v5litepod-32 切片，请使用以下命令：

   $ gcloud compute tpus tpu-vm create ${TPU_NAME}  \
       --zone=${ZONE} \
       --project=${PROJECT_ID} \
       --accelerator-type=${ACCELERATOR_TYPE}  \
       --version=${RUNTIME_VERSION} 

在切片上安装 JAX

创建 TPU 切片之后，您必须在 TPU 切片中的所有主机上安装 JAX。您可以使用 gcloud compute tpus tpu-vm ssh 命令以及 --worker=all 和 --commamnd 参数来执行此操作。

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

在切片上运行 JAX 代码

如需在 TPU 切片上运行 JAX 代码，您必须在 TPU 切片中的每个主机上运行该代码。jax.device_count() 调用在它在切片中的每个主机上被调用之前会停止响应。以下示例说明了如何在 TPU 切片上运行 JAX 计算。

准备代码

您需要 gcloud 344.0.0 版或更高版本（对于 scp 命令）。使用 gcloud --version 检查您的 gcloud 版本，并根据需要运行 gcloud components upgrade。

创建一个名为 example.py 的文件，其中包含以下代码：

import jax

# The total number of TPU cores in the slice
device_count = jax.device_count()

# The number of TPU cores attached to this host
local_device_count = jax.local_device_count()

# The psum is performed over all mapped devices across the slice
xs = jax.numpy.ones(jax.local_device_count())
r = jax.pmap(lambda x: jax.lax.psum(x, 'i'), axis_name='i')(xs)

# Print from a single host to avoid duplicated output
if jax.process_index() == 0:
    print('global device count:', jax.device_count())
    print('local device count:', jax.local_device_count())
    print('pmap result:', r)

将 example.py 复制到切片中的所有 TPU 工作器虚拟机

$ gcloud compute tpus tpu-vm scp ./example.py ${TPU_NAME}: \
  --worker=all \
  --zone=${ZONE} \
  --project=${PROJECT_ID}

如果您以前没有使用过 scp 命令，则可能会看到类似于以下内容的错误：

ERROR: (gcloud.alpha.compute.tpus.tpu-vm.scp) SSH Key is not present in the SSH
agent. Please run `ssh-add /.../.ssh/google_compute_engine` to add it, and try
again.

如需解决此错误，请运行错误消息中显示的 ssh-add 命令，并重新运行该命令。

在切片上运行代码

在每个虚拟机上启动 example.py 程序：

$ gcloud compute tpus tpu-vm ssh ${TPU_NAME} \
  --zone=${ZONE} \
  --project=${PROJECT_ID} \
  --worker=all \
  --command="python3 ./example.py"

输出（使用 v5litepod-32 切片生成）：

global device count: 32
local device count: 4
pmap result: [32. 32. 32. 32.]

清理

完成 TPU 虚拟机的操作后，请按照以下步骤清理资源。

1. 删除您的 Cloud TPU 和 Compute Engine 资源。

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

2. 通过运行 gcloud compute tpus execution-groups list 验证资源是否已删除。删除操作可能需要几分钟时间才能完成。以下命令的输出不应包含本教程中创建的任何资源：

   $ gcloud compute tpus tpu-vm list --zone=${ZONE} \
   --project=${PROJECT_ID}