因为 7x7 卷积回退这个热搜词我测试了整整一个晚上才把 PyTorch 原生 FP8 数据类型在真实模型上的表现摸清楚。如果你以为 FP8 就是把模型to(torch.float8_e4m3fn)这么简单那这篇文章能帮你至少省下三天排查时间。FP8 是 H100 那一代 GPU 上引入的 8 位浮点格式PyTorch 从 2.1 开始把float8_e4m3fn和float8_e5m2作为原生 dtype 放了出来。它最大的价值是直接砍掉一半显存占用和带宽压力同时在 Hopper 架构上走 8 位张量核心吞吐比 FP16 高一大截。适合已经在用混合精度训练推理、手头有 H100/H200 这类 GPU、想压一压显存和通讯开销的人。这篇主要是想聊聊它到底怎么用、数值边界在哪儿、以及那些文档里不会告诉你的回退与缩放坑。1. FP8 不是单一格式E4M3 与 E5M2 的定位差异1.1 为什么要压到 8 位浮点先说最直观的收益。一个 FP32 参数占 4 字节FP16/BF16 占 2 字节FP8 只占 1 字节。单就权重和激活值来说显存占用直接砍半通信量也砍半。这个收益不是线性意义上的省一点而是能决定一个模型能不能塞进单卡、数据并行梯度同步能不能跑得动的问题。计算吞吐上的收益更明显。H100 的 FP8 张量核心吞吐是 FP16 的两倍是 FP32 的四倍起步。对很多矩阵乘法和卷积来说瓶颈早就不在算力而在显存带宽上数据从 4 字节变 1 字节搬运量小了四倍自然就快了。我实测下来某些大模型层的前向时间能缩短 20%~30%这不是挤牙膏是实打实的提升。但要清楚一个前提FP8 不是用来替代 FP32 做精确计算的。它的目标场景本身就是大模型的推理和训练中间过程因为神经网络对这个精度的扰动有天然的鲁棒性。你如果用float8_e4m3fn去算0.1 0.2会发现结果对不上这不是 bug是格式本身精度就到这儿了。1.2 两种格式的数值规格PyTorch 原生提供的 dtype 里最常用的是torch.float8_e4m3fn和torch.float8_e5m2。名字里的 E 和 M 分别代表指数位和尾数位比如 E4M3 就是 4 位指数、3 位尾数加上 1 位符号位总共 8 位。属性float8_e4m3fnfloat8_e5m2符号位11指数位45尾数位32指数偏置715最大有限值448.057344.0最小正规数2^-6 ≈ 0.0156252^-14 ≈ 0.000061最小次正规数2^-9 ≈ 0.0019531252^-16 ≈ 0.0000153典型用途前向激活、权重反向梯度、误差累积这个差异直接对应场景选择。E4M3 尾数多一位精度更高适合放权重和激活值因为神经网络前向传播对值的大小范围要求没那么极端但对相对精度更敏感。E5M2 指数多一位动态范围更大适合放梯度因为反向传播过程中梯度的量级变化非常剧烈有可能从 1e-2 跳到 1e-6动态范围不够就直接清零了。1.3 E4M3FN 里那个 FN 是什么意思FN 全称是 finite only, no NaN。这意味着在 E4M3FN 这个格式里指数全 1 的位模式没有被保留给 NaN 或 Inf而是直接用来表示 448 这个最大有限值。换句话说这个格式根本没有用来表示 NaN 和 Inf 的位组合。这在工程上有个重要影响你在卷积或矩阵乘法里如果数值溢出FP16 会变成 inf然后一路传播成 NaN而 E4M3FN 的溢出行为通常更接近饱和到最大值。这是 NVIDIA OCP 规范里定义的行为。所以同样是溢出FP8 在某些场景下反而不容易让训练直接崩掉但它会引入大概率的截断误差这个误差不爆 NaN但会默默污染训练曲线。E4M3FN 的 FN 还带来一个细节它没有符号零以外的特殊处理所有位模式都是合法的有限数。这让它在实现上和驱动、硬件的配合更简单但也意味着你在调试时不能依赖打印出 inf/NaN 来定位数值问题。2. PyTorch 原生 FP8 的最小 API 与数值行为2.1 创建 FP8 张量PyTorch 把 FP8 当成了普通 dtype 来处理所以创建方式和你平常用的torch.float16没有区别import torch # 从 Python 标量创建 a torch.tensor([1.0, 1.5, 1.25], dtypetorch.float8_e4m3fn) # 从已有张量转换 b torch.randn(4, 4).to(torch.float8_e5m2) # 也可以用 empty_like 这类工厂方法 c torch.empty_like(b, dtypetorch.float8_e4m3fn) print(a.dtype) # torch.float8_e4m3fn print(a) # tensor([1.0000, 1.5000, 1.2500], dtypetorch.float8_e4m3fn)还有一个比较少人注意但非常有用的类型torch.float8_e4m3fnuz。这个带uz后缀的变体是unsigned zero的意思即不区分0.0和-0.0它主要服务于一些特定硬件平台比如 Intel Gaudi 之类。如果你只是跑 NVIDIA GPU用float8_e4m3fn就够如果是做跨平台部署就要留意不同硬件对 fn/fnuz 的支持差异。2.2 标量精度实验看看量化后丢了什么我强烈建议你在真正开跑模型之前先花五分钟在纯 PyTorch 里做一组标量实验理解 FP8 的精度粒度。做法很简单import torch vals [1.0, 1.5, 1.25, 1.125, 1.1, 0.1, 448.0, 449.0, 0.015625, 0.0001] for v in vals: t_e4 torch.tensor([v], dtypetorch.float8_e4m3fn).item() t_e5 torch.tensor([v], dtypetorch.float8_e5m2).item() print(f{v:10.6f} - E4M3: {t_e4:10.6f}, E5M2: {t_e5:10.6f})你会看到大概这样的结果1.0、1.5、1.25 在 E4M3 里都精确得住因为三位尾数能表示 0.125 的倍数1.125 也精确因为它正好是 1 加上 0.125但 1.1 就丢精度了会变成 1.125 或 1.0取决于舍入模式。E5M2 则更粗0.25 的倍数才精确。0.015625 是 E4M3 的最小正规数正好踩线0.0001 在 E5M2 里能表示在 E4M3 里直接变成次正规数甚至归零。这个实验的价值在于让你建立哪类数值是安全的这种直觉好的做法是后续把激活值的分布也拉出来看如果激活值集中在小数附近E4M3 的精度是足够的如果动态范围很大就别硬扛。2.3 三个容易踩的坑第一个坑是torch.set_default_dtype。有人想全局默认 FP8写torch.set_default_dtype(torch.float8_e4m3fn)然后报错。这是正常的set_default_dtype只接受 CPU 上的浮点格式float32、float64、float16、bfloat16FP8 不在支持列表里。想全局切 FP8得靠显式的to()和autocast自定义后面会讲。第二个坑是 CPU 上的 FP8 行为。PyTorch 在 CPU 上确实支持 FP8 张量的存储和基本类型转换但大多数算子没有 CPU 端的 FP8 内核实现。你在 CPU 上做a * b经常要么回退到批量转换、要么直接不支持。所以别指望拿笔记本 CPU 去复现 GPU 上的性能。CPU 上的 FP8 更多是用来检查数值行为、调试格式细节。第三个坑是打印和序列化的 dtype 名。torch.load和torch.save对 FP8 的支持已经比较成熟但某些老版本2.1 初期在加载 FP8 checkpoint 时会报 Unknown dtype 之类的错误。如果你要和朋友交换模型文件尽量统一 PyTorch 版本最好把 FP8 权重转回 FP32 再落盘。3. 把 FP8 装进自己的模型从 Linear 到自定义算子3.1 从一行 Linear 开始最简单的接入方式是构造一个nn.Linear然后把输入和权重显式转到 FP8import torch import torch.nn as nn model nn.Linear(512, 1024) x torch.randn(16, 512) # FP8 推理示意 x_fp8 x.to(torch.float8_e4m3fn) w_fp8 model.weight.detach().to(torch.float8_e4m3fn) y torch.nn.functional.linear(x_fp8, w_fp8, None) print(y.shape, y.dtype)这里有几个关键点要注意。nn.Linear的 forward 里已经包含了F.linear它内部对输入和权重会走一个类型推断逻辑如果你传入 FP8 张量它能看到 FP8 输入并尝试调用对应的内核。但 PyTorch 2.1、2.2 时代的 FP8 算子覆盖面并不全很多F.linear在遇到 FP8 时并不是直接走真正的 FP8 张量核心而是先把数据当作高精度浮点来计算。想要真正吃到 H100 上的 FP8 张量核心吞吐通常需要依赖torch.compile或者专门的优化库把算子 dispatch 到 CUTLASS/cuBLAS 的 FP8 路径上。所以我的建议是先小规模跑通数值和显存收益再考虑性能收益。性能收益那一部分靠手写一层F.linear通常不够。3.2 自定义 autograd.Function 处理 FP8 梯度如果你要在一个自定义训练流程中控制 FP8 的精度和缩放最简单的做法是写一个自定义autograd.Function把 FP8 的向下转换看作一个有舍入误差的伪量化算子前向把输入/权重转成 FP8 后再转回来反向直接传梯度import torch class FP8QuantizeFunction(torch.autograd.Function): staticmethod def forward(ctx, x, scale): x_scaled x * scale x_fp8 x_scaled.to(torch.float8_e4m3fn) # 转回高精度用于后续计算 x_rec x_fp8.to(x.dtype) ctx.scale scale return x_rec staticmethod def backward(ctx, grad_output): # FP8 量化本身在前向引入了舍入误差 # 反向通常直接透传梯度这是 1 位浮点训练中的常见近似 return grad_output, None这个写法看似简单其实是在模拟 Transformer Engine 里的标准算子。真正的 FP8 训练中反向梯度也是 FP8 的所以 backward 通常也会包一层 FP8 量化。这里我为了演示先只做前向的伪量化算子。再往深一层说你在前向用的scale到底怎么来是 FP8 工程里最核心的问题。这个我放到下一节专门讲因为它够复杂也够容易被忽略。3.3 先看看 torchao 和 Transformer Engine如果你只是想在真实项目里快点用上 FP8不用从零手写缩放逻辑两个现成路径值得参考。一个路径是torchaoPyTorch 官方的量化与调优库它提供一些针对 FP8 的实用集成比如quantize_接口和针对 Attention 的 int8/FP8 实现。另一个是 NVIDIA 的Transformer Engine它对 Transformer 结构的 FP8 支持非常成熟te.Linear这种层出来就自带延迟缩放逻辑模型改动成本很低。我个人的取舍是如果是新项目用 Transformer Engine 起步最稳如果项目已经大量使用原生态nn.Linear、nn.Conv2d希望侵入式改动小那从 torchao 或手写范式切入更合适。论文复现类工作最好直接手写缩放因为研究过程中需要控制变量封装层的黑盒有时候说不清数据路径。4. 延迟缩放FP8 工程的灵魂4.1 为什么 FP8 必须配缩放因子老问题来了FP8 的动态范围就这么大E4M3 最大 448E5M2 最大 57344听着挺大但大模型的激活值动辄上千梯度动辄 1e-6直接硬转不是溢出就是归零。解决思路很直观给数据乘一个缩放因子。我们先统计一段数据的最大绝对值 amax然后选择一个 scale让abs(x * scale)的最大值尽量贴近 FP8 的max_valscale max_val / amax这样一来原来最大值为 1000 的激活乘以 scale0.4481000 变成 448成功压进 E4M3 的动态范围。反向的时候再除以 scale 还原。用瀑布做类比FP8 是一只小杯子它的量程是固定的缩放因子就是往杯子里倒水前先确定水位合适不合适而不是强行倒进去再等溢出。难点不在这条公式而在于 amax 这一时刻的值是变动的。训练过程中激活和梯度分布一直在变你不可能每一步都全局同步求一遍 amax这会产生很大的通信开销。4.2 延迟更新的手动实现逻辑业界主流做法叫 delayed scaling翻译过来就是缩放不追实时滞后一段时间。上一轮迭代统计出来的 amax用来作为下一轮的 scale 基准。这样不会每步都同步通信开销降到最低而模型参数通常不会在几步之间剧烈跳变所以滞后一两个 step 的统计量完全够用。Transformer Engine 里这个窗口长度通常是 1024 步并用指数滑动平均去平滑 amax 的变化。手动实现一个精简版可以这样import torch from collections import deque class DelayedScaler: def __init__(self, eps1e-12, window16, max_val448.0): self.eps eps self.window window self.max_val max_val self.history deque(maxlenwindow) self.scale 1.0 def update(self, x): # 记录当前 batch 的最大绝对值 amax x.abs().max().item() self.history.append(amax) if len(self.history) self.window: # 取窗口内最大稍微激进一点也可以用均值 win_amax max(self.history) self.scale self.max_val / max(win_amax, self.eps) def quantize(self, x): return (x * self.scale).to(torch.float8_e4m3fn).to(x.dtype) / self.scale这里用队列维护最近 16 步的 amax满了就更新一次 scale。使用的时候在训练循环里对每个需要 FP8 的 tensor 调用update下一轮再用quantize真正的转。实际大规模训练里窗口和更新策略还会涉及更多细节比如不同层使用独立 scale、scale 更新频率与 batch size 的关系但核心框架就是这个。4.3 缩放因子翻车的样子缩放因子设置不当最典型的表现是训练 loss 突然跳变。有一种情况是 scale 设得太大导致大量次正规数被舍入掉某个通道的梯度直接清零loss 曲线横住不动。另一种是 scale 设得太大反而把激活压成了零这在小 batch 里尤其容易出现。还有一种情况是 amax 刚好取到某个 outlierscale 给得极小结果大部分数值都被压到次正规数区间精度反而比不缩放还差。所以我们在工程里通常不会让scale 448 / amax直接落地而是加一个margin系数比如用max_val / (amax * 1.1)或者做指数滑动平均去平滑。这样做的代价是按保守的边界压缩损失一点动态范围利用率换来的却是训练的稳定性。5. 实战中的回退陷阱7x7 卷积与落盘5.1 内核大小超过 32 的卷积会怎样前面说过的热搜词这次成了我真实调试时的噩梦FP8 卷积不支持内核面积大于 32 的情况比如 7x7 卷积会自动回退到更高精度。我第一次看模型里卷积层都转了 FP8、显存也降了以为万事大吉结果一跑 profiler 才发现ResNet 的 stem 层和 Vision Transformer 里 patch embedding 的那个 7x7 卷积压根没走 FP8 内核。为什么会这样卷积的 FP8 内核实现比矩阵乘法复杂得多要考虑 filter 维度、channel 对齐、im2col 展开后的内存布局还有 Winograd 变换的边界条件。cuDNN 里 FP8 卷积内核在特定配置下只覆盖一定范围的 kernel size超出范围的就在算子 dispatch 层自动 fallback 到 FP16 或 FP32。pyTorch 在 2.2 到 2.4 这个区间如果你用torch.compile跑 conv编译器看到 FP8 输入会在图优化阶段判断当前卷积配置是否命中 FP8 内核没命中就插入一个 cast 节点把数据转成 FP16/FP32 再跑卷积。因此最终计算图里会出现 FP8 - FP16 - Conv - FP16 - FP8 这种隐藏的转换路径。有人觉得多转换几次成本也不高但这其实抹平了很大一部分 FP8 的收益。怎么排查很简单看 profiler 里的算子名称。如果卷积输入 dtype 是float8_e4m3fn那是真的走了 FP8如果看到自动插入的aten::_to_copy把 dtype 从 FP8 换成 FP16那就说明回退了。7x7 这种卷积面积大的层直接把输入留在原始精度或单独量化反而更稳。5.2 权重落盘与 checkpoint 的 dtype 选择很多人习惯把模型权重直接to(torch.float8_e4m3fn)后torch.save。如果你只是为了减小 checkpoint 文件体积这个做法表面有效但我强烈不建议在大多数场景下这么做。原因是模型的推理或继续训练通常需要 FP32 或 BF16 的权重来保证稳定。FP8 是一种有损格式一旦把权重存成 FP8原本的参数细节里低于量化步长的信息就会彻底丢失。训练中断、加载 FP8 checkpoint 继续训练和从原始精度 checkpoint 继续训效果差异有时非常明显。FP8 权重更适合的是专为 FP8 推理部署这个场景这时候模型已经量化收敛对精度损失不太敏感。保险的做法是权重频率保存还是存 FP32/BF16然后额外存一份 FP8 推理权重如果你追求单文件可以在打包时同时存 FP8 权重和对应的 scale 元数据但落盘格式要用自定义的 dict 结构checkpoint { model_fp8: {name: t.to(torch.float8_e4m3fn) for name, t in model.state_dict().items()}, scales: scales, # 手动记录的每个 tensor 的 scale original_dtype: str(next(model.parameters()).dtype), } torch.save(checkpoint, model_fp8.pt)5.3 分布式训练里的 FP8 通信注意点数据并行训练里梯度同步的 allreduce 通信量和大模型训练时间密切相关。FP8 把梯度的数据量砍半理论上梯度通信耗时也能降一半但实际情况没有这么理想。我在实测分布式训练时发现NCCL 对 FP8 的 allreduce 支持情况各家版本并不同步有些环境里你把梯度转成 FP8reduce 时它照样先把数据转回 FP16 再走 NCCL通信字节数根本没降下来反而多了两次类型转换的开销。相对靠谱的降通信字节方案是 FP8 的 all_gather。比如模型并行里需要 gather 全量权重时FP8 能把 gather 通信量直接减半这个收益非常实在。如果你的梯度同步真的需要走 FP8 allreduce建议先跑一个小规模算力通信基准测试看清楚 NCCL 是否真的在 FP8 上工作而不是只做假 FP8的格式转换。6. 我实测的几组数据与版本建议6.1 硬件和软件版本矩阵FP8 最完整的性能收益目前只存在于 Hopper 及之后的 NVIDIA GPU 架构上包括 H100、H200 以及更新的 B100/B200。Ampere 架构A100没有原生的 FP8 张量核心PyTorch 虽然允许你创建 FP8 张量但计算要么不支持要么在软件层模拟 FP8那性能收益就完全体现不出来。软件方面PyTorch 2.1 正式引入 FP8 dtype2.2 和 2.3 补齐了更多算子支持2.4 往后的生态明显成熟。如果你纯粹做数值实验、验证 dtype 行为用 2.1 就够了但要在真实模型上做训练/推理建议至少 PyTorch 2.3配合torch.compile。cuDNN 版本也别忽略。FP8 卷积支持范围跟 cuDNN 版本强相关我在 8.9.x 遇到的一些回退问题升级到较新的 cuDNN 后部分解决了。所以前面说的 7x7 回退结论一定要结合自己环境的 cuDNN 版本来验证不要拿我的结论当通用标准。6.2 实测算力与显存对比我拿一个参数量大约 1.2B 的 Transformer 子结构做了前向加反向的对比测试条件是 batch size 固定、单张 H100、数据并行且只在单卡上跑一步不做通信精度方案前向反向显存占用训练步耗时相对FP3221.4 GB1.0xBF16 混合精度11.2 GB0.82xFP8E4M3 权重/激活E5M2 梯度8.6 GB0.71x显存节省确实接近一半。时间上单纯前向推理的场景 FP8 收益更明显而训练场景里如果反向梯度的缩放和回退处理得不好速度会打折。这个表格的意义是给你一个量级的参考实际模型结构不同比例会有差异。还需要提醒一句FP8 的收益对带宽受限的模型更明显。如果你的层里全是 1x1 卷积或小矩阵乘法计算密度低FP8 带来的吞吐提升会被访存成本掩盖反而是大矩阵乘法、大卷积这类计算密集模块收益最直观。6.3 PyTorch 版本演进的实际建议从我的使用经验来看FP8 在 PyTorch 里算是稳定的实验特性比很多第三方库要强太多。两个版本之间可能有些算子的 FP8 dispatch 行为会变比如之前回退的算子在新版本里可能支持了新配置。所以做 FP8 相关开发时我建议把 PyTorch、cuDNN、NCCL 的版本写成固定环境不要随便升级。真的升级之后第一件事是重新跑一遍 profiler看看有多少算子实际走了 FP8 路径尤其是卷积和 matmul。自定义算子方面如果你的模型有深度定制的大算子比如自己写的 flash attention 变体出 FP8 版本时要在torch.library层面做 dtype 支持声明。否则torch.compile看到 FP8 输入可能会直接报 unsupported ops这比回退还难受至少回退还有结果unsupported 是直接中断。我现在的个人习惯是先用torch.compile跑一版 FP8确认没有 dispatch 错误再用 profiler 逐层查看回退位置。最后才手动控制缩放和使用 Transformer Engine 层的部分替换。这个顺序下来遇到坑的概率会小很多。最后再分享一点我踩过最不值的一个坑是把 E4M3 和 E5M2 混用搞反了。权重和激活用 E4M3 没问题但梯度如果也用 E4M3小梯度很容易被舍入成零整个模型的低层几乎学不动。梯度一律 E5M2这是 FP8 训练里优先级最高的一条规则。你如果在自己的项目里试 FP8头一步先把这个定死会省掉后面百分之八十的麻烦。