Ai2（AI2）のAIインフラチームは、GPUクラスタ向けのスケジューラを改善したと発表しました。従来の優先度ベースのスケジューラを、GPU時間予算、階層的フェアシェア割り当て、タイムスライシング契約を含むシステムに置き換えました。これにより、研究プロジェクトが受けるGPU時間の配分を、ケースバイケースの運用タスクから透明な予算編成プロセスに移行したといいます。ロールアウトは7月末から開始されました。
背景：GPU需要が供給を上回る
同社は数千基のNVIDIA H100、B200、B300 GPUを管理しており、クラスタの規模は88〜1024 GPUです。約150人の研究者が利用しており、LLMやVLMのトレーニング、ロボット工学の強化学習シミュレーション、科学エージェント向けのポストトレーニングなど多岐にわたります。需要は供給を2〜3倍上回っており、各GPU時間には2〜3のワークロードが競合しています。
従来の優先度ベースのスケジューラでは、GPUの「居座り（squatting）」や優先度のインフレーションといった問題が発生していました。ユーザーがデバッグ用のワークロードを低レイテンシで起動できず、何もしないワークロードを置いておく「居座り」が見られました。また、優先度がインフレーションを起こし、低優先度のワークロードがGPU時間を得られない状況も発生しました。
新しいスケジューラの仕組み
新しいシステムでは、各研究プロジェクトにGPU時間の予算を割り当てます。マネージャーがプロジェクトや研究者に比例配分し、これはプログラム戦略を直接反映します。例えば、プロジェクトA1は総容量の35%の権利を持ちます。GPU時間のリクエストは予算で賄われる必要があり、予算外のワークロードはプリエンプションから保護されません。
さらに、階層的フェアシェアスケジューラを導入しました。7日間のスライディングウィンドウで占有を追跡し、割り当てを下回っているグループのワークロードを優先します。また、「スケジューリング契約」として、各ワークロードは最小実行時間を宣言します。この間はプリエンプションから保護されますが、実行後は自動的に再キューイングされる可能性があります。これにより、時間スライシングが実現し、GPUの居座りを抑止します。
シミュレーションと実測結果
同社はシミュレーション環境を構築し、履歴データと構築したシナリオでテストしました。デバッグワークロード（少ないGPU数、最小実行時間15分以下）の待ち時間を予測したところ、p90で6時間から5分に短縮されると予測しました。実測では、デバッグワークロードのp90待ち時間は2時間から30秒に短縮されました。シミュレーションは方向的に正確で、実際の結果は予測を上回りました。
30日間のテスト期間では、チームは受け取るべきGPU時間（割り当てを実際の需要で上限処理したもの）の98%を受け取りました。15チームのうち13チームが95%以上、最悪のケースでも90%でした。クラスタの占有度は変更前後で98%を維持しました。配分されたGPU時間の18%は未割り当てで、需要が低い期間に占有度を維持するために使用されました。
また、この変更により、修理が必要なホストのドレインが自動化され、人手による介入が必要な修理が74%削減されました。
どう見るか
今回の発表では、スケジューラの設計思想と実測結果が示されています。一方で、実装コードは公開されていません。数値はすべて同社の測定によるもので、第三者による検証は示されていません。同社は、対話型の開発セッションが8時間の保護時間を過ぎるとプリエンプトされうるようになり、使い勝手が悪化したことも課題として挙げています。今回の結果は同社の環境に基づくものです。
