文章标题: FLUX: 通过内核融合在GPU上实现快速的基于软件的通信重叠
作者/机构: Li-Wen Chang∗1, Wenlei Bao∗1, Qi Hou∗1, Chengquan Jiang∗1, Ningxin Zheng∗1, Yinmin Zhong∗2, Xuanrun Zhang∗1, Zuquan Song1, Ziheng Jiang1, Chengji Yao1, Haibin Lin1, Xin Jin2, and Xin Liu1
1ByteDance Ltd, 2Peking University

A1 主要贡献

本文旨在解决大规模深度学习模型中张量并行（Tensor Parallelism）引入的通信开销问题。张量并行虽然能克服单个处理器内存容量限制并加速计算，但其引入的额外通信（如数据收集或重分发）会占据总运行时间的很大一部分，从而限制了其在高速互联设备组（如节点内的NVLink GPU）中的可扩展性。

核心问题: 现有用于张量并行的通信重叠技术在GPU上存在局限性，例如：当使用CUDA流时，无法精确控制GPU上的执行时机；执行多个较小的内核实例导致GPU利用率低下。

研究目标: 提出一种新颖的方法Flux，旨在通过更有效地将通信延迟与依赖的计算相重叠，来显著隐藏GPU上的通信延迟。

创新点与主要贡献:
* 识别现有技术的性能问题: 论文识别并分析了在GPU上应用现有张量并行通信重叠技术时遇到的几个性能瓶颈。
* 提出新颖的通信重叠技术Flux: 提出了一种名为Flux的新方法，该方法将通信和计算操作过度分解（overdecomposes）为更细粒度的操作，并将它们融合成一个更大的内核。这种方法能够有效地隐藏通信，同时不损害内核效率。在一个融合内核中，Flux可以重叠高达96%的通信。
* 基于CUTLASS的优化实现: 使用NVIDIA CUTLASS实现了所提出的技术，并针对不同的GPU代次（A100和H800）和节点内互连方式（PCIe和NVLink）进行了多项优化。Flux的实现是模块化的，可以轻松地在各种GPU架构和互连组合中进行自动调优。
* 全面的评估与显著的性能提升: 在多个GPU集群上对Flux进行了评估。
* 训练: 在一个包含128个GPU的集群上，与Megatron-LM相比，训练速度提升高达1.24倍。
* 推理: 在一个包含8个GPU的集群上，与vLLM相比，prefill阶段的速度提升高达1.66倍，decoding阶段的速度提升高达1.30倍。
* 与先前重叠方法的对比: 与先前的方法TransformerEngine相比，Flux在训练、prefill和decoding上分别取得了高达1.38倍、2.06倍和2.10倍的加速。

图1：在各种128-GPU集群上使用2路数据、8路流水线、8路张量并行进行训练，以及在各种8-GPU集群上使用8路张量并行进行推理时，常见LLM工作负载中张量并行内的非重叠通信部分。

A3 背景知识与关键观察

2.1 常见的划分和通信模式

MLP中的张量并行模式 本文讨论的常见划分策略是Megatron-LM【索引24，Megatron-lm: Training multi-billion parameter language models using model parallelism，2019，arXiv】的扩展，并带有分片的激活（sharded activation）【索引13，Overlap communication with dependent computation via decomposition in large deep learning models，2023，ASPLOS；索引25，Reducing activation recomputation in large transformer models，2023，MLSys】。以Transformer中的多层感知机（MLP）部分为例，第一个GEMM操作沿行方向对权重（W1）进行分片，并在GEMM前沿列方向对输入的激活进行AllGather操作。第二个GEMM操作沿列方向对权重（W2）进行分片，并在GEMM后沿列方向对输出的激活进行ReduceScatter操作。在反向传播中，AllGather和ReduceScatter的位置会互换。这两个GEMM操作的维度取决于张量并行的程度（N）。

图2：N台设备上MLP部分N路划分的前向传播。这里，B和L被展平以适应GEMM的通用表示法。

