Ai2 发布 GPU 集群调度器改造实践：用 GPU 时间预算与层级 fair-share 替代优先级调度

构建集群调度器，在保持完全占用的同时优先支持高影响力研究

在 Ai2 的 AI 基础设施团队，我们负责为研究院提供 GPU 算力，专门面向大型分布式训练工作负载。我们把这项工作看作一座由四个层层递进的指标构成的金字塔。

基础是可用性（availability）：硬件健康并可以投入工作的频率。其上是占用（occupancy）：可用时间中分配给特定工作负载的比例。接下来是影响力（impact）：最有价值的工作负载被选中获得资源的频率。金字塔的顶端是利用率（utilization）：在一个工作负载的生命周期内使用的 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 独占权分配给重要项目来解决。虽然起初我们没有意识到，但我们实际上已经建立了一个观察“公地悲剧”的完美实验室。个体们在争夺一种稀缺的共享资源，而通过追求个体利益的最大化，他们达成了一个非最优的全局结果，并滥用了底层的资源。 

我们远非第一个观察到这种相互作用的人。资源分配是一个引人入胜的研究领域，它融合了算法开发、经济学和系统管理。一个核心问题在于，用户往往比组织更了解自己作业的价值，但他们可能有动机隐瞒这种价值，或者即使这样做会损害整体性能也不愿释放资源。例如，Ghodsi 等人在 2011 年提出 Dominant Resource Fairness 的论文中讲述了一则轶事：某搜索公司只有在用户能保证高利用率的情况下才为其作业提供专用机器。他们很快发现“用户会在代码中 sprink上无限循环来人为抬高利用率”。硬件在变化，但使资源分配变得复杂的根本问题依然存在。

预算而非时间表

解决公地悲剧的经典方案是将共享资源私有化——所有者会有动力使其财产价值最大化。当我们把 GPU 集合的独占权分配给各个团队时，我们其实已经在做类似的事情，但粒度太粗。由于研究的季节性，这导致 GPU 闲置。各团队准备运行实验和训练的时间不同，因此分配独占权会导致有时没有任何作业可以执行，而另一个团队却在等待算力。 

我们是在手动求解一个 背包问题，试图把动态变化的研究需求塞进一个静态的时间表中。我们既想要所有权带来的激励，又想保持 GPU 的完全占用。

我们决定对所有权模型进行迭代。我们没有把 GPU 分发给团队，而是选择分配一定比例的 GPU 时间。预测未来的需求需要知晓新型科学实验的结果，因此无法精确预测。然而，研究工作之间的优先级是一个战略问题，可以更容易地提前讨论和决定。我们不再试图解决调度难题，而是让领导层像投资者一样思考。在工作负载出现之前，根据他们对各项研究工作可能产生的影响的判断，决定如何用 GPU 时间为其提供资金。随后，调度器可以在对到达的工作负载进行优先级排序时使用这些信息。

基于这一思路，我们设计了一个层级化系统，管理者可以按比例将 GPU 时间分配给他们负责的项目和研究人员。正如下图所示，这将项目战略直接转化为有保障的 GPU 时间份额。项目 A1 知道自己占总容量的 35%，无论其他地方还有多少项目在排队。

括号中的数值表示分配给叶子项目的集群总容量占比。

在这个系统中，每一次申请 GPU 时间的请求都必须由预算提供资金，否则就无法免于被抢占。在旧系统中，HIGH 优先级不产生任何成本，且不可抢占性允许团队无限制地填满其并发 GPU 上限，所以人人都这么做。现在，没有什么是免费的，任何获取 GPU 时间的技巧都要从受益用户的配额中扣除。占着不用的负载就是在白白消耗团队预算。我们的策略是让钻调度器的空子比诚实地参与争取更多预算的讨论成本更高。我们仍在不断迭代这个预算评审流程，但关键要求是：研究人员要有足够频繁的机会来为他们所需的时间进行申辩，并且决策由对相关权衡最有了解的管理者做出。这意味着研究项目内的分配决策由项目负责人做出，研究计划内由首席研究员（PI）做出，跨计划的由计划负责人或 CEO 做出。

