bitsandbytes

bitsandbytes é a opção mais simples para quantizar um modelo para 8 e 4 bits. A quantização de 8 bits multiplica os valores extremos em fp16 pelos valores não extremos em int8, converte os valores não extremos de volta para fp16 e, em seguida, os soma para retornar os pesos em fp16. Isso reduz o efeito degradante que os valores extremos têm no desempenho de um modelo.

A quantização de 4 bits comprime ainda mais um modelo e é comumente usada com QLoRA para fazer o ajuste fino de LLMs quantizados.

Este guia demonstra como a quantização pode permitir a execução do FLUX.1-dev com menos de 16GB de VRAM e até mesmo em uma instância gratuita do Google Colab.

Para usar o bitsandbytes, verifique se as seguintes bibliotecas estão instaladas:

pip install diffusers transformers accelerate bitsandbytes -U

Agora você pode quantizar um modelo passando um BitsAndBytesConfig para from_pretrained(). Isso funciona para qualquer modelo em qualquer modalidade, desde que ele suporte carregamento com Accelerate e contenha camadas torch.nn.Linear.

Quantizar um modelo em 8 bits reduz o uso de memória pela metade:

bitsandbytes é suportado tanto no Transformers quanto no Diffusers, permitindo quantizar tanto o FluxTransformer2DModel quanto o T5EncoderModel.

Para GPUs da série Ada e superiores, recomendamos alterar dtype para torch.bfloat16.

O CLIPTextModel e o AutoencoderKL não são quantizados porque já são pequenos em tamanho e porque o AutoencoderKL possui apenas algumas camadas torch.nn.Linear.

Copiado

from diffusers import BitsAndBytesConfig as DiffusersBitsAndBytesConfig
from transformers import BitsAndBytesConfig as TransformersBitsAndBytesConfig
import torch
from diffusers import AutoModel
from transformers import T5EncoderModel

quant_config = TransformersBitsAndBytesConfig(load_in_8bit=True,)

text_encoder_2_8bit = T5EncoderModel.from_pretrained(
    "black-forest-labs/FLUX.1-dev",
    subfolder="text_encoder_2",
    quantization_config=quant_config,
    dtype=torch.float16,
)

quant_config = DiffusersBitsAndBytesConfig(load_in_8bit=True,)

transformer_8bit = AutoModel.from_pretrained(
    "black-forest-labs/FLUX.1-dev",
    subfolder="transformer",
    quantization_config=quant_config,
    dtype=torch.float16,
)

Por padrão, todos os outros módulos, como torch.nn.LayerNorm, são convertidos para torch.float16. Você pode alterar o tipo de dados desses módulos com o parâmetro dtype.

Copiado

transformer_8bit = AutoModel.from_pretrained(
    "black-forest-labs/FLUX.1-dev",
    subfolder="transformer",
    quantization_config=quant_config,
+   dtype=torch.float32,
)

Vamos gerar uma imagem usando nossos modelos quantizados.

Definir device_map="auto" preenche automaticamente todo o espaço disponível primeiro nas GPUs, depois na CPU e, por fim, no disco rígido (a opção mais lenta) se ainda não houver memória suficiente.

Copiado

from diffusers import FluxPipeline

pipe = FluxPipeline.from_pretrained(
    "black-forest-labs/FLUX.1-dev",
    transformer=transformer_8bit,
    text_encoder_2=text_encoder_2_8bit,
    dtype=torch.float16,
    device_map="auto",
)

pipe_kwargs = {
    "prompt": "A cat holding a sign that says hello world",
    "height": 1024,
    "width": 1024,
    "guidance_scale": 3.5,
    "num_inference_steps": 50,
    "max_sequence_length": 512,
}

image = pipe(**pipe_kwargs, generator=torch.manual_seed(0),).images[0]

Quando há memória suficiente, você também pode mover diretamente o pipeline para a GPU com .to("cuda") e aplicar enable_model_cpu_offload() para otimizar o uso de memória da GPU.

Uma vez que um modelo é quantizado, você pode enviar o modelo para o Hub com o método push_to_hub(). O arquivo de configuração de quantização config.json é enviado primeiro, seguido pelos pesos do modelo quantizado. Você também pode salvar os modelos 8-bit serializados localmente com save_pretrained().

O treinamento com pesos de 8 bits e 4 bits é suportado apenas para o treinamento de parâmetros extras.

