企业官网建设流程全解析

Transformers FP8（W8A8）量化实战：基于 FBGEMM 的模型压缩与加载

【免费下载链接】transformers🤗 Transformers: the model-definition framework for state-of-the-art machine learning models in text, vision, audio, and multimodal models, for both inference and training.项目地址: https://gitcode.com/GitHub_Trending/tra/transformers

导读

本文讲解 Hugging Face Transformers 库中基于 FBGEMM（Facebook General Matrix Multiplication）的 FP8 量化方案：它面向 H100 等 compute capability ≥ 9.0 的 GPU 推理场景，将权重以 8-bit/通道（channel-wise）、激活以 8-bit/行（token-wise）压缩为 FP8（W8A8）格式。读完本文你将掌握如何安装 FBGEMM_GPU 环境、通过FbgemmFp8Config一键量化加载模型、保存并复用量化权重，以及理解其底层按行量化（per-row quantization）与 FP8 矩阵乘 kernel 的实现原理。本文核心依据为 官方量化文档，并补充了仓库源码与测试用例作为纵深佐证。

1. FBGEMM FP8 是什么

FBGEMM 是一套面向小 batch 推理的低精度矩阵乘法库，尤其擅长在保持精度的前提下压榨吞吐，内置了行级量化（row-wise quantization）与面向离群值的量化（outlier-aware quantization）等精度保护手段。

在 Transformers 中，FBGEMM FP8 量化意味着：

• 权重被量化为 8-bit，每个输出通道单独维护一个缩放因子（即文档所述weights to 8-bits/channel）；
• 激活在推理时被动态量化为 8-bit，按 token（即张量的行）单独缩放，量化前会以给定上界截断输入以抑制离群激活的影响；
• 这种 W8A8 的组合在文档中即被称为fp8。

关于数据格式，从源码看，量化后的权重以torch.float8_e4m3fn存储（4 位指数 + 3 位尾数的 8 位浮点），缩放因子保持float32，参见 fbgemm_fp8.py 集成模块。

硬件与软件前置条件

你需要一张compute capability ≥ 9.0的 GPU（如 NVIDIA H100）。之所以有硬性门槛，是因为validate_environment中会直接读取torch.cuda.get_device_capability()，当主版本号major < 9时抛错拒绝加载，见 quantizer_fbgemm_fp8.py。

除 GPU 外，官方还要求安装：

• fbgemm-gpu：提供 FP8 矩阵乘与量化 kernel；
• accelerate：负责权重加载前的空权重初始化与设备调度；
• PyTorch：配合支持 FP8 张量类型。

安装命令（原文档推荐，先升级到最新版再使用）：

pip install --upgrade accelerate fbgemm-gpu torch

如果安装遇到问题，官方建议尝试安装 fbgemm-gpu 的nightly 版本（其安装说明中对 CUDA/PyTorch 版本组合有更细的约束，按你所用的 PyTorch 构建对应 wheel 即可）。仓库测试中对运行环境也有等价约束：必须在具备加速器、安装 fbgemm-gpu（或 XPU + kernels）且安装了 accelerate 的前提下才允许执行，见 tests/quantization/fbgemm_fp8/test_fbgemm_fp8.py。

扩展说明（XPU 支持）：源码同时支持 Intel XPU 路径——在 XPU 上运行时不再依赖 fbgemm-gpu，而是通过kernels包动态拉取kernels-community/fp8-fbgemmkernel，并改用torch._scaled_mm完成 FP8 矩阵乘（该分支逻辑同样在 fbgemm_fp8.py 集成模块 与 quantizer 的环境校验中体现）。绝大多数用户走 CUDA + fbgemm-gpu 路线即可。

2. 五分钟上手：把模型量化成 FP8 并加载

原文档给出的核心用法非常简洁——实例化 [FbgemmFp8Config] 并通过AutoModelForCausalLM.from_pretrained传入quantization_config：

from transformers import FbgemmFp8Config, AutoModelForCausalLM quantization_config = FbgemmFp8Config() quantized_model = AutoModelForCausalLM.from_pretrained( "meta-llama/Meta-Llama-3-8B", dtype="auto", device_map="auto", quantization_config=quantization_config, )

