SCAPE: Accurate and Efficient LLM Training with Extreme Sparse Communication

Abstract: Communication increasingly dominates the cost of 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.

• The paper introduces a distributed optimizer that uses first-moment-driven masking to minimize communication overhead in LLM training.
• It achieves up to 3.26x per-step speedup and up to 43.3% reduction in wall-clock time while maintaining accuracy similar to dense training.
• The methodology combines delayed mask synchronization and sharded alignment to ensure stable, efficient sparse communications across large-scale models.

SCAPE: Accurate and Efficient LLM Training with Extreme Sparse Communication

Motivation and Problem Statement

LLMs present significant scalability challenges in distributed training, primarily due to the communication bottleneck imposed by synchronizing gradients and optimizer states across accelerators. As model and data scales increase, this communication cost, especially under sharded and data-parallel schemes, becomes the dominant system bottleneck (Figure 1).

Figure 1: Scaling bottleneck for pre-training Llama-500M and Llama-1.8B on NVIDIA GH200 GPUs using Megatron-LM with distributed optimizer.

Traditional approaches for mitigating communication costs include gradient sparsification with error feedback and quantization. However, these approaches either encounter instability at high sparsity rates (especially with Adam-style optimizers) or hit a ceiling due to the minimal bitwidth constraint in quantization. Moreover, many sparsification strategies are not well-integrated with modern sharded optimizer architectures, suffer from the fill-in effect in large clusters, or have performance penalties from misalignment with the optimizer's update dynamics.

SCAPE: Methodology and System Design

SCAPE (Sparse Communication via AdamS Preconditioned Events) introduces a distributed optimizer that enables communication-efficient LLM training by shifting the sparsification target to the first-moment statistics of the optimizer—specifically, the AdamS preconditioner—and aligning mask generation with optimizer sharding. Key elements include:

1. First-Moment-Driven Masking:

Instead of conventional gradient sparsification, SCAPE constructs the top-$k$ mask from the magnitude of the first-moment (momentum) buffer, leveraging its observed temporal stability and greater resilience to error-feedback staleness, especially with the AdamS optimizer.

2. Delayed Mask Synchronization:

Mask usage is delayed by one step (i.e., a mask computed at step $t$ is applied at step $t+1$), allowing the expensive all-gather mask synchronization to be fully overlapped with backward computation, significantly reducing effective communication overhead (Figure 2).

Figure 2: Communication pattern for refreshing the top-k mask, aligning sharded mask generation and overlapping mask all-gather with computation.

3. Single-Buffer Communication for Updates:

SCAPE transmits only the sparsified first-moment buffer, and, through an update reconstruction scheme, derives all quantities needed for both first- and second-moment updates from this single synchronized sparse buffer, avoiding redundant collective operations.

4. Sharded Alignment and Fill-in Immunity:

Mask computation and synchronization are partitioned in accordance with sharded optimizer state allocation (Megatron-LM style) (Figure 3). This prevents the fill-in effect observed in legacy sparse all-reduce schemes and yields communication reduction that scales linearly with imposed sparsity.

Figure 3: Megatron-LM with sharded data parallel distributed optimizer—optimizer states and parameter shards are aligned with communication responsibilities.

Empirical Findings and Observations

The paper provides empirical evidence that directly sparsifying gradients with error feedback causes rapid increases in the residual norm and diverging loss, especially for AdamW at realistic sparsity levels. In contrast, AdamS demonstrates:

• Lower residual growth—indicative of more stable error-feedback dynamics.
• A more centralized gradient distribution, reducing the risk of loss spikes when untransmitted gradients accumulate (Figure 4).

Figure 4: AdamW optimizer schematic; AdamS leverages first-moment preconditioning.

Temporal mask stability analysis further indicates that top-$k$ support derived from AdamS’s first-moment displays high overlap across consecutive steps, even at 10% density, which is crucial for enabling the system’s delayed mask usage scheme.

Experimental Evaluation

Experiments are conducted on pre-training GPT-345M and Llama-500M models using the TACC Vista supercomputer (32 NVIDIA GH200 GPUs). Key protocols:

• Llama-500M, SlimPajama-6B dataset, 100k steps, 4096 sequence length.
• GPT-345M, OpenWebText, 100k steps.

Training Loss and Efficiency

SCAPE maintains final training and validation losses that are nearly identical to, or slightly better than, those obtained by dense AdamS/AdamW, even at sparsity rates of 90% and 99%. For Llama-500M, SCAPE reduces wall-clock pre-training time by up to 43.3% without loss in model quality (Figures 8, 9).

Figure 5: Training loss trajectory for Llama-500M under different optimizers and SCAPE densities.

Figure 6: Validation loss trajectory for Llama-500M matching dense baselines despite aggressive communication sparsity.

Per-Step Time and Scaling

SCAPE significantly reduces per-step iteration time:

• Llama-1.8B achieves up to 3.26$\times$ per-step speedup relative to dense AdamS.
• Llama-500M attains 1.93$\times$ speedup at 99% sparsity (Figure 7).

Figure 7: Per-step time for Llama-500M as a function of GPU count and optimizer configuration.

Strong Scaling and Memory

SCAPE consistently improves strong scaling efficiency on up to 64 GPUs, with efficiency gains most pronounced when the communication/computation ratio is highest (i.e., for larger models and batch sizes).

Memory overhead from SCAPE's sparse error feedback and mask buffers is mitigated through CPU offloading and buffer packing, rendering the approach compatible with the limited device memory of current HBM GPUs.

Downstream Task Robustness

Zero-shot evaluation on LAMBADA, SuperGLUE, PIQA, MMLU, ARC, and other benchmarks demonstrates that SCAPE-trained models at 90% and 99% sparsity retain comparable or superior accuracy relative to dense-trained baselines. SCAPE ($d = 0.1$) outperforms or matches the dense baseline on a majority of tasks. This quality preservation is critical for practical adoption in production LLM pipelines.

Implications and Future Directions

SCAPE demonstrates that extreme communication sparsification is feasible for LLM training, provided that sparsity aligns with temporally stable, first-moment statistics (as in AdamS), and communication strategies are architected for hardware and optimizer locality. These findings contradict the common assumption that high sparsity necessarily entails instability or accuracy degradation in large-scale Adam-style training.

Practical implications for systems engineering include:

• Enabling larger batch sizes and models per node by minimizing communication.
• Lowering wall-clock and energy costs, especially for bandwidth-limited environments.
• Increasing the range and scale of feasible LLM experiments for both academic and commercial users with commodity clusters.

Theoretically, SCAPE's mask stability and the observed robustness of AdamS under stochastic error-feedback should motivate further optimizer analysis, as well as adaptation to more advanced preconditioned and block-wise optimizer structures. Integration with quantization, low-rank approaches, or alternative optimizer state compression may compound benefits.

Conclusion

SCAPE provides an effective, system-level solution to the communication bottleneck in large-scale LLM training by exploiting the statistical stability of AdamS's first-moment for extreme sparsification. It achieves substantial empirical speedups (up to 3.26$\times$ per-step and over 40% E2E time reductions) without loss of accuracy—even at 99% sparsity—across diverse LLM architectures and downstream benchmarks. SCAPE’s methodology and results have significant implications for future distributed optimization research and for the practical deployment of ever-larger LLMs (2607.01678).