Re0: 当表达细节成为 FP4 新瓶颈时,我们是否仍要坚守 E2M1?

大家好,最近我们在 arxiv 中上传了一篇文章:《Rethinking Shrinkage Bias in LLM FP4 Pretraining: Geometric Origin, Systemic Impact, and UFP4 Recipe

很高兴,也希望在这里简单聊一聊这篇技术报告背后的一些想法。

 


 

继去年 Ling Team 发布原生 FP8 训练的 Ling 2.0 系列模型,交出了我们在低精度训练下的第一份答卷之后(聊一聊我们最近开源的 Ling 2.0 原生 FP8 混合精度训练方案),我们开始探索更低 bit 表示的 FP4。

和许多低精度训练工作一样,我们最初想回答的问题很直接:

FP4 预训练能否达到接近“无损”的表现?

 

这篇技术报告讲述了我们近期的一些新发现:

  • 绝大多数 FP4 相关的工作,人们更倾向于在 rounding 策略、scaling 设计、outlier 缓解、量化感知等方面修修补补,而很少质疑 E2M1 这一数据格式本身的合理性。
  • 我们发现,E2M1 先天存在 Shrinkage Bias,即部分数值左右 rounding bin 的不对称性导致量化过程中存在有偏 bias,该 bias 在训练过程中逐渐累积拖慢了 FP4 训练的收敛速度。而在 E1M2 / INT4 这类 uniform grid 下,Hadamard Transform 带来的 bucket 利用率提升更容易转化为实际量化质量的收益,从而带来比 E2M1 更好的收敛表现。

 

这一发现是否意味着 E1M2 相比 E2M1 更具潜力作为未来 FP4 训练的默认格式?为什么当前硬件和社区工作又主要围绕 E2M1 展开?

我想结合低精度训练历史中的一些事件,以及我们这篇研究,简单聊聊自己的看法。

以下内容及相关表述仅为个人观点,如与真实情况不符,欢迎大家指正。

 


 

为什么低精度训练的发展总是慢于推理?

img

回顾过去可以发现,推理侧的低精度发展远远领先于训练侧。

在许多人还在探索如何解决 FP32 转到 FP16 后的 loss 跑飞、训练不稳定等问题时,推理侧已经迈入了 INT8 w/ per-channel 细粒度量化阶段。

 

一个很重要的原因是:训练侧每一次向更低精度探索,都必须回答一个更难的问题:

采用更低精度训练带来的加速收益,能否抹平它引入的潜在风险?

 

在推理侧,精度砍半通常直接意味着显存消耗大幅下降、吞吐提升和部署成本降低。因此,推理侧有更充足的空间尝试各种量化方案。新的量化方案往往只需要在多个模型、多个 benchmark 上完成离线评测,确认模型表现符合预期即可。

 

训练侧则完全不同。

一种新的低精度格式或 recipe 是否可靠,往往需要多种模型尺寸、数百 B 甚至 T token 级别的长跑验证。短跑可行并不意味着长跑误差不会累积;小模型可行也不意味着大模型不会暴露新的数值问题。

因此,训练侧的低精度探索成本远高于推理侧,也更容易依赖上一代格式选择中积累下来的历史经验。

 

这正是 FP4 训练面对的第一个现实背景:E2M1 并不是从零佐证的最优格式,而是承载着 FP16、BF16、FP8 的历史惯性所做出的选择。

 

历史经验:动态范围曾挽救了 16 bit 训练

技术更迭里有一种常见的 bitter lesson:

一条路线之所以成为主流,往往不是因为它永远正确,而是因为它曾经正确地解决了上一代问题。

低精度训练也是如此。

 

自 Volta V100 支持 FP16 Tensor Core 矩阵乘时,远高于 Pascal 架构 FP32 CUDA Core 算力的优势,吸引了人们逐步从 32bit 训练跨入 16bit。同期工作也基本奠定了后来 mixed-precision training 的基础 recipe:

  • Linear / GEMM 等核心计算使用低精度输入
  • 乘加累加通常保留更高精度
  • 保留 FP32 master weight
  • 梯度更新和优化器状态使用 FP32
  • 通过 loss scaling 等机制缓解 FP16 下溢问题

 

这套方案让 FP16 训练真正进入主流,但随着 GPT-2、GPT-3、T5 等模型的推出,LLM 参数量越来越大,网络也越来越深。在深层网络中,activations 随层逐渐变大的趋势,以及 backward 中存在极小 gradient 等现象,导致 FP16 很容易产生上溢/下溢之类的问题。

