Cloud TPU v5e 训练

Cloud TPU v5e 是 Google Cloud 的最新一代 AI 加速器。v5e 经过优化，每个 Pod 的芯片占用空间更小（256 个），是用于转换器、文本到图像和卷积神经网络 (CNN) 训练、微调和服务的价值最高的产品。如需详细了解如何使用 Cloud TPU v5e 进行服务器端推理，请参阅使用 v5e 进行推理。

如需详细了解 Cloud TPU v5e TPU 硬件和配置，请参阅 TPU v5e。

开始使用

以下部分介绍了如何开始使用 TPU v5e。

请求配额

您需要配额才能使用 TPU v5e 进行训练。按需 TPU、预留 TPU 和 TPU 竞价型虚拟机有不同的配额类型。如果您将 TPU v5e 用于推理，则需要单独的配额。如需详细了解配额，请参阅配额。如需申请 TPU v5e 配额，请与 Cloud 销售团队联系。

创建 Google Cloud 账号和项目

您需要拥有 Google Cloud 账号和项目才能使用 Cloud TPU。如需了解详情，请参阅设置 Cloud TPU 环境。

创建 Cloud TPU

最佳实践是使用 queued-resource create 命令将 Cloud TPU v5 预配为队列化资源。如需了解详情，请参阅管理队列中的资源。

您还可以使用 Create Node API (gcloud compute tpus tpu-vm create) 预配 Cloud TPU v5。如需了解详情，请参阅管理 TPU 资源。

如需详细了解可用于训练的 v5e 配置，请参阅用于训练的 Cloud TPU v5e 类型。

框架设置

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

如需查看推理设置说明，请参阅 v5e 推理简介。

JAX 设置

如果您的 slice 形状大于 8 个条状标签，则一个 slice 中将包含多个虚拟机。在这种情况下，您需要使用 --worker=all 标志在单个步骤中在所有 TPU 虚拟机上运行安装，而无需使用 SSH 分别登录每个虚拟机：

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

命令标志说明

变量 说明
TPU_NAME TPU 的用户分配文本 ID，在分配队列中的资源请求时创建。
PROJECT_ID Google Cloud 项目名称。在 设置 Google Cloud 项目中使用现有项目或创建新项目
ZONE 如需了解支持的区域，请参阅 TPU 区域和可用区文档。
worker 有权访问底层 TPU 的 TPU 虚拟机。

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

gcloud 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() 显示给定 slice 中的芯片总数。jax.local_device_count() 表示此 slice 中单个虚拟机可访问的条状标签数量。

# 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 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 上开始使用 PJRT 和 PyTorch/XLA，并为所有工作器提供命令。

安装依赖项

gcloud 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 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 个步骤删除 slice 和已排队的资源：

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

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

这两步也可以用于移除处于 FAILED 状态的队列化资源请求。

JAX/FLAX 示例

以下部分通过示例介绍了如何在 TPU v5e 上训练 JAX 和 FLAX 模型。

在 v5e 上训练 ImageNet

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

设置

  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 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 compute tpus queued-resources describe ${QUEUED_RESOURCE_ID}  \
   --project=${PROJECT_ID} \
   --zone=${ZONE}

    当 QueuedResource 处于 ACTIVE 状态时，输出将类似于以下内容：

     state: ACTIVE

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

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

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

    gcloud 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 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 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 compute tpus tpu-vm delete ${TPU_NAME} \
   --project=${PROJECT_ID} \
   --zone=${ZONE} \
   --quiet

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

Hugging Face 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 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 compute tpus queued-resources describe ${QUEUED_RESOURCE_ID}  \
   --project=${PROJECT_ID} \
   --zone=${ZONE}

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

     state: ACTIVE

  3. 安装 JAX 及其库：

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

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

    gcloud 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.18.0 && pip install -r examples/flax/vision/requirements.txt'

  5. 下载 Imagenette 数据集：

    gcloud 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 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 compute tpus tpu-vm delete ${TPU_NAME} \
   --project=${PROJECT_ID} \
   --zone=${ZONE} \
   --quiet

