Diffusers documentation

BFL(fp8r / nvfp4)

Hugging Face's logo
Join the Hugging Face community

and get access to the augmented documentation experience

to get started

BFL(fp8r / nvfp4)

Black Forest Labs 为其 Flux3 模型提供了两种量化方案,以单文件检查点的形式与原始 bf16 权重一同存放:

  • fp8r:行级 FP8。每个量化线性层存储一个 E4M3 格式的 weight 和一个 fp32 格式的按行 weight_scale。激活值在运行时按行量化为 E4M3。
  • nvfp4:ModelOpt 布局下的 NVFP4。每个量化线性层存储打包的 E2M1 weight(每字节两个值)、16 个 E4M3 块缩放因子(weight_scale)、一个全局 fp32 weight_scale_2,以及一个经过校准的 fp32 激活缩放因子 input_scale。

Diffusers 通过 ~FromSingleFileMixin.from_single_file 加载两者,并使用 BFLQuantizationConfig。量化方案从检查点中读取,因此相同的配置适用于任一文件。仅支持加载预量化检查点,不支持在流水线级别加载。

以下示例加载 FLUX 3 Image Klein Base 的 fp8r 检查点,该检查点位于模型仓库的 variants/fp8r/model.safetensors 路径下。传入 config 以从该仓库的 Diffusers 文件夹中读取模型配置,然后按名称加载其他组件(参见 单文件加载)。

import torch

from diffusers import BFLQuantizationConfig, Flux3Transformer2DModel, ModularPipeline

ckpt_path = "https://huggingface.co/black-forest-labs/flux-3-image-klein-base/blob/main/variants/fp8r/model.safetensors"
transformer = Flux3Transformer2DModel.from_single_file(
    ckpt_path,
    config="black-forest-labs/flux-3-image-klein-base",
    subfolder="transformer",
    quantization_config=BFLQuantizationConfig(compute_dtype=torch.bfloat16),
    dtype=torch.bfloat16,
)
pipe = ModularPipeline.from_pretrained("black-forest-labs/flux-3-image-klein-base")
pipe.load_components(dtype=torch.bfloat16)
pipe.update_components(transformer=transformer)
pipe.to("cuda")

prompt = "A wide-angle wildlife photograph of a red fox standing alert in a snow-covered forest clearing at golden hour."
image = pipe(prompt=prompt, generator=torch.Generator().manual_seed(0), output="images")[0]
image.save("flux3-klein-base-fp8r.png")

将 variants/fp8r 替换为 variants/nvfp4 以加载 nvfp4 变体。步数蒸馏版 FLUX 3 Image Klein 提供相同的两个变体;加载时,将其仓库 id 作为 config 和 pipeline 传入。

内核

量化后的权重保持在其低精度存储 dtype 中。具体执行哪种 GEMM 取决于设备:

方案融合内核要求
fp8r使用行级缩放的 torch._scaled_mm计算能力为 8.9 或更高的 CUDA GPU
nvfp4flashinfer.mm_fp4(CUTLASS)配合 flashinfer.nvfp4_quantize 用于激活值Blackwell GPU(计算能力 10.0 或更高)、包含 nvcc 的 CUDA 工具包,以及版本相同的 flashinfer-python 和 flashinfer-cubin

请同时安装相同版本的两个 FlashInfer 软件包,例如 pip install flashinfer-python==0.6.13 flashinfer-cubin==0.6.13。flashinfer-cubin 的发布版本可能会落后于 flashinfer-python,当两个包的版本不一致时,FlashInfer 会拒绝导入。

FlashInfer 在首次使用时,会使用当前环境的 CUDA 工具包编译其 NVFP4 内核,并根据可见的 GPU(或当设置了 FLASHINFER_CUDA_ARCH_LIST 时,根据该变量)选择目标架构。消费级和工作站 Blackwell GPU(计算能力 12.0)需要 12.9 或更高版本的 CUDA 工具包。如果使用较旧的工具包,FlashInfer 将找不到符合条件的架构,首次 NVFP4 调用会因 FlashInfer requires GPUs with sm75 or higher 而失败。

在其他任何设备上,权重会在每次前向传播中反量化为 compute_dtype,该层作为常规线性层运行。在这种情况下,激活值不会被量化,因此输出与融合内核的结果略有差异。

Update on GitHub