LLM 训练中的 Overlap 优化（二）Google - Intra-layer Overlapping via Kernel Fusion、AMD - T3_transparent t racking & t riggering for fine-grain-CSDN博客

本文链接：https://blog.csdn.net/star_nwe/article/details/145223748

四、Google - Intra-layer Overlapping via Kernel Fusion

4.1 摘要

对应的论文为：[ASPLOS 23] Overlap Communication with Dependent Computation via Decomposition in Large Deep Learning Models [3]

在大规模深度学习模型训练中，层内模型并行化产生的通信开销可能占据整体训练时间的显著部分，而层内模型并行化又对支持大型深度学习模型至关重要。因此作者提出了一种新颖计算，通过计算与通信的 Overlap 来有效降低通信开销。

该技术将识别出的原始集合通信与依赖的计算操作分解为一系列更细粒度的操作，通过创造更多 Overlap 机会并并行执行新生成的细粒度通信与计算操作，可以有效隐藏数据传输时延，实现更优的系统利用率。

在 TPU v4 Pod 上评估不同规模的大模型（10B - 1T 参数量），所提方案可以实现 1.14x 到 1.38x 的吞吐提升。在 1024 TPU 集群中，500B 参数量语言模型训练可以实现 72% 的峰值 FLOPs 利用率。

4.2 方法

4.2.1 背景

如下图 Figure 1 所示，不同规模模型在 128-2048 TPU 上训练，通信开销可达 22%-42%：

在这里插入图片描述

4.2.2 方案概览

如下图 Figure 4 展示了 AllGather 场景中的执行流程。假设数据 A（比如模型权重）在初始阶段已被切分，每个设备各持有 A 的一个分片，A0 位于设备 0，A1 位于设备 1。

现有系统中：两个设备均需要通过 AllGather 操作得到完整的数据 A，即 [A0, A1]，然后开始相应的计算。
所提系统中：并不用等待全部数据准备就绪再启动计算。
每个设备异步发送存储在当前设备的数据到其他设备（比如设备 1 异步发送 A1 到设备 0），同时利用已有数据开始启动计算，这样设备即在计算，也在通信。
当之前结果计算完成，并且从其他设备接收完成（比如设备 1 的 [A1, B1] 已经计算完，并且已经接收完 A0），开始启动新数据的计算（比如设备 1 上的 [A0, B1]）。
为了得到最终结果，每个部分结果需要额外执行一次 Dynamic Updata Slice 操作。
通过执行多次上述操作可以获得最终结果，确切的步骤次数取决于 A 的切片数。

在这里插入图片描述

同样地，ReduceScatter 操作可与相应的计算过程 Overlap 执行，如下图 Figure 5 所示。在此例中，C0（C00 与 C01 之和）及 C1（C10 与 C11 之和）分别为 C 在设备 0 和设备 1 上进行 ReduceScatter 后的操作分片。由于基于计算结果进行通信传输，在此情形下，各设备需要异步传输累加结果分片而非操作数，累加结果分片在各设备上初始化为 0：

每轮迭代开始时，各设备异步发送累加结果分片到另一个设备（例如，首轮迭代中设备 0 发送切片 C0 到设备 C1），并与此同时启动部分 Einsum 计算。
计算完成后，部分结果在迭代末尾被加到接收的累加结果分片，比如首轮迭代的结果 C10 与从设备 1 接收的结果分片 C1。

在这里插入图片描述

