使用 TPU v6e 训练模型

本文档将指导您在 Cloud TPU v6e（也称为 Trillium）上训练模型，其中涵盖环境设置、性能优化以及使用 JAX 和 PyTorch/XLA 的实际训练示例。

TPU v6e（也称为 Trillium）是 Google 的第 6 代 TPU。在所有技术界面（例如 API 和日志）以及本文档中，Trillium 都将称为 v6e。TPU v6e 架构的每个 Pod 有 256 个芯片，与 v5e 有很多相似之处。TPU v6e 针对 Transformer、文本转图片和卷积神经网络 (CNN) 训练、微调和部署进行了优化。如需详细了解 TPU v6e 系统架构和配置，请参阅 TPU v6e。

如需了解如何在 Cloud TPU v6e 上运行推理，请参阅以下教程：

准备工作

在开始之前，您需要：

  • 创建 Google Cloud 启用了结算功能的账号和项目
  • 安装 Google Cloud CLI Alpha 版组件
  • 启用 Cloud TPU API
  • 创建 Cloud TPU 服务代理
  • 创建 Cloud TPU 服务账号并授予权限

如需了解详情，请参阅设置 Cloud TPU 环境。

验证配额和权限

验证您的项目是否具有以下配额：

如果您将 GKE 与 XPK 搭配使用，则需要在 Google Cloud 控制台中获得其他权限。如需了解详情，请参阅Google Cloud 控制台所需的权限 。

预配 TPU

您可以使用以下方法配置和管理 TPU v6e：

  • GKE：您可以使用 GKE 将 TPU 作为加速器池进行预配和管理，以用于容器化的机器学习工作负载。如需了解详情，请参阅 GKE 中的 TPU 简介。
  • GKE 和 XPK：XPK 是一种命令行工具，可简化 GKE 上的集群创建和工作负载执行。它专为机器学习从业者设计，可用于预配 TPU 和运行训练作业，而无需具备深厚的 Kubernetes 专业知识。如需了解详情，请参阅 XPK GitHub 代码库。
  • Cloud TPU 已排队的资源：借助已排队的资源，您可以请求在容量可用时预配的 TPU 容量。非常适合可在队列中等待的批处理作业和容错工作负载。您可以为请求指定时间窗口。如需了解详情，请参阅管理已排队的资源。

使用 GKE 和 XPK 预配 v6e Cloud TPU

如果您将 GKE 命令与 v6e 搭配使用，则可以使用 Kubernetes 命令或 XPK 来预配 Cloud TPU 并训练或部署模型。如需了解如何规划 GKE 集群中的 Cloud TPU 配置，请参阅规划 GKE 中的 Cloud TPU。以下部分提供了用于创建支持单个 NIC 和支持多 NIC 的 XPK 集群的命令。

创建支持单个 NIC 的 XPK 集群

export CLUSTER_NAME=xpk-cluster-name
export ZONE=us-east1-d
export PROJECT_ID=your-project-id
export TPU_TYPE=v6e-256
export NUM_SLICES=2

export NETWORK_NAME=${CLUSTER_NAME}-mtu9k
export NETWORK_FW_NAME=${NETWORK_NAME}-fw
 
gcloud compute networks create ${NETWORK_NAME} \
   --mtu=8896 \
   --project=${PROJECT_ID} \
   --subnet-mode=auto \
   --bgp-routing-mode=regional
 
gcloud compute firewall-rules create ${NETWORK_FW_NAME} \
   --network=${NETWORK_NAME} \
   --allow tcp,icmp,udp \
   --project=${PROJECT_ID}
 
export CLUSTER_ARGUMENTS="--network=${NETWORK_NAME} --subnetwork=${NETWORK_NAME}"

python3 xpk.py cluster create --cluster=${CLUSTER_NAME} \
   --cluster-cpu-machine-type=e2-standard-8 \
   --num-slices=${NUM_SLICES} \
   --tpu-type=${TPU_TYPE} \
   --zone=${ZONE} \
   --project=${PROJECT_ID} \
   --on-demand \
   --custom-cluster-arguments="${CLUSTER_ARGUMENTS}" \
   --create-vertex-tensorboard

命令标志说明

变量 说明
CLUSTER_NAME 用户分配的 XPK 集群名称。
PROJECT_ID Google Cloud 项目名称。使用现有项目或创建新项目。如需了解详情，请参阅设置您的 Google Cloud 项目。
ZONE 如需了解支持的可用区，请参阅 Cloud TPU 区域和可用区文档。
TPU_TYPE 请参阅加速器类型。
NUM_SLICES 您要创建的切片数
CLUSTER_ARGUMENTS 要使用的网络和子网。