gcloud 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

在《精灵宝可梦》中训练扩散

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

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

设置

  1. 为模型输出设置存储桶。 

    gcloud storage buckets create gs://your_bucket 
    
--project=your_project 
    
--location=us-west1 
    
export GCS_BUCKET_NAME=your_bucket

  2. 创建环境变量

    export GCS_BUCKET_NAME=your_bucket
export PROJECT_ID=your_project_ID
export ACCELERATOR_TYPE=v5litepod-16
export ZONE=us-west1-c
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

    命令标志说明

    变量 说明
    GCS_BUCKET_NAME 显示在 Google Cloud 控制台 -> Cloud Storage -> 存储分区中
    PROJECT_ID Google Cloud 项目名称。使用现有项目或前往 设置 Google Cloud 项目创建新项目
    ACCELERATOR_TYPE 请参阅 TPU 版本页面，了解您的 TPU 版本。
    区域 如需了解支持的区域，请参阅 TPU 区域和可用区文档。
    RUNTIME_VERSION 为 RUNTIME_VERSION 使用 v2-alpha-tpuv5。
    SERVICE_ACCOUNT 这是您的服务账号的地址，您可以在 Google Cloud 控制台 -> IAM -> 服务账号中找到。例如：tpu-service-account@myprojectID。iam.gserviceaccount.com
    TPU_NAME TPU 的用户分配文本 ID，在分配队列中的资源请求时创建。
    QUEUED_RESOURCE_ID 已加入队列的资源请求的用户分配的文本 ID。如需了解已排队的资源，请参阅已排队的资源文档。
    QUOTA_TYPE 可以是 reserved 或 spot。如果未指定这两者，则 QUOTA_TYPE 默认为 on-demand。如需了解 Cloud TPU 支持的不同类型的配额，请参阅配额。
    VALID_UNTIL_DURATION 请求有效的时长。如需了解不同的有效时长，请参阅 队列化资源。

  3. 创建 TPU 资源：

    gcloud 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 compute tpus queued-resources describe ${QUEUED_RESOURCE_ID}  \
   --project=${PROJECT_ID} \
   --zone=${ZONE}

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

     state: ACTIVE

  4. 安装 JAX 及其库。

    gcloud 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'

  5. 下载 HuggingFace repository 和安装要求。

    gcloud 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.18.0 clu && pip install -U -r examples/text_to_image/requirements_flax.txt'

训练模型

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

   gcloud compute tpus tpu-vm ssh ${TPU_NAME} --zone=${ZONE} --project ${PROJECT_ID} --worker=all --command="
   git clone https://github.com/google/maxdiffusion
   cd maxdiffusion
   git reset --hard 57629bcf4fa32fe5a57096b60b09f41f2fa5c35d # This identifies the GitHub commit to use.
   pip3 install jax[tpu] -f https://storage.googleapis.com/jax-releases/libtpu_releases.html
   pip3 install -r requirements.txt
   pip3 install .
   export LIBTPU_INIT_ARGS=""
   python -m src.maxdiffusion.models.train src/maxdiffusion/configs/base_2_base.yml run_name=your_run base_output_directory=gs://${GCS_BUCKET_NAME}/ enable_profiler=False"

删除 TPU 和已排队的资源

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

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

gcloud 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 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 compute tpus queued-resources describe ${QUEUED_RESOURCE_ID} \
   --project=${PROJECT_ID} \
   --zone=${ZONE}

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

     state: ACTIVE

  3. 安装 JAX 及其库。

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

  4. 下载 HuggingFace repository 和安装要求。

    gcloud 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 compute tpus tpu-vm ssh ${TPU_NAME} \
   --project=${PROJECT_ID} \
   --zone=${ZONE} \
   --worker=all \
   --command='cd transformers/examples/flax/language-modeling && gcloud storage cp gs://cloud-tpu-tpuvm-artifacts/v5litepod-preview/jax/gpt . --recursive'

