A 4D Hybrid Algorithm to Scale Parallel Training to Thousands of GPUs

要約

大量の通信、特に集団操作は、10 億パラメータのニューラル ネットワークのトレーニングを大規模な並列システムに拡張する際に、重大なパフォーマンスのボトルネックになる可能性があります。
このペーパーでは、並列トレーニングにおけるコミュニケーションを最適化するための 4 次元 (4D) アプローチを紹介します。
この 4D アプローチは 3D テンソルとデータ並列処理のハイブリッドであり、AxoNN フレームワークに実装されています。
さらに、通信オーバーヘッドをさらに最小限に抑えるために 2 つの重要な戦略を採用しています。
まず、負荷の高い集団操作 (reduce-scatter、all-gather、および all-reduce) を計算と積極的に重ね合わせます。
次に、4D アルゴリズムによって定義された大規模な検索空間内で高性能の構成を特定するための分析モデルを開発します。
このモデルは、特定のトレーニング ワークロードの調整プロセスを簡素化することで、実践者に力を与えます。
Perlmutter の 1024 GPU で 800 億パラメータの GPT をトレーニングすると、AxoNN は最先端のフレームワークである Megatron-LM を 26% 大幅に上回りました。
さらに、理論上のピーク FLOP/s の 57% または合計 182 PFLOP/s という大幅に高い値を達成します。

要約(オリジナル)

Heavy communication, in particular, collective operations, can become a critical performance bottleneck in scaling the training of billion-parameter neural networks to large-scale parallel systems. This paper introduces a four-dimensional (4D) approach to optimize communication in parallel training. This 4D approach is a hybrid of 3D tensor and data parallelism, and is implemented in the AxoNN framework. In addition, we employ two key strategies to further minimize communication overheads. First, we aggressively overlap expensive collective operations (reduce-scatter, all-gather, and all-reduce) with computation. Second, we develop an analytical model to identify high-performing configurations within the large search space defined by our 4D algorithm. This model empowers practitioners by simplifying the tuning process for their specific training workloads. When training an 80-billion parameter GPT on 1024 GPUs of Perlmutter, AxoNN surpasses Megatron-LM, a state-of-the-art framework, by a significant 26%. Additionally, it achieves a significantly high 57% of the theoretical peak FLOP/s or 182 PFLOP/s in total.

arxiv情報

提供元, 利用サービス

arxiv.jp, Google