GKE 中的 TPU 简介

本页面介绍了 Cloud TPU，并说明了可从哪里找到有关在 Google Kubernetes Engine (GKE) 上使用 Cloud TPU 的信息。张量处理单元 (TPU) 是 Google 定制开发的专用集成电路 (ASIC)，用于加速使用框架的机器学习工作负载，例如：TensorFlow、PyTorch 和 JAX。

使用 GKE 中的 TPU 之前，我们建议您完成以下学习路线：

  1. 参阅 Cloud TPU 简介，了解机器学习加速器的工作原理。
  2. 了解 Cloud TPU 系统架构中的当前 TPU 版本可用性。

如需了解如何在 GKE 中设置 Cloud TPU，请参阅以下资源：

使用 GKE 中 TPU 的优势

GKE 全面支持 TPU 虚拟机生命周期管理，包括创建、配置和删除 TPU 虚拟机。GKE 还支持 Spot 虚拟机和使用预留的 Cloud TPU。在 GKE 中使用 TPU 的好处包括：

  • 一致的操作环境：面向所有机器学习和其他工作负载的单一平台。
  • 自动升级：GKE 自动执行版本更新，从而减少运营开销。
  • 负载均衡：GKE 会分配负载，从而缩短延迟时间并提高可靠性。
  • 自适应扩缩：GKE 会自动扩缩 TPU 资源以满足工作负载的需求。
  • 资源管理：使用 Kueue（Kubernetes 原生作业排队系统），您可以使用排队、抢占、优先级和公平共享来管理组织内多个租户的资源。

与 GKE 中的 TPU 相关的术语

本文档使用以下与 TPU 相关的术语：

  • TPU 类型：Cloud TPU 类型，如 v5e。
  • TPU 切片：一组互连的 TPU 芯片。
  • TPU 拓扑：TPU 切片中 TPU 芯片的数量和物理排列方式。
  • 单主机 TPU 切片节点：一个或多个独立 TPU 节点。连接到虚拟机的 TPU 不通过高速互连相互连接。
  • 多主机 TPU 切片节点：两个或更多个互连的 TPU 节点。连接到虚拟机的 TPU 通过高速互连进行连接。多主机 TPU 节点具有以下特征：
    • 原子性：GKE 将所有互连的节点视为一个单元。在扩缩操作期间，GKE 会将整个节点组缩减为 0 并创建新节点。如果节点组中的机器发生故障或终止，GKE 会作为一个新单元重新创建整个节点组。
    • 不可变：您无法手动向一组互连的节点添加新节点。

GKE 中 TPU 的工作原理

Kubernetes 资源管理和优先级处理 TPU 虚拟机的方式与此类其他类型的虚拟机相同。您使用资源名称 google.com/tpu 请求 TPU 芯片：

    resources:
        requests:
          google.com/tpu: 4
        limits:
          google.com/tpu: 4

使用 GKE 中的 TPU 时，必须考虑以下 TPU 特性：

  • 一个 TPU 虚拟机最多可以访问 8 个 TPU 芯片。
  • TPU 切片包含固定数量的 TPU 芯片，具体取决于您选择的 TPU 机器类型。
  • 请求的 google.com/tpu 数量必须等于 TPU 节点上的可用芯片总数。GKE Pod 中请求 TPU 的任何容器都必须使用节点中的所有 TPU 芯片。否则，您的 Deployment 会失败，因为 GKE 无法使用部分 TPU 资源。例如，请参阅以下场景：
    • 机器类型 ct5l-hightpu-8t 具有单个 TPU 节点，其中包含 8 个 TPU 芯片。一个需要 8 个 TPU 芯片的 GKE Pod 可以部署在该节点上，但两个各自需要 4 个 TPU 芯片的 GKE Pod 无法部署在该节点上。
    • 具有 2x4 拓扑的机器类型 ct5lp-hightpu-4t 包含两个 TPU 节点，每个节点具有 4 个芯片，总共有 8 个 TPU 芯片。一个需要 8 个 TPU 芯片的 GKE Pod 无法部署在此节点池中的任何节点上，但两个各自需要 4 个 TPU 芯片的 Pod 可以部署在此节点池中的两个节点上。
    • 采用拓扑 4x4 的 TPU v5e 在 4 个节点中有 16 个芯片。选择此配置的 GKE Autopilot 工作负载必须在每个副本中请求 4 个芯片，最多 4 个副本。
  • 在 Standard 集群中，一个 TPU 虚拟机上可以调度多个 Kubernetes Pod，但每个 Pod 中只有一个容器可以访问 TPU 芯片。
  • 每个 Standard 集群必须至少具有一个非 TPU 节点池，才能创建 kube-system Pod，例如 kube-dns。
  • 默认情况下，TPU 节点具有google.com/tpu污点，以防止非 TPU Pod 被安排到它们上。污点无法保证 TPU 资源得到充分利用。它可让您在非 TPU 节点上运行不使用 TPU 的工作负载，从而释放 TPU 节点上的计算资源对使用 TPU 的代码的影响。
  • GKE 会收集 TPU 虚拟机上运行的容器发出的日志。如需了解详情，请参阅日志记录。
  • Cloud Monitoring 提供 TPU 利用率指标（例如运行时性能）。如需了解详情，请参阅可观测性和指标。