几个参数的含义：

• dtype="auto"：沿用 checkpoint 自带的 torch_dtype。不过需要留意，当前实现的前向计算只支持 bfloat16——update_dtype会检测传入的 dtype，若不是torch.bfloat16则打印 warning 并强制改写为 bfloat16，参见 quantizer_fbgemm_fp8.py；
• device_map="auto"：让 accelerate 自动把模块分发到可用 GPU。由于量化 kernel 必须在 GPU/XPU 上运行，若不传 device_map，源码会警告你将模型放到 GPU/XPU 上；而在线量化阶段（权重尚未预量化）device_map 中不允许出现 CPU/disk，否则直接抛ValueError；
• quantization_config：承载量化策略的配置对象，若缺省会退化为普通加载。

在底层，Transformers 会通过量化器注册表把quant_method = "fbgemm_fp8"映射到FbgemmFp8HfQuantizer（见 quantizer_fbgemm_fp8.py）。其加载流程大体是：

1. 环境校验：确认 CUDA compute capability ≥ 9.0（或 XPU + kernels）、已装 accelerate 与 fbgemm-gpu；
2. 模块替换：在权重加载前，把模型中所有torch.nn.Linear递归替换为FbgemmFp8Linear（MoE 模型的专家层会被替换为专用的FbgemmFp8Llama4TextExperts），替换函数即replace_with_fbgemm_fp8_linear；
3. 边加载边量化：权重落地时逐参数完成 FP8 转换，并生成对应的weight_scale；
4. 后处理：把activation_scale_ub（输入激活缩放上界）写入各层。

该流程属于开箱即用的在线量化（on-the-fly），无需准备校准集——quantizer 中requires_calibration = False，这与需要激活统计的 AWQ 等方法形成鲜明对比。

2.1 推理调用

量化后的quantized_model仍是一个标准的PreTrainedModel，可以直接照常用generate/forward做文本生成或服务端推理，不必修改任何业务代码。

3. 保存与重新加载量化模型

save_pretrained与from_pretrained组合支持把量化后的权重持久化，这样下次加载时无需再次执行在线量化：

quant_path = "/path/to/save/quantized/model" quantized_model.save_pretrained(quant_path) # 之后从本地目录直接加载（自动识别其中的 fbgemm_fp8 量化配置） model = AutoModelForCausalLM.from_pretrained(quant_path, device_map="auto")

两个值得展开的细节：

1. 保存的是真正的 FP8 权重：FbgemmFp8Linear.weight以float8_e4m3fn参数存储，每个参数旁还会序列化一个fp32的weight_scale。因此导出的 checkpoint 相比 BF16 原始权重显著更小（权重部分单元素仅 1 字节，param_element_size对此有专门处理）。
2. 保存后即可预量化（pre-quantized）加载：当从已保存目录加载时，模型以pre_quantized=True路径进入，此时权重不再二次量化，device_map 的约束也放宽——仓库测试甚至把 32 层之外的lm_head、norm全部 offload 到 CPU/disk 后仍能正确加载并生成（见 test_fbgemm_fp8.py 的 save_pretrained_offload 用例）。

也就是说：在线量化适合第一次落地，保存后再加载适合生产部署，后者还允许你用加速器内存换磁盘/CPU offload 来部署超大模型。

4.FbgemmFp8Config参数详解

配置类定义在 src/transformers/utils/quantization_config.py（FbgemmFp8Config，quant_method = QuantizationMethod.FBGEMM_FP8）。它目前暴露两个核心参数：

activation_scale_ub的可调性有测试直接验证：以FbgemmFp8Config(activation_scale_ub=1000.0)重新加载保存过的模型后，检查model.model.layers[1].mlp.down_proj.input_scale_ub的值应等于1000.0，且生成结果与默认配置一致。这一机制依赖权重加载完成后对input_scale_ub缓冲区的强制回填（源码注释说明：加载与设备分发过程中持久缓冲区可能被清零或覆盖，因此_process_model_after_weight_loading在加载完成后统一fill_一次，见 quantizer_fbgemm_fp8.py）。

