Ai2 提升 GPU 叢集排程效率：以預算制和公平共享加速 AI 研究

在 Ai2 的 AI 基礎設施團隊中，我們負責提供研究所的 GPU 運算能力，特別是針對大規模分散式訓練工作負載。我們將這項任務視為一個由四個相互建立的指標組成的金字塔。

金字塔的基礎是「可用性」，即硬體健康並準備好工作的頻率。其上是「佔用率」，即分配給特定工作負載的可用時間比例。接著是「影響力」，即最有價值的工作負載被選中獲得資源的頻率。金字塔的頂點是「利用率」，即工作負載生命週期中 GPU 容量的使用比例。

這篇文章旨在說明我們如何改進排程決策的「影響力」。我們最近用一套包含 GPU 時間預算、分層公平共享分配和時間切片合約的系統，取代了舊有的優先級排程器。結果，我們將每個研究專案應得多少 GPU 時間的爭論，從逐案處理的營運任務轉變為透明的行政預算流程。

超額承諾

在 Ai2，我們管理著數千個 NVIDIA H100、B200 和 B300 GPU，這些 GPU 組成規模從 88 到 1024 個 GPU 不等的叢集。這些叢集專為 AI 模型的大規模分散式訓練而建，服務著約 150 名內部研究人員，他們的工作涵蓋多樣化的 AI 領域，包括 LLM 和 VLM 訓練的完整模型流程、機器人強化學習 (RL) 模擬，以及科學代理應用案例的後訓練。

與許多實驗室一樣，我們對 GPU 時間的需求遠超供應。根據提交的工作負載，任何時刻我們都有 2-3 倍於可用 GPU 數量的待處理請求。換句話說，我們叢集上的每個可用 GPU 小時都有 2-3 個不同的研究工作負載在競爭。

過去，我們使用基於優先級的排程器，並允許工作負載選擇不被搶佔。每個團隊都有一個同時運行 GPU 的限制，這些 GPU 上的工作負載受到搶佔保護。可搶佔的工作負載則可以在閒置 GPU 上超出該限制。這種策略產生了可預見的病態行為。

例如，我們觀察到 GPU「佔用」的情況，使用者會停放無操作的工作負載，以便在需要時連接。這發生是因為研究人員發現他們無法以足夠低的延遲啟動偵錯工作負載，以即時解決問題。我們還觀察到優先級膨脹，最終 100% 的排程工作負載都使用「高」優先級。這意味著較低優先級完全缺乏 GPU 時間。

由於可搶佔性是可選的，我們還發現值班工程師將大部分的工單回應時間花在協商有組織地關閉運行在已知維護問題主機上的不可搶佔工作負載。

共有財悲劇

當這些問題出現時，我們遲遲未能找出其根本原因。我們最初確保最重要工作獲得 GPU 時間的嘗試，主要集中在更嚴格地控制優先級設定，並最終透過明確分配 GPU 獨佔權給重要專案來繞過基於優先級的排程器。雖然我們一開始沒有意識到，但我們已經建立了一個完美的實驗室來觀察「共有財悲劇」。

個人為了稀缺的共享資源而競爭，透過追求個人利益最大化，卻導致了非最佳的整體結果並濫用了底層資源。我們遠非第一個觀察到這種互動的人。資源分配是一個引人入勝的研究領域，它融合了演算法開發、經濟學和系統管理。

一個核心問題是，使用者通常比組織更了解自己任務的價值，但他們可能有動機隱藏該價值或佔用資源，即使這樣做會損害整體性能。例如，Ghodsi 等人在 2011 年介紹 Dominant Resource Fairness 的論文中，講述了一個軼事：一家搜尋公司僅在使用者能保證高利用率的情況下才為其任務提供專用機器。

他們很快發現「使用者會在其程式碼中散佈無限迴圈，以人為地提高利用率」。硬體在變，但使資源分配複雜化的根本問題依然存在。

預算而非排程

解決共有財悲劇的經典方法是將共享資源私有化，所有者有動機最大化其財產價值。當我們為團隊分配 GPU 獨佔權時，我們已經在做這種事情，但它太粗糙了。由於研究的季節性，這導致 GPU 閒置。團隊在不同時間準備運行實驗和訓練，因此分配獨佔權會確保在某些時候沒有任務準備執行，而另一個團隊則在等待容量。

我們手動解決了一個背包問題，試圖將動態變化的研究需求納入靜態排程。我們想要所有權激勵，但我們也想保持 GPU 的完全佔用。

我們決定迭代所有權模型。我們選擇分配一部分 GPU 時間，而不是直接發放 GPU 給團隊。預測未來需求需要了解新穎科學實驗的結果，因此無法精確預測。然而，研究工作的優先級是一個策略問題，可以更容易地事先討論和決定。我們不再試圖解決排程難題，而是讓領導層像投資者一樣思考。

在工作負載存在之前，根據他們對其潛在影響的判斷，決定如何以 GPU 時間資助每項研究工作。排程器隨後可以使用這些資訊來優先處理到來的工作負載。考慮到這一點，我們設計了一個分層系統，經理可以按比例將 GPU 時間分配給他們負責的專案和研究人員。如下圖所示，這將專案策略直接轉化為保證的 GPU 時間份額。專案 A1 知道它擁有總容量的 35% 權利，無論有多少其他專案在其他地方排隊。

括號中的值代表分配給葉節點專案的總叢集容量。在這個系統中，每個 GPU 時間請求都必須由預算資助，否則它就不受搶佔保護。在舊系統中，「高」優先級沒有成本，不可搶佔性允許團隊無限期地佔用其同時運行 GPU 限制，所以每個人都使用它們。現在，沒有什麼是免費的，所以任何獲取 GPU 時間的技巧都會從受益使用者的分配中扣除。一個佔用的工作負載正在浪費團隊預算。

我們的策略是讓玩弄排程器比誠實地參與爭取更大預算的討論付出更高的代價。我們不斷迭代這個預算審查流程，但關鍵要求是研究人員有頻繁的機會爭取他們所需的時間，並且決策由對權衡取捨最有背景資訊的經理做出。這意味著研究專案內的分配決策由首席研究員做出，研究專案內的分配決策由主要研究員做出，跨專案的分配決策由首席專案經理或執行長做出。

公平共享

除了這個 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%，這大幅節省了值班工作的辛勞。

模擬

我們知道排程策略變更可能產生意想不到的後果。零