PaperlayerView exact revision on arXiv

Source-linked AI summary

BitSparse: Compact Activation Storage for ReLU Feed-Forward Networks

Max Zhu, Bohang Sun, Pietro Lio, Jialin Yu

arXiv:2610.22256v1cs.LG

TL;DR

Activation storage grows with the number of tokens, motivating lossless compression of exact zeros in ReLU-family activations without recomputation or approximation. BitSparse stores sparse activations in a GPU-resident format and reconstructs them for backward computation. It reduces per-token memory by 31% on Modded-NanoGPT and 30% on Nemotron-H, while its gains are scoped to feed-forward activations that retain exact structural zeros.

  • Problem

    Activation memory grows with the number of simultaneously processed tokens, motivating savings from exact ReLU-family zeros without recomputation or approximation.

    Source: §1 · p. 1

  • Method

    BitSparse losslessly stores sparse activations using a GPU-native bitmask and compact nonzero values, then reconstructs them during backward computation.

    Source: §1 · p. 2

  • Results

    Per-token VRAM falls by 31% on Modded-NanoGPT and 30% on Nemotron-H with BitSparse.

    Source: §6.4 · p. 8

  • Takeaways & Limitations

    BitSparse converts naturally occurring exact zeros in feed-forward activations into memory savings without approximation, with overhead typically below 1% at large sequence lengths.

    Source: §7 · p. 9

    BitSparse depends on ReLU-family activations retaining exact structural zeros and applies only to feed-forward activations, not attention or Mamba layers.

    Source: §7 · p. 9

Abstract

from arXiv · show

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.

1 Introduction

Input-dependent activation storage constrains the batch sizes and sequence lengths that fit on accelerators. BitSparse targets naturally sparse ReLU-family activations with lossless caching, reducing memory without recomputation or approximation.

1 Introduction

  • Saved activations and temporary buffers grow with the number of simultaneously processed tokens, unlike largely input-independent parameter storage.This input-dependent memory constrains the batch sizes and sequence lengths that fit on an accelerator.Source: §1 · p. 1
  • Feed-forward layers are a memory target because they expand token representations by 3–5× and retain intermediates for weight gradients.The passage notes that feed-forward layers remain a significant memory cost in architectures with linear-in-token memory attention or state-space components.Source: §1 · p. 1
  • Pretrained Nemotron-H has approximately 92% feed-forward activation sparsity, and ReLU2 is used in Nemotron-family and Modded-NanoGPT models.The passage describes this sparsity as emerging naturally during training without a sparsity-inducing objective.Source: §1 · p. 1
  • The motivating question is whether exact zeros can reduce activation memory without checkpointing recomputation or lossy approximation.Checkpointing recomputes forward intermediates, while lossy compression changes stored activations.Source: §1 · p. 1
  • BitSparse stores only nonzero ReLU-family activation values and reconstructs them with Triton kernels.The approach exploits exact zeros already present in sparse activations.Source: §1 · p. 2, Figure 1 · p. 2
  • 31% less memory per token on Modded-NanoGPT and 30% on Nemotron-H are reported alongside minimal runtime overhead.The MLP evaluation reports peak allocated memory reductions of up to 48%.Source: §1 · p. 2

2 Background

The background identifies the activation information needed for ReLU-family gradients and establishes an entropy reference for sparse activation storage. Its entropy model assumes independent, identically distributed entries with a specified zero probability and payload distribution.

2.1 ReLU activations and their gradients

  • ReLU2 backpropagation uses the positive preactivation values and their positions as the activation-dependent information for computing gradients.The section reviews ReLU2 followed by a linear projection and gives the associated forward and gradient expressions.Source: §2.1 · p. 2

2.2 Entropy of sparse ReLU activations

  • The entropy reference assumes independent, identically distributed entries, with zero probability ρ and uniformly distributed (d −1)-bit nonzero payloads.Under this assumption, entropy per entry combines support entropy h2(ρ) and expected nonzero payload cost (1 −ρ)(d −1).Source: §2.2 · p. 3
  • For actual nonnegative floating-point activations, the stated entropy expression is an upper bound because payload restrictions, nonuniformity, or positional correlations can lower entropy.Source: §2.2 · p. 3

3 Methodology

BitSparse stores sparse feed-forward activations in a tile-based format and reconstructs them on the GPU for backward computation. Its method combines packed masks and compact values with optional sign-bit packing, fused kernels, and an asynchronous shared-value buffer.

3.1 Overview

  • BitSparse compresses ReLU-like activations before caching and reconstructs them during backward, while the feed-forward projection remains dense.The cached representation preserves values in the selected storage format.Source: §3.1 · p. 3

