ORPO와 Unsloth를 이용한 미세 조정

2024년 9월 3일 · 태그: ai, python, llms · 원문 · 한국어 번역(원문에서 변경됨)

ORPO는 “ORPO: Monolithic Preference Optimization without Reference Model(ORPO: 참조 모델 없는 단일 선호도 최적화)” 논문에서 소개되었습니다. 핵심 혁신은 학습 과정에서 선호되지 않는 생성에 작은 페널티를 부여함으로써 별도의 선호도 정렬 단계가 필요 없게 만든다는 점입니다.

DPO와 같은 기존 접근 방식은 두 단계를 거쳐야 합니다:

베이스 모델에 대한 지도 미세 조정
쌍으로 구성된 데이터를 이용한 선호도 최적화

이 2단계 과정에는 중대한 단점이 있습니다. SFT 동안 모델은 선호되는 응답과 원치 않는 응답 모두를 생성할 확률을 높입니다. 이는 SFT가 더 나은 응답과 더 나쁜 응답을 구분하지 않고 모든 학습 예제를 똑같이 타당한 것으로 취급하기 때문입니다. 그 결과, 선호되는 출력과 거부된 출력 사이의 격차를 벌리기 위해 별도의 직접 선호도 정렬 단계가 필요합니다. ORPO를 사용하면 이 모든 것을 한 단계로 처리할 수 있으므로 훌륭할 뿐 아니라 더 효율적입니다.

ORPO에는 세 가지 핵심 구성 요소를 갖춘 선호도 데이터셋이 필요합니다:

프롬프트/지시문
선택된(선호되는) 응답
거부된(선호되지 않는) 응답

선호도 데이터셋은 다음과 같은 형태입니다:

{
    "prompt": "What is the meaning of life?",
    "chosen": "The meaning of life is to grow, learn, and find purpose through our experiences and connections with others.",
    "rejected": "42 lol"
}

효율적인 학습을 위해 새로운 Llama 3.2 8B 모델을 Unsloth와 함께 사용하겠습니다. 먼저 환경을 설정합시다:

# Install required dependencies
pip install flash-attn --no-build-isolation
pip install -v -U git+https://github.com/facebookresearch/xformers.git@main#egg=xformers
pip install "unsloth[colab-new] @ git+https://github.com/unslothai/unsloth.git"

이제 학습 스크립트를 설정합시다:

import torch
from datasets import load_dataset
from transformers import TextStreamer
from trl import ORPOTrainer, ORPOConfig

from unsloth import FastLanguageModel, is_bfloat16_supported
from unsloth.chat_templates import get_chat_template

# Model configuration
max_seq_length = 2048
model, tokenizer = FastLanguageModel.from_pretrained(
    model_name="unsloth/Meta-Llama-3-8B-bnb-4bit",
    max_seq_length=max_seq_length,
    load_in_4bit=True,
    dtype=None,
)

# Setup Unsloth's optimized LoRA
model = FastLanguageModel.get_peft_model(
    model,
    r=16,
    lora_alpha=16,
    lora_dropout=0,
    target_modules=[
        "q_proj", "k_proj", "v_proj",
        "o_proj", "gate_proj", "up_proj", "down_proj"
    ],
    use_rslora=True,  # Faster training
    use_gradient_checkpointing="unsloth",
)

# Setup ChatML format
tokenizer = get_chat_template(
    tokenizer,
    chat_template="chatml"
)

데이터셋으로는 비윤리적이고 유해한 응답을 피하도록 모델을 학습시키기 위한 이진화된 선호도가 담긴 ultrafeedback_binarized 데이터셋을 사용하겠습니다. trl을 사용하면 이미 선호도 형식을 이해하므로 데이터셋을 전처리할 필요가 없습니다.

dataset = load_dataset("trl-lib/ultrafeedback_binarized", split="train")

이제 Unsloth에 최적화된 설정으로 ORPO 트레이너를 구성합시다:

orpo_args = ORPOConfig(
    learning_rate=8e-6,
    beta=0.1,
    lr_scheduler_type="cosine_with_restarts",  # Better with Unsloth
    max_length=2048,
    max_prompt_length=1024,
    per_device_train_batch_size=4,  # Increased due to Unsloth optimizations
    per_device_eval_batch_size=4,
    gradient_accumulation_steps=2,
    optim="adamw_8bit",
    num_train_epochs=3,
    evaluation_strategy="steps",
    eval_steps=0.2,
    logging_steps=1,
    warmup_steps=10,
    output_dir="./results/",
    fp16=not is_bfloat16_supported(),
    bf16=is_bfloat16_supported(),
)

trainer = ORPOTrainer(
    model=model,
    args=orpo_args,
    train_dataset=dataset["train"],
    eval_dataset=dataset["test"],
    tokenizer=tokenizer,
)

trainer.train()

학습이 끝나면 모델을 다양한 형식으로 저장할 수 있습니다:

# Save LoRA weights
model.save_pretrained("lora_model")
tokenizer.save_pretrained("lora_model")

# Save merged model in 16-bit
print("Saving merged model...")
model.save_pretrained_merged("fused_model", tokenizer, save_method="merged_16bit")

# Save GGUF formats for inference
print("Saving GGUF formats...")
model.save_pretrained_gguf("ggufs", tokenizer, quantization_method="f16")
model.save_pretrained_gguf("ggufs", tokenizer, quantization_method="q4_k_m")

학습 중에 ORPO는 몇 가지 중요한 지표를 기록합니다:

rewards/chosen: 선택된 응답의 로그 확률
rewards/rejected: 거부된 응답의 로그 확률
rewards/accuracies: 선택된 보상이 거부된 보상을 초과하는 빈도
rewards/margins: 선택된 보상과 거부된 보상 간의 평균 차이
log_odds_ratio: 선택된 응답 대 거부된 응답의 평균 로그 오즈비
nll_loss: SFT 구성 요소에서 나온 음의 로그 우도 손실

Weights & Biases를 사용하여 이 지표들을 모니터링할 수 있습니다.

Keyboard shortcuts

Stephen Diehl 블로그 한국어 번역

ORPO와 Unsloth를 이용한 미세 조정