Memory-Efficient LLM Training with Dynamic Sparsity: From Stability to Practical Scaling
Abstract
Dynamic Sparse Training (DST) offers a promising paradigm for improving the training and inference efficiency of deep neural networks; however, we find that in large language model training, DST can suffer from optimization instability, manifested as loss spikes after topology updates. In this work, we show that the naive use of standard Adam-based optimizers leads to a cold-start issue for newly regrown parameters, resulting in excessively large updates and disrupted training dynamics. To address this issue, we propose Sparse Memory-Efficient Training (SMET), which stabilizes DST with optimizer warm-up and improves training progress through density-aware learning-rate scaling. SMET further reduces memory consumption by storing gradients and optimizer states only for active parameters. We provide a theoretical analysis of the update behaviors under SMET, showing improved optimization stability. Extensive experiments demonstrate that SMET enables stable, scalable, and memory-efficient sparse pre-training of LLMs, paving the way for sparse training as a practical alternative to dense training. Our code is publicly available at: https://github.com/QiaoXiao7282/SMET.
Lay Summary
Training large language models requires substantial memory and computing resources. One promising way to reduce these costs is to train sparse models, which use only a fraction of their possible connections at any given time while allowing the active connections to change during training. However, we find that introducing new connections can cause sudden spikes in training loss. We propose Sparse Memory-Efficient Training (SMET), a method that makes sparse training more stable and memory-efficient. SMET helps sparse models learn more effectively and stably, and stores training information only for currently active connections. Experiments across different model sizes and sparsity levels show that SMET reduces memory use while achieving performance close to standard full-model training, particularly for larger models. These results suggest that sparse training can be a practical approach for training large language models.