张量级FP8量化详解 | AI生成和翻译
问题：解释这个 _to_fp8() 函数以及张量级 FP8 量化背后的数学原理。
答案：
这是动态对称张量级量化。
目标是：
在尽可能保留动态范围的前提下，将浮点张量转换为 FP8，然后返回矩阵乘法中恢复原始值所需的缩放因子。
关键在于：
FP8 张量本身没有用。 它总是与一个缩放因子配对使用。
第 0 步：禁用梯度
@torch.no_grad()
量化不属于反向传播的一部分。
这告诉 PyTorch：
- 不构建 autograd 图
- 节省内存
- 略微加快速度
第 1 步
fp8_max = torch.finfo(fp8_dtype).max
假设
fp8_dtype = torch.float8_e4m3fn
则
fp8_max ≈ 448
对于另一种 FP8 格式
e5m2
max ≈ 57344
这仅仅是
FP8 能表示的最大数
第 2 步
amax = x.float().abs().max()
假设
x =
[
2,
-5,
9,
1
]
则
amax = 9
注意
只有一个标量
作用于整个张量
这就是为什么它被称为
张量级缩放 (Tensorwise scaling)
而不是
行级 (Rowwise)
列级 (Columnwise)
块级 (Blockwise)
第 3 步
scale = fp8_max / amax
假设
fp8_max = 448
amax = 9
则
scale = 49.777...
为什么？
因为我们需要
9
↓
448
所以每个值都要乘以
448 / 9
可以理解为将张量拉伸，直到其最大值恰好填满 FP8 的范围。
原始范围
[-9,9]
缩放后
[-448,448]
这最大化了精度，因为它几乎使用了全部可表示的范围。
为什么需要 clamp？
clamp(min=EPS)
假设
amax = 0
那么
448 / 0
会产生无穷大。
取而代之
amax = max(amax, EPS)
避免了除以零。
为什么使用双精度？
amax.double()
注释中提到
compile
vs
eager
问题在于
float32 除法
可能会产生微小的舍入差异。
例如
448 / 0.3
根据优化方式的不同，
1493.333374
或
1493.333252
这些微小差异可能导致数值落在 FP8 舍入边界的不同侧。
使用
float64
使得两种执行模式都具有确定性。
第 4 步
x_scaled = x.float() * scale
假设
x =
[
2,
5,
9
]
缩放因子
49.777
结果
[
99.5,
248.9,
448
]
现在张量完全占满了 FP8 的数值范围。
第 5 步
x_clamped = x_scaled.clamp(
-fp8_max,
fp8_max
)
假设数值误差给出了
449.2
FP8 无法表示该值。
如果不进行 clamp，
转换时可能溢出或产生未定义行为（取决于实现和 dtype）。Clamping 确保值在转换前保持在有效的 FP8 范围内。
第 6 步
x_fp8 = x_clamped.to(fp8_dtype)
现在转换
248.9
为 FP8。
请记住
FP8 具有
- 非常少的指数位
- 非常少的尾数位
因此
248.9
可能变成
248
或
256
具体取决于最近的 FP8 可表示值。
这就是引入量化误差的地方。
第 7 步
inv_scale = scale.reciprocal()
不是返回
scale
而是返回
1 / scale
为什么？
因为
torch._scaled_mm 期望：
real_value
=
fp8_value
×
inverse_scale
假设
原始值
9
缩放后
448
存储的
448 (FP8)
之后
448 × (9/448)
=
9
不需要显式的反量化步骤。在矩阵乘法过程中，torch._scaled_mm 会在内部将 FP8 操作数乘以各自的逆缩放因子，同时以更高精度累加。
整个流程
FP32 张量
│
▼
找出 max(abs(x))
│
▼
计算
scale
=
fp8_max / amax
│
▼
乘以张量
│
▼
Clamp
│
▼
转换为 FP8
│
▼
存储
(
FP8 张量,
逆缩放因子
)
为什么采用张量级缩放？
优点：
只有一个缩放因子
元数据极少
速度快
cuBLAS 直接支持
缺点：
假设
[
0.001,
0.002,
0.003,
100
]
缩放因子由
100
决定。
缩放后，
0.001
会变成一个非常小的 FP8 值，可能丢失大量精度，因为动态范围被大的离群值主导。
行级或块级缩放会计算多个缩放因子，为小值保留更多精度，但代价是增加额外的元数据和更复杂的内核。
为什么返回 (x_fp8, inv_scale) 而不是稍后反量化？
现代 FP8 GEMM 内核（例如 torch._scaled_mm 背后的内核）融合了计算：
FP8 A
FP8 B
↓
相乘
↓
应用逆缩放因子
↓
在 FP16/BF16/FP32 中累加
↓
输出
这避免了在内存中实例化反量化的 FP16 张量，减少了带宽并提高了性能。硬件实际上计算：
[ C = (A_{fp8} \times s_A^{-1}) \times (B_{fp8} \times s_B^{-1}) ]
同时以更高精度执行累加，这就是为什么该函数同时返回量化后的张量及其逆缩放因子，而不是先重建原始张量。