AI-LLM 파인튜닝 한방에 끝내기(gemma-2B)

오픈소스 LLM이 매일같이 나오는 상황. 파인튜닝을 해보고 싶다.

하지만, 파인튜닝을 검색해보면, 양자화, PEFT, LoRA 등등 알아야 할것도 많고, 막상 코드를 돌릴려고 하니, 알아야 할 것이 많고 복잡하다.

한번은 스스로 정리가 필요하다고 판다. 거대 언어모델 Gemma-2B(20억개) 파인튜닝하는 과정을 직접 정리해 보도록 하겠다.

 

0. 목표

구글에서 발표한 비교적 작은 모델인 Gemma-2B를 파인튜닝하여, 문자내용을 보고 spam인지 아닌지를 판별하는 모델을 만들어 보겠다.

물론 classification모델에 Generation모델을 사용하는 것은 오버 스팩일수 있다. 하지만, 파인튜닝 학습이기도 하고, 사람처럼 어려운 스팸도 걸러내는 LLM의 능력을 보고자 한다.

 

1. 셋팅

허깅페이스 로그인 소스코드. 창이 뜨면 자신의 키값을 입력하여 로그인 한다.

추후 데이터셋을 가져오고, 모델을 업로드 할때 사용예정

from huggingface_hub import notebook_login

notebook_login()

 

필요한 라이브러리를 설치하자. 구글 gemma-2b파인튜닝 코드를 참고해서 필요한 라이브러리 설치

그리고 임포트

!pip3 install -q -U bitsandbytes==0.42.0

!pip3 install -q -U peft==0.8.2

!pip3 install -q -U trl==0.7.10

!pip3 install -q -U accelerate==0.27.1

!pip3 install -q -U datasets==2.17.0

!pip3 install -q -U transformers==4.38.1

import os

import os.path as osp

import sys

import json

from typing import List, Union

import torch

from torch.nn import functional as F

import transformers

from transformers import TrainerCallback, TrainingArguments, TrainerState, TrainerControl

from transformers.trainer_utils import PREFIX_CHECKPOINT_DIR

from transformers import AutoModelForCausalLM, AutoTokenizer

from datasets import load_dataset

 

2. 모델가져오기

GPU가 없기 때문에, 기존 32bit의 모델을 4bit구조로 가져온다

약 1/8사이즈로 양자화한 모델을 가져온다

import torch

from transformers import AutoTokenizer, AutoModelForCausalLM, BitsAndBytesConfig, GemmaTokenizer

model_id = "google/gemma-2b"

bnb_config = BitsAndBytesConfig(

#입력값을 4bit로 변환

load_in_4bit=True,

#모델을 4bit로 양자화

bnb_4bit_quant_type="nf4",

#4bit계산에 사용될 데이터유형, 4비트 부동소수점(bfloat16), 4비트정수(uint8)

bnb_4bit_compute_dtype=torch.bfloat16

)

device = "auto"

tokenizer = AutoTokenizer.from_pretrained(model_id)

# 2B모델을 4비트로 양자화 하여 가져오니 1/8사이즈로 가져오는 것

model = AutoModelForCausalLM.from_pretrained(model_id, quantization_config=bnb_config, device_map=device)

 

양자화 설정값을 bnb_config로 설정(주석참고)

사용하는 모델가 매칭되는 토크나이저를 가져오고,

양자화된 gemma-2b를 가져온다

 

3. 토크나이저-사전 체크

사용하는 모델은 자신만의 사전(vocab)을 가지고 있다.

그 사전은 형태소-숫자가 1:1로 매칭되어 있다.

예를 들면 name - 100, korea - 101, ah - 102... 이런식으로.

토크나이저-사전체크의 이유는 예전에 다른 모델을 사용할때, 'spam'이라는 단어가 막상 모델의 토크나이저에 없어서,

효과적인 학습이 안된경우가 있기 때문이다. 막상 구분은 하지만, spam이라는 단어를 못 뱉어내는 경우가 있게된다.

아래 코드를 통해서, 내가 사용할 모델의 토큰과 숫자의 매칭을 확인해보자. 그리고, 학습에 중요한 단어가 들어있는지 체크해보자

word = "spam"

token_id = tokenizer.convert_tokens_to_ids(word)

