规划 GKE 中的 TPU

本页面介绍了如何规划 Google Kubernetes Engine (GKE) 中的张量处理单元 (TPU) 用量，以降低 TPU 配置错误、不可用错误或配额超出中断的风险。

在 GKE 中使用 TPU 之前，请确保您熟悉 GKE 中的 TPU 定义和术语。

规划 TPU 配置

如需在 GKE 集群中使用 TPU，您必须规划其配置。我们建议您按以下步骤操作：

  1. 选择 GKE 操作模式：在 GKE Autopilot 或 Standard 集群中的 TPU 上运行工作负载。我们建议您使用 Autopilot 集群获得全托管式 Kubernetes 体验。
  2. 选择 TPU 版本：不同的 TPU 类型具有不同的功能，例如性价比、训练吞吐量和服务延迟时间。TPU 类型会影响可用的 CPU 和内存容量。
  3. 验证 TPU 可用性：TPU 可在特定 Google Cloud 区域使用。如需在 GKE 工作负载中使用 TPU 类型，您的集群必须位于该类型支持的区域。
  4. 选择 TPU 拓扑：TPU 在 TPU 切片中的物理排列方式。选择符合模型并行性要求的拓扑。

我们建议您使用本页上的参考表格来确定您的节点池是单主机 TPU 切片节点还是多主机 TPU 切片节点。

选择 GKE 操作模式

您可以在 GKE 为集群提供的操作模式中使用 TPU：

  • Autopilot 模式（推荐）：GKE 管理底层基础架构，例如节点配置、自动扩缩、自动升级、基准安全配置和基准网络配置。 在 Autopilot 中，您需要选择 TPU 类型和拓扑，然后在 Kubernetes 清单中指定这些内容。GKE 会为节点预配 TPU 并调度工作负载进行管理。
  • Standard 模式：您负责管理底层基础架构，包括配置各个节点。

如需选择最适合您的工作负载的 GKE 操作模式，请参阅选择 GKE 操作模式。

选择 TPU 版本

TPU 切片中的虚拟机具有以下技术特征：

Autopilot

TPU 版本 机器类型 vCPU 的数量 内存 (GiB) NUMA 节点的数量 TPU 切片节点中的 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

Standard

TPU 版本 机器类型 vCPU 的数量 内存 (GiB) NUMA 节点的数量 被抢占的可能性
TPU v5p
 ct5p-hightpu-4t 208 448 2
TPU v5e
 ct5l-hightpu-1t 24 48 1 较高
TPU v5e
 ct5l-hightpu-4t 112 192 1 中等
TPU v5e
 ct5l-hightpu-8t 224 384 2 较低
TPU v5e
 ct5lp-hightpu-1t 24 48 1 较高
TPU v5e
 ct5lp-hightpu-4t 112 192 1 中等
TPU v5e
 ct5lp-hightpu-8t 224 384 1 低
TPU v4
 ct4p-hightpu-4t 240 407 2

请查看 Cloud TPU 价格文档中的 TPU 规格和价格，以决定要使用的 TPU 配置。

限制

在选择要使用的 TPU 时，请考虑以下限制：

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

验证 GKE 中的 TPU 可用性

TPU 可在特定 Google Cloud 区域使用。如需在 GKE 集群中使用 TPU 类型，您的集群必须位于该类型支持的区域。

Autopilot

请参阅 Cloud TPU 文档中的 TPU 区域和可用区。

Standard

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

TPU 版本 机器类型开头为 GKE 最低版本 可用性 可用区
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
TPU v4 ct4p- 1.26.1-gke.1500 正式版 us-central2-b

  1. 您可以在特定可用区（europe-west4-a、us-east5-b 和 us-west4-b）中创建机器类型以 ct5lp- 开头但不以 ct5l- 开头的单主机 TPU v5e 节点池。您可以在这些可用区中使用拓扑至少为 2x4 或更大的 ct5lp-hightpu-4t。如需在 us-west4 区域中创建单主机 TPU v5e，请选择可用区 us-west4-a，并使用以 ct5lp- 开头的机器类型，例如 ct5lp-hightpu-1t。如需在前面表格中列出的其他区域创建单主机 TPU v5e，请使用以 ct5l- 开头的机器类型（例如 ct5l-hightpu-1t、ct5l-hightpu-4t 或 ct5l-hightpu-8t）。请注意，以 ct5l- 开头的机器类型需要的配额与以 ct5lp- 开头的机器类型不同。↩

