Ai2 如何用 GPU 时间预算与分层公平分配改造集群调度

构建集群调度器，优先处理高影响力研究，同时保持满占用率

在 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 独占权来绕开基于优先级的调度器。虽然我们起初并未意识到，但我们实际上搭建了一个观察“公地悲剧”的完美实验室。个体在争夺一种稀缺的共享资源，并且通过追求个体结果最大化，达成了非最优的全局结果，并滥用了底层资源。 

我们远非最早观察到这类相互作用的人。资源分配是一个引人入胜的研究领域，融合了算法开发、经济学和系统管理。一个核心问题在于，用户往往比组织更清楚自己任务的价值，但他们可能有动机隐藏这一价值，或在这样做会损害整体性能时仍然占住资源。例如，Ghodsi 等人在其 2011 年提出 Dominant Resource Fairness 的论文中讲述了一则轶事：一家搜索公司只向那些用户能保证高利用率的任务提供专用机器。他们很快发现，“用户会在代码中塞入无限循环，以人为抬高利用率水平。”硬件在变，但使资源分配变得复杂的基本问题依然存在。

预算而非排期

公地悲剧的经典解决方案是将共享资源私有化——所有者有动力最大化其财产的价值。当我们把多组 GPU 的独占权分配给团队时，我们已经在做某种版本的这个方案，但它过于粗放。由于研究的季节性，它导致 GPU 闲置。团队准备运行实验和训练的时间各不相同，因此分配独占权会确保有时没有任务可执行，而另一个团队却在等待容量。 

我们当时在手动解决一个背包问题，试图把动态变化的研究需求塞进一个静态排期中。我们想要所有权激励，但也想保持 GPU 的充分利用。

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

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

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

在这个系统中，每个 GPU 时间请求都必须由预算来支撑，否则就不受抢占保护。在旧系统中，HIGH 优先级不产生任何成本，而不可抢占性又允许团队无限期地占满其并发 GPU 上限，因此所有人都去使用它们。现在，没有什么是免费的，所以任何获取 GPU 时间的技巧都会消耗受益用户的配额。一个占着资源却不干活的工作负载，就是在白白消耗团队预算。我们的策略是让钻调度器空子比诚实地参与争取更大预算的讨论代价更高。我们一直在迭代这个预算评审流程，但关键要求是：研究人员要有频繁的机会为他们所需的时间进行争取，并且决策由对相关权衡最有上下文的管理者做出。这意味着研究项目内的分配决策由首席研究员做出，研究计划内的分配决策由项目负责人做出，跨计划的分配决策由首席计划经理或 CEO 做出。

公平份额

配合这个 GPU 时间预算工具，我们构建了一个分层公平份额调度器，用于管理整个计划树中配额的实际占用情况。这里的算法并不新鲜——基于时间窗口的分层公平份额属于一条可追溯到 2009 年 Hadoop Fair Scheduler 的谱系，同样的方法如今仍在 SLURM 的 Fair Tree 和 YARN 的 Fair Scheduler 中积极使用。对我们来说新鲜的是输入：这棵树镜像了研究计划的结构，而权重是由管理者设定的预算，而非静态配额。

调度器在一个滑动回看窗口（我们默认 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 分钟。

较小规模的模拟器可视化：基线（左）与新的“allocations”调度器（右）。每一行代表一块 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 小时）。

与我们最初要解决的三个问题对比：

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

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

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

挑战

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

这是一个关键时刻，因为它标志着从早期充满挫败感和民间理论的阶段，转向了当前模式：研究团队更频繁、更广泛地就其实验的 GPU 需求进行沟通。研究人员现在参与预算讨论时，能更清楚地了解为满足任何新请求所做的权衡。

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

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

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

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

未来

当我们展望此处描述的调度工作之外时，我们的目标是金字塔的顶点：利用率。我们需要确保引导、检查点和训练应用本身都尽可能高效地完成，最大化每个工作负载所获得的调度时间的价值。

如果你希望与研究人员密切合作应对此类挑战，我们鼓励你探索 Ai2 的开放工程职位。