Cloud TPU v5p 训练

Cloud TPU v5p 是 Google Cloud 的第五代 Cloud TPU， v4 TPU 的继任者。v5p 针对大规模训练进行了优化 并成为发展基础模型的 LLM、扩散模型和生成式 AI。概括地说，v5p 的性能最高可达到 v4 的 2 倍，同时还可将 2 倍多的 TPU 打包到一个 Pod 中（最大 slice 为 6,000，而 v4 为 3,000），从而在 Pod 级别将性能提高到最高 4 倍。它还采用更高的时钟频率（1.75Ghz 对比 1.05Ghz），添加了 SparseCore 以实现大规模嵌入，并将高带宽内存 (HBM) 容量提高了三倍。

Cloud TPU v5p 概念

如果您是 Cloud TPU 新手，请参阅 TPU 文档首页。

Cloud TPU 概念（例如切片、主机、 和 TensorCore）及 Cloud TPU 系统架构。 请参阅 Cloud TPU 系统架构 页面。

每个 Cloud TPU 版本都需要特定的加速器类型才能进行训练或推理。v5p 配置中介绍了这些加速器类型。

管理 TPU 资源

如需了解用于管理 TPU 虚拟机的所有命令，请参阅管理 TPU或有序资源用户指南（用于管理有序资源）。

框架设置

本部分介绍模型的常规设置过程。 使用 JAX 或 PyTorch 与 TPU v5p 进行训练。

JAX 设置

如果您的 slice 形状大于 4 个条状标签，则一个 slice 中将包含多个虚拟机。在这种情况下，您需要使用 --worker=all 标志。 使用单个命令在所有 TPU 虚拟机上运行安装：

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'

您可以运行以下命令来检查设备数量（此处显示的输出是使用 v5p-32 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.]

使用 --node=all 在所有多切片工作器上运行该命令。

gcloud compute tpus queued-resources ssh ${QUEUED_RESOURCE_ID} \
--project ${PROJECT_ID} --zone ${ZONE} --node=all --worker=all \
--command='python3 -c "import jax; print(jax.device_count()); print(jax.local_device_count())"'

尝试学习本文档中的 JAX 教程，开始使用 JAX 进行 v5p 训练。

设置 PyTorch

PJRT 运行时 是 v5p 唯一支持的运行时，而 PyTorch 2.1+ 使用 PJRT 作为所有 TPU 版本的默认运行时。本部分介绍如何开始在 v5p Pod 上使用 PJRT，并为所有工作器安装 PyTorch/XLA 2.2.0。

安装依赖项

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

使用带有 PJRT 的 Python 脚本验证安装 显示可用的 TPU 设备（此处显示的输出 。

gcloud compute tpus tpu-vm ssh ${TPU_NAME} \
--project ${PROJECT_ID} --zone ${ZONE} --worker=all \
--command='
PJRT_DEVICE=TPU python3 -c "import torch_xla.core.xla_model as xm; print(xm.get_xla_supported_devices(\"TPU\"))"
'
 
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...
['xla:0', 'xla:1', 'xla:2', 'xla:3']
['xla:0', 'xla:1', 'xla:2', 'xla:3']
['xla:0', 'xla:1', 'xla:2', 'xla:3']
['xla:0', 'xla:1', 'xla:2', 'xla:3']

使用 --node=all 在所有多切片工作器上运行该命令。

gcloud compute tpus queued-resources ssh ${QUEUED_RESOURCE_ID} \
--project ${PROJECT_ID} --zone ${ZONE} --node=all --worker=all \
--command='
PJRT_DEVICE=TPU python3 -c "import torch_xla.core.xla_model as xm; print(xm.get_xla_supported_devices(\"TPU\"))"
'

尝试学习本文档中的 PyTorch 教程，开始使用 PyTorch 进行 v5p 训练。

监控和分析

Cloud TPU v5p 支持使用与上一代 Cloud TPU 相同的方法进行监控和性能分析。您可以参阅使用 Cloud TPU 工具剖析模型的性能，详细了解性能分析；也可以参阅监控 Cloud TPU 虚拟机，详细了解监控。

