Hugging Face로 LoRA 적용하기

LoRA란?

Low-Rank Adaptaion으로, PEFT 기법의 일종이다. downstream task 수행을 위해 모델 전체를 tuning하는 것이 아닌, 대부분의 파라미터를 freeze하고 일부 파라미터만 tuning하는 것이다.
Hugging Face는 peft라는 라이브러리를 통해 다양한 PEFT 기법을 적용할 수 있게 했는데, 그 중 LoRA를 적용하는 법에 대해 정리하고자 한다.
한 가지 유의할 점은, LoRA는 일부 파라미터를 tuning하여 다시 원래 모델에 얹는 격이기 때문에 학습 시 리소스를 절약해주기는 하지만, 추론 시에는 모델을 전부 사용하기 때문에 리소스 절약 효과가 없다. 이 부분에 대해서는 Quantized 개념이 적용된 QLoRA 기법을 사용해야 한다.

Hugging Face로 LoRA 적용하기

1. 사용할 모델과 tokenizer 불러오기

나는 STS task를 수행할 것이기 때문에 AutoModelForSequenceClassification 모듈을 사용해서 모델을 불러올 것이다. 또한 tokenizer는 사용할 모델에 맞게 같은 uid를 사용해서 불러오도록 한다. 이때 tokenizer는 이후 데이터 전처리에도 사용할 것이므로 잘 정의해두도록 한다.

from transformers import AutoTokenizer, AutoModelForSequenceClassification

model_uid = 'klue/roberta-small'

tokenizer = AutoTokenizer(model_uid)
model = AutoModelForSequenceClassificaion.from_pretrained(model_uid)

2. peft config 정의하기

peft 모듈의 LoraConfig와 TaskType 모듈을 사용해서 모델에 적용할 LoRA config를 정의한다. 이때 여러 매개변수가 있는데, 중요한 것 일부만 설명하고 넘어가겠다! 나머지는 Hugging Face API를 참고해보세요.

from peft import LoraConfig, TaskType

peft_config = LoraConfig(
    task_type=TaskType.SEQ_CLS,
    inference_mode=False,
    r=8,
    lora_alpha=16,
    lora_dropout=0.1
)

r
LoRA의 Low-Rank에 해당하는 rank를 결정하는 매개변수이다. rank란, 행렬에서 독립인 행 또는 열의 수인데, 지정한 r값의 크기로 줄인 가중치 행렬을 학습하는 것이 바로 LoRA이다.
따라서 이 값이 작을 수록 fine-tuning 시 리소스를 절약할 수 있지만 추론 시에는 전체 모델을 사용하기 때문에 리소스 절약이 없음을 유의한다.

lora_alpah
LoRA를 scaling하는 alpha 파라미터이다. 수식과 함께 이해해야 할 것 같은데 일반적으로 $(\text{r}, \alpha)$를 (16, 32) 또는 (8, 16)으로 지정한다고 한다.

lora_dropout
LoRA layer에 적용할 dropout 비율을 의미한다.

target_modules
adapter를 적용할 module을 String List로 전달한다. target_module이 될 수 있는 것들은 아래와 같다.
target_modules = ['q_proj','k_proj','v_proj','o_proj','gate_proj','down_proj','up_proj','lm_head']
all-linear라고 지정하면 모든 linear/Conv1D layer가 선택된다. 아무것도 전달하지 않으면 모델 구조에 맞춰서 적합한 모듈이 선택된다.
여기서 지정하는 값에 따라 성능이 좌우되기도 한다는 글을 봐서 각각 어떤 영향이 있는지 알아봐야 할 것 같다.

3. 데이터셋 준비하기

원본 데이터의 형식이 모두 다를 것이므로 특별히 코드를 작성하지는 않았다.
Hugging Face의 Trainer를 사용할 것이므로, input 데이터의 형식은 아래와 같은 dictionary로 준비한다.

input_vectors = {
    "input_ids": [],
    "tokey_type_ids": [],
    "attention_mask": [],
    "labels": []
}

이때 추론 시 사용할 labels가 없다면 labels는 빼도 된다. 나머지 값들은 1번에서 불러온 AutoTokenizer를 사용하면 쉽게 구할 수 있다.

4. Trainer 정의하기

import numpy as np
from scipy.stats import pearsonr
from transformers import TrainingArguments, Trainer

def compute_metrics(eval_pred):
    logits, labels = eval_pred
    
    predictions = logits.squeeze()
    labels = np.array(labels)

    pearson_corr, _ = pearsonr(predictions, labels)

    return {"pearson_corr": pearson_corr}

training_arguments = TrainingArguments(
    output_dir='./peft/',
    overwrite_output_dir=True,
    num_train_epochs=10,
    learning_rate=1e-3,
    weight_decay=0.01,
    eval_strategy="epoch",
    per_device_train_batch_size=16,
    per_device_eval_batch_size=16,
    lr_scheduler_type='linear',
)

trainer = Trainer(
    model=model,
    args=training_arguments,
    train_dataset=train_data,
    eval_dataset=valid_data,
    compute_metrics=compute_metrics,
)

필요한 매개변수는 Hugging Face API에 잘 나와있으므로 참고하고, 여기서 중요한 점은 Trainer에서 사용할 compute_metrics를 정의해야 한다는 것이다. 나는 STS를 수행하기 때문에 pearson correlation 수식을 사용했지만, 경우에 따라 정확도를 검증할 수식을 정의 및 로드해주면 된다.
optimizer는 자동으로 AdamW로 지정이 되고, loss 또한 불러온 모델에 적합한 loss를 자동으로 지정해준다. 추가로 설정하고 싶다면 optimizer와 scheduler를 정의 및 선언해주고 Trainer의 변수로 입력한다.

5. 학습 시작

trainer.train()

학습 시작! 자동으로 wandb에 연결해주므로 회원가입하고 로그인해서 API 키를 받아놓도록 한다.
로그인 > 우측 상단 프로필 클릭 > User settings > 하단의 Danger zone에서 Reveal 클릭하고 키를 복사해서 붙여넣기해주면 된다. 학습하며 보여지는 링크에 접속하면 모델 학습 현황, gpu 사용량 등을 파악할 수 있다.