例如：--network=${NETWORK_NAME} --subnetwork=${NETWORK_NAME}

NUM_SLICES 要创建的切片数。
NETWORK_NAME 要使用的次要网络的名称。
NETWORK_FW_NAME 要使用的次要网络防火墙的名称。

创建支持多 NIC 的 XPK 集群

export CLUSTER_NAME=xpk-cluster-name
export REGION=your-region
export ZONE=us-east1-d
export PROJECT_ID=your-project-id
export TPU_TYPE=v6e-256
export NUM_SLICES=2

export NETWORK_NAME_1=${CLUSTER_NAME}-mtu9k-1-${ZONE}
export SUBNET_NAME_1=${CLUSTER_NAME}-privatesubnet-1-${ZONE}
export NETWORK_FW_NAME_1=${NETWORK_NAME_1}-fw-1-${ZONE}
export FIREWALL_RULE_NAME=${CLUSTER_NAME}-privatefirewall-1-${ZONE}
export ROUTER_NAME=${CLUSTER_NAME}-network-1-${ZONE}
export NAT_CONFIG=${CLUSTER_NAME}-natconfig-1-${ZONE}
 
gcloud compute networks create ${NETWORK_NAME_1} \
   --mtu=8896 \
   --bgp-routing-mode=regional \
   --subnet-mode=custom \
   --project=${PROJECT_ID}
 
gcloud compute networks subnets create ${SUBNET_NAME_1} \
   --network=${NETWORK_NAME_1} \
   --range=10.11.0.0/18 \
   --region=${REGION} \
   --project=${PROJECT_ID}
 
gcloud compute firewall-rules create ${FIREWALL_RULE_NAME} \
   --network=${NETWORK_NAME_1} \
   --allow tcp,icmp,udp \
   --project=${PROJECT_ID}
 
gcloud compute routers create ${ROUTER_NAME} \
   --project=${PROJECT_ID} \
   --network=${NETWORK_NAME_1} \
   --region=${REGION}
 
gcloud compute routers nats create ${NAT_CONFIG} \
   --router=${ROUTER_NAME} \
   --region=${REGION} \
   --auto-allocate-nat-external-ips \
   --nat-all-subnet-ip-ranges \
   --project=${PROJECT_ID} \
   --enable-logging
 
# Secondary subnet for multi-nic experience.
# Need custom IP routing to be different from the first network's subnet.

export NETWORK_NAME_2=${CLUSTER_NAME}-privatenetwork-2-${ZONE}
export SUBNET_NAME_2=${CLUSTER_NAME}-privatesubnet-2-${ZONE}
export FIREWALL_RULE_NAME=${CLUSTER_NAME}-privatefirewall-2-${ZONE}
export ROUTER_NAME=${CLUSTER_NAME}-network-2-${ZONE}
export NAT_CONFIG=${CLUSTER_NAME}-natconfig-2-${ZONE}

gcloud compute networks create ${NETWORK_NAME_2} \
   --mtu=8896 \
   --bgp-routing-mode=regional \
   --subnet-mode=custom \
   --project=${PROJECT_ID}
 
gcloud compute networks subnets create ${SUBNET_NAME_2} \
   --network=${NETWORK_NAME_2} \
   --range=10.10.0.0/18 \
   --region=${REGION} \
   --project=${PROJECT_ID}
 
gcloud compute firewall-rules create ${FIREWALL_RULE_NAME} \
   --network=${NETWORK_NAME_2} \
   --allow tcp,icmp,udp \
   --project=${PROJECT_ID}
 
gcloud compute routers create ${ROUTER_NAME} \
   --project=${PROJECT_ID} \
   --network=${NETWORK_NAME_2} \
   --region=${REGION}
 
gcloud compute routers nats create ${NAT_CONFIG} \
   --router=${ROUTER_NAME} \
   --region=${REGION} \
   --auto-allocate-nat-external-ips \
   --nat-all-subnet-ip-ranges \
   --project=${PROJECT_ID} \
   --enable-logging
 
export CLUSTER_ARGUMENTS="--enable-dataplane-v2 --enable-ip-alias --enable-multi-networking --network=${NETWORK_NAME_1} --subnetwork=${SUBNET_NAME_1}"
export NODE_POOL_ARGUMENTS="--additional-node-network network=${NETWORK_NAME_2},subnetwork=${SUBNET_NAME_2}"

