Buckets:

mishig's picture
|
download
raw
4.5 kB

BFL (fp8r / nvfp4)

Black Forest Labs disponibiliza variantes quantizadas de seus transformers Flux3 em dois esquemas, armazenadas como checkpoints de arquivo único ao lado dos pesos originais em bf16:

  • fp8r: FP8 por linha. Cada linear quantizado armazena um weight E4M3 e uma weight_scale fp32 por linha. As ativações são quantizadas para E4M3 por linha em tempo de execução.
  • nvfp4: NVFP4 no layout do ModelOpt. Cada linear quantizado armazena um weight E2M1 empacotado (dois valores por byte), escalas de bloco E4M3 de 16 (weight_scale), uma weight_scale_2 global fp32 e uma escala de ativação fp32 calibrada input_scale.

O Diffusers carrega ambos por meio de ~FromSingleFileMixin.from_single_file com um BFLQuantizationConfig. O esquema é lido do checkpoint, portanto, a mesma configuração funciona para qualquer um dos arquivos. Apenas o carregamento de checkpoints pré-quantizados é suportado; o carregamento em nível de pipeline não é.

O exemplo abaixo carrega o checkpoint fp8r do FLUX 3 Image Klein Base, que está em variants/fp8r/model.safetensors no repositório do modelo. Passe config para que a configuração do modelo seja lida da pasta Diffusers desse repositório e, em seguida, carregue os outros componentes pelo nome (veja Carregamento de arquivo único).

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")

Substitua variants/fp8r por variants/nvfp4 para carregar a variante nvfp4. O FLUX 3 Image Klein com destilação por etapas inclui as mesmas duas variantes; carregue-as usando o id do repositório dele como config e pipeline.

Kernels

Os pesos quantizados permanecem em seu dtype de armazenamento de baixa precisão. Qual GEMM é executado depende do dispositivo:

Esquema Kernel fundido Requisitos
fp8r torch._scaled_mm com escalas por linha GPU CUDA com capacidade de computação 8.9 ou superior
nvfp4 flashinfer.mm_fp4 (CUTLASS) com flashinfer.nvfp4_quantize para as ativações GPU Blackwell (capacidade de computação 10.0 ou superior), um toolkit CUDA com nvcc e flashinfer-python com flashinfer-cubin na mesma versão

Instale os dois pacotes do FlashInfer na mesma versão, por exemplo pip install flashinfer-python==0.6.13 flashinfer-cubin==0.6.13. A liberação do flashinfer-cubin pode ficar para trás em relação ao flashinfer-python, e o FlashInfer se recusa a importar quando as versões diferem.

FlashInfer compila seus kernels NVFP4 no primeiro uso com o toolkit CUDA do ambiente e seleciona a arquitetura de destino a partir da GPU visível (ou de FLASHINFER_CUDA_ARCH_LIST, quando definido). GPUs Blackwell de consumo e de estação de trabalho (capacidade de computação 12.0) exigem um toolkit CUDA 12.9 ou mais recente. Com um toolkit mais antigo, o FlashInfer não encontra nenhuma arquitetura elegível e a primeira chamada NVFP4 falha com FlashInfer requires GPUs with sm75 or higher.

Em qualquer outro dispositivo, os pesos são desquantizados para compute_dtype em cada passagem de frente e a camada é executada como uma linear regular. Nesse caso, as ativações não são quantizadas, portanto as saídas diferem ligeiramente dos kernels fundidos.

Xet Storage Details

Size:
4.5 kB
·
Xet hash:
f1f743916c5e1a51f8015ccbbc40ec6998202b09a9f7a8027fdcb3c825e5dd40

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