
在 Allen Institute for AI(Ai2)的 AI 基础设施团队,我们负责为研究所提供 GPU 计算能力,主要支撑大规模分布式训练工作负载。我们将这一使命构建为由四个相互关联指标组成的金字塔:底层是可用性(硬件健康且就绪的频率),其上是占用率(分配给特定工作负载的可用时间比例),接着是影响力(高价值工作负载获取资源的频率),顶端则是利用率(工作负载生命周期内 GPU 容量的实际使用比例)。
本文重点探讨如何提升调度决策的“影响力”。近期,我们以一套包含 GPU 时间预算、分层公平共享分配和时间切片合约的新系统,取代了传统的基于优先级的调度器。这一变革将关于“每个研究项目应得多少 GPU 时间”的争论,从逐个处理的运营琐事,转变为透明的行政预算流程。
过度承诺与“公地悲剧”
Ai2 管理着数千块 NVIDIA H100、B200 和 B300 GPU,这些算力被组织成规模从 88 到 1024 个节点不等的集群,服务于约 150 名内部研究人员。他们的研究涵盖大语言模型(LLM)、多模态视觉语言模型(VLM)的全流程训练、机器人强化学习仿真以及科学智能体后训练等多个领域。
与许多实验室一样,我们的 GPU 需求远超供给。任何时刻,待处理请求所需的 GPU 数量均为可用数量的 2-3 倍。这意味着集群中每一个可用的 GPU 小时,都面临着来自两三个不同研究工作负载的竞争。
历史上,我们使用基于优先级的调度器,并允许工作负载选择退出抢占机制。每个团队拥有并发 GPU 使用上限,受保护的工作负载不会因抢占而中断,而可抢占负载则可在空闲 GPU 上超出该限制。这种策略引发了可预测的问题:
- GPU “霸占”:用户停放无实际操作负载,以便随时连接调试。由于无法以低延迟启动新负载,研究人员倾向于长期持有资源。
- 优先级膨胀:最终 100% 的调度工作负载均标记为“高”优先级,导致低优先级队列完全缺乏算力。
- 运维负担:由于抢占是可选的,值班工程师需花费大量时间协商关闭那些在存在维护问题的主机上运行的不可抢占负载。
我们最初试图通过更严格控制优先级或为重要项目显式分配 GPU 独占权来解决问题,但这实际上加剧了“公地悲剧”:个人为最大化自身利益争夺稀缺共享资源,导致全局结果非优化及底层资源滥用。
从调度到预算:迭代所有权模式
解决公地悲剧的经典方案是将共享资源私有化。虽然赋予团队 GPU 独占权类似于此,但这过于粗糙,导致 GPU 因研究的季节性波动而闲置。我们希望保留所有权的激励,同时保持 GPU 的高占用率。
为此,我们决定迭代所有权模式:不再向团队发放物理 GPU,而是分配GPU 时间预算。研究努力的优先级是一个战略问题,更适合由领导层像投资者一样,在工作负载存在之前,根据其潜在影响决定如何为每个研究项目“注资”。随后,调度器利用这些信息优先处理到达的工作负载。
我们设计了一个分层系统,经理可按比例分配其负责项目的 GPU 时间份额。例如,项目 A1 明确知晓其对总容量拥有 35% 的请求权,无论其他队列如何拥堵。在此系统中,每个 GPU 时间请求必须由预算资助,否则不受抢占保护。由于没有任何资源是免费的,任何“霸占”行为都将直接消耗团队预算,使得操纵调度器的成本高于诚实参与预算辩论的成本。
分层公平共享与调度合约
配合 GPU 时间预算,我们构建了一个分层公平共享调度器。该算法源自 2009 年的 Hadoop Fair Scheduler,并在 SLURM 和 YARN 中广泛应用。创新之处在于,树结构反映了研究程序的组织架构,权重由管理者设定的预算动态决定,而非静态配额。
调度器跟踪滑动回溯窗口(默认 7 天)内的占用情况,将未充分利用分配的工作负载排在充分利用者之前。这确保了在一周时间内,只要各组积极提交需求,即可获得分配的 GPU 时间。
此外,针对长周期训练任务,我们引入了“调度合约”。工作负载必须声明其最小运行时间(即取得有意义进展所需的短占用时间)。在此期间,负载免受抢占。一旦达到最小运行时间,调度者有权重新平衡,自动重新排队可恢复负载。若用户将最小运行时间设为零,则视为“未分配占用”,这类负载免费但始终面临抢占。
这一机制带来了显著的运维自动化收益:当主机出现健康问题时,系统可在工作负载达到最小运行时间时自动排空负载。这将需要人工介入的维修工作减少了 74%。
模拟验证与实际成效
在正式部署前,我们构建了模拟环境,以预测策略变更的后果。模拟结果显示,小规模“调试工作负载”的 p90 等待时间可从约 6 小时降至 5 分钟。
7 月底,我们开始逐集群部署新系统。30 天的测试期数据显示:
- 公平性:团队平均收到了 98% 的应得 GPU 小时。15 个团队中有 13 个获得了 95% 或更多的份额,最差情况也达到了 90%。
- 高占用率:集群占用率在变更前后均稳定在 98%。其中 18% 的交付时间为“未分配占用”,有效填补了资助用例尚未就绪时的算力空缺。
- 延迟降低:实际效果优于预测。调试工作负载的 p90 队列等待时间从 2 小时骤降至 30 秒。在最大的 H100 集群上,中位数队列等待时间从 5 分钟降至 24 秒,p90 等待时间下降了约三分之一(从 2.8 小时降至 1.8 小时)。
新系统有效解决了三大痛点:短调试任务的快速启动降低了“霸占”的价值;优先级仅影响团队内部排序,遏制了膨胀;自动化排空机制大幅减少了运维琐事。
挑战与未来展望
变革伴随学习曲线。初期,界面术语的滞后含义引发了一些混淆,我们通过现场解释会议和新可视化工具(展示预算使用情况及排队指标)逐步解决了这一问题。
并非所有用例都得到了改善。研究人员依赖的交互式会话因 8 小时最小运行时间限制而频繁被抢占,导致状态重建困难。为此,我们计划建立一个仅限 CPU 的集群用于数据准备和开发会话,并构建可恢复会话机制,以兼顾调度优势与用户体验。
目前,我们正关注潜在的容量碎片化问题,即最小运行时间保护可能使大型作业难以找到连续的空闲资源。我们正利用模拟器在生产环境中重现并测量这一问题。
展望未来,我们的目标直指金字塔顶端——利用率。我们将进一步优化引导、检查点和训练应用程序的效率,确保每个工作负载都能最大化其获得的调度时间价值。