一句话总结

蚂蚁集团研究者发现，基于E2M1的FP4训练中存在一种名为Shrinkage Bias的系统性负向舍入误差，该误差源于几何bin的不对称性。他们提出UFP4，一种使用E1M2/INT4网格的统一4-bit训练方案，将Random Hadamard Transform应用于所有三个训练GEMM，并将随机舍入限制在dY上，在Dense 1.5B、MoE 7.9B和MoE 124B预训练中实现了更低的BF16相对损失退化。

核心贡献

• 论文识别出Shrinkage Bias，一种由非均匀E2M1网格的几何不对称性引起的系统性负向舍入误差，并表明该误差在层间乘性累积，且会被Random Hadamard Transform放大，解释了现有基于E2M1的FP4方案中观察到的训练不稳定性。
• 提出UFP4，一种统一4-bit训练方案，将RHT应用于所有三个训练GEMM，将随机舍入限制在dY梯度路径，并使用E1M2/INT4风格均匀网格，将RHT带来的改善的bucket利用率转化为更高的量化质量。
• 在Dense 1.5B、MoE 7.9B和MoE 124B长程预训练中，UFP4相比强E2M1基线，始终实现更低的BF16相对损失退化，并得到缩放定律分析和消融研究的支持。

引言

作者研究了使用4-bit浮点格式训练大语言模型的技术，该技术有望显著降低内存和计算成本。当前硬件和训练方案主要依赖非均匀E2M1格式，但端到端FP4预训练仍不稳定，且相比更高精度基线存在精度损失。先前工作尝试通过Random Hadamard Transform (RHT) 和随机舍入等技术缓解这种不稳定性，但仍保留了有问题的E2M1数据格式。

作者发现了一个根本性局限：像E2M1这样的非均匀网格固有地存在Shrinkage Bias，即由其可表示bin的几何不对称性引起的系统性负向舍入误差。他们证明该偏差在神经网络层间乘性累积，导致信号衰减，并矛盾地被RHT放大，因为RHT将张量质量移入该格式最不对称的区域。作者的主要贡献是UFP4，一种用E1M2或INT4等统一4-bit格式替代E2M1网格的训练方案。该方法绕过了网格几何误差，使得RHT能够有效应用于所有关键矩阵乘法，同时将随机舍入限制在单个梯度张量上，从而在规模上实现了更优的训练稳定性和精度。

数据集

作者描述了用于构建低精度训练数据的量化格式和分块量化过程。

• 格式码本：考虑了两种4-bit浮点格式E2M1和E1M2，以及一种4-bit整数格式(INT4)。每种格式定义了一个具有固定可表示级别集合的归一化码本 $G$G。
• 分块量化：张量被划分为连续块。对每个块，共享缩放因子 $s_B$sB​ 由块内最大绝对值除以码本中最大幅度级别计算得出。然后每个元素通过舍入规则缩放并映射到码本级。
• 舍入规则：论文支持确定性的Round-To-Nearest-Even (RTNE) 和随机舍入 (SR)，后者以与相邻级别间距离成比例的概率采样一个级别。RTNE被指出是近期低精度方案（如NVFP4方法）中广泛采用的选择。
• 模型中的使用：量化张量 $Q_G(\mathbf{T})$QG​(T) 在整个训练管线中使用；在算术表达式中，它表示反量化后的数值张量，确保低精度表示直接集成到前向和后向传递中。

方法

作者解决了量化训练张量的挑战，这些张量通常包含导致大多数码本级利用率不足的离群坐标。为缓解此问题，他们利用了Random Hadamard Transforms (RHT)。RHT在量化前应用一种保范旋转，将离群能量分散到所有坐标上。作者使用递归定义的Sylvester Hadamard矩阵：

$$\mathbf { H } _ { 2 n } = \frac { 1 } { \sqrt { 2 } } \left[ \begin{array} { l l } { \mathbf { H } _ { n } } & { \mathbf { H } _ { n } } \\ { \mathbf { H } _ { n } } & { - \mathbf { H } _ { n } } \end{array} \right] , \qquad \mathbf { H } _ { 1 } = [ 1 ] ,$$H2n​=2​1​[Hn​Hn​​Hn​−Hn​​],H1​=[1],

其中 $\mathbf{H}_{n}^{\top}\mathbf{H}_{n}=I_{n}$Hn⊤​Hn​=In​ 对于 $n=2^{k}$n=2k。RHT在Hadamard变换前额外应用一个随机符号矩阵 $\mathbf{S}_{n}=\operatorname{diag}(\epsilon_{1},\ldots,\epsilon_{n})$Sn​=diag(ϵ1​,…,ϵn​)，$\epsilon_{i}\in\{-1,+1\}$ϵi​∈{−1,+1}。由于 $\mathbf{H}_{n}^{\prime}=\mathbf{S}_{n}\mathbf{H}_{n}$Hn′​=Sn​Hn​ 是正交的，对共享GEMM维度应用相同变换可保留全精度结果：

