arXiv:2505.11170cs.LG2025-05

用高斯采样实现高效稳定的伪量化训练,支持极低精度参数

Gaussian Weight Sampling for Scalable, Efficient and Stable Pseudo-Quantization Training

  • 采用高斯噪声分布实现轻量级伪量化,通过加法与浮点转换完成模拟
  • 在A100上仅增1.4%计算开销,每参数仅需2字节显存,支持FP6至BF16
  • 训练稳定,10亿参数模型3000亿词训练下性能媲美甚至超越BF16基准

大型语言模型规模持续增长,推动更高效的训练方式,全量化训练(FQT)相比BF16更具优势。但FQT存在一致性挑战,且需在超指数级的组合中搜索,每种配置需超过2000亿词才能保证稳定。伪量化训练(PQT)缓解了这些问题,但研究仍不充分。本文深入探讨了PQT的实际意义,提出一种面向浮点数(FP)友好、具备随机精度退火特性的噪声分布 $R$。该方法为通过PQT实现低精度浮点参数提供了有效理论基础,利用加法与后续浮点类型转换实现高效伪量化。实验表明,基于高斯权重采样的方法具有:(1) 可扩展性——支持低至FP6的参数精度与高达9位的噪声精度,同时兼容BF16算子;(2) 高效性——在A100 GPU上计算开销仅为1.40\\(\%\),每参数仅占2字节显存;(3) 稳定性——在预训练GPT2和Llama2模型(最大10亿参数,3000亿词)时,性能接近或超越BF16基线。

原文摘要 · Abstract (English)

Ever-growing scale of large language models (LLMs) is pushing for improved efficiency, favoring fully quantized training (FQT) over BF16. While FQT accelerates training, it faces consistency challenges and requires searching over an exponential number of cases, each needing over 200B tokens to ensure stability. Pseudo-quantization training (PQT) addresses the issues of FQT, although it is not well-studied. We explore the practical implications of PQT in detail and propose a noise distribution $R$ that is floating-point (FP)-friendly, with ideal properties including stochastic precision annealing. As a result, the proposed method serves as an effective theoretical foundation for low-precision FP parameters through PQT, utilizing efficient fake quantization via an addition and subsequent FP casting. We demonstrate that Gaussian weight sampling is (1) scalable: supports low-precision FP parameters down to FP6 and high-precision noise up to 9-bit with BF16 operator. The proposed method is (2) efficient: incurring computational overhead as low as 1.40\% on the A100 GPU in terms of Llama2 training tokens per second, and requiring 2 bytes per parameter in GPU memory. We demonstrate that PQT with Gaussian weight sampling is (3) stable: closely following or even surpassing performance of the BF16 baseline while pre-training GPT2 and Llama2 models with up to 1B parameters and 300B tokens.

伪量化低精度训练高斯采样显存优化

Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。