使用FP8进行大模型量化原理及实践

近年来，随着Transformer和MOE架构的相继提出，深度学习模型的参数量轻松突破了万亿大关。模型越大，部署成本越高，推理性能也越受挑战。怎么办？这就引出了大模型压缩技术——我们总得想点办法，让这些庞然大物既能干活，又不至于吃掉所有的算力资源。模型压缩通常能分为这么几大类：剪枝（Pruning）、知识蒸馏（Knowledge Distillation），以及今天我们要重点聊的量化（Quantization）。 这个系列会逐一拆解一些主流的大模型量化方案，比如GPTQ、LLM.int8()、SmoothQuant、AWQ等等。本文就先从FP8量化说起，聊聊它背后的原理和实际应用。 --- ## FP8 简介 FP8，全称是8位浮点数据格式，由Nvidia、Arm、Intel三家联合推出，专门为了加速深度学习的训练和推理而生。他们当时还发了一篇论文，就叫《FP8 Formats for Deep Learning》。现在，Nvidia的Hopper和Ada Lovelace架构的GPU，都已经原生支持这种格式了。 ## 补充：浮点表示法 在正式聊FP8之前，有必要先回顾一下浮点数的表示方法。根据IEEE 754标准，一个浮点数在计算机里由三部分组成： - **符号位**：1bit，0表示正数，1表示负数。 - **指数部分**：决定了数值的范围。 - **尾数部分**：决定了数值的精度。 浮点数的计算公式其实很简单，但根据指数位的取值，又可分为规格化数、非规格化数以及特殊值（无穷和NaN）。对于规格化数，指数位不全为0或全为1；非规格化数则是指数位全为0；特殊值嘛，指数全为1，尾数全0表示无穷，尾数不全为0就是NaN。 ## FP8 表示形式 在硬件层面，FP8通常有两种变体：E4M3和E5M2。这两种格式各有各的用武之地。 - **E4M3**：1个符号位，4个指数位，3个尾数位。它能表示的最大值是448，也能表示NaN。 - **E5M2**：1个符号位，5个指数位，2个尾数位。它能表示到±57344，还能表示无穷和NaN。 简单来说，E4M3精度更高，E5M2动态范围更大——范围大了，代价就是精度少一些。 有意思的是，E5M2的指数部分其实和FP16是一致的，这意味着它的动态范围和FP16相当，因此常被用在训练的反向传播阶段。而E4M3则更适合前向传播。 在H100的第四代Tensor Core上，FP8的矩阵乘法（无论是E4M3还是E5M2）都能做，甚至混合搭配也行，计算结果最终会累加到FP32或FP16中。格式间的转换，硬件也一并支持。 E5M2完全遵循IEEE 754标准；E4M3则没那么死板——它允许指数位全为1时仍然表示规格化数，只有当指数和尾数全为1时才表示NaN，而且它本身就不能表示无穷。 举个例子： - 二进制 (0 1111 110) 在E4M3里，表示 2^(15-7) * (1+1/2 + 1/4) = 448 - 二进制 (1 1111 100) 则代表 -384 - 二进制 (1 1111 111) 就是NaN 换个角度看，我们可以把浮点数看作2的幂之间的若干样本。比如在E5M2中，2和4之间只有4个样本；而在E4M3中，同样区间却有8个样本。这意味着，E5M2的量化误差比E4M3大。另外，FP8表示的数值本身是非均匀的——越靠近0，分布越密；越远离0，越稀疏。所以，浮点量化的误差会随着数值的增大而增大。 ## INT8 与 FP8 对比 ### 精度对比 到底是INT8好还是FP8强？这得看数据分布。高通AI研究院在这方面有过专门的研究。 - 对于**均匀分布**，INT8表现最好，FP8-E4和FP8-E5都差一些。 - 对于**正态分布**，反而是FP8-E2（2个指数位、5个尾数位）最准，INT8紧随其后。 - 而面对带有异常值的**t分布**，大家精度都一般，不过FP8-E4稍微好一点。 具体到模型权重和激活值的量化，情况就更有趣了。有一篇论文统计了LLaMA-65B各层权重的量化偏差，结果发现：**对于权重，INT8明显好于FP8-E4**。但对于激活值，情况就反过来了——**FP8-E4在大多数层上的精度优于INT8**。 为什么会这样？原因是激活值是动态的，每次输入都不一样，需要靠校准集来决定量化尺度。校准过程通常是取所有batch的最大值来计算尺度，这样一来，非最大值的batch数值普遍偏小，而FP8-E4对小数值的精度更好。所以，**INT8和FP8并没有谁绝对占优**，完全可以根据不同层的实际损失去灵活选择。 ### 性能对比 以Nvidia L40（Ada Lovelace）为例，它的FP8 Tensor Core和INT8 Tensor Core在计算能力上是一致的。H800、L20、L40S也类似。 但如果深入到硬件层面——乘法累加运算单元（MAC）——就有差别了。对于8位MAC，FP运算通常比INT占用更多的芯片面积。不过有意思的是，当位宽降到8位时，FP8和INT8 MAC单元的面积几乎一样。这意味着，两者的硬件成本和推理性能其实很接近。 累加器是另一个关键因素。定点累加器对INT8来说既精确又高效；而浮点累加器虽然不那么精确，但对FP8-E4这种动态范围大的格式更友好——它不需要大幅增加累加器的大小就能覆盖更大的乘积范围。 从综合数据看，INT8加上定点累加器是最节省面积的方案。FP8-E4如果搭配浮点累加器，需要的门电路数量比INT8高出53%；如果是FP32累加器，效率甚至要低183%。所以，纯看硬件效率，FP8-E4比INT8更贵。但这只是其中一面。 ### 综合对比 1. INT8是均匀量化，所有数值等间隔映射，容易损失边界特征值；FP8是非均匀量化，小数值更精确，大数值则误差更大。 2. 混合精度训练时，FP8转成FP16/FP32/BF16更直接；INT8则需要额外的乘法和加法操作，开销更大。 3. 在CPU上，FP8比INT8慢得多。 4. 量化误差和步长正相关——FP8对小数值友好，但对大数值就不太友好了。 ## FP8 量化方案 ### FP8 量化模拟 高通AI研究院在研究中提出了一种在FP32硬件上模拟FP8量化的方法，可以快速完成实验，并且很容易在深度学习框架里实现。结论是：FP8格式在多种网络模型的PTQ（训练后量化）中，通常优于INT8；其中E5M2和E4M3效果最好。而对于像Transformer这类含有较多异常值的模型，增加指数位数效果更好。不过，一旦做QAT（量化感知训练），FP8格式的优势会缩小——网络模型在训练中学会了在INT8下也能表现不错。 ### MoFQ 混合格式量化方法 不同层的数据分布差异很大，没有一种量化格式能一路通吃。上海交大、北大和微软亚洲研究院联合提出的MoFQ方法，采用逐层选择最优格式的策略，在W8A8量化上取得了接近全精度的结果。而且，MoFQ不会增加硬件开销——位宽还是8比特，换的只是格式。 ## FP8 推理流程 训练时，为了梯度计算的准确性，权重通常保持在高精度（BF16或FP32）。但推理时权重是固定的，所以可以在模型加载时就把它转成FP8。 操作融合也很重要。比如把LayerNorm和后续的数据格式转换合并，确保kernel的输入输出尽可能地是FP8，这样能有效提升GPU内存带宽利用率。GeLU激活函数也一样，尽量融合进去。 目前，少量的输出还是会保持FP16——因为NCCL只支持高精度的reduce操作。所以，流程大概是：LayerNorm输出和权重都是FP8，矩阵运算结果暂时保持FP16，完成后再转成FP8。整个推理链路里，绿线代表FP8的I/O，红线代表高精度的I/O。 ## FP8 在 TensorRT-LLM 中的应用 TensorRT-LLM从9.0版本开始支持FP8推理。想启用FP8，需要几个步骤：设置FP8标志、添加GEMM的缩放因子（scale）、编写FP8模型。具体来说，就是把FP16输入量化为FP8再反量化，权重也一样操作。TensorRT会自动把量化/反量化操作和前后的kernel融合，最终的计算图就是量化的X和W直接做FP8计算，输出也是FP8。 TensorRT-LLM还封装了FP8 linear和FP8 row linear函数，直接用线性层就行，不用自己写代码。流程大致是：权重以FP8存储，计算前先反量化（变成FP16甚至FP32），然后做矩阵乘法，累加用FP32完成，最后再乘以scale，输出FP8的结果。 ## FP8 在 vLLM 中的应用 vLLM在H100和AMD MI300x等GPU上支持FP8量化，目前只支持Hopper和Ada Lovelace架构。使用FP8可以让模型内存占用减半，吞吐量提升1.6倍，精度影响很小。 ### 在线动态量化 不需要校准数据，直接在启动推理引擎时指定 `--quantization="fp8"` 就行。所有Linear模块（除了lm_head）的权重都按per-tensor量化为FP8_E4M3，激活值则在每次前向传播时动态计算。缺点是延迟提升有限，而且需要足够的内存先加载原始精度的模型。 ### 离线动态量化 用AutoFP8库提前把权重量化为FP8（E4M3）并生成检查点，vLLM在运行时只处理激活的动态缩放。这样精度更高，内存占用也更低。 ### 离线静态量化 为了最佳推理性能，可以用AutoFP8加上校准数据，为权重和激活都生成静态的per-tensor缩放因子。然后在vLLM中直接加载量化后的模型检查点。 ### KV Cache FP8 量化 vLLM还支持对KV Cache进行FP8量化，分为E5M2和E4M3两种。E4M3精度更高，但数据范围小，通常需要每个量化tensor旁边放一个FP32的缩放因子。如果指定了量化参数JSON，缩放因子可以自动加载；没指定的话默认就是1.0。研究表明，FP8 E4M3的KV Cache量化带来的精度损失很小，但推理速度能显著提升。 ## FP8 在不同推理框架及硬件的性能对比 在TensorRT-LLM中，FP16的max值能做到75，而FP8能提升到85——主要因为权重内存省了。如果KV Cache也转成FP8，内存再减半，吞吐量甚至可以接近FP16的两倍。 另外，用vLLM对Qwen1.5-Chat在不同硬件（4090、H800、H20）上做了FP8静态量化测试。结果很直观：相比BF16，FP8在所有硬件上都有10%~30%的性能提升。 ## 总结 本文从FP8的基本概念讲起，对比了它和INT8在精度与性能上的差异，也介绍了几种主流的FP8量化方案。虽然FP8还在不断演进中，但主流的推理框架，比如TensorRT-LLM和vLLM，都已经提供了成熟的支持。实测数据也说明，在不同硬件上，FP8相对于BF16都能带来实实在在的性能提升。量化这件事，选对了路，效果立竿见影。