python3 xpk.py cluster create \
   --cluster=${CLUSTER_NAME} \
   --cluster-cpu-machine-type=e2-standard-8 \
   --num-slices=${NUM_SLICES} \
   --tpu-type=${TPU_TYPE} \
   --zone=${ZONE}  \
   --project=${PROJECT_ID} \
   --on-demand \
   --custom-cluster-arguments="${CLUSTER_ARGUMENTS}" \
   --custom-nodepool-arguments="${NODE_POOL_ARGUMENTS}" \
   --create-vertex-tensorboard

命令标志说明

变量 说明
CLUSTER_NAME 用户分配的 XPK 集群名称。
PROJECT_ID Google Cloud 项目名称。使用现有项目或创建新项目。如需了解详情，请参阅设置您的 Google Cloud 项目。
ZONE 如需了解支持的可用区，请参阅 Cloud TPU 区域和可用区文档。
TPU_TYPE 请参阅加速器类型。
NUM_SLICES 您要创建的切片数
CLUSTER_ARGUMENTS 要使用的网络和子网。

例如：--enable-dataplane-v2 --enable-ip-alias --enable-multi-networking --network=${NETWORK_NAME_1} --subnetwork=${SUBNET_NAME_1}

NODE_POOL_ARGUMENTS 要使用的其他节点网络。

例如：--additional-node-network network=${NETWORK_NAME_2},subnetwork=${SUBNET_NAME_2}

NUM_SLICES 要创建的切片数（只有多切片需要）。
NETWORK_NAME 要使用的次要网络的名称。
NETWORK_FW_NAME 要使用的次要网络防火墙的名称。

设置 JAX 或 PyTorch

以下资源介绍了如何根据您使用的预配和管理方法在 Cloud TPU 上设置 JAX 或 PyTorch：

如需设置和运行 MaxText 的 XPK，请参阅使用 XPK 大规模运行 MaxText 。

优化网络性能

本部分介绍如何通过配置最大传输单元 (MTU)、在多切片环境中使用多 NIC 以及改进 TCP 设置来优化网络性能。

配置 MTU

为获得最佳网络性能，请使用 MTU（最大传输单元）为 8,896 的网络。

默认情况下，Virtual Private Cloud (VPC) 仅提供 1,460 字节的 MTU，这会使网络性能欠佳。您可以将 VPC 网络的 MTU 设置为 1,300 字节到 8,896 字节之间（含边界值）的任何值。常见的自定义 MTU 大小为 1,500 字节（标准以太网）或 8,896 字节（可能的最大值）。如需了解详情，请参阅有效的 VPC 网络 MTU 大小。

如需详细了解如何更改现有网络或默认网络的 MTU 设置，请参阅更改 VPC 网络的 MTU 设置。

以下示例会创建一个 MTU 为 8,896 的网络，以及一条允许网络内 TCP、ICMP 和 UDP 流量的相应防火墙规则。

export RESOURCE_NAME=your-resource-name
export NETWORK_NAME=${RESOURCE_NAME}-privatenetwork
export NETWORK_FW_NAME=${RESOURCE_NAME}-privatefirewall
gcloud compute networks create ${NETWORK_NAME} --mtu=8896 --project=${PROJECT_ID} \
    --subnet-mode=auto --bgp-routing-mode=regional
gcloud compute firewall-rules create ${NETWORK_FW_NAME} --network=${NETWORK_NAME} \
    --allow tcp,icmp,udp --project=${PROJECT_ID}

将 your-resource-name 替换为网络和防火墙的基本名称。

使用多切片的多 NIC 选项

如果您使用的是多切片环境，请设置以下环境变量，这些变量是次要子网所必需的：

export NETWORK_NAME_2=${RESOURCE_NAME}
export SUBNET_NAME_2=${RESOURCE_NAME}
export FIREWALL_RULE_NAME=${RESOURCE_NAME}
export ROUTER_NAME=${RESOURCE_NAME}-network-2
export NAT_CONFIG=${RESOURCE_NAME}-natconfig-2
export REGION=your-region

使用以下命令可为网络和子网创建自定义 IP 路由。

  1. 创建辅助网络。

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

  2. 为辅助网络创建子网。

    gcloud compute networks subnets create ${SUBNET_NAME_2} \
--network=${NETWORK_NAME_2} \
--range=10.10.0.0/18 --region=${REGION} \
--project=${PROJECT_ID}

  3. 创建防火墙规则，以允许在新子网内传递流量。

    gcloud compute firewall-rules create ${FIREWALL_RULE_NAME} \
--network=${NETWORK_NAME_2} --allow tcp,icmp,udp \
--source-ranges 10.10.0.0/18 --project=${PROJECT_ID}

  4. 为辅助网络创建 Cloud Router 路由器。

    gcloud compute routers create ${ROUTER_NAME} \
