Ai2 介绍 GPU 集群调度改造：以时间预算和公平份额分配资源

构建集群调度器，在保持满负载的同时优先支持高影响力研究

在 Ai2 的 AI 基础设施团队，我们负责提供研究所的 GPU 算力，特别是面向大型分布式训练工作负载。我们将这项任务视为由四项层层递进指标构成的金字塔。

基础是可用性：硬件处于健康状态并可投入工作的频率。其上是占用率：可用时间中分配给特定工作负载的比例。接下来是影响力：最有价值的工作负载获选并得到资源的频率。金字塔顶端是利用率：工作负载整个生命周期内实际使用的 GPU 算力比例。

这篇文章将介绍如何提升我们调度决策的影响力。我们最近用一套包含 GPU 时间预算、分层公平份额分配和时间片契约的系统，取代了基于优先级的调度器。因此，关于每个研究项目应获得多少 GPU 时间的争论，从逐个案例处理的运维任务转变为透明的行政预算流程。

超额承诺

在 Ai2，我们管理着数千块 NVIDIA H100、B200 和 B300 GPU，这些 GPU 被部署在规模从 88 到 1024 块 GPU 不等的集群中。这些集群专为 AI 模型的大规模分布式训练而构建，服务于约 150 名内部研究人员；他们的工作涵盖多个 AI 领域，包括 LLM 和 VLM 训练的完整模型流程、机器人强化学习（RL）仿真，以及面向科学智能体应用场景的后训练。

和许多实验室一样，我们对 GPU 时间的需求远远超过了供给。根据已提交的工作负载，在任何时刻，尚未满足的 GPU 请求量都是可用数量的 2-3 倍。换个角度看，集群中的每个可用 GPU 小时，都有 2-3 个不同的研究工作负载在争抢。

过去，我们使用基于优先级的调度器，并允许工作负载选择不接受抢占。每个团队都有一个并发 GPU 使用上限，受保护、不受抢占的工作负载不得超过该上限。可抢占的工作负载则可以在空闲 GPU 上超出这一上限。这种策略带来了可预见的问题。例如，我们发现有人会让 GPU 工作负载“占坑”——这些工作负载实际上什么也不做，用户可以在有需要时连接到它们。出现这种情况，是因为研究人员发现，他们无法以足够低的延迟启动调试工作负载，无法实时解决问题。我们还观察到优先级膨胀，最终所有已调度的工作负载都使用 HIGH 优先级。这意味着较低优先级的工作负载完全得不到 GPU 时间。由于是否接受抢占是可选的，我们还发现，值班工程师响应工单的大部分时间都花在协商有序关闭那些运行在已知存在维护问题的主机上、且不接受抢占的工作负载上。

公地悲剧

当这些问题出现时，我们未能及时找出根本原因。最初，为确保最重要的工作获得 GPU 时间，我们试图更严格地控制优先级的设定；最终，我们绕过了基于优先级的调度器，直接将 GPU 独占分配给重要项目。起初我们并未意识到，但我们其实搭建了一个观察“公地悲剧”的绝佳实验室。人们争夺稀缺的共享资源，为了最大化个人收益，最终却造成全局结果不理想，并滥用底层资源。

我们绝非第一个观察到这类互动的人。资源分配是一个引人入胜的研究领域，融合了算法开发、经济学和系统管理。一个核心问题是，用户往往比组织更了解自身作业的价值，但他们可能有动机隐瞒这种价值，或在损害整体性能的情况下仍占用资源。例如，在介绍 Dominant Resource Fairness 的2011年论文中，Ghodsi 等人讲述了一则轶事：一家搜索公司只有在用户能够保证高利用率时，才会为作业提供专用机器。他们很快发现，“用户会在代码中加入无限循环，人为抬高利用率。”硬件会变，但令资源分配变得复杂的根本问题始终存在。

分配预算，而非排期

公地悲剧的经典解决方案是将共享资源私有化——所有者会受到激励，去最大化其财产的价值。我们把一组 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 时间分配预算。这类工作负载始终可能被抢占，但它们也是免费的，因为不会计入任何预算。

工作负载的生命周期如下：

提交工作负载时，需指定最短运行时间，并说明是否可恢复。
调度器根据公平共享算法安排工作负载，并依据回看窗口内实际占用时间与分配时间的比率加权。
工作负载运行至其最短运行时间，这段时间计入其分配额度。
只要相关分配额度仍优先支持该工作负载，它就可以继续运行。这段时间也计入其分配额度。
工作负载可能被抢占并重新排入队列，然后回到第 2 步。
工作负载完成，并释放其对所有资源的占用。

