FP4 All the Way: Fully Quantized Training of LLMs

标题: FP4 一路通行:LLM 的全量化训练
作者/机构: Brian Chmiel †∗, Maxim Fishman †∗, Ron Banner ∧, Daniel Soudry ◦
†Intel, Israel ∧Nvidia, Israel ◦Technion 电气与计算机工程系, Haifa, Israel


A1 主要贡献

随着大型语言模型(LLM)的快速发展,模型参数已扩展至数千亿,这带来了巨大的计算能力和内存带宽需求,使得硬件限制成为主要瓶颈。为提高计算效率和减少内存需求,研究从主流的 BF16 精度转向了更低的精度,如 FP8 已被证明是可行且有效的,并迅速成为新的标准。在此背景下,探索更进一步的 FP4 精度成为合乎逻辑的下一步,它有望在保持训练稳定性和模型准确性的同时,实现更高的效率。特别是 NVIDIA 新推出的 Blackwell GPU 架构首次在硬件上支持 FP4 矩阵乘法,相比 FP8 可将吞吐量提升近两倍。本文旨在系统性地研究 FP4 训练的各种设计选择,并首次实现了在权重、激活和梯度上均使用 FP4 的全量化训练。

本文的主要贡献如下:

  1. FP4 格式优化。 本文对 FP4 训练中的块大小和缩放格式进行了全面的实验,发现块大小小于 16 个元素时,准确性收益递减。通过比较不同的指数-尾数配置(从 E1M6 到 E8M0),发现 E4M3 格式(NVFP4 中使用的格式)性能最佳,从而验证了 NVIDIA 的硬件设计选择。
  2. 分离式舍入策略。 本文确定了在反向传播中使用随机舍入(Stochastic Rounding, SR),而在前向传播中使用四舍五入到最近值(Round-to-Nearest, RtN)的组合策略。与在整个训练过程中单独使用任一方法相比,该策略能显著提高 FP4 训练的稳定性和最终模型的准确性。
  3. 精度转换分析。 本文提供了一个理论框架,用于识别 FP4 精度在何时对持续训练不再有效的临界点。具体来说,当全精度梯度的标准差降至低于量化噪声标准差的 3 倍时,训练效果会减弱。因此,本文建议在训练末期应用更高精度的量化感知微调(Quantization Aware Finetuning, QAF),以提高信噪比,从而快速收敛至基线水平。
  4. 端到端的大规模 FP4 训练。 本文使用 256 个 Intel Gaudi2 加速器,成功地在万亿(trillion)级别的 tokens 上完全以 FP4 精度训练了一个 7B 参数的 LLM。虽然与 BF16 基线相比,最终训练损失存在微小差距,但通过短暂的 QAF 阶段可以完全弥补。最终,模型在下游任务上的性能与 BF16 基线相当,证实了 FP4 在实际大规模应用中的可行性。

A2 相关工作

量化作为神经网络压缩的关键技术。 量化是压缩神经网络以实现更高效部署和减少资源消耗的关键研究领域。该领域主要有两种方法:训练后量化(Post-Training Quantization, PTQ)【22, Guangxuan Xiao et al., Smoothquant: Accurate and efficient post-training quantization for large language models, 2022, ArXiv】【11, Ji Lin et al., Awq: Activation-aware weight quantization for on-device llm compression and acceleration, 2023, GetMobile Mob. Comput. Commun.】【9, Elias Frantar et al., Gptq: Accurate post-training quantization for generative pre-trained transformers, 2022, ArXiv】和量化感知训练(Quantization-Aware Training, QAT)【7, Tim Dettmers et al., Qlora: Efficient finetuning of quantized llms, 2023, ArXiv】【3, Wenhua Cheng et al., Optimize weight rounding via signed gradient descent for the quantization of llms, 2023, ArXiv】【12, Andrei Panferov et al., Quest: Stable training of llms with 1-bit weights and activations, 2025, ArXiv】。PTQ 专注于将预训练模型转换为低比特表示而无需额外训练,适用于快速部署,尤其是在推理场景中。相比之下,QAT 在模型训练或微调期间就引入量化效应,使网络能够适应并在低精度约束下保持准确性。两种方法都取得了显著进展,近期工作已证明在 4-bit 精度【9, Elias Frantar et al., Gptq: Accurate post-training quantization for generative pre-trained transformers, 2022, ArXiv】及以下【20, Hongyu Wang et al., Bitnet: Scaling 1-bit transformers for large language models, 2023, ArXiv】也能达到有竞争力的性能。

全量化训练(FQT)的挑战与进展。 全量化训练(FQT)比 PTQ 或 QAT 更具挑战性,因为它要求从头开始使用低精度权重、激活和梯度进行训练,以加速所有矩阵乘法。直到最近,将 FQT 应用于 16-bit 以下的精度仍被认为很困难,因为它存在不稳定和收敛问题。然而,近期的研究证明了其可行性。【13, Houwen Peng et al., Fp8-lm: Training fp8 large language models, 2023, ArXiv】首次展示了在高达 1000 亿 tokens 的数据集上以 FP8 精度进行大型语言模型的 FQT。【8, Maxim Fishman et al., Scaling fp8 training to trillion-token llms, 2024, ArXiv】将其扩展到 2 万亿 tokens,揭示了训练后期的稳定性问题,并提出了一种改进的激活函数来解决。【6, DeepSeek-AI et al., Deepseek-v3 technical report, 2024, ArXiv】通过使用更细粒度的量化来减轻不稳定性,进一步推动了该领域的发展,成功地以 FP8 FQT 训练了一个大型专家混合(MoE)模型。

