LLM을 QLoRA로 파인 튜닝 하기

허깅페이스에서는 각종 AI모델을 쉽게 사용하도록 라이브러리를 제공하고 있는데, 개인용 PC(?) 수준에서도 LLM 모델을 테스트할 수 있도록 QLoRA로 LLM을 PEFT(Parameter Efficient Fine Tuning) 할 수 있는 방법을 제공하고 있다.

 

QLoRA는 아래의 이미지로 설명된다. LoRA가 Base 모델의 네트워크는 그대로 둔 채, 추가 데이터 만을 학습하여 본래의 Network에 concatenation하는 것이라면, QLoRA는 여기에 더해 16bit Network Node를 4bit로 양자화하고, 부족한 메모리로 큰 모델을 Handling할 수 있도록 Paging(바이너리를 나누어 2차 저장 장치에 Swapping하는 것)을 추가한 것이다. 16비트를 4비트로 바꾸었으니 정보량 손실이 일어날 수 밖에 없지만, 그 정도는 생각보다 크지 않다라는 것이 저자들의 주장이다.

 

아무튼 HuggingFace의 Transformer Package에 이를 핸들링할 수 있는 방법이 제공되었고, 

아래는 고맙게도 koalpaca 12.8b 모델 기반으로 이준범(https://github.com/Beomi)님이 작성방법을 공유해 주신 내용이다.

 

https://colab.research.google.com/gist/Beomi/a3032e4eaa33b86fdf8de1f47f15a647/2023_05_26_bnb_4bit_koalpaca_v1_1a_on_polyglot_ko_12_8b.ipynb

 

이준범님의 설명은 아래를 참고해도 좋다.

Transformers Trainer 뜯어보기

 

본 포스팅은 위 내용을 이해(됐나?)하면서 정리하는 것이다.

 

1) 필요 패키지 설치

!pip install -q -U bitsandbytes
!pip install -q -U git+https://github.com/huggingface/transformers.git 
!pip install -q -U git+https://github.com/huggingface/peft.git
!pip install -q -U git+https://github.com/huggingface/accelerate.git
!pip install -q datasets

 

bitsandbytes 패키지가 양자화 관련한 configuration을 하고, peft 패키지로 train한다.

 

2) Data Load

from datasets import load_dataset

data = load_dataset("beomi/KoAlpaca-v1.1a")

 

data의 형식은 transformer의 DatasetDict 형식을 따른다. 형식과 내용은 아래와 같다.

DatasetDict({
    train: Dataset({
        features: ['instruction', 'output', 'url'],
        num_rows: 21155
    })
})

data['train'][0]

{'instruction': '양파는 어떤 식물 부위인가요? 그리고 고구마는 뿌리인가요?',
 'output': '양파는 잎이 아닌 식물의 줄기 부분입니다. 고구마는 식물의 뿌리 부분입니다. \n\n식물의 부위의 구분에 대해 궁금해하는 분이라면 분명 이 질문에 대한 답을 찾고 있을 것입니다. 양파는 잎이 아닌 줄기 부분입니다. 고구마는 다른 질문과 답변에서 언급된 것과 같이 뿌리 부분입니다. 따라서, 양파는 식물의 줄기 부분이 되고, 고구마는 식물의 뿌리 부분입니다.\n\n 덧붙이는 답변: 고구마 줄기도 볶아먹을 수 있나요? \n\n고구마 줄기도 식용으로 볶아먹을 수 있습니다. 하지만 줄기 뿐만 아니라, 잎, 씨, 뿌리까지 모든 부위가 식용으로 활용되기도 합니다. 다만, 한국에서는 일반적으로 뿌리 부분인 고구마를 주로 먹습니다.',
 'url': 'https://kin.naver.com/qna/detail.naver?d1id=11&dirId=1116&docId=55320268'}

 

DatasetDict는 train/validation/test 데이터를 한 곳에 모아 놓은 Dictionary로 아래를 참조한다.

huggingface🤗 datasets 라이브러리 다루기

 

3) Model Load

import torch
from transformers import AutoTokenizer, AutoModelForCausalLM, BitsAndBytesConfig

model_id = "beomi/polyglot-ko-12.8b-safetensors"  # safetensors 컨버팅된 레포
bnb_config = BitsAndBytesConfig(
    load_in_4bit=True,
    bnb_4bit_use_double_quant=True,
    bnb_4bit_quant_type="nf4",
    bnb_4bit_compute_dtype=torch.bfloat16
)

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

 

중요한 함수는 BitsAndBytesConfig로 양자화를 위한 정보를 저장하여 모델 Load(AutoModelForCausalLM.from_pretrained) 시에 활용한다.  bnb_4bit_quant_type이 4Bit 양자화를 지정하는 것이고, bnb_4bit_compute_dtype은 이를 계산하기 위한 데이터 타입 지정이라고 하는데, 저장은 4비트라고 하더라도 16비트로 계산하는 것은 가능하다고 한다. 

tokenizer는 데이터를 token화하는데 사용하는 것으로 기반 모델의 tokenizer를 그대로 사용한다. tokenizer를 사용해 나오는 형식은 아래와 같다.

{'input_ids': [2, 19017, 8482, 3], 'token_type_ids': [0, 0, 0, 0], 'attention_mask': [1, 1, 1, 1]}

 

4) 학습 데이터 형식 구성하기

data = data.map(
    lambda x: {'text': f"### 질문: {x['instruction']}\n\n### 답변: {x['output']}<|endoftext|>" }
)

data['train']['text'][0]

