BitSparse reduces memory use in large language model training
BitSparse: Compact Activation Storage for ReLU Feed-Forward Networks
Machine Learning
Summary
Training large language models uses a lot of computer memory, especially to store intermediate data called activations. The authors introduce BitSparse, a way to shrink this stored data without losing any information, making training use less memory. BitSparse works well with certain common activation types and runs directly on GPUs with little slowdown. This lets bigger or longer trainings fit into the same hardware more easily.
What this means in practice
- •For machine learning infrastructure teams: Reduce GPU memory usage during training of large language models by compressing sparse activations without slowdown.
- •For cloud service providers: Offer more cost-efficient LLM training instances by enabling denser GPU workloads with BitSparse activation compression.$Commercial implications: BitSparse enables cloud providers to sell higher capacity GPU training services that handle larger LLMs with the same hardware.
Authors
Max Zhu, Bohang Sun, Pietro Lio, Jialin Yu
Abstract
Reducing memory use in large language model (LLM) training is an ongoing challenge. Activation checkpointing saves memory by recomputation, while lossy compression changes the stored activations. The reemergence of sparse $\relu$ and $\relu^2$ activations creates a third opportunity: losslessly compressing the activation cache itself. We present BitSparse, a GPU-resident lossless compression format for sparse activations. Compression and reconstruction run in custom Triton kernels fused into an autograd function, shrinking the activation cache while recovering it exactly for the backward pass. Bfloat16 (BF16) and FP8 activations are supported. Optional sign-bit packing further increases compression up to close to the theoretical limit. We analyze its compression ratio and compare it to the theoretical optimum. BitSparse lowers per-token VRAM by 31\% and 30\% Modded-NanoGPT and Nemotron-H respectively, with overheads of no more than 1\% at practical sequence lengths. Additionally, a new preallocated buffer system avoids host synchronization during training when variable size tensors are allocated.