QLoRA-Efficient Finetuning of Quantized LLMs

지난번 포스팅에서는 LoRA에 대한 포스팅을 진행했습니다. LoRA가 등장하게 된 배경은 LLM모델의 크기가 시간이 지날수록 커지고 있어 finetuning과정에서 모든 파라미터들을 학습시키기에는 컴퓨팅 자원에 대한 부담이 커진다라는 것이였습니다. LoRA를 통해 속도와 메모리에 대한 부담을 줄일 수 있는것은 사실이였지만 중요한건 LoRA에서도 원래의 모델에 대한 가중치를 사용한다는 것입니다.(자세한 내용은 지난 포스팅을 참고해주세요)
 
예를들어서  GPT-3를 finetuning하는 과정을 생각해보겠습니다. 150B의 파라미터를 갖고 있습니다. 이 모델 자체를 일반적인 GPU에 올리기란 쉽지 않습니다. 순수하게 fully finetuning을 진행한다고 하면 A100 80G GPU가 약 70개가 되는 자원이 필요하고, Inference는 7개의 A100 GPU가 필요할 것입니다. 물론 LoRA를 통한다면 finetuning과정은 크게 축소될 수 있습니다. 하지만 Inference과정에서는 그 영향을 받지 못합니다. 이러한 어려움을 해결하고자 나온 방법이 QLoRA입니다. 대표적으로 QLoRA는 65B의 파라미터를 갖고 있는 모델을 finetuning하는데 고작 싱글 GPU(48G)만으로 훈련을 성공했다고 합니다. 단순히 하나의 모델이 아니라 1000개가 되는 모델들을 통해 실험을 진행하였다고합니다. Quantization(양자화) + LoRA로 그 방법에 대해서 알아보겠습니다.
 
제가 이해한 내용을 바탕으로 해당 글은 작성되었으며 좀 더 자세한 내용은 아래 링크를 참고해주세요.
https://github.com/artidoro/qlora?tab=readme-ov-file

GitHub - artidoro/qlora: QLoRA: Efficient Finetuning of Quantized LLMs

 

Model Quantization?

QLoRA에 대한 설명에 앞서 모델 양자화(Model Quantization)에 대해서 간단하게 짚고 넘어가보려고 합니다. 사실 이 부분만 알면 이를 응용하고 LoRA와 결합했다는게  QLoRA의 전부이기는합니다...
 
양자화(Quantization)는 딥러닝 모델의 파라미터를 더 작은 비트 수로 표현하는 기술을 말합니다. 부동 소수점 수를 표현하는데 사용되는 비트 수를 감소시킴으로써 모델을 효과적으로 압축하고 메모리의 우위를 확보하여 추론 속도를 향상시킵니다. 이런 양자화는 본질적으로 한 data type을 다른 data type으로 반올림하며 수행됩니다. 모델의 양자화에 대해서 이야기 해보았는데 좀 추상적이니 예를들어서 그 과정을 이해해보겠습니다. 양자화가 진행되는 과정은 크게 정규화, 양자화, 역양자화 3가지 단계로 구분됩니다.

1. 정규화
   -> 모델의 파라미터 확은 가중치 혹은 활성화 값을 정규화합니다. 예를들어, 0과 1사이의 값으로 스케일링을 진행한다고 해보겠습니다.
   
   
2. 양자화
   -> 정규화된 값을 8비트 양자화로 표현한다고 가정해보겠습니다. 8비트 양자화는 2의 8승으로 256개의 값으로 표현하게 됩니다. 양자화 하려는 값이 정규화를 통해 0.2라는 값을 갖고 있었다면 8비트 양자화는 이와 가장 가까운 값인 0.19921875로 양자화가 진행됩니다. 이를 8비트로 표현하여 값을 갖고 있게 됩니다.
   
   
3. 역양자화
   8비트화해서 갖고 있던 값은 역양자화를 통해 다시 부동 소수점으로 변환하게 되는데 당연히 가장 가까운 값이었던 0.19921875로 값이 변환되게 됩니다. 원래 정규화를 통해 얻은 값은 0.2였지만 양자화를 통해 변화된 값은 작지만 오차를 발생 시키는 모습을 볼 수 있습니다. 

즉, 양자화는 역전파 과정이 많이 일어날수록 이러한 오차가 누적되어 그 의미가 퇴색된다고 볼 수 있습니다. 이러한 문제점을 극복하여 제안한 방법이 QLoRA입니다. 그 방법을 지금부터 살펴보겠습니다.

QLoRA