更细粒度量化成为 FP4 FQT 的关键方向。 更细粒度的量化正成为推动全量化训练(FQT)超越 FP8 的关键方向,尤其是在 FP4 精度下。通过减小量化块大小,这些方法旨在更好地捕捉数据分布的局部变化,从而提高稳定性和准确性。典型的例子包括 MXFP4【15, Bita Darvish Rouhani et al., Microscaling data formats for deep learning, 2023, ArXiv】和 NVFP4【1, Nvidia blackwell architecture, URL https://resources.nvidia.com/ en-us-blackwell-architecture】。

与本文最相关的工作及其局限性。 与我们工作最相关的两项研究是【21, Ruizhe Wang et al., Optimizing large language model training using fp4 quantization, 2025, ArXiv】和【19, Albert Tseng et al., Training llms with mxfp4, 2025, ArXiv】。【21】提出使用一种逐向量的 FP4 格式训练大型语言模型,并结合了两种关键技术:可微分梯度估计器(DGE)以替代标准的直通估计器(STE),以及离群点裁剪补偿(OCC)通过一个额外的稀疏残差矩阵来处理激活离群值。他们的模型在高达 1000 亿 tokens 的数据上进行训练,但他们只量化了权重和激活,将梯度保持在较高精度,因此只加速了训练中涉及的三个矩阵乘法中的一个。【19】则专注于使用 MXFP4 格式进行梯度量化,并应用随机舍入和哈达玛变换来稳定训练。他们训练了高达 400 亿 tokens 的模型。然而,与【21】类似,他们也只加速了部分矩阵乘法。相比之下,我们的工作是第一个展示对大型 LLM 进行完整 FP4 全量化训练(FQT)的,从而能够在训练期间加速所有矩阵乘法。

A3 方法细节

3 FP4 训练

FP4 训练的挑战与微缩放浮点格式(MXFP)。 从 FP8 转向 FP4 训练带来了巨大挑战,因为 FP4 的动态范围有限,难以在不过度产生量化误差的情况下捕捉激活和梯度的全部变化。然而,最近引入的微缩放浮点(MXFP)系列【15, Bita Darvish Rouhani et al., Microscaling data formats for deep learning, 2023, ArXiv】通过在更细的粒度上动态调整缩放因子,提供了一种有前景的替代方案以减轻精度损失。

MXFP4 和 NVFP4 格式的比较。 MXFP4 是一种浮点格式,包含 1 个符号位、2 个指数位和 1 个尾数位(E2M1),它通过将数据分成大小为 32 的块来增强低精度训练,每个块共享一个共同的缩放因子。每个块的缩放因子使用 E8M0 格式存储,这是一种 8 位纯指数表示,提供宽动态范围而没有尾数和符号。另一种潜在的 FP4 训练格式是 NVFP4,它使用与 MXFP4 相同的 E2M1 数据表示,但在块大小和缩放格式上有所不同。与 MXFP4 的 32 相比,NVFP4 将数据分成更小的块(大小为 16),从而允许进行更细粒度的缩放调整。此外,NVFP4 采用 E4M3 格式存储缩放因子,在动态范围和精度之间取得了平衡。NVIDIA 的 Blackwell 架构【1, Nvidia blackwell architecture, URL https://resources.nvidia.com/ en-us-blackwell-architecture】同时支持 MXFP4 和 NVFP4。表 1 对这两种格式进行了比较。

表 1:MXFP4 和 NVFP4 格式的比较
表 1:MXFP4 和 NVFP4 格式的比较

3.1 探索块大小和缩放格式

研究动机。 观察到 Blackwell 支持的两种 FP4 格式——MXFP4 和 NVFP4——在块大小和缩放格式方面存在差异,我们决定进一步研究它们的影响。具体来说,我们的目标是比较所有可能的缩放格式(保持 FP8 数据格式),并探索不同的块大小,以了解它们对数值稳定性、训练效率和模型准确性的影响。这项分析将为低精度训练中动态范围、精度和计算效率之间的权衡提供更深入的见解。

实验与发现。 在图 1 中,我们使用 FP4 格式(E2M1)、块大小为 16 和不同的缩放格式来训练一个 350M 的 Llama 风格模型。值得注意的是,与 NVFP4 类似,大多数配置(除了对应于 MXFP4 的 E8M0)在缩放因子中并未使用符号位,这可能代表了未来工作可以利用的潜在低效率。我们观察到,使用 E3M4 和 E4M3 取得了最好的结果,后者正是 NVFP4 格式所采用的。在图 2 中,我们比较了使用 E8M0 和 E4M3 两种缩放格式(分别是 MXFP4 和 NVFP4 使用的缩放格式)下不同块大小的效果。需要注意的是,选择合适的缩放格式对最终准确性有显著影响——例如,E1M6 导致了完全的发散——而块大小的影响则较为温和,通常较小的块大小会产生更好的结果。这使我们决定继续使用 NVFP4 格式,即块大小为 16,缩放格式为 E4M3。

