构建集群调度器，优先处理高影响力研究，同时保持满占用率
在 Ai2 的 AI Infrastructure 团队，我们负责提供研究院的 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 时间的伎俩都会从受益用户的分配额度中扣除。一个占着不走的工作负载只是在白白消耗团队预算。我们的策略是让钻调度器空子比诚实地参与争取更大预算的讨论更昂贵。我们一直在迭代这个预算审查流程，但关键要求是研究人员有频繁的机会为他们所需的时间进行争取，并且决策由对相关权衡拥有最多背景信息的管理者做出。这意味着研究项目内部的分配决策由首席研究员做出，研究计划内部由首席研究员（PI）做出，跨计划则由首席项目经理或 CEO 做出。
公平份额
与这个 GPU 时间预算工具相配套，我们构建了一个分层公平份额调度器，用于管理整个计划树中分配额度的实际占用情况。这里的算法并不新鲜——基于时间窗口的分层公平份额属于一条可追溯至2009 年的 Hadoop Fair Scheduler，并且同样的方法至今仍在积极使用于SLURM 的 Fair Tree和YARN 的 Fair Scheduler中。对我们来说新鲜的是输入：这棵树镜像了研究计划的结构，而权重是由管理者设定的预算，而不是静态配额。
调度器在一个滑动回看窗口（默认 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 小时上限的限制，超过分配额度后，会话就变为可抢占的。我们没有充分意识到研究人员对这些会话易失状态的依赖程度。被抢占意味着要等待获得新会话，还要手动重建他们的状态。在调研研究人员以了解这一问题的广度后，我们创建了两个新的路线图项目。我们正在本地存储旁边投资一个仅 CPU 集群，用于专注于数据准备任务的开发会话。这将为真正需要它的工作负载保留我们的训练集群容量。此外，我们计划为这些仅 CPU 工作负载构建可恢复的会话。这将使我们能够在工作负载达到最短运行时间后，为维护或时间切片而继续抢占它们，同时还能在其他地方恢复会话，而无需研究人员重建。我们既能保留这个新系统在运维和调度方面的优势，又能改善用户体验。
我们仍在持续留意新出现的问题。我们正在调查的一个潜在问题是容量碎片化，这可能会导致最大工作负载的队列等待时间增加。我们的直觉是，最短运行时间保护正被应用于那些过去依赖可抢占机制来超出其团队并发 GPU 限制的作业类型。以前，这些作业可以随时被中断，可能浪费那段时间，但也使调度大型作业更容易。现在，调度器可能更少有機會一次性中断许多作业，以放置一个大型待处理工作负载。我们目前正在使用模拟器工具来复现这个问题，同时也在生产环境中测量真实情况。
未来
当我们把目光投向本文所述的调度工作之外时，我们的目标是这座金字塔的顶点：利用率。我们需要确保引导启动、检查点保存以及训练应用本身都尽可能高效地完成，从而最大化每个工作负载所获得的调度时间的价值。
如果你想与研究人员密切合作来应对这类挑战，我们鼓励你探索 Ai2 的开放工程职位。