通过稀疏通信提升大模型训练效率,90%以上通信量减少仍保质量。
SCAPE: Accurate and Efficient LLM Training with Extreme Sparse Communication

- 用一阶动量统计生成稀疏掩码,避免高稀疏度下优化不稳。
- 在32张GH200上训练Llama-500M,通信量减99%仍保持准确率。
- 适合大规模分布式训练场景,尤其对通信瓶颈敏感的团队。
通信成本正日益成为大语言模型预训练的主要开销,尤其在数据并行与分片训练中,梯度同步和参数重建开销随模型规模与系统规模增长。现有方法或稀疏化原始梯度(高稀疏时对Adam类优化器不稳定),或量化通信(节省受限于位宽且增加运行开销)。我们提出SCAPE,一种面向大模型训练的高效分布式优化器,利用AdamS一阶动量的稳定性,实现无损的激进稀疏化。SCAPE不基于原始梯度构造掩码,而是从一阶动量统计中生成,并将掩码生成分摊至各工作节点以匹配优化器分片,延迟掩码使用一步,使掩码同步可与计算重叠。同时,仅通过一个同步的稀疏缓冲区重构二阶动量更新所需量,避免额外集体通信。我们在Megatron-LM中实现SCAPE,评估其在TACC Vista上使用32张NVIDIA GH200 GPU预训练GPT-345M(OpenWebText)与Llama-500M(SlimPajama-6B)的收敛性。两个模型在90%与99%稀疏度下均保持训练稳定性、验证损失与下游任务准确率。对Llama-500M,SCAPE将端到端预训练耗时最多减少43.3%,模型质量媲美稠密AdamW与AdamS;对Llama-1.8B,SCAPE相比稠密AdamS每步最高提速3.26倍。
原文摘要 · Abstract (English)
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.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。