BitSparse:面向 ReLU 稀疏激活的无损压缩格式
BitSparse: Compact Activation Storage for ReLU Feed-Forward Networks
训练 LLM 显存吃紧的可以看看:BitSparse 靠 ReLU 稀疏性无损压缩激活缓存,NanoGPT 和 Nemotron-H 上实测省三成显存,几乎不拖慢速度。
arXiv 论文提出 BitSparse,一种针对 ReLU/ReLU² 稀疏激活的 GPU 无损压缩格式,通过定制 Triton 内核融合进 autograd 函数来缩减激活缓存。它支持 BF16 和 FP8 激活,可选符号位打包使压缩比接近理论上限。实测在 Modded-NanoGPT 和 Nemotron-H 上分别将每 token 的 VRAM 占用降低 31% 和 30%,开销不超过 1%。论文还引入预分配缓冲区机制,在训练中分配可变大小张量时避免主机同步。
BitSparse: Compact Activation Storage for ReLU Feed-Forward Networks
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.