图 1:格式 E4M3(用于 NVFP4)和 E3M4 取得了最佳结果。比较了在训练一个 350M Llama 模型时,使用 FP4 格式(E2M1)和块大小 16 的不同缩放格式(E1M6, E2M5, E3M4, E4M3, E5M2, E6M1, E8M0)。格式 E3M4 和 E4M3 取得了最佳结果(注意 E4M3 用于 NVFP4),而 E1M6 导致了完全的发散。
图 2:块大小 16 是最佳选择。我们使用缩放格式 (a) E8M0(用于 MXFP4)和 (b) E4M3(用于 NVFP4)检验了不同块大小(8, 16, 32, 64, 128)对训练准确性的影响。较小的块大小在准确性上带来适度改善,但收益在每块 16 个元素以下时递减。因此,块大小 16 在性能和计算开销之间提供了最优的折衷。

3.2 探索舍入模式

转向舍入模式的研究。 在完成了对 FP4 格式中各种块大小和缩放因子的探索之后,我们现在将注意力转向另一个关键方面:舍入模式的选择。在下一节中,我们将深入研究 FP4 可用的不同舍入策略,并分析它们对数值稳定性和模型性能的影响。

全量化训练中的六个量化点。 全量化训练包含了对三个关键的通用矩阵乘法(GEMM)的量化:前向(forward)、反向(backward)和更新(update)。每个 GEMM 都涉及两个量化操作数,从而在整个训练流程中产生了六个不同的量化点:

$\textbf{[Forward]} \quad z_l = Q(W_l)Q(a_{l-1}); \quad a_l = f_l(z_l)$

$$ \textbf{[Backward]}\ g_{l-1}=Q(W_l^T)Q(\delta_l);\ \delta_l=f'(z_l) \odot g_l $$

$$[Update] \frac{\partial C}{\partial W_l} = Q(\delta_l)Q(a_{l-1}^T),$$

其中 C 是损失函数,Q 是量化操作,⊙ 是逐元素乘积,并且在每一层 l 中,$f_l$ 是激活函数,$W_l$ 是权重矩阵,$z_l$ 是预激活值,$g_l \triangleq \frac{\partial C}{\partial a_l}$。

不同舍入模式的实验与结论。 值得注意的是,我们可以为这六个元素中的每一个独立选择舍入模式。在图 3 中,我们展示了使用 NVFP4 训练 350M Llama 模型的结果,其中对这六个元素中的每一个分别应用随机舍入(SR),而其余元素使用四舍五入到最近值(RtN),从而评估其各自的贡献。在更新和反向 GEMM 中对神经梯度应用 SR,以及在更新 GEMM 中对激活应用 SR,有助于降低训练损失;而在其他组件中使用 SR 则会导致损失增加。在附录的图 7 中,我们展示了支持相同结论的额外实验。因此,我们采用以下选择性舍入方案:

图 3:使用 NVFP4 格式训练 350M Llama 模型时不同舍入方案的比较。在每个图中,我们将 SR 应用于其中一个 GEMM 的六个元素之一,而其余部分使用四舍五入到最近值(RtN)。注意到,在‘更新’和‘反向’ GEMM 中对神经梯度应用 SR,以及在‘更新’ GEMM 中对激活应用 SR,可以降低训练损失,而对其他组件应用 SR 则会产生相反效果,增加损失。
图 3:使用 NVFP4 格式训练 350M Llama 模型时不同舍入方案的比较。在每个图中,我们将 SR 应用于其中一个 GEMM 的六个元素之一,而其余部分使用四舍五入到最近值(RtN)。注意到,在‘更新’和‘反向’ GEMM 中对神经梯度应用 SR,以及在‘更新’ GEMM 中对激活应用 SR,可以降低训练损失,而对其他组件应用 SR 则会产生相反效果,增加损失。

$$ \textbf{[Forward]} \quad z_l = Q_{\text{RtN}}(W_l) Q_{\text{RtN}}(a_{l-1}); \quad a_l = f_l(z_l) $$


$[\text{Backward}] \quad g_{l-1} = Q_{\text{RtN}}(W_l^T)Q_{\text{SR}}(\delta_l); \quad \delta_l = f_l'(z_l) \odot g_l$
[Update] $\frac{\partial C}{\partial W_l} = Q_{\text{SR}}(\delta_l)Q_{\text{SR}}(a_{l-1}^T)$,

4 对使用随机舍入的量化 SGD 的分析

本节分析了在使用随机舍入(SR)时,低精度梯度训练在何时会停止有效。虽然 SR 消除了偏差并在训练的大部分时间里实现了稳定的下降,但一旦梯度相对于量化噪声变得过小,其优势就会减弱。我们推导出一个梯度与噪声比的阈值,该阈值标志着进一步的进展停滞,从而促使在反向和更新传播中切换精度。这种切换可以改善收敛性,而无需改变前向传播或部署的模型。

分析的核心要点:
* SR 实现无偏更新: SR 实现了无偏更新,即使在激进的 FP4 量化下也能稳定下降。相比之下,确定性舍入引入了持续的偏差,导致不可约的误差下限并阻碍收敛(见附录 B.2)。
* 存在一个临界阈值: 通过 SR,我们在 4.1 节(并在附录 B.1 中更详细地)证明,当平均每个坐标的梯度幅值下降到大约低于量化噪声标准差的 3 倍时,训练不再产生有效的损失减少。该阈值是在简化假设下推导的——例如,使用最优步长的梯度下降、损失的泰勒近似以及集中的 Hessian 谱。
* 经验证据支持理论: 在 4.2 节中,对于合成和真实模型设置,我们凭经验表明性能在该阈值以下急剧下降。由该理论指导的精度切换可以恢复收敛。