其他划分模式 另一种常见的划分模式是将所有权重（W1和W2）通过数据并行进一步分片到设备上，并在GEMM之前通过AllGather收集权重【索引20，Zero++: Extremely efficient collective communication for giant model training，2023，arXiv；索引27，Zero: Memory optimizations toward training trillion parameter models，2020，SC20；索引28，{ZeRO-Offload}: Democratizing {Billion-Scale} model training，2021，USENIX ATC；索引29，Zero-infinity: Breaking the gpu memory wall for extreme scale deep learning，2021，SC21；索引30，Pytorch fsdp: experiences on scaling fully sharded data parallel，2023，arXiv】。在这种模式下，由于所有权重在其消费者GEMM操作之前没有数据依赖，这些AllGather操作可以很容易地被预取并与独立的操作重叠。因此，本文主要讨论第一种模式。

2.2 传统的通信重叠策略

基于分块的调度 传统方法【索引12，Breaking the computation and communication abstraction barrier in distributed machine learning workloads，2022，ASPLOS；索引13，Overlap communication with dependent computation via decomposition in large deep learning models，2023，ASPLOS；索引14，TransformerEngine，2022，NVIDIA；索引31，Megascale: Scaling large language model training to more than 10,000 gpus，2024，NSDI】将原始的计算和通信操作分解成多个块（chunks），然后通过精心调度这些操作来重叠通信与计算。分解的分区数量通常与张量并行中的设备数量（或其两倍，以更好地利用双向数据传输）对齐。这种限制分区数量的做法可以避免复杂的调度并减少可能的调度开销。图3展示了一个理想情况下，通信可以被GEMM计算完全隐藏的ReduceScatter重叠场景。

在GPU上的局限性 这些方法在GPU上可能效果不佳，原因如下：
1. 执行时机不可控：这些方法的性能严重依赖于独立分区的执行顺序、并发执行和执行时机。虽然可以通过流（streams）和事件（events）来控制GPU内核间的执行顺序和并发，但执行时机在大多数GPU编程模型中不易控制，在涉及许多流和事件的生产环境中通常变得不可预测。
2. 数据依赖阻止并发：ReduceScatter的重叠通常需要在GEMM操作之间执行额外的计算（如图3中的加法操作），这会产生数据依赖，从而阻止了多个GEMM内核通过GPU多路复用（multiplexing）并发执行。尽管加法操作可以进一步与通信融合【索引14，TransformerEngine，2022，NVIDIA】，但它们仍然阻碍了多个GEMM内核的并发执行。
3. GPU利用率低：将一个大的GEMM内核拆分成多个小的GEMM内核，很可能会导致GPU流处理器（SMs）利用率不足，即使分区数量与设备数量相同。当张量并行规模扩大时，这个问题尤其严重。

2.3 有效通信时间与重叠效率

性能指标定义 评估通信重叠方法的性能并非易事。为了公平比较不同方法并突出通信时间，论文提出了两个指标：
1. 有效通信时间 (Effective Communication Time, ECT)：定义为总时间（OverallTime）减去最优的、未拆分的GEMM计算时间（GEMM_non-split）。对于非重叠方法，ECT等于常规的通信时间；对于重叠方法，任何低效因素导致的减速时间和未重叠的通信部分都计入ECT。一个完美的重叠方法其ECT为零。

重叠效率 (Overlap Efficiency, E_overlap)：定义为1减去重叠方法的ECT与非重叠基线的ECT之比。非重叠基线的重叠效率为零，而完美的重叠方法效率为100%。负的重叠效率意味着该方法比非重叠基线还要慢。

Eq 2