规划 TPU 配置

如需在 GKE 集群中使用 TPU，您必须确定以下参数：

  • TPU 类型：不同的 TPU 类型具有不同的功能，例如性价比、训练吞吐量和服务延迟时间。连接 TPU 的虚拟机具有不同的 CPU 和内存容量。
  • 拓扑：TPU 类型中芯片的布局。每种 TPU 类型都支持 2D 或 3D 芯片布局和特定排列方式。选择与模型的并行性要求匹配的拓扑。

  • TPU 互连：TPU 类型和拓扑决定了您是否在具有高速互连的多个节点中获得 TPU。这决定了您需要单主机还是多主机 TPU 切片节点。

    • 对于大规模模型，请使用多主机 TPU 切片节点
    • 对于小规模模型，请使用单主机 TPU 切片节点

为 GKE Autopilot 模式选择 TPU 配置

在 Autopilot 中，您需要选择 TPU 类型和拓扑，并在 Kubernetes 清单中指定这些内容。GKE 会为节点预配 TPU 并调度工作负载进行管理。

GKE Autopilot 中的 TPU 可用性

TPU 仅在特定 Google Cloud 区域提供。如需在 GKE 工作负载中使用 TPU 类型，您的集群必须位于该类型支持的区域。如需了解详情，请参阅 Cloud TPU 文档中的 TPU 区域和可用区。

在 Autopilot 中选择 TPU 类型

TPU 类型 vCPU 的数量 内存 (GiB) NUMA 节点的数量 切片中的 TPU 芯片数量上限
TPU v5p（预览版）
tpu-v5p-slice		 208 448 2 6144
TPU v5e
tpu-v5-lite-podslice		 24 到 224 48 到 384 1 256
TPU v5e（仅限单主机）
tpu-v5-lite-device		 24 到 224 48 到 384 1 到 2 8
TPU v4
tpu-v4-podslice		 240 407 2 4096

请查看 Cloud TPU 文档中的芯片规格和价格，以帮助您决定要使用的版本。

为 Autopilot 选择拓扑

拓扑是 TPU 切片中芯片的物理布局。拓扑是二维或三维拓扑，具体取决于 TPU 类型。确定 TPU 类型后，请选择该 TPU 类型支持的拓扑。模型的并行性要求可帮助您确定拓扑。拓扑的乘积是切片中的芯片数量。例如：

  • 2x2x2 是一个 8 芯片多主机 TPU v4 切片
  • 2x2 是一个 4 芯片单主机 TPU v5e 切片

您为 TPU 类型选择的拓扑决定了您获得单主机还是多主机 TPU 切片节点。如果特定拓扑同时支持单主机和多主机切片，则工作负载请求的 TPU 芯片的数量决定您获得的主机类型。例如，TPU v5e (tpu-v5-lite-podslice) 支持 2x4 拓扑同时作为单主机和多主机切片。如果您在工作负载中请求 4 个芯片，则会获得一个具有 4 个芯片的多主机节点。如果您在工作负载中请求 8 个芯片，则会获得一个具有 8 个芯片的单主机节点。

下表介绍了每种 TPU 类型支持的拓扑，以及芯片数量、节点数量和主机类型：

TPU 类型 拓扑 切片中的 TPU 芯片数量 节点数量 主机类型 Notes
TPU v5p（预览版）
tpu-v5p-slice		 2x2x1 4 1 单个主机

支持超过 64 个芯片的自定义拓扑。 以下条件适用：

  • 如果芯片数量超过 64 个，{A}、{B} 和 {C} 必须是 4 的倍数
  • 最大拓扑为 16x16x24
  • 值必须是 {A}≤{B}≤{C}，例如 8x12x16。
