arXiv CSopen access

SCAPE: Accurate and Efficient LLM Training with Extreme Sparse Communication

SCAPE: Accurate and Efficient LLM Training with Extreme Sparse Communication

I. I NTRODUCTION Large Language Models (LLMs) with ever-increasing sizes have achieved unparalleled performance across many fields, including math reasoning [1], code generation [2], and autonomous laboratory [3]. State-of-the-art LLM training requires a massive amount of graphics processing units (GPUs) to achieve faster training and accommodate enormous model parameters and optimizer states. The legacy data parallel (DP) strategy replicates the model and optimizer states (e.g., firstand second-momentum in AdamW [4]) across GPUs and distributes a mini-batch of data among them. Modern sharded data parallel strategies, such as ZeRO [5], FSDP [6], and Megatron-LM [7], reduce spatial redundancy by partitioning the model and optimizer states across GPUs. In either case, communication (i.e., all-reduce and all-gather) is the scaling bottleneck given the dependency on model sizes and the scale, as shown in Figure 1.

100.0%

llama-500M llama-1.8B

Computation Communication llama-500M efficiency llama-1.8B efficiency

12500 100.0% 89.9% 80.8% 75.2% 10000 7500 2500 0

60

37.8% 35.1% 21.5%

5000 4

8

16

100 80

58.6% 55.8%

32

Number of GPUs

64

Strong Scaling Efficiency (%)

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 communicationreduction 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× speedup per step compared to dense AdamS. Index Terms—large language model training, distributed training, sparsified communication

Time (ms)

arXiv:2607.01678v1 [cs.LG] 2 Jul 2026

Mingkai Zheng, Junlin Chen, Haotian Xie, and Zhao Zhang Rutgers University {mz687, junlin.chen110, haotian.xie, zhao.zhang}@rutgers.edu

40 20

Fig. 1: Scaling bottleneck for pre-training Llama-500M (sequence length 4K) and Llama-1.8B (sequence length 2K) using Megatron-LM with distributed optimizer on NVIDIA GH200 GPUs. Llama-1.8B uses a sequence length of 2K due to the limited memory of GH200.

Lowering communication volume is an effective way to reduce communication overhead. Researchers have explored several approaches. DGC [8], DeMo [9], EDGC [10], oktopk [11], and Radius [12] communicates top-k gradients with error feedback to preserve model performance. A second line of research exploits quantization (low-bit representation) techniques. QSDP [13], ZeRO++ [14], and SDP4bit [15] reduce the number of bits for model parameters and gradients in the sharded data parallel strategy. Existing top-k methods are limited in applicability to modern optimizers or require low sparsity due to constraints on model performance. Legacy top-k methods, such as DGC [8] and DeMo [9], are designed to sparsify the momentum in the SGD [16] optimizer, which is rarely used in today’s LLM training. Radius [12] proposes an error-correction mechanism for the AdamW [4] optimizer, and can achieve 40% sparsity without degrading downstream task performance. The communication volume reduction capability of quantization methods, including QSDP [13], ZeRO++ [14], and SDP4bit [15], is inherently limited by the number of bits required to represent each gradient. Furthermore, without NCCL backend support, these quantization-based methods rely on expensive all-to-all communication, imposing high memory and communication overhead. In this work, we ask the following research question: Can we reduce communication volume to the extreme without

harming the model performance? To this end, we propose SCAPE, a communication-efficient distributed optimizer that sparsifies the optimizer states rather than gradients. SCAPE is built upon AdamS [17] and is inspired with two important insights: • Compared to AdamW, AdamS exhibits substantially smaller residual growth under high sparsity in LLM pretraining, indicating improved robustness to stale errorfeedback dynamics. • The top-k structure of the first-moment is temporally stable across two adjacent steps, suggesting that delayed mask reuse is practical. Based on these observations, SCAPE introduces a partitioned mask-refresh mechanism aligned with optimizer sharding. SCAPE computes masks from momentum-based statistics rather than raw gradients, and delays their use by one step so that mask synchronization can overlap with computation. In addition, SCAPE reconstructs the quantities needed for second-moment updates using a single synchronized sparse buffer, avoiding a second all-reduce. When sharded data parallel is enabled, with the model and optimizer states distributed across all GPUs, SCAPE compresses the volume using sparsity for both reduce-scatter for gradient synchronization and allgather for reconstructing the model. We implement SCAPE using Megatron-LM [7] and empirically verify its convergence through pre-training GPT345M on the OpenWebText [18] dataset and Llama-500M on the SlimPajama-6B [19] dataset. We run experiments on the TACC Vista supercomputer with evaluation on an extensive suite of downstream tasks. Our experiment results show that SCAPE with 90% and 99% sparsity can reduce the wall-clock time for Llama-500M pretraining on 32 GH200 GPUs by 35.6% and 43.3%, respectively. More importantly, under such high sparsities, SCAPE maintains final training and validation losses close to those of dense AdamS and AdamW, without affecting their performance on downstream benchmarks, such as LAMBADA, SuperGLUE, PIQA, MMLU, and ARC. SCAPE is expected to be effective across GPU clusters with various GPU and interconnect configurations, though the improvements it achieves may vary depending on the hardware. SCAPE leverages NVLink-C2C with 900 GB/s throughput between the CPU and GPU on the GH200 superchips for offloading error feedback buffers. Buffer offloading may introduce additional overhead on other GPU clusters with PCIe connections. However, the impact of buffer offloading is minimal compared to the overall improvement with SCAPE. II. BACKGROUND In this section, we provide background information on distributed training with Megatron-LM [7] with distributed optimizers, top-k sparsification for gradient, and the AdamS [17] optimizer. A. Training with Sharded DP Distributed Optimizer In the conventional distributed data parallel (DDP) training, each worker (e.g., a GPU) in a data parallel (DP) group

Fig. 2: Megatron-LM with sharded data parallel distributed optimizer

duplicates all the model parameters, gradient, and optimizer states, and has a split of a micro-batch of training samples. After all workers finish computing their gradient using the local samples, they synchronize the gradient with a ring allreduce operation before updating the optimizer states and model parameters. Assume there are N workers and the gradient size is W (same as the model size). This communication cost can be decomposed into a reduce-scatter following an all-gather operation, which is modeled as Tring-all-reduce = Treduce-scatter + Tall-gather   W , = 2(N − 1) α + β N

(1)

where α is the fixed latency to initialize a message, and β is the inverse of bandwidth. As the number of workers N increases, 2W β(N − 1)/N is dominated by the gradient size W and the inverse of the bandwidth β. Duplicating the whole model and optimizer states for all workers in a DP group guarantees that the forward and backward computation can finish without communication. However, the optimizer states usually require full precision (FP32) to guarantee the training stability, which becomes a main memory bottleneck in training large language models (LLMs) with increasingly more parameters. Therefore, many frameworks have been proposed to reduce this bottleneck. For instance, Megatron-LM with sharded DP distributed optimizers partitions the optimizer states and the model parameters evenly among all DP workers. Thus, before the forward pass, workers need to perform all-gather to collect the model shards from each other. Note that the model parameters can be prefetched to overlap the communication of all-gather with computation. After the backward pass finishes, each DP worker has the complete gradient. Since the optimizer states have been partitioned evenly, gradients’ all-reduce synchronization can be replaced by reduce-scatter and does not require gathering (i.e., avoid all-gather). This process is illustrated in Figure 2. B. Top-k Gradient Sparsification Top-k sparsification is a technique for reducing Tring-all-reduce by reducing W in Equation 1, given a fixed network bandwidth 1/β. Specifically, before synchronizing the gradient g, we

