[LLM] 오픈소스 LLM 파인튜닝 - 자체 데이터셋으로 LLaMA2 기반 한국어 LLM 파인튜닝하기 (feat. LoRA)

자체 데이터셋으로 LLaMA2 파인튜닝하기

 

너무 귀여워 라마..

 

 

지난 포스팅에서 LLaMA2를 한국어 데이터셋으로 파인튜닝한 

한국어 LLM 모델 (Kor-Orca-Platypus-13B)을 다운받아서 inference 해보고, 성능을 평가해봤습니다.

 

 

[이전글] : [LLM] Ko-LLM 리뷰, LLaMA2 기반 한국어 파인튜닝 모델 인퍼런스

[LLM] Ko-LLM 리뷰, LLaMA2 기반 한국어 파인튜닝 모델 인퍼런스

 

 

이번에는 자체 데이터셋을 활용해 해당 모델을 파인튜닝 해본 내용을 공유하려고 합니다.

 

LaMMA2 모델은 7B, 13B, 13B 모델이 성능이 가장 좋지만, 컴퓨팅 자원이나 학습에 소요되는 시간 등

아무래도 full fine-tuning은 힘든 상황이죠.

 

효율적인 학습을 위해 PEFT(Parameter Efficient Fine-Tuning) 기법 중 하나인

LoRA(Low-Rank Adaptation)를 이용해서 학습을 진행했습니다.

 

[LLM] LoRA (Low-Rank Adaptation) 를 이용한 LLM (Large Language Model) 최적화

 

 

파인튜닝 코드는 beomi/KoAlpaca에서 진행한 내용과  mete 공식 llama-recipes gihub 문서, tloen/alpaca-lora 의 finetune.py 코드를 참고했습니다.

 

• beomi/KoAlpaca : https://github.com/Beomi/KoAlpaca
• llama-recipes : https://github.com/facebookresearch/llama-recipes
• tloen/alpaca-lora :  https://github.com/tloen/alpaca-lora/tree/main

GitHub - tloen/alpaca-lora: Instruct-tune LLaMA on consumer hardware

 

저는 폐쇄망에 있는 데이터셋을 이용해 파인튜닝을 진행해서,

tloen/alpaca-lora github에 있는 파일 통째로 zip 으로 말아서 반입한 뒤 코드를 수정했습니다.

 

우선 동작이 잘 되는지 확인해 보기 위해 모델 경로 및 데이터 경로만 바꿔주고

해당 코드를 거의 그대로 사용했습니다.

 

라이브러리 설치

 

inference 때와 마찬가지로, LLaMA를 사용하기 위해 transformer는 허깅페이스에서 최신 버전을 받아와 설치해줬습니다.

그리고 requirment.txt 에 있는 라이브러리 들을 설치해줍니다.

이때도 폐쇄망 환경에 설치된 버전이 없는 파일들은 직접 whl 파일을 받아서 설치해줍니다.

 

pip install git+https://github.com/huggingface/transformers # 최신 transformers 설치
pip install alpaca-lora/requirements.txt

 

 

그리고 아주 간단하게는

그냥 아래 코드를 자기가 원하는 모델과 데이터 경로로 바꿔서 실행시키면 됩니다. ㅎㅎ

 

python finetune.py \
    --base_model 'KoR-Orca-Platypus-13B' \ # 돌리고 싶은 LLM 모델 경로
    --data_path 'HR/dataset/train.json' \ # 학습 데이터 경로
    --output_dir './lora-alpaca_1102' # 모델 학습 결과 저장 경로

 

 

저는 내부 코드도 볼 겸 실행 결과를 하나씩 확인해보고 싶어서 

주피터 노트북으로 학습을 진행해봤습니다. (그러다가 train 함수를 통째로 돌렸지만..)

 

아래는 주피터 노트북에서 학습을 진행할 경우에 참고하시면 좋을 것 같네요 ㅎㅎ

 

라이브러리 임포트 

finetune.py 코드를 뜯어서 단계 별로 실행해보죠. 우선 필요 라이브러리들을 임포트 합니다.

 

import os
import sys
from typing import List

import fire
import torch
import transformers
from datasets import load_dataset

"""
Unused imports:
import torch.nn as nn
import bitsandbytes as bnb
"""

from peft import (
    LoraConfig,
    get_peft_model,
    get_peft_model_state_dict,
    prepare_model_for_int8_training,
    set_peft_model_state_dict,
)
from transformers import LlamaForCausalLM, LlamaTokenizer, AutoTokenizer, AutoModelForCausalLM

from utils.prompter import Prompter

 

 

모델 학습

 

그 다음 실제 학습에 쓰이는 train 함수를 정의합니다.

 

 

