Cloud TPU v5e 训练

Cloud TPU v5e 是 Google Cloud 最新一代 AI 加速器。每个 Pod 占用的 256 个芯片更少，v5e 经过优化，是 Transformer、文本到图像和卷积神经网络 (CNN) 训练、微调和传送的最具价值的产品。

Cloud TPU v5e 概念、系统架构和配置

如果您刚开始接触 Cloud TPU，请查看 TPU 文档首页。

Cloud TPU 系统架构页面介绍了一般 Cloud TPU 概念（例如切片、主机、芯片和 TensorCore）以及 Cloud TPU 系统架构。

每个 Cloud TPU 版本都需要特定的加速器类型来进行训练和推断。v5e 配置部分介绍了这些加速器类型。

推理

推断是使用训练好的模型对新数据进行预测的过程。供传送进程使用。

切片

一个切片表示全部位于同一个 Pod 内通过高速芯片间互连 (ICI) 连接的芯片的集合。v5e 具有 2D 切片形状。如需了解支持的切片形状，请参阅 v5e 配置部分中的表格。

条状标签形状和条状标签拓扑也是指切片形状。

正在处理

服务是将经过训练的机器学习模型部署到生产环境的过程，在该环境中，模型可用于进行预测或决策。延迟时间和服务级可用性对于服务很重要。

单个主机与多主机

主机是运行虚拟机的物理计算机 (CPU)。一个主机可以同时运行多个虚拟机。

使用少于 8 个条状标签的切片最多使用一个主机。超过 8 个芯片的切片可以访问多个主机，并且可以使用多个主机运行分布式训练。如需详细了解切片和芯片，请参阅 TPU 系统架构页面。

v5e 支持多主机训练和多主机推理（使用 SAX）。

TPU 虚拟机

运行 Linux 的虚拟机，能够访问底层 TPU。对于 v5e TPU，每个 TPU 虚拟机可以直接访问 1、4 或 8 个芯片，具体取决于用户指定的加速器类型。TPU 虚拟机也称为“工作器”。

工作器

请参阅 TPU 虚拟机。

开始使用

如需了解 v5e TPU 硬件，请参阅系统架构。

保障容量

联系 Cloud 销售团队，开始使用 Cloud TPU v5e 处理 AI 工作负载。

准备 Google Cloud 项目

  1. 登录您的 Google 账号。 如果您尚未注册新帐号，请先注册。
  2. 在 Google Cloud 控制台中，从项目选择器页面选择或创建 Google Cloud 项目。

  3. 所有 Google Cloud 用量都需要结算设置，因此请确保您的项目已启用结算功能。

  4. 安装 gcloud alpha components。

  5. 如果您是重复使用现有 gcloud alpha 组件的 TPU 用户，请更新这些组件以确保相关命令和标志受支持：

    gcloud components update

  6. 在 Cloud Shell 中通过以下 gcloud 命令启用 TPU API。（您也可以通过 Google Cloud 控制台启用该功能。）

    gcloud services enable tpu.googleapis.com

  7. 启用 TPU 服务帐号。

    服务帐号允许 Cloud TPU 服务访问其他 Google Cloud 服务。建议使用用户代管式服务帐号。请按照相关指南创建和授予角色。以下角色必不可少：

    • TPU Admin
    • Storage Admin：访问 Cloud Storage 时需要
    • Logs Writer：使用 Logging API 写入日志时需要
    • Monitoring Metric Writer：将指标写入 Cloud Monitoring 时需要用到

  8. 配置项目和区域。

    您的项目 ID 是 Cloud 控制台上显示的项目名称。

    export PROJECT_ID=your-project-id
export ZONE=us-west4-a

gcloud alpha compute tpus tpu-vm service-identity create --zone=${ZONE}

gcloud auth login
gcloud config set project ${PROJECT_ID}
gcloud config set compute/zone ${ZONE}

预配 Cloud TPU 环境

最佳实践是使用 queued-resource create 命令将 Cloud TPU v5 预配为已加入队列的资源。但是，您也可以使用 Create Node API (gcloud alpha compute tpus tpu-vm create) 预配 Cloud TPU v5。

创建环境变量

设置创建 TPU 所需的环境变量。

将以下列表中的变量（红色）替换为您将用于训练或推理作业的值。

export PROJECT_ID=your_project_ID
export ACCELERATOR_TYPE=v5litepod-16
export ZONE=us-west4-a
export RUNTIME_VERSION=v2-alpha-tpuv5-lite
export SERVICE_ACCOUNT=your_service_account
export TPU_NAME=your_tpu_name
export QUEUED_RESOURCE_ID=your_queued_resource_id
export QUOTA_TYPE=quota_type
export VALID_UNTIL_DURATION=1d
变量 说明
PROJECT_ID Google Cloud 项目名称
ACCELERATOR_TYPE 如需了解支持的加速器类型，请参阅加速器类型部分。
区间 所有容量均位于 us-west4-a。
RUNTIME_VERSION 请使用 v2-alpha-tpuv5-lite。
SERVICE_ACCOUNT 这是您的服务帐号的地址，您可以在 Google Cloud 控制台 ->“IAM”->“服务帐号”中找到它。例如：tpu-service-account@myprojectID.iam.gserviceaccount.com
TPU_NAME TPU 的文本 ID，由用户分配的文本 ID，在分配排队的资源请求时创建。
QUEUED_RESOURCE_ID 已加入队列的资源请求且由用户分配的文本 ID。如需了解已加入队列的资源，请参阅已加入队列的资源文档。
QUOTA_TYPE 可以是 reserved 或 best-effort。如果这两种类型都未指定，则 QUOTA_TYPE 默认为 on-demand。 如需了解 Cloud TPU 支持的不同类型的配额，请参阅quotas。
VALID_UNTIL_DURATION 请求的有效时长。如需了解不同的有效期，请参阅 排入队列的资源。

创建 TPU 资源

gcloud alpha compute tpus queued-resources create ${QUEUED_RESOURCE_ID} \
   --node-id=${TPU_NAME} \
   --project=${PROJECT_ID} \
   --zone=${ZONE} \
   --accelerator-type=${ACCELERATOR_TYPE} \
   --runtime-version=${RUNTIME_VERSION} \
   --valid-until-duration=${VALID_UNTIL_DURATION} \
   --service-account=${SERVICE_ACCOUNT} \
   --${QUOTA_TYPE}