Verifique a sua pegada de memória com o método get_memory_footprint:

print(model.get_memory_footprint())

Observe que isso apenas informa a pegada de memória dos parâmetros do modelo e não estima os requisitos de memória para inferência.

Modelos quantizados podem ser carregados a partir do método from_pretrained() sem a necessidade de especificar os parâmetros de quantization_config:

from diffusers import AutoModel, BitsAndBytesConfig

quantization_config = BitsAndBytesConfig(load_in_4bit=True)

model_4bit = AutoModel.from_pretrained(
    "hf-internal-testing/flux.1-dev-nf4-pkg", subfolder="transformer"
)

8-bit (algoritmo LLM.int8())

Saiba mais sobre os detalhes da quantização de 8 bits neste post de blog!

Esta seção explora alguns dos recursos específicos dos modelos de 8 bits, como limites de outliers e a conversão de módulos ignorada.

Limiar de outlier

Um “outlier” é um valor de estado oculto maior que um determinado limiar, e esses valores são calculados em fp16. Embora os valores geralmente sigam uma distribuição normal ([-3.5, 3.5]), essa distribuição pode ser muito diferente para modelos grandes ([-60, 6] ou [6, 60]). A quantização de 8 bits funciona bem para valores ~5, mas, além disso, há uma penalidade significativa de desempenho. Um bom valor de limiar padrão é 6, mas um limiar menor pode ser necessário para modelos mais instáveis (modelos pequenos ou finetuning).

Para encontrar o melhor limiar para o seu modelo, recomendamos experimentar com o parâmetro llm_int8_threshold em BitsAndBytesConfig:

from diffusers import AutoModel, BitsAndBytesConfig

quantization_config = BitsAndBytesConfig(
    load_in_8bit=True, llm_int8_threshold=10,
)

model_8bit = AutoModel.from_pretrained(
    "black-forest-labs/FLUX.1-dev",
    subfolder="transformer",
    quantization_config=quantization_config,
)

Pular a conversão de módulos

Para alguns modelos, não é necessário quantizar todos os módulos para 8 bits, o que pode, na verdade, causar instabilidade. Por exemplo, para modelos de difusão como Stable Diffusion 3, o módulo proj_out pode ser ignorado usando o parâmetro llm_int8_skip_modules em BitsAndBytesConfig:

from diffusers import SD3Transformer2DModel, BitsAndBytesConfig

quantization_config = BitsAndBytesConfig(
    load_in_8bit=True, llm_int8_skip_modules=["proj_out"],
)

model_8bit = SD3Transformer2DModel.from_pretrained(
    "stabilityai/stable-diffusion-3-medium-diffusers",
    subfolder="transformer",
    quantization_config=quantization_config,
)

4-bit (algoritmo QLoRA)

Saiba mais sobre os detalhes em este post de blog.

Esta seção explora alguns dos recursos específicos dos modelos de 4 bits, como a alteração do tipo de dados de computação, o uso do tipo de dados Normal Float 4 (NF4) e o uso de quantização aninhada.

Tipo de dados de computação

Para acelerar o cálculo, você pode alterar o tipo de dados de float32 (o valor padrão) para bf16 usando o parâmetro bnb_4bit_compute_dtype em BitsAndBytesConfig:

import torch
from diffusers import BitsAndBytesConfig

quantization_config = BitsAndBytesConfig(load_in_4bit=True, bnb_4bit_compute_dtype=torch.bfloat16)

Float 4 Normal (NF4)

NF4 é um tipo de dados de 4 bits do paper QLoRA, adaptado para pesos inicializados a partir de uma distribuição normal. Você deve usar NF4 para treinar modelos base de 4 bits. Isso pode ser configurado com o parâmetro bnb_4bit_quant_type no BitsAndBytesConfig:

from diffusers import BitsAndBytesConfig as DiffusersBitsAndBytesConfig
from transformers import BitsAndBytesConfig as TransformersBitsAndBytesConfig

from diffusers import AutoModel
from transformers import T5EncoderModel

quant_config = TransformersBitsAndBytesConfig(
    load_in_4bit=True,
    bnb_4bit_quant_type="nf4",
)

text_encoder_2_4bit = T5EncoderModel.from_pretrained(
    "black-forest-labs/FLUX.1-dev",
    subfolder="text_encoder_2",
    quantization_config=quant_config,
    dtype=torch.float16,
)