--project=${PROJECT_ID} \
--network=${NETWORK_NAME_2} \
--region=${REGION}

  5. 为 Cloud Router 创建 NAT 配置。

    gcloud compute routers nats create ${NAT_CONFIG} \
--router=${ROUTER_NAME} \
--region=${REGION} \
--auto-allocate-nat-external-ips \
--nat-all-subnet-ip-ranges \
--project=${PROJECT_ID} \
--enable-logging

创建多网络切片后，您可以通过设置 XPK 集群并将 --command ifconfig 标志添加到 XPK 工作负载创建命令中来验证是否使用了两个网络接口卡 (NIC)。

  1. 使用以下 workload create 命令可在 Google Cloud 控制台日志中显示 ifconfig 命令的输出，并检查 eth0 和 eth1 的 MTU 是否均设置为 8,896。

    python3 xpk.py workload create \
    --cluster CLUSTER_NAME \
    {--base-docker-image maxtext_base_image | --docker-image your-cloud-image-name} \
    --workload=${USER}-xpk-${ACCELERATOR_TYPE}-${NUM_SLICES} \
    --tpu-type=${ACCELERATOR_TYPE} \
    --num-slices=${NUM_SLICES}  \
    --on-demand \
    --zone=${ZONE} \
    --project=${PROJECT_ID} \
    --command "ifconfig"

    如果您想启用调试日志或使用 Vertex AI TensorBoard，请将以下可选参数添加到该命令中：

    --enable-debug-logs \
--use-vertex-tensorboard

  2. 通过检查 Google Cloud 控制台日志中 XPK 工作负载的输出，验证 eth0 和 eth1 的 MTU 是否均设置为 8,896。

改进了 TCP 设置

如果您使用已排队的资源预配了 Cloud TPU，则可以运行以下命令，通过增加 TCP 接收缓冲区限制来提高网络性能。

gcloud alpha compute tpus queued-resources ssh "${QUEUED_RESOURCE_ID}" \
    --project "${PROJECT_ID}" \
    --zone "${ZONE}" \
    --node=all \
    --worker=all \
    --command='
    sudo sh -c "echo \"4096 41943040 314572800\" > /proc/sys/net/ipv4/tcp_rmem"'

优化内存分配性能

Cloud TPU 虚拟机默认使用 tcmalloc 库来提升具有较大规模、高频率内存分配的模型的性能。这是通过 LD_PRELOAD 环境变量配置的。

不过，对于某些工作负载（例如嵌入表分配非常大的 DLRM），tcmalloc 可能会造成运行缓慢。在这种情况下，您可以在运行训练脚本之前，通过在 shell 会话中取消设置 LD_PRELOAD 变量来恢复为标准 malloc 函数：

unset LD_PRELOAD

使用 SkyPilot

您可以将 Cloud TPU v6e 与 SkyPilot 搭配使用。SkyPilot 是一个开源框架，可简化运行、管理和伸缩 AI 工作负载的过程。您可以向 SkyPilot 添加与 v6e 相关的位置和价格信息。如需了解详情，请参阅 SkyPilot TPU v6e 示例。

训练样本

以下部分提供了在 Cloud TPU v6e 上训练 MaxText、MaxDiffusion 和 PyTorch 模型的样本。

这些示例已通过以下软件版本测试：

  • Python 3.10 或更高版本
  • 夜间版软件版本：
    • 夜间版 JAX 0.4.32.dev20240912
    • 夜间版 LibTPU 0.1.dev20240912+nightly
  • 稳定版软件版本：
    • JAX + JAX Lib（版本 0.4.37）

在 Cloud TPU v6e 上训练 MaxText 和 MaxDiffusion

以下部分介绍了 MaxText 和 MaxDiffusion 模型的训练生命周期。

一般来说，简要步骤如下：

  1. 构建工作负载基础映像。
  2. 使用 XPK 运行工作负载。
    1. 为工作负载构建训练命令。
    2. 部署工作负载。
  3. 跟踪工作负载并查看指标。
  4. 如果不需要 XPK 工作负载，请将其删除。
  5. 如果不再需要 XPK 集群时，请将其删除。

构建基础映像