如果成功创建了已加入队列的资源，response 字段中的状态将为 WAITING_FOR_RESOURCES 或 FAILED。如果已加入队列的资源处于 WAITING_FOR_RESOURCES 状态，则意味着已加入队列的资源已通过初步验证，正在等待容量。当容量可用后，请求将转换为 PROVISIONING。处于 WAITING_FOR_RESOURCES 状态并不一定意味着您能获得分配的配额，并且系统可能需要一些时间才能从 WAITING_FOR_RESOURCES 状态更改为 ACTIVE 状态。如果已加入队列的资源处于 FAILED 状态，则输出中会显示失败原因。如果未在 --valid-until-duration 中填充请求，并且状态变为“FAILED”，该请求将过期。

待排队的资源处于 ACTIVE 状态后，您就可以使用 SSH 访问 TPU 虚拟机。

使用 [list](/tpu/docs/managing-tpus-tpu-vm) 或 [describe](/tpu/docs/managing-tpus-tpu-vm) 命令查询已加入队列的资源的状态。

gcloud alpha compute tpus queued-resources describe ${QUEUED_RESOURCE_ID} \
   --project=${PROJECT_ID} \
   --zone=${ZONE}

state 表示已加入队列的资源的状态。状态的定义如下：

状态 说明
WAITING_FOR_RESOURCES 已收到已加入队列的资源 create 命令，该命令将在容量可用时立即开始预配。
PROVISIONING 正在预配 TPU 切片。
ACTIVE 所有 TPU 均已预配，随时可以使用。如果给定了启动脚本，当排队的资源状态转换为 ACTIVE 时，该脚本将在所有 TPU 上开始执行。
FAILED 无法预配这些切片。
SUSPENDING 正在删除一个或多个切片。
SUSPENDED 系统会删除所有底层切片，但排队的资源将保持不变，直到明确删除为止。目前，已暂停队列中的资源无法恢复，应将其删除。
DELETING 正在删除已加入队列的资源。

使用 SSH 连接到 TPU 虚拟机

以下部分介绍如何在 TPU 切片中的每个 TPU 虚拟机上安装二进制文件并运行代码。在此上下文中，TPU 虚拟机也称为“工作器”。

如需确定您的切片将拥有的虚拟机数量，请参阅虚拟机类型部分。

如需安装二进制文件或运行代码，请使用 tpu-vm ssh 命令连接到 TPU 虚拟机。

gcloud alpha compute tpus tpu-vm ssh ${TPU_NAME}

如需使用 SSH 访问特定的 TPU 虚拟机，请使用 --worker 标志，该标志后跟索引（从 0 开始）：

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

如果切片形状超过 8 个芯片，则一个切片中会有多个虚拟机。在这种情况下，请使用 --worker=all 标志在所有 TPU 虚拟机上运行安装，而不必分别连接到每个虚拟机。例如：

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

管理

管理 TPU 中介绍了可用于管理 TPU 虚拟机的所有命令。

框架设置

本部分介绍了搭配使用 JAX 或 PyTorch 与 TPU v5e 训练自定义模型的常规设置过程。tpu-vm-tf-2.15.0-pjrt 和 tpu-vm-tf-2.15.0-pod-pjrt TPU 运行时版本支持 TensorFlow。

如需了解推断设置说明，请参阅 v5e 推断简介。

JAX 设置

如果切片形状超过 8 个芯片，则一个切片中会有多个虚拟机。在这种情况下，您需要使用 --worker=all 标志在所有 TPU 虚拟机上运行安装步骤，无需使用 SSH 分别登录每个虚拟机：

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

您可以运行以下命令来检查设备数量（此处显示的输出是使用 v5litepod-16 切片生成的）。此代码通过检查 JAX 是否能够看到 Cloud TPU TensorCore，以及是否可以运行基本操作来测试所有内容是否均已正确安装：

gcloud alpha compute tpus tpu-vm ssh ${TPU_NAME} \
   --project=${PROJECT_ID} \
   --zone=${ZONE} \
   --worker=all \
   --command='python3 -c "import jax; print(jax.device_count()); print(jax.local_device_count())"'

输出将如下所示：

SSH: Attempting to connect to worker 0...
SSH: Attempting to connect to worker 1...
SSH: Attempting to connect to worker 2...
SSH: Attempting to connect to worker 3...
16
4
16
4
16
4
16
4

jax.device\_count() 显示给定切片中的条状标签总数。jax.local\_device\_count() 表示此切片中单个虚拟机可访问的芯片数量。

# Check the number of chips in the given slice by summing the count of chips
# from all VMs through the
# jax.local_device_count() API call.
gcloud alpha compute tpus tpu-vm ssh ${TPU_NAME} \
   --project=${PROJECT_ID} \
   --zone=${ZONE} \
   --worker=all \
   --command='python3 -c "import jax; xs=jax.numpy.ones(jax.local_device_count()); print(jax.pmap(lambda x: jax.lax.psum(x, \"i\"), axis_name=\"i\")(xs))"'

输出将如下所示：

SSH: Attempting to connect to worker 0...
SSH: Attempting to connect to worker 1...
SSH: Attempting to connect to worker 2...
SSH: Attempting to connect to worker 3...
[16. 16. 16. 16.]
[16. 16. 16. 16.]
[16. 16. 16. 16.]
[16. 16. 16. 16.]

尝试参阅本文档中的 JAX 教程，开始使用 JAX 进行 v5e 训练。

PyTorch 设置

请注意，v5e 仅支持 PJRT 运行时，PyTorch 2.1 及更高版本会将 PJRT 用作所有 TPU 版本的默认运行时。

本部分介绍如何开始在 v5e 上通过 PyTorch/XLA 通过命令使用 PJRT，为所有工作器使用 PJRT。

安装依赖项

gcloud alpha compute tpus tpu-vm ssh ${TPU_NAME}  \
   --project=${PROJECT_ID} \
   --zone=${ZONE} \
   --worker=all \
   --command='
      sudo apt-get update -y
      sudo apt-get install libomp5 -y
      pip3 install mkl mkl-include
      pip3 install tf-nightly tb-nightly tbp-nightly
      pip3 install numpy
      sudo apt-get install libopenblas-dev -y
      pip3 install torch~=2.1.0 torchvision torch_xla[tpu]~=2.1.0 -f https://storage.googleapis.com/libtpu-releases/index.html
      pip3 install torch_xla[tpu] -f https://storage.googleapis.com/libtpu-releases/index.html'

