FLUX: Fast Software-based Communication Overlap On GPUs Through Kernel Fusion

要約

大規模な深層学習モデルは、幅広いアプリケーションにわたって多くのタスクを解決する強力な能力を実証しています。
これらの大規模なモデルは通常、トレーニングと推論を分散する必要があります。
テンソル並列処理は、単一プロセッサのメモリ容量制限を克服したり、特定のレイテンシー要件を満たすために計算を高速化したりするために、デバイス間でオペレーションまたはレイヤーの計算を分割する一般的な手法です。
ただし、この種の並列処理により追加の通信が発生し、全体の実行時間のかなりの部分に影響を与える可能性があります。
したがって、ノード内の NVLink を備えた GPU など、高速相互接続を備えたデバイスのグループ内でのこの手法のスケーラビリティは制限されます。
この論文では、GPU の依存計算による通信遅延を大幅に隠すための新しい方法 Flux を提案します。
Flux は、通信と計算の操作をさらに細かい操作に分解し、さらにそれらをより大きなカーネルに融合して、カーネルの効率を損なうことなく通信を効果的に隠蔽します。
融合カーネルを使用すると、Flux は通信の最大 96% と重複する可能性があります。
全体として、さまざまな GPU 世代と相互接続を備えた 128 GPU のクラスター上の Megatron-LM と比較してトレーニングでは最大 1.24 倍の高速化を実現でき、8 つの GPU を備えたクラスターでは vLLM と比較してプリフィルとデコード推論では最大 1.66 倍および 1.30 倍の高速化を実現できます。
さまざまな GPU 世代と相互接続を備えた GPU。

要約(オリジナル)

Large deep learning models have demonstrated strong ability to solve many tasks across a wide range of applications. Those large models typically require training and inference to be distributed. Tensor parallelism is a common technique partitioning computation of an operation or layer across devices to overcome the memory capacity limitation of a single processor, and/or to accelerate computation to meet a certain latency requirement. However, this kind of parallelism introduces additional communication that might contribute a significant portion of overall runtime. Thus limits scalability of this technique within a group of devices with high speed interconnects, such as GPUs with NVLinks in a node. This paper proposes a novel method, Flux, to significantly hide communication latencies with dependent computations for GPUs. Flux over-decomposes communication and computation operations into much finer-grained operations and further fuses them into a larger kernel to effectively hide communication without compromising kernel efficiency. Flux can potentially overlap up to 96% of communication given a fused kernel. Overall, it can achieve up to 1.24x speedups for training over Megatron-LM on a cluster of 128 GPUs with various GPU generations and interconnects, and up to 1.66x and 1.30x speedups for prefill and decoding inference over vLLM on a cluster with 8 GPUs with various GPU generations and interconnects.

arxiv情報

提供元, 利用サービス

arxiv.jp, Google