### 질문: 양파는 어떤 식물 부위인가요? 그리고 고구마는 뿌리인가요?\n\n
### 답변: 양파는 잎이 아닌 식물의 줄기 부분입니다. 고구마는 식물의 뿌리 부분입니다. \n\n
식물의 부위의 구분에 대해 궁금해하는 분이라면 분명 이 질문에 대한 답을 찾고 있을 것입니다. 
양파는 잎이 아닌 줄기 부분입니다. 고구마는 다른 질문과 답변에서 언급된 것과 같이 뿌리 부분입니다. 
따라서, 양파는 식물의 줄기 부분이 되고, 고구마는 식물의 뿌리 부분입니다.\n\n 
덧붙이는 답변: 고구마 줄기도 볶아먹을 수 있나요? \n\n
고구마 줄기도 식용으로 볶아먹을 수 있습니다. 하지만 줄기 뿐만 아니라, 잎, 씨, 뿌리까지 모든 부위가 식용으로 
활용되기도 합니다. 다만, 한국에서는 일반적으로 뿌리 부분인 고구마를 주로 먹습니다.<|endoftext|>

data = data.map(lambda samples: tokenizer(samples["text"]), batched=True)

 

질문과 답변을 하나의 데이터('text')로 구성하고, 이를 tokenizer를 통해 위의 token 형식(input_ids, token_type_ids, attention_mask)을 만들어 data에 덧붙인다. 학습시에는 token 형식의 데이터를 활용하게 된다.

 

5) 학습 시작

 

학습 파라미터에 대한 내용은 많아서 아래 inline으로 작성하였다. 양자화를 위한 준비 -> LoraConfig 작성 -> peft 모델 가져오기 -> 모델 학습의 순서로 이루어진다고 보면 될 것 같다.

from peft import prepare_model_for_kbit_training

model.gradient_checkpointing_enable()
model = prepare_model_for_kbit_training(model)

from peft import LoraConfig, get_peft_model

# Lora Configuration
# target_modules는 model의 형태에 따라 다름
# https://junbuml.ee/lora-ckpt-size-mismatch 참고
config = LoraConfig(
    r=8, 
    lora_alpha=32, 
    target_modules=["query_key_value"], 
    lora_dropout=0.05, 
    bias="none", 
    task_type="CAUSAL_LM"
)

# 학습할 model을 configuration 적용하여 준비
model = get_peft_model(model, config)

import transformers

# needed for gpt-neo-x tokenizer
tokenizer.pad_token = tokenizer.eos_token

# 학습 진행. batch는 2, step은 50 정도로 짧게 수행
# data_collator는 data batch를 형성하는 개체라고 함
# https://huggingface.co/docs/transformers/main_classes/data_collator 참조
trainer = transformers.Trainer(
    model=model,
    train_dataset=data["train"],
    args=transformers.TrainingArguments(
        per_device_train_batch_size=2,
        gradient_accumulation_steps=1,
        max_steps=50, ## 초소량만 학습: 50 step만 학습. 약 4분정도 걸립니다.
        learning_rate=1e-4,
        fp16=True,
        logging_steps=10,
        output_dir="outputs",
        optim="paged_adamw_8bit"
    ),
    data_collator=transformers.DataCollatorForLanguageModeling(tokenizer, mlm=False),
)
model.config.use_cache = False  # silence the warnings. Please re-enable for inference!
trainer.train()

 

4~5분 정도의 학습으로 나만의 LLM 모델을 갖게 된다 @.@

 

6) 테스트

 

model.generate 함수를 통해 생성형 AI 결과를 확인할 수 있는데, 이준범님은 아래와 같은 간단한 테스트 함수를 만들었다. 

def gen(x):
    gened = model.generate(
        **tokenizer(
            f"### 질문: {x}\n\n### 답변:", 
            return_tensors='pt', 
            return_token_type_ids=False
        ), 
        max_new_tokens=256,
        early_stopping=True,
        do_sample=True,
        eos_token_id=2,
    )
    print(tokenizer.decode(gened[0]))

 

따로 설명이 필요 없는 것 같아서 바로 테스트 진행.

gen('어린이 동화에 늑대가 자주 등장하는 이유는?')

### 질문: 어린이 동화에 늑대가 자주 등장하는 이유는?

### 답변: 어린이들 사이에서 늑대가 사람의 성격, 말투나 행동에 따라 의인화되어 표현되기도 합니다. 
이런 이유에서 동화 속의 늑대가 어린이들에게 두려움의 대상인 것이 아니라 오히려 친근한 것으로 인식되기도 합니다. 
또 옛날 우리나라 시골의 밤을 무서워하는 어린아이에게 어미 늑대가 새끼를 물고 와서 달을 보고 
"이제부터 내가 너를 잘 지키마." 하고 밤새워 아이를 지킨 일로 인하여 '효심이 지극함', '믿음직스러움', '효도' 와 같은 덕목을 상징합니다. 
늑대가 이런 의인화된 이미지로 사람들에게 인식되는 것 같이 어린이들에게 이런 인식이 되어 있는 것은 어린이의 심리의 흐름에 있다고 합니다. 
어린이가 늑대가 자기를 위협하고 있다고 느끼지 않고 늑대가 자기가 생각했던 것보다 강하지 않는 것으로 
느끼게 되는 것은, 늑대를 보았을 때 어른은 늑대를 두려워했다면, 아이는 그 자체를 인정하고 무서움과 두려움을 
느낄 수 없게 된다는 것입니다. 여기서의 생각의 흐름을 '늑대는 약하다'라는 인식의 흐름으로 받아들이는 것입니다. 
동화에서

 

토큰 수 때문에 대답이 잘렸지만 제법 그럴싸한 대답을 내놓는다. 다만, PEFT가 내가 추가로 입력한 데이터의 성격과 내용에 맞게 결과를 뱉는거라고 한다면, 그리 좋은 결과를 내 놓는 것 같지는 않다. 학습 데이터와는 다소 상이한 내용으로 대답하기도 한다.