SCAPE：利用极端稀疏通信实现准确高效的LLM训练

# SCAPE: 基于极端稀疏通信的准确高效大语言模型训练
来源：https://arxiv.org/html/2607.01678
Mingkai Zheng, Junlin Chen, Haotian Xie, and Zhao Zhang Rutgers University \{mz687, junlin\.chen110, haotian\.xie, zhao\.zhang\}@rutgers\.edu

###### 摘要

通信在大语言模型（LLM）预训练中的成本占比日益增加，尤其是在数据并行和分片训练方案下，梯度同步和参数重建的开销随模型规模和系统规模增大而增长。现有通信压缩方法要么对原始梯度进行稀疏化（在高稀疏度下对现代 Adam 类优化器可能不稳定），要么对通信进行量化（其节省量从根本上受位宽限制，且通常引入额外的运行时开销）。我们提出 SCAPE，一种面向 LLM 训练的通信高效分布式优化器，它利用 AdamS 一阶矩的稳定性，在不损失 LLM 质量的前提下实现激进稀疏化。SCAPE 不直接从原始梯度构建掩码，而是基于一阶矩统计量推导掩码，将掩码生成任务按优化器分片方式分布到各工作节点，并将掩码使用延迟一个步骤，使掩码同步能与计算重叠。SCAPE 还从单个同步的稀疏缓冲区重建二阶矩更新所需的量，避免了一次额外的集合通信。我们在 Megatron-LM 中实现 SCAPE，并在 TACC Vista 上使用 32 块 NVIDIA GH200 GPU，通过预训练 OpenWebText 上的 GPT-345M 和 SlimPajama-6B 上的 Llama-500M 来评估其收敛性。在两种模型中，SCAPE 在 90% 和 99% 稀疏度下均保持了训练稳定性、验证损失和下游任务准确率。对于 Llama-500M，SCAPE 将端到端预训练时间最多减少 43.3%，同时保持与密集 AdamW 和 AdamS 相当的模型质量。对于 Llama-1.8B，与密集 AdamS 相比，SCAPE 每步加速比最高达 3.26 倍。

## I 引言

规模不断增长的大语言模型（LLM）在数学推理[1]、代码生成[4]和自动化实验室[30]等多个领域取得了无与伦比的性能。最先进的 LLM 训练需要大量图形处理单元（GPU）以实现更快的训练并容纳庞大的模型参数和优化器状态。传统的数据并行（DP）策略将模型和优化器状态（例如 AdamW 中的一阶和二阶动量）复制到各 GPU 上，并在它们之间分配一个小批次数据。现代分片数据并行策略，如 ZeRO、FSDP 和 Megatron-LM，通过将模型和优化器状态分区到各 GPU 上来减少空间冗余。在两种情况下，通信（即 all-reduce 和 all-gather）都是扩展瓶颈，因为它依赖于模型规模和系统规模，如图 1 所示。

降低通信量是减少通信开销的有效方法。研究人员探索了多种方法。DGC、DeMo、EDGC、oktopk 和 Radius 通过错误反馈传递 top-k 梯度以保持模型性能。另一条研究路线利用量化（低位表示）技术。QSDP、ZeRO++ 和 SDP4bit 在分片数据并行策略中减少了模型参数和梯度的位数。

参照图注图 1：使用 Megatron-LM 的分布式优化器在 NVIDIA GH200 GPU 上预训练 Llama-500M（序列长度 4K）和 Llama-1.8B（序列长度 2K）的扩展瓶颈。由于 GH200 内存有限，Llama-1.8B 使用序列长度 2K。现有的 top-k 方法受限于对现代优化器的适用性或需要低稀疏度以保证模型性能。传统的 top-k 方法，如 DGC 和 DeMo，旨在稀疏化 SGD 优化器中的动量，而 SGD 在现代 LLM 训练中已很少使用。Radius 为 AdamW 优化器提出了一种纠错机制，可实现 40% 的稀疏度而不影响下游任务性能。量化方法（包括 QSDP、ZeRO++ 和 SDP4bit）的通信量降低能力从根本上受限于表示每个梯度所需的位数。此外，如果没有 NCCL 后端支持，这些基于量化的方法依赖昂贵的 all-to-all 通信，带来了高的内存和通信开销。

