Diffusers documentation
bitsandbytes
bitsandbytes
bitsandbytes는 모델을 8비트 및 4비트로 양자화하는 가장 쉬운 옵션입니다. 8비트 양자화는 fp16의 아웃라이어를 int8의 비아웃라이어와 곱한 후, 비아웃라이어 값을 다시 fp16으로 변환하고, 이를 더하여 fp16의 가중치를 반환합니다. 이를 통해 아웃라이어가 모델 성능에 미치는 부정적인 영향을 줄일 수 있습니다.
4-bit 양자화는 모델을 더욱 압축하며, 양자화된 LLM을 파인튜닝하기 위해 QLoRA와 함께 널리 사용됩니다.
이 가이드는 양자화(quantization)를 통해 FLUX.1-dev를 16GB 미만의 VRAM에서, 심지어 무료 Google Colab 인스턴스에서도 실행하는 방법을 보여줍니다.

bitsandbytes를 사용하려면 다음 라이브러리가 설치되어 있는지 확인하세요:
pip install diffusers transformers accelerate bitsandbytes -U
이제 from_pretrained()에 BitsAndBytesConfig를 전달하여 모델을 양자화할 수 있습니다. Accelerate를 통한 로딩을 지원하고 torch.nn.Linear 레이어를 포함하는 한, 모든 모달리티의 모든 모델에 적용됩니다.
모델을 8-bit로 양자화하면 메모리 사용량이 절반으로 줄어듭니다:
bitsandbytes는 Transformers와 Diffusers 모두에서 지원되므로 FluxTransformer2DModel과 T5EncoderModel을 모두 양자화할 수 있습니다.
Ada 및 상위 시리즈 GPU의 경우 dtype을 torch.bfloat16(으)로 변경하는 것을 권장합니다.
CLIPTextModel과 AutoencoderKL은 이미 크기가 작고, AutoencoderKL에는torch.nn.Linear레이어가 몇 개밖에 없기 때문에 양자화되지 않습니다.
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,
)기본적으로 torch.nn.LayerNorm과 같은 다른 모든 모듈은 torch.float16으로 변환됩니다. dtype 매개변수를 사용하여 이러한 모듈의 데이터 유형을 변경할 수 있습니다.
transformer_8bit = AutoModel.from_pretrained(
"black-forest-labs/FLUX.1-dev",
subfolder="transformer",
quantization_config=quant_config,
+ dtype=torch.float32,
)양자화된 모델을 사용하여 이미지를 생성해 보겠습니다.
device_map="auto"를 설정하면 먼저 GPU의 모든 사용 가능한 공간을 채운 다음 CPU, 그리고 여전히 메모리가 부족하면 하드 드라이브(가장 느린 옵션)에 배치합니다.
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]
메모리가 충분할 경우, .to("cuda")를 사용하여 파이프라인을 GPU로 직접 이동하고 enable_model_cpu_offload()를 적용하여 GPU 메모리 사용량을 최적화할 수도 있습니다.
Once a model is quantized, you can push the model to the Hub with the push_to_hub() method. The quantization config.json file is pushed first, followed by the quantized model weights. You can also save the serialized 8-bit models locally with save_pretrained().
모델을 4-bit로 양자화하면 메모리 사용량이 4분의 1로 줄어듭니다:
bitsandbytes는 Transformers와 Diffusers 모두에서 지원되므로 FluxTransformer2DModel과 T5EncoderModel을 모두 양자화할 수 있습니다.
Ada 및 상위 시리즈 GPU의 경우 dtype을 torch.bfloat16(으)로 변경하는 것을 권장합니다.
CLIPTextModel과 AutoencoderKL은 이미 크기가 작고, AutoencoderKL에는torch.nn.Linear레이어가 몇 개밖에 없기 때문에 양자화되지 않습니다.
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_4bit=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,)
transformer_4bit = AutoModel.from_pretrained(
"black-forest-labs/FLUX.1-dev",
subfolder="transformer",
quantization_config=quant_config,
dtype=torch.float16,
)기본적으로 torch.nn.LayerNorm과 같은 다른 모든 모듈은 torch.float16으로 변환됩니다. dtype 매개변수를 사용하여 이러한 모듈의 데이터 유형을 변경할 수 있습니다.
transformer_4bit = AutoModel.from_pretrained(
"black-forest-labs/FLUX.1-dev",
subfolder="transformer",
quantization_config=quant_config,
+ dtype=torch.float32,
)양자화된 모델을 사용하여 이미지를 생성해 보겠습니다.
device_map="auto"를 설정하면 먼저 GPU의 모든 사용 가능한 공간을 채운 다음 CPU, 그리고 여전히 메모리가 부족하면 하드 드라이브(가장 느린 옵션)에 배치합니다.
from diffusers import FluxPipeline
pipe = FluxPipeline.from_pretrained(
"black-forest-labs/FLUX.1-dev",
transformer=transformer_4bit,
text_encoder_2=text_encoder_2_4bit,
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]
메모리가 충분할 경우, .to("cuda")를 사용하여 파이프라인을 GPU로 직접 이동하고 enable_model_cpu_offload()를 적용하여 GPU 메모리 사용량을 최적화할 수도 있습니다.
Once a model is quantized, you can push the model to the Hub with the push_to_hub() method. The quantization config.json file is pushed first, followed by the quantized model weights. You can also save the serialized 4-bit models locally with save_pretrained().
8-bit 및 4-bit 가중치를 사용한 학습은 추가 파라미터 학습만 지원합니다.
get_memory_footprint 메서드로 메모리 사용량을 확인하세요:
print(model.get_memory_footprint())이것은 모델 파라미터의 메모리 사용량만 알려줄 뿐, 추론에 필요한 메모리 요구사항을 추정하지는 않습니다.
양자화된 모델은 quantization_config 매개변수를 지정할 필요 없이 from_pretrained() 메서드로 로드할 수 있습니다:
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 (LLM.int8() 알고리즘)
8-bit 양자화의 세부 사항은 이 블로그 게시물에서 자세히 알아보세요!
이 섹션에서는 8-bit 모델의 특정 기능, 예를 들어 아웃라이어 임계값 및 모듈 변환 건너뛰기 등을 살펴봅니다.
이상치 임계값
“outlier”는 특정 임계값을 초과하는 hidden state 값이며, 이러한 값은 fp16로 계산됩니다. 값은 일반적으로 정규 분포([-3.5, 3.5])를 따르지만, 대규모 모델의 경우 이 분포가 매우 달라질 수 있습니다([-60, 6] 또는 [6, 60]). 8-bit 양자화는 ~5인 값에 잘 작동하지만, 그 이상에서는 상당한 성능 저하가 발생합니다. 좋은 기본 임계값은 6이지만, 더 불안정한 모델(소형 모델 또는 finetuning)의 경우 더 낮은 임계값이 필요할 수 있습니다.
모델에 가장 적합한 임계값을 찾기 위해 BitsAndBytesConfig에서 llm_int8_threshold 매개변수를 실험해 보시는 것을 권장합니다:
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,
)모듈 변환 건너뛰기
일부 모델의 경우 모든 모듈을 8-bit로 양자화할 필요가 없으며, 오히려 불안정성을 유발할 수 있습니다. 예를 들어 Stable Diffusion 3와 같은 diffusion 모델에서는 BitsAndBytesConfig의 llm_int8_skip_modules 매개변수를 사용하여 proj_out 모듈을 건너뛸 수 있습니다:
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(QLoRA 알고리즘)
이 블로그 게시물에서 세부 사항을 자세히 알아보세요.
이 섹션에서는 4-bit 모델의 특정 기능, 예를 들어 계산 데이터 타입 변경, Normal Float 4 (NF4) 데이터 타입 사용, 중첩 양자화 사용 등을 살펴봅니다.
Compute data type
계산 속도를 높이기 위해 BitsAndBytesConfig의 bnb_4bit_compute_dtype 매개변수를 사용하여 데이터 유형을 float32(기본값)에서 bf16으로 변경할 수 있습니다:
import torch
from diffusers import BitsAndBytesConfig
quantization_config = BitsAndBytesConfig(load_in_4bit=True, bnb_4bit_compute_dtype=torch.bfloat16)일반 Float 4 (NF4)
NF4는 QLoRA 논문에서 제시된 4비트 데이터 타입으로, 정규 분포에서 초기화된 가중치에 맞게 조정되었습니다. 4비트 베이스 모델을 학습할 때는 NF4를 사용해야 합니다. 이는 BitsAndBytesConfig의 bnb_4bit_quant_type 파라미터로 설정할 수 있습니다:
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,
)추론 시 bnb_4bit_quant_type은 성능에 큰 영향을 미치지 않습니다. 그러나 모델 가중치와 일관성을 유지하려면 bnb_4bit_compute_dtype과 dtype 값을 사용해야 합니다.
중첩 양자화
Nested quantization은 추가적인 성능 저하 없이 추가 메모리를 절약할 수 있는 기법입니다. 이 기능은 이미 양자화된 가중치에 대해 두 번째 양자화를 수행하여 파라미터당 0.4비트를 추가로 절약합니다.
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,
)bitsandbytes 모델의 역양자화
양자화 후 원래 정밀도로 역양자화할 수 있지만, 이로 인해 품질이 약간 저하될 수 있습니다. 역양자화된 모델을 수용할 수 있을 만큼 충분한 GPU RAM이 있는지 확인하세요.
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
torch.compile을 사용하여 추론을 가속화하세요. 최신 bitsandbytes가 설치되어 있는지 확인하고, 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)컴파일링을 적용한 RTX 4090에서 4-bit Flux 생성은 25.809초에 완료되었으며, 적용하지 않은 경우 32.570초보다 빨랐습니다.
자세한 내용은 벤치마킹 스크립트를 확인하세요.