Stochastic Rounding for LLM Training: Theory and Practice
作者: Kaan Ozkara, Tao Yu, Youngsuk Park
分类: cs.LG
发布日期: 2025-02-27
备注: AISTATS 2025
💡 一句话要点
提出基于随机舍入的BF16训练策略，提升LLM训练效率与稳定性
🎯 匹配领域: 支柱九：具身大模型 (Embodied Foundation Models)
关键词: 随机舍入 低精度训练 大型语言模型 混合精度 分布式训练
📋 核心要点
- 现有混合精度训练策略缺乏理论支撑，需要手动调整，限制了LLM训练效率的进一步提升。
- 论文核心在于利用随机舍入（SR）缓解低精度训练中的数值误差，并从理论上分析了其隐式正则化和收敛性。
- 实验结果表明，提出的BF16 + SR策略在模型困惑度、吞吐量和内存使用方面均优于传统混合精度策略。
📝 摘要（中文）
随着大型语言模型（LLM）的参数扩展到数千亿级别，对高效训练方法的需求变得前所未有地重要，这些方法需要在不牺牲准确性的前提下，平衡更快的计算速度和更低的内存使用量。近年来，各种混合精度策略被提出，它们为优化组件采用不同的精度级别，以在最小化精度下降的同时提高训练速度。然而，这些策略通常需要手动调整，并且缺乏理论依据。在这项工作中，我们利用随机舍入（SR）来解决低精度表示训练中的数值误差。我们提供了在使用Adam优化器时，SR下的隐式正则化和收敛性的理论分析。基于这些分析的见解，我们扩展了先前的BF16 + SR策略，使其可用于分布式环境，从而增强了大规模训练的稳定性和性能。首次对高达67亿参数的模型进行预训练的经验结果表明，我们的BF16 + SR策略优于（BF16，FP32）混合精度策略，实现了更好的验证困惑度，高达1.54倍的吞吐量提升，以及30％的内存使用量减少。
🔬 方法详解
问题定义：论文旨在解决大规模语言模型（LLM）训练过程中，使用低精度（如BF16）表示时出现的数值误差问题。现有的混合精度训练策略虽然能加速训练并减少内存占用，但往往缺乏理论依据，需要手动调整，且在大规模分布式训练中稳定性不足。
核心思路：论文的核心思路是利用随机舍入（Stochastic Rounding, SR）来缓解低精度表示带来的数值误差。SR通过引入随机性，使得舍入误差的期望为零，从而减少了偏差，并具有隐式正则化的效果，有助于提高模型的泛化能力和训练稳定性。
技术框架：论文主要包含以下几个部分：1) 对Adam优化器下使用SR的收敛性进行理论分析，揭示其隐式正则化作用；2) 将BF16 + SR策略扩展到分布式训练环境，解决大规模训练中的稳定性问题；3) 通过实验验证BF16 + SR策略在LLM预训练中的有效性。整体流程是先进行理论分析，然后基于分析结果改进现有策略，最后通过实验验证改进后的策略。
关键创新：论文的关键创新在于将随机舍入（SR）应用于LLM的低精度训练，并从理论上证明了SR的隐式正则化和收敛性。与传统的舍入方法相比，SR能够更有效地缓解低精度带来的数值误差，提高训练的稳定性和模型的泛化能力。此外，论文还将BF16 + SR策略成功扩展到分布式训练环境，解决了大规模训练中的实际问题。
关键设计：论文的关键设计包括：1) 使用Adam优化器进行训练；2) 在BF16精度下进行前向和反向传播，并使用随机舍入；3) 针对分布式训练，优化了数据并行策略，以提高训练效率和稳定性。具体的参数设置和网络结构与所使用的LLM模型相关，论文主要关注的是优化算法和精度策略。
🖼️ 关键图片
📊 实验亮点
实验结果表明，在预训练高达67亿参数的LLM时，提出的BF16 + SR策略优于传统的（BF16，FP32）混合精度策略。具体而言，验证困惑度更低，吞吐量提高了高达1.54倍，内存使用量减少了30%。这些结果表明，BF16 + SR策略在提高LLM训练效率和降低资源消耗方面具有显著优势。
🎯 应用场景
该研究成果可广泛应用于大规模语言模型的预训练和微调，尤其是在资源受限的环境下。通过降低内存占用和提高训练速度，可以加速LLM的开发和部署，并降低训练成本。此外，该方法也有潜力应用于其他深度学习模型的训练，提升模型性能和训练效率。
📄 摘要（原文）
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.