효율적인 LLM 학습 전략

introduction

최근에는 LLM을 적은 GPU 내에서 Fine-tuning을 진행하려는 연구 분야가 있다.

대기업이나, 대규모 연구실 등, LLM을 연구/개발하는 사람들은

고객들에게 서비스를 하기 위해 모델을 효율적으로 압축(양자화)하고, 효율적으로 Fine-tuning 하는 방법을 사용한다.

이런 연구분야는 어떻게 발전되어 왔는지 살펴보자.

Method

이 연구분야는 크게 두가지 갈래로 발전이 되어왔다.

첫번째는, 모델을 압축(양자화) 하는 방법이다.

모델을 압축(양자화) 함으로써, 모델의 성능은 유지하고 GPU에 올라가는 모델의 메모리를 줄일 수 있다.

두번째는, Fine-tuning을 진행할 때, 모델의 전체 파라미터를 학습시키는 것이 아니라,

학습 가능한 적은양의 파라미터를 모델 내부에 삽입해 이 부분만 학습을 시킨다.

이렇게 하면, 적은양의 파라미터로 모델을 Fine-tuning 할 수 있고,

모델의 성능은 유지하고 학습 시 소모되는 메모리 소모량을 감소 시킬 수 있다.

Large Lanuage Model Quantization

Quantization이란?

모델의 파라미터를 Low bit로 표현하고 저장함으로써, 모델을 압축하는 방법.

기존의 언어모델의 파라미터는 Float32로 표현되었는데, 이를 FP16, BF16, FP8, NF4 등 같은 Low bit로 압축을 한다.

다음은 기본적으로 fp32로 모델 업로드하는 방법이다

from transformers import AutoModelForCausalLM, GPTNeoXConfig

model_config = GPTNeoXConfig.from_pretrained("EleutherAI/polyglot-ko-12.8b")

model = AutoModelForCausalLM.from_pretrained(
  "EleutherAI/polyglot-ko-12.8b",
  config = model_config)

다음은 간단하게 pytorch를 사용해 fp16으로 양자화 하는 방법이다.

from transformers import AutoModelForCausalLM, GPTNeoXConfig

model_config = GPTNeoXConfig.from_pretrained("EleutherAI/polyglot-ko-12.8b")
model_config.torch_dtype = torch.float16

model = AutoModelForCausalLM.from_pretrained(
  "EleutherAI/polyglot-ko-12.8b",
  config = model_config)

다음은 간단하게 Bitsandbytes 사용해 nf4으로 양자화 하는 방법이다.

from transformers import GPTNeoXConfig, BitsAndBytesConfig, AutoModelForCausalLM, 

model_config = GPTNeoXConfig.from_pretrained("EleutherAI/polyglot-ko-12.8b")
model_config.torch_dtype = torch.float16

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(
  "EleutherAI/polyglot-ko-12.8b",
  config = model_config
  quantization_config=bnb_config,
  device_map = "auto"
)

Parameter-Efficient Fine-Tuning (PEFT)

PEFT란?

Pre-trained Language model(PLM)을 Fine-tuning 하는 단계에서

기존 언어모델이 가지고 있는 파라미터는 freezing하고

학습 가능한 적은양의 파라미터를 모델 내부에 삽입해 이 부분만 학습을 시킨다.

다음은 간단하게 LoRA를 모델에 삽입하는 방법이다.

from transformers import AutoModelForCausalLM, GPTNeoXConfig
from peft import LoraConfig, get_peft_model

model_config = GPTNeoXConfig.from_pretrained("EleutherAI/polyglot-ko-12.8b")
model_config.torch_dtype = torch.float16

model = AutoModelForCausalLM.from_pretrained(
  "EleutherAI/polyglot-ko-12.8b",
  config = model_config
)

lora_config = LoraConfig(
        r=8,
        lora_alpha=32,
        target_modules=["query_key_value"],
        lora_dropout= 0.05,
        task_type="CAUSAL_LM",
)

model = get_peft_model(model, lora_config)

다음은 nf4로 양자화한 모델에 LoRA를 모델에 삽입하는 방법이다. -> QLoRA 구현

from transformers import GPTNeoXConfig, BitsAndBytesConfig, AutoModelForCausalLM, 

import bitsandbytes as bnb

from peft import LoraConfig, get_peft_model, prepare_model_for_kbit_training


model_config = GPTNeoXConfig.from_pretrained("EleutherAI/polyglot-ko-12.8b")
model_config.torch_dtype = torch.float16

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(
  "EleutherAI/polyglot-ko-12.8b",
  config = model_config
  quantization_config=bnb_config,
  device_map = "auto"
)

lora_config = LoraConfig(
        r=8,
        lora_alpha=32,
        target_modules=["query_key_value"],
        lora_dropout= 0.05,
        task_type="CAUSAL_LM",
)

model = get_peft_model(model, lora_config)
model = prepare_model_for_kbit_training(model)

## QLoRA는 Paged Optimizer도 제안했음. 이를 사용

optimizer = bnb.optim.PagedAdam32bit(model.parameters(), lr=2e-4, betas=(0.9, 0.999)) # equivalent

Conclusion

LLM을 적은 GPU 내에서 Fine-tuning을 진행하기 위해서는

PEFT, Quantization을 진행 해야한다.

이전 연구들은 PEFT, Quantization 따로 진행하며 연구 되었지만,

최근에는 PEFT와 Quantization을 합쳐 좋은 성능을 내고있다.

앞으로 서비스 단계에서 고객에게 LLM을 온프레미스로 제공하고 데이터 구축 방법을 제공한다면

고객이 직접 간단하게 fine-tuning이 가능해질 것이다.

Reference

QLoRA

Quantization survey

LQ-LoRA

Bitsandbytes

PEFT

https://towardsdatascience.com/introduction-to-weight-quantization-2494701b9c0c

https://huggingface.co/blog/4bit-transformers-bitsandbytes