3.2 Tile-wise sparse activation format

  • The tile format partitions activations into 64 × 64 tiles and represents them with a bitmask, compact nonzero-value array, and prefix offsets.The mask uses one bit per padded position; each tile has 512 bytes of mask metadata and 32 bits of prefix metadata.Source: §3.2 · p. 3

3.3 Sign-bit packing

  • Optional sign-bit packing removes the redundant sign bit from nonnegative values, reducing BF16 values to 15 bits and FP8 values to 7 bits.The custom packing and decoding kernels operate without atomics.Source: §3.3 · p. 4

3.4 GPU compression and reconstruction

  • GPU compression uses tile kernels to build masks and counts, a scan for offsets, and a scatter of positive values; reconstruction uses a Triton kernel to gather values and rebuild activations.Backward kernels can apply the ReLU support mask without reconstructing activations, and blockwise backward matrix multiplication avoids materializing the full activation.Source: §3.4 · p. 4

3.5 Shared value buffers and implementation details

  • Appending activation values to a preallocated TensorBuffer avoids host synchronization when variable-sized value arrays are allocated.The buffer is reset each training step and must be sized to avoid overflow; its runtime reduction is 5.2% at 32 tokens and becomes negligible beyond 128.Source: §3.5 · p. 4

3.6 Sparse activation checkpointing for feed-forward layers

  • A custom autograd function combines ReLU or ReLU2 with a linear matrix multiplication and supports BF16 and FP8; for ReLU2, the unsquared activation is stored.The compression overhead relative to the linear layer is O(1/D), where D is the hidden dimension in the cited analysis.Source: §3.6 · pp. 4–5

4 Theoretical compression analysis

The theoretical analysis compares BitSparse storage with an ideal entropy-based rate under an assumed sparse-payload model. It derives the implemented format’s storage costs and identifies sparsity levels at which storage falls below dense representation.

4 Theoretical compression analysis

  • For B = 64 with negligible padding and per-tensor overhead, the compression ratios are derived from mask bits, prefix bits, and nonzero values stored at d or d −1 bits.Sign-bit packing uses d −1 bits per nonzero value, compared with d bits without packing.Source: §4 · p. 5
  • At ρ = 1/2, BitSparse-sb’s excess storage over the ideal rate is exactly the 1/128 prefix cost.Source: §4 · p. 5
  • Storage is reduced above sparsity thresholds of roughly 6% and 13% for unpacked BF16 and FP8, respectively, and below 0.2% with sign-bit packing.The thresholds are given as (1+1/128)/d without packing and 1/[128(d −1)] with packing.Source: §4 · p. 5

5 Related work

Prior work spans lossy activation compression, low-precision and structured-sparsity methods, and lossless activation storage; BitSparse is distinguished by its GPU-training focus and implementation features. Other approaches target checkpointing or inference-time sparsity.

Lossy activation compression and low-precision training

  • ActNN uses randomized mixed-precision quantization, GACT dynamically selects compression ratios for generic architectures, and CompAct stores low-rank activations using random projections.These methods are lossy activation-compression approaches.Source: §5 · p. 5
  • FP8 halves per-value storage relative to BF16, COAT targets activation and optimizer-state memory with FP8 quantization, and 2:4 methods prune ReLU2 activations for training and inference.Haziza et al. (2025) investigate hardware-accelerated 2:4 sparsity, retaining two of every four activations as nonzero.Source: §5 · p. 5

Lossless memory reduction techniques

  • Prior lossless methods target CNN edge training, host-cached GPU transfers, or 8-bit CSR, while BitSparse adds GPU training without host synchronization, fused low-overhead kernels, low-precision support, and sign-bit packing.The authors describe these features as essential for adoption in modern LLM training.Source: §5 · p. 5
  • Activation checkpointing trades speed for lower memory by recomputing activations during backpropagation, and has been used in large LLM training setups with thousands of GPUs.The cited LLM training work is Korthikanti et al. (2023).Source: §5 · pp. 5–6
  • Table 2 covers encoding and decoding timings for a 1 GiB sparse BF16 tensor, comparing sign-bit packing modes and reporting read-plus-write bandwidth against attainable and nameplate bandwidth.The caption specifies a tensor shape of 65,536 × 8,192 and a nameplate bandwidth of 936 GB/s.Source: Table 2 · p. 6

Activation sparsity for efficient computation

  • Inference methods exploit sparse activations by skipping weights linked to zeros, but do not generalize to training because different tokens activate different channels.The cited discussion includes CNN inference work and ReLU-based LLM inference methods.Source: §5 · p. 6