如果您无法为 Torch/torch_xla/torchvision 安装车轮，并看到 pkg_resources.extern.packaging.requirements.InvalidRequirement: Expected end or semicolon (after name and no valid version specifier) torch==nightly+20230222 这样的错误，请使用以下命令降级您的版本：

pip3 install setuptools==62.1.0

使用 PJRT 运行脚本：

unset LD_PRELOAD

以下是使用 Python 脚本在 v5e 虚拟机上进行计算的示例：

gcloud alpha compute tpus tpu-vm ssh ${TPU_NAME} \
   --project=${PROJECT_ID} \
   --zone=${ZONE} \
   --worker all \
   --command='
      export LD_LIBRARY_PATH=$LD_LIBRARY_PATH:$HOME/.local/lib/
      export PJRT_DEVICE=TPU_C_API
      export PT_XLA_DEBUG=0
      export USE_TORCH=ON
      unset LD_PRELOAD
      export TPU_LIBRARY_PATH=$HOME/.local/lib/python3.10/site-packages/libtpu/libtpu.so
      python3 -c "import torch; import torch_xla; import torch_xla.core.xla_model as xm; print(xm.xla_device()); dev = xm.xla_device(); t1 = torch.randn(3,3,device=dev); t2 = torch.randn(3,3,device=dev); print(t1 + t2)"'

这将生成如下所示的输出：

SSH: Attempting to connect to worker 0...
SSH: Attempting to connect to worker 1...
xla:0
tensor([[ 1.8611, -0.3114, -2.4208],
[-1.0731, 0.3422, 3.1445],
[ 0.5743, 0.2379, 1.1105]], device='xla:0')
xla:0
tensor([[ 1.8611, -0.3114, -2.4208],
[-1.0731, 0.3422, 3.1445],
[ 0.5743, 0.2379, 1.1105]], device='xla:0')

参阅本文档中的 PyTorch 教程，开始使用 PyTorch 进行 v5e 训练。

在会话结束时删除 TPU 和已加入队列的资源。如需删除已加入队列的资源，请按照以下 2 个步骤删除切片，然后删除已加入队列的资源：

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

gcloud alpha compute tpus queued-resources delete ${QUEUED_RESOURCE_ID} \
   --project=${PROJECT_ID} \
   --zone=${ZONE} \
   --quiet

这两个步骤还可用于移除处于 FAILED 状态的已加入队列的资源请求。

监控和性能分析

Cloud TPU v5e 支持使用与上一代 Cloud TPU 相同的方法进行监控和分析。您可以阅读使用 Cloud TPU 工具分析模型，详细了解如何对 Cloud TPU 虚拟机进行性能分析和监控，从而详细了解如何进行监控。

JAX/FLAX 示例

在 v5e 上训练 ImageNet

本教程介绍如何使用虚构输入数据在 v5e 上训练 ImageNet。如果您想使用真实数据，请参阅 GitHub 上的 README 文件。

设置

  1. 创建环境变量：

    export PROJECT_ID=your_project_ID
export ACCELERATOR_TYPE=v5litepod-16
export ZONE=us-west4-a
export RUNTIME_VERSION=v2-alpha-tpuv5-lite
export SERVICE_ACCOUNT=your_service_account
export TPU_NAME=your_tpu_name
export QUEUED_RESOURCE_ID=your_queued_resource_id
export QUOTA_TYPE=quota_type
export VALID_UNTIL_DURATION=1d

  2. 创建 TPU 资源：

    gcloud alpha compute tpus queued-resources create ${QUEUED_RESOURCE_ID} \
   --node-id=${TPU_NAME} \
   --project=${PROJECT_ID} \
   --zone=${ZONE} \
   --accelerator-type=${ACCELERATOR_TYPE} \
   --runtime-version=${RUNTIME_VERSION} \
   --valid-until-duration=${VALID_UNTIL_DURATION} \
   --service-account=${SERVICE_ACCOUNT} \
   --${QUOTA_TYPE}

    待排队的资源处于 ACTIVE 状态后，您可以通过 SSH 连接到 TPU 虚拟机：

    gcloud alpha compute tpus queued-resources describe ${QUEUED_RESOURCE_ID}  \
   --project=${PROJECT_ID} \
   --zone=${ZONE}

    当 QueuedResource 处于 ACTIVE 状态时，输出将如下所示：

     state: ACTIVE

  3. 安装最新版本的 JAX 和 jaxlib：

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

  4. 克隆 ImageNet 模型并安装相应的要求：

    gcloud alpha compute tpus tpu-vm ssh ${TPU_NAME} \
   --project=${PROJECT_ID} \
   --zone=${ZONE} \
   --worker=all \
   --command='git clone https://github.com/google/flax.git && cd flax/examples/imagenet && pip install -r requirements.txt && pip install flax==0.7.4'

  5. 为了生成虚构数据，模型需要有关数据集维度的信息。这可以从 ImageNet 数据集的元数据中收集：

    gcloud alpha compute tpus tpu-vm ssh ${TPU_NAME} --project=${PROJECT_ID} --zone=${ZONE} --worker=all --command='mkdir -p $HOME/flax/.tfds/metadata/imagenet2012/5.1.0 && curl https://raw.githubusercontent.com/tensorflow/datasets/v4.4.0/tensorflow_datasets/testing/metadata/imagenet2012/5.1.0/dataset_info.json --output $HOME/flax/.tfds/metadata/imagenet2012/5.1.0/dataset_info.json'

训练模型

完成前面的所有步骤后，您就可以训练模型了。

gcloud alpha compute tpus tpu-vm ssh ${TPU_NAME} \
   --project=${PROJECT_ID} \
   --zone=${ZONE} \
   --worker=all \
   --command='cd flax/examples/imagenet && JAX_PLATFORMS=tpu python3 imagenet_fake_data_benchmark.py'

删除 TPU 和已加入队列的资源

在会话结束时删除 TPU 和已加入队列的资源。

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

gcloud alpha compute tpus queued-resources delete ${QUEUED_RESOURCE_ID} \
   --project=${PROJECT_ID} \
   --zone=${ZONE} \
   --quiet

拥抱脸 FLAX 模型

采用 FLAX 实现的 Hugging Face 模型可在 Cloud TPU v5e 上开箱即用。本部分介绍了如何运行热门模型。

在 Imagenette 上训练 ViT

本教程介绍了如何在 Cloud TPU v5e 上使用 Fast AI imagenette 数据集通过 HuggingFace 训练 Vision Transformer (ViT) 模型。