4.1 理论推导

推导过程。 我们从损失函数在当前参数矢量 $θ_t$ 附近的二阶泰勒展开开始,这揭示了一个与 $-∇L^T ∆θ$ 成比例的下降项和一个涉及 Hessian 矩阵 H 的曲率项。然后,我们将全精度梯度 $∇L$ 替换为其量化版本 $g_q = ∇L + ε$,其中 $ε$ 是由随机舍入引入的零均值噪声。

期望损失变化。 使用更新规则 $∆θ = -ηg_q$ 并在 SR 条件下取期望($E[ε] = 0, E[εε^T] = σ_q^2 I$),我们得到期望的损失变化:

$$ \mathbb{E}[\Delta L] = -\eta\|\nabla L\|^2 + \frac{1}{2}\eta^2 \left(\nabla L^T H \nabla L + \sigma_q^2 \mathrm{tr}(H)\right). $$

最优步长。 平衡下降项和噪声项,得到最优步长:

$$\eta^* = \frac{\|\nabla L\|^2}{\nabla L^T H \nabla L + \sigma_q^2 \text{tr}(H)}$$

代入最优步长后的期望损失。 将 $η^∗$ 代入期望损失中得到:

$$\mathbb{E}[\Delta L] = -\frac{\|\nabla L\|^4}{2(\nabla L^T H \nabla L + \sigma_q^2 \text{tr}(H))}.$$

临界点推导。 对 $σ_q$ 求导,并使用一些简化假设,揭示了对量化噪声的敏感度在以下情况达到峰值:

$$\sigma_{\text{critical}} = \frac{\left\| \nabla L \right\|}{\sqrt{3d}}.$$

结论。 一旦每个坐标的梯度幅值下降到 $3 σ_q$ 以下,下降变得微不足道,需要更高精度的梯度来继续改善损失。

4.2 经验验证

为了验证这个理论阈值,我们提出了两种类型的实验。

简单二次损失实验。 首先,我们模拟在一个简单的二次损失上进行训练,并采用自适应噪声调度。我们将量化噪声缩放为 $σ_q = k · σ_{critical}$,其中 k = 2, 1, 0.5。如图 4 所示,在高噪声水平下(例如,k = 2),收敛完全停滞;在临界阈值附近(k = 1)时,收敛减慢;当噪声降低到阈值以下(k = 0.5)时,训练过程与全精度训练紧密相符。

图 4:在简单二次损失中对理论推导的验证。玩具二次模型中,噪声水平为 σq = k · σcrit,k = 2, 1, 0.5 时的训练损失。高噪声阻碍下降;低噪声允许持续进展。
图 4:在简单二次损失中对理论推导的验证。玩具二次模型中,噪声水平为 σq = k · σcrit,k = 2, 1, 0.5 时的训练损失。高噪声阻碍下降;低噪声允许持续进展。

真实 Llama 模型实验。 其次,我们在图 5 中测试了一个真实的 60M 参数 Llama 模型上的阈值。在训练过程中,我们监控比率 $∥∇L∥/(σ_q \sqrt{d})$,并在第 1000 次迭代时切换到更高精度的梯度。当这个比率超过 3 时,与全精度基线的损失差距立即消失,验证了该阈值的预测能力。

图 5:在 Llama 60M 模型中验证理论预测。(左):基线损失曲线与在训练中途(第 1000 次迭代,垂直虚线橙色线)增加精度的模型损失曲线之间的差异。增加精度后,损失差异完全消除。(右):梯度与噪声比率及 3 的阈值(黑色虚线)。
图 5:在 Llama 60M 模型中验证理论预测。(左):基线损失曲线与在训练中途(第 1000 次迭代,垂直虚线橙色线)增加精度的模型损失曲线之间的差异。增加精度后,损失差异完全消除。(右):梯度与噪声比率及 3 的阈值(黑色虚线)。

A4 实验环境

  • 模型架构: 使用 Llama2 模型【18, Hugo Touvron et al., Llama 2: Open foundation and fine-tuned chat models, 2023, ArXiv】作为基线,这是一个 decoder-only 的 Transformer【2, Tom Brown et al., Language models are few-shot learners, 2020, Advances in Neural Information Processing Systems】,采用了 pre-normalization RMSNorm【23, Biao Zhang and Rico Sennrich, Root mean square layer normalization, 2019, Advances in Neural Information Processing Systems】、Smooth-SwiGLU 激活函数【8, Maxim Fishman et al., Scaling fp8 training to trillion-token llms, 2024, ArXiv】和旋转位置嵌入【17, Jianlin Su et al., Roformer: Enhanced transformer with rotary position embedding, 2024, Neurocomput.】。
  • 数据集: 在开源的 Red Pajama 数据集【5, Together Computer, Redpajama: an open dataset for training large language models, 2023, URL https://github.com/togethercomputer/RedPajama-Data】上训 练 1 万亿(1T)个 tokens。
  • 硬件配置: 所有训练均在 256 个 Intel Gaudi2 设备上进行,耗时约 30 天。
  • 软件与超参数配置:
    • 优化器: AdamW,其中 $β_1 = 0.9, β_2 = 0.95$。
    • 学习率调度: 使用余弦学习率调度,包含 2000 步的 warmup,峰值学习率为 $3 × 10^{−4}$,并衰减至峰值学习率的 0.1。
    • 批处理大小: 全局批处理大小为 400 万个 tokens。
    • 其他: 保持与【18】一致的超参数,包括训练-测试集划分和初始化。

