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

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

论文试图解决使用Tensor parallelism技术时，通信延迟对于模型训练和推理速度的影响问题。

论文提出了一种名为Flux的方法，将通信和计算操作分解成更细粒度的操作，并将它们融合成更大的内核，以有效地隐藏通信而不影响内核效率。

Flux可以潜在地重叠高达96％的通信，从而实现高效的通信隐藏。实验表明，在使用各种GPU代数和互连的128个GPU集群上，Flux可以比Megatron-LM获得高达1.24倍的训练速度提升，并且在使用各种GPU代数和互连的8个GPU集群上，Flux可以比vLLM获得高达1.66倍和1.30倍的预填充和解码推理速度提升。

目前在这个领域中，还有一些相关的研究，例如《Megatron-LM: Training Multi-Billion Parameter Language Models Using Model Parallelism》和《vLLM: The Hidden Challenge of Very Large Language Models》。