ViT 模型是第一个在 ImageNet 上成功训练 Transformer 编码器并比卷积网络取得出色结果的模型。如需了解详情，请参阅以下资源：

设置

  1. 创建环境变量：

    export PROJECT_ID=your_project_ID
export ACCELERATOR_TYPE=v5litepod-16
export ZONE=us-west4-a
export RUNTIME_VERSION=v2-alpha-tpuv5-lite
export SERVICE_ACCOUNT=your_service_account
export TPU_NAME=your_tpu_name
export QUEUED_RESOURCE_ID=your_queued_resource_id
export QUOTA_TYPE=quota_type
export VALID_UNTIL_DURATION=1d

  2. 创建 TPU 资源：

    gcloud alpha compute tpus queued-resources create ${QUEUED_RESOURCE_ID} \
   --node-id=${TPU_NAME} \
   --project=${PROJECT_ID} \
   --zone=${ZONE} \
   --accelerator-type=${ACCELERATOR_TYPE} \
   --runtime-version=${RUNTIME_VERSION} \
   --valid-until-duration=${VALID_UNTIL_DURATION} \
   --service-account=${SERVICE_ACCOUNT} \
   --${QUOTA_TYPE}

    待排队的资源处于 ACTIVE 状态后，您将能够通过 SSH 连接到 TPU 虚拟机：

    gcloud alpha compute tpus queued-resources describe ${QUEUED_RESOURCE_ID}  \
   --project=${PROJECT_ID} \
   --zone=${ZONE}

    当已加入队列的资源处于 ACTIVE 状态时，输出将类似于以下内容：

     state: ACTIVE

  3. 安装 JAX 及其库：

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

  4. 下载 Hugging Face 代码库和安装要求：

    gcloud alpha compute tpus tpu-vm ssh ${TPU_NAME} \
   --project=${PROJECT_ID} \
   --zone=${ZONE} \
   --worker=all \
   --command='git clone https://github.com/huggingface/transformers.git && cd transformers && pip install . && pip install -r examples/flax/_tests_requirements.txt && pip install --upgrade huggingface-hub urllib3 zipp && pip install tensorflow==2.16.1 && pip install -r examples/flax/vision/requirements.txt'

  5. 下载 Imagenette 数据集：

    gcloud alpha compute tpus tpu-vm ssh ${TPU_NAME} \
   --project=${PROJECT_ID} \
   --zone=${ZONE} \
   --worker=all \
   --command='cd transformers && wget https://s3.amazonaws.com/fast-ai-imageclas/imagenette2.tgz && tar -xvzf imagenette2.tgz'

训练模型

使用 4GB 的预映射缓冲区来训练模型。

gcloud alpha compute tpus tpu-vm ssh ${TPU_NAME} \
   --project=${PROJECT_ID} \
   --zone=${ZONE} \
   --worker=all \
   --command='cd transformers && JAX_PLATFORMS=tpu python3 examples/flax/vision/run_image_classification.py --train_dir "imagenette2/train" --validation_dir "imagenette2/val" --output_dir "./vit-imagenette" --learning_rate 1e-3 --preprocessing_num_workers 32 --per_device_train_batch_size 8 --per_device_eval_batch_size 8 --model_name_or_path google/vit-base-patch16-224-in21k --num_train_epochs 3'

删除 TPU 和已加入队列的资源

在会话结束时删除 TPU 和排入队列的资源。

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

gcloud alpha compute tpus queued-resources delete ${QUEUED_RESOURCE_ID} \
   --project=${PROJECT_ID} \
   --zone=${ZONE} \
   --quiet

ViT 基准比较结果

训练脚本在 v5litepod-4、v5litepod-16 和 v5litepod-64 上运行，下表显示了使用不同加速器类型时的吞吐量。

加速器类型 v5litepod-4 v5litepod-16 V5litepod-64
纪元 3 3 3
全局批量大小 32 128 512
吞吐量（示例/秒） 263.40 429.34 470.71

《Pokémon》中的 Train Diffusion

本教程介绍如何使用 Cloud TPU v5e 上的 Pokémon 数据集通过 HuggingFace 训练 Stable Diffusion 模型。

稳定扩散模型是一种潜在的文本到图像模型，可根据任何文本输入生成逼真的图像。如需了解详情，请参阅以下资源：

设置

  1. 创建环境变量：

    export PROJECT_ID=your_project_ID
export ACCELERATOR_TYPE=v5litepod-16
export ZONE=us-west4-a
export RUNTIME_VERSION=v2-alpha-tpuv5-lite
export SERVICE_ACCOUNT=your_service_account
export TPU_NAME=your_tpu_name
export QUEUED_RESOURCE_ID=queued_resource_id
export QUOTA_TYPE=quota_type
export VALID_UNTIL_DURATION=1d

  2. 创建 TPU 资源：

    gcloud alpha compute tpus queued-resources create ${QUEUED_RESOURCE_ID} \
   --node-id=${TPU_NAME} \
   --project=${PROJECT_ID} \
   --zone=${ZONE} \
   --accelerator-type=${ACCELERATOR_TYPE} \
   --runtime-version=${RUNTIME_VERSION} \
   --valid-until-duration=${VALID_UNTIL_DURATION} \
   --service-account=${SERVICE_ACCOUNT} \
   --${QUOTA_TYPE}

    待排队的资源处于 ACTIVE 状态后，您将能够通过 SSH 连接到 TPU 虚拟机：

    gcloud alpha compute tpus queued-resources describe ${QUEUED_RESOURCE_ID}  \
   --project=${PROJECT_ID} \
   --zone=${ZONE}

    当已加入队列的资源处于 ACTIVE 状态时，输出将类似于以下内容：

     state: ACTIVE

  3. 安装 JAX 及其库。

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

  4. 下载 HuggingFace 代码库和安装要求。

    gcloud alpha compute tpus tpu-vm ssh ${TPU_NAME} \
   --project=${PROJECT_ID} \
   --zone=${ZONE} \
   --worker=all \
   --command='git clone https://github.com/RissyRan/diffusers.git && cd diffusers && pip install . && pip install tensorflow==2.16.1 clu && pip install -U -r examples/text_to_image/requirements_flax.txt'

训练模型

使用 4GB 的预映射缓冲区来训练模型。