2x2x2 8 2 多主机
2x2x4 16 4 多主机
2x4x4 32 8 多主机
4x4x4 64 16 多主机
{A}x{B}x{C} A*B*C (A*B*C/4)1 多主机
TPU v5e
tpu-v5-lite-podslice		 1x1 1 1 单个主机 不支持自定义拓扑。
2x2 4 1
2x4 8 1
2x4 2 1 多主机
4x4 16 4
4x8 32 8
8x8 64 16
8x16 128 32
16x16 256 64
TPU v5e（仅限单主机）
tpu-v5-lite-device		 1x1 1 1 单个主机 不支持自定义拓扑
2x2 4 1
2x4 8 1
TPU v4
tpu-v4-podslice		 2x2x1 4 1 单个主机

支持超过 64 个芯片的自定义拓扑。 以下条件适用：

  • 如果芯片数量超过 64 个，{A}、{B} 和 {C} 必须是 4 的倍数
  • 最大拓扑为 12x16x16
  • 值必须是 {A}≤{B}≤{C}，例如 8x12x16。
2x2x2 8 2 多主机
2x2x4 16 4 多主机
2x4x4 32 8 多主机
4x4x4 64 16 多主机
{A}x{B}x{C} A*B*C (A*B*C/4)1 多主机

  1. 计算方法是用拓扑产品除以四。 ↩

选择 TPU 类型和拓扑后，请在工作负载清单中指定这些内容。如需查看相关说明，请参阅在 GKE Autopilot 上部署 TPU 工作负载。

为 GKE Standard 模式选择 TPU 配置

以下几个部分介绍了您在 GKE 中规划和设置 TPU 工作负载时要评估的 TPU 特征。如需详细了解可用版本、机器类型、有效拓扑及其芯片数量，请参阅本文档中的 TPU 配置映射。

GKE Standard 模式下的 TPU 可用性

下表列出了不同机器类型和版本的 TPU 可用性：

TPU 版本 机器类型开头为 GKE 最低版本 可用情况 可用区
TPU v4 ct4p- 1.26.1-gke.1500 正式版 us-central2-b
TPU v5e ct5l- 1.27.2-gke.2100 正式版 europe-west4-b
us-central1-a
TPU v5e ct5lp- 1.27.2-gke.2100 正式版 europe-west4-a1
us-central1-a1
us-east1-c
us-east5-b1
us-west1-c
us-west4-a
us-west4-b1
TPU v5p ct5p- 1.28.3-gke.1024000 预览 us-east1-d
us-east5-a
us-east5-c

  1. 在 europe-west4-a、us-central1-a、us-east5-b 或 us-west4-b 可用区中的任何一个中创建机器类型以 ct5lp- 开头的 TPU v5e 时，不支持单主机 TPU v5e 节点池。换句话说，在上述任何可用区中创建 TPU v5e 节点池时，仅支持拓扑至少为 2x4 或更大的机器类型 ct5lp-hightpu-4t。如需在 us-central1 或 europe-west4 中创建单主机 TPU v5e，请分别选择可用区 us-central1-a 或 europe-west4-b，并使用以 ct5l- 开头的机器类型，例如 ct5l-hightpu-1t、ct5l-hightpu-4t 或 ct5l-hightpu-8t。如需在 us-west4 区域中创建单主机 TPU v5e，请选择可用区 us-west4-a，并使用以 ct5lp- 开头的机器类型，例如 ct5lp-hightpu-1t。请注意，以 ct5l- 开头的机器类型与以 ct5lp- 开头的机器类型需要不同的配额。↩

以下几个部分介绍了您在 GKE 中规划和设置 TPU 工作负载时要评估的 TPU 特征。如需详细了解可用版本、机器类型、有效拓扑及其芯片数量，请参阅本文档中的 TPU 配置映射。

机器类型

支持 TPU 资源的机器类型遵循包含 TPU 版本和每个节点的芯片数（例如 ct<version>-hightpu-<node-chip-count>t）的命名惯例。例如，机器类型 ct5lp-hightpu-1t 支持 TPU v5e，且总共包含一个 TPU 芯片。

拓扑

拓扑定义了 TPU 切片中 TPU 的物理排列方式。GKE 会根据 TPU 版本在二维或三维拓扑中预配 TPU 切片。您可以将拓扑指定为每个维度中的 TPU 芯片的数量：

  • 对于在多主机 TPU 切片节点池中调度的 TPU v4 和 v5p，您可以使用 3 元组 ({A}x{B}x{C}) 来定义拓扑，例如 4x4x4。{A}x{B}x{C} 的乘积定义了节点池中的芯片数。例如，您可以使用 2x2x2、2x2x4 或 2x4x4 等拓扑形式定义少于 64 个芯片的小型拓扑。如果您使用多于 64 个芯片的拓扑，则分配给 {A}、{B} 和 {C} 的值必须满足以下条件：

    • {A}、{B} 和 {C} 是四的倍数。
    • v4 支持的最大拓扑为 12x16x16，v5p 为 16x16x24。
    • 分配的值保持 A ≤ B ≤ C 模式。例如 4x4x8 或 8x8x8。