培训教程

本部分重点介绍单个 Slice 的训练教程。如需将这些教程调整为适用于多 Slice 训练，只需向 SSH 命令添加 --node=all 标志即可。有关详情和最佳实践，请参阅 多切片简介。

JAX 教程

模块序列扩散 2.1

本教程将向您展示如何使用 使用 Pokémon 的 HuggingFace 数据集。

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

设置

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

    gcloud storage buckets create gs://your_bucket \
 --project=your_project \
 --location=us-east5-a

  2. 创建环境变量：

    export GCS_BUCKET_NAME=your-bucket
export PROJECT_ID=your-project-ID
export ACCELERATOR_TYPE=v5p-32
export ZONE=us-east5-a
export RUNTIME_VERSION=v2-alpha-tpuv5
export SERVICE_ACCOUNT=your-service-account
export TPU_NAME=your-tpu-name
export QUEUED_RESOURCE_ID=your-qr-name
export QUOTA_TYPE=spot
export VALID_UNTIL_DURATION=1d

    命令标志说明

    变量 说明
    PROJECT_ID Google Cloud 项目名称
    ACCELERATOR_TYPE 请参阅 TPU 版本页面，了解您的 TPU 版本。
    ZONE 如需了解受支持的可用区，请参阅 TPU 区域和可用区文档。
    RUNTIME_VERSION 对于 v5p，请为 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}

    排队的资源准备就绪后，您就可以通过 SSH 连接到 TPU 虚拟机 处于 ACTIVE 状态。 运行以下命令检查队列中资源的状态：

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

    当已加入队列的资源处于 ACTIVE 状态时， 类似如下：

    state: ACTIVE

  4. 训练模型

    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 # Install the latest version of JAX
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=my_run base_output_directory=gs://$GCS_BUCKET_NAME enable_profiler=False"

清理

在会话结束时删除 TPU 和排队的资源请求，或 移除“FAILED”中已加入队列的资源请求状态。 如需删除已排队的资源，请按以下 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

或者，使用 --force 删除 Slice 和已加入队列的资源请求 即可：

# With --force
gcloud compute tpus queued-resources delete ${QUEUED_RESOURCE_ID}
--project ${PROJECT_ID} --zone ${ZONE} --quiet --force

基准测试结果

Stable Diffusion 训练脚本在 v5p-8、v5p-32 和 v5p-128 上运行。 下表显示了吞吐量。

v5p-8

v5p-32

v5p-128

训练步骤

150

150

150

全局批量大小

32

64

64

吞吐量（示例/秒）

12.10

8 月 18 日

19 月 10 日

MaxText

本教程介绍如何在 Cloud TPU 上使用合成数据集训练 MaxText 模型。

MaxText 是一个高性能、任意可扩缩、开源且经过充分测试的 LLM，以纯 Python/JAX 编写，以 Cloud TPU 为目标平台。MaxText 是一款简单易用且可自定义的工具，可帮助研究人员和开发者推进自然语言处理 (NLP) 研究和开发的前沿。

在运行本教程之前，您需要设置 Cloud TPU 环境。

  1. 设置环境变量

    export PROJECT_ID=your_project_ID
export TPU_NAME=your_tpu_name # user defined TPU name
export ACCELERATOR_TYPE=v5p-256
export ZONE=us-east5-a
export RUNTIME_VERSION=v2-alpha-tpuv5
export RUN_NAME=your_experiment_run_name # user defined name for this run
export GCS_BUCKET_NAME=your_bucket_name # Output cloud folder. Should start with gs://
export MAXTEXT_OUTPUT_PATH=${GCS_BUCKET_NAME}/your_experiment_output_path
export NUM_SLICES=1 # Update the value to a number >1 for Multislice.

    命令标志说明

    变量 说明
    PROJECT_ID Google Cloud 项目名称
    TPU_NAME 为 TPU 指定的用户定义的名称。
    ACCELERATOR_TYPE 请参阅 TPU 版本页面，了解您的 TPU 版本。
    可用区 如需了解受支持的可用区，请参阅 TPU 区域和可用区文档。
    RUNTIME_VERSION 对于 v5p，请为运行时版本使用 v2-alpha-tpuv5。
    RUN_NAME 用户提供的实验运行名称。

    建议针对多切片进行可选设置：

    export NETWORK_NAME=your_network_name