gcloud alpha compute tpus tpu-vm ssh ${TPU_NAME} \
   --project=${PROJECT_ID} \
   --zone=${ZONE} \
   --worker=all \
   --command='cd diffusers/examples/text_to_image && JAX_PLATFORMS=tpu,cpu python3 train_text_to_image_flax.py --pretrained_model_name_or_path=duongna/stable-diffusion-v1-4-flax --dataset_name=lambdalabs/pokemon-blip-captions --resolution=128 --center_crop --random_flip --train_batch_size=4 --mixed_precision=fp16 --max_train_steps=1500 --learning_rate=1e-05 --max_grad_norm=1 --output_dir=sd-pokemon-model'

删除 TPU 和排入队列的资源

在会话结束时删除 TPU 和已加入队列的资源。

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

gcloud alpha compute tpus queued-resources delete ${QUEUED_RESOURCE_ID} \
   --project=${PROJECT_ID} \
   --zone=${ZONE} \
   --quiet

扩散的基准测试结果

训练脚本在 v5litepod-4、v5litepod-16 和 v5litepod-64 上运行，下表显示了吞吐量。

加速器类型 v5litepod-4 v5litepod-16 V5litepod-64
训练步数 1500 1500 1500
全局批量大小 32 64 128
吞吐量（示例/秒） 36.53 43.71 49.36 岁

使用 OSCAR 数据集训练 GPT2

本教程介绍了如何在 Cloud TPU v5e 上使用 OSCAR 数据集通过 HuggingFace 训练 GPT2 模型。

GPT2 是一种基于原始文本进行预训练的转换器模型，无需人工标记。经训练，它能够预测句子中的下一个字词。如需了解详情，请参阅以下资源：

设置

  1. 创建环境变量：

    export PROJECT_ID=your_project_ID
export ACCELERATOR_TYPE=v5litepod-16
export ZONE=us-west4-a
export RUNTIME_VERSION=v2-alpha-tpuv5-lite
export SERVICE_ACCOUNT=your_service_account
export TPU_NAME=your_tpu_name
export QUEUED_RESOURCE_ID=queued_resource_id
export QUOTA_TYPE=quota_type
export VALID_UNTIL_DURATION=1d

  2. 创建 TPU 资源：

    gcloud alpha compute tpus queued-resources create ${QUEUED_RESOURCE_ID} \
   --node-id=${TPU_NAME} \
   --project=${PROJECT_ID} \
   --zone=${ZONE} \
   --accelerator-type=${ACCELERATOR_TYPE} \
   --runtime-version=${RUNTIME_VERSION} \
   --valid-until-duration=${VALID_UNTIL_DURATION} \
   --service-account=${SERVICE_ACCOUNT} \
   --${QUOTA_TYPE}

    排队的资源处于 ACTIVE 状态后，您可以通过 SSH 连接到 TPU 虚拟机：

    gcloud alpha compute tpus queued-resources describe ${QUEUED_RESOURCE_ID} \
   --project=${PROJECT_ID} \
   --zone=${ZONE}

    当已加入队列的资源处于 ACTIVE 状态时，输出将类似于以下内容：

     state: ACTIVE

  3. 安装 JAX 及其库。

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

  4. 下载 HuggingFace 代码库和安装要求。

    gcloud alpha compute tpus tpu-vm ssh ${TPU_NAME} \
   --project=${PROJECT_ID} \
   --zone=${ZONE} \
   --worker=all \
   --command='git clone https://github.com/huggingface/transformers.git && cd transformers && pip install . && pip install -r examples/flax/_tests_requirements.txt && pip install --upgrade huggingface-hub urllib3 zipp && pip install tensorflow && pip install -r examples/flax/language-modeling/requirements.txt'

  5. 下载配置以训练模型。

    gcloud alpha compute tpus tpu-vm ssh ${TPU_NAME} \
   --project=${PROJECT_ID} \
   --zone=${ZONE} \
   --worker=all \
   --command='cd transformers/examples/flax/language-modeling && gsutil cp -r gs://cloud-tpu-tpuvm-artifacts/v5litepod-preview/jax/gpt .'

训练模型

使用 4GB 的预映射缓冲区来训练模型。

gcloud alpha compute tpus tpu-vm ssh ${TPU_NAME} \
   --project=${PROJECT_ID} \
   --zone=${ZONE} \
   --worker=all \
   --command='cd transformers/examples/flax/language-modeling && TPU_PREMAPPED_BUFFER_SIZE=4294967296 JAX_PLATFORMS=tpu python3 run_clm_flax.py --output_dir=./gpt --model_type=gpt2 --config_name=./gpt --tokenizer_name=./gpt --dataset_name=oscar --dataset_config_name=unshuffled_deduplicated_no --do_train --do_eval --block_size=512 --per_device_train_batch_size=4 --per_device_eval_batch_size=4 --learning_rate=5e-3 --warmup_steps=1000 --adam_beta1=0.9 --adam_beta2=0.98 --weight_decay=0.01 --overwrite_output_dir --num_train_epochs=3 --logging_steps=500 --eval_steps=2500'

删除 TPU 和已加入队列的资源

在会话结束时删除 TPU 和排入队列的资源。

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

gcloud alpha compute tpus queued-resources delete ${QUEUED_RESOURCE_ID} \
   --project=${PROJECT_ID} \
   --zone=${ZONE} \
   --quiet

GPT2 的基准化分析结果

训练脚本在 v5litepod-4、v5litepod-16 和 v5litepod-64 上运行，下表显示了吞吐量。

v5litepod-4 v5litepod-16 V5litepod-64
纪元 3 3 3
全局批量大小 64 64 64
吞吐量（示例/秒） 74.60 72.97 72.62

PyTorch/XLA

使用 PJRT 运行时训练 ResNet

PyTorch/XLA 正在从 PyTorch 2.0+ 迁移到 PjRt。以下是为 PyTorch/XLA 训练工作负载设置 v5e 的更新说明。

设置

  1. 创建环境变量：

    export PROJECT_ID=your_project_ID