perform a top-k operation to select the k largest-magnitude gradient values, and then synchronize those top-k values using a compressed buffer. This effectively reduces the communication volume from W to kW . Note that the sparsified g is biased, since the top-k operator drops all but the largest-magnitude entries in gradient g and therefore does not preserve the original gradient in expectation. Consequently, directly applying top-k sparsification at every iteration may introduce accumulated compression error and degrade convergence. To mitigate this issue, prior work usually employs error feedback [16], which accumulates the unsent residual from the current step and adds it back to the gradient before the next top-k selection. Let et denote the residual buffer at step t, then the communicated gradient is g̃t = TopK(gt + et ), and the residual is updated as et+1 = gt + et − g̃t . In this way, the information discarded in one iteration is not permanently lost, but is gradually re-injected into later communications, improving optimization stability while retaining the communication savings of sparsification. Nevertheless, this requires a dedicated error feedback buffer with the same size as g and usually in FP32 for preserving the accuracy. Since each worker’s data samples for computing the gradient gt at step t are distinct, their top-k gradient indices are distinct as well. Thus, directly using all-reduce to compute average of the sparsified gradient leads to wrong results. There are two approaches for tackling this issue. The first one is to perform two all-gather collectives to collect both the top-k mask and the corresponding values, and the second method is to first use all-gather to form a synchronized top-k mask and use all-reduce to synchronize the gradient entries selected by the globally synchronized mask. However, the second approach has the fill-in effect. As described in [11], when the number of workers N is large enough, the synchronized global top-k mask is dense rather than sparse. C. AdamS Optimizer AdamS [17] is a new Adam-like optimizer which, instead of keeping a second-moment vt as an exponential moving average (EMA) of gt2 , uses the first-moment mt itself as the preconditioner:

using the same set of hyperparameters, AdamS is more robust to noise than AdamW and avoids large spikes in its training loss curve. Since mt is the EMA of gradient gt , it is a smoother preconditioner than AdamW’s vt , which is an EMA of gradient squared gt2 . When large gradient occurs in the stable region where gradient norms are usually small, it is more robust to such noise. Second, AdamS achieves model quality comparable to, and in some cases slightly better than, AdamW. In GPT-2 [21] pretraining, AdamS closely mirrors AdamW across model scales and even attains slightly lower validation perplexity. In post-training with GRPO [22] on the Countdown task, AdamS yields score curves that closely align with those of AdamW and occasionally surpass its validation performance. Finally, AdamS eliminates the need to store second moment, reducing optimizer state memory by 50% and leading to practical system benefits such as lower memory consumption in distributed training (e.g., FSDP [6]) and higher throughput in memory-bound large-scale pretraining. III. O BSERVATIONS In this section, we present our observations when pretraining the GPT-345M model. Based on these observations, we provide insights that form the foundations for our system design in Section IV. A. Gradient Distribution We first pre-train GPT-345M with AdamW optimizer on the OpenWebText [18] dataset with top-10% sparsity on the gradient, meaning that we only synchronize the gradient with the top-10% largest magnitudes. In this setup, we switch from a dense all-reduce with no gradient compression (baseline method) to top-10% sparsification at step 10,000. Moreover, the top-k sparsification is performed in a per-layer style. We can see from Figure 3a that the difference between top-10% sparsity and the baseline method rapidly increases and then slowly decreases. We then try AdamS using the same setup, and the result is presented in Figure 3b. We notice that, using AdamS, the gap between the baseline method and the sparsified gradient is much smaller.

mt = β1 mt−1 + (1 − β1 )gt , ⊙2 vt = β2 m⊙2 t−1 + (1 − β2 )gt ,   mt + λwt−1 , wt = wt−1 − η √ vt + ϵ

(2)

where β1 and β2 are the coefficients for computing the running average, wt is the model parameter at step t, η is the learning rate, λ is the weight decay hyper-parameter, and ϵ maintains the numerical stability. Note that the update metric in Equation 2 is simplified for demonstration purposes, and in practice, mt and vt should be multiplied by their bias correction terms 1/(1 − β1t ) and 1/(1 − β2t ). Compared to AdamW [4] (i.e., Adam [20] with weight decay), AdamS has superior performance in terms of training stability, model quality, and memory consumption. First, when

(a) AdamW

(b) AdamS

Fig. 3: Comparison between AdamW and AdamS after switching from dense all-reduce to top-10% sparsified gradient at step 10,000 of pre-training GPT-345M To analyze these phenomena, we plot the norms of the top-k sparsified gradient and the residual buffer, which stores the non-top-k gradient accumulated during training. From the

curves of top-k gradient norm and the residual norm plotted in Figure 3a and Figure 3b, we see that AdamW’s residual norm is much higher than AdamS’s. With a high sparsity rate, such as 10% in this experiment, the staleness effect of error feedback can significantly influence AdamW’s training process. The reason is that the chance of being selected as top-k is so low that many gradient values with very large magnitudes (but still smaller than the top-10% threshold) have to wait until they have accumulated for several steps. When such large stale gradient values are added back to the optimizer states, including the first- and second-moment, they cause the optimizer states to drift from the baseline’s direction and thus lead to a gradually larger gap between the baseline and top-k method. This aligns well with the findings in [23] and [24] (i.e., the loss spike is triggered by suddenly having a large gradient when the training regime enters a region where the square of the gradient and the second-moment are both small). However, for AdamS, the norm of the top-k gradient and the norm of the residual values are small, so when adding them back, it does not lead to a significant shift in the optimization direction. We plot the distribution of AdamW’s and AdamS’s gradient for different layers at two different steps in Figure 4 for GPT-345M model. We also provide the gradient distribution for Llama-500M in Figure 5. From these results, we can see that AdamS’s gradient is more centralized around 0, while AdamW’s gradient distribution spreads out more evenly and is flatter. This suggests that compared to AdamW, AdamS is more suitable for adapting to top-k gradient sparsification.

(a) Decoder 0 - Self attention (b) Decoder 10 - MLP dmodel to dense at step 1,000 4dmodel at step 1,000

(c) Decoder 0 - Self attention (d) Decoder 2 - QKV projection dense at step 50,000 weight at step 50,000

(a) Decoder 0 - gate and up (b) Decoder 8 - Self attention project in SwiGLU at step 1,000 dense at step 1,000

(c) Decoder 0 - gate and up (d) Word embedding at step project in SwiGLU at step 50,000 50,000

Fig. 5: Gradient distribution of different layers in Llama-500M at different steps

Nevertheless, in our setting, we observe that this property also holds when top-k masks are constructed from AdamS’s first-moment rather than raw gradients, for both GPT-345M and Llama-500M. In Figure 6, we show the overlap between adjacent-step top-10% masks selected according to the magnitude of AdamS’s first-moment across different layers and training stages. As shown in Figure 6a, many layers exhibit high mask overlap at the beginning of training. Although the overlap gradually decreases as training proceeds, it remains substantial throughout pre-training, indicating that the top-k mask from step t is a useful heuristic for step t + 1. For Llama-500M, the overlap curves in Figure 6c and Figure 6d are noisier, but the average overlap for each layer remains high across training stages.

(a) GPT-345M from step 1

(b) GPT-345M from step 10,000

(c) Llama-500M from step 1

(d) Llama-500M from step 10,000