QLoRA 논문의 내용을 바탕으로 이야기 해보겠습니다. 앞서 이야기했지만 QLoRA는 4비트 양자화 + LoRA를 활용하여 65B 모델을 48GB 싱글 GPU로 finetuning작업에 성공하였습니다. 일반적으로는 65B모델을 사용하려면 VRAM 260GB+α가 필요합니다. 그런데 QLoRA는 자원을 엄청난 이점을 취했음에도 성능면에서 크게 뒤쳐지는 모습이 없었습니다. 과연 어떻게 이렇게 모델 사이즈를 줄일 수 있었을까요? 이에 대한 테크닉은 3가지를 통해 이뤄질 수 있었습니다. 이제 그 세 가지를 알아보겠습니다.

1) 4-bit NormalFloat Quantization(NF4)

NF4의 이 데이터 형식은 Quantile Quantization(분위수 양자화)을 기반으로 하며, 각 양자화 구간이 입력 텐서에서 동일한 수의 값을 갖도록 하는 정보 이론적으로 최적인 데이터 형식입니다. 그런데 이 분위수 양자화의 단저믄 분위수 추정 과정에서의 비용이 크게 들어간다는 점입니다. 이를 해결하기 위해 SRAM quantiles와 같은 빠른 분위수 근사 알고리즘이 사용됩니다. 그런데 이렇게 하면 또 문제가 발생할 수 있는게 이상치에 대해 큰 양자화 오차를 범할 수 있다는 것입니다. 이 오차를 피하기 위한 방법중 하나는 양자화를 위한 상수가 고정된 분포에서 나오는 경우입니다. 그런데 훈련된 신경망 모델의 가중치는 표준 정규 분포의 형태를 갖고 있습니다. 따라서 이를 통해 하나의 고정된 분포로 변화할 수 있고 이를 표준편차를 이용해 [-1, 1] 범위로 설정하여 정규화를 진행합니다. 이를 통해 얻게 되는 것은 양자화 속도와 오차의 감소 두 효과 모두 얻을 수 있는 방법이 되었습니다.

2) Double Quantization(DQ)

양자화를 진행하게 되면 앞서 이야기한 양자화 상수라는게 존재합니다. DQ는 1차로 진행한 양자화에서의 상수를 다시 양자화를 진행합니다. 이렇게 함으로서 메모리에 대한 이점을 얻으려고 했습니다. 실제로 이 방법을 통해 파라미터당 0.373 bit의 이득을 취할 수 있다고 합니다. 이러한 DQ가 가능한 이유는 다음 논문에서 증명한 내용을 바탕으로 8비트 양자화가 결과에 큰 영향을 주지 못한다는 가정을 바탕으로 진행하였습니다.

DQ예시 논문 중 발췌

3) Paged Optimizers

Nvidia의 통합 메모리 기능을 사용하여 GPU가 메모리가 부족한 상황이 생긴다면 CPU RAM 혹은 Disk로 작업을 전송하여 작업을 처리할 수 있도록 한다고 합니다. 이는 RAM도 작업량이 메모리보다 클 경우 가상메모리로 HDD,SSD를 활용하는 방식으로 똑같이 활용되기도 합니다.

 

 

IN CODE

QLoRA는 huggingface에서 쉽게 호출이 가능합니다. 아래는 직접 구현보다는 QLoRA에서 설명한 요소들을 설정하여 모델을 불러오는 간단한 예시로 글을 마무리하겠습니다.

 

아래는 huggingface에서 모델을 불러올때 설정할 수 있는 Config 클래스입니다. LLM.int8, FP4, NF4를 선택하여 모델을 불러올 수 있도록 현재는 지원하고 있습니다. 이후에 추가적으로 불러올 수 있는 arguments가 추가된다면 이 클래스에 추가 될 것이라고 합니다. 

