SCAPE：通过极其稀疏的通信进行准确高效的 LLM 训练

通信越来越主导 大语言模型 (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$\times$。