4.1 排除特定模块不量化

若某个模型包含不能量化成 FP8 的层，可以显式排除。测试中给出了一个直观示例：把 OPT 模型的所有torch.nn.Linear换成FbgemmFp8Linear后二者数量应完全相等；当传入modules_to_not_convert=["fc1"]时，量化后的FbgemmFp8Linear数量恰好减少 24（OPT-350M 的 12 个 block × 每个 block 2 个fc1层），证明排除规则逐模块生效。对应实现为replace_with_fbgemm_fp8_linear中的should_convert_module过滤逻辑，参见 fbgemm_fp8.py 集成模块。

5. 源码视角：FP8 前向究竟如何计算

理解集成实现（src/transformers/integrations/fbgemm_fp8.py）能帮你判断该方案适合哪些模型、出现精度问题时应调什么。

5.1 权重：静态 per-channel 量化

FbgemmFp8Linear.__init__中：

• weight参数以torch.float8_e4m3fn存储，形状(out_features, in_features)；
• weight_scale为(out_features, 1)的fp32参数——即每个输出通道一个缩放值，这就是"权重 8-bit/channel"；
• bias若存在则保持fp32；
• input_scale_ub是非持久的Buffer，即前文所述激活缩放上界。

5.2 激活：动态 per-row 量化

每次forward时，输入x会被view成 2D 后调用quantize_fp8_per_row（内部以input_scale_ub约束上界）得到量化激活与逐行缩放。代码特别加了设备对齐保护：由于量化 kernel 的产物不一定与输入同设备，须先把当前设备切到输入张量所在设备再启动 kernel，避免设备错位。

5.3 矩阵乘：FBGEMM rowwise FP8 kernel

在 CUDA 上，核心计算落到一行 kernel：

output = torch.ops.fbgemm.f8f8bf16_rowwise( x_quantized, self.weight, x_scale, weight_scale_float32, use_fast_accum=True ) output = output + self.bias if self.bias is not None else output

• f8f8bf16_rowwise：两个 FP8 输入矩阵相乘、结果累积为 BF16，两端各自携带 per-row 缩放；
• use_fast_accum=True：启用快速累加路径；
• 在 XPU 上则等价改走torch._scaled_mm并显式传入scale_a/scale_b与out_dtype；
• 计算完成后按(*x.shape[:-1], -1)还原输出形状，x_quantized、x_scale随即释放以省显存。

由此可见，每次推理前向都要动态量化一次激活，其开销由 FBGEMM 高度优化的 CUDA kernel 承担，从而换来 W8A8 的小 batch 高吞吐推理。

5.4 MoE 模型的专门处理：Llama4 专家层

对包含 MoE 的 Llama4 系列，实现中不再走普通 Linear，而是整层替换为FbgemmFp8Llama4TextExperts：

• gate_up_proj、down_proj均以 FP8 参数形式整体存于(num_experts, hidden_size, …)张量中，每专家一组；
• forward中对每个专家单独做激活量化与f8f8bf16_rowwise，并按 MoE 路由逻辑完成up * act_fn(gate)后再过一次down_proj；
• 权重转换阶段对gate_up_proj按行、down_proj按列分块执行quantize_fp8_per_row，并相应重塑缩放张量维度。

同时 quantizer 为 Llama4 定制了张量并行（TP）切分计划：由于 FBGEMM 操作不便像 dtensor 那样并行，注意力与共享专家采用colwise/rowwise本地切分后加 gather，专家gate_up_proj则采用packed_rowwise等策略（update_tp_plan）。也就是说，FBGEMM FP8 与 Llama4 的组合在 Tensor Parallel 场景下同样可用，但切分维度与普通模型不同，属于实现层的高级细节。

6. 精度、速度与适用边界（谨慎声明）

以下结论均以当前仓库源码与文档可验证的内容为限：