export FIREWALL_RULE_NAME=your_firewall_rule_name

    如果您正在运行多 slice 工作负载并希望获得最佳网络性能，不妨考虑创建一个最大传输单元 (MTU) 为 8896 字节的专用网络，并配置适当的防火墙规则。虽然此步骤并非必需，但可以显著提升性能，尤其是在通过数据中心网络 (DCN) 扩容 slice 数量时。正在创建记事 网络需要项目中的 compute.networks.create 权限。 以下示例展示了如何创建专用网络和防火墙规则。

    创建专用网络：

    gcloud compute networks create ${NETWORK_NAME} \
--mtu=8896 \
--project=${PROJECT_ID} \
--subnet-mode=auto \
--bgp-routing-mode=regional

    创建防火墙规则：

    gcloud compute firewall-rules create ${FIREWALL_RULE_NAME} \
--network ${NETWORK_NAME} --allow tcp,icmp,udp --project=${PROJECT_ID}

  2. 克隆 MaxText 代码库

    git clone https://github.com/google/maxtext.git

  3. 训练模型

    以下部分介绍了训练 MaxText 的两个选项。

    选项 1

    如果您希望使用脚本来管理整个工作流（包括配置、 Cloud TPU 以及安装依赖项以运行模型 以及拆解资源，您可以使用 multihost_job.py。

    cd maxtext && python3 multihost_job.py --PROJECT=${PROJECT_ID} --ZONE=${ZONE} \
--NUM_SLICES=${NUM_SLICES} --TPU_TYPE=${ACCELERATOR_TYPE} \
--VERSION=${RUNTIME_VERSION} --RUN_NAME=${RUN_NAME} #user defined run name \
--BUCKET_NAME=${GCS_BUCKET_NAME} \ #used to store logs and configs
--COMMAND="bash setup.sh && bash MaxText/configs/experimental/64b.sh RUN_NAME=${RUN_NAME} OUTPUT_PATH=${MAXTEXT_OUTPUT_PATH} PLATFORM=gce"

    启动脚本后，您应该会在日志中看到类似于以下内容的消息。输出消息中会引用日志位置。TPU 预配完成后，点击第一个链接即可访问所有工作器的日志。

    ------------------------------------

multihost_job finished running, TPUs are starting up to run your job remotely.

Logs for your job are displayed here:
https://console.cloud.google.com/logs/query;query=resource.type%3D%22gce_instance%22%20AND%0Alog_id%2528%22_log%22%2529;?project=PROJECT_ID

To see the output of a single host, you may edit the slice and worker
number in the `log_file_path` property here:

https://console.cloud.google.com/logs/query;query=resource.type%3D%22gce_instance%22%20AND%0Alog_id%2528%22RUN_NAME_log%22%2529%20AND%0Alabels.%22agent.googleapis.com%2Flog_file_path%22%3D%20%22%2FRUN_NAME%2Fmain_command_log_slice_0_worker_0%22;?project=PROJECT_ID

When your job is finished, the main command log is in your Cloud Storage
bucket:
https://console.cloud.google.com/storage/browser/YOUR_BUCKET_NAME/RUN_NAME?project=PROJECT_ID

View the status of the created TPUs using:
gcloud compute tpus queued-resources list --filter=RUN_NAME
--zone=ZONE --project=PROJECT_ID
选项 2

如需在预配的 Cloud TPU 上多次运行训练脚本，请使用 multihost_runner.py 脚本使用该资源。

  1. 设置变量以创建 TPU。

    export SERVICE_ACCOUNT=your_service_account