公平共享

与这套 GPU 时间预算工具配套，我们构建了一个层级化公平共享调度器，用于管理整个程序树中各配额的实际占用。这里的算法并不新颖——基于时间窗口的层级化公平共享可以追溯到 2009 年的 Hadoop 公平调度器，同样的方法如今仍在 SLURM 的 Fair Tree 和 YARN 的公平调度器中被积极使用。对我们而言，新的是输入：树结构对应研究计划的结构，权重由管理者设定的预算决定，而非静态配额。

调度器在一个滑动的回溯窗口内（我们默认为 7 天）跟踪占用情况，并将来自低利用率配额的负载排在来自高利用率配额的负载之上。这样，在一周的时间范围内，只要每个团队在持续提交有足够需求的负载，我们就能确保他们获得所分配的 GPU 时间。 

“新调度器让我们感觉多了 30% 的算力。在旧调度器下，如果某些时候我们没用满自己的槽位上限，那部分算力基本就浪费了。现在有了新调度器，如果出现这种情况，我们之后可以突发性地超出配额上限，任务依然能被快速调度且不被抢占，相当于把那部分算力拿了回来。我们的负载常常是突发式的，所以这为我们节省了大量算力。” —— Chris Clark

调度器区分两种占用。已分配占用是指负载被计入某个预算的时间。这会消耗负载所有者的配额，影响公平共享预算的计算，且这些负载在其最短运行时间窗口内受保护、不被抢占。未分配占用不计入任何预算，从一开始就不受保护，并且可能被任何已分配的请求抢占。这使我们即使在配额与需求不匹配时也能让 GPU 保持满载，同时防止团队拒绝免费的 GPU 算力。

调度契约

分布式训练中还有一个让公平资源分配变得困难的特点：工作负载可能运行非常长的时间。训练任务通常会运行数小时、数天，有时甚至数周。一旦被调度，一个工作负载可能会在其分配到的 GPU 上停留一周或更久，其他人就没有机会获得他们应得的时间份额。正是这一系统特性使得“GPU 占坑”成为可能，也迫使值班工程师不得不与长时间运行任务的所有者协商，以解决持续的维护问题。

为了解决这些问题，我们引入了“调度契约”。作为使用集群的交换条件，工作负载必须声明其最短运行时间，即取得有意义进展所需的最短占用时长。在此期间，工作负载会受到保护，不被抢占。这为研究人员提供了进展保证，同时赋予调度器在进展落袋后重新平衡的权利，自动将可恢复的工作负载重新排队。或者，用户也可以将最短运行时间设为零，这表示 GPU 时间应保持未分配状态。这类工作负载始终可能被抢占，但它们也是“免费”的，因为不会计入任何预算。

工作负载的生命周期遵循以下模式：

1. 工作负载提交时会带有最短运行时间，并标明其是否可恢复。
2. 工作负载根据公平共享算法被调度，权重为回看窗口内实际占用时间与已分配时间之比。
3. 工作负载运行其最短运行时间，这段时间计入其分配额度。
4. 只要相关分配额度继续将其排在其他工作负载之前，工作负载就可以继续运行。这段时间同样计入其分配额度。
5. 它可能被抢占并重新排队，回到第 2 步。
6. 工作负载完成，释放其对任何资源的占用。

这些约定共同为我们的调度器增加了时间分片。运行中的工作负载可以被自动移除并重新排队，使公平共享得以收敛，并抑制占坑行为。它们还允许不健康的主机在工作负载达到最短运行时间后将其排空，从而使修复活动可以完全自动化。这一点比我们在规划这项工作时意识到的更为重要。它将需要人工介入的修复减少了 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 倍。交付的 GPU 时间中有 18% 是未分配的，这使我们在被资助的用例尚未准备好运行时仍保持了高占用率。