TPU 配置的映射

根据您的用例使用下表定义要使用的 TPU 机器类型和拓扑：

  • 对于小规模模型训练或推理，请使用 TPU v4 或 TPU v5e 和单主机 TPU 切片节点池。
  • 对于大规模模型训练或推理，请将 TPU v4 或 TPU v5e 与多主机 TPU 切片节点池搭配使用。
TPU 版本 机器类型 拓扑 TPU 芯片数量 虚拟机数量 节点池类型
TPU v4 ct4p-hightpu-4t 2x2x1 4 1 单个主机
2x2x2 8 2 多主机
2x2x4 16 4 多主机
2x4x4 32 8 多主机
{A}x{B}x{C} A*B*C (A*B*C/4)1 多主机
TPU v5p ct5p-hightpu-4t 2x2x1 4 1 单个主机
2x2x2 8 2 多主机
2x2x4 16 4 多主机
2x4x4 32 8 多主机
{A}x{B}x{C} A*B*C (A*B*C/4)1 多主机
TPU v5e ct5l-hightpu-1t 1x1 1 1 单个主机
ct5l-hightpu-4t 2x2 4 1 单个主机
ct5l-hightpu-8t 2x4 8 1 单个主机
ct5lp-hightpu-1t 1x1 1 1 单个主机
ct5lp-hightpu-4t 2x2 4 1 单个主机
ct5lp-hightpu-8t 2x4 8 1 单个主机
ct5lp-hightpu-4t 2x4 8 2 多主机
4x4 16 4 多主机
4x8 32 8 多主机
8x8 64 16 多主机
8x16 128 32 多主机
16x16 256 64 多主机

  1. 计算方法是用拓扑产品除以四。↩

TPU v5e 特征

TPU v5e 机器具有以下技术特征：

机器类型 vCPU 的数量 内存 (GB) NUMA 节点的数量 被抢占的可能性
ct5l-hightpu-1t 24 48 1 较高
ct5l-hightpu-4t 112 192 1 中
ct5l-hightpu-8t 224 384 2 较低
ct5lp-hightpu-1t 24 48 1 较高
ct5lp-hightpu-4t 112 192 1 中
ct5lp-hightpu-8t 224 384 1 低

TPU v4 和 v5p 特征

TPU v4p 和 v5p 机器具有以下技术特征：

机器类型 vCPU 的数量 内存 (GB) NUMA 节点的数量
ct4p-hightpu-4t 240 407 2
ct5p-hightpu-4t 208 448 2

TPU 预留

购买承诺时，TPU 预留可供使用。任何 TPU 预留都可以与 GKE 搭配使用。

创建 TPU 节点池时，请使用 --reservation 和 --reservation-affinity=specific 标志 来使用预留的 TPU 实例。

在 GKE 中自动扩缩 TPU

GKE 支持张量处理单元 (TPU) 来加速机器学习工作负载。单主机 TPU 切片节点池和多主机 TPU 切片节点池都支持自动扩缩和自动预配。

通过在 GKE 集群上使用 --enable-autoprovisioning 标志，GKE 可以创建或删除具有 TPU 版本和拓扑的单主机或多主机 TPU 切片节点池，以满足待处理工作负载的要求。

使用 --enable-autoscaling 时，GKE 会根据节点池的类型扩缩节点池，如下所示：

  • 单主机 TPU 切片节点池：GKE 会在现有节点池中添加或移除 TPU 节点。节点池可以包含零到节点池最大大小（由 --max-nodes 和 --total-max-nodes 标志确定）之间的任意数量的 TPU 节点。当节点池扩缩时，节点池中的所有 TPU 节点具有相同的机器类型和拓扑。如需详细了解如何创建单主机 TPU 切片节点池，请参阅创建节点池。

  • 多主机 TPU 切片节点池：GKE 以原子方式将节点池从零扩容到满足 TPU 拓扑所需的节点数。例如，对于机器类型为 ct5lp-hightpu-4t 且拓扑为 16x16 的 TPU 节点池，节点池包含 64 个节点。GKE 自动扩缩器可确保此节点池正好包含 0 个或 64 个节点。缩容时，GKE 会逐出所有调度的 Pod，并将整个节点池排空为零。如需详细了解如何创建多主机 TPU 切片节点池，请参阅创建节点池。