传统方法的性能问题分析 为了最小化基线的影响，评估中均使用广泛使用的标准GPU通信库NCCL【索引32，NCCL，2016，NVIDIA】作为非重叠基线。图4展示了PyTorch（非重叠）和传统重叠技术（TransformerEngine实现）的性能对比。结果显示，由于上述局限性，传统重叠技术可能性能不佳，甚至比非重叠方法更差。例如，当矩阵的m维度较小时，TransformerEngine性能劣于PyTorch，这证实了拆分GEMM导致GPU利用率低的观点。此外，TransformerEngine在AllGather中的表现优于ReduceScatter，因为AllGather中拆分的GEMM操作可能通过GPU多路复用并发执行，而ReduceScatter则不能，这证实了数据依赖阻止并发的观点。

图4：在具有NVLink互连的8-H800集群上，m从1024到8192，AllGather的(n, k)为(49152, 12288)，ReduceScatter的(n, k)为(12288, 49152)时，PyTorch（非重叠）和TransformerEngine（先前的重叠方法）之间的性能比较。

A2 方法细节

3. Fused GEMM与通信概览

Flux方法概述 本文提出的Flux是一种比传统方法【索引12，Breaking the computation and communication abstraction barrier in distributed machine learning workloads，2022，ASPLOS；索引14，TransformerEngine，2022，NVIDIA】更高效的通信重叠方法。与现有方法将计算和通信划分为设备数量（或其两倍）个分区不同，Flux将计算和通信过度分解（overdecompose）为更细粒度的瓦片（tiles）。由于高性能GEMM内核（如线程块瓦片或warp瓦片）本身就是基于瓦片化编写的，Flux的分解可以自然地映射到内核中已有的瓦片结构。Flux将依赖的通信和/或等待逻辑融合到一个GEMM内核中，并且只启动一个融合后的内核，而不是像先前方法那样启动多个拆分的GEMM内核。由于Flux的粒度远细于先前方法，本文将先前方法称为中等粒度分解，而将Flux称为细粒度分解。

3.1 ReduceScatter 重叠

通过Epilogue融合实现 在Flux中，ReduceScatter是通过将通信融合到GEMM内核的epilogue（结尾部分）来实现的。算法1展示了融合了ReduceScatter（或AlltoAll）的GEMM伪代码（用于C = A × B）。与标准GEMM内核只有一个输出指针不同，融合后的GEMM内核的输出指针数量（Cs）增加到张量并行中的设备数量（NTP），这些指针可以在相应PyTorch操作的初始化阶段通过进程间通信收集。输出坐标（m和n）通过TileCoord函数根据线程块索引和本地rank索引（rank_id）计算得出。融合GEMM中输出指针的选择（GetOutput）基于输出坐标（m和n）和设备数量（NTP）。例如，在图2的两个GEMM操作中，选择是基于行索引的。

// 算法1：一个简化的ReduceScatter-GEMM重叠内核
参数: 输入矩阵指针 A, B
参数: 输出矩阵指针列表 Cs
参数: 整型标量 rank_id, NTP
[m, n] ← TileCoord(threadblock_id, rank_id, NTP);
acc ← 0;
标准GEMM prologue(A, B, m, n);
标准GEMM mainloop(A, B, m, n) 更新 acc;
// epilogue 开始
C ← GetOutput(Cs, NTP, m, n);
if 融合 reduction then
    Reduce(C, acc);
else
    Write(C, acc);
end
// epilogue 结束

AlltoAll与Reduction的解耦 一个ReduceScatter操作可以进一步解耦为一个AlltoAll操作和一个reduction操作。其中，AlltoAll仅指设备间的通信，而reduction则在各个设备上本地进行。因此，将AlltoAll（伪代码中的Write分支）融合到GEMM的epilogue中通常足以重叠通信，而融合reduction（Reduce分支）仅提供边际性能增益。该算法要求GPU支持点对点（P2P）访问，现代NVIDIA GPU在节点内（无论是NVLink还是PCIe）都已具备此功能。NVSHMEM【索引33，NVSHMEM，2020，NVIDIA】将P2P功能扩展到了跨节点GPU。

3.2 AllGather 重叠