먼저 앞부분에서 살펴볼 부분은 base_model, data_path, output_dir로 

아까 스크립트 파일로 실행했을때 인자로 넣어줬던 부분입니다.

저는 주피터 노트북에서 순차적으로 실행하므로 ㅎㅎ 그냥 따로 변수를 선언해줬습니다. 

 

base_model = "KoR-Orca-Platypus-13B",  # 파인튜닝 하고 싶은 모델 경로
data_path = "HR/dataset/train.json", # 학습 데이터 경로
output_dir = "./lora-alpaca_1102", # 모델 학습 결과 저장 경로

 

그 train 함수 내에서 model과 tokenizer를 불러오는 부분을 보니,

LlamaForCausalLM.from_pretrained()와 LlamaTokenizer.from_pretrained()를 이용해 받아오는 모습을 확인할 수 있네요.

device_map은 위에서 "auto"로 정의되어 있습니다.

(아래 코드는 그냥 참고용)

 

# base model 로드
model = LlamaForCausalLM.from_pretrained(
        base_model,
        load_in_8bit=True,
        torch_dtype=torch.float16,
        device_map=device_map,
    )

# 토크나이저 로드
tokenizer = LlamaTokenizer.from_pretrained(base_model)

 

 

또 LoRA를 적용해 학습을 진행하기 위해서, model에 LoRA config를 적용하는 부분을 보죠.

(아래 코드도 그냥 참고만)

 

model = prepare_model_for_int8_training(model)

# config로 LoreConfig 할당
config = LoraConfig(
    r=lora_r,
    lora_alpha=lora_alpha,
    target_modules=lora_target_modules,
    lora_dropout=lora_dropout,
    bias="none",
    task_type="CAUSAL_LM",
)
# model에 LoreConfig 적용
model = get_peft_model(model, config)

 

 

이렇게 LoRA까지 적용해주고 (너무 대충 보는 것 아닌가요...)

Transformer의 Trainer를 이용해 모델을 학습합니다.  (아래 코드는 참고용)

 

trainer = transformers.Trainer(
    model=model,
    train_dataset=train_data,
    eval_dataset=val_data,
    args=transformers.TrainingArguments(
        per_device_train_batch_size=micro_batch_size,
        gradient_accumulation_steps=gradient_accumulation_steps,
        warmup_steps=100,
        num_train_epochs=num_epochs,
        learning_rate=learning_rate,
        fp16=True,
        logging_steps=10,
        optim="adamw_torch",
        evaluation_strategy="steps" if val_set_size > 0 else "no",
        save_strategy="steps",
        eval_steps=200 if val_set_size > 0 else None,
        save_steps=200,
        output_dir=output_dir,
        save_total_limit=3,
        load_best_model_at_end=True if val_set_size > 0 else False,
        ddp_find_unused_parameters=False if ddp else None,
        group_by_length=group_by_length,
        report_to="wandb" if use_wandb else None,
        run_name=wandb_run_name if use_wandb else None,
    ),
    data_collator=transformers.DataCollatorForSeq2Seq(
        tokenizer, pad_to_multiple_of=8, return_tensors="pt", padding=True
    ),
)
model.config.use_cache = False

old_state_dict = model.state_dict
model.state_dict = (
    lambda self, *_, **__: get_peft_model_state_dict(
        self, old_state_dict()
    )
).__get__(model, type(model))

if torch.__version__ >= "2" and sys.platform != "win32":
    model = torch.compile(model)

# 학습 진행시켜!
trainer.train(resume_from_checkpoint=resume_from_checkpoint)

# 학습한 모델 저장
model.save_pretrained(output_dir)

 

 

모든게 다 합쳐진 전체 train 코드는 다음과 같습니다.

 