Fig. 4: Gradient distribution of different layers in GPT-345M at different steps B. Temporal Stability of first-moment-Derived Top-k Masks Prior work has shown that sparse support patterns can exhibit temporal correlation across adjacent training steps. For example, the temporal stability of AdamW’s top-1% gradient has been studied in [12], and [10] reports strong gradient correlation during the early stage of LLM pre-training.

Fig. 6: The overlap rate of AdamS’s first-moment-derived top10% mask between two adjacent steps for different stages of pre-training GPT-345M and Llama-500M

IV. S YSTEM D ESIGN Motivated by the observations in Section III, we design our system in Algorithm 1, which is mainly composed of three important components: 1) new mask synchronization, 2) topk mask computation, and 3) optimizer state update. Algorithm 1 SCAPE Require: model parameters θ, dataset D, loss function f Require: step t, worker rank n, total number of workers N Require: density ρ, weight decay λ, learning rate η Require: number of layers in model L 1: M0 ← 1 ▷ Initialize sparse mask 2: en ← 0 ▷ Initialize EF buffer 0 3: for t = 1, 2, . . . do 4: bnt ← 0 ▷ Re-initialize payload buffer 5: gtn ← ∇f (Dtn , θt ) 6: m̃nt ← β1 mt−1 + (1 − β1 )gtn + ent−1 7: ent [¬Mt−1 ] ← m̃t [¬Mt−1 ] 8: ent [Mt−1 ] ← 0 9: bnt [Mt−1 ] ← m̃nt [Mt−1 ]P N 10: All-reduce b̄nt : b̄t ← N1 n=1 bnt ▷ Check Figure 7 11: Mt ← refresh mask(ρ, |m̃nt |) 12: mt ← b̄t ▷ First-moment’s nontopk all 0 13: b̄t [Mt−1 ] ← (b̄t [Mt−1 ] − β1 mt−1 [Mt−1 ])/(1 − β1 ) 14: clip grad(b̄t ) 15: vt ← β2 m⊙2 + (1 − β2 )b̄⊙2 t t−1 √ 16: ut ← mt /( vt + ϵ) ▷ Compute update metrics 17: θt+1 ← θt − η(ut + λθt ) 18: end for

B. Top-k Mask Computation The top-k mask can be computed from the gradient and moment. SCAPE selects the first-moment to construct the topk mask. The reason is that, as the gradient contains noise inherited from the data, the top-k values generated from the gradient of the current step can have a low overlap rate with the real top-k gradient’s indices in the next step. Therefore, we use the first-moment to generate the top-k mask, as it maintains a running average of the gradient, which is more stable and has much lower noise. Specifically, we use |m̃nt | = |β1 mt−1 +(1− β1 )gtn + ent−1 | to construct the top-k masks. Because gtn and ent−1 are not synchronized and contain local information, each worker’s top-k mask guarantees that its local largest values in |m̃nt | are immediately selected and then used to update the optimizer states and model parameters. C. Optimizer State Update In Algorithm 1, SCAPE transmits a sparsified buffer bnt , which has the values in the temporal, semi-updated firstmoment buffer m̃nt on positions selected by the top-k mask Mt−1 . For DGC [8] and DeMo [9], transmitting the sparsified first-moment is sufficient to update the optimizer state, because they use SGD optimizer [25]. However, from Equation 2, to update the second-moment, we need both the synchronized first-moment and the gradient. A naive implementation would perform two all-reduce operations to synchronize both. Nevertheless, we noticed that one all-reduce suffices. Given that mt is guaranteed to be synchronized, we can compute the globally averaged gradient and residual as ḡt + ēt−1 = (b̄t − β1 mt−1 )/(1 − β1 ).

(3)

Then, we use ḡt + ēt−1 as the synchronized top-k to update the second-moment vt and the model parameters wt . A. Refreshed Mask Synchronization

D. Optimizations for Distributed Optimizers

We propose a new mask refresh strategy illustrated in Figure 7. This method has two important differences from the existing work. First, instead of having all workers compute their own top-k masks for all layers, each worker only computes a portion. This modification aligns with the design of Megatron-LM’s distributed optimizer, where the optimizer states are partitioned evenly onto all workers, and at each step, each worker updates its own portion only. Moreover, when the distributed optimizer is not used, this design reduces each worker’s workload of performing the top-k operations by N (i.e., the total number of workers in a DP group). This design is also immune to the fill-in effect mentioned in [11]: as the number of workers grows, the gradient ‘sparsified’ by the collected top-k indices from all workers is nearly dense. Since each worker in our design treats the top-k masks for layers computed by its peers as all zeros, this avoids the fill-in effect. Second, the top-k masks computed from step t are later used at step t + 1. This delayed use of the top-k masks provides an opportunity to use asynchronous communication to hide the synchronization communication in heavy computations, such as the forward- and backward-pass.

We can see from Algorithm 1 that since b̄t only has nonzero values for indices selected by top-k mask Mt−1 , we have mt [¬Mt−1 ] = b̄t [¬Mt−1 ] = 0.

(4)

Therefore, when computing the update metrics, we have p ut [¬Mt−1 ] = mt [¬Mt−1 ]/( vt [¬Mt−1 ] + ϵ) = 0. (5) Hence, θt+1 [¬Mt−1 ] is updated as θt+1 [¬Mt−1 ] = θt [¬Mt−1 ] − η(ut [¬Mt−1 ] + λθt [¬Mt−1 ]) = θt [¬Mt−1 ] − η(0 + λθt [¬Mt−1 ]) = (1 − ηλ)θt [¬Mt−1 ]. (6) This suggests that if the model parameters θ is partitioned among all workers, we can avoid the expensive all-gather operation for θt+1 [¬Mt−1 ] by computing it from each worker’s local FP32 copy of the complete model (see Figure 2). To reduce the memory overhead of keeping full model parameters on every worker, we offload the full-parameter replica to CPU memory. During parameter updates, we use a double-buffered pipeline that overlaps asynchronous CPU-to-GPU prefetch and

Fig. 7: Communication in refreshing top-k mask. Each work computes a sharded top-k mask and then uses all-gather to construct the full mask. Since the usage of top-k mask is delayed by one step, asynchronously all-gather can be hidden by expensive backward computation.

sparse payload communication with local non-topk weightdecay updates, then writes back the updated buffer and offloads it to the host memory for the next iteration.

tion, because the LayerNorm, RMSNorm, and bias layers are usually not sparsified as they are less stable than the weights, leading to uneven buffers for reduce-scatter on workers.

Algorithm 2 Optimize Megatron-LM distributed optimizer with SCAPE Require: model parameter full FP32 copy w, parameter for forward and backward computation ŵ, parameter for communication w̃, local parameter w′ updated by partitioned optimizer states, top-k mask M , learning rate η, weight decay λ, worker rank n, total number of workers N , step t, dataset D 1: function F ORWARD PASS ▷ Compress comm volume 2: w̃tn ← wt′ [Mt−1 ] 3: ŵt [Mt−1 ] ← AllGather(w̃tn ) 4: ŵt [¬Mt−1 ] ← (1 − ηλ)ŵt [¬Mt−1 ] ▷ Use FP32 wt to compute untransmitted params 5: ForwardPass(Dtn , ŵt ) 6: end function 7: function BACKWARD PASS 8: Start AsyncMaskRefresh(Mt ) ▷ Hide mask sync with BW computation 9: gtn ←BackwardPass(Dtn , ŵt ) 10: end function 11: ReduceScatter(gtn ) n 12: bn t ←ComputeUpdatedMoments(gt ) n 13: b̄t ←AllReduce(bt [Mt−1 ]) 14: UpdateLocalOptimizerStatesAndParams(b̄t )

