bitsandbytes 작동 원리 - 100-hours-a-week/5-yeosa-wiki GitHub Wiki
1. bitsandbytes
양자화 동작 원리 및 단계별 설명
a. 개요
bitsandbytes
는 NVIDIA GPU에서 8bit / 4bit 양자화 연산을 효율적으로 처리할 수 있도록 지원하는 CUDA 기반 라이브러리
- 이를 통해 PyTorch 모델을 매우 작은 메모리 footprint로 추론에 사용할 수 있음.
b. 단계별 흐름
가. 모델 로딩 요청
from_pretrained()
메서드를 통해 모델을 로드할 때, quantization_config=BitsAndBytesConfig(...)
을 전달
- 이 설정에 따라
transformers
또는 diffusers
는 로딩 시 자동으로 bitsandbytes.nn.Linear8bitLt
등으로 레이어를 바꿔치기
나. 가중치 로딩 및 변환
- 원래
float32
또는 float16
이던 가중치를 int8
또는 int4
포맷으로 변환
scaling factor
를 함께 저장하여 추론 시 approximate float
로 복원 가능
다. 모델을 VRAM에 로드
load_in_8bit=True
옵션에 따라 양자화된 가중치만 VRAM에 올라감
- 일반적으로 가중치만 8bit/4bit로 줄어들고, 활성값 및 일부 연산은 여전히
float16
또는 float32
- 참고 :
torch_dtype=torch.float16
이 함께 사용되면
- 레이어 입력/출력은 float16
- 내부 곱셈(매트릭스 연산)은 8bit
- 결과는 float16으로 다시 전환됨
라. 추론 연산
- forward 과정에서
bitsandbytes
의 Linear8bitLt
, Linear4bit
, 혹은 QLoRA 스타일의 adapter를 사용해 연산
- CUDA 커널이 내부적으로 효율적으로 작동하여 VRAM 소모 적고 속도 빠름
2. 예시 코드의 양자화 전략 분석
# text encoder 양자화
quant_config = BitsAndBytesConfig(load_in_8bit=True)
text_encoder_8bit = T5EncoderModel.from_pretrained(
"black-forest-labs/FLUX.1-dev",
subfolder="text_encoder_2",
quantization_config=quant_config,
torch_dtype=torch.float16,
)
# transformer 양자화
quant_config = DiffusersBitsAndBytesConfig(load_in_8bit=True)
transformer_8bit = FluxTransformer2DModel.from_pretrained(
"black-forest-labs/FLUX.1-dev",
subfolder="transformer",
quantization_config=quant_config,
torch_dtype=torch.float16,
)
# 파이프라인 구성
pipeline = FluxPipeline.from_pretrained(
"black-forest-labs/FLUX.1-dev",
text_encoder_2=text_encoder_8bit,
transformer=transformer_8bit,
torch_dtype=torch.float16,
device_map="balanced",
)
a. 분석 요약
구성 요소 |
전략 |
설명 |
text_encoder_2 |
BitsAndBytesConfig(load_in_8bit=True) + torch_dtype=torch.float16 |
transformers 용 T5 모델을 8bit로 로드하되, 입력/출력은 float16 |
transformer |
DiffusersBitsAndBytesConfig(load_in_8bit=True) + torch_dtype=torch.float16 |
diffusers 용 Transformer 구조를 동일하게 8bit로 로드 |
FluxPipeline |
위에서 로드한 두 모델을 구성요소로 전달 |
메모리 최적화를 위한 맞춤형 파이프라인 구성 |
device_map="balanced" |
자동으로 GPU 사용량을 분산 배치 |
큰 모델을 여러 GPU 또는 GPU+CPU에 분산 가능 |
b. 이점
- VRAM 절약: 일반적으로
float16
대비 50% 이상 절약
- 속도 개선: 일부 연산에서 CUDA kernel 최적화로 인해 빠름
- 활용성 확대: T4 또는 A10G처럼 VRAM 16GB 수준의 GPU에서도 고성능 모델 실행 가능
c. 주의사항
bitsandbytes
는 CUDA 환경 필수
- 일부 레이어(예: normalization, cross-attention 등)는 양자화가 어려워 float16으로 남아야 함
- 최적 정확도를 원할 경우, QLoRA 등의 양자화 후 미세조정 방법도 고려 가능