LoRA (Low-Rank Adaptation)에 대해 알아보자

 

최근 Large Language Model (LLM)이 다양한 태스크에서 뛰어난 성능을 보임에 따라, LLM을 원하는 태스크에 맞게 fine-tuning하여 사용하려는 니즈가 많다.

그런데 LLM은 이름에 나와 있듯이 굉장히 많은 파라미터 (10B은 기본..)를 가지고 있기 때문에, full fine-tuning을 하려면 많은 메모리와 시간이 든다.

그래서 사전학습된 LLM의 모든 파라미터를 튜닝하지 않고, 수행하고자 하는 특정 태스크를 위한 소량의 파라미터만을 추가적으로 학습하여 adaptation하는 lightweight fine-tuning이 주목받고 있다.

그 중에서 최근 가장 많이 사용되는 LoRA (Low-Rank Adaptation)에 대해 알아보려고 한다.

 

글 작성에 도움을 받은 참고 링크들이다.

• Paper
• Databricks' Post (KOR) → 설명이 잘 되어있고 번역도 잘 되어있어 읽어보기를 매우 추천한다 :)
• Databricks' Github

 

LoRA는 Low-rank factorization을 기반으로 한 기술이다.

Low-rank factorization은 무엇일까?

이는 사전학습된 모델의 가중치 행렬을 low-rank 행렬로 근사화하는 기법이다. 좀 더 자세히 살펴보자.

• factorization
  • decomposition과 동일한 '분해'의 뜻을 갖는다.
• rank
  • rank가 갖는 의미를 이해하려면 선형대수학을 떠올려야 한다. (핵심만 설명할 것이니 겁먹지 않아도 된다)
  • rank는 쉽게 말해 '행렬에서 독립인 행 또는 열의 수'이다.
  • 예시를 보면 이해될 것이다.
    아래의 A 행렬에 대해서, 열 벡터들의 독립/종속을 살펴보자.
    1열과 2열은 동일하므로 종속이고, 여기에 2배를 하면 3열이 되기 때문에 3열도 종속이다.
    단 4열은 이들과 독립이다.
    그러면 A 행렬에서 독립인 열의 수는 2이므로 rank(A)=2가 된다.
    (자세한 설명을 원한다면 이 블로그를 참고)

그러면 Low-rank factorization의 의미는 행렬을 작은 rank로 분해시킨다는 것을 알 수 있다.

LoRA paper의 다음 그림을 보면 명확히 와닿을 것이다.

왼쪽의 파란색 행렬이 사전학습된 가중치 행렬이고, 오른쪽의 주황색이 이 가중치 행렬을 Low-rank (=r)로 factorization한 것이다. 이제 명확히 이해되지 않는가?

 

LoRA는 사전학습된 가중치 행렬은 freezing하고, low-rank factorization한 두 행렬을 사용해 fine-tuning하는 기법이다.

LoRA의 효과는 trainable parameter의 수를 크게 줄여 학습 속도를 높이고 메모리 사용량을 낮추며, full fine-tuning과 유사하거나 더 높은 성능을 낼 수 있다는 점이다.

• trainable parameter의 수 $|Θ| = 2×\hat{L}_{LoRA}×d_{model}×r$​

추론 시에는 fine-tuning된 LoRA adapter를 pre-trained model (PTM)에 load하여 사용한다.

• 유의: 추론 시에는 PTM의 파라미터도 모두 불러와 사용하기 때문에, 메모리 사용량은 기존 PTM을 사용할 때와 비슷하다.

LoRA는 HuggingFace의 PEFT 라이브러리에 구현되어 있어 간편하게 사용할 수 있다.

 

최근에는 LoRA의 더 업그레이드된 버전인 QLoRA가 많이 사용되는 추세이다.

QLoRA는 한마디로 LoRA에 양자화(Quantization)를 결합한 기법으로, 이를 통해 메모리 효율성을 더욱 높였다.

이는 HuggingFace의 bitsandbytes 라이브러리를 통해 간편히 사용 가능하다.

 

 

이제 LLM에 QLoRA를 활용해 fine-tuning하는 코드를 살펴보자. 총 7단계이다.

 

1단계: 학습 데이터를 불러와 train/test 분할

from datasets import Dataset


dataset = Dataset.from_pandas(df).train_test_split(test_size=0.05, seed=42)

 

2단계: LoRA config 정의

from peft import LoraConfig


#If only targeting attention blocks of the model
target_modules = ["q_proj", "v_proj"]

#If targeting all linear layers
#target_modules = ['q_proj','k_proj','v_proj','o_proj','gate_proj','down_proj','up_proj','lm_head']

lora_config = LoraConfig(
    r=8, # or r=16
    target_modules=target_modules,
    lora_alpha=8,
    lora_dropout=0.05,
    bias="none",
    task_type="CAUSAL_LM",
)

• r
  • r이 커지면, 학습 시간과 학습 파라미터 수가 증가
    r이 작아지면, 학습 시간과 연산량은 줄어들지만 성능이 낮아질 수 있음
  • 그러나 r을 특정 값 이상으로 키워도 모델 성능이 향상되지 않을 수 있음
