使用 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}