在这项工作中，我们提出以下研究问题：能否在不损害模型性能的前提下将通信量减少到极致？为此，我们提出 SCAPE，一种通信高效的分布式优化器，它稀疏化优化器状态而非梯度。SCAPE 基于 AdamS 构建，并受两个重要见解启发：

- 与 AdamW 相比，AdamS 在 LLM 预训练的高稀疏度下表现出显著更小的残差增长，表明其对陈旧错误反馈动态具有更好的鲁棒性。
- 一阶矩的 top-k 结构在相邻两步之间是时间稳定的，表明延迟掩码复用是可行的。

基于这些观察，SCAPE 引入了一种与优化器分片对齐的分区掩码刷新机制。SCAPE 从基于动量的统计量而非原始梯度计算掩码，并将其使用延迟一个步骤，使掩码同步能与计算重叠。此外，SCAPE 使用单个同步的稀疏缓冲区重建二阶矩更新所需的量，避免了一次额外的 all-reduce。当启用分片数据并行时，模型和优化器状态分布到所有 GPU，SCAPE 使用稀疏性压缩梯度同步的 reduce-scatter 和模型重建的 all-gather 的通信量。

我们在 Megatron-LM 中实现 SCAPE，并通过在 OpenWebText 数据集上预训练 GPT-345M 和在 SlimPajama-6B 数据集上预训练 Llama-500M 来经验性地验证其收敛性。我们在 TACC Vista 超级计算机上运行实验，并在广泛的下游任务套件上进行评估。实验结果表明，在 90% 和 99% 稀疏度下，SCAPE 可将 Llama-500M 在 32 块 GH200 GPU 上的预训练时间分别减少 35.6% 和 43.3%。更重要的是，在这种高稀疏度下，SCAPE 保持最终训练和验证损失接近密集的 AdamS 和 AdamW，且不影响其在 LAMBADA、SuperGLUE、PIQA、MMLU 和 ARC 等下游基准测试上的性能。

SCAPE 预计在各种 GPU 和互连配置的 GPU 集群上都能有效，尽管其改进程度可能因硬件而异。SCAPE 利用 GH200 超级芯片上 CPU 与 GPU 之间 900 GB/s 的 NVLink-C2C 来卸载错误反馈缓冲区。缓冲区卸载可能会在具有 PCIe 连接的其他 GPU 集群上引入额外开销。然而，与 SCAPE 的整体改进相比，缓冲区卸载的影响很小。

## II 背景

本节提供关于使用 Megatron-LM 的分布式训练（含分布式优化器）、梯度的 top-k 稀疏化以及 AdamS 优化器的背景信息。

### II-A 使用分片数据并行分布式优化器的训练

在传统分布式数据并行（DDP）训练中，数据并行（DP）组内的每个工作节点（例如 GPU）复制所有模型参数、梯度和优化器状态，并拥有一份微批训练样本。所有工作节点使用本地样本完成梯度计算后，它们通过环形 all-reduce 操作同步梯度，然后更新优化器状态和模型参数。假设有 N 个工作节点，梯度大小为 W（与模型大小相同）。此通信成本可分解为 reduce-scatter 后接 all-gather 操作，其模型为

T_环形-all-reduce = T_reduce-scatter + T_all-gather
= 2(N-1)[α + β W/N],

其中 α 是初始化消息的固定延迟，β 是带宽的倒数。随着工作节点数 N 增加，2Wβ(N-1)/N 主要由梯度大小 W 和带宽倒数 β 决定。