A4 实验结果

  • FP4 全量化训练:

    • 实验内容: 采用提出的 FP4 方案训练 Llama2 7B 模型。该方案使用 NVFP4 格式(块大小 16,缩放格式 E4M3),并在神经梯度(更新+反向 GEMM)和激活(更新 GEMM)中应用随机舍入(SR),而在权重(前向+反向 GEMM)和激活(前向 GEMM)中应用四舍五入到最近值(RtN)。
    • 实验结果: 如图 6a 所示,FP4 训练的损失曲线与 BF16 基线非常接近,仅存在一个微小的差距。
    • 分析结论: 本文首次实现了对所有矩阵乘法进行加速的全量化 FP4 训练,与之前仅部分量化的工作(如表 2 所示)相比,这是一个重要突破。此外,不同随机种子下的训练损失非常相似(见附录表 4),表明该方法具有噪声鲁棒性。
    图 6:完整的 FP4 训练 7B 模型达到了与 BF16 基线相同的损失。(a):使用提出的 FP4 方案训练 Llama2 7B 的训练损失,该方案包括 NVFP4 格式(块大小 16,缩放格式 E4M3),并在神经梯度(更新+反向 GEMM)和激活(更新 GEMM)中使用 SR。注意到训练损失存在微小差距,但这并未导致下游任务性能出现显著差异(表 3)。(b):使用 NVFP4 格式在前向 GEMM 中、BF16 在反向和更新 GEMM 中进行 Llama2 7B 的量化感知微调(QAF)训练损失。注意到这个简短的 QAF 能够完全弥补与 BF16 基线的差距。
    图 6:完整的 FP4 训练 7B 模型达到了与 BF16 基线相同的损失。(a):使用提出的 FP4 方案训练 Llama2 7B 的训练损失,该方案包括 NVFP4 格式(块大小 16,缩放格式 E4M3),并在神经梯度(更新+反向 GEMM)和激活(更新 GEMM)中使用 SR。注意到训练损失存在微小差距,但这并未导致下游任务性能出现显著差异(表 3)。(b):使用 NVFP4 格式在前向 GEMM 中、BF16 在反向和更新 GEMM 中进行 Llama2 7B 的量化感知微调(QAF)训练损失。注意到这个简短的 QAF 能够完全弥补与 BF16 基线的差距。

    表 2:不同 FP4 训练工作的比较。 DGE 指可微分梯度估计器,OCC 指离群点裁剪补偿,SR 指随机舍入,RHT 指随机哈达玛变换。虽然之前的工作仅将部分 GEMM 量化到 FP4,并带有残差稀疏矩阵(OCC [21])或哈达玛变换(RHT [19])的额外开销,但我们的工作是第一个展示将所有 GEMM 操作数都量化到 FP4,且没有增加显著开销的工作。
    表 2:不同 FP4 训练工作的比较

  • 量化感知微调 (QAF):

    • 实验内容: 为了弥补 FP4 训练与 BF16 基线之间微小的损失差距,引入了一个简短的 QAF 阶段。在此阶段,前向传播保持 FP4 精度(以确保模型与低精度推理兼容),而反向传播切换到 BF16 精度,以提高梯度信噪比。
    • 实验结果: 如图 6b 所示,这个简短的 QAF 阶段可以完全消除与 BF16 基线的损失差距。在整个训练(FP4 FQT + QAF)过程中,GEMM 的平均比特数为 4.3 bits。附录表 5 的消融研究表明,所需 QAF 的时长比例随着数据集的增大而减少。
    • 分析结论: QAF 是一种有效的方法,可以在不影响推理效率的前提下,弥补全 FP4 训练带来的微小性能损失。
  • 零样本性能评估:

    • 实验内容: 在多个下游任务上比较了 BF16 基线模型、FP4 训练后模型(1T tokens)以及 QAF 后模型(+40B tokens)的零样本准确率和困惑度。
    • 实验结果: 如表 3 所示,FP4 训练后的模型在部分任务上与 BF16 基线存在微小差距。然而,经过 QAF 阶段后,这个差距被完全弥补。
    • 分析结论: QAF 后的 FP4 模型在下游任务上的性能与 BF16 基线相当,证明了 FP4 全量化训练在实际应用中的可行性,其效率提升并未以牺牲模型质量为代价。

    表 3:BF16 基线与提出的 FP4 模型在 FP4 训练阶段(1T)和 QAF 阶段后(+40B)的零样本准确率和困惑度比较。 HS 指 HellaSwag,WG 指 Winogrande,BQ 指 BoolQ,AC 指 Arc-C,PQ 指 PiQA,LA 指 Lambada,GQ 指 GPQA,IE 指 If Eval,MB 指 MBPP,TQ 指 TrivialQA,XS 指 XSum。注意到在 QAF 之后,提出的 FP4 模型取得了与 BF16 基线相当的结果。
    表 3:零样本准确率和困惑度比较