quant_config = DiffusersBitsAndBytesConfig(
    load_in_4bit=True,
    bnb_4bit_quant_type="nf4",
)

transformer_4bit = AutoModel.from_pretrained(
    "black-forest-labs/FLUX.1-dev",
    subfolder="transformer",
    quantization_config=quant_config,
    dtype=torch.float16,
)

Para inferência, o bnb_4bit_quant_type não tem um impacto enorme no desempenho. No entanto, para manter a consistência com os pesos do modelo, você deve usar os valores de bnb_4bit_compute_dtype e dtype.

Quantização aninhada

A quantização aninhada é uma técnica que pode economizar memória adicional sem custo de desempenho extra. Este recurso executa uma segunda quantização dos pesos já quantizados para economizar mais 0,4 bits/parâmetro.

from diffusers import BitsAndBytesConfig as DiffusersBitsAndBytesConfig
from transformers import BitsAndBytesConfig as TransformersBitsAndBytesConfig

from diffusers import AutoModel
from transformers import T5EncoderModel

quant_config = TransformersBitsAndBytesConfig(
    load_in_4bit=True,
    bnb_4bit_use_double_quant=True,
)

text_encoder_2_4bit = T5EncoderModel.from_pretrained(
    "black-forest-labs/FLUX.1-dev",
    subfolder="text_encoder_2",
    quantization_config=quant_config,
    dtype=torch.float16,
)

quant_config = DiffusersBitsAndBytesConfig(
    load_in_4bit=True,
    bnb_4bit_use_double_quant=True,
)

transformer_4bit = AutoModel.from_pretrained(
    "black-forest-labs/FLUX.1-dev",
    subfolder="transformer",
    quantization_config=quant_config,
    dtype=torch.float16,
)

Desquantizando modelos bitsandbytes

Uma vez quantizado, você pode desquantizar um modelo para sua precisão original, mas isso pode resultar em uma pequena perda de qualidade. Certifique-se de ter RAM de GPU suficiente para acomodar o modelo desquantizado.

from diffusers import BitsAndBytesConfig as DiffusersBitsAndBytesConfig
from transformers import BitsAndBytesConfig as TransformersBitsAndBytesConfig

from diffusers import AutoModel
from transformers import T5EncoderModel

quant_config = TransformersBitsAndBytesConfig(
    load_in_4bit=True,
    bnb_4bit_use_double_quant=True,
)

text_encoder_2_4bit = T5EncoderModel.from_pretrained(
    "black-forest-labs/FLUX.1-dev",
    subfolder="text_encoder_2",
    quantization_config=quant_config,
    dtype=torch.float16,
)

quant_config = DiffusersBitsAndBytesConfig(
    load_in_4bit=True,
    bnb_4bit_use_double_quant=True,
)

transformer_4bit = AutoModel.from_pretrained(
    "black-forest-labs/FLUX.1-dev",
    subfolder="transformer",
    quantization_config=quant_config,
    dtype=torch.float16,
)

text_encoder_2_4bit.dequantize()
transformer_4bit.dequantize()

torch.compile

Acelere a inferência com torch.compile. Certifique-se de ter a versão mais recente do bitsandbytes instalada e também recomendamos instalar o PyTorch nightly.

torch._dynamo.config.capture_dynamic_output_shape_ops = True

quant_config = DiffusersBitsAndBytesConfig(load_in_8bit=True)
transformer_4bit = AutoModel.from_pretrained(
    "black-forest-labs/FLUX.1-dev",
    subfolder="transformer",
    quantization_config=quant_config,
    dtype=torch.float16,
)
transformer_4bit.compile(fullgraph=True)

quant_config = DiffusersBitsAndBytesConfig(load_in_4bit=True)
transformer_4bit = AutoModel.from_pretrained(
    "black-forest-labs/FLUX.1-dev",
    subfolder="transformer",
    quantization_config=quant_config,
    dtype=torch.float16,
)
transformer_4bit.compile(fullgraph=True)

Em uma RTX 4090 com compilação, a geração de Flux em 4 bits foi concluída em 25,809 segundos, em comparação com 32,570 segundos sem compilação.

Confira o script de benchmarking para mais detalhes.

Recursos

• Nota de demonstração de ponta a ponta mostrando a inferência do Flux.1 Dev em um Colab de nível gratuito
• Treinamento