通过Prologue融合实现 与ReduceScatter不同，AllGather是通过将信号检查融合到GEMM内核的prologue（起始部分）来实现的。算法2展示了融合了AllGather的GEMM伪代码（用于C = Aagg × B），其中Aagg是通过AllGather聚合输入A矩阵得到的聚合矩阵缓冲区。算法3则展示了在主机端发生的相应通信过程。

// 算法2：一个简化的AllGather-GEMM重叠内核
参数: 输入矩阵指针 Aagg, B
参数: 输出矩阵指针 C
参数: 标量信号列表 signal_list
参数: 整型标量 rank_id, NTP
[m, n] ← TileCoord(threadblock_id, rank_id, NTP);
signal ← GetSignal(signal_list, NTP, m, n);
WaitSignal(signal);
标准GEMM(A, B, C, m, n);

内核端与主机端的协同工作 在内核端，GEMM瓦片计算会被WaitSignal函数阻塞，直到signal中的值被设置成true。signal的选择基于输出坐标（m和n）以及设备数量（NTP）。主机端（可以是pull-based或push-based）异步执行瓦片化的通信操作（DataTransfer）并设置相应的信号（SetSignal）。在pull-based方法中，通过GetRemotePtr从远程设备拉取数据瓦片，并设置本地信号。在push-based方法中，则是将数据瓦片推送到远程设备，并设置远程信号。这两种变体的选择是一个调优选项。

// 算法3：一个用于AllGather-GEMM重叠的主机函数
参数: 输入矩阵指针列表 A_list
参数: 输出矩阵指针列表 Aagg_list
参数: 标量信号列表 signal_list
参数: 整型标量 rank_id, NTP
参数: 瓦片信息列表 tilescomm
for tile from tilescomm do
    if pull then // pull-based
        Aremote ← GetRemotePtr(A_list, tile);
        Alocal ← GetLocalPtr(Aagg_list, tile);
        DataTransfer(Aremote, Alocal, tile.size);
    else // push-based
        Aremote ← GetRemotePtr(Aagg_list, tile);
        Alocal ← GetLocalPtr(A_list, tile);
        DataTransfer(Alocal, Aremote, tile.size);
    end
    signal ← GetSignalHost(signal_list, tile);
    SetSignal(signal);
end

设计的灵活性 在AllGather中，Flux只将通信的等待逻辑融合到GEMM内核中，而不是整个通信操作，因此不一定需要P2P支持。同时，通信的瓦片化策略（tilescomm）与GEMM计算的瓦片化策略是解耦的。这种设计提供了一种灵活的方式来权衡重叠机会和通信效率，而不会影响GEMM的效率。

3.3 分解策略的比较

ReduceScatter中的策略对比 图5展示了不同重叠技术在ReduceScatter中的主要区别。尽管现有重叠方法（Tm）可能比原始的粗粒度方法（Tc）快，但由于将一个大GEMM内核拆分成多个小GEMM内核会导致GPU效率下降，它通常仍比原始GEMM时间（Tg）慢。此外，小GEMM操作之间的数据依赖性阻碍了它们通过GPU多路复用并发运行。相比之下，Flux（Tf）没有这些限制，其性能可以与原始GEMM操作（Tg）相当，只有很小的开销。Flux的细粒度分解策略完美地适应了现代GPU的设计，通过在上下文切换的warps和SM上数百个并发活动的warps之间隐藏延迟。最终，Flux只在执行的尾部暴露一小部分通信，而不影响GEMM的计算效率。

图5：在2路张量并行的GEMM-ReduceScatter模式中，非重叠方法和不同重叠方法之间的差异示意图。

AllGather中的策略对比 图6展示了在AllGather中的类似比较。同样，现有方法（Tm）可能比原始方法（Tc）快，但由于GPU GEMM效率较低，仍慢于原始GEMM时间（Tg）。而Flux（Tf）可以提供与原始GEMM操作（Tg）相似的性能。在Flux中，长延迟指令来自WaitSignal，其延迟根据相应数据传输的到达时间而变化。对于数据已到达的瓦片，延迟接近于零；对于数据未就绪的瓦片，warps之间的上下文切换可以隐藏延迟。值得注意的是，本地瓦片的信号总是预设为true，因此总有一些warps无需等待。最终，Flux只在执行的头部暴露一小部分通信，而不影响GEMM的计算效率。