A5 结论

  • 核心结论: 本文首次展示了在大型语言模型(Llama2 7B)上进行的全量化 FP4 训练,覆盖了权重、激活和梯度。实验证明,虽然与 BF16 相比初始训练损失有微小差距,但通过一个简短的量化感知微调(QAF)阶段(其中前向传播保持 FP4,反向传播切换到 BF16),可以完全弥补这一差距,并且在下游任务上达到与 BF16 相当的性能,证实了 FP4 的实际可行性。

  • 关键发现与贡献:

    1. FP4 格式设计: 研究发现 NVFP4 格式(E4M3 缩放,块大小 16)在动态范围和精度之间提供了最佳权衡,这与 NVIDIA Blackwell 的硬件设计决策一致。
    2. 分离式舍入策略: 提出了一种分离式舍入策略,即仅在反向传播中使用随机舍入,这显著提高了训练的稳定性。
    3. 理论指导: 理论分析确定了一个临界点:当全精度梯度的标准差降至量化噪声的 3 倍以下时,训练会停滞。这一发现为设计 QAF 阶段提供了理论依据,通过提高信噪比来匹配 BF16 的收敛效果。
  • 局限性与未来工作:

    • 硬件限制: 由于当前的 Gaudi 硬件缺乏对 FP4 的专门支持,所有实验都是通过 FP4 模拟进行的,这无法直接衡量原生 FP4 执行带来的速度提升和能效优势。基于 FP8 的经验粗略估计,相比 FP8,训练时间可加速约 35-40%,相比 BF16 则可加速约 85%。
    • 架构范围: 本工作主要集中在 LLaMA 架构上。初步实验表明该方法可直接应用于专家混合(MoE)架构,但对 MoE 模型和视觉任务的进一步分析将留待未来工作。

A6 附录

A 更广泛的影响

加速 LLM 运行时的重要性。 加速大型语言模型(LLM)的运行时间在塑造现代数字体验中扮演着关键角色,特别是随着像 ChatGPT 和 Gemini 这样的系统日益融入日常应用。提高它们的速度和效率解决了大规模 AI 部署中的一个主要瓶颈。除了提升性能和减少内存需求外,更快的模型还扩大了可及性,使更多用户能够根据自身特定需求调整和创新 LLM。然而,这种增加的可及性也引发了对潜在滥用的担忧,强调了负责任开发和监督的必要性。

B 量化 SGD 的分析

B.1 对使用随机舍入的量化 SGD 的分析

研究目标。 我们研究量化噪声在梯度下降过程中如何影响期望的损失减少。

前提。 设 $L(θ)$ 是一个在 $R^d$ 上的二次可微标量损失函数。

步骤 1:从损失的泰勒展开开始。 我们考虑从 $θ_t$ 开始的一个小步长 $∆θ$。L 在 $θ_t$ 的二阶泰勒展开为:
$L(\theta_t + \Delta \theta) = L(\theta_t) + \nabla L(\theta_t)^T \Delta \theta + \frac{1}{2} \Delta \theta^T H(\theta_t) \Delta \theta + \cdots$
其中 $H(θ_t) = ∇^2L(θ_t)$。

步骤 2:应用量化梯度更新。 由于量化,我们使用梯度的带噪估计:
$g_q = \nabla L(\theta_t) + \varepsilon,$
其中 $ε$ 是量化噪声。更新规则变为:

$$\theta_{t+1} = \theta_t - \eta g_q \Rightarrow \Delta\theta = -\eta g_q.$$

步骤 3:代入泰勒展开。 将 $∆θ = -η g_q$ 代入得:
$L(\theta_{t+1}) \approx L(\theta_t) - \eta \nabla L(\theta_t)^T g_q + \frac{1}{2}\eta^2 g_q^T H(\theta_t) g_q.$

步骤 4:对量化噪声取期望。 在随机舍入下,噪声 $ε$ 满足:
$\mathbb{E}[\varepsilon]=0, \quad \mathbb{E}[\varepsilon \varepsilon^T]=\sigma_q^2 I.$
在泰勒展开中取期望得到:

$$ \mathbb{E}[L(\theta_{t+1})] \approx L(\theta_t) - \eta \nabla L(\theta_t)^T \mathbb{E}[g_q] + \frac{1}{2}\eta^2 \mathbb{E}[g_q^T H(\theta_t) g_q] $$


因为:
$\mathbb{E}[g_q] = \mathbb{E}[\nabla L(\theta_t) + \epsilon] = \nabla L(\theta_t),$
线性项简化为:

$$-\eta \nabla L(\theta_t)^T \nabla L(\theta_t) = -\eta \|\nabla L(\theta_t)\|^2.$$
接下来:
$$\mathbb{E}[g_{q}g_{q}^{T}] = \mathbb{E}[(\nabla L + \varepsilon)(\nabla L + \varepsilon)^{T}] = \nabla L \nabla L^{T} + \sigma_{q}^{2}I,$$
所以:
$$\mathbb{E}[g_q^T H(\theta_t) g_q] = \mathrm{tr}(H(\theta_t) \mathbb{E}[g_q g_q^T]) = \nabla L(\theta_t)^T H(\theta_t) \nabla L(\theta_t) + \sigma_q^2 \mathrm{tr}(H(\theta_t)).$$
综上所述:
$$ \mathbb{E}[L(\theta_{t+1})] = L(\theta_t) - \eta \|\nabla L(\theta_t)\|^2 + \frac{1}{2}\eta^2 \left( (\nabla L(\theta_t))^T H(\theta_t) \nabla L(\theta_t) + \sigma_q^2 \mathrm{tr}(H(\theta_t)) \right). $$

步骤 5:SR 的收敛动态。 从步骤 4,期望的损失变化为:

