人工智能工作量的资源和基础设施挑战

核心观点与关键数据

本文探讨了云原生AI工作负载调度中涉及的资源与基础设施挑战，主要包括拓扑感知、资源异构性、GPU利用与共享、可扩展性、I/O瓶颈、容错性与弹性，以及预算与成本约束等方面。

拓扑感知

• 节点级拓扑：NUMA架构下，将工作负载的线程和内存放置在同一NUMA节点上可避免跨插槽访问的延迟。
• 簇级拓扑结构：网络架构对分布式训练性能影响显著，同一机架顶层的交换机下的节点通信速度远快于跨数据中心的节点。
• 重要性：对于需要大量通信的分布式训练任务，节点放置直接影响任务完成时间，而默认Kubernetes调度器对GPU互连或网络基础设施感知有限。

资源异构性

• 问题：集群中存在混合硬件，包括不同型号、代次、内存大小的GPU，以及TPU、FPGA和CPU等。
• 解决方案：通过节点池将相似硬件节点分组，并使用资源标签标记节点能力（如gpu-type=高内存，gpu-generation=当前世代）。
• 挑战：工作负载可能存在偏好而非硬性要求，如“优先考虑当前代GPU，但也可在上一代GPU上运行”，表达和处理这些偏好增加了复杂性。

GPU利用率与共享

• 问题：GPU昂贵但利用率低，推理服务或交互式会话可能仅使用GPU计算能力的20%。
• 解决方案：
  • GPU分区：将物理GPU分割为多个虚拟GPU，每个虚拟GPU拥有独立内存和计算能力。
  • 时间分片：快速在多个工作负载间切换共享GPU，提高利用率但增加延迟，无内存隔离。
  • MPS（多进程服务）：允许多个容器同时使用一个GPU，提高利用率但无硬件隔离。
  • vGPU（虚拟GPU）：软件级别共享GPU，适用于虚拟化环境。
  • 动态资源分配（DRA）：Kubernetes功能，支持按需GPU分区，需供应商DRA驱动实现（GA于1.34版本）。
• 注意事项：大型语言模型会消耗全部GPU内存，内存共享不适用于大型模型推理。

可扩展性

• 问题：大规模AI训练需要数万GPU，调度系统需处理此类规模。
• 挑战：
  • 集群自动扩展：如Karpenter和Kubernetes集群自动扩展器，但添加GPU节点需几分钟，可能无法满足大型作业需求。
  • 负载自动扩展：弹性训练和数据处理作业可动态调整规模，推理按请求负载驱动，需监控每秒令牌数、队列深度等指标。
  • 调度器吞吐量：调度器需在数万Pod上高效决策，避免成为瓶颈。
  • 协调开销：大规模集群调度比独立Pod调度更复杂。

I/O瓶颈

• 问题：GPU速度快但数据传输慢，I/O瓶颈浪费GPU时间。
• 数据加载：训练中数据加载占执行时间的10-70%，GPU在此期间空闲。
• 检查点（Checkpointing）：训练任务周期性保存模型状态，保存大型模型需数分钟，需权衡频率与性能。
• 推理模型加载：
  • 冷启动：大型模型加载需数分钟，可使用Run:ai模型流式传输器等工具优化。
  • 存储性能：MLCommons存储基准测试显示，存储吞吐量不足会卡住训练流程。

容错性与弹性

• 问题：长时间训练任务（数天至数周）易遇故障，如硬件故障、节点宕机、抢占。
• 解决方案：
  • 检查点/重启：定期保存模型状态，从最后一个检查点恢复，权衡检查点频率与开销。
  • 弹性任务：如PyTorch Elastic支持动态调整工作者数量，失败时继续训练但减少工作者，需支持工人数范围、角色感知调度、任务重新分配等功能。
• 重要性：Meta报告称Llama 3预训练54天期间经历466次任务中断，其中47次为计划内，其余为计划外故障。

预算和成本限制

• 问题：GPU资源昂贵，失控任务可能消耗大量预算。
• 解决方案：
  • GPU小时预算：限制团队或作业可消耗的GPU时间。
  • 成本感知调度：考虑不同资源分配成本（按量付费与预留实例、不同GPU类型）。
  • 自动终止：预算用尽时自动终止作业。
  • 成本可见性：帮助团队理解并优化支出。
• 额外考量：能源预算在能源密集型AI工作负载中同样重要，可利用遥测数据优化调度决策。

基础设施考量

• GPU直接通信：
  • 机箱内：高速GPU互连提供高带宽、低延迟连接（如NVLink），调度器需理解拓扑结构。
  • 服务器间：InfiniBand（支持RDMA）和RoCE（以太网上的RDMA）提供高带宽、低延迟连接，GPUDirect RDMA实现GPU间直接数据传输。
• AI工作负载的网络拓扑：
  • 地域因素：频繁通信的工作者应放置在网络相邻位置（如同一叶交换机）。
  • 带宽需求：大规模分布式训练需避免网络链路饱和。
  • 故障域：集中部署可减少网络跳数但降低容错性，分散部署提高容错性但增加延迟。
  • 权衡：延迟敏感型训练可能牺牲部分容错能力，长时间任务可接受更高延迟以避免重启。

研究结论

本文系统梳理了AI工作负载调度中的资源与基础设施挑战，并提出了相应的解决方案。关键结论包括：

1. 拓扑感知对分布式训练性能至关重要，但默认Kubernetes调度器支持有限。
2. GPU资源需通过分区、共享等技术提高利用率，但需注意内存限制。
3. 大规模训练需考虑集群和负载自动扩展，但调度器吞吐量和协调开销是挑战。
4. I/O瓶颈（数据加载、检查点、模型加载）需通过优化工具和存储性能解决。
5. 容错性通过检查点和弹性任务实现，但需支持角色感知调度和任务重新分配。
6. 预算和成本限制需通过预算控制、成本感知调度和能源管理实现。
7. 基础设施（GPU直接通信、网络拓扑）对调度决策影响显著，需权衡地域、带宽和故障域。

下一步

本系列最终论文将梳理应对这些挑战的工具和Kubernetes功能，提供解决方案参考表及实践指南。