• target_modules
  • target_modules를 늘리면 학습 시간과 학습 파라미터 수가 증가
  • 그래서 관행적으로는 Attention blocks만 target_modules로 사용
  • 하지만 모든 linear layers를 target_modules 로 사용하면 성능이 향상되는 경향
• 모든 linear layers를 target_modules에 추가하는 코드 (더보기 클릭)
  • import re
    model_modules = str(model.modules)
    pattern = r'\((\w+)\): Linear'
    linear_layer_names = re.findall(pattern, model_modules)
    
    names = []
    # Print the names of the Linear layers
    for name in linear_layer_names:
        names.append(name)
    target_modules = list(set(names))

 

3단계: Training arguments 정의 (hyperparameter 설정)

from transformers import TrainingArguments


# per_device_train_batch_size = 4
# gradient_accumulation_steps = 4
# optim = 'adamw_hf'
# learning_rate = 1e-5
# max_grad_norm = 0.3
# warmup_ratio = 0.03
# lr_scheduler_type = "linear"

training_args = TrainingArguments(
    output_dir=base_dir,
    save_strategy="epoch",
    evaluation_strategy="epoch",
    num_train_epochs = 3.0,
    per_device_train_batch_size=per_device_train_batch_size,
    gradient_accumulation_steps=gradient_accumulation_steps,
    optim=optim,
    learning_rate=learning_rate,
    fp16=True,
    max_grad_norm=max_grad_norm,
    warmup_ratio=warmup_ratio,
    group_by_length=True,
    lr_scheduler_type=lr_scheduler_type,
)

 

4단계: 모델을 4-bit로 GPU 메모리에 load하고, PEFT 모델을 생성

from transformers import BitsAndBytesConfig
from transformers import LlamaTokenizer, LlamaForCausalLM # 사용하는 모델에 맞게 변경한다.
from peft import get_peft_model


nf4_config = BitsAndBytesConfig(
    load_in_4bit=True, # 학습 시에는 4bit로 load
    bnb_4bit_quant_type="nf4",
    bnb_4bit_use_double_quant=True,
    bnb_4bit_compute_dtype=torch.bfloat16
)

model = LlamaForCausalLM.from_pretrained(
    MODEL_PATH,
    quantization_config=nf4_config, # 학습 시에는 4bit로 load
    device_map='auto'
)

model = get_peft_model(model, lora_config) # lora_config: 2단계에서 정의함

 

5단계: SFTTrainer 인스턴스 생성 후 학습(fine-tuning) 진행

from trl import SFTTrainer


trainer = SFTTrainer(
    model,
    train_dataset=dataset['train'],
    eval_dataset = dataset['test'],
    dataset_text_field="text",
    max_seq_length=256,
    args=training_args,
)

trainer.train()

 

6단계: PTM을 새로 불러오고, PTM에 LoRA adapter의 weights를 load하여 PEFT 모델을 얻음

from peft import PeftModel, get_peft_model
from transformers import LlamaForCausalLM # 사용하는 모델에 맞게 변경한다.

model = LlamaForCausalLM.from_pretrained(
    MODEL_PATH,
    load_in_8bit=True, # 추론 시에는 8bit로 load
    device_map='auto'
)

peft_model = PeftModel.from_pretrained(model, ADAPTER_PATH) # ADAPTER_PATH는 학습 완료된 LoRA adapter의 경로

 

7단계: 추론하기

# load tokenizer
tokenizer = LlamaTokenizer.from_pretrained(MODEL_PATH)
tokenizer.add_special_tokens({'pad_token': '[PAD]'})

# 입력 prompt
# instruction 데이터를 사용해 fine-tuning한 것으로 가정하여, Alpaca의 prompt format을 사용한다.
prompt = """Below is an instruction that describes a task. Write a response that appropriately completes the request.

### Instruction:
Create a detailed description for the following product: Corelogic Smooth Mouse, belonging to category: Optical Mouse

### Response:"""


input_ids = tokenizer(prompt, return_tensors="pt").input_ids.to('cuda')

generation_output = peft_model.generate( # peft_model: 6단계에서 정의함
    input_ids=input_ids,
    max_new_tokens=256,
)
# generation_output은 입력 토큰들 + 새로 생성된 토큰들에 대한 id 리스트
# shape: (batch_size, max_seq_len)

# token id 리스트 -> 문자열로 변환
generated_text = tokenizer.decode(generation_output[0])
print(generated_text)

 

• 추가 정보
  • 4단계에서의 get_peft_model() 과 6단계에서의 PeftModel.from_pretrained()의 차이점
    • get_peft_model(): 학습이 가능한 PEFT 모델을 생성
    • PeftModel.from_pretrained(): 사전학습된 모델 및 학습 완료된 PEFT adapter를 함께 load하며, 이때 파라미터는 frozen되므로 학습 가능한 상태가 아님
  • 일반적으로는 PTM의 가중치와 adapter의 가중치를 병합하지 않지만, 병합하면 추론 시간을 조금 더 단축할 수 있다고 함
    • 병합하는 코드: merged_model = peft_model.merge_and_unload()