$$\mathbb{E}[L(\theta_{t+1}) - L(\theta_t)] \approx -\eta\|\nabla L(\theta_t)\|_2^2 + \frac{1}{2}\eta^2 \left( \nabla L(\theta_t)^T H(\theta_t)\nabla L(\theta_t) + \sigma_q^2 \mathrm{tr}(H(\theta_t)) \right).$$


这可以写成:

$$\mathbb{E}[L(\theta_{t+1}) - L(\theta_t)] \approx - \left( \eta\|\nabla L(\theta_t)\|_2^2 - \frac{1}{2}\eta^2 \nabla L(\theta_t)^T H(\theta_t)\nabla L(\theta_t) \right) + \frac{1}{2}\eta^2 \sigma_q^2 \mathrm{tr}(H(\theta_t))$$
当满足以下条件时,有用的下降分量为负:
$$ \eta\|\nabla L(\theta_t)\|_2^2 > \frac{1}{2}\eta^2\nabla L(\theta_t)^T H(\theta_t)\nabla L(\theta_t) \Rightarrow \eta < \frac{2\|\nabla L(\theta_t)\|_2^2}{\nabla L(\theta_t)^T H(\theta_t)\nabla L(\theta_t)}. $$
一个更保守的条件是:
$$\eta < \frac{2}{\lambda_{\max }(H(\theta_t))}.$$

步骤 6:最优步长 $η^∗$。 为了找到最优步长,我们定义:

$$U(\eta)=\mathbb{E}[L(\theta_{t+1}) - L(\theta_t)] = -\eta\|\nabla L(\theta_t)\|_2^2 + \frac{1}{2}\eta^2\left(\nabla L(\theta_t)^T H(\theta_t)\nabla L(\theta_t)\right) + \sigma_q^2\text{tr}(H(\theta_t)).$$


将其导数设为 0:

$$\frac{dU}{d\eta} = -\|\nabla L(\theta_t)\|_2^2 + \eta\left(\nabla L(\theta_t)^T H(\theta_t)\nabla L(\theta_t) + \sigma_q^2\text{tr}(H(\theta_t))\right) = 0.$$
解出 $η^∗$:
$$\eta^{*} = \frac{\left\|\nabla L\left(\theta_{t}\right)\right\|_{2}^{2}}{\nabla L\left(\theta_{t}\right)^{T} H\left(\theta_{t}\right) \nabla L\left(\theta_{t}\right)+\sigma_{q}^{2} \operatorname{tr}\left(H\left(\theta_{t}\right)\right)}$$

步骤 7:使用最优步长 $η^∗$ 进行训练。 将 $η^∗ = \frac{A}{B}$ 代入:

$$U(\eta)=-\eta A+\frac{1}{2} \eta^{2} B,$$


其中:

$$A = \|\nabla L(\theta_t)\|_2^2, \quad B = \nabla L(\theta_t)^T H(\theta_t) \nabla L(\theta_t) + \sigma_q^2 \text{tr}(H(\theta_t)).$$
那么:
$$U(\eta^*) = -\frac{A}{B}A + \frac{1}{2}\left(\frac{A}{B}\right)^2B = -\frac{A^2}{B} + \frac{1}{2}\frac{A^2}{B} = -\frac{1}{2}\frac{A^2}{B},$$
即:
$$U(\eta^*) = - \frac{\left\| \nabla L(\theta_t) \right\|_2^4}{2 \left( \nabla L(\theta_t)^T H(\theta_t) \nabla L(\theta_t) + \sigma_q^2 \operatorname{tr}(H(\theta_t)) \right)}.$$
设 $X = ∇L^T H∇L$, $Y = tr(H)$, $Z = ∥∇L∥_2^4$。则:
$$U(\eta^{*})=-\frac{Z}{2(X+\sigma_{q}^{2}Y)}$$

步骤 8:对噪声的最大敏感度。 我们从敏感度开始:

$$f(\sigma_q) = \frac{\partial U(\eta^*)}{\partial \sigma_q} = \frac{ZY \sigma_q}{(X + Y \sigma_q^2)^2}.$$


为了找到其最大值,计算关于 $σ_q$ 的导数:

$$\frac{df}{d\sigma_q} = ZY \frac{(X + Y\sigma_q^2)^2 - \sigma_q \cdot 2(X + Y\sigma_q^2)(2Y\sigma_q)}{(X + Y\sigma_q^2)^4}.$$
我们分解并简化这个表达式:
$$\frac{df}{d\sigma_q} = \frac{ZY(X+Y\sigma_q^2)[(X+Y\sigma_q^2)-4Y\sigma_q^2]}{(X+Y\sigma_q^2)^4} = \frac{ZY[X-3Y\sigma_q^2]}{(X+Y\sigma_q^2)^3}.$$
将其设为零以找到最小值:
$$X - 3Y \sigma_{q}^{2} = 0 \implies \sigma_{q}^{2} = \frac{X}{3Y}.$$
因此,临界噪声水平是:
$\sigma_{\text{critical}}^2 = \frac{X}{3Y}$
即:
$$\sigma_{\text{critical}}^2 = \frac{\nabla L(\theta_t)^T H(\theta_t) \nabla L(\theta_t)}{3 \operatorname{tr}(H(\theta_t))}.$$
最后,假设:
$$\frac{\nabla L(\theta_t)^T H(\theta_t) \nabla L(\theta_t)}{\|\nabla L(\theta_t)\|_2^2} \approx \frac{\text{tr}(H(\theta_t))}{d},$$
我们得到:
$\|\nabla L(\theta_t)\|_2^2 \approx 3d\sigma_{\text{critical}}^2 \implies \|\nabla L(\theta_t)\|_2 \approx \sqrt{3d}\sigma_{\text{critical}}.$
因此:
$$\sigma_{\text{critical}} = \frac{\left\|\nabla L(\theta_t)\right\|_2}{\sqrt{3d}}.$$
换句话说,一旦平均每个坐标的梯度下降到量化噪声标准差的 3 倍,FP4 梯度就失去效力,是时候切换到更高精度了。如图 4 所示,将噪声标准差设置为 $k · σ_{critical}$,其中 k = 2.0, 1.0, 0.5,在该阈值附近产生了显著不同的收敛行为。在附录 B.2 中,我们展示了没有 SR 的类似分析,该分析表明,随着我们训练,“有用的下降”逐渐消失为零,而有偏噪声项即使经过长时间训练仍然存在。