我们的模拟器结果在方向上是准确的，真实结果甚至超出了预测。在新调度器下，调试工作负载的 p90 排队等待时间从 2 小时降至 30 秒，而基于手工构建测试场景的模拟预测是从 6 小时降至 5 分钟。值得注意的是，基线中调试工作负载样本量较小，意味着这些测量的方差更高。作为时间切片的副作用，排队延迟整体上也有所改善：在我们最大的 H100 集群上，排队等待时间中位数从 5 分钟降至 24 秒，p90 等待时间下降了约三分之一（从 2.8 小时降至 1.8 小时）。

对照我们着手解决的三个问题：

1. 占坑： 短调试工作负载在一分钟内即可启动，降低了占坑的价值。这种行为的代价会计入占坑者的预算，使其在真正需要时无法获得时间。

2. 优先级通胀：我们仍然允许工作负载声明优先级，但它只影响团队内部的排序。管理者有动力去监控整个组的优先级，以优化其预算的使用。

3. 值班负担：当工作负载达到最短运行时间后，不健康的主机会自动排空。需要人工介入的修复工作减少了 74%。

挑战

学习曲线比我们预想的更陡峭。我们是逐步推出这一变更的，因此在早期，研究人员根据所针对的集群不同会体验到不同的行为。此外，我们的界面保留了一些旧术语（如工作负载优先级），但其含义已经改变。仅靠文档无法消除这些困惑。真正有效的是举办现场讲解会，为研究人员提供一个提问的论坛，也让工程团队通过实际例子更深入地说明调度器如何以及为何做出优先级决策。

这是一个关键时刻，它标志着从早期充满挫败感和民间理论的状态，转变为如今研究小组更频繁、更广泛地交流其实验 GPU 需求的模式。研究人员现在在参与预算讨论时，对为满足任何新请求而需要做的权衡有了更清晰的认识。

除了现场会议之外，我们还在上线后引入了新的可视化功能，让用户更清楚地了解其分配的 GPU 时间与预期配额的匹配程度，并直接展示用于工作负载队列排序的指标。这为工作负载被抢占时提供了一个简单的地方来了解原因。这些可视化还帮助了预算负责人，他们可以看到其管理的各个项目的 GPU 时间使用情况。 

随时间变化的配额使用情况示例可视化。

并非所有用例都得到了改善。除了分布式训练之外，我们的研究人员还会启动交互式会话，在其中进行数据分析并在编写训练代码的同时进行测试。在旧系统中，研究人员可以让这样的会话持续长达一周。采用时间切片后，他们受 8 小时受保护运行时间上限的约束，超过该时间后，如果会话超出其配额就会变为可抢占。我们没有充分意识到研究人员对这些会话易失状态的依赖程度。被抢占意味着等待获取新会话，还要手动重建其状态。在调研了研究人员以了解这一问题的广泛程度之后，我们创建了两个新的路线图项目。我们将在本地存储旁投资建设一个仅供 CPU 使用的集群，用于专注于数据准备任务的开发会话。这将把我们的训练集群容量留给真正需要的工作负载。此外，我们计划为这些仅供 CPU 的工作负载构建可恢复的会话。这将使我们能够在最短运行时间结束后继续为维护或时间切片而抢占工作负载，同时还能在别处恢复会话，而无需研究人员手动重建。这样我们既保留了新系统在运维和调度方面的优势，又改善了用户体验。

我们持续关注新出现的问题。我们正在调查的一个潜在问题是容量碎片化，它可能导致最大规模的训练任务的排队等待时间增加。我们的直觉是，最低运行时间保护正被应用于那些过去依赖可抢占机制来超出其团队并发 GPU 限额的任务类型。以前，这些任务随时可能被中断，虽然可能浪费已运行的时间，但也让大型任务更容易被调度。现在，调度器可能有更少的机会一次性中断许多任务，以放置一个大型待处理任务。我们目前正在使用模拟器工具来重现这个问题，同时也在生产环境中测量真实情况。

未来

在展望本文所述调度工作之外的未来时，我们的目标是金字塔的顶点：利用率。我们需要确保引导启动、检查点保存以及训练应用本身都尽可能高效地完成，最大化每个任务所获得的调度时间的价值。

如果您想与研究人员密切合作，迎接这样的挑战，我们鼓励您了解 Ai2 的开放工程职位。