bitsandbytes

bitsandbytes adalah opsi termudah untuk mengkuantisasi model ke 8 dan 4-bit. Kuantisasi 8-bit mengalikan outlier dalam fp16 dengan non-outlier dalam int8, mengkonversi nilai non-outlier kembali ke fp16, dan kemudian menjumlahkannya untuk mengembalikan bobot dalam fp16. Hal ini mengurangi efek degradatif nilai outlier terhadap kinerja model.

Kuantisasi 4-bit mengompresi model lebih jauh lagi, dan umumnya digunakan bersama QLoRA untuk melakukan finetuning pada LLM yang telah dikuantisasi.

Panduan ini menunjukkan bagaimana kuantisasi dapat memungkinkan FLUX.1-dev dijalankan dengan VRAM kurang dari 16GB, bahkan pada instance Google Colab gratis.

Untuk menggunakan bitsandbytes, pastikan Anda telah menginstal pustaka berikut:

pip install diffusers transformers accelerate bitsandbytes -U

Sekarang Anda dapat mengkuantisasi model dengan meneruskan BitsAndBytesConfig ke from_pretrained(). Ini berlaku untuk model apa pun dalam modalitas apa pun, selama model tersebut mendukung pemuatan dengan Accelerate dan berisi lapisan torch.nn.Linear.

Mengkuantisasi model dalam 8-bit mengurangi penggunaan memori menjadi setengahnya:

bitsandbytes didukung di Transformers dan Diffusers, sehingga Anda dapat mengkuantisasi FluxTransformer2DModel dan T5EncoderModel.

Untuk GPU seri Ada dan yang lebih tinggi, kami merekomendasikan untuk mengubah dtype menjadi torch.bfloat16.

CLIPTextModel dan AutoencoderKL tidak dikuantisasi karena ukurannya sudah kecil dan karena AutoencoderKL hanya memiliki beberapa lapisan torch.nn.Linear.

Disalin

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

Secara default, semua modul lain seperti torch.nn.LayerNorm dikonversi ke torch.float16. Anda dapat mengubah tipe data dari modul-modul ini dengan parameter dtype.

Disalin

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

Mari kita generate sebuah gambar menggunakan model terkuantisasi kita.

Menetapkan device_map="auto" secara otomatis mengisi semua ruang yang tersedia di GPU terlebih dahulu, kemudian CPU, dan akhirnya hard drive (opsi yang paling lambat secara absolut) jika masih tidak cukup memori.

Disalin

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]

Ketika memori cukup, Anda juga dapat langsung memindahkan pipeline ke GPU dengan .to("cuda") dan menerapkan enable_model_cpu_offload() untuk mengoptimalkan penggunaan memori GPU.

Setelah model dikuantisasi, Anda dapat mendorong model ke Hub dengan metode push_to_hub(). File konfigurasi kuantisasi config.json didorong terlebih dahulu, diikuti oleh bobot model yang telah dikuantisasi. Anda juga dapat menyimpan model 8-bit yang telah diserialisasi secara lokal dengan save_pretrained().

Pelatihan dengan bobot 8-bit dan 4-bit hanya didukung untuk melatih parameter ekstra.

Periksa jejak memori Anda dengan metode get_memory_footprint:

print(model.get_memory_footprint())

Perhatikan bahwa ini hanya menunjukkan jejak memori dari parameter model dan tidak menaksir kebutuhan memori untuk inferensi.

Model terkuantisasi dapat dimuat dari metode from_pretrained() tanpa perlu menentukan parameter 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 (algoritma LLM.int8())

Pelajari lebih lanjut tentang detail kuantisasi 8-bit dalam artikel blog ini!

Bagian ini membahas beberapa fitur spesifik dari model 8-bit, seperti ambang batas outlier dan pengabaian konversi modul.

Ambang batas outlier

“Outlier” adalah nilai hidden state yang lebih besar dari ambang batas tertentu, dan nilai-nilai ini dihitung dalam fp16. Meskipun nilai-nilai tersebut biasanya berdistribusi normal ([-3.5, 3.5]), distribusi ini bisa sangat berbeda untuk model besar ([-60, 6] atau [6, 60]). Kuantisasi 8-bit bekerja dengan baik untuk nilai ~5, tetapi melampaui itu, terdapat penurunan kinerja yang signifikan. Nilai ambang batas default yang baik adalah 6, tetapi ambang batas yang lebih rendah mungkin diperlukan untuk model yang lebih tidak stabil (model kecil atau finetuning).

Untuk menemukan ambang batas terbaik untuk model Anda, kami merekomendasikan untuk bereksperimen dengan parameter llm_int8_threshold di 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,
)

Lewati konversi modul

Untuk beberapa model, Anda tidak perlu mengkuantisasi setiap modul ke 8-bit karena hal itu justru dapat menyebabkan ketidakstabilan. Misalnya, untuk model difusi seperti Stable Diffusion 3, modul proj_out dapat dilewati menggunakan parameter llm_int8_skip_modules di 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 (algoritma QLoRA)

Pelajari lebih lanjut tentang detailnya di artikel blog ini.

Bagian ini membahas beberapa fitur spesifik dari model 4-bit, seperti mengubah tipe data komputasi, menggunakan tipe data Normal Float 4 (NF4), dan menggunakan kuantisasi bersarang.

Jenis data komputasi

Untuk mempercepat komputasi, Anda dapat mengubah tipe data dari float32 (nilai default) menjadi bf16 menggunakan parameter bnb_4bit_compute_dtype di 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 adalah tipe data 4-bit dari paper QLoRA, yang diadaptasi untuk bobot yang diinisialisasi dari distribusi normal. Anda sebaiknya menggunakan NF4 untuk melatih model dasar 4-bit. Hal ini dapat dikonfigurasi dengan parameter bnb_4bit_quant_type di 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,
)

Untuk inferensi, bnb_4bit_quant_type tidak memiliki dampak besar pada kinerja. Namun, untuk tetap konsisten dengan bobot model, Anda harus menggunakan nilai bnb_4bit_compute_dtype dan dtype.

Kuantisasi bersarang

Kuantisasi bersarang adalah teknik yang dapat menghemat memori tambahan tanpa biaya kinerja tambahan. Fitur ini melakukan kuantisasi kedua pada bobot yang sudah dikuantisasi untuk menghemat 0,4 bit/parameter tambahan.

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

Mendequantisasi model bitsandbytes

Setelah dikuantisasi, Anda dapat mendekuantisasi model ke presisi aslinya, tetapi hal ini mungkin menyebabkan penurunan kualitas yang kecil. Pastikan Anda memiliki RAM GPU yang cukup untuk menampung model yang telah didekuantisasi.

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

Percepat inferensi dengan torch.compile. Pastikan Anda telah menginstal bitsandbytes versi terbaru dan kami juga menyarankan untuk menginstal 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)

Pada RTX 4090 dengan kompilasi, generasi Flux 4-bit selesai dalam 25,809 detik dibandingkan 32,570 detik tanpa kompilasi.

Lihat skrip benchmarking untuk detail lebih lanjut.

Sumber daya

• Notebook end-to-end yang menunjukkan inferensi Flux.1 Dev di Colab tier gratis
• Pelatihan