[論文レビュー] TACCL: Guiding Collective Algorithm Synthesis using Communication Sketches

Aashaka Shah, Vijay Chidambaram+7 more|arXiv (Cornell University)|Nov 8, 2021

Ferroelectric and Negative Capacitance Devices被引用数 6

ひとこと要約

TACCL は、人間が提供する通信スケッチを用いて、自動合成された効率的な GPU コラレクティブ通信アルゴリズムをガイドする画期的なツールであり、探索空間を顕著に削減し、マルチノード GPU トポロジにおけるスケーラブルで高性能なアルゴリズム生成を可能にする。TACCL は NCCL を最大 6.7× 上回り、BERT や Transformer-XL の学習を 11%–2.3× 速くする。

ABSTRACT

Machine learning models are increasingly being trained across multiple GPUs and servers. In this setting, data is transferred between GPUs using communication collectives such as AlltoAll and AllReduce, which can become a significant bottleneck in training large models. Thus, it is important to use efficient algorithms for collective communication. We develop TACCL, a tool that enables algorithm designers to guide a synthesizer into automatically generating algorithms for a given hardware configuration and communication collective. TACCL uses a novel communication sketch abstraction to get crucial information from the designer to significantly reduce the search space and guide the synthesizer towards better algorithms. TACCL also uses a novel encoding of the problem that allows it to scale beyond single-node topologies. We use TACCL to synthesize algorithms for three collectives and two hardware topologies: DGX-2 and NDv2. We demonstrate that the algorithms synthesized by TACCL outperform the Nvidia Collective Communication Library (NCCL) by up to 6.7x. We also show that TACCL can speed up end-to-end training of Transformer-XL and BERT models by 11%--2.3x for different batch sizes.

研究の動機と目的

• 大規模かつ非均一な GPU トポロジ（帯域幅と遅延が異なる特性を併せ持つ）における効率的なコラレクティブ通信アルゴリズムの設計という課題に取り組む。
• マルチノードシステムにおけるコラレクティブアルゴリズムのルーティングおよびスケジューリング選択肢の全探索空間を計算的に実行不可能にする問題を克服する。
• 通信スケッチと呼ばれる新規な抽象化を通じてアルゴリズム設計者の高レベルな直感を組み込むことで、合成時間の短縮とスケーラビリティの向上を実現する。
• DGX-2 や NDv2 などの特定のハードウェア構成に最適化された Allreduce、Allgather、Alltoall アルゴリズムの自動生成を可能にする。
• リラックスした MILP 定式化とヒューリスティック最適化手法を組み合わせることで、NCCL より優れたパフォーマンスを達成する。

提案手法

• コラレクティブ通信合成問題を、ルーティングを緩和し、連続時間スケジューリングを採用した混合整数線形計画問題（MILP）として定式化することで、スケーラビリティを向上させる。
• 3段階の合成パイプラインを適用する：緩和されたルーティングによる妥当な経路の探索、帯域幅の高いリンクを優先するヒューリスティック順序付け、連続性制約を含む正確なスケジューリングによるアイドル時間の最小化。
• 対称性とトポロジに依存するエンコーディングを活用して状態空間を削減し、2次元トーラスや階層的クラスタ（例：DGX-2、Azure NDv2）などの複雑なトポロジでも合成を可能にする。
• プロファイル済みのネットワークトポロジと入力サイズの情報を入力として用い、ハードウェアに適した、遅延および帯域幅最適化されたアルゴリズムへの合成をガイドする。
• 低レベルの GPU 通信プリミティブにマッピングするバックエンドを実装し、既存の分散ディープラーニングフレームワークとの互換性を確保する。

実験結果

リサーチクエスチョン

• RQ1 通信スケッチは、大規模な GPU トポロジにおけるコラレクティブ通信アルゴリズムの自動合成を、探索空間の削減とともに効果的にガイドできるか？
• RQ2 帯域幅が異なるリンクを持つ非均一なネットワークトポロジにおいて、単一ノードを超えたスケーラブルな効率的アルゴリズム合成はどのように達成できるか？
• RQ3 直感的なスケッチを用いた人間とシステムの共同作業（human-in-the-loop）アプローチは、完全に自動化されたまたはヒューリスティックベースの手法に比べて、高性能な通信アルゴリズムの生成においてどの程度優位性を示せるか？
• RQ4 リラックスした定式化とヒューリスティック最適化を組み合わせた本研究の MILP 定式化は、大規模な GPU クラスタにおいても実行可能であり、近似的に最適なパフォーマンスを達成できるか？
• RQ5 本研究で生成されたアルゴリズムは、さまざまなワークロードとハードウェア構成において、最先端のライブラリ（例：NCCL）と比較して実際の性能でどの程度優れているか？

主な発見

• TACCL は 8 台の Azure NDv2 ノード（64 GPU）向けに Allgather アルゴリズムを 5 分未満で合成し、さまざまなデータサイズで NCCL より最大 1.7× の高いアルゴリズム帯域幅を達成した。
• 8 台の DGX-2 ノード（128 GPU）では、TACCL が約 11 時間で有効な Allgather アルゴリズムを生成し、大規模マルチノードシステムへのスケーラビリティを実証した。
• テストされたコラレクティブおよびトポロジにおいて、TACCL は NCCL より最大 6.7× の高いエンドツーエンド通信パフォーマンスを達成した。
• バッチサイズに応じて、BERT や Transformer-XL モデルのエンドツーエンド学習が 11% から 2.3× 速くなった。これは通信ボトルネックの低減によるものである。
• TACCL は、2D $6\times8$ トーラスのような非階層的トポロジに対しても、通信スケッチの対称性を活用して効率的なアルゴリズムを生成した。
• 通信スケッチの使用により、入力サイズの範囲に応じて最適化された複数のアルゴリズムバージョンの探索が可能となり、開発者の設計直感を反映した。