$$\begin{array} { r } { \mathbf { H } _ { n } ^ { \prime } = \mathbf { S } _ { n } \mathbf { H } _ { n } , } \\ { \mathbf { Y } = \mathbf { X } \mathbf { W } ^ { \top } = ( \mathbf { X } \mathbf { H } _ { n } ^ { \prime } ) ( \mathbf { W } \mathbf { H } _ { n } ^ { \prime } ) ^ { \top } . } \end{array}$$Hn′​=Sn​Hn​,Y=XW⊤=(XHn′​)(WHn′​)⊤.​

相应的低精度GEMM量化旋转后的操作数：

$$\widehat { \mathbf { Y } } = Q _ { G } ( \boldsymbol { \chi } \mathbf { H } _ { n } ^ { \prime } ) \, Q _ { G } ( \mathbf { W } \mathbf { H } _ { n } ^ { \prime } ) ^ { \top } .$$Y=QG​(χHn′​)QG​(WHn′​)⊤.

虽然RHT减少了离群值主导，但其有效性严重依赖于底层量化网格。作者识别出主流E2M1格式的一个根本性问题，他们称之为Shrinkage Bias。该偏差源于非均匀网格中Round-to-Nearest-Even (RTNE) 舍入bin的几何不对称性。对于内部量化级别 $q_i$qi​，RTNE舍入bin具有左右宽度 $\ell_{i}$ℓi​ 和 $r_{i}$ri​。如果bin内密度局部均匀，则条件期望误差为：

$$\mathbb { E } [ \rho _ { G } ( t ) - t \mid t \in \mathcal { B } _ { i } ] = \frac { \ell _ { i } - r _ { i } } { 2 } = \frac { 2 q _ { i } - q _ { i - 1 } - q _ { i + 1 } } { 4 } .$$E[ρG​(t)−t∣t∈Bi​]=2ℓi​−ri​​=42qi​−qi−1​−qi+1​​.

一个满足 $r_{i} > \ell_{i}$ri​>ℓi​ 的不对称bin固有地产生负期望误差，导致幅度收缩。如下图所示，E2M1网格在间距转换点处表现出这些不对称bin，而像E1M2和INT4这样的均匀网格则保持对称bin（$\ell_{i}=r_{i}$ℓi​=ri​），从而消除了这种几何偏差来源。

为解决此问题，作者提出UFP4，一种基于统一网格（E1M2/INT4）的4-bit训练方案。一旦RHT将张量从动态范围受限状态转变为局部分辨率受限状态，4-bit网格优先考虑局部幅度保持，而非极端动态范围。完整训练管线参见框架图。

在UFP4方案中，对每个线性层GEMM，作者应用RHT并将操作数量化到均匀间隔网格。关键的是，虽然E2M1方案通常将RHT限制在权重梯度路径以避免复合几何误差，但UFP4的无偏特性允许他们安全地在所有三个GEMM上启用RHT：FPROP (fwd_y)、DGRAD (bwd_dx) 和 WGRAD (bwd(dw))。随机舍入 (SR) 应用于 $dY$dY 以保持梯度期望。缩放层次设计与此格式级解决方案正交，允许根据硬件效率选择单层、两层或块缩放。

实验

实验考察了FP4网格几何和Random Hadamard Transform (RHT) 范围如何在局部量化和端到端训练中相互作用。单张量和GEMM输出诊断显示，RHT在离群值密集的张量上逆转了首选格式排名，使旋转后的均匀E1M2网格性能优于E2M1。在多达124B参数的Dense和MoE模型长程预训练中，基于E1M2的UFP4方案相比E2M1参考，始终减少了BF16相对损失差距，且缩放定律验证确认此优势在模型规模上持续存在。消融研究揭示，完整RHT覆盖和输出梯度上的随机舍入均有助于增益，而范围受限的E2M1无法替代原生均匀网格支持，且内核基准测试表明将RHT融合到FP4量化中仅引入适度开销。

作者进行了一项消融研究，评估Random Hadamard Transform范围和随机舍入对FP4 E1M2训练的影响。结果表明，在所有GEMM输出上应用完整变换覆盖可获得最低的语言建模损失，且添加随机舍入进一步提升了性能。完整变换覆盖优于部分应用，证明在E1M2网格下旋转所有GEMM输出是有益的。当与完整变换覆盖结合时，随机舍入提供了额外的损失降低。UFP4方案相比范围受限或无随机舍入的配置，实现了最佳整体性能。

消融研究考察了Random Hadamard Transform的范围和随机舍入的使用如何影响FP4 E1M2训练。将变换应用于所有GEMM输出始终比部分覆盖产生更低的语言建模损失，确认了在此数据格式下完整旋转是有益的。在完整变换覆盖之上添加随机舍入进一步降低了损失，组合后的UFP4方案实现了最佳整体性能。