export TPU_NAME=your_tpu_name
export QUEUED_RESOURCE_ID=your_queued_resource_id
export VALID_DURATION=1d
export QUOTA_TYPE=quota_type
     
    --node-count ${NODE_COUNT} \
--node-prefix ${NODE_PREFIX} # optional, the default is QUEUED_RESOURCE_ID

  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_DURATION} \
--service-account ${SERVICE_ACCOUNT} \
--${QUOTA_TYPE}

    当 QueuedResource 处于 ACTIVE 状态后，您就可以使用 SSH 连接到 TPU 虚拟机：

    使用 describe 命令查询队列中资源的状态。

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

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

     state: ACTIVE

  3. 使用 SSH 连接到 TPU

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

  4. 安装依赖项

    export TPU_NAME=your_tpu_name
export MAXTEXT_OUTPUT_PATH=output-path
     
    cd maxtext && python3 multihost_runner.py --TPU_PREFIX=${TPU_NAME} \
--COMMAND='bash setup.sh'

  5. 使用各种不同的 配置脚本，例如 32b.sh、64b.sh。如果您是从 TPU 虚拟机运行脚本， 您需要添加 --INTERNAL_IP=true 标志。

    python3 multihost_runner.py --TPU_PREFIX=${TPU_NAME} \
--COMMAND="bash MaxText/configs/experimental/64b.sh RUN_NAME=${RUN_NAME}
OUTPUT_PATH=${MAXTEXT_OUTPUT_PATH} PLATFORM=gce"

清理

删除您的 TPU 和已排队的资源。

基准测试结果

MaxText 训练脚本的运行速度从 320 亿到 11600 亿次，且精确率为 bf16。 运行结果如下表所示。

参数数量

加速器类型

TFLOP/芯片/秒

模型 FLOP 利用率

（MFU）

32B

v5p-128

3.28E+02

71.47%

64B

v5p-128

3.23E+02

70.31%

1280 亿

v5p-256

3.15E+02

68.68%

128B

v5p-512

3.15E+02

68.53%

2560 亿

v5p-1024

3.16E+02

68.82%

5,120 亿

v5p-1024

2.94E+02

63.99%

1024B

v5p-2048

2.49E+02

64.05%

1,0240 亿

5p-4096

2.97E+02

64.80%

1,1600 亿

5p-7680

2.95E+02

64.27%

1160B

v5p-12288

3.04E+02

66.23%

256B 参数模型已在 v5p-512 和 v5p-1024 上进行了测试 使用 bf16 和 int8 权重。下表显示了这些测试的结果。

v5p-512

v5p-512

v5p-1024

v5p-1024

全局批量大小

(tokens)

5.24 欧洲 + 05

5.24E+05

1.05E+06

1.05E+06

精确率

bf16

int8

bf16

int8

TFLOP/芯片/秒

307

408

308

414

模型 FLOP 利用率

（MFU）

66.98%

88.85%

67.09%

90.23%

TensorFlow 教程

在单个主机 v5p 上训练 ResNet

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

设置

  1. 创建环境变量：

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

    在本教程中，请使用 v5p-8 作为 ACCELERATOR_TYPE。

  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

训练模型

  1. 运行训练脚本。

    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 和已排队的资源。

在多主机 v5p 上训练 ResNet

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

  1. 创建环境变量：

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

    ACCELERATOR_TYPE 可以是 v5p-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 虚拟机了。

    使用 describe 命令查询队列中资源的状态：

    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 TPU_NAME=your_tpu_name
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

  5. 切换到模型代码库目录并安装要求。

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

训练模型

  1. 运行训练脚本。

    python3 official/vision/train.py \
  --tpu=${TPU_NAME} \
  --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 和排队的资源。

PyTorch/XLA

Llama 2

本教程将介绍如何使用分支在 v5p 上训练 Llama 2 7B 模型 支持通用和可扩缩的 PyTorch/XLA 上的 HuggingFace 代码库 机器学习计算图并行处理 (GSPMD)。

