BitSnap: Checkpoint Sparsification and Quantization in LLM Training
Yanxin Peng, Qingping Li, Baodong Wu, Shigang Li, Guohao Dai, Shengen Yan, Yu Wang
TL;DR
BitSnap tackles the growing cost of large-scale LLM checkpointing by introducing an asynchronous checkpoint engine that combines bitmask-based sparsification for model states with cluster-based quantization for optimizer states. The system leverages in-memory redundancy to accelerate recovery and integrates with Megatron-LM to minimize disruption to training. Empirical results show substantial compression (up to $16\times$ for model states and $2\times$ for optimizer states) and dramatic reductions in checkpointing time, with minimal impact on convergence ($\text{MRE}/\text{MSE}$ metrics remaining small and loss curves stable). The approach offers a practical path to faster, fault-tolerant training of ever-larger language models by optimizing storage, memory, and IO bottlenecks in a scalable, scheme-aware manner.
Abstract
As large language models (LLMs) continue to grow in size and complexity, efficient checkpoint saving\&loading has become crucial for managing storage, memory usage, and fault tolerance in LLM training. The current works do not comprehensively take into account the optimization of these several aspects. This paper proposes a novel checkpoint sparsification and quantization method that adapts dynamically to different training stages and model architectures. We present a comprehensive analysis of existing lossy and lossless compression techniques, identify current limitations, and introduce our adaptive approach that balances compression ratio, speed, and precision impact throughout the training process. Experiments on different sizes of LLMs demonstrate that our bitmask-based sparsification method achieves 16x compression ratio without compromising model accuracy. Additionally, the cluster-based quantization method achieves 2x compression ratio with little precision loss.