安装 MaxText 或 MaxDiffusion 并构建 Docker 映像：

  1. 克隆您要使用的仓库，并切换到该仓库的目录：

    MaxText：

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

    MaxDiffusion：

    git clone https://github.com/google/maxdiffusion.git && cd maxdiffusion && git checkout 4a8155ec0129512812b31930f0a91c6d5a141103

  2. 配置 Docker 以使用 Google Cloud CLI：

    gcloud auth configure-docker

  3. 使用以下命令或使用 JAX AI 映像构建 Docker 映像。如需详细了解 JAX AI 图片，请参阅 JAX AI 图片。

    MaxText：

    bash docker_build_dependency_image.sh MODE=stable JAX_VERSION=0.4.35

    MaxDiffusion：

    bash .github/workflows/build_and_upload_images.sh CLOUD_IMAGE_NAME=maxdiffusion_jax_stable_stack MODE=jax_ai_image PROJECT=${PROJECT_ID} LOCAL_IMAGE_NAME=maxdiffusion_jax_stable_stack BASEIMAGE=us-docker.pkg.dev/cloud-tpu-images/jax-ai-image/tpu:latest

  4. 在活跃的 gcloud CLI 配置中设置项目 ID：

    gcloud config set project ${PROJECT_ID}

  5. 如果您要从未在本地构建映像的机器启动工作负载，请上传映像。

    1. 设置 CLOUD_IMAGE_NAME 环境变量：

      export CLOUD_IMAGE_NAME=${USER}_runner

    2. 上传映像：

      bash docker_upload_runner.sh ${CLOUD_IMAGE_NAME}

使用 XPK 运行工作负载

  1. 如果您不使用由 MaxText 设置的默认值或由 MaxDiffusion 设置的默认值，请设置以下环境变量：

    export BASE_OUTPUT_DIR=gs://YOUR_BUCKET
export PER_DEVICE_BATCH_SIZE=2
export NUM_STEPS=30
export MAX_TARGET_LENGTH=8192

  2. 构建模型脚本。此脚本将在后续步骤中复制为训练命令。

    暂时不要执行模型脚本。

    MaxText

    MaxText 是一种高性能、高度可伸缩的开源 LLM，采用纯 Python 和 JAX 编写，以 Google Cloud TPU 和 GPU 为目标，用于训练和推理。

    JAX_PLATFORMS=tpu,cpu \
ENABLE_PJRT_COMPATIBILITY=true \
TPU_SLICE_BUILDER_DUMP_CHIP_FORCE=true \
TPU_SLICE_BUILDER_DUMP_ICI=true && \
python3 -m MaxText.train MaxText/configs/base.yml \
     base_output_directory=${BASE_OUTPUT_DIR} \
     dataset_type=synthetic \
     per_device_batch_size=${PER_DEVICE_BATCH_SIZE} \
     enable_checkpointing=false \
     gcs_metrics=true \
     profiler=xplane \
     skip_first_n_steps_for_profiler=5 \
     steps=${NUM_STEPS}  # attention='dot_product'"

    Gemma2

    Gemma 是 Google DeepMind 基于 Gemini 研究和技术开发的一系列开放权重 LLM。

    python3 -m MaxText.train MaxText/configs/base.yml \
    model_name=gemma2-27b \
    run_name=gemma2-27b-run \
    base_output_directory=${BASE_OUTPUT_DIR} \
    max_target_length=${MAX_TARGET_LENGTH} \
    per_device_batch_size=${PER_DEVICE_BATCH_SIZE} \
    steps=${NUM_STEPS} \
    enable_checkpointing=false \
    use_iota_embed=true \
    gcs_metrics=true \
    dataset_type=synthetic \
    profiler=xplane \
    attention=flash

    Mixtral 8x7b

    Mixtral 是 Mistral AI 开发的一款先进的 AI 模型，采用稀疏的混合专家 (MoE) 架构。

    python3 -m MaxText.train MaxText/configs/base.yml \
    base_output_directory=${BASE_OUTPUT_DIR} \
    per_device_batch_size=${PER_DEVICE_BATCH_SIZE} \
    model_name=mixtral-8x7b \
    steps=${NUM_STEPS} \
    max_target_length=${MAX_TARGET_LENGTH} \
    tokenizer_path=assets/tokenizer.mistral-v1 \
    attention=flash \
    dtype=bfloat16 \
    dataset_type=synthetic \
    profiler=xplane

    Llama3-8b

    Llama 是 Meta 开发的一系列开放权重 LLM。

    如需查看如何在 PyTorch 上运行 Llama3 的示例，请参阅 torchprime 仓库中的 torch_xla 模型。

    MaxDiffusion

    MaxDiffusion 是一系列参考实现，其中包含在 XLA 设备（包括 Cloud TPU 和 GPU）上运行的各种潜在 diffusion 模型，采用纯 Python 和 JAX 编写。Stable Diffusion 是一种潜在文本转图片模型，可根据任何文本输入生成逼真图片。

    您需要安装特定的 Git 分支才能运行 MaxDiffusion，如以下训练脚本所示。

    git clone https://github.com/google/maxdiffusion.git