def train(
    # model/data params
    base_model: str = "",  # the only required argument
    data_path: str = "yahma/alpaca-cleaned",
    output_dir: str = "./lora-alpaca",
    # training hyperparams
    batch_size: int = 128,
    micro_batch_size: int = 4,
    num_epochs: int = 3,
    learning_rate: float = 3e-4,
    cutoff_len: int = 256,
    val_set_size: int = 2000,
    # lora hyperparams
    lora_r: int = 8,
    lora_alpha: int = 16,
    lora_dropout: float = 0.05,
    lora_target_modules: List[str] = [
        "q_proj",
        "v_proj",
    ],
    # llm hyperparams
    train_on_inputs: bool = True,  # if False, masks out inputs in loss
    add_eos_token: bool = False,
    group_by_length: bool = False,  # faster, but produces an odd training loss curve
    # wandb params
    wandb_project: str = "",
    wandb_run_name: str = "",
    wandb_watch: str = "",  # options: false | gradients | all
    wandb_log_model: str = "",  # options: false | true
    resume_from_checkpoint: str = None,  # either training checkpoint or final adapter
    prompt_template_name: str = "alpaca",  # The prompt template to use, will default to alpaca.
):
    if int(os.environ.get("LOCAL_RANK", 0)) == 0:
        print(
            f"Training Alpaca-LoRA model with params:\n"
            f"base_model: {base_model}\n"
            f"data_path: {data_path}\n"
            f"output_dir: {output_dir}\n"
            f"batch_size: {batch_size}\n"
            f"micro_batch_size: {micro_batch_size}\n"
            f"num_epochs: {num_epochs}\n"
            f"learning_rate: {learning_rate}\n"
            f"cutoff_len: {cutoff_len}\n"
            f"val_set_size: {val_set_size}\n"
            f"lora_r: {lora_r}\n"
            f"lora_alpha: {lora_alpha}\n"
            f"lora_dropout: {lora_dropout}\n"
            f"lora_target_modules: {lora_target_modules}\n"
            f"train_on_inputs: {train_on_inputs}\n"
            f"add_eos_token: {add_eos_token}\n"
            f"group_by_length: {group_by_length}\n"
            f"wandb_project: {wandb_project}\n"
            f"wandb_run_name: {wandb_run_name}\n"
            f"wandb_watch: {wandb_watch}\n"
            f"wandb_log_model: {wandb_log_model}\n"
            f"resume_from_checkpoint: {resume_from_checkpoint or False}\n"
            f"prompt template: {prompt_template_name}\n"
        )
    assert (
        base_model
    ), "Please specify a --base_model, e.g. --base_model='huggyllama/llama-7b'"
    gradient_accumulation_steps = batch_size // micro_batch_size

    prompter = Prompter(prompt_template_name)

    device_map = "auto"
    world_size = int(os.environ.get("WORLD_SIZE", 1))
    ddp = world_size != 1
    if ddp:
        device_map = {"": int(os.environ.get("LOCAL_RANK") or 0)}
        gradient_accumulation_steps = gradient_accumulation_steps // world_size

    # Check if parameter passed or if set within environ
    use_wandb = len(wandb_project) > 0 or (
        "WANDB_PROJECT" in os.environ and len(os.environ["WANDB_PROJECT"]) > 0
    )
    # Only overwrite environ if wandb param passed
    if len(wandb_project) > 0:
        os.environ["WANDB_PROJECT"] = wandb_project
    if len(wandb_watch) > 0:
        os.environ["WANDB_WATCH"] = wandb_watch
    if len(wandb_log_model) > 0:
        os.environ["WANDB_LOG_MODEL"] = wandb_log_model

    model = LlamaForCausalLM.from_pretrained(
        base_model,
        load_in_8bit=True,
        torch_dtype=torch.float16,
        device_map=device_map,
    )

    tokenizer = LlamaTokenizer.from_pretrained(base_model)

    tokenizer.pad_token_id = (
        0  # unk. we want this to be different from the eos token
    )
    tokenizer.padding_side = "left"  # Allow batched inference

    def tokenize(prompt, add_eos_token=True):
        # there's probably a way to do this with the tokenizer settings
        # but again, gotta move fast
        result = tokenizer(
            prompt,
            truncation=True,
            max_length=cutoff_len,
            padding=False,
            return_tensors=None,
        )
        if (
            result["input_ids"][-1] != tokenizer.eos_token_id
            and len(result["input_ids"]) < cutoff_len
            and add_eos_token
        ):
            result["input_ids"].append(tokenizer.eos_token_id)
            result["attention_mask"].append(1)

        result["labels"] = result["input_ids"].copy()

        return result

    def generate_and_tokenize_prompt(data_point):
        full_prompt = prompter.generate_prompt(
            data_point["instruction"],
            data_point["input"],
            data_point["output"],
        )
        tokenized_full_prompt = tokenize(full_prompt)
        if not train_on_inputs:
            user_prompt = prompter.generate_prompt(
                data_point["instruction"], data_point["input"]
            )
            tokenized_user_prompt = tokenize(
                user_prompt, add_eos_token=add_eos_token
            )
            user_prompt_len = len(tokenized_user_prompt["input_ids"])

            if add_eos_token:
                user_prompt_len -= 1

            tokenized_full_prompt["labels"] = [
                -100
            ] * user_prompt_len + tokenized_full_prompt["labels"][
                user_prompt_len:
            ]  # could be sped up, probably
        return tokenized_full_prompt

    model = prepare_model_for_int8_training(model)

    config = LoraConfig(
        r=lora_r,
        lora_alpha=lora_alpha,
        target_modules=lora_target_modules,
        lora_dropout=lora_dropout,
        bias="none",
        task_type="CAUSAL_LM",
    )
    model = get_peft_model(model, config)

    if data_path.endswith(".json") or data_path.endswith(".jsonl"):
        data = load_dataset("json", data_files=data_path)
    else:
        data = load_dataset(data_path)

    if resume_from_checkpoint:
        # Check the available weights and load them
        checkpoint_name = os.path.join(
            resume_from_checkpoint, "pytorch_model.bin"
        )  # Full checkpoint
        if not os.path.exists(checkpoint_name):
            checkpoint_name = os.path.join(
                resume_from_checkpoint, "adapter_model.bin"
            )  # only LoRA model - LoRA config above has to fit
            resume_from_checkpoint = (
                False  # So the trainer won't try loading its state
            )
        # The two files above have a different name depending on how they were saved, but are actually the same.
        if os.path.exists(checkpoint_name):
            print(f"Restarting from {checkpoint_name}")
            adapters_weights = torch.load(checkpoint_name)
            set_peft_model_state_dict(model, adapters_weights)
        else:
            print(f"Checkpoint {checkpoint_name} not found")

    model.print_trainable_parameters()  # Be more transparent about the % of trainable params.

    if val_set_size > 0:
        train_val = data["train"].train_test_split(
            test_size=val_set_size, shuffle=True, seed=42
        )
        train_data = (
            train_val["train"].shuffle().map(generate_and_tokenize_prompt)
        )
        val_data = (
            train_val["test"].shuffle().map(generate_and_tokenize_prompt)
        )
    else:
        train_data = data["train"].shuffle().map(generate_and_tokenize_prompt)
        val_data = None

    if not ddp and torch.cuda.device_count() > 1:
        # keeps Trainer from trying its own DataParallelism when more than 1 gpu is available
        model.is_parallelizable = True
        model.model_parallel = True

    trainer = transformers.Trainer(
        model=model,
        train_dataset=train_data,
        eval_dataset=val_data,
        args=transformers.TrainingArguments(
            per_device_train_batch_size=micro_batch_size,
            gradient_accumulation_steps=gradient_accumulation_steps,
            warmup_steps=100,
            num_train_epochs=num_epochs,
            learning_rate=learning_rate,
            fp16=True,
            logging_steps=10,
            optim="adamw_torch",
            evaluation_strategy="steps" if val_set_size > 0 else "no",
            save_strategy="steps",
            eval_steps=200 if val_set_size > 0 else None,
            save_steps=200,
            output_dir=output_dir,
            save_total_limit=3,
            load_best_model_at_end=True if val_set_size > 0 else False,
            ddp_find_unused_parameters=False if ddp else None,
            group_by_length=group_by_length,
            report_to="wandb" if use_wandb else None,
            run_name=wandb_run_name if use_wandb else None,
        ),
        data_collator=transformers.DataCollatorForSeq2Seq(
            tokenizer, pad_to_multiple_of=8, return_tensors="pt", padding=True
        ),
    )
    model.config.use_cache = False

    old_state_dict = model.state_dict
    model.state_dict = (
        lambda self, *_, **__: get_peft_model_state_dict(
            self, old_state_dict()
        )
    ).__get__(model, type(model))

    if torch.__version__ >= "2" and sys.platform != "win32":
        model = torch.compile(model)

    trainer.train(resume_from_checkpoint=resume_from_checkpoint)

    model.save_pretrained(output_dir)

    print(
        "\n If there's a warning about missing keys above, please disregard :)"
    )

 

길다..

하나씩 자세히 뜯어보면 더 좋겠지만, 우선 모델 학습에 핵심적인 부분 위주로 먼저 살펴보았습니다.

 

이제 fire를 이용해 해당 함수를 돌려주면 끝입니다.

 

fire 패키지는 Python에서의 모든 객체를 command line interface로 만들어 주는데,

python 객체(함수, 클래스, dictionary, list, tuple) 면 모두 다 호출이 가능하다고 합니다.

 

fire.Fire를 이용하면 마치 폭탄 심지에 불을 붙이는 것처럼, 해당 함수에 작성된 코드를 순차적으로 실행합니다.

 

fire.Fire(train) # 학습 시작

 

 

자 이제 학습이 시작되었습니다!!! 

 

주피터 노트북으로 실행하면 학습 과정이 실시간으로 화면에 출력됩니다.

저는 어느 정도 학습을 시키다가,

train loss는 0.1 이하로 계속 떨어지는데 validation loss 가 더이상 줄어들지 않고 발산하는 모습을 확인하고

오버피팅 되었다고 생각해서 학습을 멈췄습니다.

 

 

이제 평가 데이터셋을 활용해서 성능을 평가해봅시다.

 

화이팅!