UE5M3 FP4 Block Scaling for Stable Language Model Pretraining

TL;DR

提出UE5M3 FP4块尺度方案,结合周期性张量缩放,优化语言模型预训练,模型达190B参数,训练损失优于NVFP4。

cs.LG 🔴 高级 2026-09-03 81 次浏览
Robert Hu Carlo Luschi Paul Balanca
FP4量化 大规模预训练 块尺度 张量缩放 量子化算法

核心发现

方法论

本文提出一种基于UE5M3无符号块尺度的FP4预训练方案,结合周期性张量缩放和选择性随机舍入,省略RHT变换,全部内部线性层采用FP4。采用多阶段预训练,模型规模达8B参数,训练近1900亿tokens。通过对比NVFP4方案,验证其在训练损失和验证负对数似然上的优越性。引入块尺度范围界限,结合定期张量最大值刷新机制,有效控制尺度饱和。实现软件模拟的端到端UEFP4预训练,简化流程,提升吞吐量。

关键结果

  • 所提出方案在Nemotron-H 8B模型上,训练近1900亿tokens后,最终训练窗口损失低于NVFP4方案,验证集负对数似然指标更优,且量化推理的点估计值在三项指标中均优于NVFP4。
  • 在原生NVFP4执行中,联合移除RHT和BF16最终块豁免,模型的每秒处理Token数提升21.2%,验证性能保持或改善。
  • 块尺度范围的扩展(T从448到2048)显著提升了小梯度的表示能力,减少下溢和饱和事件,验证了UE5M3尺度的优势。

研究意义

该研究突破了FP4预训练的稳定性瓶颈,提出更简洁高效的块尺度方案,降低实现复杂度,推动FP4在大规模语言模型中的应用。其端到端软件模拟验证了UE5M3方案的可行性,为未来原生硬件支持提供理论基础,具有重要的工业和学术价值。

技术贡献

技术创新在于引入UE5M3无符号块尺度,结合周期性张量缩放和选择性随机舍入,省略了复杂的RHT和全局缩放机制,简化了FP4训练流程。提出软件模拟的端到端预训练流程,验证了在不依赖BF16的情况下,模型性能优于NVFP4方案,展示了块尺度范围扩展带来的稳定性提升。

新颖性

本研究首次系统性引入UE5M3无符号块尺度,结合周期性张量缩放,显著扩大动态范围,简化预训练流程,突破了FP4训练的稳定性限制。与现有NVFP4方案相比,省略了RHT和BF16豁免,验证了更宽范围块尺度的实用性,推动FP4硬件友好型预训练技术的发展。

局限性

  • 目前方案依赖软件模拟,尚未在硬件原生支持下验证性能,存在一定的实现差异风险。
  • 块尺度范围的调整需精细调节目标参数T,可能影响训练的稳定性与收敛速度,需进一步自动化调优策略。
  • 模型在极端尺度变化或特殊任务中仍可能出现饱和或下溢,未来需结合动态尺度调节机制进一步优化。

未来方向

未来将结合硬件原生支持,优化块尺度调节策略,探索更高阶的量子化方案,提升训练效率和模型性能。同时,扩展到更大模型规模和多任务场景,验证方案的普适性和鲁棒性,为FP4在实际工业应用中铺平道路。

AI 总览摘要

随着大规模语言模型的发展,低精度量化技术成为提升训练效率的关键。FP4格式因其极低的存储和计算成本而备受关注,但其动态范围有限,导致训练不稳定。本文提出一种基于UE5M3无符号块尺度的FP4预训练方案,结合周期性张量缩放和选择性随机舍入,有效缓解了尺度饱和和下溢问题。通过在Nemotron-H 8B模型上进行近1900亿tokens的训练,结果显示新方案在训练损失和验证指标上优于NVFP4方案,且推理点估计值更高。关键在于扩展块尺度范围(T从448到2048),提升了小梯度的表示能力,减少了量化误差。实验中,联合移除RHT和BF16豁免,模型吞吐量提升21.2%,验证了方案的实用性和高效性。该研究简化了FP4预训练流程,为未来硬件原生支持提供理论基础,推动低精度大模型训练的产业化。未来工作将结合硬件支持,优化尺度调节策略,扩大模型规模,验证方案的广泛适用性。整体来看,本文在低比特量化训练领域实现了技术突破,为大规模高效预训练提供了新思路。

深度分析

研究背景

近年来,深度学习模型规模不断扩大,推动了对高效低精度训练的需求。FP8、INT4等格式已在硬件中实现,但FP4因其极低的存储成本和潜在的计算优势受到关注。早期研究如Sun等的INT4梯度方案和Chmiel的无偏舍入技术,解决了梯度范围和偏差问题。NVIDIA的Transformer Engine引入RHT、二维缩放和BF16层,提升了FP4训练的稳定性,但复杂度较高。近年来,学界开始探索无符号块尺度(UE5M3)和周期性张量缩放,试图简化流程同时保持性能。尽管如此,FP4在大规模预训练中的稳定性仍是挑战,尤其在动态范围有限的情况下。本文在此背景下提出更宽范围的块尺度方案,旨在突破现有瓶颈。

核心问题

FP4预训练面临的核心问题是动态范围有限,导致激活、权重和梯度在训练过程中易饱和或下溢,影响模型收敛和性能。现有方案如RHT和BF16层虽能缓解,但增加了实现复杂度和计算开销。如何在简化流程的同时,确保训练的稳定性和模型性能,是亟待解决的问题。此外,块尺度范围的限制限制了梯度的表达能力,影响模型的训练效果。本文旨在通过引入更宽的块尺度(UE5M3)和周期性张量缩放,解决动态范围不足的问题,提升FP4训练的稳定性和效率。

核心创新