&& cd maxdiffusion
&& git checkout 4a8155ec0129512812b31930f0a91c6d5a141103
&& pip install -r requirements.txt && pip install .
&& pip install huggingface_hub==0.30.2 && OUT_DIR=${BASE_OUTPUT_DIR}
&& python src/maxdiffusion/train_sdxl.py \
    src/maxdiffusion/configs/base_xl.yml \
    revision=refs/pr/95 \
    activations_dtype=bfloat16 \
    weights_dtype=bfloat16 \
    resolution=1024 \
    per_device_batch_size=1 \
    output_dir=${OUT_DIR} \
    jax_cache_dir=${OUT_DIR}/cache_dir/ \
    max_train_steps=200 \
    attention=flash \
    run_name=sdxl-ddp-v6e

  3. 导出以下变量：

    export CLUSTER_NAME=CLUSTER_NAME
export ACCELERATOR_TYPE=ACCELERATOR_TYPE
export NUM_SLICES=NUM_SLICES
export YOUR_MODEL_SCRIPT=YOUR_MODEL_SCRIPT

    环境变量说明

    变量 说明
    CLUSTER_NAME XPK 集群的名称。
    ACCELERATOR_TYPE 加速器类型用于指定您要创建的 Cloud TPU 的版本和大小。如需详细了解每个 TPU 版本支持的加速器类型，请参阅 TPU 版本。
    NUM_SLICES TPU 切片数量。
    YOUR_MODEL_SCRIPT 要作为训练命令执行的模型脚本。

  4. 使用您在上一步中创建的脚本运行模型。您必须指定 --base-docker-image 标志以使用 MaxText 基础映像，或者指定 --docker-image 标志和要使用的映像。

    您可以选择添加以下可选标志：

    • 您可以通过添加 --enable-debug-logs 标志来启用调试日志记录。如需了解详情，请参阅在 MaxText 上调试 JAX。
    • 您可以创建 Vertex AI Experiments 实验，通过添加 --use-vertex-tensorboard 标志将数据上传到 Vertex AI TensorBoard。如需了解详情，请参阅使用 Vertex AI 在 MaxText 上监控 JAX。
    python3 xpk.py workload create \
  --cluster ${CLUSTER_NAME} \
  {--base-docker-image maxtext_base_image | --docker-image gcr.io/${PROJECT_ID}/${CLOUD_IMAGE_NAME}:latest} \
  --workload=${USER}-xpk-${ACCELERATOR_TYPE}-${NUM_SLICES} \
  --tpu-type=${ACCELERATOR_TYPE} \
  --num-slices=${NUM_SLICES}  \
  --on-demand \
  --zone=${ZONE} \
  --project=${PROJECT_ID} \
  --command="${YOUR_MODEL_SCRIPT}"

    输出包含一个链接，您可以通过该链接跟踪工作负载。 打开该链接，然后点击日志标签页，以实时跟踪您的工作负载。

在 MaxText 上调试 JAX

使用补充性 XPK 命令诊断集群或工作负载未运行的原因：

使用 Vertex AI 在 MaxText 上监控 JAX

如需使用 TensorBoard，您的 Google Cloud 用户账号必须具有 aiplatform.user 角色。运行以下命令可授予此角色：

gcloud projects add-iam-policy-binding your-project-id \
   --member='user:your-email' \
   --role='roles/aiplatform.user'

通过 Vertex AI 托管式 TensorBoard 查看标量和性能分析数据。

  1. 将您使用的可用区的资源管理 (CRUD) 请求数量从 600 增加到 5,000。对于使用少于 16 个虚拟机的较小工作负载，这可能不是问题。

  2. 安装 Vertex AI 的依赖项（例如 cloud-accelerator-diagnostics）：

    # xpk dependencies will install cloud-accelerator-diagnostics for Vertex AI
cd ~/xpk
pip install .

  3. 按照创建 Vertex AI TensorBoard 中所述，使用 --create-vertex-tensorboard 标志创建 XPK 集群。您也可以在现有集群上运行此命令。

  4. 在运行 XPK 工作负载时，使用 --use-vertex-tensorboard 标志和可选的 --experiment-name 标志创建 Vertex AI Experiments 实验。如需查看完整的步骤列表，请参阅创建 Vertex AI Experiments 实验以将数据上传到 Vertex AI TensorBoard。

日志包含指向 Vertex AI TensorBoard 的链接，类似于以下内容：

View your TensorBoard at https://us-central1.tensorboard.googleusercontent.com/experiment/project_id+locations+us-central1+tensorboards+hash+experiments+name

