Buckets:

mishig's picture
|
download
raw
4.68 kB

BFL (fp8r / nvfp4)

Black Forest Labs는 Flux3 transformers의 양자화 변형을 두 가지 방식으로 제공하며, 원래의 bf16 가중치 옆에 단일 파일 체크포인트로 저장합니다:

  • fp8r: rowwise FP8. 각 양자화 선형 레이어는 E4M3 weight와 행별 fp32 weight_scale를 저장합니다. 활성화 값은 런타임에 행별로 E4M3로 양자화됩니다.
  • nvfp4: ModelOpt 레이아웃의 NVFP4. 각 양자화 선형 레이어는 패킹된 E2M1 weight(바이트당 두 값), 16개의 E4M3 블록 스케일(weight_scale), 전역 fp32 weight_scale_2, 그리고 보정된 fp32 활성화 스케일 input_scale을 저장합니다.

Diffusers는 BFLQuantizationConfig와 함께 ~FromSingleFileMixin.from_single_file를 통해 두 가지를 모두 로드합니다. 스킴은 체크포인트에서 읽히므로 동일한 구성이 두 파일 모두에 적용됩니다. 사전 양자화(checkpoint) 로드만 지원되며, 파이프라인 수준의 로드는 지원하지 않습니다.

아래 예시는 모델 저장소의 variants/fp8r/model.safetensors에 있는 FLUX 3 Image Klein Base의 fp8r 체크포인트를 로드합니다. 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 및 파이프라인으로 사용하여 로드합니다.

커널

양자화된 가중치는 저정밀도 저장 dtype에 그대로 유지됩니다. 어떤 GEMM이 실행되는지는 장치에 따라 달라집니다:

스키마 Fused kernel 요구 사항
fp8r rowwise 스케일을 사용하는 torch._scaled_mm compute capability 8.9 이상의 CUDA GPU
nvfp4 flashinfer.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 toolkit으로 NVFP4 커널을 컴파일하며, 가시적인 GPU(또는 설정된 경우 FLASHINFER_CUDA_ARCH_LIST)에서 대상 아키텍처를 선택합니다. 소비자용 및 워크스테이션용 Blackwell GPU(compute capability 12.0)에는 12.9 이상의 CUDA toolkit이 필요합니다. 구버전 toolkit을 사용하는 경우 FlashInfer는 적합한 아키텍처를 찾지 못하며, 첫 NVFP4 호출이 FlashInfer requires GPUs with sm75 or higher 오류로 실패합니다.

다른 모든 장치에서는 각 forward pass마다 가중치가 compute_dtype으로 역양자화되고, 해당 레이어는 일반적인 linear로 실행됩니다. 이 경우 활성화 값은 양자화되지 않으므로, 출력 값이 fused kernel의 출력과 약간 다를 수 있습니다.

Xet Storage Details

Size:
4.68 kB
·
Xet hash:
a4d2cee6d62ea6152cdc4b75415dee94ec5c584d88d61e96a20e73f366b71e7d

Xet efficiently stores files, intelligently splitting them into unique chunks and accelerating uploads and downloads. More info.