## Memo
## Memo with LLM
### 論文情報
- **タイトル**: CASSINI: Network-Aware Job Scheduling in Machine Learning Clusters
- **著者**: Sudarsanan Rajasekaran(MIT)、Manya Ghobadi(MIT)、Aditya Akella(UT Austin)
- **カンファレンス**: 21st USENIX Symposium on Networked Systems Design and Implementation (NSDI '24)
- **発表年**: 2024年(Santa Clara, CA, USA, April 16–18, 2024, pp. 1403–1420)
### 論文概要
CASSINIは、機械学習(ML)クラスタ向けのネットワーク対応ジョブスケジューラである。DNN訓練ジョブの周期的な通信パターン(Up-Down フェーズ)に着目し、幾何学的抽象化とAffinityグラフを用いて、同一ネットワークリンクを共有するジョブの通信フェーズを時間シフトによってインターリーブさせる。24サーバのテストベッドで13種のMLモデルを用いた実験において、既存スケジューラと比較して平均および末尾完了時間をそれぞれ最大1.6×・2.5×改善し、ECNマークパケット数を最大33×削減することを示した。
### 詳細解説
![[Pasted image 20260401180637.png]]
![[Pasted image 20260401180732.png]]
#### 問題設定
- **入力**: MLクラスタ上で動作する複数のDNN訓練ジョブのセット。各ジョブは複数のGPUサーバにまたがり、データ並列・モデル並列・ハイブリッド並列など様々な並列化戦略を使用する。
- **出力**: 各ジョブのGPUサーバへの配置(placement)と、各ジョブの訓練イテレーション開始を遅延させる時間シフト値 {t_j}。
- **前提条件**:
- DNN訓練ジョブの通信パターンはイテレーション間で繰り返され、定常的である。
- ネットワークスイッチ/NICに予約・優先度などの特殊サポートは不要。
- 輻輳制御プロトコルの変更は不要。
- **課題**: 既存のMLスケジューラ(Themis、Pollux等)はジョブの通信パターンを無視してGPU配置を決定するため、同一リンクを共有するジョブ間でネットワーク輻輳が発生し、訓練時間が増大する。
#### 提案手法
##### アーキテクチャ
CASSINIは既存MLスケジューラへの「プラガブルモジュール」として動作する。主な構成要素は以下の2つである。
1. **Cassini Module(中央スケジューラ側)**: 複数の配置候補を受け取り、幾何学的抽象化とAffinityグラフに基づいてcompatibility scoreを計算し、最適配置と時間シフトを決定する。
2. **Agent(各サーバ側)**: 受け取った時間シフト値を適用し、ドリフト(ノイズやストラグラー起因のズレ)を継続的に監視・補正する。
##### アルゴリズム/手法の詳細
**1. 幾何学的抽象化(Geometric Abstraction)**
各ジョブの訓練イテレーション時間 T を円の円周とし、通信需要の時系列データをその円上にマッピングする。Up フェーズ(高帯域)とDown フェーズ(低帯域)が円の特定の弧(arc)として表現される。
(Figure 3: CASSINIの幾何学的抽象化。VGG16の訓練イテレーション時間255msを円周として、通信需要を円上に表現する。)
- 異なるイテレーション時間を持つジョブ間の比較には、全ジョブのイテレーション時間のLCM(最小公倍数)を円周とする**統一円(Unified Circle)**を使用する。
- 2つ以上のジョブが同一リンクを共有する場合、各ジョブの統一円を重ね合わせ、回転角度 Δ_j を最適化することで帯域需要のインターリーブを実現する。
**2. 最適化定式化(Table 1)**
- **目的関数**: `score = 1 - Σ_α Excess(demand_α) / (|A| × C_l)` を最大化
- `Excess(demand_α) = max(demand_α - C_l, 0)`:リンク容量超過分の帯域需要
- **制約**: `demand_α = Σ_j bw_circle_j(α - Δ_j)`(全ジョブの帯域需要の和)
- **出力**: 各ジョブの回転角度 Δ_j(時間シフト `t_j = (Δ_j / 2π × p_l) mod iter_time_j` に変換)
**3. Affinityグラフによるクラスタレベルへの拡張**
大規模クラスタでは各ジョブが複数リンクを通過し、異なるリンク上で異なるジョブと競合する。この際、単一ジョブに対して複数の時間シフト値が算出されるという問題が生じる。
CASSINIは二部グラフ G = (U, V, E) を導入する:
- U:他ジョブとリンクを共有するジョブの頂点集合
- V:複数ジョブが通過するリンクの頂点集合
- E:ジョブ j がリンク l を通過する場合のエッジ(重み = t_j^l)
**BFSベースのグラフ探索アルゴリズム(Algorithm 1)**: ランダムに1つのジョブを基準点(t_j = 0)として選択し、BFS探索で他の全ジョブに対して一意な時間シフトを計算する。ループを含むAffinityグラフの配置候補は除外される(Theorem 1により、ループなしのグラフに対して正確かつ一意な時間シフトが保証される)。
##### 実装上の工夫
- **並列処理**: Algorithm 2では、各配置候補および各リンクに対する最適化計算をマルチスレッドで並列実行する。
- **ドリフト補正**: 各サーバのエージェントがイテレーション時間の変動を継続的に監視し、時間シフトのズレを補正する(実験では補正が稀であることを確認)。
- **既存スケジューラへの統合**: ThemisとPolluxへの統合はそれぞれ約1,000行のコード追加で実現。スケジューラのアービタを、単一の配置決定ではなくN個の候補配置を返すよう変更する(Themisでは約300行の修正)。
#### 新規性
**既存手法の課題**:
- Themis・Polluxなどの既存MLスケジューラは、ジョブの通信パターンを無視したGPU配置を行う。
- Themisはラック内/ラック間のスローダウンペナルティでネットワーク局所性を考慮するが、ジョブ間のトラフィックパターンの干渉は考慮しない。
- 単純に全ジョブのリンクレベルの制約を1つの最適化問題に統合するアプローチは、LCMが巨大になるため指数的な計算量増加を招き非現実的である。
**CASSINIの解決策**:
- 幾何学的抽象化により、任意の並列化戦略(データ並列・モデル並列・ハイブリッド並列)のトラフィックパターンを統一的に扱える。
- AffinityグラフとBFS探索により、クラスタレベルで一意な時間シフトをスケーラブルに計算する。
- スイッチ/NICの特殊サポートや輻輳制御の変更が不要で、プラガブルモジュールとして既存システムに統合可能。
#### 実験設定
- **実験環境**: 24台のASUS ESC4000A-E10サーバ(各1台のNVIDIA A100 GPU・40GB HBM2、1台の50Gbps Mellanox ConnectX5 NIC)。通信はRoCEv2(RDMA over Converged Ethernet)。DCB・PFCによるロスレスファブリック。Tofino スイッチで13個の論理スイッチ・48本の双方向リンクを持つ2:1オーバーサブスクライブドトポロジーをエミュレート。Ubuntu 18.04 LTS、PyTorch 1.8.0、CUDA 11.1、NCCL 2.11.4。
- **データセット**: 13種のDNNモデル(VGG11・VGG16・VGG19・ResNet50・WideResNet101・BERT・RoBERTa・XLM・CamemBERT・GPT-1・GPT-2・GPT-3・DLRM)。各モデルは均等な確率で出現し、訓練期間は200〜1,000イテレーションからランダム選択。
- **比較対象(Baseline)**: Themis(finish-time fairness)、Pollux(goodput最大化)、Ideal(専有クラスタ)、Random(ランダム配置)。
- **評価指標**: 訓練イテレーション時間のCDF(平均・99パーセンタイル尾端)、ECNマークパケット数(ネットワーク輻輳の指標)。
- **トレース種別**: Poissonトレース(ジョブ到着率を変化、80〜100%負荷)、Dynamicトレース(クラスタ動的変化)、Snapshotトレース(特定時点でのクラスタ状態)。
#### 実験結果
(Figure 1: GPT-1(データ並列)、GPT-2(パイプライン並列)、GPT-3(テンソル並列・ハイブリッド並列)の訓練時のトラフィックパターン。各並列化戦略でUp-Downフェーズの形状が異なる。)
##### 定量的評価
- **データ並列ジョブ(Poissonトレース)**: Th+CASSINIはThemisと比較して、平均イテレーション時間を1.6×、99パーセンタイル尾端を1.8×改善。Th+CASSINIはIdeal(専有クラスタ)とほぼ同等の性能を達成。
- **モデル並列ジョブ(Poissonトレース)**: Th+CASSINIはThemisと比較して平均1.2×、99パーセンタイル尾端1.6×改善。
- **Dynamicトレース(ECN削減)**:
- DLRM:ThemisとPolluxは各々Th+CASSINIおよびPo+CASSINIの27×・33×のECNマークを記録。
- RoBERTa:ThemisはTh+CASSINIの3.6×のECNマーク。
- GPT-2:Themisは29.1×のECNマーク。
- 平均・尾端イテレーション時間:Th+CASSINIはThemisより1.5×・2.2×改善、Po+CASSINIはPolluxより1.6×・2.5×改善。
##### アブレーションスタディ
- **部分的互換性(Partial Compatibility)**: Snapshotトレースでcompatibility score 0.6〜1.0の様々なシナリオを検証。スコアが高いほどインターリーブが良好で輻輳が少なく、スコアが低くても完全非互換より改善される。
- **複数GPU/サーバ構成**: 複数GPU/サーバ環境でも同様の性能向上を確認。
- **時間シフト補正の頻度**: 実験全体を通じて時間シフトの補正はごく稀(実用上の問題なし)。
#### 考察 (Discussion)
**結果の解釈**:
DNN訓練の通信パターンがイテレーション間で高い規則性を持つという性質を活用することで、ジョブ配置時に将来の通信競合を予測・回避できる。幾何学的抽象化は、円の回転という直感的な操作で複数ジョブのインターリーブを定量化するため、複雑な通信パターンにも適用可能である。
**優位性の根拠**:
既存スケジューラがジョブのGPU数決定(リソース配分)とサーバ配置を同時に最適化しようとするのに対し、CASSINIはリソース配分はThemis/Polluxに委ね、サーバ配置の選択のみを担当する分離設計により、既存システムへの統合コストを最小化しつつネットワーク性能を改善する。
**限界と例外**:
- Affinityグラフにループがあるジョブの組み合わせは、一意な時間シフトの保証ができないため除外される。
- LCMが非常に大きくなる場合(イテレーション時間が互いに素な関係)、統一円が非常に大きくなり最適化の計算量が増大する可能性がある。
- 時間シフトのドリフト(ノイズ、ストラグラー)が大きい環境では補正コストが増大する可能性がある。
- クラスタ全体での最適配置保証はなく、上位N候補から選択するヒューリスティック的アプローチである。
#### 強み (Strengths)
- 既存スケジューラへの変更を最小限(約1,000行)に抑えたプラガブル設計で実用性が高い。
- 幾何学的抽象化がデータ並列・モデル並列・ハイブリッド並列のあらゆる通信パターンを統一的に扱える汎用性。
- スイッチ/NICの特殊機能が不要で、既存のRDMAインフラに直接適用可能。
- 実機24サーバ・13モデルの包括的な評価で実用的な性能改善を実証。
#### 弱点・課題 (Weaknesses / Limitations)
- AffinityグラフのループによりN候補が全て除外される場合、CASSINIはベーススケジューラの選択にフォールバックする(compatibility scoreが考慮されない)。
- 訓練ジョブのプロファイリング(通信パターンの計測)が事前に必要であり、新規モデルや動的なハイパーパラメータ変更に対応するコストがある。
- 最適化問題(Table 1)の離散化精度(デフォルト5度)によって解の品質にトレードオフがある。
- 24サーバ規模での評価に限定されており、数百〜数千サーバの実運用規模でのスケーラビリティ検証が課題として残る。
## Abstract
我々はCASSINI、機械学習(ML)クラスタ向けのネットワーク対応ジョブスケジューラを提案する。CASSINIは、ジョブをネットワークリンクに配置する際に、異なるジョブの通信パターンを考慮するための新しい幾何学的抽象化を導入する。そのために、CASSINIはAffinityグラフを用いて一連の時間シフト値を求め、ジョブのサブセットの通信フェーズを調整することで、同一ネットワークリンクを共有するジョブの通信パターンが互いにインターリーブされるようにする。24サーバのテストベッド上で13の一般的なMLモデルを用いた実験により、最先端のMLスケジューラと比較して、CASSINIはジョブの平均および末尾完了時間をそれぞれ最大1.6×および2.5×改善することを示す。さらに、CASSINIはクラスタ内のECNマークパケット数を最大33×削減することを示す。