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

2026-05-30

AI Summary

This work addresses the optimization instability, loss spikes, and cold-start challenges commonly encountered in dynamic sparse training of large language models. To overcome these issues, the authors propose Sparse Memory-Efficient Training (SMET), the first method to systematically tackle the cold-start problem. SMET enhances training stability through optimizer warm-up and density-aware learning rate scaling, while significantly reducing memory overhead by maintaining gradients and optimizer states only for active parameters. Experimental results demonstrate that SMET achieves comparable model performance to dense training, yet offers substantially improved training stability and memory efficiency, establishing dynamic sparse pretraining as a practical alternative to conventional dense approaches.

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.

Problem

• Dynamic Sparse Training
• Large Language Models
• Optimization Instability
• Loss Spikes
• Training Efficiency

Innovation

• Dynamic Sparse Training
• Memory-Efficient Training
• Optimizer Warm-up
• Density-Aware Learning Rate
• Large Language Models