B.2 非零均值噪声(确定性舍入)的影响

有偏量化方案的问题。 在本节中,我们通过一个简单的标量优化问题来举例说明有偏量化方案(如 RtN)的问题,该问题具有二次损失:

$$L(\theta) = \frac{1}{2}\lambda(\theta-\theta^{*})^{2} \implies \nabla L(\theta) = \lambda(\theta-\theta^{*}),$$


和一个带有量化噪声 $ε$ 的步长更新,其均值为 $µ_ε = E[ε] \neq 0$:

$$\theta_{t+1}=\theta_t-\eta\left(\nabla L\left(\theta_t\right)+\varepsilon\right) \implies \mathbb{E}\left[\theta_{t+1}\right]=\mathbb{E}\left[\theta_t\right]-\eta\left(\lambda\left(\mathbb{E}\left[\theta_t\right]-\theta^*\right)+\mu_{\varepsilon}\right).$$
定义误差:
$e_t \triangleq \mathbb{E}[\theta_t] - \theta^*$
那么:
$e_{t+1}=\mathbb{E}[\theta_{t+1}]-\theta^* = [\mathbb{E}[\theta_t] - \eta(\lambda e_t + \mu_\varepsilon)] - \theta^* = e_t - \eta\lambda e_t - \eta\mu_\varepsilon.$
因此:
$e_{t+1} = (1 - \eta\lambda) e_t - \eta \mu_{\varepsilon}.$
展开这个递归,对于 $a = 1 - ηλ$:
$$e_n = a^n e_0 - \eta \mu_\varepsilon \sum_{k=0}^{n-1} a^k.$$
由于 $\sum_{k=0}^{n-1} a^k = \frac{1-a^n}{1-a}$ 且 $1 - a = ηλ$,我们得到:
$e_n = a^n e_0 - \frac{\eta \mu_\varepsilon}{\eta \lambda} (1 - a^n) = a^n e_0 - \frac{\mu_\varepsilon}{\lambda} (1 - a^n).$
第 n 步的损失是:
$L_n = L(\mathbb{E}[\theta_n]) = \frac{1}{2}\lambda e_n^2 = \frac{\lambda}{2}\left(a^n e_0 - \frac{\mu_\varepsilon}{\lambda}(1-a^n)\right)^2.$
当 $n → ∞$,$a^n → 0$(对于 $a < 0$,这是成功优化所必需的),得到稳态误差和残余损失:
$e_\infty = -\frac{\mu_\varepsilon}{\lambda}, \quad L_\infty = \frac{\mu_\varepsilon^2}{2\lambda}.$
因此,有偏的 SGD 并不会收敛到 $θ^∗$ 且损失为零,而是稳定在:
$\mathbb{E}[\theta_\infty] = \theta^* - \frac{\mu_\varepsilon}{\lambda},$
并留下一个残余损失:
$$L(\mathbb{E}[\theta_{\infty}]) = \frac{\mu_{\varepsilon}^2}{2\lambda}.$$

C 额外的实验结果

图 7:使用 NVFP4 格式训练 350M Llama 模型时不同舍入方案的比较。在每个图中,我们将 RtN 应用于其中一个 GEMM 的六个元素之一,而其余部分使用 SR。注意到,在‘更新’和‘反向’ GEMM 中对神经梯度应用 RtN,以及在‘更新’ GEMM 中对激活应用 RtN,会导致更高的训练损失,而对其他组件应用 RtN 则会产生相反效果,降低损失。
图 7:使用 NVFP4 格式训练 350M Llama 模型时不同舍入方案的比较。在每个图中,我们将 RtN 应用于其中一个 GEMM 的六个元素之一,而其余部分使用 SR。注意到,在‘更新’和‘反向’ GEMM 中对神经梯度应用 RtN,以及在‘更新’ GEMM 中对激活应用 RtN,会导致更高的训练损失,而对其他组件应用 RtN 则会产生相反效果,降低损失。

表 4:在 30B tokens 上使用不同种子训练 Llama 125M 的训练损失,得到相似的损失,标准差为 0.001。
表 4:不同种子下的训练损失

表 5:确定需要多少 tokens 才能达到与 BF16 基线相同损失的消融研究。注意到对于更大的数据集,QAF 的比例会下降。在所有 QAF 实验中,我们使用的峰值学习率等于 FP4 训练中的最后一个学习率。
表 5:QAF 长度消融研究