• 设计目标：FBGEMM 面向小 batch 服务器端推理做极致优化，FP8 由 GPU 原生算子加速。仓库并未在本文件或测试中给出任何量化前后精度对比的量化数值，请勿臆造 benchmark 数据；
• 数值约束：激活量化带activation_scale_ub（默认 1200.0）上界截断，这是控制精度与离群值的主要旋钮；权重保持 fp32 bias、lm_head类敏感模块可经modules_to_not_convert排除；
• 仅推理：quantizer 的is_trainable返回False，即该路径不支持在量化权重上继续训练/微调；
• 前向 dtype：仅 bfloat16；
• 序列化：is_serializable为True，支持save_pretrained持久化。

7. 结合测试用例验证一切

若想亲手验证整套流程，仓库在 tests/quantization/fbgemm_fp8/test_fbgemm_fp8.py 中提供了端到端用例（标记为slow，且需要 fbgemm-gpu 或 XPU+kernels），覆盖：

• test_quantized_model_conversion：OPT-350M 的 Linear 全量替换与modules_to_not_convert排除计数；
• test_quantized_model：Meta-Llama-3-8B在线量化后贪心生成的输出与预期一致；
• test_save_pretrained/test_save_pretrained_multi_gpu：保存后再加载（含多卡device_map="auto"）仍可复现输出；
• test_change_loading_attributes：验证activation_scale_ub能写穿到层内input_scale_ub；
• test_quantized_model_offload：验证在线量化阶段把设备映射到 CPU/disk 会正确报错，而保存为预量化 checkpoint 后（test_save_pretrained_offload）offload 合法。

这些用例即是最佳"运行说明书"：它们在功能层面证明"在线量化 → 保存 → 预量化加载（含 offload/多卡）"是完整闭环。

8. 更完整的应用：从量化到文本生成

把以上要点拼装起来，一个典型的端到端落地脚本长这样：

from transformers import FbgemmFp8Config, AutoModelForCausalLM, AutoTokenizer quantization_config = FbgemmFp8Config( activation_scale_ub=1200.0, modules_to_not_convert=["lm_head"], # 保留输出头精度 ) model = AutoModelForCausalLM.from_pretrained( "meta-llama/Meta-Llama-3-8B", dtype="auto", device_map="auto", quantization_config=quantization_config, ) tokenizer = AutoTokenizer.from_pretrained("meta-llama/Meta-Llama-3-8B") output = model.generate( **tokenizer("What are we having for dinner?", return_tensors="pt"), max_new_tokens=50, do_sample=False, ) print(tokenizer.decode(output[0], skip_special_tokens=True)) # 保存量化权重用于生产复用 model.save_pretrained("/path/to/save/quantized/model")

再次提醒运行前提：CUDA GPU 的 compute capability ≥ 9.0（如 H100）、已pip install --upgrade accelerate fbgemm-gpu torch、推理 dtype 为 bfloat16。若希望在 Transformers 量化方案中横向对比其他后端（如 AWQ、GPTQ、bitsandbytes 等），可参阅 量化总览 与 量化核心概念指南，二者对"weight-only 与 weight+activation 量化"的取舍有系统说明，有助于把 FBGEMM FP8（W8A8）放到正确的方案坐标中。

附：主要参考文件索引

• FBGEMM FP8 官方量化文档：本文核心骨架，含安装、量化加载、保存/重载三段最简用法；
• FbgemmFp8 集成实现：FbgemmFp8Linear、FbgemmFp8Llama4TextExperts、replace_with_fbgemm_fp8_linear与 per-row 量化/矩阵乘 kernel 调用链；
• FbgemmFp8HfQuantizer 量化器：环境校验（CC ≥ 9.0、bf16、device_map 约束）、模块替换、input_scale_ub回填、TP 计划；
• FbgemmFp8Config 配置类：activation_scale_ub、modules_to_not_convert两个参数的定义与默认值；
• FBGEMM FP8 端到端测试：覆盖转换、生成、保存/重载、多卡与 offload 的全部可运行示例。

【免费下载链接】transformers🤗 Transformers: the model-definition framework for state-of-the-art machine learning models in text, vision, audio, and multimodal models, for both inference and training.项目地址: https://gitcode.com/GitHub_Trending/tra/transformers