将整个模型和优化器状态复制到 DP 组中的所有工作节点上，保证了前向和后向计算无需通信即可完成。然而，优化器状态通常需要全精度（FP32）来保证训练稳定性，这成为训练参数日益增长的大语言模型（LLM）的主要内存瓶颈。因此，提出了许多框架来缓解这一瓶颈。例如，Megatron-LM 的分片数据并行分布式优化器将优化器状态和模型参数均匀地分区到所有 DP 工作节点上。因此，在前向传播之前，工作节点需要执行 all-gather 以收集彼此拥有的模型分片。注意，模型参数可以预取，使得 all-gather 的通信与计算重叠。后向传播完成后，每个 DP 工作节点拥有完整的梯度。由于优化器状态已被均匀分区，梯度的 all-reduce 同步可以替换为 reduce-scatter，而不需要收集（即避免 all-gather）。此过程在图 2 中说明。

参照图注图 2：带有分片数据并行分布式优化器的 Megatron-LM

### II-B 梯度 Top-k 稀疏化

Top-k 稀疏化是一种通过减小公式 1 中的 W 来减少 T_环形-all-reduce 的技术，前提是网络带宽 1/β 固定。具体来说，在同步梯度 g 之前，我们执行 top-k 操作选择 k 个最大幅度的梯度值，然后使用压缩缓冲区同步这些 top-k 值。这有效地将通信量从 W 减少到 kW。

注意，稀疏化后的 g 是有偏的，因为 top-k 操作丢弃了梯度 g 中除最大幅度条目之外的所有内容，因此不保留原始梯度的期望。因此，在每个迭代直接应用 top-k 稀疏化可能会累积压缩误差并降低收敛性。为了缓解此问题，先前的工作通常使用误差反馈，该方法累积当前步骤未发送的残差，并在下次 top-k 选择之前将其加回到梯度中。令 e_t 表示步骤 t 的残差缓冲区，则通信的梯度为 g̃_t = TopK(g_t + e_t)，残差更新为 e_{t+1} = g_t + e_t - g̃_t。这样，一次迭代中丢弃的信息不会永久丢失，而是逐渐重新注入后续通信中，提高了优化稳定性，同时保留了稀疏化的通信节省。然而，这需要一个与 g 大小相同且通常为 FP32 的专用误差反馈缓冲区，以保持精度。

由于每个工作节点用于计算步骤 t 梯度 g_t 的数据样本是不同的，它们的 top-k 梯度索引也不同。因此，直接使用 all-reduce 计算稀疏化梯度的平均值会导致错误结果。有两种方法可以解决此问题。第一种是执行两次 all-gather 集合通信来收集 top-k 掩码和相应的值；第二种方法是先使用 all-gather 形成同步的 top-k 掩码，然后使用 all-reduce 同步由全局同步掩码选择的梯度条目。然而，第二种方法存在填充效应。如文献所述，当工作节点数 N 足够大时，同步后的全局 top-k 掩码是密集的而非稀疏的。

### II-C AdamS 优化器

AdamS 是一种新的类 Adam 优化器，它不保留二阶矩 v_t 作为 g_t^2 的指数移动平均（EMA），而是使用一阶矩 m_t 本身作为预条件器：

m_t = β1 m_{t-1} + (1-β1) g_t,
v_t = β2 m_{t-1}^{⊙2} + (1-β2) g_t^{⊙2},
w_t = w_{t-1} - η ( m_t / (√v_t + ε) + λ w_{t-1} ),

其中 β1 和 β2 是计算运行平均的系数，w_t 是步骤 t 的模型参数，η 是学习率，λ 是权重衰减超参数，ε 保持数值稳定性。注意，公式 2 中的更新度量已简化用于演示目的，实践中 m_t 和 v_t 应乘以它们的偏差修正项 1/(1-β1^t) 和 1/(1-β2^t)。

与 AdamW（即带权重衰减的 Adam）相比，AdamS 在训练稳定性、模型质量和内存消耗方面具有优越性能。首先，使用相同的超参数集合时，AdamS 对噪声比 AdamW 更鲁棒，避免了训练损失曲线中的大尖峰。由于 m_t 是梯度 g_t 的 EMA，它是比 AdamW 的 v_t（梯度平方 g_t^2 的 EMA）更平滑的预条件器。当在梯度范数通常较小的稳定区域出现大梯度时，它对这种噪声更鲁棒。其次，AdamS 达到了与 AdamW 相当甚至在某些情况下更优的模型质量。