Ai2、GPUクラスタのスケジューラを優先度ベースから時間予算制へ刷新

何が変わったのか

Ai2のAIインフラチームは、大規模な分散学習向けにGPU計算資源を提供している。同チームは運用指標を4段階のピラミッドで捉えており、土台が「可用性」(ハードウェアが健全で作業可能な割合)、その上が「占有度」(稼働可能時間のうち特定ワークロードに割り当てられた割合)、次が「インパクト」(最も価値の高いワークロードに資源が渡る頻度)、頂点が「利用率」(ワークロードの寿命に対してGPU容量が使われた割合)である。記事はこのうち「インパクト」、つまりスケジューリング判断の質を高めた取り組みを扱っている。

同チームは優先度ベースのスケジューラを置き換え、GPU時間の予算、階層型フェアシェア割り当て、タイムスライス契約を組み合わせた仕組みを導入した。結果として、各研究プロジェクトにどれだけGPU時間を割くかという議論が、個別案件ごとの運用対応から透明性のある管理会計的なプロセスへ変わったという。対象はNVIDIAのH100、B200、B300で構成される数千台のGPUで、クラスタ規模は88台から1024台。約150名の社内研究者が、LLMやVLMの学習、ロボティクスの強化学習シミュレーション、科学エージェント向けのポストトレーニングなどに利用している。

旧方式が抱えた問題

需要は供給を大きく上回っており、提出されたワークロードに基づくと常時2〜3倍のGPUが要求されている状態だ。稼働可能なGPU1時間に対して、2〜3件の研究ワークロードが競合していることになる。旧方式では優先度ベースのスケジューラを使い、ワークロード側がプリエンプション(実行中断による資源の譲渡)の対象外を選べるようにしていた。保護されたワークロードには同時利用GPU数の上限があり、保護されないワークロードは空いているGPUでその上限を超えられた。

この設計は「GPUの占有」を招いた。必要になったときに接続できるよう、何もしないワークロードを置き続ける行為である。背景には、デバッグ用ワークロードを低遅延で起動できないという研究者側の事情があった。さらに優先度インフレが起き、最終的にはスケジュールされたワークロードの100%がHIGH優先度を使うようになり、下位の優先度はGPU時間を得られなくなった。プリエンプション除外が任意だったため、オンコール担当エンジニアのチケット対応時間の大半が、保守対象ホスト上で動く中断不可ワークロードをどう止めるかの交渉に費やされたという。

予算とフェアシェア

共有資源が枯渇する「コモンズの悲劇」への古典的な対処は、資源を私有化することだ。Ai2は以前、チームにGPU群の独占を割り当てていたが粒度が粗すぎた。研究には季節性があり、チームごとに実験や学習を回す時期が異なるため、独占を割り当てると遊休する時間帯が生じ、別のチームが待たされる。動的に変わる研究ニーズを静的なスケジュールにはめ込む、手作業のナップサック問題になっていた。

そこでGPUそのものではなくGPU時間の一部を割り当てる方式に変えた。将来の需要予測は新規の科学実験の結果次第で正確にはできないが、研究努力間の優先順位は戦略の問いであり、事前に議論して決められる。経営層が投資家のように振る舞い、ワークロードが生まれる前にGPU時間を配分する。管理者が配下のプロジェクトや研究者へ比例配分する階層構造で、例えばプロジェクトA1は他プロジェクトの待ち状況にかかわらず総容量の35%を主張できる。

実際の占有管理には階層型フェアシェアスケジューラを組み合わせた。アルゴリズム自体は新しくなく、2009年のHadoop Fair Schedulerに遡る系譜で、現在もSLURMのFair TreeやYARNのFair Schedulerで使われている。新しいのは入力であり、ツリーが研究プログラム構造を反映し、重みが静的なクォータではなく管理者が設定する予算になっている。既定7日のスライディングウィンドウで占有状況を追跡し、割り当てを使い切っていないグループのワークロードを、使い過ぎのグループより上位に並べる。

契約と運用効果

分散学習のワークロードは数時間から数週間動き続けることがあり、それが占有の温床になっていた。そこで「スケジューリング契約」を導入する。クラスタへのアクセスの対価として、ワークロードは最低実行時間、つまり意味のある進捗に必要な最短の占有時間を宣言する。その間はプリエンプションから保護され、進捗が確保された後はスケジューラが再配置でき、再開可能なワークロードは自動で再キューイングされる。最低実行時間をゼロに設定すれば、その分は予算に課金されず、常にプリエンプション対象となる。

これによりタイムスライスが実現し、実行中のワークロードが自動で退避・再キューされてフェアシェアが収束する。健全でないホストは、ワークロードが最低実行時間に達した時点でドレインでき、修理作業を自動化できる。この効果は想定以上だったとされ、人手を介する修理が74%減少し、オンコールの負荷が大幅に減った。記事では、政策変更が意図しない結果を生みうるとして、ゼロサム性を踏まえたシミュレーションも実施したと述べられている。