E. Optimization for Memory consumption SCAPE holds one additional residual buffer for error feedback, which has the same size as the model and uses FP32 for accuracy. It also needs two buffers for the masks of the current and next steps. Thus, the total memory overhead can be expressed as Memoverhead = Memcur mask + Memnext mask + Memresidual = W + W + 4W = 6W bytes, where W is the number of parameters, and the dtype for top-k masks and residual are INT8 and FP32. To solve additional memory bottlenecks, we used CPU offloading with double buffering for residual and full parameters, minimizing the influence on both per-step time and memory. Furthermore, each value in a mask is represented as one bit, and we pack eight of them into one byte. Note that the smallest dtype supported by PyTorch operations is INT8, and this single-bit format is not supported. Therefore, we store both masks in packed single-bit format and use double buffers to convert when needed. We also use Triton kernels for packing and unpacking operations to reduce their computation overhead. After applying these memory optimization techniques, the memory overhead is 2 × W/8 + 2P + 8P = W/4 + 10P bytes, where P represents the number of parameters in the largest layer, 2P represents the memory usage for double buffers for masks in INT8, and 8P represents the double buffers for the residual offloading.

Algorithm 2 shows the optimized Megatron-LM with a distributed optimizer. If the model parameters are partitioned and need to be gathered across all workers before the forward- and backward-pass, we can reduce such communication volume by the inverse of the sparsity rate 1/k. Additionally, to hide the communication cost of synchronizing and updating Mt , we launch the asynchronous all-gather for Mt after the backwardpass’s all-gather finishes, guaranteeing no communications interference. After the backward pass finishes, when workers synchronize their local temporal, semi-updated first-moment bnt , we replace the reduce-scatter with the all-reduce opera-

V. E XPERIMENT R ESULTS In this section, we evaluate SCAPE from three complementary perspectives. We first present the pre-training, downstream, and end-to-end wall-clock results for Llama-500M, our primary modern LLM workload for demonstrating practical systems benefit. We then report GPT-345M results to show that SCAPE generalizes beyond a single architecture family and training configuration. Finally, we analyze per-step time and strong-scaling efficiency for Llama-500M and Llama1.8B under Megatron-LM’s distributed optimizer, including the SCAPE-specific optimizations, from 4 to 64 GPUs.

TABLE I: Final training and validation loss and end-to-end (E2E) wall-clock time of pre-training Llama-500M M ETHOD

E2E T IME (DAYS )

T RAIN LOSS

VAL . LOSS

A DAM W ( DENSE all-reduce) A DAM S ( DENSE all-reduce) SCAPE (d = 0.1) SCAPE (d = 0.01)

2.47 (1×) 2.47 (1×) 1.59 (1.55×) 1.40 (1.76×)

2.10 2.11 2.13 2.17

2.30 2.32 2.30 2.31

(a) Training loss

(b) Validation loss

Fig. 8: Pre-training loss curves for Llama-500M

A. Experiment Setup We evaluate SCAPE by pre-training GPT-345M and Llama500M on 32 NVIDIA GH200 GPUs of the Vista supercomputer [26] at the Texas Advanced Computing Center (TACC). Each Vista node consists of a Grace-Hopper architecture with one GH200 GPU, 96 GB of HBM3 memory, and an NVLinkC2C interconnect between the Grace CPU and Hopper GPU. The nodes are connected through a 400 Gbps NVIDIA NDR InfiniBand network. To ensure a consistent and reproducible software environment, all experiments were conducted using the NVIDIA NGC PyTorch container (v26.01). We evaluate the pre-trained models using zero-shot downstream benchmarks using lm-evaluation-harness [27], including ARC (Easy and Challenge) [28], LAMBADA [29], HellaSwag [30], MMLU [31], PIQA [32], WinoGrande [33], OpenBookQA [34], and SuperGLUE [35]. B. Llama-500M Pre-training 1) Model Architecture and Hyperparameters: We use the same model architecture of H2O-Danube3-500M [36] to define our Llama-500M model and pre-trained it on SlimPajama6B, a subdataset sampled from SlimPajama [19]. We used the same tokenizer as Llama-7B [37] from HuggingFace. We pretrain Llama-500M for 100,000 steps and use a global batch size of 1024 and a sequence length of 4096 for each step. The hyperparameters for the optimizers are: peak learning rate η = 3 × 10−4 , the minimum learning rate ηmin = 3 × 10−5 , cosine learning rate decay, learning rate warmup for 2,000 steps, (β1 , β2 ) = (0.9, 0.95), RMSNorm ϵ = 10−5 , weight decay λ = 0.1, and the gradient clip of 1. We use BF16 for parameters and FP32 for the gradient. For SCAPE, we use sparsity warmup to exponentially decrease density d from 1 to 0.1 or 0.01, and we do not compress RMSNorm layers.

2) Pre-training Results: The pre-training loss curves for Llama-500M using AdamW, AdamS, SCAPE (d = 0.1), and SCAPE (d = 0.01) are shown in Figure 8. Their final endto-end (E2E) wall-clock times, training losses, and validation losses are summarized in Table I. Since AdamW is still the de facto optimizer for pre-training LLMs, we include it as a reference baseline. Compared with dense AdamS, SCAPE (d = 0.1) reduces wall-clock time from 2.47 to 1.59 days (1.55× speedup, 35.6% reduction) and lowers validation loss from 2.32 to 2.30. SCAPE (d = 0.01) further reduces time to 1.40 days (1.76× speedup, 43.3% reduction) with validation loss 2.31, which is still below AdamS. These results demonstrate SCAPE as a quality-preserving communication-efficient method for large-scale pre-training: it delivers substantial E2E wall-clock speedups while maintaining, and in validation loss slightly improving, model quality similar to dense AdamS. 3) Downstream Task Evaluation: Table II shows that SCAPE largely preserves downstream task performance for the pre-trained Llama-500M despite using aggressive communication sparsity. Specifically, SCAPE (d = 0.1) outperforms dense AdamW on 5 of 16 tasks and dense AdamS on 9 of 16 tasks. In particular, SCAPE (d = 0.1) surpasses both dense baselines on RTE and ReCoRD, while also improving over AdamW on ARC-E and OpenBookQA, and over AdamS on ARC-C, ARC-E, LAMBADA, OpenBookQA, PIQA, WSC, and MultiRC. Even at d = 0.01, SCAPE remains competitive, exceeding AdamW on 4 of 16 tasks and AdamS on 7 of 16 tasks, including 3 SuperGLUE improvements over AdamW and 4 SuperGLUE improvements over AdamS. These results indicate that SCAPE retains strong downstream generalization for Llama-500M under 90% and 99% sparsity, with d = 0.1 offering the best balance between compression and quality preservation. C. GPT-345M Pre-training 1) Model Architecture and Hyperparameters: We pre-train GPT-345M model using the same model architecture definition described in [21] on OpenWebText dataset [18] for 100,000 steps. We set the global batch size to 512 and the sequence length to 1024. As for the hyperparameters of AdamW, AdamS, and SCAPE, we set (β1 , β2 ) = (0.9, 0.999), the peak learning rate η = 1.5 × 10−4 , the minimum learning rate ηmin = 10−5 , the learning rate scheduler to cosine, the learning rate warmup steps to 5,000 (5% of the pre-training steps), the weight decay λ = 0.01, and gradient clip to 1. Similar to Llama-500M training, mixed-precision training is used: the model parameters are in BF16, and the gradient is accumulated in FP32. Similarly, we use density warmup to gradually decrease d from 1 to 0.1 or 0.01. The LayerNorm and bias layers in GPT-345M were not compressed. 2) Pre-training Results: The training and validation loss curves for pre-training GPT-345M model with AdamW, AdamS, SCAPE (d = 0.1), and SCAPE (d = 0.01) are shown in Figure 9, and the final training and validation loss are listed in Table III.

