Track papers across labs, topics, and trends.
Search papers, labs, and topics across Lattice.
BitSparse introduces a GPU-resident, lossless compression format for sparse ReLU and ReLU^2 activations, enabling significant memory reduction in LLM training by exploiting activation sparsity. By fusing custom Triton kernels for compression and reconstruction within an autograd function, it achieves substantial VRAM savings (30-31%) with minimal overhead, supporting BF16 and FP8 formats and optional sign-bit packing. This approach offers a novel avenue for memory optimization beyond activation checkpointing and lossy compression, complemented by a preallocated buffer system to mitigate host synchronization bottlenecks.
Lossless compression of sparse ReLU activations slashes LLM activation VRAM by up to 31% with near-zero overhead.
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.