Kernel Fusion 是一种有效减少慢速主内存访问和 Kernel 启动开销的方案，作为最重要的优化手段之一，Kernel Fusion 在 XLA 中通过启发式方法自动执行。因此，针对本文的方案作者也会进一步应用 Kernel Fusion。然而，基于默认启发式构建的某些融合操作可能损害 Overlap 性能。如下图 Figure 11 所示，11a 展示了一个简化的图结构，为默认的融合策略，其中的灰色方框为融合节点，白色方框表示一个或多个融合的高阶算子指令。其中两个 Einsum，Einsum_1 有一个异步的 CollectivePermuteDone 输入，由于 Einsum_0 与 CollectivePermuteDone 相互独立，预期其能与异步数据通信并行执行，以实现 Overlap。然而，与 Einsum_0 融合的加法操作在 Fusion_0 与 CollectivePermuteDone 之间引入了数据依赖，导致第三个节点顺序执行。为了避免这种不良融合，启发式策略调整为先将 Add 操作与具有异步 CollectivePermuteDone 操作的 Einsum 进行融合，新生成的图结构如图 11b 所示，数据通信得以成功与 Fusion_0 Overlap。

在这里插入图片描述

4.3 结果

如下图 Figure 12 所示为不同模型优化前后可以达到的峰值 TFLOPS，可以看出，优化后有比较明显的提升：

在这里插入图片描述

五、AMD - T3

5.1 摘要

对应的论文为：[2401.16677] T3: Transparent Tracking & Triggering for Fine-grained Overlap of Compute & Collectives [4]

LLM 在训练与推理过程中，尽管某些分布式技术能够通过计算与通信 Overlap 来隐藏通信开销，但诸如 TP 等技术，其通信与模型执行本质上具有序列化特性。为隐藏这种序列化通信，常见做法是以细粒度方式将其数据生成操作交错进行。然而，在软件层面实现通信与计算的细粒度交错颇为复杂，此外，并发执行通常都要求计算与通信共享计算和存储资源，导致资源竞争，从而削弱 Overlap 的效能。

为应对这些挑战，作者提出 T3 方案，通过硬件与软件协同设计，透明地 Overlap 序列化通信，同时最小化与计算的资源竞争。

在软件层面，T3 通过简单配置生产者的输出地址空间，透明地将生产操作与后续通信融合，仅需少量软件改动。
在硬件层面，T3 引入轻量级的追踪与触发机制，以协调生产者的计算与通信。同时，利用计算增强型内存处理通信伴随的计算任务。由此，T3 减少了资源竞争，并高效地实现了序列化通信与计算的 Overlap。

对于 T-NLG 等重要 Transformer 模型，T3 使通信密集型子层的平均加速比可以达到 30%（最高47%），平均减少 22% 的传输量（最高 36%）。此外，随着模型规模的扩大，T3 的优势依然显著：在约 500B 参数的 PALM 和 MT-NLG 模型中，子层的平均加速比为 29%。

PS：T3 依赖于特定的硬件特性，如计算增强型存储器（NMC），这可能需要新的硬件支持或者对现有硬件的修改，这也许是其落地应用的最大挑战。

5.2 方法

5.2.1 背景

如下图 Figure 2 所示，Transformer 模型通常会采用 TP 来切分模型，以便能训练更大规模模型。然而，这一过程要求在层与层之间执行 AllReduce 操作。其中 Figure 2b 为未切分的操作，而 Figure 2c 为切分后跨两个设备的操作。在 2b 中，每个设备仅持有部分权重。连续两个矩阵乘可以先按列切，再按行切，之后通过一个 AllReduce 操作获取完整结果。

在这里插入图片描述

而这些串行的 AllReduce 操作可能成为性能瓶颈。如下图 Figure 4 展示了多种常见 Transformer 模型中使用 TP 后各操作的执行时间占比。可以看出，其 AllReduce（ReduceScatter + AllGather）的通信占比甚至比 GEMM 还长，比如 Megatron-GPT2 和 T-NLG 在训练与推理（Prefill）中，通信时间分别高达 34% 和 43%。而且往往算力增长比网络带宽增长更快，这也会进一步加大通信的占比。

在这里插入图片描述

5.2.2 对比

在 T3 中，作者也提到上述 Microsoft - CoCoNet 和 Google Intra-layer Overlapping via Kernel Fusion 两篇论文：

Microsoft：需要昂贵的细粒度同步机制。
Google：需要对矩阵乘法 Kernel 进行改动，并且可能对 GPU 软件基础设施造成干扰。
此外，上述两个工作中计算与通信的 Overlap 会竞争计算与内存带宽，从而削弱 Overlap 的实际效果。