TABLE II: Zero-shot evaluation results for Llama-500M pre-trained using different approaches M ETHOD

ARC-C (ACC)

ARC-E (ACC)

H ELLA S WAG (ACC)

LAMBADA (ACC)

MMLU (ACC)

O PEN B OOK QA (ACC)

PIQA (ACC)

W INO G RANDE (ACC)

A DAM W ( DENSE all-reduce) A DAM S ( DENSE all-reduce) SCAPE (d = 0.1) SCAPE (d = 0.01)

22.53 20.31 20.90 21.93

48.86 46.97 48.95 47.47

34.70 34.57 34.38 33.93

38.60 36.66 38.39 36.64

24.08 25.77 23.28 23.28

18.80 18.80 19.40 18.60

67.30 65.94 66.76 66.32

50.67 53.51 52.33 51.07

S UPER GLUE M ETHOD

B OOL Q (ACC)

WIC (ACC)

RTE (ACC)

CB (ACC)

R E C O RD (EM)

WSC (ACC)

M ULTI RC (ACC)

COPA (ACC)

AVG .

A DAM W ( DENSE all-reduce) A DAM S ( DENSE all-reduce) SCAPE (d = 0.1) SCAPE (d = 0.01)

57.65 57.55 51.90 57.92

49.69 50.63 49.06 50.00

53.43 53.07 57.40 52.71

35.71 41.07 28.57 32.14

69.81 69.54 70.00 69.66

41.35 34.62 39.42 36.54

57.03 49.86 53.11 56.48

65.00 68.00 68.00 68.00

53.71 53.04 52.18 52.93

AdamS AdamS w/ dist-optm SCAPE (d=0.1) SCAPE (d=0.1) w/ dist-optm

SCAPE (d=0.1) w/ dist-optm & CPUOffload SCAPE (d=0.01) SCAPE (d=0.01) w/ dist-optm SCAPE (d=0.01) w/ dist-optm & CPUOffload

Time per Step (ms)

8000 6000 4000 2000

(a) Training loss

0

(b) Validation loss

4

8

Fig. 9: Pre-training loss curves for GPT-345M

16

Number of GPUs

32

64

32

64

(a) Llama-500M

M ETHOD

T RAIN LOSS

VAL . LOSS

A DAM W ( DENSE all-reduce) A DAM S ( DENSE all-reduce) SCAPE (d = 0.1) SCAPE (d = 0.01)

2.80 2.77 2.77 2.81

2.76 2.73 2.73 2.76

Surprisingly, given the same token budget, when using d = 0.1 (90% sparsity), SCAPE achieves lower training and validation loss than the dense AdamS. When the sparsity rate is increased to 99% (i.e., d = 0.01), the final training and validation losses of SCAPE differ only slightly from those of AdamS, with differences below 0.04. Additionally, the difference between AdamW and SCAPE with d = 0.01 is so small that it can be considered negligible. 3) Downstream Task Evaluation: Table IV presents the zero-shot evaluation results of GPT-345M on the same downstream benchmark suite used for Llama-500M. Using dense AdamW as the reference baseline, SCAPE achieves strong task-level gains under aggressive sparsity. In particular, SCAPE (d = 0.1) improves accuracy on 9 of 16 tasks, including ARC-C, HellaSwag, LAMBADA, MMLU, OpenBookQA, BoolQ, WiC, CB, and ReCoRD. Even when the sparsity is increased to d = 0.01, SCAPE still outperforms AdamW on 8 of 16 tasks. These results indicate that GPT-345M is highly tolerant to SCAPE’s sparse synchronization, especially at d

12500

Time per Step (ms)

TABLE III: Final training and validation loss of pre-training GPT-345M

10000 7500 5000 2500 0

4

8

16

Number of GPUs

(b) Llama-1.8B

Fig. 10: Per-step time comparison between different methods used for training llama-500M (sequence length 4K) and Llama-1.8B (sequence length 2K)

= 0.1, where sparse training delivers improvements on a majority of the reported zero-shot evaluations while maintaining competitive overall downstream performance. D. Per-step Time Analysis We profile the time per iteration for pre-training Llama500M (global batch 1024 and sequence length 4K) and Llama1.8B (global batch size 1024 and sequence length 2K) with eight configurations. These configurations include AdamS, AdamS with distributed optimizer, SCAPE (d = 0.1), SCAPE (d = 0.1) with distributed optimizer, SCAPE (d = 0.1) with distributed optimizer and CPU offloading, SCAPE (d = 0.01), SCAPE (d = 0.01) with distributed optimizer, and SCAPE (d

TABLE IV: Zero-shot evaluation results for GPT-345M pre-trained using different methods M ETHOD

ARC-C (ACC)

ARC-E (ACC)

H ELLA S WAG (ACC)

LAMBADA (ACC)

MMLU (ACC)

O PEN B OOK QA (ACC)

PIQA (ACC)

W INO G RANDE (ACC)

A DAM W ( DENSE all-reduce) A DAM S ( DENSE all-reduce) SCAPE (d = 0.1) SCAPE (d = 0.01)

20.39 19.88 20.99 19.71

48.82 47.77 48.82 46.93

30.63 31.45 31.28 31.06

39.57 39.51 39.71 39.06

22.88 22.94 22.91 22.93

17.00 18.00 17.80 18.60

64.91 65.18 64.80 64.42

51.38 51.54 51.38 50.67

S UPER GLUE WIC (ACC)

RTE (ACC)

CB (ACC)

R E C O RD (EM)

WSC (ACC)

M ULTI RC (ACC)

COPA (ACC)

AVG .

A DAM W ( DENSE all-reduce) A DAM S ( DENSE all-reduce) SCAPE (d = 0.1) SCAPE (d = 0.01)

53.12 55.72 59.36 56.42

49.53 50.00 51.72 49.84

52.71 53.43 52.35 50.18

28.57 50.00 32.14 35.71

76.24 78.06 77.38 76.70

59.62 36.54 41.35 36.54

57.24 54.02 55.18 54.83

74.00 70.00 73.00 69.00

56.38 55.97 55.31 53.65

AdamS AdamS w/ dist-optm SCAPE (d=0.1) SCAPE (d=0.1) w/ dist-optm

SCAPE (d=0.1) w/ dist-optm & CPUOffload SCAPE (d=0.01) SCAPE (d=0.01) w/ dist-optm SCAPE (d=0.01) w/ dist-optm & CPUOffload

100 80 60 40 20 0 4

8

16

Number of GPUs

32

64

32

64

(a) Llama-500M

Strong Scaling Efficiency (%)

= 0.01) with distributed optimizer and CPU offloading, where we define the model architecture of Llama-1.8B by following the definition of H2O-Danube3-1.8B [36]. The profiling results ranging from 4 to 64 GH200 GPUs are plotted in Figure 10. Note that SCAPE (d = 0.1) and SCAPE (d = 0.01) refer to the DDP setting, and each worker has the full replicated model and optimizer states in GPU memory. As we can see from the profiling results, SCAPE can efficiently reduce the time per iteration for pre-training both Llama-500M and Llama-1.8B under different training configurations. For Llama-500M training, the per-step time reduction using SCAPE (d = 0.1) and SCAPE (d = 0.01) is less pronounced than Llama-1.8B. For Llama-500M on 64 GPUs, SCAPE with d = 0.1 and d = 0.01 reduces the per-step time from 1753.66 ms under dense AdamS to 1077.40 ms and 908.19 ms, corresponding to 1.63× and 1.93× speedup, respectively. For Llama-1.8B, SCAPE reduces the per-step time from dense AdamS’s 4804.44 ms to 1983.08 ms and 1473.45 ms, yielding 2.42× and 3.26× speedup, respectively. The reason is that with such a high number of tokens per step and a smaller model size than Llama-1.8B, the per-step training time for Llama-500M is dominated by computation rather than communication, which can be found from Figure 1. We notice that when the distributed optimizer is used, AdamS’s per-step time is reduced, while SCAPE’s is slightly increased for Llama-500M. The reason is that by sharding the model and optimizer states to all workers for AdamS, each worker’s computation time for updating the model and states is reduced. Moreover, the composition of reduce-scatter for synchronizing gradient and all-gather for gathering the model shards in the distributed optimizer setting has the same cost as Tring-all-reduce (see Equation 1). Therefore, using distributed optimizer reduces the computation for each worker without introducing additional communication overhead. However, for SCAPE with distributed optimizer, since each worker still needs to update the full model, the computation is not reduced. Furthermore, the use of all-reduce introduces additional communication overhead compared to reduce-scatter (see Equation 1), although its communication volume is significantly reduced. Hence, using SCAPE with distributed optimizer has

