Memory-Efficient LLM Training with Dynamic Sparsity: From Stability to Practical Scaling

Technical Summary: Memory-Efficient LLM Training with Dynamic Sparsity

Problem Statement
Dynamic Sparse Training (DST) offers a paradigm for improving the efficiency of deep neural networks by maintaining a sparse topology throughout training. However, applying DST to Large Language Models (LLMs) has proven challenging due to optimization instability. The authors identify that naive application of standard Adam-based optimizers to DST results in severe "loss spikes" immediately following topology updates (pruning and regrowth).

The root cause is identified as a "cold-start" effect: newly regrown parameters lack accumulated optimizer states (first and second moment estimates). In Adam, the bias correction terms for these new parameters result in disproportionately large update magnitudes relative to mature weights. This disrupts the training trajectory, causing transient loss spikes and hindering convergence, particularly as model scales increase or topology updates become more aggressive. Furthermore, conventional DST implementations often fail to realize memory savings because they allocate full-sized optimizer states for the entire parameter space, even though only a subset of parameters is active.

Methodology: Sparse Memory-Efficient Training (SMET)
To address these challenges, the authors propose SMET, a framework designed to stabilize DST and reduce memory consumption for LLMs. The method comprises three core components:

1. Optimizer-State Warm-up: To mitigate the cold-start effect, SMET employs a two-stage warm-up strategy for newly regrown parameters:

   • Implicit Warm-up (Timestep Reset): The internal timestep counter ($t_i$) for regrown parameters is reset. This reactivates Adam's bias correction mechanism, inflating the denominator of the update rule and naturally constraining the initial update magnitude.
   • Explicit Warm-up (Learning Rate Scaling): A local linear warm-up is applied where the learning rate for regrown parameters increases linearly from zero to the target rate over $W$ steps.
   • These mechanisms jointly prevent update explosions, smoothing the optimization trajectory immediately after topology changes.

2. Density-Aware Learning Rate Scaling: To compensate for the conservative updates introduced by warm-up and to account for the reduced effective fan-in of sparse layers, SMET scales the learning rate. Motivated by Maximal Update Parameterization ($\mu$P), the learning rate is scaled by $1/\sqrt{d}$, where $d$ is the current sparsity density. This ensures that update magnitudes remain comparable to dense training regimes.

3. Memory-Efficient Optimization: SMET implements a sparse storage strategy where gradients and optimizer states (moment estimates $m$ and $v$) are stored only for active (non-zero) parameters.

   • Gradients are retained exclusively at active indices during backpropagation.
   • To manage the overhead of tracking active indices and timesteps, a block-wise storage strategy is used, grouping parameters into blocks of size $B$ to amortize metadata costs.
   • This approach preserves sparsity throughout the entire training lifecycle, unlike methods that compress states but produce dense models.

Key Contributions

1. Analysis of Instability: The paper identifies and theoretically analyzes the loss spiking phenomenon in LLM pre-training with DST, attributing it to the cold-start effect of newly regrown weights in Adam optimizers.
2. Stabilization Framework: The authors propose SMET, which integrates optimizer-state warm-up and density-aware learning-rate scaling to stabilize topology updates and improve training performance.
3. Memory-Efficient Implementation: A sparse optimization implementation is developed that stores gradients and optimizer states only for active parameters, supporting both unstructured and block-wise sparsity patterns while maintaining sparsity throughout training.
4. Systematic Evaluation: The work provides a comprehensive evaluation of DST across various model scales and sparsity levels, demonstrating that SMET outperforms existing DST methods and approaches dense training performance.

Experimental Results
Experiments were conducted on LLaMA-based models (ranging from 60M to 1B parameters) pre-trained on the C4 and OpenWebText datasets.

• Performance: SMET consistently outperforms other memory-efficient baselines (LoRA, ReLoRA, GaLore) and existing DST methods (SET, RigL). At larger model scales (e.g., 1B parameters), SMET achieves perplexity close to dense training, even at high sparsity levels (e.g., density $d=0.1$).
• Stability: The method effectively eliminates the loss spikes observed in vanilla DST, resulting in smooth training curves.
• Memory Efficiency: By storing states only for active parameters, SMET significantly reduces the memory footprint of gradients and optimizer states. For a 1B model at density $d=0.4$, substantial memory savings are observed.
• Scalability: The performance gap between sparse and dense training narrows as model size increases, suggesting sparse training is particularly viable for large-scale models.
• Generalization: The benefits of SMET generalize across different regrowth strategies (random and gradient-based) and different model architectures (LLaMA and Qwen).

Significance and Claims
The paper claims that SMET paves the way for sparse training as a practical alternative to dense training for LLMs. By solving the optimization instability inherent in dynamic sparse training and realizing actual memory savings through sparse state storage, SMET enables stable, scalable, and memory-efficient pre-training. The authors note that while SMET improves training stability and reduces optimizer memory, challenges remain regarding activation memory and the need for efficient hardware support for fine-grained sparsity to fully realize inference speedups. Nevertheless, the work demonstrates that with appropriate stabilization and memory management, dynamic sparsity can be successfully scaled to foundation models.