Bias Compounds, Variance Washes Out：随机舍入如何打破 BF16 训练的瓶颈

背景

在深度学习模型训练中，为了节省显存并加速计算，混合精度训练（Mixed Precision Training）已成为标准实践。其中，使用 BFloat16（BF16）格式存储模型参数和优化器状态（如 Adam 的一阶矩和二阶矩）可以显著减少内存占用。然而，BF16 的精度较低（1 位符号位、8 位指数位、7 位尾数位），在处理极小的梯度更新时，容易遇到“精度丢失”的问题。

传统的舍入模式是“向最近值舍入”（Round-to-Nearest, RNE）。在这种模式下，如果数值更新量小于 BF16 能表示的最小间隔，更新值会被舍入回原值，导致参数实际上没有发生变化。这种现象在长周期的训练中会导致损失函数停滞，模型无法收敛。

本文源自 Hacker News 讨论的一篇技术笔记，通过实验对比了“向最近值舍入”与“随机舍入”（Stochastic Rounding, SR）在 BF16 优化器状态中的表现，揭示了误差累积机制对模型训练的根本性影响。

核心内容

误差的本质：偏差与方差

文章首先从数学角度剖析了两种舍入模式的误差特性：

1. 向最近值舍入（RNE）：

   • 确定性误差：每次舍入产生的误差是固定的、相同的。
   • 偏差累积（Bias Compounds）：当相同的误差重复出现时，它们会线性叠加。例如，在 BF16 精度下，将 0.001 累加到 1.0 上一千次，由于每次更新量都小于 BF16 的分辨率，且舍入后仍指向 1.0，因此参数值永远无法改变。误差以 $O(n)$ 的速度增长，其中 $n$ 是步数。

2. 随机舍入（SR）：

   • 随机性误差：每次舍入产生的误差是随机的，但其期望值为零（Zero-mean）。
   • 误差抵消（Variance Washes Out）：由于误差围绕零分布，正负误差在多次迭代中会相互抵消。虽然单次更新可能向上舍入，另一次向下舍入，但在期望意义上，总和是精确的。
   • 方差扩散：在 $n$ 步之后，无偏误差（方差）的增长速度为 $O(\sqrt{n})$。虽然方差随步数增加而扩散，但其增长速度远慢于偏差的线性增长。在长时间的小更新过程中，这种差异至关重要。

实验设置与结果

为了验证上述理论，作者设计了一个小型的多层感知机（MLP）实验：

• 任务：基于 Teacher-Student 架构的回归任务。
• 优化器：使用 HeavyBall 版本的 AdamW。
• 配置：所有配置均将参数存储在 BF16 格式中。
• 变量：仅改变优化器状态中的舍入模式。

对比结果：

1. BF16 + SR（随机舍入）：

   • 内存占用：参数、一阶矩和二阶矩各占 6 字节（BF16 为 2 字节，但优化器状态通常需额外空间，此处指代整体优化器状态的紧凑表示或特定实现下的字节对比，原文意指 SR 方案在 6 字节宽度下表现优异，而对比基准 FP32 状态为 10 字节）。注：原文表述 "6 bytes for parameters, first moment, and second moment" 可能指代某种特定的量化或打包策略，或者是指 BF16 参数(2B) + 优化器状态(4B/8B) 的某种等效对比。根据上下文 "matches FP32 state (10 bytes)"，这里主要强调 SR 在低精度下达到了高精度训练的效果。
   • 性能：尽管每一步的更新引入了噪声，但由于噪声是无偏的，它在训练过程中逐渐抵消。最终性能与使用 FP32 状态（10 字节）的训练相匹配。

2. BF16 + RNE（向最近值舍入）：

   • 内存占用：6 字节。
   • 性能：损失函数在远高于 FP32 的水平上停滞（Plateaus）。由于相同的舍入误差不断重复，模型无法有效更新参数，导致训练失败。

关键结论

• 随机舍入的优势：通过在优化器内核中用随机舍入替换向最近值舍入，不仅没有增加额外的内存开销和带宽需求，还解决了 BF16 在低精度下因舍入误差导致的训练停滞问题。
• 精度与内存的权衡：消除偏差后，6 字节的 BF16 配置可以达到 10 字节 FP32 配置的性能水平。如果保留偏差（即使用 RNE），6 字节的配置将触及性能天花板，无法进一步收敛。

注：原文在 2026-03-16 进行了更正。原始实验中存在一个 torch.compile 融合问题，该问题意外消除了 BF16 的往返舍入操作，导致结果失真。修正后重新运行实验，确认了上述结论。

关键要点

• 误差累积机制不同：RNE 产生固定偏差，误差随步数线性累积（$O(n)$）；SR 产生零均值随机误差，误差随步数平方根增长（$O(\sqrt{n})$）。
• 小更新量的困境：在 BF16 中，当梯度更新量小于分辨率时，RNE 会导致参数不变（死区效应），而 SR 通过概率性舍入确保期望值的准确性。
• 性能对等：BF16 + SR 的训练效果可与 FP32 优化器状态相媲美，尽管单步更新存在噪声。
• 零成本优化：引入随机舍入无需增加内存占用或带宽开销，仅通过改变舍入逻辑即可提升低精度训练的稳定性。
• 实验严谨性：修正后的实验排除了编译器优化带来的干扰，证实了随机舍入在打破 BF16 精度瓶颈中的核心作用。

意义与影响

这一发现对大规模语言模型（LLM）和深度学习基础设施具有深远影响：

1. 降低硬件门槛：证明在优化器状态中使用 BF16 配合随机舍入，可以达到接近 FP32 的训练精度。这意味着在显存受限的硬件上，可以使用更低精度的格式存储优化器状态，从而允许更大的批量大小（Batch Size）或更深的模型架构。
2. 优化器效率提升：Adam 等优化器需要存储一阶和二阶矩，这通常占用了模型参数内存的数倍。如果优化器状态可以安全地保持在 BF16 甚至更低精度，将大幅降低内存足迹，提高训练吞吐量。
3. 重新评估低精度训练：长期以来，业界倾向于在优化器状态中使用 FP32 以保持精度。本文表明，通过算法层面的改进（随机舍入），可以更安全、更高效地使用低精度格式，无需牺牲模型收敛性。
4. 对编译器优化的警示：实验修正部分提醒开发者，在使用 torch.compile 等自动优化工具时，需警惕其可能改变数值计算路径（如消除舍入操作），从而导致实验结果与真实硬件行为不符。

总之，随机舍入作为一种简单而有效的技术，为在资源受限环境下进行高精度深度学习训练提供了新的解决方案，平衡了内存效率与模型性能。