Strong Scaling Efficiency (%)

M ETHOD

B OOL Q (ACC)

100 80 60 40 20 0 4

8

16

Number of GPUs (b) Llama-1.8B

Fig. 11: Strong scaling efficiency for training Llama-500M (sequence length 4K) and Llama-1.8B (sequence length 2K) with different methods

slightly higher per-step time than SCAPE when communication does not dominate the per-step time, such as training Llama-500M on four GPUs. From Figure 10, we can see that adding CPU offload to SCAPE with distributed optimizer does not add overhead for transmitting between CPU and GPU for each step. Thanks to the high bandwidth provided by GH200’s NVLink-C2C of 900 GB/s, our double-buffering scheme for CPU offload only has minimal impact on the per-step time. Nevertheless, this could add extra overhead for systems that use PCIe for connecting CPU and GPU, which has a lower bandwidth than NVLink-C2C.

76,344 74,760 77,722 76,620

71,812 70,914

Reference (AdamS) SCAPE (no res. offload) 80000 60000 60,024

60000

68,290

SCAPE (res. offload) Dist-optm enabled 75,832

69,474

61,930

47,980

40000

40000

20000

20000

(a) Llama-500M

es.

DO

+r

E+

DO

AP

E+

SC

SC

AP

E

res .

AP

E+

SC

AP SC

am S am S+ DO Ad

es. +r

DO E+

DO

AP

E+

SC

SC

AP

E AP

E+

SC

AP SC

Ad

res .

0

Ad am S am S+ DO

0

Ad

Memory Usage (MB)

80000

SCAPE (res. offload) Dist-optm enabled

Memory Usage (MB)

Reference (AdamS) SCAPE (no res. offload)

(b) Llama-1.8B

Fig. 12: Memory usage for training Llama-500M (sequence length 4K) and Llama-1.8B (sequence length 2K) E. Strong Scaling Efficiency The strong scaling efficiency results for training Llama500M and Llama-1.8B using all aforementioned methods on 4, 8, 16, 32, and 64 GH200 GPUs are plotted in Figure 11. Using the four-GPU setup as the baseline, the strong scaling efficiency improves with SCAPE for both models: for Llama500M at 64 GPUs, efficiency is improved from AdamS’s 30.98% to 46.11% for SCAPE (d = 0.1) and 54.56% for SCAPE (d = 0.01), and up to 63.69% when using SCAPE (d = 0.01) with distributed optimizer and CPU offloading; for Llama-1.8B on 64 GPUs which has communication dominating the per-step time (see Figure 1), the same pattern is more obvious, improving from 17.58% for AdamS and 21.48% for AdamS with distributed optimizer to 37.15% for SCAPE (d = 0.1), 47.83% for SCAPE (d = 0.01), and 55.04% for SCAPE (d = 0.01) with distributed optimizer and CPU offloading. F. Memory Consumption We profile the memory consumption for training Llama500M and Llama-1.8B under different settings, and the results are shown in Figure 12. Since the memory overhead of residual buffer and full model parameters is determined by the number of model parameters and does not depend on SCAPE’s d or the number of workers, only the memory usage for d = 0.1 on 4 GPUs is reported. We can see that with residual offloading enabled, the memory usage for training Llama-1.8B is reduced from 68,290 MB to 61,930 MB and from 75,832 MB to 69,474 MB before and after distributed optimizer is enabled. For Llama-500M training, because its model size is small and most of GPU memory is consumed by long-sequence activations, using residual offloading does not significantly reduce its memory usage. VI. R ELATED W ORK Many research works have been proposed to achieve communication-efficient pre-training through compression on gradient and model weights. Specifically, they can be classified into three main categories: gradient sparsification with error feedback, low-rank approximation, and quantization. For gradient sparsification, DGC [8] proposes a momentumbased sparsification method to reduce the communication volume of synchronizing gradient in DDP. DeMo [9] proposes a decoupled momentum optimization method which

first performs DCT to orthonormalize the momentum and then uses top-k sparsification with error feedback to reduce communication volume. Ok-Topk [11] proposes a novel sparse all-reduce algorithm to achieve near asymptotically optimal communication volume. EDGC [38] proposes an entropydriven adaptive gradient sparsification framework to dynamically specify the sparsity rate for each model layer. Radius [12] exploits the temporal stability of the top-k gradient values’ indices to amortize the computation cost of top-k operations, and thus achieves lower computation overhead in gradient sparsification and higher throughput. For low-rank compression, ATOMO [39] proposes to use singular value decomposition (SVD) to express gradients as atomic components and transmits only a subset of them to reduce communication overhead. PowerSGD [40] proposes a low-rank approximation of the gradient by representing it with two much smaller factor matrices. Optimus-CC [41] proposes a novel framework, combining 3D parallel training with PowerSGD. For quantization-based methods, QSDP [13] proposes to extend FSDP [6] with both model weight and gradient quantization to reduce the communication cost with convergence guaranteed. ZeRO++ [14] improves the communication efficiency of ZeRO [5] through the combination of block-wise quantized all-gather, communication-aware data remapping, and quantized gradient averaging built on all-to-all communication SDP4Bit [15] proposes to use 4-bit quantization to reduce both weight and gradient communication via quantizing the weight differences and a two-level smooth quantization scheme for gradients, while also introducing runtime optimizations to mitigate quantization overhead. VII. C ONCLUSION AND F UTURE W ORKS We present SCAPE, a communication-efficient distributed optimizer for LLM pre-training implemented based on Megatron-LM that achieves aggressive sparsification while preserving model quality on downstream evaluation tasks. Instead of sparsifying and transmitting the raw gradients, SCAPE leverages the stability of AdamS’s first-moment to construct top-k masks for sparse communication, aligns the top-k mask generation with optimizer sharding, and applies the top-k masks with a one-step delay to overlap the mask synchronization with computation. SCAPE also reconstructs the quantities required for second-moment updates from a single synchronized sparse buffer, thereby avoiding an additional collective. Through extensive evaluations, we show that SCAPE preserves the training loss, the validation loss, and the scores of the downstream task under both 90% and 99% sparsity. For Llama-500M, SCAPE reduces end-to-end pretraining wall-clock time by up to 43.3% relative to dense baselines while maintaining comparable model quality. For Llama1.8B, SCAPE achieves 3.26× speedup on 64 GPUs compared to dense AdamS. In future work, we will explore extending SCAPE’s moment-based sparsification to other optimizers, including Muon, which is gaining increasing popularity in production LLM training.

