[paper]
概要
どんなもの？
- 与えられたトポロジに適した集団通信アルゴリズムを求める
- 全域木パッキング問題を解くことで、帯域を有効活用するbroadcastの通信スケジュールを生成
- broadcastの通信スケジュールからその他の集団通信も生成
- 求めた通信スケジュールを元にCUDAコードも生成
- NCCL2のリングベースのプロトコルと比較して、Blinkはモデル同期を最大8倍高速化、訓練時間を最大40%短縮
背景
AI向けハードウェア > DGX-1
深層モデル学習用に専用のハードウェアが開発されています。 例えばNVIDIA DGX-1*1では下図のように8つのGPUをNVLINKで結合させています。 ちなみにDGX-1のトポロジでは6つの片方向Ringに分解可能であり、Ringベースのアルゴリズムもbroadcastベースのアルゴリズムのどちらにも強そうな見た目をしていますね。
スケジューラ
課題: 状態の良い（均一的な）リソースをジョブに割り当てられない
DGX-1を使用した環境を考えます。この環境は複数のDGX-1が集まったもので、パブリッククラウドやデータセンターなどを想定しています。ジョブがこの環境に送られると、スケジューラが動作し、必要とされるマシンリソース（具体的にはGPU）の数に応じて、利用可能なGPUを割り当てます。
理想的には、DGX-1のトポロジを最大限に活用し、物理的に近接したGPU群をジョブに割り当てることが望ましいです。しかし、ジョブの待機時間（キューイング遅延）短縮のためにGPUへの割り当てがえいやと行われ、徐々に割り当て可能なGPUが断片化する状態が生じているとのこと。
図は実際にジョブに割り当てられた、1 つのDGX-1上における GPU数の分布。3、5、7といった、いかにもDGX-1のトポロジを活かすのが難しそうな数の割り当ても行われてるようです。
不均一なトポロジにおける通信アルゴリズムを動的に見つける
そのような不均一なトポロジ（言い換えると変な形のトポロジ）において、Ringアーキテクチャ向けにチューニングされたNCCLのような通信アルゴリズムは性能を発揮できません。 なぜならGPU→GPU間の帯域差により、全体のリンク帯域使用率の低下するから。 不均一なトポロジに対してRingベースのアルゴリズムよりも、より良い通信アルゴリズムがあるはず。それを動的に見つけたい、というのがこの論文のモチベーションです。
手法
集団通信を次の二つに分類します。
- one-to-many: ひとつのrootノードから他のノードに情報を配る通信
- many-to-many: 複数の送信元をもち、複数のノードに情報を配る通信
| type | collectives |
|---|---|
| one-to-many | broadcast, scatter, reduce |
| many-to-many | all-reduce, reduce-scatter, ... |
この論文ではまず、
- one-to-many集団通信のアルゴリズムを求める
具体的にはbroadcastの集団通信を求めます。broadcastをベースに他の通信スケジュールを生成します。
- many-to-many集団通信はone-to-manyの結果を結合させて生成(many-to-one(収集) → one-to-many(分配))
と、通信アルゴリズムを求めています。
one-to-many
- 与えられたGPUリソースの接続情報をグラフGで表現(a)
- グラフGのspanning-tree: 全域木を生成(c)
- 複数の木($T_i$)に従ってデータを流す
- 木によって流すデータ量($w_i$)を変える
- リンクをデータ量の総量($\sum_i w_i$)がリンク容量$c_e$を超えてはならない
- sourceから流れるデータ量を最大化させる
これは次の最適化問題で表現することもできます。
近似解法
ただ、全域木は一般に与えられたグラフ上に数多く存在するため、それらを列挙するだけでも大変です。 そのため上記の最適化問題を直接扱うのは難しくなります。そこで、multiplicative weight update (MWU)テクニックを用いて、(1-ε)-approximation解を見つける、という工夫を提案しています。計算量はO(m ln m/ ε2), ここで m は枝数です。
アルゴリズムの流れはこんな感じ。
- 流れているフロー値を枝の重みとする
- 初期値は全ての枝の重み0
- minimum weight spanning treeを見つける
- εフローをその全域木に流して枝のフロー値を更新
- （終了条件を満たすと終了[1]）
[1] Chekuri, Chandra, and Kent Quanrud. "Near-linear time approximation schemes for some implicit fractional packing problems." Proceedings of the Twenty-Eighth Annual ACM-SIAM Symposium on Discrete Algorithms. Society for Industrial and Applied Mathematics, 2017.
全域木の削減
例えば8GPU(DGX-1)のトポロジで計算したところ181全域木がMWUの探索の結果発見され、解に使用されていたとのこと。 このとき、全域木の個数が多いと、ひとつの全域木に流すデータ量が小さくなり、スループットが下がる要因になります。
そこで全域木を減らすために、木に流すデータ量を整数値に限定させてILPを解くという処理も行なっています。 この処理により8GPU(DGX-1)のケースで181全域木を6全域木まで減らすことができたようです。
many-to-many集団通信
リンクは双方向であることを前提とします。まずone-to-manyによって、全域木とその木に流すデータ量を求め、 それを反転させてmany-to-oneによりデータを集積、そして同じ全域木を逆向きに辿ってデータを分配する、 という流れでmany-to-many集団通信を構成しています。これで本当に良いのか？というのが直感的な疑問ですが、後述する実験ではそこそこ良い結果が出ていますし、リンクが双方向であることを前提としてパイプライン処理をゴリゴリ使えば、これでも良い通信アルゴとその実装が行えそうな気がします。
コード生成
得られた通信スケジュールをもとにCUDAコードを生成
パイプライン処理の工夫もそのときに入れる。ワークロードとして同じような集団通信が周期的に行われることを想定。なので、最初の数エポックは情報収集とパイプラインサイズの調整用に使用する。
- チャンクサイズを自動調整、最初、数iterationはチャンクサイズ調整用に使用
- チャンクサイズを小さな値で初期化
- スループットが増加している → チャンクサイズを乗法的増加
- スループットが低下した場合 → チャンクサイズを加算的に減少
実験
不均一なトポロジに対して、提案手法で求めた通信アルゴリズムがどれほど有効かを検証しています。 実験としてはユニークなものを設定します。
DGX-1 8GPUのユニークな割り当てを列挙し、それぞれについて性能をNCCL2と比較しています。例えばGPU(5, 6, 7)を使用した場合のトポロジを抽出し、そのトポロジに対してNCCL通信ライブラリと、この手法で作成した通信アルゴリズムとの比較をしています。通信のtotal data sizeは 500MB。図はBroadcastの結果。
図はAllReduceの結果。
議論はある？
- 今回の提案は8GPU以下のトポロジを扱う、ローカルな範囲の通信アルゴリズムの最適化。実際これが現場でどれほど使えるか？
- 全域木のトポロジからなるデータフローを足し合わせて、通信アルゴリズムを作るのは面白い。
- また得られた通信アルゴリズムの結果をもとにCUDAコードの生成を行なっていますが、そこの部分もきちんと理解すれば勉強になりそう。
- SCCL*2による評価
- TACCL*3による評価
*1:https://www.nvidia.com/ja-jp/data-center/dgx-1/
*2:Cai, Zixian, et al. "Synthesizing optimal collective algorithms." Proceedings of the 26th ACM SIGPLAN Symposium on Principles and Practice of Parallel Programming. 2021.
*3:Shah, Aashaka, et al. "Synthesizing collective communication algorithms for heterogeneous networks with taccl." arXiv preprint arXiv:2111.04867 (2021).