虽然 loss scaling 能部分缓解上溢和下溢,但它本身也引入了新的工程问题:以什么策略 scaling?如何保障庞大网络下每一处 FP16 GEMM 都不出问题?如何在训练吞吐和稳定性之间权衡?

 

后来,Google TPU 带来了 BF16。BF16 的指数位与 FP32 一致,尾数位少于 FP16。它本质上是一次 bit budget 的重新分配:牺牲一部分 mantissa,换取更大的 exponent range。

BF16 胜出的关键在于,它显著降低了动态范围不足带来的 overflow、underflow 以及 loss scaling 管理的难度,使大模型训练变得更加稳健。

 

从 FP16 到 BF16 的历史中,社区自然形成了一个很强的经验:

对训练来说,动态范围非常重要。
如果要在 exponent 和 mantissa 之间取舍,很多时候应该优先保护 exponent。

 

这条经验在 16 bit 时代非常正确,在早期 FP8 时代也很有解释力。

例如,NVIDIA 为 Hopper 架构同时支持了 FP8 E4M3 与 E5M2。在早期 per-tensor wise FP8 训练中,社区通常反馈对 WGRAD、DGRAD 使用 E5M2 往往能够取得更好的收敛效果。这进一步强化了一个直觉:训练中的 gradient 很难量化,动态范围必须优先保障。

于是,当训练继续走向 FP4 时,E2M1 非常自然的成为了优先选择。

 

历史惯性:成功的经验更易被延续

站在 FP16、BF16、FP8 的历史之后再看 FP4,E2M1 成为当下主流 GPU 的选择并不难理解。

 

在 FP4 下,紧缺的 4 bit 的预算只能衍生出以下三种 format:

  • E3M0:有着最高规格的动态范围,但局部分辨率极差
  • E2M1:保留更多动态范围,同时给尾数留下 1 bit,是相对折中的选择
  • E1M2:牺牲更多动态范围,换取更均匀、更细的数值 grid

结合真实 tensor 分布的分析,LLM 中很多权重、激活和梯度具有明显的 outlier-heavy 的特征,少量 outlier 会显著影响 absmax scaling。毫无疑问,选择 E3M0/E1M2 面临着极大的适用性风险,而 E2M1 则是 FP4 下的 sweet point。

 

因此,E2M1 的选择,一定程度上也是历史经验的自然延续。

但这真的是当下的最佳选择么?怀疑的念头一旦被种下,便挥之不去。

 

当量化粒度越来越细,问题的主导矛盾发生了改变

历史经验之所以成为惯性,是因为外在条件发生变化之后,我们仍然沿用旧问题的答案。

 

在 per-tensor scaling 时代,一个 tensor 共享一个 scale。此时 element format 必须自己承担 tensor 内部很大一部分动态范围。outlier 一旦出现,就会影响全局 scale,导致大量小值被压缩到很少的 bucket 中(codebook 数值)。

在这种情况下,给 element 更多 exponent bit 是合理的。但随着低精度训练继续发展,scaling 的粒度变得越来越细,当每个很小的 block 都有独立 scale 时,动态范围的控制正在逐渐从 data element 转移到 scale。此时,每个 block 内的数值能否以“更高分辨率”的 format 来表达变得越来越重要。

 

从 BF16 到 FP8 的演变中,也能看到类似趋势。早期 per-tensor FP8 训练中,WGRAD、DGRAD 往往需要使用更高动态范围的 E5M2;而在更细粒度的 blockwise FP8 路线中,E4M3 用于 backward 相比于 E5M2 能带来更好的收敛效果。

 

一个自然的问题随之出现:

量化粒度越细,是否意味着我们可以在 format 层面少分配一个 exponent bit,而多分配一个 mantissa bit?

这便是我们重新考虑 E1M2 / INT4 的起点。

 

多给一位 mantissa bit 更好?E1M2 的首轮实验但以失败告终

在 FP4 下,提升 1 bit 预算给尾数的 E1M2 是否表现优于 E2M1?

有了上述直觉之后,我们立即展开了一系列实验,复刻 NVFP4 recipe 及其变体,并尝试将底层 format 替换为 E1M2。

但遗憾的是,在最初所有尝试中,E1M2 的表现并不理想,较小的动态范围所带来的劣势导致它在长跑下均会被 E2M1 所超越。

 

此后一段时期内,E2M1 仍是我们实验的 base format 未曾改变。

而真正影响我们判断的,起源于 E2M1 方案探索下的一个反直觉的现象。

 

转折点:RHT 显著降低了量化下溢,但却有损于 FPROP/DGRAD

Random Hadamard Transform(RHT)在低精度量化中常被用来缓解 outlier。它通过近似正交变换,把原本集中在少数维度上的大值打散,使张量分布变得更平滑、更集中。

 

