用随机舍入提升大模型训练效率，省内存还更快。
Stochastic Rounding for LLM Training: Theory and Practice
- 用随机舍入缓解低精度训练的数值误差问题。
- 在67亿参数模型上，速度提升1.54倍，内存减少30%。
- 理论支撑下优化策略，适合大规模分布式训练场景。
随着大语言模型参数量突破千亿级，高效训练方法的需求愈发迫切——需在加快计算、减少内存消耗的同时保持精度。近年来提出的混合精度策略虽能提升训练速度，但常依赖手动调参且缺乏理论支持。本文利用随机舍入（SR）来缓解低精度表示带来的数值误差，首次在Adam优化器下提供隐式正则化与收敛性的理论分析。基于这些分析，我们将原有的BF16 + SR策略扩展至分布式训练环境，显著提升稳定性与性能。在最大达67亿参数的预训练模型上，实验首次证明：该方法优于传统的(BF16, FP32)混合精度策略，在验证困惑度上表现更优，吞吐量最高提升1.54倍，内存使用降低30%。
原文摘要 · Abstract (English)
As the parameters of Large Language Models (LLMs) have scaled to hundreds of billions, the demand for efficient training methods -- balancing faster computation and reduced memory usage without sacrificing accuracy -- has become more critical than ever. In recent years, various mixed precision strategies, which involve different precision levels for optimization components, have been proposed to increase training speed with minimal accuracy degradation. However, these strategies often require manual adjustments and lack theoretical justification. In this work, we leverage stochastic rounding (SR) to address numerical errors of training with low-precision representation. We provide theoretical analyses of implicit regularization and convergence under the Adam optimizer when SR is utilized. With the insights from these analyses, we extend previous BF16 + SR strategy to be used in distributed settings, enhancing the stability and performance for large scale training. Empirical results from pre-training models with up to 6.7B parameters, for the first time, demonstrate that our BF16 with SR strategy outperforms (BF16, FP32) mixed precision strategies, achieving better validation perplexity, up to $1.54\times$ higher throughput, and $30\%$ less memory usage.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品；中文卡片由大模型生成，请以原文为准。