export ACCELERATOR_TYPE=v5litepod-16
export ZONE=us-west4-a
export RUNTIME_VERSION=v2-alpha-tpuv5-lite
export SERVICE_ACCOUNT=your_service_account
export TPU_NAME=tpu-name
export QUEUED_RESOURCE_ID=queued_resource_id
export QUOTA_TYPE=quota_type
export VALID_UNTIL_DURATION=1d

  2. 创建 TPU 资源：

    gcloud alpha compute tpus queued-resources create ${QUEUED_RESOURCE_ID} \
   --node-id=${TPU_NAME} \
   --project=${PROJECT_ID} \
   --zone=${ZONE} \
   --accelerator-type=${ACCELERATOR_TYPE} \
   --runtime-version=${RUNTIME_VERSION} \
   --valid-until-duration=${VALID_UNTIL_DURATION} \
   --service-account=${SERVICE_ACCOUNT} \
   --{QUOTA_TYPE}

    待 QueuedResource 处于 ACTIVE 状态后，您将能够通过 SSH 连接到 TPU 虚拟机：

    gcloud alpha compute tpus queued-resources describe ${QUEUED_RESOURCE_ID} \
   --project=${PROJECT_ID} \
   --zone=${ZONE}

    当已加入队列的资源处于 ACTIVE 状态时，输出将类似于以下内容：

     state: ACTIVE

  3. 安装 Torch/XLA 专用依赖项

    gcloud alpha compute tpus tpu-vm ssh ${TPU_NAME} \
   --project=${PROJECT_ID} \
   --zone=${ZONE} \
   --worker=all \
   --command='
      sudo apt-get update -y
      sudo apt-get install libomp5 -y
      pip3 install mkl mkl-include
      pip3 install tf-nightly tb-nightly tbp-nightly
      pip3 install numpy
      sudo apt-get install libopenblas-dev -y
      pip3 install torch~=2.1.0 torchvision torch_xla[tpu]~=2.1.0 -f https://storage.googleapis.com/libtpu-releases/index.html
      pip3 install torch_xla[tpu] -f https://storage.googleapis.com/libtpu-releases/index.html'
训练 ResNet 模型
gcloud alpha compute tpus tpu-vm ssh ${TPU_NAME} \
   --project=${PROJECT_ID} \
   --zone=${ZONE} \
   --worker=all \
   --command='
      date
      export PJRT_DEVICE=TPU_C_API
      export PT_XLA_DEBUG=0
      export USE_TORCH=ON
      export XLA_USE_BF16=1
      export LIBTPU_INIT_ARGS=--xla_jf_auto_cross_replica_sharding
      export LD_LIBRARY_PATH=/usr/local/lib:$LD_LIBRARY_PATH
      export TPU_LIBRARY_PATH=$HOME/.local/lib/python3.10/site-packages/libtpu/libtpu.so
      git clone https://github.com/pytorch/xla.git
      cd xla/
      git reset --hard caf5168785c081cd7eb60b49fe4fffeb894c39d9
      python3 test/test_train_mp_imagenet.py --model=resnet50  --fake_data --num_epochs=1 —num_workers=16  --log_steps=300 --batch_size=64 --profile'

删除 TPU 和已加入队列的资源

在会话结束时删除 TPU 和排入队列的资源。

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

gcloud alpha compute tpus queued-resources delete ${QUEUED_RESOURCE_ID} \
   --project=${PROJECT_ID} \
   --zone=${ZONE} \
   --quiet
基准测试结果

下表显示了基准吞吐量。

加速器类型 吞吐量（示例/秒）
v5litepod-4 4240 ex/s
v5litepod-16 10,810 ex/s
V5litepod-64 46,154 ex/s

在 v5e 上训练 GPT2

本教程介绍了如何使用 Wikitext 数据集在 PyTorch/XLA 上使用 HuggingFace 代码库在 v5e 上运行 GPT2。

设置

  1. 创建环境变量：

    export PROJECT_ID=your_project_ID
export ACCELERATOR_TYPE=v5litepod-16
export ZONE=us-west4-a
export RUNTIME_VERSION=v2-alpha-tpuv5-lite
export SERVICE_ACCOUNT=your_service_account
export TPU_NAME=your_tpu_name
export QUEUED_RESOURCE_ID=queued_resource_id
export QUOTA_TYPE=quota_type
export VALID_UNTIL_DURATION=1d

  2. 创建 TPU 资源：

    gcloud alpha compute tpus queued-resources create ${QUEUED_RESOURCE_ID} \
   --node-id=${TPU_NAME} \
   --project=${PROJECT_ID} \
   --zone=${ZONE} \
   --accelerator-type=${ACCELERATOR_TYPE} \
   --runtime-version=${RUNTIME_VERSION} \
   --valid-until-duration=${VALID_UNTIL_DURATION} \
   --service-account=${SERVICE_ACCOUNT} \
   --${QUOTA_TYPE}

    待 QueuedResource 处于 ACTIVE 状态后，您将能够通过 SSH 连接到 TPU 虚拟机：

    gcloud alpha compute tpus queued-resources describe ${QUEUED_RESOURCE_ID} \
   --project=${PROJECT_ID} \
   --zone=${ZONE}

    当已加入队列的资源处于 ACTIVE 状态时，输出将类似于以下内容：

    state: ACTIVE

  3. 安装 Torch/xla 依赖项。

    gcloud alpha compute tpus tpu-vm ssh ${TPU_NAME} \
   --project=${PROJECT_ID} \
   --zone=${ZONE} \
   --worker=all \
   --command='
      sudo apt-get -y update
      sudo apt install -y libopenblas-base
      pip3 install torchvision
      pip3 uninstall -y torch
      pip3 install torch~=2.1.0 torchvision torch_xla[tpu]~=2.1.0 -f https://storage.googleapis.com/libtpu-releases/index.html
      pip3 install torch_xla[tpu] -f https://storage.googleapis.com/libtpu-releases/index.html'

  4. 下载 HuggingFace 代码库和安装要求。

    gcloud alpha compute tpus tpu-vm ssh ${TPU_NAME} \
   --project=${PROJECT_ID} \
   --zone=${ZONE} \
   --worker=all \
   --command='
      git clone https://github.com/pytorch/xla.git
      pip install --upgrade accelerate
      git clone https://github.com/huggingface/transformers.git
      cd transformers
      git checkout ebdb185befaa821304d461ed6aa20a17e4dc3aa2
      pip install .
      git log -1
      pip install datasets evaluate scikit-learn
      '

  5. 下载预训练模型的配置。

    gcloud alpha compute tpus tpu-vm ssh ${TPU_NAME} \
   --project=${PROJECT_ID} \
   --zone=${ZONE} \
   --worker=all \
   --command='
      gsutil cp -r gs://cloud-tpu-tpuvm-artifacts/config/xl-ml-test/pytorch/gpt2/my_config_2.json transformers/examples/pytorch/language-modeling/
      gsutil cp gs://cloud-tpu-tpuvm-artifacts/config/xl-ml-test/pytorch/gpt2/fsdp_config.json transformers/examples/pytorch/language-modeling/'

训练模型

使用 16 的批次大小来训练 2B 模型。

