BitSparse: Compact Activation Storage for ReLU Feed-Forward Networks
作者: Max Zhu, Bohang Sun, Pietro Lio, Jialin Yu
分类: cs.LG
发布日期: 2026-10-15
💡 一句话要点
提出BitSparse以解决大语言模型训练中的内存使用问题
🎯 匹配领域: 支柱九：具身大模型 (Embodied Foundation Models)
关键词: 稀疏激活 无损压缩 大语言模型 内存优化 深度学习
📋 核心要点
- 现有的大语言模型训练方法在内存使用上存在显著挑战，尤其是在处理稀疏激活时。
- 论文提出的BitSparse通过无损压缩稀疏激活，结合自定义Triton内核实现高效的内存管理。
- 实验结果表明，BitSparse在多个模型上显著降低了内存占用，同时保持了高效的训练性能。
📝 摘要（中文）
在大型语言模型（LLM）训练中，减少内存使用一直是一个持续的挑战。激活检查点通过重新计算来节省内存，而有损压缩则改变存储的激活值。稀疏ReLU和ReLU²激活的重新出现为无损压缩激活缓存提供了新的机会。我们提出了BitSparse，这是一种GPU驻留的稀疏激活无损压缩格式。压缩和重构在自定义的Triton内核中运行，并融合到自动梯度函数中，缩小激活缓存的同时在反向传播中精确恢复。支持Bfloat16（BF16）和FP8激活。可选的符号位打包进一步提高了压缩率，接近理论极限。我们分析了其压缩比，并与理论最优值进行了比较。BitSparse在Modded-NanoGPT和Nemotron-H上分别降低了31%和30%的每个token VRAM，实际序列长度下的开销不超过1%。此外，新的预分配缓冲区系统避免了在训练过程中分配可变大小张量时的主机同步。
🔬 方法详解
问题定义：本论文旨在解决大型语言模型训练中的内存使用问题，尤其是激活存储的高开销。现有方法如激活检查点和有损压缩存在一定的局限性，无法有效利用稀疏激活的特性。
核心思路：论文提出的BitSparse通过无损压缩稀疏激活，利用GPU的计算能力和自定义内核，确保在反向传播过程中能够精确恢复激活值，从而实现内存的高效管理。
技术框架：整体架构包括激活的压缩和重构两个主要模块。压缩过程在Triton内核中实现，重构则与自动梯度计算相结合，确保训练过程的高效性。
关键创新：BitSparse的核心创新在于其无损压缩机制，能够在不损失信息的情况下显著减少激活缓存的大小。这一方法与传统的有损压缩技术有本质区别，后者会改变激活值。
关键设计：在技术细节上，BitSparse支持Bfloat16和FP8激活格式，并通过可选的符号位打包进一步提高压缩效率。预分配缓冲区系统的设计避免了在训练过程中因张量分配而导致的主机同步问题。
🖼️ 关键图片
📊 实验亮点
实验结果显示，BitSparse在Modded-NanoGPT和Nemotron-H模型上分别降低了31%和30%的每个token VRAM使用，且在实际序列长度下的开销不超过1%。这一显著的内存节省为大规模模型训练提供了新的可能性。
🎯 应用场景
BitSparse的研究成果在大型语言模型的训练中具有广泛的应用潜力，尤其是在需要高效内存管理的场景中，如自然语言处理和深度学习模型的训练。其无损压缩技术能够有效降低内存开销，提高训练效率，未来可能对多种深度学习应用产生积极影响。
📄 摘要（原文）
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.