from datasets import load_dataset
from peft import (
AutoPeftModelForCausalLM,
LoraConfig,
get_peft_model,
prepare_model_for_kbit_training,
)
from transformers import AutoTokenizer, BitsAndBytesConfig, logging
from trl import DPOConfig, DPOTrainer
# 데이터 전처리
def format_prompt(example):
"""TinyLLama가 사용하는 <|user|> 템플릿을 사용하여 프롬프트 포맷"""
system = "<|system|>\n" + example["system"] + "</s>\n"
prompt = "<|user|>\n" + example["input"] + "</s>\n<|assistant|>\n"
chosen = example["chosen"] + "</s>\n"
rejected = example["rejected"] + "</s>\n"
return {
"prompt": system + prompt,
"chosen": chosen,
"rejected": rejected,
}
# 데이터셋에 포맷 적용 및 비교적 짧은 답변 선택
dpo_dataset = load_dataset("argilla/distilabel-intel-orca-dpo-pairs", split="train")
dpo_dataset = dpo_dataset.filter(
lambda r: r["status"] != "tie" and r["chosen_score"] >= 8 and not r["in_gsm8k_train"]
)
dpo_dataset = dpo_dataset.map(format_prompt, remove_columns=dpo_dataset.column_names)
# 모델 양자화
# 4비트 양자화 설정 - QLoRA의 Q
bnb_config = BitsAndBytesConfig(
load_in_4bit=True, # 4비트 정밀도 모델 로딩 사용
bnb_4bit_quant_type="nf4", # 양자화 타입
bnb_4bit_compute_dtype="float16", # 계산 데이터 타입
bnb_4bit_use_double_quant=True, # 중첩 양자화 적용
)
# LoRA와 기본 모델 병합
model = AutoPeftModelForCausalLM.from_pretrained(
"./model/TinyLlama-1.1B-qlora",
low_cpu_mem_usage=True,
device_map="auto",
quantization_config=bnb_config,
)
merged_model = model.merge_and_unload()
# LLaMA 토크나이저 로드
model_name = "TinyLlama/TinyLlama-1.1B-intermediate-step-1431k-3T"
tokenizer = AutoTokenizer.from_pretrained(model_name, trust_remote_code=True)
tokenizer.pad_token = "<PAD>"
tokenizer.padding_side = "left"
# LoRA 설정 준비
peft_config = LoraConfig(
lora_alpha=32, # LoRA 스케일링
lora_dropout=0.1, # LoRA 레이어의 드롭아웃
r=64, # 랭크
bias="none",
task_type="CAUSAL_LM",
target_modules=[
# 대상 레이어
"k_proj",
"gate_proj",
"v_proj",
"up_proj",
"q_proj",
"o_proj",
"down_proj",
],
)
# 훈련을 위한 모델 준비
model = prepare_model_for_kbit_training(model)
model = get_peft_model(model, peft_config)
output_dir = "./model"
# 훈련 인자
training_arguments = DPOConfig(
output_dir=output_dir,
per_device_train_batch_size=2,
gradient_accumulation_steps=4,
optim="paged_adamw_32bit",
learning_rate=1e-5,
lr_scheduler_type="cosine",
max_steps=500,
logging_steps=100,
fp16=True,
gradient_checkpointing=True,
warmup_ratio=0.1,
beta=0.1, # beta 값을 여기에 추가
max_prompt_length=512,
max_length=512,
disable_tqdm=False,
)
# DPO 트레이너 생성
dpo_trainer = DPOTrainer(
model,
args=training_arguments,
train_dataset=dpo_dataset,
processing_class=tokenizer, # tokenizer 대신 processing_class 사용
peft_config=peft_config,
)
# DPO로 모델 미세조정
dpo_trainer.train()
# 어댑터 저장
dpo_trainer.model.save_pretrained("./model/TinyLlama-1.1B-dpo-qlora")