Search papers, labs, and topics across Lattice.
This paper introduces SCAPE, a novel communication-efficient distributed optimizer for Large Language Model (LLM) training that leverages the stability of AdamS's first-moment statistics to enable aggressive sparsification without compromising model quality. By generating sparsification masks from these first-moment statistics and implementing a partitioned mask generation strategy, SCAPE significantly reduces communication overhead while maintaining training stability and downstream task performance. The results show that SCAPE can achieve up to 43.3% reduction in pre-training time and a 3.26x speedup per step compared to traditional dense optimizers, all while preserving model accuracy under high sparsity levels.
SCAPE achieves up to 43.3% faster LLM pre-training with 90% and 99% sparsity, all while maintaining model quality.
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.