设置

  1. 为项目 ID、加速器类型、可用区、运行时版本和 TPU 名称创建变量。

    export PROJECT_ID=your_project_ID
export ACCELERATOR_TYPE=v5p-8
export ZONE=us-east5-a
export RUNTIME_VERSION=v2-alpha-tpuv5
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_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_DURATION} \
--service-account ${SERVICE_ACCOUNT} \
--${QUOTA_TYPE}

    完成连接后，您可以使用 SSH 连接到 QueuedResource 处于 ACTIVE 状态：

    使用 describe 命令查询已加入队列的资源的状态。

    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
sudo apt-get install libopenblas-dev -y
pip3 install numpy
pip3 install typing-extensions
pip install torch~=2.2.0 torch_xla[tpu]~=2.2.0 -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 -b llama2-google-next-training https://github.com/pytorch-tpu/transformers.git
cd transformers
pip3 install git+file://$PWD
pip3 install datasets accelerate evaluate scikit-learn'

  5. 下载 7B 模型配置。

    gcloud compute tpus tpu-vm ssh ${TPU_NAME} \
--project=${PROJECT_ID} \
--zone=${ZONE} \
--worker=all \
--command="curl https://huggingface.co/TheBloke/Llama-2-7B-fp16/raw/main/config.json --output ~/config.json"

  6. 训练模型

    gcloud compute tpus tpu-vm ssh ${TPU_NAME} \
--project=${PROJECT_ID} \
--zone=${ZONE} \
--worker=all \
--command='
export PJRT_DEVICE=TPU
export XLA_USE_BF16=1
export XLA_IR_DEBUG=1
export XLA_HLO_DEBUG=1

export LIBTPU_INIT_ARGS="--xla_enable_async_collective_permute=true
--xla_tpu_enable_async_collective_fusion_multiple_steps=true
--xla_tpu_enable_async_collective_fusion=true
--xla_tpu_overlap_compute_collective_tc=true
--xla_enable_async_all_gather=true
--xla_jf_spmd_threshold_for_windowed_einsum_mib=0"

export PROFILE_EPOCH=0
export PROFILE_STEP=3
export PROFILE_DURATION_MS=20000
export PROFILE_LOGDIR=/tmp/home/

cd transformers
python examples/pytorch/language-modeling/run_clm.py \
 --tokenizer_name hf-internal-testing/llama-tokenizer \
 --dataset_name wikitext \
 --dataset_config_name wikitext-2-raw-v1 \
 --per_device_train_batch_size 96 \
 --per_device_eval_batch_size 8 \
 --num_train_epochs 1 \
 --do_train \
 --output_dir /tmp/output \
 --overwrite_output_dir \
 --config_name ~/config.json \
 --save_strategy no \
 --logging_strategy no \
 --remove_unused_columns no \
 --optim adafactor \
 --torch_dtype bfloat16 \
 --dataloader_drop_last yes \
 --block_size 2048 \
 --spmd_2d_sharding 1 \
 --spmd_grad_chkpt
'

如果您是在多 slice 环境中运行，则需要将标志 --spmd_dcn_parallelism 设置为 slice 数量。

SPMD_USER_GUIDE 提供了更深入的用户指南，其中介绍了 HF 脚本的所有不同环境变量和切换开关。请注意， LIBTPU_INIT_ARGS 将合并到 PyTorch/XLA 中，并且 默认处于启用状态

清理

删除您的 TPU 和已排队的资源。

基准测试结果

下表列出了三种 Llama 2 模型大小的吞吐量。

v5p-8

v5p-128

v5p-128

模型大小

70 亿

130 亿

700 亿

全局批量大小

96

1024

128

分片网格形状

(4, 1)

（64、1）

（16、4）

模型 flops 利用率

（MFU）

56.67%

55.80%

51.85%

支持与反馈

我们欢迎各种反馈！如需分享反馈或请求支持， 填写 Cloud TPU 支持或反馈表单。