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 上超出该限制。这种策略导致了可预测的病态现象。例如，我们观察到 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 时间的技巧都会消耗受益用户的配额。霸占工作负载实际上是在用团队预算做无用功。我们的策略是让操纵调度的成本高于诚实地参与争取更大预算的辩论。我们不断迭代这个预算审查流程，但关键要求是研究人员要有频繁的机会为他们所需的时间进行辩护，并且决策由最了解相关权衡的管理者做出。这意味着在一个研究项目内部的分配决策由首席研究员做出，在一个研究项目群内由主要研究者（Principal Investigator）做出，而在跨项目群之间则由首席项目经理或首席执行官做出。

公平份额

配合这一 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 分钟的最小运行时间，足以让用户查看作业是否成功启动或因错误/配置不当而早期崩溃。我们想知道这些作业的队列等待时间是否比大型训练工作负载更短，后者通常需要大量 GPU 和数小时的运行时间才能取得有意义的进展。直观地说，这些较小的作业应该排在队列顶部，因为小作业能容纳的位置比大作业多。但精确的队列延迟至关重要。一两分钟的短暂等待将解锁一种新的开发实践，但十分钟的等待就变得不可行。

我们的模拟需要手工构建的测试用例数据，因为我们的历史记录中没有足够高数量的此类类似调试的工作负载。我们的结果支持了这一假设，显示 p90 调试工作负载的等待时间从约 6 小时降至仅 5 分钟。

基线（左）和新“allocations”调度器（右）的小规模模拟器可视化。每一行代表一个 GPU；每个条形代表一个作业，颜色由父级工作负载决定，每个团队一种色调；斜线标记表示作业可中断的时间段，红色边缘标记抢占事件。在基线中，长期的紧急作业永远不会被中断，低优先级工作的抢占次数也较少。新调度器在每个 GPU 上具有更多混合的颜色，说明了跨团队的占用率轮换。

结果

有了模拟结果，我们在七月底开始了逐集群的部署。我们关心的结果是：我们选择资助的工作负载是否获得了所需的时间，新系统是否保持了全占用率，以及研究人员是否能够理解调度器以做出明智的决策。

自上线以来，我们观察到用户和团队始终能获得其分配的 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 的开放工程职位。