@dataclass
class BitsAndBytesConfig(QuantizationConfigMixin):
    """
    This is a wrapper class about all possible attributes and features that you can play with a model that has been
    loaded using `bitsandbytes`.

    This replaces `load_in_8bit` or `load_in_4bit`therefore both options are mutually exclusive.

    Currently only supports `LLM.int8()`, `FP4`, and `NF4` quantization. If more methods are added to `bitsandbytes`,
    then more arguments will be added to this class.

    Args:
        load_in_8bit (`bool`, *optional*, defaults to `False`):
            This flag is used to enable 8-bit quantization with LLM.int8().
        load_in_4bit (`bool`, *optional*, defaults to `False`):
            This flag is used to enable 4-bit quantization by replacing the Linear layers with FP4/NF4 layers from
            `bitsandbytes`.
        llm_int8_threshold (`float`, *optional*, defaults to 6.0):
            This corresponds to the outlier threshold for outlier detection as described in `LLM.int8() : 8-bit Matrix
            Multiplication for Transformers at Scale` paper: https://arxiv.org/abs/2208.07339 Any hidden states value
            that is above this threshold will be considered an outlier and the operation on those values will be done
            in fp16. Values are usually normally distributed, that is, most values are in the range [-3.5, 3.5], but
            there are some exceptional systematic outliers that are very differently distributed for large models.
            These outliers are often in the interval [-60, -6] or [6, 60]. Int8 quantization works well for values of
            magnitude ~5, but beyond that, there is a significant performance penalty. A good default threshold is 6,
            but a lower threshold might be needed for more unstable models (small models, fine-tuning).
        llm_int8_skip_modules (`List[str]`, *optional*):
            An explicit list of the modules that we do not want to convert in 8-bit. This is useful for models such as
            Jukebox that has several heads in different places and not necessarily at the last position. For example
            for `CausalLM` models, the last `lm_head` is kept in its original `dtype`.
        llm_int8_enable_fp32_cpu_offload (`bool`, *optional*, defaults to `False`):
            This flag is used for advanced use cases and users that are aware of this feature. If you want to split
            your model in different parts and run some parts in int8 on GPU and some parts in fp32 on CPU, you can use
            this flag. This is useful for offloading large models such as `google/flan-t5-xxl`. Note that the int8
            operations will not be run on CPU.
        llm_int8_has_fp16_weight (`bool`, *optional*, defaults to `False`):
            This flag runs LLM.int8() with 16-bit main weights. This is useful for fine-tuning as the weights do not
            have to be converted back and forth for the backward pass.
        bnb_4bit_compute_dtype (`torch.dtype` or str, *optional*, defaults to `torch.float32`):
            This sets the computational type which might be different than the input time. For example, inputs might be
            fp32, but computation can be set to bf16 for speedups.
        bnb_4bit_quant_type (`str`,  *optional*, defaults to `"fp4"`):
            This sets the quantization data type in the bnb.nn.Linear4Bit layers. Options are FP4 and NF4 data types
            which are specified by `fp4` or `nf4`.
        bnb_4bit_use_double_quant (`bool`, *optional*, defaults to `False`):
            This flag is used for nested quantization where the quantization constants from the first quantization are
            quantized again.
        kwargs (`Dict[str, Any]`, *optional*):
            Additional parameters from which to initialize the configuration object.
    """

 

그러면 이제 위에서 논문에서 설정한 내용을 바탕으로 불러오는 코드를 작성해보겠습니다. 양자화 타입은 NF4로 설정하고 Double Quantization을 진행하도록 설정하겠습니다. 또한 추가적으로 compute type은 실질적인 속도에 영향을 줄 수 있는 요소로 타입을 fp32가 아닌 bf16으로 설정하여 모델을 불러와 보겠습니다. 아래 코드를 실행하면 큰 모델을 colab에서도 불러올 수 있게 됩니다.

import torch
from transformers import AutoTokenizer, AutoModelForCausalLM, BitsAndBytesConfig

# # 모델, 토크나이저

model_id = "EleutherAI/polyglot-ko-12.8b" #huggingFace

bnb_config = BitsAndBytesConfig(
    load_in_4bit=True,
    bnb_4bit_use_double_quant=True,
    bnb_4bit_quant_type="nf4",
    bnb_4bit_compute_dtype=torch.bfloat16
)

model = AutoModelForCausalLM.from_pretrained(model_id, quantization_config=bnb_config, device_map={"":0})

---

QLoRA를 직접적으로 바로 체감할 수 있는 부분은 Colab에서 감당할 수 없는 모델을 불러와서 finetuning하는 과정입니다. 논문의 예시처럼 너무 큰 모델을 불러올 수는 없습니다. 기본 GPU가 그렇게 크지 않기 떄문이기도 합니다. 그런데 최근의 LLM들은 small사이즈 모델로도 나오기 때문에 이러한 모델을 불러와 직접 QLoRA를 통해 LLM모델에 대한 finetuning을 진행해볼 수 있습니다. 저도 이 방법을 통해서 기존에 던전앤파이터 챗봇 만들기 프로젝트를 KoGPT가 아닌 polyglot-ko-12.8b모델을 활용해 진행해보려 합니다. 이에 대한 결과는 훈련이 성공하면 다음 포스팅에 올려보겠습니다!