R EFERENCES [1] J. Ahn, R. Verma, R. Lou, D. Liu, R. Zhang, and W. Yin, “Large language models for mathematical reasoning: Progresses and challenges,” 2024. [Online]. Available: https://arxiv.org/abs/2402.00157 [2] M. Chen, J. Tworek, H. Jun, Q. Yuan, H. P. D. O. Pinto, J. Kaplan, H. Edwards, Y. Burda, N. Joseph, G. Brockman et al., “Evaluating large language models trained on code,” 2021. [Online]. Available: https://arxiv.org/abs/2107.03374 [3] N. J. Szymanski, B. Rendy, Y. Fei, R. E. Kumar, T. He, D. Milsted, M. J. McDermott, M. Gallant, E. D. Cubuk, A. Merchant, H. Kim, A. Jain, C. J. Bartel, K. Persson, Y. Zeng, and G. Ceder, “An autonomous laboratory for the accelerated synthesis of inorganic materials,” Nature, vol. 624, no. 7990, pp. 86–91, 2023. [Online]. Available: https://doi.org/10.1038/s41586-023-06734-w [4] I. Loshchilov and F. Hutter, “Decoupled Weight Decay Regularization,” 2019. [Online]. Available: https://arxiv.org/abs/1711.05101 [5] S. Rajbhandari, J. Rasley, O. Ruwase, and Y. He, “ZeRO: Memory Optimizations Toward Training Trillion Parameter Models,” 2020. [Online]. Available: https://arxiv.org/abs/1910.02054 [6] Y. Zhao, A. Gu, R. Varma, L. Luo, C.-C. Huang, M. Xu, L. Wright, H. Shojanazeri, M. Ott, S. Shleifer, A. Desmaison, C. Balioglu, P. Damania, B. Nguyen, G. Chauhan, Y. Hao, A. Mathews, and S. Li, “PyTorch FSDP: Experiences on Scaling Fully Sharded Data Parallel,” 2023. [Online]. Available: https://arxiv.org/abs/2304.11277 [7] NVIDIA, “Megatron-lm,” 2026. [Online]. Available: https://github.com/ NVIDIA/Megatron-LM [8] Y. Lin, S. Han, H. Mao, Y. Wang, and W. J. Dally, “Deep Gradient Compression: Reducing the Communication Bandwidth for Distributed Training,” 2020. [Online]. Available: https://arxiv.org/abs/1712.01887 [9] B. Peng, L. Chen, B. Su, J. Quesnelle, D. P. Kingma, and Q. Liu, “DeMo: Decoupled Momentum Optimization,” 2026. [Online]. Available: https://arxiv.org/abs/2411.19870 [10] Q. Yi, J. Duan, H. Hu, Q. Hua, H. Zhao, S. Qian, D. Yang, J. Cao, J. Tang, Y. Yu, C. Liao, K. Wang, and L. Zhang, “EDGC: Entropy-driven Dynamic Gradient Compression for Efficient LLM Training,” 2025. [Online]. Available: https://arxiv.org/abs/2511.10333 [11] S. Li and T. Hoefler, “Near-optimal sparse allreduce for distributed deep learning,” in Proceedings of the 27th ACM SIGPLAN Symposium on Principles and Practice of Parallel Programming. ACM, mar 2022. [Online]. Available: https://doi.org/10.1145/3503221.3508399 [12] M. Zheng and Z. Zhang, “Radius: Range-based Gradient Sparsity for Large Foundation Model Pre-training,” in Proceedings of Machine Learning and Systems, M. Zaharia, G. Joshi, and Y. Lin, Eds., vol. 7. MLSys, 2025. [Online]. Available: https://proceedings.mlsys.org/paper files/paper/ 2025/file/54dd9e0cff6d9214e20d97eb2a3bae49-Paper-Conference.pdf [13] I. Markov, A. Vladu, Q. Guo, and D. Alistarh, “Quantized Distributed Training of Large Models with Convergence Guarantees,” 2023. [Online]. Available: https://arxiv.org/abs/2302.02390 [14] G. Wang, H. Qin, S. A. Jacobs, C. Holmes, S. Rajbhandari, O. Ruwase, F. Yan, L. Yang, and Y. He, “ZeRO++: Extremely Efficient Collective Communication for Giant Model Training,” 2023. [Online]. Available: https://arxiv.org/abs/2306.10209 [15] J. Jia, C. Xie, H. Lu, D. Wang, H. Feng, C. Zhang, B. Sun, H. Lin, Z. Zhang, X. Liu, and D. Tao, “SDP4Bit: Toward 4-bit Communication Quantization in Sharded Data Parallelism for LLM Training,” 2024. [Online]. Available: https://arxiv.org/abs/2410.15526 [16] S. U. Stich, J.-B. Cordonnier, and M. Jaggi, “Sparsified SGD with Memory,” 2018. [Online]. Available: https://arxiv.org/abs/1809.07599 [17] H. Zhang, B. Wang, and L. Chen, “AdamS: Momentum Itself Can Be A Normalizer for LLM Pretraining and Post-training,” in Proceedings of the 2025 Conference on Empirical Methods in Natural Language Processing, C. Christodoulopoulos, T. Chakraborty, C. Rose, and V. Peng, Eds. Suzhou, China: Association for Computational Linguistics, Nov. 2025, pp. 10 719–10 738. [Online]. Available: https://doi.org/10.18653/v1/2025.emnlp-main.543 [18] A. Gokaslan, V. Cohen, E. Pavlick, and S. Tellex, “OpenWebText Corpus,” http://Skylion007.github.io/OpenWebTextCorpus, 2019. [Online]. Available: https://doi.org/10.5281/zenodo.3834942 [19] D. Soboleva, F. Al-Khateeb, R. Myers, J. R. Steeves, J. Hestness, and N. Dey, “SlimPajama: A 627B token cleaned and deduplicated version of RedPajama,” 2023. [Online]. Available: https://huggingface. co/datasets/cerebras/SlimPajama-627B