这些约定共同为我们的调度器增加了时间片机制。运行中的工作负载可以被自动移除并重新排入队列，从而使公平共享机制逐步收敛，并减少占坑行为的动机。它们还使不健康的主机能够在工作负载达到最短运行时间后将其清空，因此维修活动可以完全自动化。规划这项工作时，我们并未意识到这一点有多重要。它使需要人工介入的维修减少了 74%，大幅降低了值班人员的负担。

模拟

我们知道，调度策略的变更可能带来意想不到的后果。由于这是一个零和问题，给一位研究人员更多时间，就意味着从另一位研究人员那里拿走时间。在这种交换中失利的用户往往会寻求新的变通办法。在推出基于预算的系统之前，我们希望找到一种快速预测等待时间可能延长之处的方法，并测试各种配置参数，例如回看窗口的长度或允许的最短运行时间上限（我们选择了 8 小时）。

我们构建了一个小型模拟环境，以一组工作负载及其提交计划为输入，并允许调度器决定是否抢占以及如何分配 GPU。借助每个工作负载请求的 GPU 数量和总运行时间，模拟器可以直接跳转到可调度的时刻，并在几秒钟内分析多天模拟运行中的队列等待时间、抢占事件以及各项目之间的 GPU 时间分配情况。我们使用历史提交数据以及自行构建、希望深入了解的场景运行了模拟器。

我们想要验证的一个假设涉及“调试工作负载”。这类任务只需要少量 GPU，运行时间最短为 15 分钟或更短，足以让用户确认任务是否成功启动，或是否因错误或配置不当而过早崩溃。我们想了解，与规模更大的训练工作负载相比，这类任务是否能缩短排队等待时间；后者通常需要许多 GPU 并运行数小时才能取得实质性进展。直觉上，这些较小的任务应该能排到队列前面，因为小任务比大任务更容易找到可运行的位置。但具体的排队延迟很重要。等待一两分钟就能开启一种新的开发实践，但等待十分钟则会让这种做法变得不可行。

由于历史记录中这类类似调试的工作负载数量不足，我们必须手动构建测试用例数据来进行模拟。结果支持了这一假设：调试工作负载的 p90 等待时间从约 6 小时降至仅 5 分钟。

基线方案（左）和新的“allocations”调度器（右）的缩小版模拟可视化。每行代表一个 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 小时）。

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

占用资源：短时调试工作负载可在一分钟内启动，从而降低了占用资源行为的收益。这种行为的代价会计入占用者的预算，导致他们在真正需要时无法获得运行时间。
优先级膨胀：我们仍允许工作负载声明优先级，但它只影响团队内部的排序。管理者有动力监控整个团队的优先级，以优化预算使用。
值班负担：随着工作负载达到最短运行时间，不健康的主机将自动排空。需要人工介入的修复工作减少了 74%。

挑战

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

这是一个关键时刻，因为它标志着我们从早期的沮丧和各种经验之谈，转向了当前的状态：研究团队更频繁、更广泛地交流实验所需的 GPU 资源。如今，研究人员对其中的权衡有了更清晰的认识，并据此参与预算讨论，以便为任何新请求做出安排。

除了现场说明会，我们还在上线后推出了新的可视化功能，让用户更清楚地了解分配给自己的 GPU 时间与预期配额的吻合程度，并直接展示用于对工作负载队列排序的指标。当工作负载被抢占时，用户可以通过这个简单的入口了解原因。这些可视化也帮助了预算负责人，让他们能够查看其管理的各个项目如何使用 GPU 时间。

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

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

我们会继续关注新出现的问题。我们正在调查的一个潜在问题是容量碎片化，这可能会导致最大型工作负载的队列等待时间增加。我们的直觉是，最低运行时间保护机制被应用到了这类任务上，而过去这些任务会依赖 可抢占 机制来突破团队并发 GPU 限额。以前，这些任务随时可能被中断，可能会浪费已运行的时间，但也让大型任务更容易得到调度。现在，调度器可能没有那么多机会一次中断多个任务，以便安排大型待处理工作负载。我们目前正在使用模拟器工具复现这一问题，同时在生产环境中测量实际情况。

未来

在展望本文所述的调度工作之后，我们的目标是攻克这座金字塔的顶点：利用率。我们需要确保引导启动、检查点保存以及训练应用本身都尽可能高效，从而最大限度地发挥每个工作负载所获调度时间的价值。

如果你想与研究人员密切合作，应对这类挑战，我们鼓励你了解 Ai2 的工程岗位。