print(f'The token ID for "{word}" is: {token_id}')

word = "not"

token_id = tokenizer.convert_tokens_to_ids(word)

print(f'The token ID for "{word}" is: {token_id}')

# "not spam" 문자열의 토큰 ID 조회

text = "not spam"

token_ids = tokenizer.encode(text, add_special_tokens=False)

print(f'The token IDs for "{text}" are: {token_ids}')

# 100부터 110까지의 인덱스 리스트 생성

indices = list(range(100, 111))

# decode() 함수를 사용하여 단어 출력

decoded_text = tokenizer.decode(indices)

print(decoded_text)

 

물론, 토크나이저에 새로운 단어/형태소를 추가할 수 있다. 하지만, 실제 수정된 토크나이저로 학습시켜봤는데... 성과가 그렇게 좋지는 않다. 아마 pretrain시 해당 단어/형태소가 학습되지 않았기 때문에, 파인튜닝 조금 학습한다고 되지 않는것 같다.

 

4. LLM모델 내용 체크

eos, bos, pad 토큰 체크

기본적으로 패딩 side는 right

###########################################

# 3-2. BOS, EOS, PAD 토큰 확인

# Check special token

bos = tokenizer.bos_token_id # 문장 시작 토큰

eos = tokenizer.eos_token_id # 문장 끝 토큰

pad = tokenizer.pad_token_id # 문장 패딩 토큰

tokenizer.padding_side = "right" # 패딩 오른쪽, to prevent warnings

print("BOS token:", bos) # 1

print("EOS token:", eos) # 2

print("PAD token:", pad) # None

 

5. 하이퍼 파라미터 셋팅

언어모델의 lr 확인해보시고,

batch_size는 GPU용량에 따라 조절하시면 됩니다.

# 하이퍼 파라미터 셋팅

# 데이터셋과 훈련 횟수와 관련된 하이퍼 파라미터

batch_size = 16

num_epochs = 1

micro_batch = 1

gradient_accumulation_steps = batch_size // micro_batch

# 훈련 방법에 대한 하이퍼 파라미터

cutoff_len = 4096

lr_scheduler = 'cosine'

warmup_ratio = 0.06 # warmup_steps = 100

learning_rate = 4e-4

optimizer = 'adamw_torch'

weight_decay = 0.01

max_grad_norm = 1.0

 

# Tokenizer에서 나오는 input값 설정 옵션

train_on_inputs = False

add_eos_token = False

 

6. LoRA셋팅

중요!!! LoRA를 어디다 붙일것인가!!

# LoRA config

lora_r = 16 #lora 가운데 차원

lora_alpha = 16 #lora 스케일링 alpha/r

lora_dropout = 0.05

lora_target_modules = ["gate_proj", "down_proj", "up_proj"]

위 설정은 내가 LoRA를 어디다 붙일것인지 설정하는 것.

LoRA 개념 대한 상세한 설명은 다른 글에서 확인하세요.

 

아래는 원래 Gemma-2B모델

Transformer의 구조인 attention(attn)과 멀티해드피드포워드네트워크(mlp)로 구성되어 있음.

여기서 LoRA설정을 통해서 어디에 LoRA를 붙일지 설정하는것.

lora_target_modules = ["gate_proj", "down_proj", "up_proj"]

mlp중 gate, down, up에 LoRA를 붙입니다

 

LoRA를 붙인후, 모델의 모습

attention은 그대로이고, mlp의 gate, down, up에 LoRA가 붙은것을 볼수 있다.

 

gemma말고 다른모델에도 붙일수 있다.

bloomz560m모델에 붙여보면, 우선 원래 모델을 확인하면 아래와 같다.

mlp의 dense_h_to_4h, dense_4h_to_h에 붙인다고 하면

lora_target_modules = ["dense_h_to_4h", "dense_4h_to_h"]

이렇게 설정하면 된다

 

LoRA붙이는 방법을 이해했으니, 최종적으로 코드로 LoRA를 붙인모델을 완성한다

from peft import LoraConfig

# LoRA옵션값 설정

lora_config = LoraConfig(

r=lora_r,

lora_alpha=lora_alpha,

#LoRA를 붙이는 위치로, attention쪽, MLP쪽 등 내가 원하는 곳에 붙일수 있다

target_modules=lora_target_modules,

lora_dropout=lora_dropout,

bias="none",

task_type="CAUSAL_LM")

 

