SCAPE: Accurate and Efficient LLM Training with Extreme Sparse Communication

2026-07-02

AI Summary

This work addresses the escalating communication overhead in large language model pretraining, which intensifies with model scale and challenges existing sparsification methods to simultaneously maintain training stability and model quality under high compression ratios. The authors propose SCAPE, a distributed optimizer that leverages the statistical properties of Adam’s first moment to generate sparse masks. SCAPE introduces three key innovations: partitioning mask generation across nodes, delaying mask application to overlap communication with computation, and reconstructing second-moment updates from a single shared sparse buffer. Evaluated on Llama models, SCAPE preserves validation loss and downstream task performance nearly on par with dense training at sparsity levels of 90%–99%, achieving a 43.3% reduction in end-to-end training time on Llama-500M and a 3.26× speedup per training step on Llama-1.8B.

Abstract

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.

Problem

• communication overhead
• large language model training
• gradient sparsification
• distributed optimization
• training efficiency

Innovation

• sparse communication
• distributed optimization
• LLM training
• gradient sparsification
• Adam optimizer