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

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

本文旨在解决分布式训练中通信延迟对性能的限制问题，提出一种新的方法Flux，以在GPU上隐藏通信延迟。

Flux将通信和计算操作细分为更小的操作，然后将它们融合成一个更大的内核，以有效地隐藏通信而不影响内核效率。Flux可以潜在地重叠高达96%的通信。通过在128个GPU的集群上对Megatron-LM进行训练，Flux可以实现高达1.24倍的加速，对于8个GPU的集群上的vLLM进行预填充和解码推理，可以分别实现高达1.66倍和1.30倍的加速。

该方法可以显著隐藏依赖计算的通信延迟，提高分布式训练的效率。实验结果表明，Flux可以在不影响内核效率的情况下，潜在地重叠高达96%的通信。此外，该论文还提供了实验细节和开源代码，方便其他研究者进行复现和扩展。

与本文相关的研究包括：GPipe、PipeDream、Megatron-LM等。