###########################################

# Model with LoRA

from peft import (

get_peft_model,

prepare_model_for_kbit_training

)

 

# 위에서 4bit로 양자한 모델을 준비

# 모델을 LoRA붙일수 있게 셋팅

model = prepare_model_for_kbit_training(model)

print(model)

# LoRA붙이기

model = get_peft_model(model, lora_config) # Applying LoRA

print(model)

 

7. 파인튜닝 학습용 데이터셋 셋팅

7-1. spam 데이터 가져오기

dataset_name = "hf/spam_data" #dataset주소 넣으세요

dataset = load_dataset(dataset_name)

print(dataset["train"])

 

7-2. 데이터셋 preprocessing

문자열 데이터의 길이를 최대 200이내로 조정

# 'sms' 열의 문자열을 200자 이내로 조정하는 함수

def truncate_sms(examples):

truncated_sms = [text[:200] if len(text) > 200 else text for text in examples['sms']]

return {"sms": truncated_sms}

# 데이터셋에 truncate_sms 함수 적용

dataset = dataset.map(truncate_sms, batched=True, num_proc=1)

dataset["train"][0]

 

7-3. instruction tuning용 데이터로 만들기

중요!!!

이게 기존의 ML과 차이점일수 있는데, pretrain된 모델이 언어의 특성을 이미 인식하고 있기 때문에

언어모델스럽게 학습을 하는 것이다.

기존 ML의 구조가 

input: 문자내용, output: 0(not spam)/1(spam)의 형태라면,

 

우리의 LLM 파인튜닝은

input: Classify if the e-mail is a spam or not:\n\n### title:{문자내용}\n\n### label:\n
output: spam/not spam

이렇게 구성된다.

마치 gpt에 물어보는 것처럼 데이터셋을 만들고, 학습한다.

그리고 LLM 파인튜닝은 instruction tuning이 잘된다고 하니, 꼭 사용하자!

###########################################
# Instruction tuning을 위한 template 작성.

instruct_template = {
    "prompt_no_input": "Classify if the sms is a spam or not:\n\n### sms:\n{instruction}\n\n### label:\n",
    "response_split": "### label:\n"
}

###########################################
# 5-3. 데이터셋 불러오는 클래스

class Prompter(object):

    def __init__(self, verbose: bool = False):
        # 문장 "스팸구분하시오 sms: {instruction}...{label}"
        self.template = instruct_template

    # 완성된 문장을 출력
    def generate_prompt(
        self,
        instruction: str,
        # input: Union[None, str] = None,
        label: Union[None, str] = None,
    ) -> str:

        res = self.template["prompt_no_input"].format(instruction=instruction)
        res = f"{res}{label}"
        
        return res

    def get_response(self, output: str) -> str:
        # 두개값(입력,label)중 label값을 출력
        return output.split(self.template["response_split"])[1].strip()

prompter = Prompter()

###########################################
# 5-4. Token generation 함수

def tokenize(prompt, add_eos_token=True):
    # 입력된 프롬프트를 토크나이징, 길면 cutoff_len 글자 짜르고
    result = tokenizer(
        prompt,
        truncation=True,
        max_length=cutoff_len,
        padding=False,
        # 파이토치 텐서 아닌, 일반 파이선객체 반환
        return_tensors=None,
    )

    if (
        # 토큰화된 것의 마지막 값이 eos가 아니고, 길이가 작으며, eos가 아닌경우
        result["input_ids"][-1] != tokenizer.eos_token_id
        and len(result["input_ids"]) < cutoff_len
        and add_eos_token
    ):

        # eos추가(완전한 문장 만들기)
        result["input_ids"].append(tokenizer.eos_token_id)
        result["attention_mask"].append(1)

    # input_ids, labels 같은문장
    result["labels"] = result["input_ids"].copy()

    return result

