构建一个集群调度器，在保持满占用率的同时优先支持高影响力研究
在Ai2的AI基础设施团队，我们负责提供研究院的GPU算力，特别针对大型分布式训练工作负载。我们把这项任务看作一个由四项相互叠加的指标构成的金字塔。
基础是 可用性：硬件健康并可供使用的频率。其上是 占用率：可用时间中被分配给特定工作负载的比例。接下来是 影响力：最有价值的工作负载被选中获得资源的频率。金字塔的顶端是 利用率：在工作负载生命周期内所使用的GPU容量比例。
这篇文章讲的是如何提升我们调度决策的影响力。我们最近用一个包含GPU时间预算、层级公平份额分配和时间切片契约的系统取代了基于优先级的调度器。这样一来，关于每个研究项目应获得多少GPU时间的争论，从逐案处理的事务性任务转变为一个透明的行政预算流程。
超售
在Ai2，我们管理着数千块NVIDIA H100、B200和B300 GPU，它们组成规模从88到1024块GPU不等的集群。这些集群专为AI模型的大规模分布式训练而建，服务于约150名内部研究人员，他们的工作覆盖多样的AI领域，包括LLM和VLM训练的完整模型流程、机器人强化学习（RL）仿真，以及面向科学智能体用例的后训练。
和许多实验室一样，我们对GPU时间的需求远超供给。根据提交的工作负载，任意时刻我们手中未完成的请求所需的GPU数量是可用数量的2-3倍。可以这样理解：集群上每一个可用的GPU小时，都有2-3个不同的研究工作负载在竞争。
历史上，我们使用基于优先级的调度器，并允许工作负载选择退出抢占机制。每个团队都有一个可被不受抢占保护的工作负载使用的并发GPU上限。可抢占的工作负载可以在空闲GPU上超出该上限。这种策略产生了可预见的病态现象。例如，我们观察到GPU“占坑”的情况，用户会挂起空转的工作负载，以便在需要时随时接入。这种情况发生的原因是研究人员发现他们无法以足够低的延迟启动调试工作负载来实时解决问题。我们还观察到优先级通胀，最终100%的调度工作负载都使用了HIGH优先级。这意味着较低的优先级层级完全得不到GPU时间。由于抢占是可选的，我们还发现值班工程师将大部分工单响应时间花在协商关闭运行在有已知维护问题主机上的不可抢占工作负载上。
公地悲剧
当这些问题出现时，我们未能及时发现其根本原因。我们最初确保最重要的工作获得GPU时间的尝试，集中于更严格地控制优先级的设定方式，并最终通过显式地将GPU垄断权分配给重要项目来绕过基于优先级的调度器。虽然起初我们没有意识到，但我们已经建立了一个观察“公地悲剧”的完美实验室。个人在争夺一种稀缺的共享资源，并且由于追求个人利益最大化，达成了非最优的全局结果并滥用了底层资源。
我们远非第一个观察到这种互动的人。资源分配是一个迷人的研究领域，融合了算法开发、经济学和系统管理。一个核心问题是，用户往往比组织更了解自己作业的价值，但他们可能有动机隐藏这种价值，或者即使这样做会损害整体性能也不愿释放资源。例如，在2011年引入 主导资源公平的论文中，Ghodsi等人讲述了一件轶事：一家搜索公司只有在用户能保证高利用率的情况下才向作业提供专用机器。他们很快发现“用户会在代码中散布无限循环来人为抬高利用率”。硬件会变化，但使资源分配变得复杂的根本问题依然存在。
要预算，不要排期
公地悲剧的经典解决方案是将共享资源私有化——所有者有动力最大化其财产的价值。当我们把GPU集合的垄断权分配给团队时，我们其实已经在做类似的事情，但粒度太粗。它导致GPU因研究的季节性而闲置。各个团队准备好运行实验和训练的时间各不相同，因此分配垄断权会导致出现没有作业准备好执行的时刻，而另一个团队却在等待算力。
我们当时在手动解决一个 背包问题，试图把动态变化的研究需求塞进静态的时间表中。我们想要所有权激励，但也想保持 GPU 的完全占用率。
我们决定对所有权模型进行迭代。我们不再向团队发放 GPU，而是选择分配一部分 GPU 时间。预测未来的需求需要知道新颖科学实验的结果，因此无法精确预测。然而，不同研究工作之间的优先级是一个战略问题，可以更容易地提前讨论和决定。我们没有去解决排班难题，而是让领导层像投资者一样思考。在工作负载出现之前，根据他们对各项工作可能产生的影响的判断，决定如何用 GPU 时间为每项研究工作提供资金。调度器随后可以在为到达的工作负载确定优先级时使用这些信息。
考虑到这一点，我们设计了一个分层系统，管理者可以在其中按比例向他们负责的项目和研究人员分配 GPU 时间。如下图所示，这将项目战略直接转化为有保障的 GPU 时间份额。项目 A1 知道自己拥有总容量的 35% 的权利，无论其他地方还有多少项目在排队。
括号中的数值表示分配给叶节点项目的集群总容量。
在这个系统中，每一个对 GPU 时间的请求都必须由预算提供资金，否则就不受抢占保护。在旧系统中，HIGH 优先级没有成本，且不可抢占允许一个团队无限期地填满其并发 GPU 上限，所以所有人都使用了它们。现在，没有什么是免费的，因此任何获取 GPU 时间的技巧都会消耗受益用户的配额。一个占着 GPU 不用的负载是在把团队预算花在毫无用处的事情上。我们的策略是让钻调度器空子的代价比诚实地参与争取更大预算的讨论更高。我们一直在对预算评审流程进行迭代，但关键要求是：研究人员有频繁的机会为他们所需的时间进行争取，并且决策由对相关权衡最了解的管理者做出。这意味着研究项目内部的分配决策由项目负责人做出，研究计划内部由首席研究员做出，跨计划则由计划主管经理或 CEO 做出。
公平份额
与这个 GPU 时间预算工具配套，我们构建了一个分层公平份额调度器，用于管理整个计划树中分配的实际占用情况。这里的算法并不新鲜——基于时间窗口的分层公平份额是源自 2009 年的 Hadoop 公平调度器的谱系的一部分，同样的方法如今仍在 SLURM 的 Fair Tree 和 YARN 的公平调度器中被积极使用。对我们来说新颖的是输入：树结构反映了研究计划的结构，而权重是由管理者设定的预算，而非静态配额。
调度器在一个滑动的回溯窗口内（我们默认为 7 天）跟踪占用情况，并将来自利用不足分配的工作负载排在来自过度利用分配的工作负载之上。这样，在长达一周的时间范围内，我们可以期望每个团队都获得其分配的 GPU 时间，只要他们持续提交有足够需求的工作负载。
“新的调度器让我们感觉多了 30% 的算力。在旧调度器中，如果我们有时不需要满额的槽位限制，那些算力基本上就浪费了。现在有了新调度器，如果发生这种情况，我们之后可以突破分配上限进行突发使用，仍然能看到我们的作业被快速调度且不被抢占，实质上让我们收回了那些算力。我们的工作负载常常是突发性的，所以这为我们找回了大量算力。” — Chris Clark
调度器区分两类占用情况。 已分配占用 是指工作负载被计入预算的时间。这会消耗工作负载所有者的配额，影响公平份额预算计算，并且这些工作负载在其最短运行时间窗口内受抢占保护。 未分配占用 不计入任何预算，从一开始就不受保护，并且可能被任何已分配的请求抢占。这使我们即使在分配与需求不匹配的情况下也能让 GPU 保持完全占用，并防止团队拒绝免费的 GPU 周期。
调度契约
分布式训练中另一个使公平资源分配变得困难的特性是，工作负载可能运行非常长的时间。训练作业通常运行数小时、数天，有时甚至数周。一旦被调度，一个工作负载可能在其分配的 GPU 上停留一周或更久，不给其他人获得其预算时间的机会。这正是使 GPU 占用不还成为可能的系统特性。这也迫使值班工程师不得不与长时间运行作业的所有者协商，以解决持续的维护问题。
为了解决这些问题，我们引入了“调度契约”。作为使用集群的交换条件，工作负载必须声明其最短运行时间，即做出有意义进展所需的最短占用时间。在此期间，工作负载受到保护，不会被抢占。这为研究人员提供了进展保障，同时让调度器在进展已被保存后获得重新平衡的权利，自动将可恢复的工作负载重新排队。或者，用户可以将最短运行时间设为零，这表示 GPU 时间应被释放。这些工作负载始终可能被抢占，但它们也是“免费”的，因为它们不计入任何预算。
工作负载的生命周期遵循以下模式：
- 工作负载被提交，带有最短运行时间，并指明其是否可恢复。
- 工作负载根据公平份额算法进行调度，其权重为回溯窗口内实际占用时间与已分配时间之比。
- 工作负载运行其最短运行时间，该时间计入其分配额度。
- 只要相关分配额度继续将其优先于其他工作负载，工作负载就可以继续运行。这段时间同样计入其分配额度。
- 它可能被抢占并重新排队，即返回步骤 2。
- 工作负载完成，释放其对任何资源的占用。
这些约定共同为我们的调度器增加了时间分片机制。正在运行的工作负载可以被自动移除并重新排队，使公平份额得以收敛，并抑制资源囤积行为。它们还让不健康的主机能够在其工作负载达到最短运行时间后将其排空，从而使修复活动可以完全自动化。在规划这项工作时，这最后一点比我们意识到的更为重要。它将需要人工介入的修复减少了74%，这在值班负担上是一笔巨大的节省。
模拟
我们知道，调度策略的更改可能带来意想不到的后果。这个问题的零和特性意味着，给一位研究人员时间就意味着从另一位那里拿走时间。在这一交换中受损的用户往往会寻找新的变通方法。在推出基于预算的系统之前，我们希望有一种快速的方法来预测更长的等待时间可能出现在哪里，并测试各种配置旋钮，例如回溯窗口的长度或最短运行时间允许的最大值（我们选择了8小时）。
我们构建了一个小型模拟环境，它以一组工作负载及其提交时间表为输入，并允许调度器做出抢占和 GPU 分配决策。凭借对每个工作负载所请求的 GPU 数量和总运行时间的了解，模拟器可以跳到可调度的时刻，并在几秒钟内提供对多天的模拟中队列等待时间、抢占事件以及各项目之间 GPU 时间分布的分析。我们针对历史提交数据和想要更好理解的构造场景运行了该模拟器。
我们想验证的一个假设涉及“调试工作负载”。这类任务需要少量 GPU 和不超过15分钟的最短运行时间，这足以让用户查看任务是成功启动，还是由于 bug 或配置错误而提前崩溃。我们想知道这些任务是否会比大型训练工作负载看到更短的队列等待时间，后者通常需要大量 GPU 和数小时的运行时间才能取得有意义的进展。直观地说，这些较小的任务应该会升到队列顶部，因为小任务可以放入的位置比大任务更多。但精确的队列延迟很重要。一两分钟的短暂等待将解锁一种新的开发实践，而十分钟的等待则变得不可行。
我们的模拟需要手工构建的测试数据，因为我们的历史记录中没有足够多的这类调试型工作负载。我们的结果支持了该假设，显示 p90 调试工作负载等待时间从约6小时降至仅5分钟。
基线（左）与新的“分配额度”调度器（右）的小规模模拟器可视化。每一行是一个 GPU；每一根柱条是一个任务，按其父工作负载着色，每个团队一种色调；阴影线标记任务可被中断的时间，红色边缘标记一次抢占。在基线中，长期紧急任务从不被中断，低优先级的工作上发生的抢占较少。新调度器在每个 GPU 上有更丰富的颜色组合，说明团队之间的占用轮转。
结果
拿到模拟结果后，我们于7月底开始逐个集群推广。我们关心的结果是：我们选择资助的工作负载是否获得了它们的时间，新系统是否保持了完全占用，以及研究人员是否能够理解调度器以做出明智的决策。
自推广以来，我们观察到用户和团队持续获得其分配的 GPU 时间。我们将欠一个团队的时间统计为逐小时按其实际需求封顶后的分配额度。在30天的测试期内，团队获得了它们应得 GPU 小时的98%，15个团队分配额度中有13个达到95%或以上，最差的情况也达到了90%。集群占用率在更改前后稳定保持在98%，两个时期的需求都超出容量2-3倍。18%的已交付 GPU 时间是未分配的，这正是我们在受资助用例尚未准备好运行时保持高占用率的方法。
我们的模拟器结果在方向上被证明是准确的，实际结果优于我们的预测。在新调度器下，调试工作负载的 p90 队列等待时间从 2 小时降至 30 秒，而基于手工构建测试场景的模拟预测为从 6 小时降至 5 分钟。值得注意的是，基线中调试工作负载的样本量较小，这意味着这些测量结果的方差更高。队列延迟总体上也因时间分片而得到改善：在我们最大的 H100 集群上，队列等待时间中位数从 5 分钟降至 24 秒，p90 等待时间下降了约三分之一（从 2.8 小时降至 1.8 小时）。
与我们着手解决的三个问题相比：
深蹲： 短小的调试工作负载会在一分钟内启动，这降低了抢占资源的价值。这种行为产生的费用会计入抢占者自己的预算，从而阻止他们在真正需要时获得时间。
优先级膨胀： 我们仍然允许工作负载声明优先级，但这只影响团队内部的排序。经理们有动力去监控整个团队范围内的优先级，以优化其预算的使用。
随时待命的琐事： 当工作负载达到最短运行时间后,不健康的主机会自动排空。需要人工介入的修复减少了74%。
挑战
学习曲线比我们预想的更为陡峭。我们是增量式推出该变更的，因此在早期，研究人员会因所针对的集群不同而体验到不同的行为。此外，我们的界面保留了一些旧术语（例如工作负载优先级），而这些术语的含义已经发生了变化。仅靠文档并不能消除困惑。什么 做了 工作正在开展实时讲解会议，为研究人员提供提问的机会，也让工程团队借助真实实例更深入地说明调度器是如何以及为何做出其优先级排序决策的。
这是一个关键时刻，因为它标志着从早期充满挫折感和民间猜想的阶段，转向了当前的模式——研究团队更频繁、更广泛地交流其实验的GPU需求。如今，研究人员在参与预算讨论时，能更清楚地了解为满足任何新请求而做出的权衡。
除现场会议外，我们还在上线后引入了新的可视化功能，让用户能更清楚地了解其分配的 GPU 时间与预期配额的接近程度，并直接展示用于工作负载队列排序的指标。当工作负载被抢占时，这为用户提供了一个简单的地方去查看原因。这些可视化还帮助预算负责人查看其管理的各个项目中 GPU 时间的使用情况。
随时间变化的分配使用情况示例可视化。
并非每个用例都有所改善。除了分布式训练之外，我们的研究人员还会启动交互式会话，在其中进行数据分析并在编写训练代码的同时进行测试。在旧系统中，研究人员可以将这样的会话保持长达一周。采用时间切片后，他们受到8小时保护运行时间上限的约束，超过分配额后，会话就会变为可抢占状态。我们没有充分意识到研究人员对这些会话的易失状态的依赖程度。被抢占意味着要等待 securing 一个新会话，并且还要手动重建他们的状态。在调研研究人员以了解这一问题的广泛程度之后，我们创建了两个新的路线图项目。我们将在本地存储旁边投资建设一个仅含CPU的集群，用于专注于数据准备任务的开发会话。这将为真正需要它的工作负载保留训练集群的容量。此外，我们计划为这些仅含CPU的工作负载构建可恢复的会话。这将使我们能够继续在其最短运行时间结束时抢占工作负载以进行维护或时间切片，同时也能够将会话恢复到其他地方，而无需研究人员重建它。我们既能保留这个新系统在运营和调度方面的优势，又能改善用户体验。
我们仍在密切关注新出现的问题。我们正在调查的一个潜在问题是容量碎片化,它可能导致最大规模的作业的队列等待时间增加。我们的直觉是,最低运行时间保护正被应用于那些过去依赖的作业类型上 可抢占的 机制以超出其团队的GPU并发限制。以前，这些作业可能随时被中断，既可能浪费时间，但也使大型作业的调度更加容易。现在，调度器可能更难一次性中断多个作业以安置一个大型待处理工作负载。我们目前正在使用模拟器工具重现此问题，同时也在生产环境中测量真实情况。
未来
在展望本文所述调度工作之外的未来时，我们的目标是金字塔的顶端：利用率。我们需要确保引导、检查点保存以及训练应用本身都尽可能高效地完成，从而最大化每个工作负载所获得的调度时间的价值。
如果您希望与研究人员紧密合作，应对这样的挑战， 我们鼓励您探索 Ai2 的开放工程职位.