本页面介绍了如何规划 Google Kubernetes Engine (GKE) 中的张量处理单元 (TPU) 用量,以降低 TPU 配置错误、不可用错误或配额超出中断的风险。
在 GKE 中使用 TPU 之前,请确保您熟悉 GKE 中的 TPU 定义和术语。
规划 TPU 配置
如需在 GKE 集群中使用 TPU,您必须规划其配置。我们建议您按以下步骤操作:
选择 GKE 操作模式:在 GKE Autopilot 或 Standard 集群中的 TPU 上运行工作负载。
最佳实践: 使用 Autopilot 集群可获得全托管式 Kubernetes 体验。
选择 TPU 版本:不同的 TPU 类型具有不同的功能,例如性价比、训练吞吐量和服务延迟时间。TPU 类型会影响可用的 CPU 和内存容量。
验证 TPU 可用性:TPU 可在特定 Google Cloud区域使用。如需在 GKE 工作负载中使用 TPU 类型,您的集群必须位于该类型支持的区域。
选择 TPU 拓扑:TPU 在 TPU 切片中的物理排列方式。选择符合模型并行性要求的拓扑。
使用本页面上的参考表来确定您的节点池是单主机 TPU 切片节点还是多主机 TPU 切片节点。
选择 GKE 操作模式
您可以在 GKE 为集群提供的操作模式中使用 TPU:
- Standard 模式:您负责管理底层基础架构,包括配置各个节点。
- Autopilot 模式(推荐):GKE 管理底层基础架构,例如节点配置、自动扩缩、自动升级、基准安全配置和基准网络配置。 在 Autopilot 中,您需要选择 TPU 类型和拓扑,然后在 Kubernetes 清单中指定这些内容。GKE 会为节点预配 TPU 并调度工作负载进行管理。
如需选择最适合您的工作负载的 GKE 操作模式,请参阅选择 GKE 操作模式。
选择 TPU 使用情况选项
在 GKE 中规划 TPU 配置时,请选择符合工作负载需求的消耗选项。您选择的消费选项会影响可用的 TPU 版本以及需要配置的配额。GKE 提供以下 TPU 使用情况选项,可帮助您在保持工作负载性能的同时优化资源分配和费用:
- 灵活启动:预配最多可运行 7 天的灵活启动虚拟机,GKE 会根据可用性尽力自动分配硬件。如需了解详情,请参阅使用灵活启动预配模式预配 GPU、TPU 和 H4D 简介。
- Spot 虚拟机:如需预配 Spot 虚拟机,您可以享受大幅折扣,但 Spot 虚拟机随时可能被抢占,并且会提前 30 秒发出警告。如需了解详情,请参阅 Spot 虚拟机。
- 未来预留(日历模式):在指定时间段内预配 TPU 资源,预留时间最长为 90 天。如需了解详情,请参阅在日历模式下请求具有未来预留的 TPU。
- TPU 预留:请求一年或更长时间的未来预留。
- 按需:无需提前安排容量即可使用 TPU。在请求资源之前,您必须拥有足够的按需配额,以用于特定类型和数量的 TPU 虚拟机。按需是最灵活的使用选项;不过,我们无法保证有足够的按需资源可用于满足您的请求。
如果您未指定其他选项,按需是 GKE 中 TPU 的默认使用模式。 如需选择满足工作负载要求的消费选项,请参阅 GKE 中 AI/机器学习工作负载的加速器使用选项简介。
选择 TPU 版本
TPU 切片中的虚拟机具有以下技术特征。
标准
| TPU 版本 | 机器类型 | cloud.google.com/gke-tpu-accelerator |
vCPU 的数量 | 每个虚拟机的芯片数量 | NUMA 节点的数量 | 被抢占的可能性 |
|---|---|---|---|---|---|---|
| Ironwood (TPU7x) | tpu7x-standard-4t |
tpu7x |
224 | 4 | 2 | 不适用 |
| TPU Trillium (v6e) | ct6e-standard-1t |
tpu-v6e-slice |
44 | 1 | 1 | 较高 |
| TPU Trillium (v6e) | ct6e-standard-4t |
tpu-v6e-slice |
180 | 4 | 1 | 中 |
| TPU Trillium (v6e) | ct6e-standard-8t |
tpu-v6e-slice |
180 | 8 | 2 | 较低 |
| TPU v5p |
ct5p-hightpu-4t |
tpu-v5p-slice |
208 | 4 | 2 | 不适用 |
| TPU v5e |
ct5lp-hightpu-1t |
tpu-v5-lite-podslice |
24 | 1 | 1 | 较高 |
| TPU v5e |
ct5lp-hightpu-4t |
tpu-v5-lite-podslice |
112 | 4 | 1 | 中 |
| TPU v5e |
ct5lp-hightpu-8t |
tpu-v5-lite-podslice |
224 | 8 | 1 | 低 |
| TPU v4 |
ct4p-hightpu-4t |
tpu-v4-podslice |
240 | 4 | 2 | 不适用 |
| TPU v3(仅限单主机) |
ct3-hightpu-4t |
tpu-v3-device |
96 | 4 | 2 | 不适用 |
| TPU v3 |
ct3p-hightpu-4t |
tpu-v3-slice |
48 | 4 | 1 | 不适用 |
多主机 ct5lp- 机器类型更适合处理大型模型或训练。多主机 ct5lp- 机器通过高速链路进行互连。
Autopilot
| TPU 版本 | cloud.google.com/gke-tpu-accelerator |
vCPU 的数量 | NUMA 节点的数量 | TPU 切片节点中的 TPU 芯片数量上限 |
|---|---|---|---|---|
| Ironwood (TPU7x) | tpu7x |
224 | 2 | 2048 |
| TPU Trillium (v6e) | tpu-v6e-slice |
44 - 180 | 1 到 2 | 256 |
| TPU v5p |
tpu-v5p-slice |
208 | 2 | 6144 |
| TPU v5e |
tpu-v5-lite-podslice |
24 到 224 | 1 | 256 |
| TPU v4 |
tpu-v4-podslice |
240 | 2 | 4096 |
| TPU v3(仅限单主机) |
tpu-v3-device |
96 | 2 | 8 |
| TPU v3 |
tpu-v3-slice |
48 | 1 | 256 |
请查看 Cloud TPU 价格文档中的 TPU 规格和价格,以决定要使用的 TPU 配置。
限制
在选择要使用的 TPU 时,请考虑以下限制:
Ironwood (TPU7x) 具有以下限制:
- 1.34.0-gke.2201000 版及更高版本的 Standard 集群。
- 1.34.1-gke.3084001 版及更高版本的 Autopilot 集群。
- 仅支持 Google Cloud Hyperdisk Balanced 和 Hyperdisk ML。
TPU Trillium 提供以下版本:
- 1.31.1-gke.1846000 版及更高版本的 Standard 集群。
- 1.31.2-gke.1115000 版及更高版本的 Autopilot 集群。
TPU Trillium 不支持在
ct6e-standard-8t上将 SMT 设置为2。控制平面至少运行 1.29.2-gke.1035000 或 1.28.7-gke.1020000 的 GKE 集群支持 TPU v5p 自动扩缩。
对于容量预留,请使用特定预留。
单个 TPU 虚拟机最多可以运行 256 个 Pod。
GKE 费用分配和用量计量不包括有关 TPU 用量或费用的任何数据。
集群自动扩缩器会取消等待状态超过 10 小时的 TPU 节点池纵向扩容操作。当资源可用时,集群自动扩缩器会重试此类纵向扩容操作。如果您不使用预留,此行为可能会降低 TPU 可获取性。
不支持 Ubuntu 节点。
TPU 节点架构已弃用。TPU v3 是唯一仍支持 GKE 中 TPU 节点架构的 TPU 版本。
验证 GKE 中的 TPU 可用性
TPU 可在特定 Google Cloud 区域使用。如需在 GKE 集群中使用 TPU 类型,您的集群必须位于该类型支持的区域。
标准
| TPU 版本 | 机器类型开头为 | GKE 最低版本 | 可用性 | 可用区 |
|---|---|---|---|---|
| TPU Ironwood (TPU7x) |
tpu7x-standard-4t
|
1.34.0-gke.2201000 | GA |
|
| TPU Trillium (v6e) |
ct6e-
|
1.31.2-gke.1115000 | GA |
|
| TPU v5e |
ct5lp-
|
1.27.2-gke.2100 | GA |
|
| TPU v5p |
ct5p-
|
1.28.3-gke.1024000 | GA |
|
| TPU v4 |
ct4p-
|
1.26.1-gke.1500 | GA |
|
| TPU v3 |
ct3p-
|
1.31.1-gke.1146000 | GA |
|
| TPU v3 |
ct3-
|
1.31.0-gke.1500 | GA |
|
Autopilot
| TPU 版本 |
cloud.google.com/gke-tpu-accelerator
|
GKE 最低版本 | 可用性 | 可用区 |
|---|---|---|---|---|
| TPU Ironwood (TPU7x) |
tpu7x
|
1.34.1-gke.3084001 | GA |
|
| TPU Trillium (v6e) |
tpu-v6e-slice
|
1.31.2-gke.1384000 | GA |
|
| TPU v5e |
tpu-v5-lite-podslice
|
1.27.2-gke.2100 | GA |
|
| TPU v5p |
tpu-v5p-slice
|
1.28.3-gke.1024000 | GA |
|
| TPU v4 |
tpu-v4-podslice
|
1.26.1-gke.1500 | GA |
|
| TPU v3 |
tpu-v3-slice
|
1.31.1-gke.1146000 | GA |
|
| TPU v3 |
tpu-v3-device
|
1.31.0-gke.1500 | GA |
|
选择拓扑
选择 TPU 版本后,请选择一个支持的拓扑。拓扑定义了 TPU 切片内 TPU 芯片的物理排列方式。更大的拓扑可提供更多 TPU 芯片,从而实现更并行化的处理,以便更快地训练大型模型或使用更大的数据集。
虽然 GKE 会自动管理虚拟机预配,但了解单主机节点池和多主机节点池之间的区别有助于明确拓扑中的芯片数量与底层虚拟机之间的关系:
- 单主机:所有芯片都位于单个虚拟机上的 TPU 切片。这种玩法涉及的筹码数量较少,通常为 4 个或更少。
- 多主机:芯片分布在多个虚拟机中的 TPU 切片。这对于大多数大规模 TPU 工作负载来说很常见。
如果您请求的拓扑的总芯片数超过了单个虚拟机上可用的芯片数(对于相应 TPU 版本而言),GKE 会自动将其配置为多主机环境。在此场景中,GKE 会启动多个节点来分配芯片。
例如,假设机器类型为 ct6e-standard-4t,拓扑为 4x4:
- 机器类型
ct6e-standard-4t为每个虚拟机配置了 4 个芯片。 - 拓扑
4x4总共需要 16 个芯片(4 * 4)。 - 由于 16(拓扑芯片数量)大于 4(每个虚拟机的芯片数量),因此此配置会生成多主机 TPU 切片节点池。GKE 会将 16 个芯片分配到 4 个虚拟机中。
请使用下表为您的用例选择 TPU 机器类型和拓扑:
标准
选择 TPU 类型和拓扑后,请在工作负载清单中指定这些内容。如需了解相关说明,请参阅在 GKE Standard 上部署 TPU 工作负载。
| TPU 版本 | 机器类型 | 节点池类型 | 技术规范 |
|---|---|---|---|
| Ironwood (TPU7x) | tpu7x-standard-4t |
单个主机 |
|
| Ironwood (TPU7x) | tpu7x-standard-4t |
多主机 |
|
| Ironwood (TPU7x) | tpu7x-standard-4t |
多主机 |
|
| Ironwood (TPU7x) | tpu7x-standard-4t |
多主机 |
|
| Ironwood (TPU7x) | tpu7x-standard-4t |
多主机 |
|
| Ironwood (TPU7x) | tpu7x-standard-4t |
多主机 |
|
| Ironwood (TPU7x) | tpu7x-standard-4t |
多主机 |
|
| Ironwood (TPU7x) | tpu7x-standard-4t |
多主机 |
|
| Ironwood (TPU7x) | tpu7x-standard-4t |
多主机 |
|
| Ironwood (TPU7x) | tpu7x-standard-4t |
多主机 |
|
| TPU Trillium (v6e) | ct6e-standard-1t |
单个主机 |
|
| TPU Trillium (v6e) | ct6e-standard-8t |
单个主机 |
|
| TPU Trillium (v6e) | ct6e-standard-4t |
单个主机 |
|
| TPU Trillium (v6e) | ct6e-standard-4t |
多主机 |
|
| TPU Trillium (v6e) | ct6e-standard-4t |
多主机 |
|
| TPU Trillium (v6e) | ct6e-standard-4t |
多主机 |
|
| TPU Trillium (v6e) | ct6e-standard-4t |
多主机 |
|
| TPU Trillium (v6e) | ct6e-standard-4t |
多主机 |
|
| TPU Trillium (v6e) | ct6e-standard-4t |
多主机 |
|
| TPU v5p | ct5p-hightpu-4t |
单个主机 |
|
| TPU v5p | ct5p-hightpu-4t |
多主机 |
|
| TPU v5p | ct5p-hightpu-4t |
多主机 |
|
| TPU v5p | ct5p-hightpu-4t |
多主机 |
|
| TPU v5p | ct5p-hightpu-4t |
多主机 |
|
| TPU v5e | ct5lp-hightpu-1t |
单个主机 |
|
| TPU v5e | ct5lp-hightpu-4t |
单个主机 |
|
| TPU v5e | ct5lp-hightpu-8t |
单个主机 |
|
| TPU v5e | ct5lp-hightpu-4t |
多主机 |
|
| TPU v5e | ct5lp-hightpu-4t |
多主机 |
|
| TPU v5e | ct5lp-hightpu-4t |
多主机 |
|
| TPU v5e | ct5lp-hightpu-4t |
多主机 |
|
| TPU v5e | ct5lp-hightpu-4t |
多主机 |
|
| TPU v5e | ct5p-hightpu-4t |
多主机 |
|
| TPU v5e | ct5p-hightpu-4t |
单个主机 |
|
| TPU v4 | ct4p-hightpu-4t |
多主机 |
|
| TPU v4 | ct4p-hightpu-4t |
多主机 |
|
| TPU v4 | ct4p-hightpu-4t |
多主机 |
|
| TPU v4 | ct4p-hightpu-4t |
多主机 |
|
| TPU v3 | ct3-hightpu-4t |
单个主机 |
|
| TPU v3 | ct3p-hightpu-4t |
多主机 |
|
| TPU v3 | ct3p-hightpu-4t |
多主机 |
|
| TPU v3 | ct3p-hightpu-4t |
多主机 |
|
| TPU v3 | ct3p-hightpu-4t |
多主机 |
|
| TPU v3 | ct3p-hightpu-4t |
多主机 |
|
| TPU v3 | ct3p-hightpu-4t |
多主机 |
|
| TPU v3 | ct3p-hightpu-4t |
多主机 |
|
-
计算方法是用拓扑产品除以四。↩
Autopilot
选择 TPU 类型和拓扑后,请在工作负载清单中指定这些内容。如需了解相关说明,请参阅在 GKE Autopilot 上部署 TPU 工作负载。
| TPU 版本 | 机器类型 | 节点池类型 | 技术规范 |
|---|---|---|---|
| Ironwood (TPU7x) | tpu7x |
单个主机 |
|
| Ironwood (TPU7x) | tpu7x |
多主机 |
|
| Ironwood (TPU7x) | tpu7x |
多主机 |
|
| Ironwood (TPU7x) | tpu7x |
多主机 |
|
| Ironwood (TPU7x) | tpu7x |
多主机 |
|
| Ironwood (TPU7x) | tpu7x |
多主机 |
|
| Ironwood (TPU7x) | tpu7x |
多主机 |
|
| TPU Trillium (v6e) | tpu-v6e-slice |
单个主机 |
|
| TPU Trillium (v6e) | tpu-v6e-slice |
单个主机 |
|
| TPU Trillium (v6e) | tpu-v6e-slice |
单个主机 |
|
| TPU Trillium (v6e) | tpu-v6e-slice |
多主机 |
|
| TPU Trillium (v6e) | tpu-v6e-slice |
多主机 |
|
| TPU Trillium (v6e) | tpu-v6e-slice |
多主机 |
|
| TPU Trillium (v6e) | tpu-v6e-slice |
多主机 |
|
| TPU Trillium (v6e) | tpu-v6e-slice |
多主机 |
|
| TPU v5p | tpu-v5p-slice |
单个主机 |
|
| TPU v5p | tpu-v5p-slice |
多主机 |
|
| TPU v5p | tpu-v5p-slice |
多主机 |
|
| TPU v5p | tpu-v5p-slice |
多主机 |
|
| TPU v5p | tpu-v5p-slice |
多主机 |
|
| TPU v5p | tpu-v5p-slice |
多主机 |
|
| TPU v5e | tpu-v5-lite-podslice |
单个主机 |
|
| TPU v5e | tpu-v5-lite-podslice |
单个主机 |
|
| TPU v5e | tpu-v5-lite-podslice |
单个主机 |
|
| TPU v5e | tpu-v5-lite-podslice |
多主机 |
|
| TPU v5e | tpu-v5-lite-podslice |
多主机 |
|
| TPU v5e | tpu-v5-lite-podslice |
多主机 |
|
| TPU v5e | tpu-v5-lite-podslice |
多主机 |
|
| TPU v5e | tpu-v5-lite-podslice |
多主机 |
|
| TPU v5e | tpu-v5-lite-podslice |
多主机 |
|
| TPU v4 | tpu-v4-podslice |
单个主机 |
|
| TPU v4 | tpu-v4-podslice |
多主机 |
|
| TPU v4 | tpu-v4-podslice |
多主机 |
|
| TPU v4 | tpu-v4-podslice |
多主机 |
|
| TPU v4 | tpu-v4-podslice |
多主机 |
|
| TPU v4 | tpu-v4-podslice |
多主机 |
|
| TPU v3 | tpu-v3-slice |
多主机 |
|
| TPU v3 | tpu-v3-slice |
多主机 |
|
| TPU v3 | tpu-v3-slice |
多主机 |
|
| TPU v3 | tpu-v3-slice |
多主机 |
|
| TPU v3 | tpu-v3-slice |
多主机 |
|
| TPU v3 | tpu-v3-device |
单个主机 |
|
-
计算方法是用拓扑产品除以四。 ↩
支持超过 64 个芯片的自定义拓扑。以下条件适用:
- 如果芯片数量超过 64 个,
{A}、{B}和{C}必须是 4 的倍数 - 最大拓扑为
16x16x24 - 值必须是
{A}≤{B}≤{C},例如8x12x16。
- 如果芯片数量超过 64 个,
-
不支持自定义拓扑。
高级配置
以下部分介绍了高级 TPU 配置的调度最佳实践。
AI 可用区
AI 可用区是专门用于 AI/机器学习训练和推理工作负载的可用区。这些可用区可提供充足的机器学习加速器容量。如需了解详情,请参阅 AI 可用区文档。
在 GKE 中使用 AI 加速区之前,请考虑以下特征:
- AI 可用区在物理上与标准可用区分开,可提供额外的存储空间和电力。这种分离可能会导致延迟时间更长,但对于 AI/机器学习工作负载而言,这通常是可以接受的。
- AI 可用区带有
ai注释的后缀。例如,us-central1区域中的 AI 可用区名为us-central1-ai1a。 - 目前仅支持 TPU 虚拟机。
- 集群的控制平面在与 AI 可用区位于同一区域内的一个或多个标准可用区中运行。
只有在满足以下要求的情况下,您才能在 AI 可用区中运行未连接 TPU 的虚拟机:
- 您已在同一可用区中运行其他使用 TPU 虚拟机的工作负载。
- 非 TPU 虚拟机可以是 Spot 虚拟机、与预留相关联的虚拟机,也可以是具有特定加速器与通用虚拟机比率的节点池的一部分。
AI 可用区与同一区域内具有相同后缀的标准可用区共享组件,例如网络连接和软件发布。对于高可用性工作负载,我们建议您使用不同的可用区。例如,避免同时使用
us-central1-ai1a和us-central1-a来实现高可用性。
默认情况下,GKE 不会将工作负载部署在 AI 可用区中。如需使用 AI 区域,您必须配置以下选项之一:
- (推荐)ComputeClass:设置最高优先级,以便在 AI 可用区中请求按需 TPU。借助 ComputeClass,您可以为工作负载定义硬件配置的优先级列表。如需查看示例,请参阅 ComputeClass 简介。
- 节点自动预配:在 Pod 规范中使用
nodeSelector或nodeAffinity,指示节点自动预配功能在 AI 可用区中创建节点池。如果工作负载未明确指定 AI 可用区,则节点自动预配功能在创建新节点池时只会考虑标准可用区或--autoprovisioning-locations中的可用区。此配置有助于确保不运行 AI/机器学习模型的工作负载保留在标准可用区中,除非您明确配置了其他可用区。如需查看使用nodeSelector的清单示例,请参阅设置自动创建的节点的默认可用区。 - GKE Standard:如果您直接管理节点池,请在创建节点池时使用
--node-locations标志指定 AI 地区。如需查看示例,请参阅在 GKE Standard 中部署 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 切片节点池,请参阅创建节点池。
为 TPU 切片预配额外存储空间
TPU 切片中的虚拟机包含一个 100 GiB 启动磁盘。如果您的 TPU 切片需要额外存储空间用于训练或预处理,或者您需要保存检查点,您可以使用 Google Cloud Hyperdisk 或平衡永久性磁盘存储空间(如果适用于您的 TPU)。如需详细了解每个 TPU 版本支持的磁盘类型,请参阅 TPU 对 Hyperdisk 和永久性磁盘的支持。
Standard 集群的 CPU
本部分不适用于 Autopilot 集群,因为 GKE 会将每个 TPU 切片放置在自己的节点上。如需了解详情,请参阅 TPU 在 Autopilot 模式下的工作原理。
对于标准集群,请考虑采用以下调度最佳实践。
如需在 TPU 切片节点中的虚拟机上调度非 TPU 工作负载,请确保您的 GKE Pod 可以容忍 google.com/tpu 污点。如果您希望将工作负载部署到特定节点,请使用节点选择器。
Kubernetes 资源管理和优先级处理 TPU 中的虚拟机的方式与此类其他类型的虚拟机相同。如需让需要 TPU 的 Pod 在同一节点上的其他 Pod 中获得更高的调度优先级,请为这些 TPU 切片请求最大 CPU 或内存。低优先级 TPU 分片应执行以下操作:
- 设置较低的 CPU 和内存请求,确保节点有足够的可分配资源来运行 TPU 工作负载。如需了解详情,请参阅 Kubernetes 如何应用资源请求和限制。
- 不设置 CPU 限制(无限制),以确保 Pod 可以激增以使用所有未使用的周期。
- 设置适当的内存限制,以确保 Pod 能够正常运行,而不会面临节点压力驱逐的风险。
如果 Kubernetes Pod 未请求 CPU 和内存(即使正在请求 TPU),则 Kubernetes 会将其视为尽力而为的 Pod,在它需要任何 CPU 和内存时无法保证会提供。只有明确请求 CPU 和内存的 Pod 才能获得此类保证。对于特定的 Kubernetes 调度,请使用明确的 CPU 和内存请求来配置 Pod 需求。如需了解详情,请参阅 Pod 和容器的资源管理。
如需了解更多最佳实践,请参阅 Kubernetes 最佳实践:资源请求和限制。
减少工作负载中断
如果您使用 TPU 训练机器学习模型且工作负载中断,则自上次检查点以来执行的所有工作都将丢失。如需降低工作负载中断的可能性,请执行以下操作:
- 为此作业设置比所有其他作业更高的优先级:如果资源不足,则 GKE 调度器会抢占优先级较低的作业,以调度优先级更高的作业。这还有助于确保优先级较高的工作负载获得所需的所有资源(最多为集群中可用的总资源)。如需了解详情,请参阅 Pod 优先级和抢占。
- 配置维护排除项:维护排除项是禁止进行自动维护的非重复性时间段。如需了解详情,请参阅维护排除项。
- 在 Autopilot 中使用延长运行时间的 Pod:在长达七天宽限期内使用延长运行时间的 Pod,然后 GKE 终止 Pod 进行缩容或节点升级。
- 在 TPU Trillium 中使用集合调度:使用集合来指明 TPU 切片节点池是服务工作负载的一部分。 Google Cloud 限制并简化对推理工作负载操作的干扰。如需了解详情,请参阅集合调度的工作原理。
这些建议有助于最大限度减少中断,但不会阻止中断。例如,由于硬件故障造成的抢占或因碎片化造成的抢占仍可能发生。同样,设置 GKE 维护排除项不会阻止 Compute Engine 维护事件。
经常保存检查点,并将代码添加到训练脚本中,以便在继续时从上一个检查点开始。
处理由于节点维护而导致的中断
托管 TPU 的 GKE 节点受维护事件或其他可能导致节点关停的中断的影响。在控制平面运行 1.29.1-gke.1425000 及更高版本的 GKE 集群中,您可以将 GKE 配置为正常终止工作负载,从而减少对工作负载的中断。
如需了解、配置和监控运行 AI/ML 工作负载的 GKE 节点上可能发生的中断事件,请参阅管理 GPU 和 TPU 的 GKE 节点中断情况。
最大限度地提高 TPU 利用率
为了最大限度地提高对 TPU 的投资,请混合调度作业优先级,并将其排入队列,以最大限度增加 TPU 的运行时间。如需实现作业级调度和抢占,您需要使用 Kubernetes 的一个插件来将作业编排到队列中。
使用 Kueue 将作业编排到队列中。
后续步骤
- 按照在 GKE 中部署 TPU 工作负载中的说明设置 GKE 上的 Cloud TPU。
- 了解将 Cloud TPU 用于机器学习任务的最佳做法。
- 利用 GKE 在 Cloud TPU 上构建大规模的机器学习服务。
- 在 TPU 上使用 KubeRay 提供大语言模型。