gcloud alpha compute tpus tpu-vm ssh ${TPU_NAME} \
   --project=${PROJECT_ID} \
   --zone=${ZONE} \
   --worker=all \
   --command='
      export PJRT_DEVICE=TPU_C_API
      cd transformers/
      export LD_LIBRARY_PATH=/usr/local/lib/
      export PT_XLA_DEBUG=0
      export USE_TORCH=ON
      python3 examples/pytorch/xla_spawn.py \
         --num_cores=4 \
         examples/pytorch/language-modeling/run_clm.py \
         --num_train_epochs=3 \
         --dataset_name=wikitext \
         --dataset_config_name=wikitext-2-raw-v1 \
         --per_device_train_batch_size=16 \
         --per_device_eval_batch_size=16 \
         --do_train \
         --do_eval \
         --logging_dir=./tensorboard-metrics \
         --cache_dir=./cache_dir \
         --output_dir=/tmp/test-clm \
         --overwrite_output_dir \
         --cache_dir=/tmp \
         --config_name=examples/pytorch/language-modeling/my_config_2.json \
         --tokenizer_name=gpt2 \
         --block_size=1024 \
         --optim=adafactor \
         --adafactor=true \
         --save_strategy=no \
         --logging_strategy=no \
         --fsdp=full_shard \
         --fsdp_config=examples/pytorch/language-modeling/fsdp_config.json'

删除 TPU 和已加入队列的资源

在会话结束时删除 TPU 和排入队列的资源。

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

gcloud alpha compute tpus queued-resources delete ${QUEUED_RESOURCE_ID} \
   --project=${PROJECT_ID} \
   --zone=${ZONE} \
   --quiet

基准测试结果

训练脚本在 v5litepod-4、v5litepod-16 和 v5litepod-64 上运行，下表显示了不同加速器类型的基准吞吐量。

v5litepod-4 v5litepod-16 V5litepod-64
纪元 3 3 3
配置 6 亿 20 亿 160 亿
全局批量大小 64 128 256
吞吐量（示例/秒） 66 77 31

在 v5e 上训练 ViT

本教程介绍了如何在 PyTorch/XLA 上使用 cifar10 数据集上的 HuggingFace 代码库在 v5e 上运行 VIT。

设置

  1. 创建环境变量：

    export PROJECT_ID=your_project_ID
export ACCELERATOR_TYPE=v5litepod-16
export ZONE=us-west4-a
export RUNTIME_VERSION=v2-alpha-tpuv5-lite
export SERVICE_ACCOUNT=your_service_account
export TPU_NAME=tpu-name
export QUEUED_RESOURCE_ID=queued_resource_id
export QUOTA_TYPE=quota_type
export VALID_UNTIL_DURATION=1d

  2. 创建 TPU 资源：

    gcloud alpha compute tpus queued-resources create ${QUEUED_RESOURCE_ID} \
   --node-id=${TPU_NAME} \
   --project=${PROJECT_ID} \
   --zone=${ZONE} \
   --accelerator-type=${ACCELERATOR_TYPE} \
   --runtime-version=${RUNTIME_VERSION} \
   --valid-until-duration=${VALID_UNTIL_DURATION} \
   --service-account=${SERVICE_ACCOUNT} \
   --${QUOTA_TYPE}

    一旦 QueuedResource 处于 ACTIVE 状态，您就可以通过 SSH 连接到 TPU 虚拟机：

     gcloud alpha compute tpus queued-resources describe ${QUEUED_RESOURCE_ID} \
   --project=${PROJECT_ID} \
   --zone=${ZONE}

    当已加入队列的资源处于 ACTIVE 状态时，输出将类似于以下内容：

     state: ACTIVE

  3. 安装 Torch/xla 依赖项

    gcloud alpha compute tpus tpu-vm ssh ${TPU_NAME} \
   --project=${PROJECT_ID} \
   --zone=${ZONE} \
   --worker=all
   --command='
      sudo apt-get update -y
      sudo apt-get install libomp5 -y
      pip3 install mkl mkl-include
      pip3 install tf-nightly tb-nightly tbp-nightly
      pip3 install numpy
      sudo apt-get install libopenblas-dev -y
      pip3 install torch~=2.1.0 torchvision torch_xla[tpu]~=2.1.0 -f https://storage.googleapis.com/libtpu-releases/index.html
      pip3 install torch_xla[tpu] -f https://storage.googleapis.com/libtpu-releases/index.html'

  4. 下载 HuggingFace 代码库和安装要求。

       gcloud alpha compute tpus tpu-vm ssh ${TPU_NAME} \
   --project=${PROJECT_ID} \
   --zone=${ZONE} \
   --worker=all \
   --command="
      git clone https://github.com/suexu1025/transformers.git vittransformers; \
      cd vittransformers; \
      pip3 install .; \
      pip3 install datasets; \
      wget https://github.com/pytorch/xla/blob/master/scripts/capture_profile.py"

训练模型

gcloud alpha compute tpus tpu-vm ssh ${TPU_NAME} \
   --project=${PROJECT_ID} \
   --zone=${ZONE} \
   --worker=all \
   --command='
      export PJRT_DEVICE=TPU_C_API
      export PT_XLA_DEBUG=0
      export USE_TORCH=ON
      export TF_CPP_MIN_LOG_LEVEL=0
      export XLA_USE_BF16=1
      export LD_LIBRARY_PATH=/usr/local/lib:$LD_LIBRARY_PATH
      export TPU_LIBRARY_PATH=$HOME/.local/lib/python3.10/site-packages/libtpu/libtpu.so
      cd vittransformers
      python3 -u examples/pytorch/xla_spawn.py --num_cores 4 examples/pytorch/image-pretraining/run_mae.py --dataset_name=cifar10 \
      --remove_unused_columns=False \
      --label_names=pixel_values \
      --mask_ratio=0.75 \
      --norm_pix_loss=True \
      --do_train=true \
      --do_eval=true \
      --base_learning_rate=1.5e-4 \
      --lr_scheduler_type=cosine \
      --weight_decay=0.05 \
      --num_train_epochs=3 \
      --warmup_ratio=0.05 \
      --per_device_train_batch_size=8 \
      --per_device_eval_batch_size=8 \
      --logging_strategy=steps \
      --logging_steps=30 \
      --evaluation_strategy=epoch \
      --save_strategy=epoch \
      --load_best_model_at_end=True \
      --save_total_limit=3 \
      --seed=1337 \
      --output_dir=MAE \
      --overwrite_output_dir=true \
      --logging_dir=./tensorboard-metrics \
      --tpu_metrics_debug=true'

