基于随机取整的直接量化语言模型训练
机器学习
2025-10-13 v3 计算与语言
摘要
虽然最近的量化大语言模型（LLM）如BitNet已为部署期间的显存使用显著减少铺平了道路（采用二进制或三值权重），但训练这些模型仍需要大量显存。这部分原因是必须在整个训练过程中维护用于直线估计的高精度（即未量化）权重。为此，我们探索在反向传播期间直接更新量化低精度权重，而无需依赖直线估计，以节省训练期间的显存使用。具体而言，我们采用随机取整技术来最小化使用低位权重期间造成的信息损失。在各种规模的LLaMA结构模型上的实验结果表明：(1)即使在权重受限为三值时，也可实现低精度训练；(2)将位宽扩展到8位可实现与BitNet b1.58持平的性能；(3)我们的模型对精度缩放和内存减少具有鲁棒性，在从FP32降至低内存环境（BF16/FP8）时性能几乎不受影响；(4)我们的模型还支持使用三值权重进行推理，显示出在部署方面的灵活性。
引用
@article{arxiv.2412.04787,
title = {Direct Quantized Training of Language Models with Stochastic Rounding},
author = {Kaiyan Zhao and Tsuguchika Tabaru and Kenichi Kobayashi and Takumi Honda and Masafumi Yamazaki and Yoshimasa Tsuruoka},
journal= {arXiv preprint arXiv:2412.04787},
year = {2025}
}
备注
Accepted to ACML 2025