高い稼働率を維持しながら、インパクトの高い研究を優先するクラスタスケジューラの構築
Ai2のAI Infrastructureチームでは、研究所のGPU計算能力の提供を担当しており、特に大規模な分散型学習ワークロードを対象としています。私たちはこのタスクを、互いに積み重なる4つの指標からなるピラミッドとして捉えています。
その土台となるのは availability(可用性)です。すなわち、ハードウェアが正常で作業可能な状態にある頻度です。その上にあるのは occupancy(占有度)です。すなわち、利用可能な時間のうち特定のワークロードに割り当てられた割合です。次は impact(インパクト)です。すなわち、最も価値のあるワークロードがリソースの割り当て先として選ばれる頻度です。ピラミッドの頂点にあるのは utilization(利用率)です。すなわち、ワークロードのライフタイム全体で使用されたGPU容量の割合です。
この記事は、スケジューリング判断のインパクトを改善することについてのものです。私たちは最近、優先度ベースのスケジューラを、GPU時間予算、階層的フェアシェア割り当て、タイムスライシング契約を含むシステムに置き換えました。その結果、各研究プロジェクトがどれだけのGPU時間に値するかという議論を、個別対応の運用タスクから、透明な管理上の予算プロセスへと移行させることができました。
オーバーコミット
Ai2では、88から1024 GPUの規模のクラスタに配置された、何千ものNVIDIA H100、B200、B300 GPUを管理しています。これらのクラスタはAIモデルの大規模分散学習のために構築されており、約150人の内部研究者にサービスを提供しています。彼らの研究は、LLMおよびVLM学習のフルモデルフロー、ロボティクス強化学習(RL)シミュレーション、科学的エージェント的ユースケース向けのポストトレーニングを含む、多様なAI領域に及んでいます。
多くの研究所と同様に、私たちのGPU時間への需要は供給をはるかに上回っています。提出されたワークロードに基づくと、任意の時点で、利用可能な数の2〜3倍のGPUに対する未処理のリクエストが存在しています。これを考える一つの方法は、クラスタ上の利用可能なGPU時間1時間ごとに、2〜3つの異なる研究ワークロードが競合しているということです。
従来、私たちは優先度ベースのスケジューラを使用し、ワークロードがpreemptability(横取り可能性)をオプトアウトすることを許していました。各チームには、横取りから保護されたワークロードが使用できる同時実行GPU数の上限がありました。横取り可能なワークロードは、アイドル状態のGPUでその上限を超えることができました。この戦略は予測可能な病理を生み出しました。たとえば、GPUの「不法占拠(squatting)」の事例を観察しました。これは、ユーザーが必要になったときに接続できるようなno-opワークロードを駐車させておくというものでした。これは、研究者が問題にリアルタイムに対処できるほど低いレイテンシでデバッグ用ワークロードを起動できないことに気づいたために発生しました。また、優先度のインフレも観察され、最終的にはスケジュールされたワークロードの100%がHIGH優先度を使用するようになりました。これは、より低い優先度レベルがGPU時間を完全に飢餓状態にされたことを意味します。横取り可能性が任意であったため、オンコールのエンジニアがチケット対応時間の大部分を、既知のメンテナンス問題を抱えるホスト上で実行されている横取り不可能なワークロードの計画的なシャットダウンの交渉に費やしていることも判明しました。
コモンズの悲劇
これらの問題が顕在化したとき、私たちはその根本原因を特定するのに時間がかかりました。最も重要な作業がGPU時間を確実に受け取るようにするための最初の試みは、優先度の設定方法をより厳格に管理すること、そして最終的には、重要なプロジェクトにGPUの独占権を明示的に割り当てることで優先度ベースのスケジューラを回避することに焦点が置かれていました。当初は認識していませんでしたが、私たちは「コモンズの悲劇」を観察するための完璧な実験室を作り上げてしまっていたのです。個々人が希少な共有リソースをめぐって競争し、個人の成果を最大化しようとすることで、全球的に最適でない結果を達成し、基盤となるリソースを濫用していました。
この種の相互作用を観察したのは私たちが最初ではありません。リソース割り当ては、アルゴリズム開発、経済学、システム管理が混ざり合う魅力的な研究領域です。中心的な問題は、ユーザーは自分のジョブの価値を組織よりもよく知っていることが多い一方で、その価値を隠したり、全体のパフォーマンスを損なうことが分かっていてもリソースを保持しようとするインセンティブを持つ可能性があることです。たとえば、2011年の論文で Dominant Resource Fairnessを紹介したGhodsiらは、ある検索企業が、ユーザーが高い利用率を保証できる場合にのみジョブに専用マシンを提供したという逸話を記しています。彼らはすぐに、「ユーザーが利用率の水準を人為的に水増しするために、コードに無限ループを散りばめる」ということを発見しました。ハードウェアは変化しますが、リソース割り当てを複雑にする根本的な問題は存続します。
スケジュールではなく予算
コモンズの悲劇に対する古典的な解決策は、共有リソースを私財化することです。所有者は自分の財産の価値を最大化するインセンティブを持ちます。私たちがチームにGPUセットの独占権を割り当てたとき、すでにこの方法の一種を行っていましたが、それは粗すぎました。研究の季節性のために、GPUが遊休状態になる事態を引き起こしました。チームごとに実験や学習を実行する準備が整う時期は異なるため、独占権を割り当てると、実行準備のできたジョブが存在しない時間帯が発生し、別のチームがキャパシティを待たされることになるのです。
私たちは ナップサック問題を手作業で解いていました。つまり、動的に変化する研究ニーズを静的なスケジュールに収めようとしていたのです。私たちは所有権によるインセンティブを望むと同時に、GPUの完全な稼働率も維持したいと考えていました。
私たちは所有権モデルを繰り返し改善することにしました。チームにGPUを支給する代わりに、GPU時間の一部を割り当てる方式を採用したのです。将来の需要を予測するには、未知の科学実験の結果を知っている必要があるため、正確に予測することはできません。しかし、研究活動間の優先順位は戦略上の問題であり、事前に議論し決定することがより容易です。スケジューリングのパズルを解こうとする代わりに、私たちはリーダーシップに投資家のように考えられるようにしました。ワークロードが存在する前に、それぞれの研究活動にもたらされるであろうインパクトについての判断に基づき、GPU時間によってどのように資金を配分するかを決定します。スケジューラはその情報を使って、到着するワークロードの優先順位を決めることができます。
この考えをもとに、マネージャーが自分の担当するプロジェクトや研究者にGPU時間を比例配分できる階層型システムを考案しました。下の図が示すように、これによりプログラムの戦略が保証されたGPU時間のシェアへと直接変換されます。Project A1は、他にいくつプロジェクトが待機していても、総容量の35%を確保できることを把握しています。
括弧内の値は、リーフプロジェクトに割り当てられたクラスタの総容量を表します。
このシステムでは、GPU時間へのすべての要求は予算によって裏付けられている必要があり、そうでない場合はプリエンプション( preempt )から保護されません。旧システムでは、HIGH優先度にはコストがなく、プリエンプション不可の設定によりチームは同時GPU上限を無制限に埋めることができたため、誰もがそれを利用していました。今では何もタダではなく、GPU時間を得るためのあらゆる手が、その恩恵を受けるユーザーの割り当てから差し引かれます。居座りワークロードは、何の生産性もなくチームの予算を消費していることになります。私たちの戦略は、スケジューラをゲーム化することを、より大きな予算を求める正当な議論に参加することよりも高くつくようにすることです。私たちはこの予算レビュープロセスを絶えず改善していますが、重要な要件は、研究者が自分に必要な時間を主張する機会が頻繁にあり、かつその決定が、問題となるトレードオフについて最も多くの文脈を知るマネージャーによって行われることです。つまり、研究プロジェクト内の割り当て決定はリード研究者が、研究プログラム内では主任研究員(PI)が、プログラム間ではリードプログラムマネージャーまたはCEOが行います。
フェアシェア
このGPU時間予算ツールと併せて、プログラムツリー全体の割り当ての実際の稼働率を管理するための階層型フェアシェアスケジューラを構築しました。ここでのアルゴリズムは新しいものではありません。時間窓上の階層型フェアシェアは 2009年のHadoop Fair Schedulerに遡る系譜の一部であり、同じアプローチが今日でも SLURMのFair Tree や YARNのFair Schedulerで活用されています。私たちにとって新しいのは入力です。ツリーは研究プログラムの構造を反映し、重みは静的なクォータではなくマネージャーが設定した予算です。
スケジューラはスライディングのルックバックウィンドウ(デフォルトは7日)にわたって稼働率を追跡し、稼働率が低い割り当てからのワークロードを、稼働率が高い割り当てからのものより上に並べます。これにより、1週間の時間範囲では、十分な需要のあるワークロードを積極的に提出している限り、すべてのグループが割り当てられたGPU時間を受け取れると期待できます。
「新しいスケジューラのおかげで、追加で30%の計算能力があるように感じます。旧スケジューラでは、スロット上限を fully 使わない瞬間があった場合、その計算能力は実質的に失われていました。新しいスケジューラでは、そのような場合、後で割り当て上限を超えてバーストできても、ジョブは迅速に、しかもプリエンプションなしでスケジュールされるため、実質的にその計算能力を取り戻せます。私たちのワークロードはしばしばバースト型なので、これによって多くの計算能力を取り戻せました。」 — Chris Clark
スケジューラは2種類の稼働率を区別します。 割り当て済み稼働率 とは、ワークロードが予算に課金されている時間のことです。これはワークロード所有者の割り当てから差し引かれ、フェアシェアの予算計算に影響し、これらのワークロードは最低実行時間ウィンドウの間プリエンプションから保護されます。 未割り当て稼働率 はどの予算にも課金されず、当初から保護されておらず、割り当て済みの要求によってプリエンプションされる可能性があります。これにより、割り当てが需要と正しく一致していない場合でもGPUを完全に稼働させ続けることができ、チームが無料のGPUサイクルを拒否することも防げます。
スケジューリングの契約
フェアなリソース割り当てを困難にする分散トレーニングのさらなる特徴は、ワークロードが非常に長時間実行されうることです。トレーニングジョブは数時間、数日、時には数週間にわたって実行されることが珍しくありません。一度スケジュールされると、ワークロードは1週間以上にわたって割り当てられたGPU上に留まり続け、他の者が予算化された時間を受け取る機会を奪ってしまいます。これこそがGPUの居座りを可能にしたシステムの特性です。また、これがオンコールエンジニアに、継続的なメンテナンス問題への対処のため、長時間実行ジョブの所有者と交渉することを強いていたのです。
これらの問題に対処するため、私たちは「スケジューリング契約」を導入しました。クラスタへのアクセスと引き換えに、ワークロードは最低実行時間、つまり意味のある進捗を得るために必要な最短の占有時間を宣言しなければなりません。この期間中、ワークロードはプリエンプションから保護されます。これにより、研究者には進捗の保証が与えられ、スケジューラにはその進捗が確保された後にもう一度再調整する権限が与えられ、再開可能なワークロードは自動的にキューに再登録されます。あるいは、ユーザーは最低実行時間を0に設定することもできます。これはGPU時間を解放すべきであることを示します。これらのワークロードは常にプリエンプションの対象となりますが、どの予算にも請求されないという意味で、無料でもあります。
ワークロードのライフサイクルは以下のパターンに従います:
- ワークロードが最低実行時間付きで提出され、再開可能かどうかを示します。
- ワークロードはフェアシェアアルゴリズムに従ってスケジュールされます。その際、ルックバックウィンドウ内の割り当て時間に対する実際の占有率の比率が重み付けに使われます。
- ワークロードは最低実行時間の間実行され、この時間はそのアロケーションに請求されます。
- ワークロードは、関連するアロケーションが他よりも優先され続ける限り、実行を続けることができます。この時間もそのアロケーションに請求されます。
- プリエンプトされて再キューイングされることがあり、その場合はステップ2に戻ります。
- ワークロードが完了し、リソースへの権利を解放します。
これらの取り決めにより、スケジューラにタイムスライシングが追加されました。実行中のワークロードは自動的に取り除かれて再キューイングされ、フェアシェアの収束が可能になり、リソースの独占が抑制されます。また、不健全なホストがワークロードを最低実行時間に達した時点で排出できるようにし、修復活動を完全に自動化できるようにします。この最後の点は、この作業を計画していた時点で私たちが認識していた以上に重要でした。これにより人間の介入を要する修復が74%削減され、オンコールの負担が大幅に軽減されました。
シミュレーション
スケジューリングポリシーの変更は意図しない結果をもたらす可能性があることを私たちは知っています。この問題のゼロサム性により、ある研究者に時間を与えることは別の研究者から取り上げることを意味します。この交換で損をしたユーザーは、新しい回避策を探しがちです。予算ベースのシステムを展開する前に、より長い待ち時間がどこで発生するかを予測し、ルックバックウィンドウの長さや最低実行時間として許容する最大値(私たちは8時間を選びました)といった構成ノブをテストする高速な方法が必要でした。
私たちは、一連のワークロードとその提出スケジュールを入力として受け取り、スケジューラがプリエンプションとGPU割り当ての決定を行える小さなシミュレーション環境を構築しました。各ワークロードの要求GPU数と総実行時間の情報をもとに、シミュレータはスケジュール可能な時点へジャンプし、多くのシミュレーション日数分のキューウェイト時間、プリエンプションイベント、プロジェクト間のGPU時間の分布についての分析を数秒で提供できました。私たちは、過去の提出データと、より深く理解したいと考えた構築シナリオの両方に対してシミュレータを実行しました。
私たちがテストしたいと考えていた仮説の一つは「デバッグワークロード」に関するものでした。これらのジョブは少数のGPUと15分以下の最低実行時間を必要とし、これはユーザーがジョブが正常に起動するか、バグや設定ミスにより早期にクラッシュするかを確認するには十分です。私たちは、これらのジョブが、意味のある進捗を得るために多くのGPUと数時間の実行時間を必要とすることが多い大規模なトレーニングワークロードよりも短いキューウェイト時間を得るかどうかを知りたかったのです。直感的には、小さなジョブは大きなジョブよりも多くの場所に収まるため、キューの上位に上がるはずです。しかし、正確なキューレイテンシが重要でした。1、2分の短い待ち時間は新しい開発手法を可能にしますが、10分の待ち時間は実現不可能になります。
私たちのシミュレーションには手作業で構築したテストケースデータが必要でした。なぜなら、過去の記録にはこの種のデバッグ的なワークロードが十分な量含まれていなかったからです。結果は仮説を裏付け、デバッグワークロードのp90待ち時間が約6時間からわずか5分に短縮されることが示されました。
ベースライン(左)と新しい「アロケーション」スケジューラ(右)の小規模なシミュレータの可視化。各行が1つのGPU、各バーが1つのジョブで、チームごとに1つの色相で親ワークロードにより色分けされています。ハッチングはジョブが中断可能な時間を示し、赤い縁はプリエンプションを示します。ベースラインでは、長い緊急ジョブは決して中断されず、低い優先度レベルの作業ではプリエンプションが少なくなっています。新しいスケジューラでは、各GPUに より多くの色が混在し、チーム間の占有率のローテーションが示されています。
結果
シミュレーション結果を手に、私たちは7月末にクラスタごとの展開を開始しました。私たちが関心を持つ結果は、資金を提供すると決めたワークロードが時間を確保できたか、新しいシステムが完全な占有率を維持したか、そして研究者がスケジューラを理解して情報に基づいた意思決定ができたかどうかです。
展開以来、ユーザーとチームが一貫して割り当てられたGPU時間を受け取っていることを確認しています。私たちはチームに支払われるべき時間を、その実際の需要によって1時間ごとに上限を設けたアロケーションとして数えます。30日間のテスト期間中、チームには支払われるべきGPU時間の98%が提供され、15チーム中13チームのアロケーションが95%以上を受け取り、最悪のケースでも90%でした。クラスタの占有率は変更前後を通じて98%で安定しており、両期間ともに需要は供給を2-3倍上回っていました。提供されたGPU時間の18%は未割り当てであり、これが資金提供されたユースケースが実行準備できていない期間に高い占有率を維持する方法でした。
シミュレータの結果は、実際の結果と比べて方向性としては正確であり、実際の成果は私たちの予測を上回るものでした。デバッグワークロードのp90キュー待ち時間は、新しいスケジューラの下で2時間から30秒に短縮されました。これは、手作業で作成したテストシナリオに基づくシミュレーション予測の6時間から5分という結果を上回っています。なお、ベースラインにおけるデバッグワークロードのサンプルサイズが小さかったため、これらの測定値にはより高い分散があったことに留意する価値があります。キューのレイテンシは全体的に、タイムスライシングの副次的効果として改善されました。最大のH100クラスタでは、キュー待ち時間の中央値は5分から24秒に短縮され、p90待ち時間は約3分の1短縮されました（2.8時間から1.8時間へ）。
私たちが解決しようとした3つの問題と比較すると:
スクワッティング（占有）: 短いデバッグワークロードは1分以内に開始されるようになり、占有行動の価値が低下しました。この行動のコストは占有者の予算に課金されるため、本当に必要なときに時間を得られなくなる可能性があります。
優先度のインフレ: 引き続きワークロードが優先度を宣言することは可能ですが、それはチーム内でのソートにのみ影響します。マネージャーは、予算の使用を最適化するために、グループ全体の優先度を監視するインセンティブを持ちます。
オンコールの負担: 異常な状態のホストは、ワークロードが最小実行時間に達すると自動的に排水されます。人間の介入を必要とする修復は74%減少しました。
課題
学習曲線は私たちが想定していたよりも急でした。変更は段階的に展開したため、初期には研究者はどのクラスタを対象にするかによって異なる動作を経験しました。さらに、私たちのインターフェースには、意味が変化した古い用語（ワークロード優先度など）が一部残っていました。ドキュメントだけでは混乱は解消されませんでした。効果が あった のは、ライブの説明セッションを実施することでした。研究者が質問できる場を提供し、エンジニアリングチームが実際の例を用いて、スケジューラがどのように、なぜその優先順位付けの決定を下すのかについてより深い説明を行う場となりました。
これは重要な瞬間でした。なぜなら、初期のフラストレーションと民俗学的理論の時代から、現在のモード、すなわち研究グループが自らの実験のGPUニーズについてより頻繁に、より広くコミュニケーションを取るモードへの転換点を示したからです。研究者たちは現在、新しい要求に応じるためにどのようなトレードオフが行われているのかをより明確に理解した上で、予算配分の議論に参加しています。
対面セッションに加えて、ローンチ後に新しい可視化機能を導入し、ユーザーが自分に割り当てられたGPU時間が期待される割り当てにどれだけ近いかを把握できるようにするとともに、ワークロードキューのソートに使用されるメトリックを直接表示するようにしました。これにより、ワークロードがプリエンプトされたときに、その理由を理解するためのシンプルな参照先が提供されました。これらの可視化は、管理下のさまざまなプロジェクトでGPU時間がどのように使用されているかを確認できる予算所有者にも役立ちました。
時間経過に伴う割り当て使用量の可視化の例。
すべてのユースケースが改善されたわけではありません。分散型トレーニングに加えて、私たちの研究者は、データ分析を行い、トレーニングコードを書きながらテストする対話型セッションを起動します。旧システムでは、研究者はそのようなセッションを最大1週間保持できました。タイムスライシングでは、保護されたランタイムの8時間の上限が適用され、その後、セッションが割り当てを超過するとプリエンプト可能になります。研究者がこれらのセッションの揮発的な状態にどれほど依存していたかを私たちは認識していませんでした。プリエンプトされるということは、新しいセッションを確保するのを待つだけでなく、状態を手作業で再構築することも意味しました。研究者に調査を行い、この問題の広がりを理解した後、2つの新しいロードマッププロジェクトを作成しました。私たちは、データ準備タスクに焦点を当てた開発セッションのために、オンプレミスストレージの隣にCPU専用クラスタへ投資しています。これにより、本当に必要なワークロードのためにトレーニングクラスタの容量を確保できます。さらに、これらのCPU専用ワークロードのために復元可能なセッションを構築する計画です。これにより、メンテナンスやタイムスライシングのために最小実行時間の終わりにワークロードをプリエンプトし続けることができると同時に、研究者が再構築することなくセッションを別の場所で復元できるようになります。私たちは、この新しいシステムの運用上およびスケジューリング上の利点を保持しながら、ユーザー体験を向上させることができます。
私たちは新たな問題にも引き続き注意を払っています。現在調査中の潜在的な問題の1つは、容量の断片化であり、これにより最大のワークロードのキュー待ち時間が増加する可能性があります。私たちの推測では、最小実行時間の保護が、以前はチームの同時GPU制限を超えるためにプリエンプト可能な メカニズム に依存していた種類のジョブに適用されているということです。以前は、それらのジョブはいつでも中断される可能性があり、時間が無駄になる可能性がありましたが、大きなジョブのスケジューリングも容易になっていました。今では、スケジューラが保留中の大きなワークロードを配置するために多くのジョブを一度に中断する機会が少なくなっている可能性があります。現在、シミュレータツールを使用してこの問題を再現しながら、本番環境での実測値も測定しています。
今後の展望
ここで説明したスケジューリングの取り組みを超えて、私たちはピラミッドの頂点、すなわち利用率を目指しています。ブートストラップ、チェックポイント、そしてトレーニングアプリケーション自体が、すべて可能な限り効率的に行われるようにし、各ワークロードが受けるスケジュールされた時間の価値を最大化する必要があります。
このような課題に研究者と緊密に協力して取り組みたい方は、 Ai2の募集中のエンジニアリング職をご覧いただくことをお勧めします.