SCAPE: Accurate and Efficient LLM Training with Extreme Sparse Communication
作者: Mingkai Zheng, Junlin Chen, Haotian Xie, Zhao Zhang
分类: cs.LG, cs.DC
发布日期: 2026-07-02
💡 一句话要点
提出SCAPE以解决大规模语言模型训练中的通信效率问题
🎯 匹配领域: 支柱九:具身大模型 (Embodied Foundation Models)
关键词: 大规模语言模型 稀疏化训练 分布式优化 通信效率 深度学习
📋 核心要点
- 现有大规模语言模型训练方法在通信效率上存在显著瓶颈,尤其是在高稀疏度下不稳定。
- SCAPE通过利用AdamS的第一矩特性,提出了一种新的稀疏化策略,增强了通信效率。
- 在GPT-345M和Llama-500M的预训练实验中,SCAPE在高达99%的稀疏度下保持了模型质量,并显著缩短了训练时间。
📝 摘要(中文)
随着大规模语言模型(LLM)预训练的规模扩大,通信成本逐渐成为主要瓶颈,尤其是在数据并行和分片训练方案中,梯度同步和参数重构的开销显著增加。现有的通信减少方法要么通过稀疏化原始梯度来降低通信量,但在高稀疏度下对现代Adam风格优化器不稳定;要么通过量化通信来节省开销,但其节省效果受到位宽的限制,并且往往会引入额外的运行时开销。本文提出了SCAPE,一种高效的分布式优化器,利用AdamS的第一矩特性实现了激进的稀疏化而不损失模型质量。SCAPE通过从第一矩统计中生成掩码,分布式生成掩码并延迟掩码的使用,从而实现了计算与同步的重叠。实验表明,SCAPE在保持模型质量的同时,显著提高了训练速度。
🔬 方法详解
问题定义:本文旨在解决大规模语言模型训练中通信成本高的问题,现有方法在高稀疏度下对优化器的稳定性影响较大,且量化通信的节省效果有限。
核心思路:SCAPE通过利用AdamS的第一矩特性,允许在不损失模型质量的情况下进行激进的稀疏化。通过从第一矩统计中生成掩码,SCAPE能够有效降低通信量。
技术框架:SCAPE的整体架构包括掩码生成、掩码同步和稀疏缓冲区重构三个主要模块。掩码生成在各个工作节点上并行进行,掩码的使用被延迟到下一步,以实现计算与通信的重叠。
关键创新:SCAPE的主要创新在于其掩码生成机制,利用第一矩统计而非原始梯度生成掩码,从而提高了稀疏化的稳定性和效率。这一设计与现有方法的根本区别在于其对优化器的适应性。
关键设计:SCAPE在参数设置上进行了优化,确保掩码生成与优化器分片对齐,同时通过单一同步稀疏缓冲区重构所需的二阶矩更新量,避免了额外的全局通信开销。
🖼️ 关键图片
📊 实验亮点
在实验中,SCAPE在Llama-500M模型上实现了高达43.3%的训练时间缩短,同时在90%和99%稀疏度下保持了训练稳定性和模型质量。与传统的AdamW和AdamS相比,SCAPE在Llama-1.8B模型上每步训练速度提升达3.26倍。
🎯 应用场景
SCAPE的研究成果在大规模语言模型的训练中具有广泛的应用潜力,尤其是在需要高效通信的分布式训练环境中。其创新的稀疏化策略能够显著降低训练时间,提高资源利用率,未来可扩展至更多深度学习任务和模型架构。
📄 摘要(原文)
Communication increasingly dominates the cost of Large Language Model (LLM) pre-training, especially under data-parallel and sharded training schemes, where gradient synchronization and parameter reconstruction overhead increase with model size and system scale. Existing communication-reduction methods either sparsify raw gradients, which can be unstable for modern Adam-style optimizers at high sparsity, or quantize communication, whose savings are fundamentally bounded by bit width and often incur additional runtime overhead. We present SCAPE, a communication-efficient distributed optimizer for LLM training that exploits the stability of AdamS's first-moment to enable aggressive sparsification without loss of LLM quality. Instead of constructing masks from raw gradients, SCAPE derives them from first-moment-based statistics, partitions mask generation across workers to align with optimizer sharding, and delays mask usage by one step so that mask synchronization can overlap with computation. SCAPE also reconstructs the quantities required for second-moment updates from a single synchronized sparse buffer, avoiding an additional collective. We implement SCAPE in Megatron-LM and evaluate its convergence by pre-training GPT-345M on OpenWebText and Llama-500M on SlimPajama-6B using 32 NVIDIA GH200 GPUs on TACC Vista. In both models, SCAPE preserves training stability, validation loss, and downstream task accuracy under 90\% and 99\% sparsity. For Llama-500M, SCAPE reduces end-to-end pre-training wall-clock time by up to 43.3\% while maintaining model quality comparable to dense AdamW and AdamS. For Llama-1.8B, SCAPE achieves up to 3.26$\times$ speedup per step compared to dense AdamS.