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

大型深度学习模型已经展示了在广泛应用范围内解决许多任务的强大能力。这些大型模型通常需要分布式的训练和推断。张量并行是一种常见的技术，将一个操作或层的计算在设备之间分配，以克服单个处理器的内存容量限制，或加速计算以满足某种延迟要求。然而，这种并行性引入了额外的通信，可能对整体运行时间产生重要影响。因此，这种技术的可扩展性在高速互连的设备组中受到限制，例如具有NVLink的GPU节点。本文提出了一种新方法Flux，可以显著地隐藏依赖计算的通信延迟。Flux将通信和计算操作超分解为更细粒度的操作，并将它们进一步融合到一个更大的内核中，以有效地隐藏通信而不影响内核效率。在融合内核的情况下，Flux可以潜在地重叠高达96％的通信。总体而言，Flux在由各种GPU世代和互连组成的128个GPU集群上，可以实现高达1.24倍的训练加速，而在由各种GPU世代和互连组成的8个GPU集群上，可以实现高达1.66倍和1.30倍的预填充和解码推断加速。

论文旨在解决分布式训练中通信延迟对性能的影响问题，提出了一种名为Flux的新方法。

Flux方法将通信和计算操作分解为更细粒度的操作，并将它们融合成更大的内核，以有效地隐藏通信，同时不影响内核效率。

论文在128个GPU的集群上，相比于Megatron-LM，Flux可以实现高达1.24倍的训练速度提升。在8个GPU的集群上，相比于vLLM，Flux可以实现高达1.66倍和1.30倍的预填充和解码推理速度提升。论文还开源了Flux的代码。

与该论文相关的研究包括：GPipe、PipeDream和Megatron等。