[20] D. P. Kingma and J. Ba, “Adam: A Method for Stochastic Optimization,” 2017. [Online]. Available: https://arxiv.org/abs/1412.6980 [21] A. Radford, J. Wu, R. Child, D. Luan, D. Amodei, and I. Sutskever, “Language Models are Unsupervised Multitask Learners,” 2019. [22] Z. Shao, P. Wang, Q. Zhu, R. Xu, J. Song, X. Bi, H. Zhang, M. Zhang, Y. K. Li, Y. Wu, and D. Guo, “DeepSeekMath: Pushing the Limits of Mathematical Reasoning in Open Language Models,” 2024. [Online]. Available: https://arxiv.org/abs/2402.03300 [23] I. Molybog, P. Albert, M. Chen, Z. DeVito, D. Esiobu, N. Goyal, P. S. Koura, S. Narang, A. Poulton, R. Silva, B. Tang, D. Liskovich, P. Xu, Y. Zhang, M. Kambadur, S. Roller, and S. Zhang, “A Theory on Adam Instability in Large-Scale Machine Learning,” 2023. [Online]. Available: https://arxiv.org/abs/2304.09871 [24] Z. Bai, Z. Zhou, J. Zhao, X. Li, Z. Li, F. Xiong, H. Yang, Y. Zhang, and Z.-Q. J. Xu, “Adaptive preconditioners trigger loss spikes in adam,” 2025. [Online]. Available: https://arxiv.org/abs/2506.04805 [25] H. Robbins and S. Monro, “A Stochastic Approximation Method,” The Annals of Mathematical Statistics, vol. 22, no. 3, pp. 400–407, 1951. [Online]. Available: https://doi.org/10.1214/aoms/1177729586 [26] A. Ruhela, J. Cazes, J. D. McCalpin, C. Del-Castillo-Negrete, J. Li, H. Liu, H. Chen, C.-Y. Lu, K. F. Milfeld, W. Zhang, I. Wang, L. Koesterke, J. DeSantis, N. Lewis, S. Hempel, and D. Stanzione, “Performance Analysis of Scientific Applications on an NVIDIA Grace System,” in SC24-W: Workshops of the International Conference for High Performance Computing, Networking, Storage and Analysis. IEEE, 2024, pp. 558–566. [Online]. Available: https://doi.org/10.1109/SCW63240.2024.00078 [27] L. Gao, J. Tow, B. Abbasi, S. Biderman, S. Black, A. DiPofi, C. Foster, L. Golding, J. Hsu, A. Le Noac’h, H. Li, K. McDonell, N. Muennighoff, C. Ociepa, J. Phang, L. Reynolds, H. Schoelkopf, A. Skowron, L. Sutawika, E. Tang, A. Thite, B. Wang, K. Wang, and A. Zou, “The Language Model Evaluation Harness,” 07 2024. [Online]. Available: https://doi.org/10.5281/zenodo.12608602 [28] P. Clark, I. Cowhey, O. Etzioni, T. Khot, A. Sabharwal, C. Schoenick, and O. Tafjord, “Think you have Solved Question Answering? Try ARC, the AI2 Reasoning Challenge,” 2018. [Online]. Available: https://arxiv.org/abs/1803.05457 [29] D. Paperno, G. Kruszewski, A. Lazaridou, Q. N. Pham, R. Bernardi, S. Pezzelle, M. Baroni, G. Boleda, and R. Fernández, “The LAMBADA dataset: Word prediction requiring a broad discourse context,” 2016. [Online]. Available: https://arxiv.org/abs/1606.06031 [30] R. Zellers, A. Holtzman, Y. Bisk, A. Farhadi, and Y. Choi, “HellaSwag: Can a Machine Really Finish Your Sentence?” 2019. [Online]. Available: https://arxiv.org/abs/1905.07830 [31] D. Hendrycks, C. Burns, S. Basart, A. Zou, M. Mazeika, D. Song, and J. Steinhardt, “Measuring Massive Multitask Language Understanding,” 2021. [Online]. Available: https://arxiv.org/abs/2009.03300 [32] Y. Bisk, R. Zellers, R. L. Bras, J. Gao, and Y. Choi, “PIQA: Reasoning about Physical Commonsense in Natural Language,” 2019. [Online]. Available: https://arxiv.org/abs/1911.11641 [33] K. Sakaguchi, R. Le Bras, C. Bhagavatula, and Y. Choi, “WinoGrande: An Adversarial Winograd Schema Challenge at Scale,” in Proceedings of the AAAI Conference on Artificial Intelligence, vol. 34, no. 05, 2020, pp. 8732–8740. [Online]. Available: https://doi.org/10.1609/aaai. v34i05.6399 [34] T. Mihaylov, P. Clark, T. Khot, and A. Sabharwal, “Can a Suit of Armor Conduct Electricity? A New Dataset for Open Book Question Answering,” in Proceedings of the 2018 Conference on Empirical Methods in Natural Language Processing. Association for Computational Linguistics, 2018, pp. 2381–2391. [Online]. Available: https://doi.org/10.18653/v1/D18-1260 [35] A. Wang, Y. Pruksachatkun, N. Nangia, A. Singh, J. Michael, F. Hill, O. Levy, and S. R. Bowman, “Superglue: A stickier benchmark for general-purpose language understanding systems,” 2020. [Online]. Available: https://arxiv.org/abs/1905.00537 [36] P. Pfeiffer, P. Singer, Y. Babakhin, G. Fodor, N. Dhankhar, and S. S. Ambati, “H2O-Danube3 Technical Report,” 2024. [Online]. Available: https://arxiv.org/abs/2407.09276 [37] H. Touvron, L. Martin, K. Stone, P. Albert, A. Almahairi, Y. Babaei, N. Bashlykov, S. Batra, P. Bhargava, S. Bhosale, D. Bikel, L. Blecher, C. C. Ferrer, M. Chen, G. Cucurull, D. Esiobu, J. Fernandes, J. Fu, W. Fu, B. Fuller, C. Gao, V. Goswami, N. Goyal, A. Hartshorn, S. Hosseini, R. Hou, H. Inan, M. Kardas, V. Kerkez, M. Khabsa, I. Kloumann, A. Korenev, P. S. Koura, M.-A. Lachaux, T. Lavril, J. Lee, D. Liskovich, Y. Lu, Y. Mao, X. Martinet, T. Mihaylov,

P. Mishra, I. Molybog, Y. Nie, A. Poulton, J. Reizenstein, R. Rungta, K. Saladi, A. Schelten, R. Silva, E. M. Smith, R. Subramanian, X. E. Tan, B. Tang, R. Taylor, A. Williams, J. X. Kuan, P. Xu, Z. Yan, I. Zarov, Y. Zhang, A. Fan, M. Kambadur, S. Narang, A. Rodriguez, R. Stojnic, S. Edunov, and T. Scialom, “Llama 2: Open Foundation and Fine-Tuned Chat Models,” 2023. [Online]. Available: https://arxiv.org/abs/2307.09288 [38] Q. Yi, J. Duan, H. Hu, Q. Hua, H. Zhao, S. Qian, D. Yang, J. Cao, J. Tang, Y. Yu, C. Liao, K. Wang, and L. Zhang, “EDGC: Entropy-driven Dynamic Gradient Compression for Efficient LLM Training,” 2025. [Online]. Available: https://arxiv.org/abs/2511.10333

[39] H. Wang, S. Sievert, Z. Charles, S. Liu, S. Wright, and D. Papailiopoulos, “ATOMO: Communication-efficient Learning via Atomic Sparsification,” 2018. [Online]. Available: https: //arxiv.org/abs/1806.04090 [40] T. Vogels, S. P. Karimireddy, and M. Jaggi, “PowerSGD: Practical Low-Rank Gradient Compression for Distributed Optimization,” 2020. [Online]. Available: https://arxiv.org/abs/1905.13727 [41] J. Song, J. Yim, J. Jung, H. Jang, H.-J. Kim, Y. Kim, and J. Lee, “Optimus-CC: Efficient Large NLP Model Training with 3D Parallelism Aware Communication Compression,” 2023. [Online]. Available: https://arxiv.org/abs/2301.09830