限制

  • 对于容量预留，您必须使用特定预留。
  • GKE 费用分配和用量计量不包括有关预留 TPU v4 用量或费用的任何数据。
  • TPU v5p 和 v5e 不支持 us-east5 中的 riptide/映像流式传输。
  • 控制平面至少运行 1.29.2-gke.1035000 或 1.28.7-gke.1020000 的 GKE 集群支持 TPU v5p 自动扩缩。

工作负载安排注意事项

TPU 具有独特的特征，需要 Kubernetes 中的特殊工作负载调度和管理。以下部分介绍了调度最佳实践。

CPU

本部分不适用于 Autopilot 集群，因为 GKE 将每个 TPU Pod 放在自己的节点上。

如需在 TPU 虚拟机的板载 CPU 上调度工作负载，请确保您的 GKE Pod 可以容忍 google.com/tpu 污点。如果您希望将工作负载部署到特定节点，请使用节点选择器。

Kubernetes 资源管理和优先级处理 TPU 虚拟机的方式与此类其他类型的虚拟机相同。如需让需要 TPU 安排的 Pod 优先于同一节点上的其他 Pod，请请求这些 TPU Pod 的最大 CPU 或内存。低优先级 Pod 应执行以下操作：

  1. 设置较低的 CPU 和内存请求，以确保节点为 TPU 工作负载具有足够的可分配资源。如需了解详情，请参阅 Kubernetes 如何应用资源请求和限制。
  2. 不设置 CPU 限制（无限制），以确保 Pod 可以激增以使用所有未使用的周期。
  3. 设置较高的内存限制，以确保 pod 可以在保持节点稳定性的同时使用大多数未使用的内存。

如果 Kubernetes Pod 未请求 CPU 和内存（即使正在请求 TPU），则 Kubernetes 会将其视为尽力而为的 Pod，在它需要任何 CPU 和内存时无法保证会提供。只有明确请求 CPU 和内存的 Pod 才能获得此类保证。如需了解详情，请参阅 Pod 和容器的资源管理。

如需了解详情，请参阅 Kubernetes 最佳做法：资源请求和限制。

减少工作负载中断

如果您使用 TPU 训练机器学习模型且工作负载中断，则自上次检查点以来执行的所有工作都将丢失。如需降低工作负载中断的可能性，请执行以下操作：

  • 为此作业设置比所有其他作业更高的优先级：如果资源不足，则 GKE 调度器会抢占优先级较低的作业，以调度优先级更高的作业。这样还可以确保优先级更高的工作负载接收它所需的所有资源（不超过集群中可用的总资源）。如需了解详情，请参阅 Pod 优先级和抢占。
  • 配置维护排除项：维护排除项是禁止进行自动维护的非重复性时间段。如需了解详情，请参阅维护排除项。
  • 在 Autopilot 中使用延长运行时间的 Pod：在长达七天宽限期内使用延长运行时间的 Pod，然后 GKE 终止 Pod 进行缩容或节点升级。

处理由于节点维护而导致的中断

所有 GKE 节点（包括托管 TPU 虚拟机的 GKE 节点）都会受到可能导致节点关停的维护事件或其他中断的影响。您可以减少在控制平面运行 1.29.1-gke.1425000 或更高版本的 GKE 集群中运行的工作负载的中断。GKE 会在逐出之前最多五分钟向节点发送 SIGTERM 信号，以提醒节点即将关停。如果您的工作负载将 MaxText、Pax 或 JAX 等机器学习框架与 Orbax 搭配使用，则工作负载可以捕获 SIGTERM 信号并启动检查点过程。

您可以将 GKE 配置为在最长通知时间内正常终止机器学习工作负载。在 Pod 清单中，将 spec.terminationGracePeriodSeconds 字段设置为 300 秒（五分钟）。GKE 会尽最大努力正常终止这些 Pod 并执行您定义的终止操作，例如保存训练状态。对于 PodDisruptionBudget 或 terminationGracePeriodSeconds 设置，GKE 遵循任何不超过五分钟的配置。

spec.terminationGracePeriodSeconds 字段仅处理由于非抢占式虚拟机上发生的维护和碎片整理事件而导致的中断。GKE 不会处理非自愿中断，例如硬件故障。

如需了解详情，请参阅配置 TPU 节点正常终止。

最大限度地提高 TPU 利用率

为了最大限度地提高对 TPU 的投资，请混合调度作业优先级，并将其排入队列，以最大限度增加 TPU 的运行时间。如果您需要作业级调度和抢占，则需要使用 Kubernetes 的一个插件来将作业编排到队列中。我们建议您将 Kueue 用于该应用场景。

后续步骤