更平滑、更集中的分布有利于量化,RHT 显著降低了量化的下溢比例,并提升了 E2M1 仅 15 个元素 bucket 的利用率。按直觉来讲,这应该能带来更好的量化质量。

但实验结果出乎意料,RHT 放置于 FPROP 以及 DGRAD GEMM 中甚至是有损的(与 NVIDIA 论文现象一致),这一现象一度非常困惑我们。

RHT 缓解了 outlier、降低了下溢比例、提高了 bucket 利用率,但却无法转化为更好的计算质量。

 

直到我们发现了 Shrinkage Bias 的存在,即:

在 E2M1 部分数值中存在左右 rounding bin 不对称的现象,RHT 对矩阵分布的处理使更多数值更易落入该区域,从而加剧了有偏的 bias,进而表现为更慢的 loss 收敛速度。(更直观点说:在 E2M1 下 RHT 所带来的 Shrinkage Bias 抵消了原有的收益)

img

 

这给了我们一个新的方向:

如果问题的根源来自 E2M1 非均匀 grid,那么使用 E1M2/INT4 这样 uniform grid 之后,RHT 所带来更高的 bucket 利用率能否转化为更佳的量化/计算质量?

 

有以下测试,以 linear_fc2 (mlp.down_proj) 的输入 tensor 来举例:

  • 在执行 RHT 之前,E2M1 raw 量化信噪比明显优于 E1M2 raw
  • 在执行 RHT 之后,结论发生了反转:RHT 为这两种 Format 同时带来了 bucket 利用率的提升,但由于 Shrinkage Bias 的存在导致该增益无法转换为 E2M1 更好的量化质量,而 E1M2 享有该增益后反超了 E2M1 raw

img

 

这也是我们对 FP4 format 的判断发生变化的关键节点。

 

UFP4:让 RHT 的收益真正转化为量化质量

基于这些观察,我们的 UFP4 recipe 十分简单:

在 E1M2/INT4 这类 uniform 4-bit format 下,为 FPROP、DGRAD、WGRAD 路径下的三个 GEMM 的操作数全部启用 RHT;将 SR 应用于 dy 量化过程中。

img

 

UFP4 更偏向于 format-level 的探索,我们在细粒度量化 + RHT 的条件下,找到了一条 4-bit 潜在更佳的 format 路线。在该设计下,FP4 训练“无惧”outlier、“无惧”Shrinkage Bias,也潜在具备更纯净的未来探索空间,更少的 patch。

此外,类似的思路也可应用于推理、RL 训推以及 4bit 压缩/解压缩中。

在 E2M1 路线下,有一类工作专注于让模型感知到“量化”过程的存在,通过梯度补偿等方式让模型在训练过程中自行纠偏,这类探索以巧妙的设计来弥补 E2M1 的天然缺陷。

但只要最终的量化还是投影到 E2M1 grid,便或多或少面临着 grid geometry 本身所存在的一些问题。持续的叠加更多 patch,本质上依然是在修补一个底层格式带来的结构性问题。

 

从“动态范围主导”转向“局部分辨率主导”,我们是否仍要坚守 E2M1?

在以上实验支撑下,我们开始反思:E2M1 是否值得作为未来 FP4 训练的默认 Format,E1M2/INT4 是否更具潜力来适用于更多场景。

FP4 训练暴露了许多在 FP8 中甚至可忽略不计的新问题需要权衡:

  • 如何确保 Rounding 的无偏性,保障正确梯度方向的同时,又能避免过多 noise 的引入
  • 如何既能保障 weight 具有足够的更新量,也能缓解 Rounding 临界点的权值震荡所造成的影响
  • 如何设计更好的量化前处理操作,既能降低量化误差,也不引入新的有偏 bias
  • 如何权衡 weight 2D 量化会引入更高的量化误差,但却能保障更佳的前反向计算的一致性

诸多因素之间互相影响,使得 FP4 Recipe 的优化困难重重。为 E2M1 定制一系列纠偏 patch 的潜在成本又让整体系统变得更加复杂。

 

我们倡导,未来的训练硬件不应只押注于 E2M1。

E2M1 作为“动态范围主导”的问题设计的折中;而当 细粒度量化 与 RHT 大幅改变张量分布后,FP4 训练可能进入了一种“局部分辨率主导”的新范式。

在这一新范式下,uniform 4-bit format 的价值也应该被重新评估。

 


 

🚀 最后打个广告,团队持续招人中,可邮件联系:zq317110#antgroup.com (# 替换为 @)


我想对千千说~