本研究的核心创新在于引入UE5M3无符号块尺度,结合周期性张量缩放机制,省略了复杂的RHT变换和BF16豁免,简化了FP4预训练流程。具体包括:• 采用每50步刷新一次的张量最大值机制,控制尺度饱和;• 利用更宽的块尺度范围(T从448提升至2048),增强小梯度的表达能力;• 仅对上游梯度采用选择性随机舍入,减少偏差;• 全部内部线性层采用FP4,降低硬件实现难度。这些创新使得训练流程更为简洁高效,同时保持甚至超越现有方案的性能。

方法详解

  • �� 设计UE5M3无符号块尺度,结合周期性张量最大值刷新机制,控制尺度饱和风险;
  • �� 采用样本保持策略,每50步更新一次全局最大值,避免频繁计算,提高效率;
  • �� 通过调节目标尺度T(默认448,特殊层2048),扩展块尺度范围,提升小梯度表示能力;
  • �� 只在上游梯度中使用随机舍入,其他操作采用确定性舍入,以减少偏差;
  • �� 全部内部线性层采用FP4,省略RHT和BF16层,简化实现;
  • �� 在模型训练中,采用多阶段预训练策略,监控训练损失和验证指标,进行逐步调优。

实验设计

在Nemotron-H 8B模型上,训练近1900亿tokens,采用12个检查点,比较不同尺度目标T(448和2048)对训练损失和验证性能的影响。通过软件模拟的端到端预训练流程,验证了方案在模型收敛性和推理性能上的优越性。对比NVFP4方案,本文方案在最终训练窗口损失和验证负对数似然上均表现更优,吞吐量提升显著。还进行了原生NVFP4执行的消融实验,联合移除RHT和BF16豁免,模型吞吐量提升21.2%。此外,调节块尺度范围,验证了更宽范围带来的梯度表达优势。

结果分析

实验结果显示,采用UE5M3块尺度方案后,模型在训练损失和验证指标上均优于NVFP4方案,验证集负对数似然降低,点估计值更高。扩大尺度范围(T从448到2048)显著改善了梯度表达,减少了下溢和饱和事件,验证了块尺度扩展的有效性。消融实验中,去除RHT和BF16豁免后,模型吞吐量提升21.2%,说明方案在硬件实现和效率方面具有潜力。整体结果证明了宽范围块尺度在FP4预训练中的优势,为未来硬件支持提供了理论基础。

应用场景

该方案适用于大规模语言模型的高效预训练,特别是在硬件支持有限的场景。可应用于工业界的模型训练,加速模型开发周期,降低算力成本。未来,结合硬件原生支持,将实现更高效的低比特量化训练流程,推动大模型的普及。

局限与展望

目前方案仍依赖软件模拟,硬件原生支持尚未实现,存在性能差异风险。尺度调节参数需手动调优,可能影响训练稳定性。极端尺度变化可能引发饱和或下溢,需进一步研究动态尺度调节机制。未来需在硬件层面验证方案效果,并优化自动调参策略。

通俗解读 非专业人士也能看懂

想象你在厨房里做饭,使用不同大小的锅来煮不同的食材。小锅代表小的数值,能装少量食材,但不适合大火或大份量。大锅则可以装更多食材,但需要调节火力,否则容易烧焦或煮不熟。传统做法用一个锅(全局尺度)来处理所有食材,但这样容易出错,特别是当食材差异很大时。本文提出一种新方法,像是给每个锅配上可调节的火力(块尺度),还能定期检查火力大小(周期性张量缩放),确保每个食材都能煮得恰到好处。这样,不仅做饭更快,还能保证味道一致。通过这种方式,厨房里的“火力调节”变得更智能、更简单,最终做出更美味的菜肴(模型更好、更稳定)。

简单解释 像给14岁少年讲一样

想象你在玩一个超级复杂的拼图游戏,你需要把很多碎片拼在一起,拼得越快越好。可是,有时候碎片太小,拼起来会很困难,因为它们容易掉或拼错。传统的方法就像用一个大框,把所有碎片都放进去,但如果碎片差异太大,就会出现问题。这个研究就像发明了一种新型的拼图框,可以根据碎片的大小自动调节,把每个碎片都放得刚刚好。它还能定期检查碎片的大小,确保没有碎片掉出来或拼错。这样一来,拼图变得更快、更稳,也更容易完成。科学家用这种方法训练了一个超级大的语言模型,模型表现比以前更好,训练也更高效。就像用这个新拼图框,拼图变得更快更漂亮,未来还能用在很多其他拼图游戏中!

原文摘要

Stable 4-bit floating-point (FP4) pretraining is difficult because the E2M1 payload represents only a narrow range of magnitudes. NVIDIA's Transformer Engine \nv{} recipe addresses this with current-tensor scaling, a randomized Hadamard transform (RHT), and bfloat16 (BF16) final layers, adding work outside the FP4 matrix multiplications. We instead pair E2M1 payloads with unsigned E5M3 (\ue{}) block scales. Their wider range permits periodic tensor scaling, while our recipe applies selective stochastic rounding to backward gradients, omits RHT, and uses FP4 in all eligible internal linears. We pretrain a Nemotron-H 8B model for nearly 190 billion tokens. Compared with Transformer Engine \nv{}, the proposed block-16 recipe finishes with lower final-window training loss and, under their respective quantized-inference policies, lower validation loss measured as held-out negative log-likelihood. Its quantized-inference downstream point estimates are also higher on all three reported aggregates. A native \nv{} execution ablation that jointly removes RHT and the BF16 final-block exemption increases measured model-body token throughput by 21.2\%. These results demonstrate end-to-end software-emulated \uefp{} pretraining with a simpler recipe and motivate native support for \ue{} block scaling.

cs.LG