본문으로 건너뛰기

[vLLM] BitsAndBytes (QLoRA): 4비트 NormalFloat 양자화

들어가며

BitsAndBytes는 QLoRA 논문에서 제안된 NormalFloat(NF4) 양자화를 구현한 라이브러리다. 4비트 양자화를 통해 대형 모델을 단일 GPU에 올릴 수 있게 해주며, vLLM에서도 추론 시 이를 지원한다. 이번 글에서는 vLLM이 BitsAndBytes를 어떻게 통합하고 있는지를 코드 수준에서 분석한다.

관련 논문: QLoRA: Efficient Finetuning of Quantized LLMs (arxiv 2305.14314)

공식 문서

vLLM 공식 문서: BitsAndBytes Quantization

핵심 구조/코드 분석

Config 클래스

BitsAndBytesConfig는 8비트와 4비트 모드를 모두 지원한다:

class BitsAndBytesConfig(QuantizationConfig):
    def __init__(
        self,
        load_in_8bit: bool = False,
        load_in_4bit: bool = True,
        bnb_4bit_compute_dtype: str = "float32",
        bnb_4bit_quant_storage: str = "uint8",
        bnb_4bit_quant_type: str = "fp4",
        bnb_4bit_use_double_quant: bool = False,
        llm_int8_threshold: float = 6.0,
    ) -> None:

주요 설정 중 bnb_4bit_quant_type은 "fp4" 또는 "nf4"를 지원하며, llm_int8_threshold는 8비트 혼합 정밀도에서 outlier를 FP16으로 처리하는 임계값이다.

4비트 가중치 생성

4비트 모드에서 가중치는 uint8로 패킹된다. quant_ratio로 패킹 비율을 계산한다:

def create_qweight_for_4bit():
    quant_ratio = calculate_quant_ratio(params_dtype)
    total_size = input_size_per_partition * sum(output_partition_sizes)
    qweight = torch.nn.Parameter(
        torch.empty(total_size // quant_ratio, 1, dtype=torch.uint8),
        requires_grad=False,
    )

bf16(16비트)을 uint8(8비트)로 저장할 때 quant_ratio = 2가 되므로, 4비트 값 2개가 하나의 uint8에 들어가는 구조다.

추론 경로: Custom Op 등록

vLLM은 torch.ops.vllm.apply_bnb_4bit라는 커스텀 연산자를 등록하여 torch.compile과의 호환성을 확보한다:

direct_register_custom_op(
    op_name="apply_bnb_4bit",
    op_func=_apply_bnb_4bit,
    mutates_args=["out"],
    fake_impl=_apply_bnb_4bit_fake,
    dispatch_key=current_platform.dispatch_key,
)

실제 연산은 bitsandbytes.matmul_4bit를 호출하되, 각 shard별로 분리 실행하여 텐서 병렬 처리를 지원한다.

MoE 지원: Dequant 후 fused_experts

MoE에서는 4비트 가중치를 먼저 역양자화한 뒤 fused_experts 커널에 전달한다:

def _apply_4bit_dequnt(self, layer):
    from bitsandbytes.functional import dequantize_4bit
    w13 = dequantize_4bit(
        layer.w13_weight.reshape(-1, 1),
        layer.w13_weight.bnb_quant_state,
    )
    w2 = dequantize_4bit(
        layer.w2_weight.reshape(-1, 1),
        layer.w2_weight.bnb_quant_state,
    )
    w13 = w13.reshape(layer.w13_weight.experts_shape)
    w2 = w2.reshape(layer.w2_weight.experts_shape)
    return w13, w2

이 방식은 매 forward마다 역양자화를 수행하므로 성능 오버헤드가 있지만, 메모리 절감 효과가 크다.

왜 이 설계인가

  1. 호환성 우선: BitsAndBytes 라이브러리의 양자화/역양자화 함수를 그대로 활용하되, vLLM의 custom op 시스템으로 감싸서 torch.compile과 CUDA Graph를 지원한다.

  2. Skip 모듈 지원: llm_int8_skip_modules 설정으로 특정 레이어(예: lm_head)를 양자화에서 제외할 수 있다. 출력 레이어의 정밀도를 유지하는 것이 모델 품질에 중요하기 때문이다.

  3. MoE dequant 전략: 전용 양자화 커널이 없으므로 역양자화 후 기존 fused_experts를 재활용한다. 성능은 다소 손해보지만, 구현 복잡도를 크게 줄이는 실용적인 선택이다.

논문 핵심 내용

QLoRA: Efficient Finetuning of Quantized LLMs (2305.14314) 논문은 4비트 양자화된 모델 위에 LoRA를 적용하는 기법을 제안했다.

핵심 아이디어: 세 가지 혁신을 통해 메모리를 극적으로 줄였다. (1) NormalFloat4(NF4): 정규분포를 따르는 가중치에 대해 정보 이론적으로 최적인 4비트 데이터 타입. (2) Double Quantization: 양자화 상수 자체를 다시 양자화하여 파라미터당 약 0.37비트를 추가 절약. (3) Paged Optimizers: 메모리 스파이크를 CPU 메모리로 오프로드.

이 조합으로 65B 파라미터 모델을 단일 48GB GPU에서 파인튜닝할 수 있게 됐고, 16비트 풀 파인튜닝과 동일한 태스크 성능을 유지했다. 논문에서 1,000개 이상의 모델을 8개 데이터셋, 다양한 아키텍처(LLaMA, T5), 다양한 규모(33B, 65B)에서 학습시켰다.

메트릭 수치
65B 모델 파인튜닝 GPU 메모리 48GB (단일 GPU)
Guanaco-65B ChatGPT 대비 성능 99.3% (Vicuna 벤치마크)
학습 시간 24시간 GPU 학습
NF4 vs FP4 NF4가 정규분포 가중치에서 정보 이론적 최적
Double Quantization 절약량 파라미터당 ~0.37비트

Guanaco 모델 패밀리는 Vicuna 벤치마크에서 **ChatGPT 성능의 99.3%**를 달성했고, 이는 단 24시간의 GPU 학습만으로 가능했다. 기존에는 33B, 65B 규모의 모델 파인튜닝이 메모리 제약으로 불가능했는데, QLoRA가 이 장벽을 허물어버린 거다.

참고 자료

댓글

관련 포스트

vLLM 의 다른글