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

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

论文提出一种名为Flux的方法，旨在解决Tensor并行技术中通信延迟对模型训练和推理性能的限制问题。

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

Flux方法可以潜在地重叠高达96％的通信，从而显着隐藏通信延迟。在128个GPU的集群上，Flux在Megatron-LM模型上实现了高达1.24倍的加速比，在8个GPU的集群上，对vLLM模型进行预填充和解码推理时，分别实现了1.66倍和1.30倍的加速比。

与Flux相关的研究包括TensorFlow和PyTorch等框架的并行计算优化，以及类似于GPipe和PipeDream等的分层并行方法。