您也可以在 Google Cloud 控制台中找到 Vertex AI TensorBoard 链接。前往 Google Cloud 控制台中的 Vertex AI Experiments。从下拉菜单中选择相应区域。

TensorBoard 目录也会写入您使用 ${BASE_OUTPUT_DIR} 指定的 Cloud Storage 存储桶。

删除 XPK 工作负载

使用 xpk workload delete 命令根据作业前缀或作业状态删除一个或多个工作负载。如果您发送的 XPK 工作负载不再需要运行，或者您的作业卡在队列中，则此命令可能会很有用。

删除 XPK 集群

使用 xpk cluster delete 命令删除集群：

python3 xpk.py cluster delete --cluster ${CLUSTER_NAME} \
    --zone=${ZONE} --project=${PROJECT_ID}

MaxDiffusion 基准测试结果

我们在 v6e-4、v6e-16 和两个 v6e-16 上运行了 MaxDiffusion 的训练脚本。下表显示了测得的吞吐量。

v6e-4 v6e-16 两个 v6e-16
训练步骤 0.069 0.073 0.13
全局批量大小 8 32 64
吞吐量（样本/秒） 115.9 438.4 492.3

在 Cloud TPU v6e 上使用 PyTorch/XLA 训练 Llama 模型

本部分介绍了如何使用 WikiText 数据集在 Cloud TPU v6e 上使用 PyTorch/XLA 训练 Llama 模型。

获取对 Hugging Face 和 Llama 3 模型的访问权限

您需要拥有 Hugging Face 用户访问令牌才能运行本示例。如需了解如何创建用户访问令牌，请参阅 Hugging Face 文档中的用户访问令牌。

您还需要有权访问 Hugging Face 上的 Llama-3-8B 模型。如需获得访问权限，请前往 HuggingFace 上的 Meta-Llama-3-8B 模型并申请访问权限。

创建 Cloud TPU 虚拟机

在此示例中，创建一个具有 8 个芯片的 Cloud TPU v6e。

  1. 设置环境变量：

    export PROJECT_ID=your-project-id
export TPU_NAME=your-tpu-name
export ZONE=us-east1-d
export ACCELERATOR_TYPE=v6e-8
export RUNTIME_VERSION=v2-alpha-tpuv6e

    环境变量说明

    变量 说明
    PROJECT_ID 您的 Google Cloud 项目 ID。使用现有项目或创建新项目。
    TPU_NAME TPU 的名称。
    ZONE 要在其中创建 TPU 虚拟机的可用区。如需详细了解支持的可用区，请参阅 TPU 区域和可用区。
    ACCELERATOR_TYPE 加速器类型用于指定您要创建的 Cloud TPU 的版本和大小。如需详细了解每个 TPU 版本支持的加速器类型，请参阅 TPU 版本。
    RUNTIME_VERSION Cloud TPU 软件版本。

  2. 创建 Cloud TPU 虚拟机：

    gcloud alpha compute tpus tpu-vm create ${TPU_NAME} --version=${RUNTIME_VERSION} \
   --accelerator-type=${ACCELERATOR_TYPE} \
   --zone=${ZONE} \
   --project=${PROJECT_ID}

安装

安装 Hugging Face Transformer 和依赖项的 pytorch-tpu/transformers 分支。此示例是使用以下依赖项版本进行测试的：

  • torch：与 2.5.0 兼容
  • torch_xla[tpu]：与 2.5.0 兼容
  • jax：0.4.33
  • jaxlib：0.4.33
gcloud alpha compute tpus tpu-vm ssh ${TPU_NAME} \
   --project=${PROJECT_ID} \
   --zone ${ZONE} \
   --worker=all \
   --command='git clone -b flash_attention https://github.com/pytorch-tpu/transformers.git
   cd transformers
   sudo pip3 install -e .
   pip3 install datasets
   pip3 install evaluate
   pip3 install scikit-learn
   pip3 install accelerate
   pip install torch~=2.6.0 torch_xla[tpu]~=2.6.0 -f https://storage.googleapis.com/libtpu-releases/index.html -f https://storage.googleapis.com/libtpu-wheels/index.html
   pip install jax==0.4.38 jaxlib==0.4.38 -i https://us-python.pkg.dev/ml-oss-artifacts-published/jax/simple/'

设置模型配置文件