训练模型

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

gcloud 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 compute tpus tpu-vm delete ${TPU_NAME} \
   --project=${PROJECT_ID} \
   --zone=${ZONE} \
   --quiet

gcloud 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

以下部分介绍了如何在 TPU v5e 上训练 PyTorch/XLA 模型的示例。

使用 PJRT 运行时训练 ResNet

PyTorch/XLA 将从 PyTorch 2.0 及更高版本从 XRT 迁移到 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 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 compute tpus queued-resources describe ${QUEUED_RESOURCE_ID} \
   --project=${PROJECT_ID} \
   --zone=${ZONE}

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

     state: ACTIVE

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

    gcloud 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 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 compute tpus tpu-vm delete ${TPU_NAME} \
   --project=${PROJECT_ID} \
   --zone=${ZONE} \
   --quiet

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

下表显示了基准吞吐量。

加速器类型 吞吐量（示例/秒）
v5litepod-4 4240 亿次/秒
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 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 compute tpus queued-resources describe ${QUEUED_RESOURCE_ID} \
   --project=${PROJECT_ID} \
   --zone=${ZONE}

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

    state: ACTIVE

  3. 安装 PyTorch/XLA 依赖项。

    gcloud 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 repository 和安装要求。

    gcloud 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 compute tpus tpu-vm ssh ${TPU_NAME} \
   --project=${PROJECT_ID} \
   --zone=${ZONE} \
   --worker=all \
   --command='
      gcloud storage cp gs://cloud-tpu-tpuvm-artifacts/config/xl-ml-test/pytorch/gpt2/my_config_2.json transformers/examples/pytorch/language-modeling/ --recursive
      gcloud storage cp gs://cloud-tpu-tpuvm-artifacts/config/xl-ml-test/pytorch/gpt2/fsdp_config.json transformers/examples/pytorch/language-modeling/'

训练模型

使用批次大小为 16 的 2B 模型进行训练。

gcloud 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 compute tpus tpu-vm delete ${TPU_NAME} \
   --project=${PROJECT_ID} \
   --zone=${ZONE} \
   --quiet

gcloud 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
config 6 亿 2B 160 亿
全局批量大小 64 128 256
吞吐量（示例/秒） 66 77 31

在 v5e 上训练 ViT

本教程将介绍如何在 PyTorch/XLA 上使用 HuggingFace 代码库在 CIFAR10 数据集上运行 VIT 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 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 compute tpus queued-resources describe ${QUEUED_RESOURCE_ID} \
   --project=${PROJECT_ID} \
   --zone=${ZONE}

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

     state: ACTIVE

  3. 安装 PyTorch/XLA 依赖项

    gcloud 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 repository 和安装要求。

       gcloud 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 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 compute tpus tpu-vm delete ${TPU_NAME}
   --project=${PROJECT_ID}
   --zone=${ZONE}
   --quiet

gcloud 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 2,844

TensorFlow 2.x

以下部分举例介绍了如何在 TPU v5e 上训练 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 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 compute tpus queued-resources describe ${QUEUED_RESOURCE_ID} \
   --project=${PROJECT_ID} \
   --zone=${ZONE}

  3. 使用 SSH 连接到 TPU

    gcloud 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 compute tpus tpu-vm delete ${TPU_NAME} \
   --project=${PROJECT_ID} \
   --zone=${ZONE} \
   --quiet

  2. 删除已排队的资源请求

    gcloud 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 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 compute tpus queued-resources describe ${QUEUED_RESOURCE_ID} \
   --project=${PROJECT_ID} \
   --zone=${ZONE}

  3. 使用 SSH 连接到 TPU（工作器 0）

    gcloud 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 compute tpus tpu-vm delete ${TPU_NAME} \
   --project=${PROJECT_ID} \
   --zone=${ZONE} \
   --quiet

  2. 删除已排队的资源请求

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