# 문장 만들고 > input, label만들고 > 마스킹하기
def generate_and_tokenize_prompt(data_point):
    # Full text로 전체 문장 생성
    full_prompt = prompter.generate_prompt(data_point["sms"], data_point["text_label"])
    # 전체 문장 토큰화
    tokenized_full_prompt = tokenize(full_prompt, add_eos_token=True)

    # 레이블 토큰 길이 계산
    label_tokens = tokenizer.encode(data_point["text_label"], add_special_tokens=True)
    # spam = [spam, 1] =2 , not spam = [not spam, 1]=3
    label_len = len(label_tokens)

    # 전체 토큰 길이
    total_length = len(tokenized_full_prompt["input_ids"])

    # 레이블 부분을 제외한 나머지는 -100으로 패딩
    labels_masked = [-100] * (total_length - label_len) + tokenized_full_prompt["input_ids"][-label_len:]

    # 레이블에 -100 패딩 적용
    tokenized_full_prompt["labels"] = labels_masked

    return tokenized_full_prompt

주석을 잘 살펴보면 되는데 간단히 정리하면,

- 데이터셋을 > instruction tuning형태의 데이터셋으로 바꾸고,

- 각 데이터를 토큰화하고, label은 -100으로 채우고, label값만 토큰값으로 구성

- input_ids는 입력값을 토큰화 한것

- attention_mask는 각 토큰에 대해 attention할지 말지(Transformer)

- labels는 입력값과 동일한 문장이지만, label만 있고, 나머지는 -100

이렇게 구성되면 LLM 모델의 학습이 labels값을 확인할때 이뤄집니다.

이 말은 입력값이 spam인지 not spam인지 출력할때마다, LoRA파라미터가 조정되는 것

 

7-4. 학습데이터셋 구성

###########################################

# 훈련 셋 만들기 (~3분)

val_data = dataset["test"].shuffle() # random

val_data = val_data.map(generate_and_tokenize_prompt)

train_data = dataset["train"].shuffle() # random

train_data = train_data.map(generate_and_tokenize_prompt)

 

8. 훈련하기

Trainer설정을 하고, 학습을 시작합니다

###########################################

# 7-2. Trainer class 정의

trainer = transformers.Trainer(

model=model,

train_dataset=train_data,

eval_dataset=val_data,

args=transformers.TrainingArguments( # 훈련에 이용될 하이퍼파라미터

# per_device_train_batch_size = micro_batch, #각 디바이스(예: GPU)에서 훈련할 배치 크기

per_device_train_batch_size = 4, #각 디바이스(예: GPU)에서 훈련할 배치 크기

# gradient_accumulation_steps = gradient_accumulation_steps, #실제로 업데이트 되기 전에 여러 스텝에 걸쳐 그라디언트를 누적

gradient_accumulation_steps = 2, #실제로 업데이트 되기 전에 여러 스텝에 걸쳐 그라디언트를 누적

warmup_ratio=warmup_ratio, #워밍업 기간 동안 학습률이 증가하는 비율

# num_train_epochs=3,

num_train_epochs=1,

learning_rate=learning_rate,

fp16=True, #16비트 부동 소수점을 사용할지 여부

logging_steps=1,

optim=optimizer,

evaluation_strategy="no",

save_strategy="steps",

max_grad_norm = max_grad_norm,

save_steps = 500, # you can change!

lr_scheduler_type=lr_scheduler,

output_dir=output_dir,

save_total_limit=2,

load_best_model_at_end=False,

ddp_find_unused_parameters=False,

group_by_length = False

),

data_collator=transformers.DataCollatorForSeq2Seq(

tokenizer, pad_to_multiple_of=8, return_tensors="pt", padding=True

),

)

model.config.use_cache = False

model.print_trainable_parameters() # 훈련하는 파라미터의 % 체크

if torch.__version__ >= "2" and sys.platform != "win32":

model = torch.compile(model)

 

###########################################

# 7-3. Training (fine-tuning)

torch.cuda.empty_cache()

trainer.train(resume_from_checkpoint=resume_from_checkpoint)

 

9. 학습이 완료된 모델 업로드

학습이 완료된 모델을 허깅페이스에 올립니다

model_name = f"여러분의 계정/gemma-2b_LoRA"

model.push_to_hub(model_name, use_auth_token=True)

 

10. 정리

LoRA파인 튜닝을 정리해보면

- 양자화된 모델을 가져오고,

- 토크나이저 함 체크하고

- LoRA설정해서, 붙이고

- instruction tuning용 데이터셋 구성하고

- 파인튜닝하고,

- 학습끝난 모델 업로드!