카테고리 없음

LoRA를 활용한 GPT-2 미세 조정 코드

aimler 2025. 3. 7. 17:58

아래 코드는 Hugging Face의 transformers, peft를 사용하여 GPT-2 모델을 LoRA로 미세 조정하는 예제입니다.

✅ 필요한 패키지 설치

먼저 필요한 라이브러리를 설치합니다.

pip install torch transformers datasets peft accelerate

 

📌 Python 코드: LoRA를 이용한 GPT-2 미세 조정

import torch
from transformers import AutoModelForCausalLM, AutoTokenizer, TrainingArguments, Trainer
from peft import LoraConfig, get_peft_model
from datasets import load_dataset

# 1️⃣ 모델 및 토크나이저 불러오기
model_name = "gpt2"
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModelForCausalLM.from_pretrained(model_name)

# 2️⃣ LoRA 설정
lora_config = LoraConfig(
    r=8,  # LoRA 랭크 (낮을수록 연산량 감소)
    lora_alpha=16,  # LoRA 스케일링 계수
    target_modules=["c_attn", "c_proj"],  # 적용할 모델의 특정 레이어
    lora_dropout=0.1,  # 드롭아웃 비율
    bias="none",
)

# 3️⃣ LoRA 적용
model = get_peft_model(model, lora_config)
model.print_trainable_parameters()  # 학습 가능한 파라미터 확인

# 4️⃣ 데이터셋 로드 및 전처리 (예제: Hugging Face "wikitext" 데이터)
dataset = load_dataset("wikitext", "wikitext-2-raw-v1", split="train")

def tokenize_function(examples):
    return tokenizer(examples["text"], padding="max_length", truncation=True)

tokenized_datasets = dataset.map(tokenize_function, batched=True)
tokenized_datasets = tokenized_datasets.shuffle().select(range(1000))  # 데이터 샘플링

# 5️⃣ 학습 설정
training_args = TrainingArguments(
    output_dir="./lora-gpt2-model",
    evaluation_strategy="epoch",
    save_strategy="epoch",
    learning_rate=5e-4,
    per_device_train_batch_size=8,
    per_device_eval_batch_size=8,
    num_train_epochs=3,
    weight_decay=0.01,
    logging_dir="./logs",
)

# 6️⃣ Trainer 객체 생성 및 학습 시작
trainer = Trainer(
    model=model,
    args=training_args,
    train_dataset=tokenized_datasets,
)

trainer.train()

# 7️⃣ LoRA 모델 저장
model.save_pretrained("./lora-gpt2-model")
tokenizer.save_pretrained("./lora-gpt2-model")

print("🎉 LoRA 학습 완료 및 저장 완료!")

 

🔹 코드 설명

1️⃣ GPT-2 모델과 토크나이저 로드
2️⃣ LoRA 설정

  • r=8: LoRA 랭크(낮을수록 메모리 사용량 감소)
  • lora_alpha=16: 학습할 가중치 크기 조정
  • target_modules=["c_attn", "c_proj"]: LoRA 적용할 Transformer 레이어 선택
    3️⃣ LoRA 적용 및 학습 가능한 파라미터 확인
    4️⃣ Hugging Face 데이터셋 로드 및 전처리
    5️⃣ 훈련 하이퍼파라미터 설정 (Trainer 사용)
    6️⃣ 모델 학습 진행 (trainer.train())
    7️⃣ 학습된 모델 저장

 

🔹 결과 및 활용

원본 GPT-2 모델을 수정하지 않고 LoRA를 활용해 미세 조정
메모리를 절약하면서 특정 도메인(예: 위키텍스트)에 맞춘 GPT-2 생성 모델 학습 가능
학습된 LoRA 모델을 저장하고 다시 불러와 사용 가능

🎯 이렇게 하면 LoRA를 사용하여 특정 데이터셋에 최적화된 GPT-2 모델을 만들 수 있습니다! 🚀