BitSparse stores sparse feed-forward activations with less memory
In plain words
Training a language model requires keeping some values from its forward pass so the system can calculate updates afterward. ReLU-family feed-forward layers often produce exact zeros, which still take up space when stored normally.
BitSparse saves the nonzero values and a map of their positions, then rebuilds the saved values when needed. In tested workloads, it reduced memory use, including in two language models; the measurements do not cover every part of a full training run.
The research
Training systems must keep feed-forward activations for backpropagation, and this input-dependent memory grows with the number of tokens processed together. Can the exact zeros produced by ReLU-family activations be used to reduce that storage without recomputing values or changing them?
BitSparse stores nonzero feed-forward activations with position maps and reconstructs them for backpropagation, reducing the memory used by their cache.
The paper asks whether exact zeros in saved ReLU-family activations can reduce training memory without recomputation or approximation. Saved activations grow with the number of tokens processed at once.
BitSparse stores nonzero values and their positions, then reconstructs them for backpropagation. In tested forward-and-backward language-model runs, fitted memory per token fell by 31% for Modded-NanoGPT and 30% for Nemotron-H versus dense activation storage. These runs did not update parameters or retain optimizer state.
The memory measurements cover selected ReLU-family feed-forward layers and workloads. The full training comparison also used TensorBuffer, and the source does not specify whether the two runs started from the same state.
Analysis
Deep
- Original paper authors
- Max Zhu, Bohang Sun, Pietro Lio, Jialin Yu
- Submitted to arXiv
- License
- CC BY 4.0