6 Experiments

Experiments find substantial ReLU2 activation sparsity in both evaluated language models and show that BitSparse can reduce memory use while keeping runtime close to dense storage. Benchmarks also compare its kernels and memory–runtime trade-offs with alternatives, and a full training run preserves validation loss while supporting a larger microbatch.

6.1 Sparsity in LLMs

  • Modded-NanoGPT activation sparsity rose from approximately ρ = 0.5 to ρ = 0.82, while Nemotron-H mean sparsity was 92.73% and its minimum batch-wise mean was 91.76%.Nemotron-H's minimum batch-wise mean was only 0.97 percentage points below its mean.Source: §6.1 · p. 6

6.2 Benchmarking and Ablations

  • On a 65,536 × 8,192 BF16 tensor, Triton encoded 7.59× faster and decoded 7.91× faster than pure PyTorch, while peak allocated memory fell by 37.5% and 21.9%, respectively.For a 1 GiB tensor at ρ = 0.5 without sign-bit packing, decoding reached 97.4% of attainable bandwidth and encoding reached 60.9%.Source: §6.2 · p. 6
  • At 50% sparsity on a 16,000 × 21,504 BF16 tensor, the fused backward gradient kernel cut peak memory by 46.8% for ReLU and 60.9% for ReLU2, and runtime by 17.2% and 39.2%, respectively.Source: §6.2 · p. 6

6.3 Comparisons on training

  • At ρ = 0.9, BitSparse was the most memory-efficient lossless method, reducing VRAM by up to 48% relative to dense storage; its memory ranking varied by dtype and sparsity.For BF16, checkpointing used the least memory at ρ = 0.5, but BitSparse overtook it at ρ = 0.9; for FP8, BitSparse-sb used the least memory at both sparsity levels.Source: Table 3 · p. 7, §6.3 · p. 7
  • As batch size increased from 32 to 32,000, memory savings rose from 6% to 55%, while runtime stayed within 1% of dense storage from 512 rows upward.BitSparse's advantage grows with naturally occurring activation sparsity while avoiding recomputation.Source: Table 4 · p. 7, §6.3 · p. 7
  • For ReLU2 at ρ = 0.9, BitSparse dominated the memory–runtime Pareto frontier, particularly at higher sparsity; the final point's VRAM drop reflected lower transient memory from the fused pointwise gradient kernel.Source: §6.3 · p. 7, Figure 2 · p. 8

6.4 Memory and runtime in LLMs

  • 31% and 30%: asymptotic VRAM per token falls from 0.402 MiB to 0.277 MiB with BitSparse on Modded-NanoGPT and from 2.719 MiB to 1.894 MiB on Nemotron-H.BitSparse-sb reaches 0.276 MiB and 1.890 MiB, respectively; BitSparse incurs a negligible performance penalty, while checkpointing generally has a notable runtime cost.Source: §6.4 · p. 8, Figure 3 · p. 8, Figure 4 · p. 9

6.5 Full training comparison

  • Validation loss ended at 3.2778 for standard training and 3.2780 for BitSparse-sb, while BitSparse supported microbatch size 64 when standard training ran out of memory.At microbatch size 64, BitSparse step time was 7.416 seconds, slightly faster than standard training at size 32 (7.428 seconds); at size 32, BitSparse took 7.668 seconds per step.Source: §6.5 · pp. 8–9

7 Conclusion

BitSparse losslessly stores sparse ReLU and ReLU2 feed-forward activations, reducing memory with optimized fused kernels. The conclusion also notes its activation-dependent scope, while the accompanying statements describe AI assistance, reproducibility materials, and ethics.

7 Conclusion

  • BitSparse is a lossless, GPU-resident, tile-wise drop-in representation for ReLU and ReLU2 feed-forward activations, typically adding under 1% overhead at large sequence lengths while reducing memory by up to 30%.It supports BF16 and FP8 and can operate without host synchronization.Source: §7 · p. 9
  • The method depends on ReLU-family activations producing exact structural zeros and does not affect attention or Mamba layers.The paper identifies extending the representation to gated activations and co-designing it with checkpointing, offloading, and distributed training as future work.Source: §7 · p. 9

AI use statement

  • AI tools assisted proofreading and Triton-kernel optimization, while the remaining supporting software was written manually with minor AI help and experiments were run manually.The paper says its Autoresearch prompt is included in the supplementary materials.Source: AI use statement · p. 9