下一部分运行模型中的训练命令使用两个 JSON 配置文件来定义模型参数和完全分片数据并行 (FSDP) 配置。借助 FSDP 分片，您可以在训练期间使用更大的批次大小，方法是将模型权重分片到多个 TPU 上。使用较小的模型进行训练时，使用数据并行处理并在每个设备上复制权重可能就足够了。如需详细了解如何在 PyTorch/XLA 中跨设备对张量进行分片，请参阅 PyTorch/XLA SPMD 用户指南。

  1. 创建模型参数配置文件。以下是 Llama-3-8B 的模型参数配置。对于其他型号，请在 Hugging Face 上查找配置文件。例如，请参阅 Llama-2-7B 配置。

    cat > llama-config.json << EOF
{
  "architectures": [
    "LlamaForCausalLM"
  ],
  "attention_bias": false,
  "attention_dropout": 0.0,
  "bos_token_id": 128000,
  "eos_token_id": 128001,
  "hidden_act": "silu",
  "hidden_size": 4096,
  "initializer_range": 0.02,
  "intermediate_size": 14336,
  "max_position_embeddings": 8192,
  "model_type": "llama",
  "num_attention_heads": 32,
  "num_hidden_layers": 32,
  "num_key_value_heads": 8,
  "pretraining_tp": 1,
  "rms_norm_eps": 1e-05,
  "rope_scaling": null,
  "rope_theta": 500000.0,
  "tie_word_embeddings": false,
  "torch_dtype": "bfloat16",
  "transformers_version": "4.40.0.dev0",
  "use_cache": false,
  "vocab_size": 128256
}
EOF

  2. 创建 FSDP 配置文件：

    cat > fsdp-config.json << EOF
{
  "fsdp_transformer_layer_cls_to_wrap": [
    "LlamaDecoderLayer"
  ],
  "xla": true,
  "xla_fsdp_v2": true,
  "xla_fsdp_grad_ckpt": true
}
EOF

    如需详细了解 FSDP，请参阅使用 SPMD 的完全分片数据并行处理 。

  3. 使用以下命令将配置文件上传到 Cloud TPU 虚拟机：

    gcloud alpha compute tpus tpu-vm scp llama-config.json fsdp-config.json ${TPU_NAME}:. \
   --worker=all \
   --project=${PROJECT_ID} \
   --zone=${ZONE}

运行模型

使用您在上一个部分中创建的配置文件，运行 run_clm.py 脚本以在 WikiText 数据集上训练 Llama-3-8B 模型。训练脚本在 Cloud TPU v6e-8 上运行大约需要 10 分钟。

  1. 使用以下命令在 Cloud TPU 上登录 Hugging Face：

    gcloud alpha compute tpus tpu-vm ssh ${TPU_NAME} \
   --project=${PROJECT_ID} \
   --zone ${ZONE} \
   --worker=all \
   --command='
   pip3 install "huggingface_hub[cli]"
   huggingface-cli login --token HUGGING_FACE_TOKEN'

  2. 运行模型训练：

    gcloud alpha compute tpus tpu-vm ssh ${TPU_NAME} \
   --project=${PROJECT_ID} \
   --zone ${ZONE} \
   --worker=all \
   --command='
   export PJRT_DEVICE=TPU
   export XLA_USE_SPMD=1
   export ENABLE_PJRT_COMPATIBILITY=true
   # Optional variables for debugging:
   export XLA_IR_DEBUG=1
   export XLA_HLO_DEBUG=1
   export PROFILE_EPOCH=0
   export PROFILE_STEP=3
   export PROFILE_DURATION_MS=100000
   # Set PROFILE_LOGDIR to a local VM path or gs://my-bucket/profile_path
   export PROFILE_LOGDIR=PROFILE_PATH
   python3 transformers/examples/pytorch/language-modeling/run_clm.py \
     --dataset_name wikitext \
     --dataset_config_name wikitext-2-raw-v1 \
     --per_device_train_batch_size 16 \
     --do_train \
     --output_dir /home/$USER/tmp/test-clm \
     --overwrite_output_dir \
     --config_name /home/$USER/llama-config.json \
     --cache_dir /home/$USER/cache \
     --tokenizer_name meta-llama/Meta-Llama-3-8B \
     --block_size 8192 \
     --optim adafactor \
     --save_strategy no \
     --logging_strategy no \
     --fsdp "full_shard" \
     --fsdp_config /home/$USER/fsdp-config.json \
     --torch_dtype bfloat16 \
     --dataloader_drop_last yes \
     --flash_attention \
     --max_steps 20'

排查 PyTorch/XLA 问题

如果您在上一部分中设置了可选变量以进行调试，则模型的配置文件将存储在由变量 PROFILE_LOGDIR 指定的位置。您可以提取存储在此位置的 xplane.pb 文件，并使用 tensorboard 按照 TensorBoard 说明在浏览器中查看性能分析文件。

如果 PyTorch/XLA 未按预期运行，请参阅问题排查指南，其中包含有关对模型进行调试、性能分析和优化的建议。