删除 TPU 和已加入队列的资源

在会话结束时删除 TPU 和排入队列的资源。

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

gcloud alpha compute tpus queued-resources delete ${QUEUED_RESOURCE_ID}
   --project=${PROJECT_ID}
   --zone=${ZONE}
   --quiet

基准测试结果

下表显示了不同加速器类型的基准吞吐量。

v5litepod-4 v5litepod-16 V5litepod-64
纪元 3 3 3
全局批量大小 32 128 512
吞吐量（示例/秒） 201 657 2844

TensorFlow 2.x

在单个主机 v5e 上训练 Resnet

本教程介绍如何使用虚构数据集在 v5litepod-4 或 v5litepod-8 上训练 ImageNet。如果要使用其他数据集，请参阅准备数据集。

设置

  1. 创建环境变量：

    export PROJECT_ID=your-project-ID
export ACCELERATOR_TYPE=v5litepod-4
export ZONE=us-east1-c
export RUNTIME_VERSION=tpu-vm-tf-2.15.0-pjrt
export TPU_NAME=your-tpu-name
export QUEUED_RESOURCE_ID=your-queued-resource-id
export QUOTA_TYPE=quota-type

    ACCELERATOR_TYPE 可以是 v5litepod-4 或 v5litepod-8。/

  2. 创建 TPU 资源：

    gcloud alpha compute tpus queued-resources create ${QUEUED_RESOURCE_ID} \
   --node-id=${TPU_NAME} \
   --project=${PROJECT_ID} \
   --zone=${ZONE} \
   --accelerator-type=${ACCELERATOR_TYPE} \
   --runtime-version=${RUNTIME_VERSION} \
   --${QUOTA_TYPE}

    待排队的资源处于 ACTIVE 状态后，您可以通过 SSH 连接到 TPU 虚拟机。如需检查已加入队列的资源的状态，请使用以下命令：

    gcloud alpha compute tpus queued-resources describe ${QUEUED_RESOURCE_ID} \
   --project=${PROJECT_ID} \
   --zone=${ZONE}

  3. 使用 SSH 连接到您的 TPU

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

  4. 设置一些环境变量

    export MODELS_REPO=/usr/share/tpu/models
export PYTHONPATH="${MODELS_REPO}:${PYTHONPATH}"
export MODEL_DIR=gcp-directory-to-store-model
export DATA_DIR=gs://cloud-tpu-test-datasets/fake_imagenet
export NEXT_PLUGGABLE_DEVICE_USE_C_API=true
export TF_PLUGGABLE_DEVICE_LIBRARY_PATH=/lib/libtpu.so

  5. 更改模型代码库目录和安装要求。

    cd ${MODELS_REPO} && git checkout r2.15.0
pip install -r official/requirements.txt

训练模型

运行训练脚本。

python3 official/vision/train.py \
   --tpu=local \
   --experiment=resnet_imagenet \
   --mode=train_and_eval \
   --config_file=official/vision/configs/experiments/image_classification/imagenet_resnet50_tpu.yaml \
   --model_dir=${MODEL_DIR} \
   --params_override="runtime.distribution_strategy=tpu,task.train_data.input_path=${DATA_DIR}/train*,task.validation_data.input_path=${DATA_DIR}/validation*,task.train_data.global_batch_size=2048,task.validation_data.global_batch_size=2048,trainer.train_steps=100"

删除 TPU 和已加入队列的资源

  1. 删除 TPU

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

  2. 删除已加入队列的资源请求

    gcloud alpha compute tpus queued-resources delete ${QUEUED_RESOURCE_ID} \
   --project=${PROJECT_ID} \
   --zone=${ZONE} \
   --quiet

在多主机 v5e 上训练 Resnet

本教程介绍如何使用虚构数据集在 v5litepod-16 或更高版本上训练 ImageNet。如果要使用其他数据集，请参阅准备数据集。

  1. 创建环境变量：

    export PROJECT_ID=your_project_ID
export ACCELERATOR_TYPE=v5litepod-16
export ZONE=us-east1-c
export RUNTIME_VERSION=tpu-vm-tf-2.15.0-pod-pjrt
export TPU_NAME=your_tpu_name
export QUEUED_RESOURCE_ID=your-queued-resource-id
export QUOTA_TYPE=quota-type

    ACCELERATOR_TYPE 可以大于 v5litepod-16。

  2. 创建 TPU 资源：

    gcloud alpha compute tpus queued-resources create ${QUEUED_RESOURCE_ID} \
   --node-id=${TPU_NAME} \
   --project=${PROJECT_ID} \
   --zone=${ZONE} \
   --accelerator-type=${ACCELERATOR_TYPE} \
   --runtime-version=${RUNTIME_VERSION} \
   --${QUOTA_TYPE}

    待排队的资源处于 ACTIVE 状态后，您可以通过 SSH 连接到 TPU 虚拟机。如需检查已加入队列的资源的状态，请使用以下命令：

    gcloud alpha compute tpus queued-resources describe ${QUEUED_RESOURCE_ID} \
   --project=${PROJECT_ID} \
   --zone=${ZONE}

  3. 使用 SSH 连接到您的 TPU（工作器零）

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

  4. 设置一些环境变量

    export MODELS_REPO=/usr/share/tpu/models
export PYTHONPATH="${MODELS_REPO}:${PYTHONPATH}"
export MODEL_DIR=gcp-directory-to-store-model
export DATA_DIR=gs://cloud-tpu-test-datasets/fake_imagenet
export TPU_LOAD_LIBRARY=0
export TPU_NAME=your_tpu_name

  5. 更改模型代码库目录和安装要求。

     cd $MODELS_REPO && git checkout r2.15.0
 pip install -r official/requirements.txt

训练模型

运行训练脚本。

python3 official/vision/train.py \
   --tpu=${TPU_NAME} \
   --experiment=resnet_imagenet \
   --mode=train_and_eval \
   --model_dir=${MODEL_DIR} \
   --params_override="runtime.distribution_strategy=tpu,task.train_data.input_path=${DATA_DIR}/train*, task.validation_data.input_path=${DATA_DIR}/validation*"

删除 TPU 和已加入队列的资源

  1. 删除 TPU

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

  2. 删除已加入队列的资源请求

    gcloud alpha compute tpus queued-resources delete ${QUEUED_RESOURCE_ID} \
   --project=${PROJECT_ID} \
   --zone=${ZONE} \
   --quiet