选择拓扑

确定 TPU 版本后，请选择该 TPU 类型支持的拓扑。拓扑是二维或三维拓扑，具体取决于 TPU 类型。模型的并行性要求可帮助您确定拓扑。您可以通过计算拓扑中各个尺寸的乘积来确定切片中的 TPU 芯片数量。例如：

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

如果特定拓扑同时支持单主机和多主机 TPU 切片节点，则工作负载请求的 TPU 芯片的数量决定主机类型。

例如，TPU v5e (tpu-v5-lite-podslice) 支持 2x4 拓扑同时作为单主机和多主机切片。如果您：

  • 在工作负载中请求 4 个芯片，您会获得一个具有 4 个 TPU 芯片的多主机节点。
  • 在工作负载中请求 8 个芯片，您会获得一个具有 8 个 TPU 芯片的单主机节点。

请使用下表为您的用例选择 TPU 机器类型和拓扑：

  • 对于小规模模型训练或推理，请使用 TPU v4 或 TPU v5e 和单主机 TPU 切片节点池。

  • 对于大规模模型训练或推理，请将 TPU v4 或 TPU v5e 与多主机 TPU 切片节点池搭配使用。

Autopilot

TPU 版本 机器类型 拓扑 切片中的 TPU 芯片数量 节点数 节点池类型
TPU v5p tpu-v5p-slice 2x2x1 4 1 单个主机
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-podslice2 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-podslice2 2x2x1 4 1 单个主机
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. 计算方法是用拓扑产品除以四。 ↩

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

    • 如果芯片数量超过 64 个，{A}、{B} 和 {C} 必须是 4 的倍数
    • 最大拓扑为 16x16x24
    • 值必须是 {A}≤{B}≤{C}，例如 8x12x16。

  2. 不支持自定义拓扑。

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

Standard

TPU 版本 机器类型 拓扑 TPU 芯片数量 虚拟机数量 节点池类型
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 多主机
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 多主机

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

高级配置

以下部分介绍了高级 TPU 配置的调度最佳实践。

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 切片节点池，请参阅创建节点池。

Standard 集群的 CPU

本部分不适用于 Autopilot 集群，因为 GKE 会将每个 TPU 切片放置在自己的节点上。如需了解详情，请参阅 TPU 在 Autopilot 模式下的工作原理。

对于标准集群，请考虑采用以下调度最佳实践。

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

Kubernetes 资源管理和优先级处理 TPU 中的虚拟机的方式与此类其他类型的虚拟机相同。如需让需要 TPU 的 Pod 在同一节点上的其他 Pod 中获得更高的调度优先级，请为这些 TPU 切片请求最大 CPU 或内存。低优先级 TPU 分片应执行以下操作：

  1. 设置较低的 CPU 和内存请求，确保节点有足够的可分配资源来运行 TPU 工作负载。如需了解详情，请参阅 Kubernetes 如何应用资源请求和限制。
  2. 不设置 CPU 限制（无限制），以确保 Pod 可以激增以使用所有未使用的周期。
  3. 设置适当的内存限制，以确保 Pod 能够正常运行，而不会面临节点压力驱逐的风险。

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

如需了解更多最佳实践，请参阅 Kubernetes 最佳实践：资源请求和限制。

减少工作负载中断

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

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

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

托管 TPU 的 GKE 节点受维护事件或其他可能导致节点关停的中断的影响。在控制平面运行 1.29.1-gke.1425000 及更高版本的 GKE 集群中，您可以将 GKE 配置为正常终止工作负载，从而减少对工作负载的中断。

如需了解、配置和监控运行 AI/ML 工作负载的 GKE 节点上可能发生的中断事件，请参阅管理 GPU 和 TPU 的 GKE 节点中断情况。

最大限度地提高 TPU 利用率

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

后续步骤