5.2.3 方案概览

现代 GPU 首先执行生产者 GEMM 并将结果存储于本地内存中。随后，启动集合通信操作。相比之下，T3 系统在 GEMM 生成数据的同时立即启动集合通信操作，以实现细粒度的 Overlap 执行。如下图 Figure 1 所示：

T3 采用追踪与触发机制监控 GEMM 与集合通信操作的进展，并协调通信流程，无需额外计算单元（CU）参与。
此外，T3 利用近内存计算（Near Memory Computing, NMC）进行 Reduce 操作，以减少因通信产生的内存访问。
最终，这些优化可以在几乎不修改 Kernel 代码的情况下透明地实现。

在这里插入图片描述

如下图 Figure 7 展示了 4 个 GPU 下，ReduceScatter（RS）操作与 GEMM 的 Overlap 执行情况。该 GEMM 根据输入数据和 Kernel 实现分为多个工作组（WG）阶段执行，而 RS 则依据参与设备数量分为多个 Step 进行。为了简化图示，图中 GEMM 阶段数比所需的 Ring Step 数多 1。在每一个 Step 中，GEMM 某一阶段的执行和其输出的 Reduce 与前一 Step 输出的通信并行进行。在第一个 Step 中，GEMM 直接将输出传输到远程设备（remote_update）。后续的稳态 Step 则需通过 DMA（dma_update）完成。对于 N 台设备，稳态 Step 需执行 N-2 次，以处理不同的数据块。

以稳态下的 GPU 0 为例，对于其在 Step 2 的操作，如下图 Figure 7 所示，GPU 0 在执行并生成GEMM Stage 3 输出的同时，通过 DMA 接收来自邻近设备 GPU 1 的 Stage 3 输出副本（蓝色）。这一过程与 GPU 0 将 GEMM 第二阶段数据（黄色）Reduce 副本通过 DMA 传输至 GPU 3 的操作并行进行，从而实现了通信重叠。在 local_update 和 dma_update 中，T3 利用近内存计算（NMC）进行原子内存位置更新。为生成 Stage 3 数据块的部分 Reduce 副本，无需额外读取或占用 GPU 计算单元。一旦操作完成，GPU 0 即启动对数据块的 dma_update，将其传输到临近设备 GPU 3 的内存中，如下图的 Step 3 所示。

5.2.4 硬件设计

通过一款轻量级且可编程的硬件追踪单一，可以实现上述自动更新追踪及 DMA 触发机制，从而进一步降低对 GPU 计算单元（CU）的依赖。远程/DMA 更新操作则通过配置 GEMM 输出地址映射，辅以微小的应用程序和 Kernel 修改，得以透明执行。

为了提升 T3 的性能，作者还对运行时和硬件进行了细微调整。为实现如图 Figure 7 中 GEMM 与 RS 的完美 Overlap，作者还对跨 GPU 的 GEMM 工作组调度进行了错位安排。此外，还通过引入一种简单而有效的内存控制仲裁（Memory Controller Arbitration，MCA）策略来增强内存系统，以管理计算与通信之间的内存竞争问题。

在这里插入图片描述

如下图 Figure 8 展示了搭载 T3 增强功能（以橙色标注）的 GPU 执行上述稳态步骤的情况。GPU 执行 GEMM 运算，为某一 Stage 生成 local 更新（L1）。同时，GPU 接收针对同一 Stage 的 DMA 更新（D1a），并向上一阶段发送 DMA 更新（D1b）。在内存控制器处，经过改进的MCA 策略对 local 与 DMA 流量进行仲裁，以避免争用。随后，更新数据被传输至经过 NMC 增强的 DRAM（L2a，D2a），同时 Tracker 记录其进度（L2b，D2b）。一旦 Tracker 检测到内存区域所需的 local 和 DMA 更新，便会触发这些更新通过 DMA 传输至相邻 GPU（L3）。

在这里插入图片描述