重新思考 LLM FP4 预训练中的收缩偏差:几何起源、系统性影响与 UFP4 配方 Rethinking Shrinkage Bias in LLM FP4 Pretraining: Geometric Origin, Systemic Impact, and UFP4 Recipe
揭示 E2M1 格式的 Shrinkage Bias 问题,提出均匀网格 UFP4 配方改善 FP4 训练质量
前置知识
FP4 格式(E2M1/E1M2)
FP4 是 4-bit 浮点格式,用 1 位符号位,剩余 3 位分为指数和尾数。E2M1 用 2 位指数和 1 位尾数,非负值为 {0, 0.5, 1, 1.5, 2, 3, 4, 6},区间宽度不均匀(如 0 到 0.5 宽 0.5,但 1.5 到 2 宽 0.5,2 到 3 宽 1)。E1M2 用 1 位指数和 2 位尾数,值均匀分布。这种几何差异导致量化误差特性不同。
本文核心发现 E2M1 的非均匀区间几何导致系统性舍入偏差(Shrinkage Bias),理解这种格式差异是读懂论文的基础。
Random Hadamard Transform (RHT)
随机 Hadamard 变换是一种正交变换,通过将 Hadamard 矩阵 H 与随机符号矩阵 S 结合,保持张量的 L2 范数不变但分散能量。数学表达为 H' = S·H,其中 H 是 Sylvester Hadamard 矩阵。对张量 X 施加 RHT 后再量化,可以将异常值能量分散到所有维度,提高码本利用率。公式:Y = XWᵀ = (XH')(WH')ᵀ。
本文发现 RHT 将张量从动态范围受限变为局部分辨率受限,这改变了量化瓶颈,导致 E2M1 和 E1M2 的相对性能发生逆转。
Round-to-Nearest-Even (RTNE)
RTNE 是一种确定性舍入规则,当值恰好位于两个量化级别中间时,舍入到偶数级别。对于 FP4 码本 G = {g₀, g₁, ..., gₙ},区间 Bᵢ = [(gᵢ₋₁+gᵢ)/2, (gᵢ+gᵢ₊₁)/2]。当输入 t ∈ Bᵢ 时,量化为 gᵢ。E2M1 的区间宽度 ℓᵢ = (gᵢ-gᵢ₋₁)/2 和 rᵢ = (gᵢ₊₁-gᵢ)/2 不相等,导致系统性偏差。
Shrinkage Bias 的几何来源正是 RTNE 区间的不对称性,理解 RTNE 才能理解为什么均匀网格不存在这种偏差。
SQNR(Signal-to-Quantization-Noise Ratio)
SQNR 衡量量化质量,定义为 10·log₁₀(NMSE_raw/NMSE_RHT),其中 NMSE = ∥Q(T)-T∥²/∥T∥²。正值表示 RHT 提高了保真度,负值表示降低。例如论文中 E2M1 在 linear_fc2/fwd_x 上从 21.90 dB 降至 20.00 dB(ΔSQNR = -1.90 dB),而 E1M2 从 19.94 dB 提升至 23.19 dB(ΔSQNR = +3.24 dB)。
这是论文中量化质量的主要评估指标,用于展示 RHT 在不同格式下的相反效果。
Blockwise Quantization(块量化)
块量化将张量 T 分割为连续块 {B},每个块共享一个缩放因子 sᵦ = maxₓⱼ∈B|xⱼ|/g_max。元素 xᵢ 被归一化并舍入到码本级别:qᵢ = ρ_G(xᵢ/sᵦ)。块大小影响动态范围和精度的权衡。论文使用 1×16 块大小,与 RHT 块大小对齐以支持融合内核。
这是所有现代 FP4 训练配方(包括 NVFP4)的基础,理解块量化才能理解论文的实验设置和硬件实现。
研究动机
当前主流的 FP4 训练配方(如 NVIDIA 的 NVFP4)都基于 E2M1 格式,但在实际训练中遇到收敛问题和损失退化问题。例如,在 Dense 1.5B 模型上,E2M1 配方的 BF16 相对损失误差达到 1.2570%,MoE 7.9B 达到 2.3596%,MoE 124B 达到 1.7308%。虽然 NVFP4 引入 RHT 和随机舍入来缓解,但问题仍然存在。论文指出,这是因为 E2M1 的非均匀量化区间存在几何不对称性,在 RTNE 舍入下产生系统性负误差(Shrinkage Bias)。例如,对于 E2M1 的量化级别 gᵢ = 2,其舍入区间为 (1.75, 2.5),左半宽度 ℓᵢ = 0.25,右半宽度 rᵢ = 0.5。当区间内密度均匀时,条件期望误差为 (ℓᵢ - rᵢ)/4 = -0.125,即数值被系统性地缩小。
本文的目标是本文的目标是重新审视 E2M1 作为 FP4 训练默认格式的合理性,并提出更优的配方。具体来说,论文希望:(1) 从理论上形式化 Shrinkage Bias 的几何来源和系统性传播机制;(2) 解释为什么 RHT 在 E2M1 格式下会放大而非缓解训练不稳定性;(3) 基于 E1M2/INT4 均匀网格设计新的训练配方 UFP4,验证其在 Dense 1.5B、MoE 7.9B 和 MoE 124B 长期预训练中的优势;(4) 推动未来硬件加速器将 E1M2/INT4 等均匀网格作为一流的训练原语。
与已有工作不同的是,本文的独特切入点是从格式网格的几何性质出发,而不是从分布特性或算法技巧入手。现有工作主要在三个方向改进:(1) 改进数据格式和块缩放设计(如 NVFP4 的两级缩放);(2) 改进量化器或训练感知方法(如 Microsoft FP4 的可微量化估计器、FAAR 的格式感知自适应舍入);(3) 张量预处理方法(如 RHT、QuaRot、SpinQuant)。这些方法都假设 E2M1 是唯一可行的 FP4 格式,试图通过算法技巧来补偿其几何缺陷。本文则指出,问题的根源在于 E2M1 的网格几何本身,一旦 RHT 将张量转变为局部分辨率受限状态,均匀网格就比非均匀网格更合适。这是首次从网格几何角度系统分析 FP4 训练稳定性问题的工作。
核心方法
本文采用理论分析与实证验证相结合的方法。首先,论文从量化区间的几何不对称性出发,推导 Shrinkage Bias 的数学形式。对于量化级别 gᵢ,其 RTNE 舍入区间为 Bᵢ = [(gᵢ₋₁+gᵢ)/2, (gᵢ+gᵢ₊₁)/2],左宽度 ℓᵢ = (gᵢ-gᵢ₋₁)/2,右宽度 rᵢ = (gᵢ₊₁-gᵢ)/2。在区间密度均匀的假设下,条件期望误差为 (ℓᵢ - rᵢ)/4。当 rᵢ > ℓᵢ 时,出现负误差,即 Shrinkage Bias。论文进一步证明这种偏差在网络中通过乘积累积传播,对于 K 个量化的 GEMM,累积衰减因子约为 exp(-∑δₖ),其中 δₖ = 1 - ηₖ 是每次 GEMM 的乘积损失。即使每次 δₖ 很小,累积效应也会导致可见的信号衰减。其次,论文分析 RHT 如何恶化 E2M1 的偏差。RHT 通过分散异常值能量,将张量从动态范围受限转变为局部分辨率受限,使数据聚集到 E2M1 最不对称的区间,加剧量化噪声。最后,基于均匀网格不存在几何偏差的洞察,论文提出 UFP4 配方并在多个模型规模上验证。
本文的核心创新点是揭示了量化网格几何性质与训练稳定性的深层联系。传统观点认为 E2M1 的宽动态范围适合处理异常值,但论文指出,在 RHT 将张量转变为局部分辨率受限状态后,关键不再是覆盖极端异常值,而是精确保存典型幅度的密集概率质量。此时,E2M1 的非均匀区间反而成为劣势,而均匀网格(E1M2/INT4)能更好地将 RHT 带来的平化分布转化为更高的保真度。另一个关键洞察是,Shrinkage Bias 不是零均值的噪声,而是系统性衰减,在深度网络中通过乘积累积。对于序列的 K 个量化 GEMM,初始干净信号被累积缩放为 ∏ηₖ ≈ exp(-∑δₖ),这与零均值噪声的相互抵消行为本质不同。这两个洞察共同构成了 UFP4 的理论基础:使用均匀网格消除几何偏差源,从而安全地将 RHT 扩展到所有三个训练 GEMM 而不累积系统性衰减。
方法步骤详情
论文的方法步骤分为三个部分。第一部分是理论分析:(1) 定义 Shrinkage Bias 为 RTNE 舍入在归一化幅度空间的负期望误差,公式为 b_G(P) = E_{t∼P}[ρ_G(t) - t];(2) 推导单个区间的不对称性导致的条件期望误差,公式为 (ℓᵢ - rᵢ)/4;(3) 证明量化 GEMM Z = ABᵀ 的信号衰减机制,通过将量化操作数投影到精确 BF16 对应物上,定义缩放因子 α_A = ⟨b̃A, A⟩_F/∥A∥²_F,衰减因子为 η ≈ α_Aα_B < 1;(4) 推导多层累积衰减公式为 exp(-∑δₖ),其中 δₖ = 1 - ηₖ;(5) 分析 RHT 将张量从动态范围受限转变为局部分辨率受限的机制,通过有效存储桶比 B_eff(G,T) = exp(E(G,T))/K 来衡量码本利用率,其中 E(G,T) = -∑p_i log(p_i+ε)。第二部分是实证诊断:(1) 在真实的 MLP 层张量(fwd_w, fwd_x, bwd_dy)上测量 SQNR 和 B_eff,比较 E2M1 和 E1M2 在有无 RHT 下的表现;(2) 在单 GEMM 输出(fwd_y, bwd_dx, bwd_dw)上测量输出 SQNR,验证格式排名的逆转。第三部分是 UFP4 配方设计:(1) 选择 E1M2/INT4 风格的均匀网格;(2) 对所有三个训练 GEMM(FPROP, DGRAD, WGRAD)应用 RHT;(3) 仅对上游梯度 dY 使用随机舍入,保留梯度期望;(4) 保持块大小(1×16)、缩放层级(FP32 单级)等辅助配置与 E2M1 基线一致。
技术新颖性
本文的技术新颖性体现在多个层面。从问题识别角度,论文首次形式化了 Shrinkage Bias 概念,将其作为非均匀量化网格固有的几何性质导致的系统性误差。这不同于以往将量化误差视为分布引起的问题的认知。从理论贡献角度,论文证明了量化偏差的乘积累积机制,给出了累积衰减因子的指数近似公式 exp(-∑δₖ),并推导了二阶项 O(∑δₖ²) 可忽略的条件。从方法设计角度,论文首次在 FP4 训练中系统性比较非均匀(E2M1)和均匀(E1M2/INT4)网格在 RHT 下的表现差异,揭示了格式依赖的逆转现象:在无 RHT 时 E2M1 可能更好(如 linear_fc2/fwd_x 上 E2M1 为 21.90 dB vs E1M2 为 19.94 dB),但在 RHT 后 E1M2 反而更好(E2M1 降至 20.00 dB,E1M2 升至 23.19 dB)。从配方创新角度,UFP4 首次在 FP4 训练中实现了全 RHT 覆盖(三个 GEMM),而传统 NVFP4 配方为了避免复合几何误差而限制 RHT 仅用于 bwd_dw。论文还评估了融合 RHT+量化的效率,在 SM90 和 SM100 上分别仅增加 1.06× 和 1.07× 延迟,验证了实用性。
实验结果
论文通过一系列实验验证了均匀网格在 FP4 训练中的优势。在 Q1 实验中,论文在单张量量化和单 GEMM 输出层面证明了 RHT 改变了首选 4-bit 网格。对于表现良好的 linear_fc1/fwd_x 张量,RHT 影响中性:E1M2 的 ΔSQNR 平均为 -0.008 dB,E2M1 为 +0.007 dB,有效存储桶比变化微小。但对于包含异常值的 linear_fc2/fwd_x,RHT 导致格式排名逆转:旋转前 E2M1 领先(21.90 vs 19.94 dB),旋转后 E1M2 领先(23.19 vs 20.00 dB),且有效存储桶比从 0.56 提升至 0.97。同样的格式依赖逆转在 GEMM 输出中也出现,E1M2 在主要路径上保持或提高输出 SQNR,而 E2M1 在旋转后常退化。在 Q2 实验中,论文在 Dense 1.5B、MoE 7.9B 和 MoE 124B 长期预训练上比较了 BF16、E2M1 参考和 E1M2-based UFP4。结果显示 UFP4 在所有三个设置中都更接近 BF16:Dense 1.5B 上最新 1000 步相对误差从 1.2570% 降至 0.9673%,MoE 7.9B 从 2.3596% 降至 1.8469%,MoE 124B 从 1.7308% 降至 1.3863%。在 Q3 实验中,论文通过缩放定律分析验证优势跨模型规模持续存在,E1M2 曲线在测量范围和拟合外推中都低于 E2M1 参考曲线,拟合的 FP4 与 BF16 差距随计算增加而减小。在 Q4 实验中,论文进行了组件消融研究,发现全 RHT 覆盖最佳:相对于无 RHT,bwd_dw-only RHT 降低损失 0.00481,fwd_y+bwd_dw 降低 0.00644,bwd_dx+bwd_dw 降低 0.00290,全 RHT 降低 0.01123。在固定全 RHT 后,dY 的 SR 再降低损失 0.00456。论文还测试了范围受限的 E2M1(如 max_fpx = 2.0 仅保留 {0, 0.5, 1.0, 1.5, 2.0}),发现所有测试变体在 Dense 1.5B 和 MoE 7.9B 上都低于 E2M1 参考,表明 E2M1 范围限制不是均匀网格的满意替代。在 Q5 实验中,论文评估了融合 RHT+量化的效率,在块大小 16 时,融合的 RHT+量化在 SM90 和 SM100 上分别为独立量化的 1.06× 和 1.07× 延迟,而未融合的 RHT+量化则是融合延迟的 1.62× 和 1.41×,表明全 RHT 量化具有较小的融合内核开销。
查看结构化数据
| 任务 | 指标 | 本文 | 基线 | 提升 |
|---|---|---|---|---|
| Dense 1.5B 长期预训练 BF16 相对损失误差 | BF16-relative LM loss error | 0.9673% | 1.2570% | 相对降低 23.1% |
| MoE 7.9B 长期预训练 BF16 相对损失误差 | BF16-relative LM loss error | 1.8469% | 2.3596% | 相对降低 21.7% |
| MoE 124B 长期预训练 BF16 相对损失误差 | BF16-relative LM loss error | 1.3863% | 1.7308% | 相对降低 19.9% |
| linear_fc2/fwd_x RHT 后单张量 SQNR | SQNR (dB) | 23.19 dB (E1M2) | 20.00 dB (E2M1) | +3.24 dB |
| 全 RHT vs 无 RHT(Dense 1.5B, E1M2) | 平均 LM 损失降低 | 0.01123 | 0.00481 (仅 bwd_dw RHT) | 额外降低 0.00642 |
| 融合 RHT+量化延迟(块大小 16) | 相对于独立量化的延迟比 | 1.06× (SM90) / 1.07× (SM100) | 1.62× / 1.41× (未融合) | 延迟降低 34.6% / 24.1% |
局限与改进
论文的局限性主要体现在几个方面。首先,UFP4 虽然在所有测试模型上都优于 E2M1 参考,但与 BF16 相比仍存在可测量的损失差距,表明 FP4 训练尚未完全达到 FP16/BF16 的等效质量。在 Dense 1.5B 上剩余差距为 0.9673%,MoE 7.9B 为 1.8469%,MoE 124B 为 1.3863%。这表明即使是均匀网格的 FP4 训练,也需要进一步的技术改进来完全弥合精度差距。其次,论文假设训练张量在 RHT 后的局部密度是均匀的,这是推导 Shrinkage Bias 条件期望误差的前提。真实训练中的分布可能不完全满足这一假设,实际偏差可能略有不同。第三,论文的实验主要在 MoE 架构(SwiGLU MLP)上进行,虽然也评估了注意力模块,但其他架构(如标准 Transformer、Conformer、Diffusion 模型)中的表现尚未充分验证。第四,论文的硬件效率测试在 SM90 和 SM100 上进行,其他架构(如 AMD ROCm、Ascend NPU)上的效率可能不同,尤其是需要原生支持 E1M2/INT4 数据元素才能充分发挥优势。第五,论文的缩放定律分析范围从 10M 到 324M,外推到更大模型(如 1T+ 参数)的准确性需要进一步验证。最后,论文未探讨 E2M1 在推理或非 RHT 场景中的应用,这些场景中 E2M1 的宽动态范围可能仍有优势。
独立分析的弱点
本文的第一个潜在弱点是 Shrinkage Bias 的理论分析依赖于局部密度均匀的假设。真实训练张量在 RHT 后的分布可能呈现复杂结构(如多峰、长尾等),不完全均匀。这意味着条件期望误差公式 (ℓᵢ - rᵢ)/4 可能只是近似。改进方向可以是对实际分布进行精确积分,或者使用蒙特卡洛模拟量化误差。第二个弱点是论文的实验主要在 MoE 架构上进行,虽然 MoE 的 SwiGLU MLP 更容易产生异常值,但这可能限制了发现的普适性。改进方向是在更多架构(如标准 Transformer、Conformer、ViT、Diffusion 模型)上验证 UFP4 的优势,并分析不同架构对量化格式的敏感性。第三个弱点是论文未深入探讨量化误差对训练动力学的影响,如梯度估计偏差、优化轨迹偏差等。改进方向是分析 Shrinkage Bias 如何影响优化器的收敛行为,以及是否需要调整学习率或优化算法来补偿。第四个弱点是论文的硬件实现依赖于融合内核,但不同硬件厂商的实现差异可能影响实际性能。改进方向是在更多硬件平台(AMD ROCm、Ascend NPU、Intel Gaudi)上评估 UFP4 的效率和准确性。
未来方向
论文提出了几个有前景的未来研究方向。从作者角度,一个方向是探索 UFP4 与其他量化方法的组合。论文指出张量预处理方法(如 RHT、QuaRot、SpinQuant)和量化器改进方法(如 FAAR 的格式感知自适应舍入、Quartet II 的 EDEN 微缩放舍入)与 UFP4 是互补的。未来可以将 UFP4 与这些方法结合,进一步缩小与 BF16 的差距。另一个方向是探索更复杂的均匀网格变体,如对数均匀网格、自适应均匀网格,或者在不同层使用不同网格的混合策略。从成果延伸角度,一个方向是研究 Shrinkage Bias 在更低精度(如 FP2、INT2)或混合精度场景中的表现。更低精度的格式更加受限,几何偏差可能更加严重,均匀网格的优势可能更加明显。另一个方向是研究量化网格几何与其他数值格式(如 FP8、FP6)的交互。论文的发现可能也适用于这些格式,值得进一步研究。第三个方向是研究量化格式对模型最终性能(如下游任务准确率、鲁棒性、泛化能力)的影响,而不仅是预训练损失。这需要完整的预训练和微调实验,成本更高但更有实际价值。第四个方向是研究硬件设计,推动 GPU、TPU、NPU 等加速器原生支持 E1M2/INT4 数据元素,而不是仅支持 E2M1。论文提到华为 Ascend 960 的 HiFloat4 已经采用均匀 S1P2 数据元素,这是一个积极的信号。
复现评估
论文的复现性评估如下。开源方面,论文没有明确声明代码和数据的开源状态。根据标准学术实践,期望作者会发布代码、训练脚本和实验日志,但未在论文中说明。数据方面,论文使用的训练数据集未详细描述,但根据 LLM 训练的惯例,可能是公开的数据集(如 The Pile、C4 等)或内部数据集。算力方面,论文在 Dense 1.5B、MoE 7.9B 和 MoE 124B 模型上进行长期预训练,需要大量计算资源。例如,MoE 124B 模型训练到 800B tokens,如果使用 8K H100,可能需要数周甚至数月。这种规模的实验对大多数研究机构来说是不可行的。难度方面,论文的方法相对独立于特定框架,可以在 PyTorch、JAX、TensorFlow 等框架中实现。关键挑战是实现高效的 FP4 量化和 RHT 融合内核,这需要 CUDA/ROCm/HIP 等底层优化技能。论文未提供融合内核的实现细节,这可能增加复现难度。另一个挑战是获取合适的硬件支持,论文的测试在 SM90(Hopper)和 SM100(Blackwell)上进行,这些是 NVIDIA 的最新 GPU,可能不易获取。综合评估,中等复现难度:论文提供了充分的理论推导和实验细节,但大规模训练实验和底层内核实现需要显著资源和技术投入。
论文图表
Figure 1a 使用红色标记 RTNE 舍入区间,柱状条显示每个区间的期望舍入误差。E2M1 的非均匀区间呈现系统性向零偏倚(Shrinkage Bias),而 E1M2/INT4 的均匀区间保持无偏。Figure 1b 展示了这种网格级别的优势转化为训练质量:在 124B MoE 长期预训练上,基于 E1M2 的 UFP4 在 BF16 相对损失退化方面显著优于基于 E2M1 的基线。
这张图是论文的核心可视化,直观展示了 Shrinkage Bias 的几何来源和实际训练效果,是理解论文问题定义和贡献的基础。
该表展示了在 Dense 1.5B E1M2 FP4 运行上的 RHT 范围和 SR 消融研究,训练超过 100B tokens。RHT 范围消融块中,Δ 是相对于启用 SR 且无 RHT 的测量的。SR 消融块中,Δ 是相对于全 RHT、无 SR 变体的测量的。论文发现全 RHT 覆盖最佳:相对于无 RHT,bwd_dw-only RHT 降低损失 0.00481,fwd_y+bwd_dw 降低 0.00644,bwd_dx+bwd_dw 降低 0.00290,全 RHT 降低 0.01123。在固定全 RHT 后,dY 的 SR 再降低损失 0.00456。
这个表格提供了关键组件的消融分析,验证了全 RHT 和 SR 的各自贡献,支持 UFP4 的设计选择。