图6：在2路张量并行的AllGather-GEMM模式中，非重叠方法和不同重叠方法之间的差异示意图。

4. 优化与实现细节

引言算法1、2和3的直接实现已经能够通过提供更好的通信重叠和GEMM效率来超越先前的方法。本节将介绍进一步提升性能的高级优化和实现细节。

4.1 瓦片坐标变换（Tile Coordinate Swizzling）

动机与实现 高效的GPU内核依赖于瓦片化来利用并行性和局部性。受调优良好的GEMM通常会通过变换（swizzling）线程块索引到瓦片坐标的映射逻辑以最大化内存效率的启发【索引21，CUTLASS，2024，NVIDIA】，Flux探索了瓦片坐标变换来进一步提高融合内核的效率。在融合的GEMM-ReduceScatter中，瓦片坐标会根据设备rank索引进行偏移，以避免来自不同设备上运行的内核的写请求冲突，从而最小化每个设备内存控制器上的潜在争用。图7说明了朴素映射中可能发生的内存争用，以及偏移映射如何避免这种争用。

图7：在4路张量并行的GEMM-ReduceScatter重叠中，朴素瓦片坐标映射在时间步Ti发生的内存争用示意图，以及所提出的解决方案。

在AllGather中的应用 在融合的AllGather-GEMM中也应用了类似的策略，以最小化线程块的等待时间，从而减少整体延迟。这里的瓦片坐标变换（TileCoord）需要与信号到达顺序对齐，而后者由主机端的通信顺序（由算法3中的tilescomm控制）决定。在实现中，这两个顺序会根据网络拓扑共同选择，以最小化整体延迟。

性能影响 图8显示了在8-A100 NVLink集群上应用瓦片坐标变换前后的性能影响。调整后的映射总是优于朴素映射，并且随着矩阵尺寸的增大，性能影响也随之增加。这主要是因为在GEMM-ReduceScatter中，朴素映射的内存争用以及在AllGather-GEMM中的等待时间会随着矩阵尺寸的增大而增加。值得注意的是，现有方法【索引13，Overlap communication with dependent computation via decomposition in large deep learning models，2023，ASPLOS】也应用了类似的变换思想，但由于Flux不拆分GEMM操作，其思想无法直接应用于Flux的算法中。

图8：在8-A100 NVLink集群上，对于小m(1024)和大m(8192)，AllGather的(n,k)为(49152, 12288)，ReduceScatter的(n,k)为(12288, 49152)时，应用或不应用瓦片坐标变换的性能对比。

4.2 ReduceScatter的实现细节

Write操作 在本地GPU或远程节点内P2P GPU上写数据，是通过以下方式实现的：1) 使用所有变体的st指令（包括向量版本）将数据从寄存器存储到全局内存；2) 使用所有变体的cp.async.bulk指令或Hopper GPU上的张量内存加速器（TMA）指令cp.async.bulk.tensor将数据从scratchpad存储到全局内存。对于向远程跨节点GPU写数据，则应用NVSHMEM并通过其所有变体的put API实现。所有方法都使用CUTLASS EVT【索引34，EVT: Accelerating deep learning training with epilogue visitor tree，2024，ASPLOS】通过模板实现，模板参数在自动调优期间选择。

Reduce操作 如果GPU支持P2P内存访问，可以使用red或atomic指令直接在设备内存上实现reduction，而无需改变代码结构或引入过多开销。这些指令仅用于支持的数据类型和有能力的GPU。在Hopper GPU上，应用warp或线程块特化（specialization）【索引21，CUTLASS，2024，NVIDIA】来实现reduction，即每个GPU将部分结果写入其本地内存，一个特化的warp或线程块则拉取已就绪的远程数据，在目标GPU上执行本地reduction。对于远程跨节点GPU，只在内核中融合AlltoAll，并执行离散的reduction。所有方法同样使用CUTLASS EVT模板实现，并通过自动调优选择参数。