Reproducibility statement

  • Supplementary materials include the BitSparse implementation, benchmark scripts, and language-model experiments; appendices document proofs, setups, workloads, measurements, and additional results.The reproducibility statement locates these details in the named sections and appendices.Source: Reproducibility statement · p. 9

Ethics statement

  • BitSparse does not collect data or introduce a new data source; privacy, bias, and misuse risks depend on its use case, and environmental impact depends on workload memory savings and runtime overhead.Source: Ethics statement · p. 9

A Memory usage during training

The appendix models training memory as parameter-related storage plus input-dependent activations, then estimates activation footprints for Nemotron-H. Its estimates are scoped by tensor lifetimes, implementation choices, and which tensors BitSparse compresses.

A Memory usage during training

  • The training-memory analysis assumes BF16 storage throughout and distinguishes static parameter and optimizer storage from activations that scale with input size.For the stated Muon setup, the passage gives a static overhead of 4P; FlashAttention and Mamba use recomputation or fused checkpointing to avoid storing full intermediate computations.Source: A Memory usage during training · p. 12

A.1 Nemotron-H memory usage

  • For Nemotron-H at 4096 tokens, retaining one expanded BF16 tensor per MLP costs 3.9375 GiB; a conventional ReLU2 backward graph may retain two such sets, or 7.875 GiB.The precise saved set depends on fusion and checkpointing.Source: §A.1 · p. 12
  • Nemotron-H implementation optimizations fuse RMS normalization and chunk the vocabulary projection, avoiding retained normalization intermediates and full FP32 logits; both dense and BitSparse baselines use them.The changes alter intermediate storage rather than the computation, with results matching the unfused formulation up to floating-point rounding.Source: §A.1 · p. 12
  • At mean observed sparsity ρ = 0.9273, one expanded BF16 MLP activation set is estimated at 0.534 GiB without sign-bit packing or 0.516 GiB with packing, versus 3.9375 GiB dense.These are estimates: sparsity was measured after squaring, while the cache stores unsquared ReLU values; reconstruction may still need a dense temporary, and attention, Mamba, residual, and output-head tensors are not compressed.Source: §A.1 · p. 13
  • At fixed precision, parameter storage is independent of input length, while listed linear-width activation footprints scale with microbatch token count; materialized attention scores scale with batch size times squared sequence length.The benchmark omits optimizer state and a retained full-model gradient buffer, and tensor lifetimes overlap only partially, so its component rows should not be summed into a predicted peak.Source: §A.1 · p. 13

B Analysis of storage usage

The storage analysis derives a reference entropy for sparse activations under a uniform, independent payload model, then relates it to actual floating-point values and the implemented representation. The reference model is a worst-case compression rate because realistic restrictions can only lower entropy.

B Analysis of storage usage

  • Under an independent model with zero probability ρ and uniformly distributed (d −1)-bit nonzero payloads, sparse-entry entropy comprises support entropy h2(ρ) plus expected payload cost (1 −ρ)(d −1).The model supplies the stated reference entropy and optimum asymptotic compression ratio relative to dense d-bit storage.Source: B Analysis of storage usage · p. 13
  • For actual nonnegative floating-point activations, the reference entropy is an upper bound: valid-value restrictions, nonuniform positive values, or correlations can only lower entropy.Subadditivity means dependence across entries does not add a penalty to the bound.Source: B Analysis of storage usage · p. 14
  • The reference bound using d −1 is strict because zero occupies one of the sign-zero patterns, leaving fewer than 2^(d−1) finite positive values.The entropy gap is at least L(1 −ρ); the passage notes equality when entries are independent and positive values are uniform over the valid set.Source: B Analysis of storage usage · p. 14
  • The implemented representation stores nonzero values, tile masks, and prefix metadata; with negligible padding and overhead, its metadata costs 1 + 1/128 bits per entry for B = 64.The cost formula distinguishes payload width d without sign-bit packing from d −1 with packing.Source: B Analysis of storage usage · p. 14

C.1 ReLU gradient derivation

The gradient derivation motivates caching positive activation values and their positions, and the fused algorithm compresses only that cache while retaining dense forward computation. The section also specifies the experimental platform and comparison baselines.

C.1 ReLU gradient derivation

  • The gradients with respect to Z and W depend on the positive activation values and their positions cached during the forward pass.The derivation defines ReLU′(0) = 0.Source: §C.1 · p. 14

C.2 Forward and backward algorithm

  • The fused autograd function compresses the activation cache after evaluating the activation in place, keeps the forward matrix multiplication dense, and decompresses for the weight gradient in backward.When the input gradient is needed, its ReLU derivative is applied from the stored bitmask without reconstructing the dense activation; BF16 or FP8 values are supported, with optional sign-bit packing.Source: §C.2 · p. 15

