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 Spot 虚拟机。有单独的 如果您要将 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。TensorFlow 支持在 tpu-vm-tf-2.17.0-pjrt 和 tpu-vm-tf-2.17.0-pod-pjrt TPU 运行时版本。

如需了解推理设置的说明，请参阅 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 个步骤删除切片，然后删除加入队列的资源：

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 上的 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 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}

    加入队列资源后，您就可以通过 SSH 连接到 TPU 虚拟机 ACTIVE 状态：

    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 模型是第一个成功训练 Transformer 编码器的模型 在 ImageNet 上使用模型训练时，结果与卷积网络相比表现出色。有关 详情，请参阅以下资源：

设置

  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}

    加入队列后，您就可以通过 SSH 连接到 TPU 虚拟机 处于 ACTIVE 状态：

    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.17.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

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

本教程介绍如何使用 使用 Pokémon 的 HuggingFace 数据集。

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 项目名称。使用现有项目或创建一个新项目 <ph type="x-smartling-placeholder"></ph> 设置您的 Google Cloud 项目
    ACCELERATOR_TYPE 请参阅 TPU 版本页面，了解您的 TPU 版本。
    ZONE 如需了解支持的区域，请参阅 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.17.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}

    加入队列后，您就可以通过 SSH 连接到 TPU 虚拟机 处于 ACTIVE 状态：

    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 代码库 和安装要求

    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

以下部分举例说明了如何在 Google Cloud 上训练 PyTorch/XLA 模型 TPU v5e。

使用 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 后，您可以通过 SSH 连接到 TPU 虚拟机 处于 ACTIVE 状态：

    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/秒

在 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 后，您可以通过 SSH 连接到 TPU 虚拟机 ACTIVE 状态：

    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 代码库 和安装要求

    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,000

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}

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

    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 (工作器零)

    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