PFNがAI計算基盤向けの超高効率スケジューラを開発
- •PFNがAI学習の実行効率を最大化するカスタムK8sスケジューラのPoCを公開
- •締切時間を考慮したLSTアルゴリズムによりデッドライン遵守率を大幅向上
- •数週間の学習シナリオを仮想時間で高速検証する評価ツールを独自に開発
Preferred Networks(PFN)が、AI計算基盤の効率を極限まで高めるためのKubernetesカスタムスケジューラと、その評価ツールの開発成果を公開しました。AIモデルの学習には膨大な計算リソースと時間が必要ですが、従来の標準的なスケジューラでは、多数のジョブが混在する複雑な環境下で個別の優先度や完了期限を最適に扱うことが困難でした。そこで同社は、ソフトウェア側からAI計算基盤と協調動作し、リソースのスケジューリングを最適化する新たなアプローチを提案しています。
今回開発されたスケジューラの核となるのは「Least Slack Time (LST)」スケジューリングです。これは各ジョブの実行予測時間とデッドラインから「余裕時間(Slack time)」を算出し、猶予の少ないジョブを優先的に実行する仕組みです。このアルゴリズムにより、標準スケジューラでは実行待ちのまま期限切れになりがちだった後続の小規模ジョブも、プリエンプション(実行中断と再開)を駆使して期限内に完了させることが可能になりました。デッドラインという定量的な指標を用いることで、主観的な優先度設定に頼らない効率的な運用を実現しています。
さらに、分散学習に不可欠な「Gang Scheduling」や、アクセラレータの断片化を防ぐ独自のビンパッキング機能も実装されています。特に興味深いのは、ノード内の故障等で利用可能なチップ数が異なる状況でも、リソースの断片化を最小限に抑えるように設計されている点です。これにより、ハードウェアの不確実性が高い大規模環境でも高い稼働率を維持できます。
また、実運用における最大の障壁となる「評価の長期化」を解決するため、評価ツール「kube-scheduler-evaluator」も開発されました。これは実際の物理リソースを消費せず、内部的な仮想時間を進めることで、数週間に及ぶジョブの実行シナリオを数分でシミュレーション可能です。Go言語による柔軟なシナリオ定義に対応し、OSSとして公開されることで、AIインフラ開発の民主化と加速に大きく寄与することが期待されます。