创建多主机 TPU 切片

创建实例模板

如需为 TPU v6e 虚拟机创建实例模板，请使用 gcloud compute instance-templates create 命令：

gcloud compute instance-templates create quickstart-tpu-instance-template \
    --machine-type=ct6e-standard-4t \
    --maintenance-policy=TERMINATE \
    --image-family=ubuntu-accel-2204-amd64-tpu-v5e-v5p-v6e \
    --image-project=ubuntu-os-accelerator-images \
    --region=us-east5

创建工作负载政策

工作负载政策定义了计算实例的物理属性。在 TPU 切片中，加速器拓扑定义了 TPU 芯片的物理排列方式。 对于互连的多主机 TPU 切片，必须指定加速器拓扑。

如需为多主机 TPU 切片创建工作负载政策，请使用带有 --accelerator-topology 标志的 gcloud compute resource-policies create workload-policy command 命令。以下命令会创建一个具有 2x4 拓扑的工作负载政策：

gcloud compute resource-policies create workload-policy quickstart-tpu-workload-policy \
    --type=high-throughput \
    --accelerator-topology=2x4 \
    --region=us-east5

创建 MIG

运行以下命令以创建构成多主机 TPU 切片的 MIG。

1. 如需创建构成多主机 TPU 切片的 MIG，请使用 gcloud compute instance-groups managed create 命令：

   gcloud compute instance-groups managed create quickstart-tpu-mig \
       --size=2 \
       --target-size-policy-mode=bulk \
       --template=quickstart-tpu-instance-template \
       --region=us-east5 \
       --target-distribution-shape=any-single-zone \
       --instance-redistribution-type=none \
       --default-action-on-vm-failure=do-nothing \
       --workload-policy=projects/PROJECT_ID/regions/us-east5/resourcePolicies/quickstart-tpu-workload-policy
   

   将 PROJECT_ID 替换为您的 Cloud de Confiance by S3NS 项目 ID。

2. （可选）使用以下命令验证代管式实例是否正在运行：

   • 如需查看 MIG 的整体状态，请使用 gcloud compute instance-groups managed describe command：

     gcloud compute instance-groups managed describe quickstart-tpu-mig \
         --region=us-east5
     

   • 如需查看每个实例的状态，请使用 gcloud compute instance-groups managed list-instances command：

     gcloud compute instance-groups managed list-instances quickstart-tpu-mig \
         --region=us-east5
     

安装 JAX

在 MIG 中所有 TPU 虚拟机实例的虚拟环境中安装依赖项和 JAX 框架。如果您的 TPU 虚拟机安装的 Python 版本低于 3.11，则必须安装 Python 3.11 才能运行最新版本的 JAX。

1. 检查 TPU 虚拟机上运行的 Python 版本：

   gcloud compute instance-groups managed list-instances quickstart-tpu-mig \
       --region=us-east5 \
       --uri \
   | xargs -I {} -P 0 gcloud compute ssh {} \
       --command='python3 --version'
   

   如果版本低于 Python 3.11，请安装 Python 3.11：

   gcloud compute instance-groups managed list-instances quickstart-tpu-mig \
       --region=us-east5 \
       --uri \
   | xargs -I {} -P 0 gcloud compute ssh {} \
       --command='sudo apt update && \
       sudo apt install -y software-properties-common && \
       sudo add-apt-repository -y ppa:deadsnakes/ppa && \
       sudo apt update && \
       sudo apt install -y python3.11 python3.11-dev'
   

2. 创建虚拟环境：

   gcloud compute instance-groups managed list-instances quickstart-tpu-mig \
       --region=us-east5 \
       --uri \
   | xargs -I {} -P 0 gcloud compute ssh {} \
       --command='sudo apt install -y python3.11-venv && \
       python3.11 -m venv ~/jax_venv'
   

3. 在虚拟环境中安装 JAX：

   gcloud compute instance-groups managed list-instances quickstart-tpu-mig \
       --region=us-east5 \
       --uri \
   | xargs -I {} -P 0 gcloud compute ssh {} \
       --command='source ~/jax_venv/bin/activate && \
       pip install --upgrade pip -q && \
       pip install jax[tpu] -f https://storage.googleapis.com/jax-releases/libtpu_releases.html -q'
   

在切片上运行 JAX 代码

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

准备代码

在每个实例上创建一个名为 example.py 的文件：

gcloud compute instance-groups managed list-instances quickstart-tpu-mig \
    --region=us-east5 \
    --uri \
| xargs -I {} -P 0 gcloud compute ssh {} \
    --command="cat << 'EOF' > ~/example.py
import jax

# Initialize the slice
jax.distributed.initialize()

# 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)
EOF"

在切片上运行代码

在切片中的每个 TPU 虚拟机上运行 example.py 程序：

gcloud compute instance-groups managed list-instances quickstart-tpu-mig \
    --region=us-east5 \
    --uri \
| xargs -I {} -P 0 gcloud compute ssh {} \
    --command='source ~/jax_venv/bin/activate && python3 ~/example.py'

输出应类似如下所示：

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