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 과정에서 bitsandbytesLinear8bitLt, 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 등의 양자화 후 미세조정 방법도 고려 가능