直接用低精度权重训练语言模型，节省显存且性能不降。
Direct Quantized Training of Language Models with Stochastic Rounding
- 训练时全程使用量化权重，不用高精度副本。
- 8比特模型性能媲美BitNet b1.58，三值也可行。
- 适合资源受限场景，支持多精度推理部署。
尽管近期的量化大语言模型（如BitNet）已显著降低部署时的内存占用，但其训练仍需大量显存，主要因为必须全程保留高精度（未量化）权重以进行直通估计。为解决此问题，我们探索在反向传播中直接更新量化后的低精度权重，避免依赖直通估计，从而节省训练内存。具体地，采用随机舍入技术以最小化低比特权重带来的信息损失。在多种规模的LLaMA结构模型上实验表明：（1）即使权重被限制为三值，仅用低精度权重训练也完全可行；（2）将位宽扩展至8比特，性能可与BitNet b1.58相当；（3）模型对精度缩放和内存减少具有鲁棒性，在从FP32降至BF16/FP8环境下性能下降极小；（4）模型还支持使用三值权重进行推理，展现部署灵活性。
原文摘要 · Abstract (English)
Although recent quantized Large Language Models (LLMs), such as BitNet, have paved the way for significant reduction in memory usage during deployment with binary or ternary weights, training these models still demands substantial memory footprints. This is partly because high-precision (i.e., unquantized) weights required for straight-through estimation must be maintained throughout the whole training process. To address this, we explore directly updating the quantized low-precision weights without relying on straight-through estimation during backpropagation, aiming to save memory usage during training. Specifically, we employ a stochastic rounding technique to minimize the information loss caused by the use of low-bit weights throughout training. Experimental results on our LLaMA-structured models of various sizes indicate that (1) training with only low-precision weights is feasible even when they are constrained to ternary values; (2) extending the bit width to 8 bits achieves performance on par with BitNet b1.58; (3) our models remain robust to precision scaling and memory reduction, showing minimal performance degradation when moving from FP32 to lower-memory environments (BF16/FP8); and (4) our models also support inference using ternary weights, showcasing their flexibility in deployment.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品；中文卡片由大模型生成，请以原文为准。