4.3 AllGather的实现细节

DataTransfer操作 对于具有P2P内存访问的GPU，可以使用cudaMemcpy API实现pull-based或push-based传输。两者的唯一区别在于指针：pull-based使用本地目标指针和远程源指针，而push-based则相反。图9显示了这两种传输方法在8-A100 PCIe和8-A100 NVLink集群上的性能差异。结果表明，不同的互连方式可能有不同的偏好，因此应用自动调优来选择合适的传输方法。对于没有P2P访问的GPU，则使用NCCL【索引32，NCCL，2016，NVIDIA】的send/recv操作。

图9：在8-A100 PCIe或NVLink集群上，AllGather中不同m值以及(n,k)为(49152, 12288)时，pull-based和push-based数据传输的性能比较。

信号（Signals）实现 使用一个常规的32位GPU内存来实现信号。所有信号都连续分配，以便于预设和重置。在主机端，通过cuStreamWriteValue API设置信号；在内核端，WaitSignal通过自旋（spinning）实现。

通信瓦片大小 在AllGather中，通信瓦片化与GEMM计算瓦片化解耦，以避免干扰对性能敏感的GEMM瓦片化。独立调整通信瓦片化可以在重叠机会和通信效率之间找到最佳权衡。调优时，从与中等粒度分区相同的瓦片大小开始（即m除以张量并行度），然后不断除以2，直到等于GEMM的瓦片大小。图10显示通信瓦片大小确实影响整体性能，但由于没有明确的趋势，因此应用自动调优来选择最佳的瓦片化因子。

图10：在8-A100 NVLink集群上，AllGather中不同m值以及(n,k)为(49152, 12288)时，不同通信瓦片大小的性能结果。

瓦片间的通信顺序 通信顺序与瓦片坐标变换对齐，并根据网络拓扑选择以最小化整体延迟。
* 节点内NVLink：采用直接通信，以本地rank之后的环形顺序开始。
* 节点内PCIe：使用基于环的通信以有效利用带宽。
* 多节点：在节点内NVLink互连中，跨节点直接通信与本地节点内通信同时发出。在节点内PCIe互连中，为避免流量冲突，跨NUMA通信先发出，然后节点内NUMA通信和跨节点通信一起发出。

4.4 GEMM实现与自动调优

GEMM内核选择与实现 Flux适用于几乎所有类型的GEMM内核。考虑到GEMM性能至关重要，Ampere GPU上通常首选工作负载平衡的GEMM【索引35，Stream-k: Work-centric parallel decomposition for dense matrix-matrix multiplication on the gpu，2023，PPoPP】，而Hopper GPU上则首选warp或线程块特化的GEMM【索引21，CUTLASS，2024，NVIDIA】。Flux使用CUTLASS【索引21，CUTLASS，2024，NVIDIA】实现，以完全控制GEMM瓦片化和相应的prologue或epilogue融合。与传统GEMM库类似，Flux通过模板编写所有prologues、epilogues、GEMM算法和调优选项，允许通过选择合适的模板参数来自动调优内核。

A4 实验环境与结果

实验环境

软件配置:
- 核心库: CUTLASS 3.4.1, NVSHMEM 2.10.1
- 编译器: NVCC 11.8 (for A100), NVCC 12.2 (for H800)
- 数据类型: bfloat16
- 基线/对比: TransformerEngine 1.4.0, Megatron-LM core r0.4.0, Megatron-LLaMA, vLLM 0.2.1
硬件配置:
- A100 PCIe集群: 每节点8个A100 PCIe (80GB) GPU，节点内PCIe互连，2个100Gbps跨节点互连。
- A100 NVLink集群: 每节点8个A100 SXM4 (80GB) GPU，节点内NVLink互连，4个200Gbps跨节点互连。
- H800 NVLink集群: 每节点8个H800 SXM5 GPU，节点内NVLink互连，8个400Gbps跨节点互连。
模型与数据集:
- 模型: GPT-3 175B, Llama-2 70B
实验设置:
- 训练: 在128-GPU集群上进行，采用2路数据并行、8路流水线并行和8路张量并行。
- 推理 (Prefill): 在8-GPU集群上进行，batch size为8，sequence length为2048。
- 推理 (Decoding): 在8-GPU集群上进行，batch size为64或512。