D.1 System setup

  • Experiments used PyTorch 2.14 and Triton 3.7.1 on an RTX 3090 24GB, except FP8 MLP and Nemotron-H experiments on an RTX 5090 32GB.The passage identifies the latter GPU as providing FP8 support and extra VRAM.Source: §D.1 · p. 15

D.2 Baselines

  • Comparisons include checkpointing, general lossless GPU codecs, CSR, and CompAct, a lossy method whose compression level trades memory against accuracy.Checkpointing recomputes activations; CSR index overhead is suboptimal at the stated 0.5–0.85 sparsity range and CSR lacks FP8 support.Source: §D.2 · pp. 15–16

D.3 MLP training benchmark setup

The benchmark setup specifies the MLP and language-model workloads, then describes how feed-forward activation sparsity is measured and how full-training memory, runtime, and validation loss are compared.

D.3 MLP training benchmark setup

  • MLP tests use eight residual blocks with 4096-to-21,504-to-4096 projections, synthetic BF16 inputs, and gradients checked against the dense baseline up to floating-point rounding.Parameters are not updated; measurements exclude optimizer work.Source: §D.3 · p. 16

D.4 Language-model memory and runtime setup

  • Language-model tests compare dense storage, checkpointing, BitSparse, and sign-bit packing on BF16 Modded-NanoGPT and Nemotron-H workloads, measuring peak memory and step time.Modded-NanoGPT sweeps sequence lengths with four sequences; Nemotron-H processes one sequence at each length.Source: §D.4 · p. 17

E.1 Activation sparsity measurement in LLMs

  • Sparsity is the exact-zero fraction in ReLU2 outputs, summarized across evaluation batches and layers; the reported minima apply only to evaluated batches.The Modded-NanoGPT measurements cover all 12 feed-forward layers, while Nemotron-H measurements exclude attention and Mamba blocks.Source: §E.1 · pp. 17–18

E.2 Full training comparison

  • The full-training comparison reports memory use and step time at two Modded-NanoGPT microbatch sizes, alongside validation loss over the run.The supplied passages identify Table 7 and Figure 6 but do not provide their measured values.Source: §E.2 · p. 19, Table 7 · p. 19, Figure 6 · p. 19

E.3 MLP benchmark results

The benchmarks examine BitSparse memory and runtime across activation types, sparsities, implementation choices, and layer counts. Results show that savings vary with sparsity and workload size, while the shared buffer’s largest measured runtime benefit occurs at the smallest input.

E.3 MLP benchmark results

  • At 32,000 rows, BitSparse cuts peak memory by 36% for ReLU2 at ρ = 0.5, 26% for ReLU at ρ = 0.5, and 50% for ReLU at ρ = 0.9, with runtime within about 1.5% of dense.Sign-bit packing becomes the smallest configuration at 8,000 rows for both ρ = 0.9 settings, but not at ρ = 0.5.Source: §E.3 · p. 19, Table 8 · p. 20, Table 9 · p. 20

E.4 Codec implementation ablation

  • The codec ablation compares pure PyTorch and custom Triton implementations on a 65,536 × 8,192 BF16 tensor, measuring peak memory and runtime.Both implementations produce bit-identical outputs.Source: §E.4 · p. 21, Table 12 · p. 21

E.5 Activation backward ablation

  • The activation-backward ablation isolates derivative computation on a 16,000 × 21,504 BF16 tensor at approximately 50% sparsity, including saved representation and backward temporaries in peak memory.BitSparse fuses reconstruction and the derivative into an in-place update.Source: §E.5 · p. 21, Table 13 · p. 21

E.6 Shared value buffer ablation

  • At 32 tokens, the shared value buffer reduces runtime from 12.563 to 11.904 ms; at 256–1024 tokens, differences remain below 0.4% in magnitude.The comparison uses ReLU2 at ρ = 0.5.Source: §E.6 · p. 21, Table 14 · p. 21

E.7 Layer-count trade-offs

  • At ρ = 0.5, checkpointing uses less VRAM in both settings; at ρ = 0.9, BitSparse uses less VRAM than checkpointing for ReLU at eight layers and both settings cross below it at the largest layer count.Moving from 7 to 8 compressed layers lowers ReLU2 VRAM by 23% at ρ = 0.9 and 15% at ρ = 0.5, versus 7% and 1% for ReLU; sign-bit packing saves slightly more memory at a small time cost.Source: §E.7 · p. 22, Figure 7 · p. 22, Figure 8 · p. 22

Related papers

Loading 2610.22256v1…