Transformers FP8W8A8量化实战基于 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 库中基于 FBGEMMFacebook General Matrix Multiplication的 FP8 量化方案它面向 H100 等 compute capability ≥ 9.0 的 GPU 推理场景将权重以 8-bit/通道channel-wise、激活以 8-bit/行token-wise压缩为 FP8W8A8格式。读完本文你将掌握如何安装 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 矩阵乘与量化 kernelaccelerate负责权重加载前的空权重初始化与设备调度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_configfrom transformers import FbgemmFp8Config, AutoModelForCausalLM quantization_config FbgemmFp8Config() quantized_model AutoModelForCausalLM.from_pretrained( meta-llama/Meta-Llama-3-8B, dtypeauto, device_mapauto, quantization_configquantization_config, )几个参数的含义dtypeauto沿用 checkpoint 自带的 torch_dtype。不过需要留意当前实现的前向计算只支持 bfloat16——update_dtype会检测传入的 dtype若不是torch.bfloat16则打印 warning 并强制改写为 bfloat16参见 quantizer_fbgemm_fp8.pydevice_mapauto让 accelerate 自动把模块分发到可用 GPU。由于量化 kernel 必须在 GPU/XPU 上运行若不传 device_map源码会警告你将模型放到 GPU/XPU 上而在线量化阶段权重尚未预量化device_map 中不允许出现 CPU/disk否则直接抛ValueErrorquantization_config承载量化策略的配置对象若缺省会退化为普通加载。在底层Transformers 会通过量化器注册表把quant_method fbgemm_fp8映射到FbgemmFp8HfQuantizer见 quantizer_fbgemm_fp8.py。其加载流程大体是环境校验确认 CUDA compute capability ≥ 9.0或 XPU kernels、已装 accelerate 与 fbgemm-gpu模块替换在权重加载前把模型中所有torch.nn.Linear递归替换为FbgemmFp8LinearMoE 模型的专家层会被替换为专用的FbgemmFp8Llama4TextExperts替换函数即replace_with_fbgemm_fp8_linear边加载边量化权重落地时逐参数完成 FP8 转换并生成对应的weight_scale后处理把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_mapauto)两个值得展开的细节保存的是真正的 FP8 权重FbgemmFp8Linear.weight以float8_e4m3fn参数存储每个参数旁还会序列化一个fp32的weight_scale。因此导出的 checkpoint 相比 BF16 原始权重显著更小权重部分单元素仅 1 字节param_element_size对此有专门处理。保存后即可预量化pre-quantized加载当从已保存目录加载时模型以pre_quantizedTrue路径进入此时权重不再二次量化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.pyFbgemmFp8Configquant_method QuantizationMethod.FBGEMM_FP8。它目前暴露两个核心参数参数默认值说明activation_scale_ub1200.0输入激活量化的缩放上界activation scale upper bound。在对激活做 per-row 量化前用于约束输入范围、抑制离群值对精度的影响同时它也会同步写入各量化层的input_scale_ub缓冲区作为激活量化 kernel 的截断上界参与前向计算。modules_to_not_convertNone保持原精度的模块名列表。用于某些必须保留原始精度的层实践中通常保留lm_head等输出头以保证数值稳定性。在参数替换时该列表还会与模型自带的需要 fp32 保留的模块集合取并集。activation_scale_ub的可调性有测试直接验证以FbgemmFp8Config(activation_scale_ub1000.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数量恰好减少 24OPT-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/channelbias若存在则保持fp32input_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 上核心计算落到一行 kerneloutput torch.ops.fbgemm.f8f8bf16_rowwise( x_quantized, self.weight, x_scale, weight_scale_float32, use_fast_accumTrue ) output output self.bias if self.bias is not None else outputf8f8bf16_rowwise两个 FP8 输入矩阵相乘、结果累积为 BF16两端各自携带 per-row 缩放use_fast_accumTrue启用快速累加路径在 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而是整层替换为FbgemmFp8Llama4TextExpertsgate_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 或 XPUkernels覆盖test_quantized_model_conversionOPT-350M 的 Linear 全量替换与modules_to_not_convert排除计数test_quantized_modelMeta-Llama-3-8B在线量化后贪心生成的输出与预期一致test_save_pretrained/test_save_pretrained_multi_gpu保存后再加载含多卡device_mapauto仍可复现输出test_change_loading_attributes验证activation_scale_ub能写穿到层内input_scale_ubtest_quantized_model_offload验证在线量化阶段把设备映射到 CPU/disk 会正确报错而保存为预量化 checkpoint 后test_save_pretrained_offloadoffload 合法。这些用例即是最佳运行说明书它们在功能层面证明在线量化 → 保存 → 预量化加载含 offload/多卡是完整闭环。8. 更完整的应用从量化到文本生成把以上要点拼装起来一个典型的端到端落地脚本长这样from transformers import FbgemmFp8Config, AutoModelForCausalLM, AutoTokenizer quantization_config FbgemmFp8Config( activation_scale_ub1200.0, modules_to_not_convert[lm_head], # 保留输出头精度 ) model AutoModelForCausalLM.from_pretrained( meta-llama/Meta-Llama-3-8B, dtypeauto, device_mapauto, quantization_configquantization_config, ) tokenizer AutoTokenizer.from_pretrained(meta-llama/Meta-Llama-3-8B) output model.generate( **tokenizer(What are we having for dinner?, return_tensorspt), max_new_tokens50, do_sampleFalse, ) print(tokenizer.decode(output[0], skip_special_tokensTrue)) # 保存量化权重用于生产复用 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 与 weightactivation 量化的取舍有系统说明有助于把 FBGEMM FP8W8A8放到正确的方案坐标中。附主要参考文件索引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创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考