实验结果

操作级性能评估

大m维度 (1024-8192) (图11, 12, 13):
- 与TransformerEngine对比:
  - A100 PCIe: Flux 提速 1.20x - 3.25x。
  - A100 NVLink: Flux 提速 1.01x - 1.33x。
  - H800 NVLink: Flux 提速 1.10x - 1.51x。
- 重叠效率:
  - A100 PCIe: Flux 达到 41%-57%，而TransformerEngine为-125%-36%。
  - A100 NVLink: Flux 达到 36%-96%，而TransformerEngine为-99%-74%。
  - H800 NVLink: Flux 达到 37%-93%，而TransformerEngine为-40%-80%。
    (注：负效率表示比非重叠基线更慢)

小m维度 (64, 512) (图14):
- 与TransformerEngine对比: Flux在大多数情况下提速1.03x - 4.68x，仅在H800 NVLink上m=64的ReduceScatter场景下有0.95x的降速。
- 重叠效率: Flux在A100 NVLink上最高达到88%，但在H800 NVLink上表现不佳（-165% 到 -82%），而TransformerEngine在所有小m场景下都表现出非常差的负效率。

16路张量并行 (图15):
- 与PyTorch基线对比: TransformerEngine不支持多节点重叠，因此只与PyTorch对比。
  - A100 PCIe: Flux 提速高达 1.32x，重叠效率 18%。
  - A100 NVLink: Flux 提速高达 1.57x，重叠效率 74%。
  - H800 NVLink: Flux 提速高达 1.55x，重叠效率 56%。

模型级性能评估

训练、Prefill和Decoding性能 (图16, 17):
- 与TransformerEngine对比:
  - A100 PCIe: 训练提速高达1.37x，prefill高达2.06x，decoding高达1.69x。
  - A100 NVLink: 训练提速高达1.04x，prefill高达1.14x，decoding高达2.10x。
  - H800 NVLink: 训练提速高达1.05x，prefill高达1.18x，decoding高达1.76x。
- 与Megatron-LM/vLLM基线对比:
  - A100 PCIe: 训练提速高达1.24x，prefill高达1.46x，decoding高达1.28x。
  - A100 NVLink: 训练提速高达1.05x，prefill高达1.45x，decoding高达1.30x。
  - H800 NVLink: 训练提速高达1.10x，prefill高达1.66x，decoding无加速。

图16：各种128-GPU集群上GPT-3 175B和Llama-2 70B的训练，以及各种8-GPU集群上prefill的端到端结果。

图17：各种8-GPU集群上GPT-3 175B和Llama-2 70B的decoding端到端结果。

讨论

小m尺寸和Decoding: 当m极小时，GEMM内核的warps数量较少，导致通过上下文切换隐藏延迟的效率降低。这解释了为什么Flux在某些极小m的情况下性能可能劣于非重叠基线。在H800上，8路张量并行使得TMA指令的存储大小进一步减小，导致效率下降。尽管如此，Flux在小m场景下仍显著优于TransformerEngine，后者因拆分内核而导致GPU利用率严重不足。
重叠效率与加速比: 模型级性能提升与两个因素高度相关：1) 张量并行的通信占比（如图1所示）；2) 重叠效率。当通信占比较低时（如A100 NVLink训练中为8%-11%），即使Flux有63%的平均重叠效率，整体加速比也仅为1.04x-1.05x。当通信占比较高时（如A100 PCIe训练和prefill中为40%-75%），Flux能带来显著的1.16x-1.46x加速。
高通信占比场景: 在A100 PCIe（慢速互连）和H800 NVLink（快速计算）等高通信占比的集群上，Flux的优势尤为明显。这证明了其算法、优化和实现的鲁棒性。在A100 PCIe集群上，Flux的性能有时甚至超过了仅通信的基线（NCCL），显示了其及其自动调优机制对互连环境的良好适应性。

