# TACCL: Guiding Collective Algorithm Synthesis using Communication Sketches > [!abstract] 概要(arXiv abstract の日本語訳) > 機械学習モデルは複数のGPUとサーバをまたいで訓練されることが増えている。この設定では、A LLTOALLやA LLREDUCEのような通信集合演算(communication collectives)を使ってGPU間でデータが転送されるが、これは大規模モデルの訓練において重大なボトルネックになりうる。したがって、集合通信のために効率的なアルゴリズムを使うことが重要である。私たちはTACCLを開発した。これはアルゴリズム設計者が、与えられたハードウェア構成と通信集合演算に対して合成器(synthesizer)を導き、アルゴリズムを自動生成させることを可能にするツールである。TACCLは新規の通信スケッチ(communication sketch)抽象を用いて、設計者から重要な情報を取得し、探索空間を大幅に削減して合成器をより良いアルゴリズムへ導く。TACCLはまた、単一ノードのトポロジを超えてスケールすることを可能にする、問題の新規エンコーディングを使う。私たちはTACCLを使って、3つの集合演算と2つのハードウェアトポロジ(DGX-2とNDv2)に対するアルゴリズムを合成した。TACCLによって合成されたアルゴリズムが、Nvidia Collective Communication Library(NCCL)を最大6.7倍上回ることを示す。また、TACCLがTransformer-XLとBERTモデルのエンドツーエンド訓練を、バッチサイズに応じて11%〜2.3倍高速化できることも示す。 ## 論文情報 - タイトル: TACCL: Guiding Collective Algorithm Synthesis using Communication Sketches - 著者: Aashaka Shah(UT Austin)・Vijay Chidambaram(UT Austin and VMware Research)・Meghan Cowan・Saeed Maleki・Madan Musuvathi・Todd Mytkowicz・Jacob Nelson・Olli Saarikivi(いずれも Microsoft Research)・Rachee Singh(Microsoft and Cornell University) - 媒体: NSDI 2023(USENIX Symposium on Networked Systems Design and Implementation) - arXiv ID: 2111.04867(v4, 2022年10月5日改訂) - コード: https://github.com/microsoft/taccl(ランタイム部分は https://github.com/microsoft/msccl) ## 概要 GPU間の通信集合演算(collective communication)を実装するアルゴリズムは、レイテンシとバンド幅のトレードオフやトポロジの異質性を扱う必要があるため設計が難しい。TACCLは、アルゴリズム設計者から4種類の低労力な入力(論理トポロジ・switch-hyperedge方針・アルゴリズムの対称性・入力サイズ)を「通信スケッチ」として受け取り、それによって探索空間を絞ったうえでMILP(混合整数線形計画)ベースの合成器にルーティングとスケジューリングを解かせる。この三段階(ルーティング緩和・ヒューリスティックな順序付け・連続化と厳密スケジューリング)の分割によって、単一ノードに限られていた先行研究SCCLを多ノードのGPUトポロジ(DGX-2・NDv2)へスケールさせた。 ## 問題設定 入力は通信集合演算の種別(A LLGATHER・A LLTOALL・A LLREDUCE)と対象ハードウェアトポロジ、およびアルゴリズム設計者が与える通信スケッチである。出力は、各データチャンクがどの経路をどの順序で転送されるかを規定する具体的な集合通信アルゴリズムである。 **Figure 2: AllGather / AllToAll / AllReduce の初期・最終データバッファ** ![[_attachments/arxiv-2111.04867/fig02-collectives.png]] (Figure 2. 4つのGPU間で行われる各集合演算の初期・最終データバッファを図示する。A LLGATHERでは各GPUが他の全GPUのデータバッファを受け取り、A LLTOALLでは各GPUが全GPU上のバッファの異なるチャンクを受け取ることでバッファインデックスとGPUインデックスが転置され、A LLREDUCEでは各GPUが全GPUの同一データインデックスに対する点ごとの計算結果(例: 和)を受け取る。) 制約となる背景として、GPU内はNVLink(最大300 GBps)、GPU間はInfiniBand(12.5〜25 GBps/NIC)という帯域の異質性があり、NCCLのようなライブラリはトポロジに関わらずRing/Tree等の事前定義テンプレートを使うため、この異質性を活かせない。 ## 提案手法 - **通信スケッチ**: プログラムスケッチング(Solar-Lezamaの技法に着想)を通信集合演算合成に応用した抽象。設計者は次の4点だけを指定する: (1) 物理トポロジの部分集合として論理トポロジを指定する(低帯域リンクの過剰使用を避ける)、(2) スイッチ(NVSwitch・IBSwitchなど)をswitch-hyperedgeとして注釈し、接続数の最大化(uc-max)・最小化(uc-min)・自由選択のいずれかの方針を与える、(3) トポロジと集合演算の対称性に基づくアルゴリズム対称性を与える(automorphismと分割の指定)、(4) 想定データサイズを指定する。ルーティングは人間の直感が働きやすい一方、スケジューリングはルートが交差する場合のみ関係が生じ煩雑なため合成器に委ねる、という役割分担がスケッチ設計の核心である。 **Figure 1: TACCLの合成器パイプライン概要** ![[_attachments/arxiv-2111.04867/fig01-overview.png]] (Figure 1. 通信スケッチ・プロファイル済みトポロジ・対象集合演算とハイパーパラメータを入力とし、ルーティング→ヒューリスティックな順序付け→連続化と厳密スケジューリングの3段の合成器を経て、アルゴリズムをバックエンド(TACCLランタイム)に実装する。) - **switch-hyperedgeと対称性**: 完全な二分バンド幅を持つスイッチファブリック(DGX-2/NDv2のNVSwitch・IBSwitch)では、単一GPU/NICからの接続数が増えるほどキューイング遅延でレイテンシが増す(Figure 4)。switch-hyperedgeはスイッチを直接リンクの集合に置き換え、接続数を制御する。小データサイズではuc-max(輻輳の可能性が低い)、大データサイズではuc-min(輻輳を避けるためRing相当に縮退)が望ましい(Figure 3)。 **Figure 3: switch-hyperedge方針の効果** ![[_attachments/arxiv-2111.04867/fig03-switch-hyperedge.png]] (Figure 3. (a) スイッチで接続された3GPUの物理トポロジ。(b) 最大接続数方針では全GPU間の接続を維持する論理トポロジになる。(c) 最小接続数方針ではRing相当の論理トポロジに縮退する。) **Figure 4: GPU近傍数とデータ量を変えたマルチ接続の帯域** ![[_attachments/arxiv-2111.04867/fig04-multiconnection.png]] (Figure 4. 左: DGX-2ノード内のNVSwitch経由、右: 4つのDGX-2ノード間のIBSwitch経由で、接続数を1〜16まで変えて計測したバンド幅。接続数が増えるほど帯域が低下する傾向が両ケースで見られるが、小データサイズでは差が小さい。) - **物理トポロジとα-βコストモデル**: TACCLはリンクのレイテンシαとバンド幅の逆数βをα-βコストモデル(チャンク転送コスト = α + β·s)で表現し、プロファイラで実測する(Table 1)。Azure NDv2はPCIe接続の詳細が公開されていないため、TACCLのプロファイラがCPU-NIC間・GPU間・CPU-NIC間の帯域/レイテンシプローブでPCIeトポロジを推定する。 **Figure 5: 各GPUシステムの物理トポロジの諸側面** ![[_attachments/arxiv-2111.04867/fig05-physical-topologies.png]] (Figure 5. (a) NDv2のNVLink接続(2ノードがInfiniBandで接続)。(b) NDv2のPCIe接続(各CPUに2つのPCIeスイッチが接続し、NICは片方のスイッチに接続)。(c) DGX-2のNVSwitch経由のNVLink接続(16GPUが全結合)。) | Link | NDv2 α (µs) | NDv2 β (µs/MB) | DGX-2 α (µs) | DGX-2 β (µs/MB) | |---|---|---|---|---| | NVLink | 0.7 | 46 | 0.7 | 8 | | InfiniBand | 1.7 | 106 | 1.7 | 106 | (Table 1. Azure NDv2とNvidia DGX-2ノードで実験的に得られたαとβのコスト。) - **合成アルゴリズム(3段階)**: (1) **ルーティング**: 各チャンクの経路のみを決めるMILP(チャンク間の順序は無視、リンク重なりは許容)を解く。symmetry制約・switch-hyperedge制約もここで適用され、リンクあたりのバイナリ変数を$O(C^2)$から$O(C)$に削減する。(2) **ヒューリスティックな順序付け**: MILPではなく貪欲法で、各リンク上のチャンク送信順を「最終GPUまでの残り経路が長いチャンクを優先」「同順位なら初期GPUからの既移動距離が短いチャンクを優先」の2条件で決める。(3) **連続化と厳密スケジューリング**: 複数チャンクを1回のsendで連続転送するか個別に送るかをMILPで決め、厳密なstart_time/send_timeを確定する。連続転送はα(レイテンシ)コストを1回分に抑えるが、後続の依存する送信を早期にスケジュールできなくなるトレードオフがある。 - **合成される集合演算の組み合わせ**: R EDUCESCATTERはA LLGATHERの「逆」(送信の代わりに受信+縮約)として構築するが、単純な逆転では正しく動作しないため、ヒューリスティックな順序付けと連続化エンコーディングを逆送信に適用する。A LLREDUCEはR EDUCESCATTERとA LLGATHERを連結して合成する。 - **バックエンド(TACCLランタイム)**: 抽象アルゴリズムはTACCL-EF(XML形式)にローワリングされ、NCCLを拡張したインタプリタ(TACCLランタイム)上で単一カーネル起動として実行される。スレッドブロック割り当て・複数インスタンス(帯域飽和のためチャンクをサブチャンクに分割して並列実行)などの実装上の工夫を持つ。PyTorchへの統合はNCCLライブラリをTACCLランタイムに差し替える1行変更で済む。 ## 新規性 先行研究SCCLはSMTソルバでcollective algorithm合成をpareto最適に解くが単一ノードに限定される。TACCLはMILPエンコーディングをルーティング緩和・ヒューリスティック順序付け・連続化の3段に分割することで多ノードへスケールし、かつ異種リンク(NVLink・InfiniBand)のper-message overheadをモデル化できる点でSCCLの制約を克服する。Blinkはヒューリスティックなスパニングツリー詰め込みでノード内帯域を最大化するがNCCLがノード全体を跨ぐRingを組めない場合にのみ優位であり、TACCLはノード全体を使う場合にもNCCLを上回る。Plinkは既知のプリミティブから階層的縮約を構築するのみでアルゴリズム空間を探索しない。 ## 実験設定 - 実験環境: Nvidia DGX-2ノード(V100 GPU 16基、NVSwitch経由でNVLink全結合、PCIeスイッチ経由でNICと接続)2台、Azure NDv2ノード(8GPU、NVLink直結、PCIe経由の1本のInfiniBand NIC)最大4台。合計最大32GPU。 - 比較対象: Nvidia Collective Communication Library(NCCL v2.8.4-1)。理由: 分散ML訓練で最も広く使われる集合通信ライブラリであり、state-of-the-artとして扱われるため。 - 評価指標: アルゴリズム帯域(algorithm bandwidth、入力バッファサイズ÷実行時間)。NCCLに対する速度向上率も併記。 - ソルバ: Gurobi。合成時間の計測にはコンティギュイティエンコーディングに30分のタイムリミットを設けた。 - エンドツーエンド評価: Transformer-XLとBERTの分散訓練(torch.distributed経由)、および内部のMixture-of-Expertsワークロード(社内クラスタ)。 ## 実験結果 ### 単体実験(A LLGATHER・A LLTOALL・A LLREDUCE) **Figure 6: A LLGATHERのNCCL対TACCL最良アルゴリズムの比較** ![[_attachments/arxiv-2111.04867/fig06-allgather-comparison.png]] (Figure 6. (i) DGX-2 2ノード: 送受信GPUを固定した dgx2-sk-1 で256MB-1GBの大バッファに20%-25%の高速化、両GPUが共有NICを使う dgx2-sk-2 で1KB-1MBに4.9倍-6.7倍、2MB-64MBに10%-3.8倍の高速化。(ii) NDv2 2ノード: 1KB-1MBに12%-35%、1MB超に61%-3.4倍の高速化。) **Figure 7: A LLTOALLのNCCL対TACCL最良アルゴリズムの比較** ![[_attachments/arxiv-2111.04867/fig07-alltoall-comparison.png]] (Figure 7. (i) DGX-2 2ノード: dgx2-sk-2の再利用で2MB以上に最大15%、専用スケッチ dgx2-sk-3で1KB-16KBに最大55%の高速化。(ii) NDv2 2ノード: ndv2-sk-1で16MB-1GBに53%-66%、ndv2-sk-2で1KB-128KBに最大12%の高速化。4ノードではndv2-sk-1で1MB超に最大46%の高速化(付録Figure 11)。) **Figure 8: A LLREDUCEのNCCL対TACCL最良アルゴリズムの比較** ![[_attachments/arxiv-2111.04867/fig08-allreduce-comparison.png]] (Figure 8. (i) DGX-2 2ノード: dgx2-sk-2から合成したアルゴリズムが1KB-4MBに49%-6.4倍、dgx2-sk-1から合成したアルゴリズムが16MB-256MBに2%-37%高速化。512MB以上ではNCCLの融合済み通信命令(receive-reduce-copy-send)がTACCLのローワリングにない機能を使うため、最大9%遅くなる。(ii) NDv2 2ノード: 単一インスタンスが1MBまで最大28%、8インスタンスが1MB超に28%-2.7倍の高速化。4ノードでは小バッファに最大34%、大バッファに1.9倍-2.1倍の高速化(付録Figure 11)。) ### 合成器入力を変えたアブレーション(DGX-2上のA LLGATHER) **Figure 9: TACCLへの各種入力を変えたA LLGATHERアルゴリズムの帯域** ![[_attachments/arxiv-2111.04867/fig09-ablation.png]] (Figure 9. (a) 論理トポロジ: IB接続数を増やすほど小チャンク(1KB)で有利だが、チャンクサイズが32KB・1MBと増えるにつれ最適な接続数は4・1に減る(β-costがα-costより優勢になるため)。(b) チャンクサイズ: 合成時に指定したサイズに近いデータサイズで最良の性能が出る。(c) データ分割: 1GBの大バッファではチャンクを2分割した方が帯域利用が良い。(d) switch-hyperedge方針: 小バッファではuc-max、大バッファではuc-minが優位。(e) ランタイムインスタンス数: インスタンス数を増やすと帯域利用が改善するが、小バッファではレイテンシコストが増す。) ### エンドツーエンド訓練 **Figure 10: Transformer-XLとBERTの訓練スループット比較** ![[_attachments/arxiv-2111.04867/fig10-training-throughput.png]] (Figure 10. Azure NDv2の2ノード・4ノード(16・32GPU)で、Transformer-XLはデータ並列(A LLREDUCEサイズ20-40MB)、BERTはモデル並列(A LLREDUCEサイズ約2MB)を使う。Transformer-XLは2ノードで11%-1.94倍、4ノードで2%-1.44倍。BERTは2ノードで12%-2.36倍、4ノードで7%-1.74倍の高速化。バーの上に示された数値がNCCL対比の速度向上率。) 社内Mixture-of-Expertsワークロード(NDv2 2ノード、A LLTOALLが約6MB・A LLREDUCEが約256MB)ではTACCLがエンドツーエンドスループットを17%改善した。 ### 合成時間 | 集合演算 | スケッチ | 時間(s) | |---|---|---| | AllGather | dgx2-sk-1 | 35.8 | | AllGather | dgx2-sk-2 | 11.3 | | AllGather | ndv2-sk-1 | 2.6 | | AllToAll | dgx2-sk-2 | 92.5 | | AllToAll | ndv2-sk-1 | 1809.8 | | AllToAll | ndv2-sk-2 | 8.4 | | AllReduce | dgx2-sk-1 | 6.1 | | AllReduce | dgx2-sk-2 | 127.8 | | AllReduce | ndv2-sk-1 | 0.3 | (Table 2. 各集合演算・通信スケッチの組み合わせでTACCLが合成に要した総時間。多くは数秒〜数分で、人間参加型のワークフローに適う。ndv2-sk-1でのA LLTOALLの連続化エンコーディングは30分のタイムリミットに達するが、実行可能解自体は4分14秒で得られていた。80GPU(NDv2 10ノード)のA LLGATHER合成は8分未満で完了した。) ## 考察 - スケッチが異なれば最適な入力サイズ帯域も異なる(Figure 9参照)。単一のスケッチが全サイズ帯を支配するわけではなく、複数スケッチを試すこと自体がTACCLの運用モデルの一部になっている。 - スケーラビリティ: TACCLはAzure NDv2 8ノードでのA LLGATHER合成を5分未満で完了し、NCCL比最大1.7倍の帯域を得た。Nvidia DGX-2 8ノード(128GPU)でも約11時間で解を発見した。合成技術自体はNP-hard問題に基づき規模に対し二乗的に増大するため、著者らは合成アルゴリズムの階層的合成を将来課題として挙げる。 - 汎用性: 階層的トポロジ(DGX-2・NDv2)だけでなく、2D 6×8トーラスのような非階層トポロジにもTACCLは適用可能で、対称性属性を活用してA LLGATHERを合成できた。ただし非階層トポロジではスケッチの探索余地がより限定的になる可能性がある。 ## 強み / 弱点・課題 - **強み**: 通信スケッチという低労力な人間入力でMILP探索空間を絞ることで、SCCLが単一ノードに限られていた合成手法を32GPU(最大128GPUの評価あり)規模の多ノードGPUトポロジへスケールさせた。3段階のMILP/ヒューリスティック分割によって合成時間を秒〜分単位に抑え、人間参加型ワークフローを現実的にした。 - **弱点・課題**: 合成技術自体は依然NP-hard問題に基づき、GPU数に対し二乗的に計算量が増大するため、大規模化には限界がある(DGX-2 128GPUで約11時間)。通信スケッチの選択自体は依然として人間の直感に依存し、どのスケッチがどの入力サイズ帯に最適かを自動的に決定する仕組みは無い(著者ら自身が将来課題として挙げる)。A LLREDUCEの一部の大バッファサイズ(512MB以上)ではNCCLの融合済み通信命令に対しTACCLが最大9%遅い場合がある。