A5 结论

对于使用张量并行的大型深度学习模型，通信重叠技术至关重要。传统的重叠技术在GPU上表现不佳。本文提出了一种新颖的技术Flux，通过将通信与相应计算过度分解并融合成一个单一的大内核，成功地解决了这些问题。实验证明，无论是在训练还是推理中，Flux都能显著减少暴露的通信时间，并有效提高系统的FLOPS利用率。

方法细节中的引用汇总

[12] Breaking the computation and communication abstraction barrier in distributed machine learning workloads, 2022, ASPLOS:
- 引用段落: Introduction, Section 2.2, Section 7
- 引用描述: 作为现有/先前（prior）的通信重叠技术被引用。具体地，它被描述为一种将粗粒度操作分解为一系列中等粒度操作，并通过精心调度来重叠通信与计算的方法，但该方法在GPU上存在局限性。在相关工作中，它被归类为一种将通信与逐元素操作融合，并以中等粒度方式与GEMM重叠的技术。
[13] Overlap communication with dependent computation via decomposition in large deep learning models, 2023, ASPLOS:
- 引用段落: Introduction, Section 2.1, Section 2.2, Section 4.1, Section 7
- 引用描述: 作为现有/先前的通信重叠技术被引用，与[12]和[14]类似，被描述为中等粒度的分解调度方法。在背景介绍中，其提出的“分片激活（sharded activation）”被作为本文讨论的常见划分策略之一。在讨论瓦片坐标变换时，提到该工作也应用了类似思想（改变拆分GEMM的执行顺序），但无法直接应用于不拆分GEMM的Flux。
[14] TransformerEngine, 2022, NVIDIA:
- 引用段落: Introduction, Section 2.2, Section 2.3, Section 3, Section 5, Section 7
- 引用描述: 作为现有/先前的通信重叠技术的代表实现被多次引用，并且是本文实验部分主要的对比对象。它被描述为一种中等粒度的分解方法，其局限性（如数据依赖阻止并发、拆分GEMM导致GPU利用率低）在背景部分被详细分析，并用图4的数据进行了说明。
[21] CUTLASS, 2024, NVIDIA:
- 引用段落: Introduction, Section 4.1, Section 4.2, Section 4.4, Section 5
- 引用描述: 作为Flux实现的核心底层库被引用。论文提到Flux采用NVIDIA CUTLASS构建，以实现模块化和自动调优。在方法细节中，提到从CUTLASS中得到启发（如瓦片坐标变换），并使用其特性（如warp/线程块特化）和扩展（CUTLASS EVT）来实现具体的优化。
[32] NCCL, 2016, NVIDIA:
- 引用段落: Section 2.3, Section 4.3, Section 5.1
- 引用描述: 作为广泛使用的标准GPU通信库，在本文中被用作非重叠通信的基线，用于计算“有效通信时间”和“重叠效率”。在AllGather的实现细节中，提到对于不支持P2P的GPU，使用NCCL的send/recv进行数据传输。
[33] NVSHMEM, 2020, NVIDIA:
- 引用段落: Section 3.1, Section 4.2, Section 5
- 引用描述: 作为支持跨节点P2P GPU内存访问的库被引用。论文指出，ReduceScatter的实现需要P2P支持，而NVSHMEM可以将此能力从节点内扩展到跨节点。在实现细节中，提到使用NVSHMEM的put API来实现跨节点的数据写入。

FLUX: FAST SOFTWARE-BASED COMMUNICATION OVERLAP ON GPUS THROUGH KERNEL FUSION