8️⃣TRL: ORPO Fine-tuning with Llama3-8B

ORPO는 기존의 감독 미세 조정과 선호도 정렬 단계를 단일 프로세스로 결합한 새롭고 흥미로운 미세 조정 기법입니다. 따라서 트레이닝에 필요한 컴퓨팅 리소스와 시간이 줄어듭니다. 또한, 경험적 결과에 따르면 ORPO는 다양한 모델 크기와 벤치마크에서 다른 얼라인먼트 방법보다 성능이 뛰어납니다.

명령어 튜닝 및 기본 설정 조정은 특정 작업에 LLM(대규모 언어 모델)을 적용하는 데 필수적인 기술입니다. 전통적으로 이 작업에는 여러 단계의 프로세스가 포함됩니다:

1. 목표 도메인에 맞게 모델을 조정하기 위한 명령어에 대한 Supervised Fine-Tuning(SFT),

2. 거부된 응답보다 선호되는 응답을 생성할 가능성을 높이기 위한 Reinforcement Learning with Human Feedback (RLHF) 또는 Direct Preference Optimization (DPO)와 같은 선호도 조정 방법.

SFT는 모델을 원하는 영역에 효과적으로 적용하지만, 실수로 선호되는 답변과 함께 바람직하지 않은 답변이 생성될 확률이 높아집니다. 그렇기 때문에 선호도 정렬 단계가 선호되는 결과와 거부되는 결과의 가능성 간격을 넓히기 위해 필요합니다.

Hong et al.2024가 소개한 ORPO(Optimization without Reference Model)는 인스트럭션 튜닝과 선호도 조정을 하나의 단일화된 훈련 프로세스로 결합하여 이 문제에 대한 우아한 해결책을 제시합니다. ORPO는 표준 언어 모델링 목표를 수정하여 음의 로그 확률 손실을 확률 비율(OR) 용어와 결합합니다. 이 OR 손실은 거부된 응답에는 약하게 불이익을 주고 선호하는 응답에는 강하게 보상함으로써 모델이 목표 작업을 학습하는 동시에 사람의 선호도에 맞게 조정할 수 있도록 합니다.

$L_{ORPO} = \mathbb{E}_{(x, y_{w}, y_{l})} [L_{SFT} + \lambda \cdot L_{OR}]$

RPO는 TRL, Axolotl, LLaMA-Factory와 같은 Python 라이브러리에서 구현할 수 있습니다. ORPO 알고리즘으로 SFT와 환경설정 정렬 단계를 단일 프로세스로 통합하는 방법으로 예제를 실행하겠습니다. TRL을 사용해 사용자 지정 환경설정 데이터 세트에서 Llama3 8B 모델을 Fine-tuning 합니다.

---

Fine-tune ORPO: Llama3-8B

ORPO에는 prompt/chosen answer/rejected answer을 포함한 기본 설정 데이터 세트가 필요합니다. Dataset은 고품질 DPO 데이터 세트의 조합인 mlabonne/orpo-dpo-mix-40k를 사용하겠습니다.

Setup Environment

%pip install -U transformers datasets accelerate peft trl bitsandbytes wandb

import gc
import os

import torch
import wandb
from datasets import load_dataset
from peft import LoraConfig, PeftModel, prepare_model_for_kbit_training
from transformers import (
    AutoModelForCausalLM,
    AutoTokenizer,
    BitsAndBytesConfig,
    TrainingArguments,
    pipeline,
)
from trl import ORPOConfig, ORPOTrainer, setup_chat_format

hf_token = "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"
wb_token = "b6ff13b6d78bff6512f1112beb3832758ff2f09a"
wandb.login(key=wb_token)

wandb: Currently logged in as: kubwai. Use `wandb login --relogin` to force relogin
wandb: WARNING If you're specifying your api key in code, ensure this code is not shared publicly.
wandb: WARNING Consider setting the WANDB_API_KEY environment variable, or running `wandb login` from the command line.
wandb: Appending key for api.wandb.ai to your netrc file: /home/kubwa/.netrc

True

Model, Tokenizer & LoRa Config

bitsandbytes를 사용하여 Llama3 8B 모델을 4비트 정밀도로 로드합니다. 그런 다음 QLoRA용 PEFT를 사용하여 LoRA 구성을 설정합니다. 또한 setup_chat_format() 함수를 사용하여 ChatML 지원을 위해 model과 tokenizer를 수정합니다. 이 함수는 이 chat template을 자동으로 적용하고, 특수 토큰을 추가하고, 새 어휘 크기에 맞게 모델의 임베딩 레이어 크기를 조정합니다.

# Model
base_model = "meta-llama/Meta-Llama-3-8B"
new_model = "Llama-3-8B-loudai-orpo"

# QLoRA config
bnb_config = BitsAndBytesConfig(
    load_in_4bit=True,
    bnb_4bit_quant_type="nf4",
    bnb_4bit_compute_dtype=torch.float16,
    bnb_4bit_use_double_quant=True,
)

# LoRA config
peft_config = LoraConfig(
    r=16,
    lora_alpha=32,
    lora_dropout=0.05,
    bias="none",
    task_type="CAUSAL_LM",
    target_modules=['up_proj', 'down_proj', 'gate_proj', 'k_proj', 'q_proj', 'v_proj', 'o_proj']
)

# Load tokenizer
tokenizer = AutoTokenizer.from_pretrained(base_model)

# Load model
model = AutoModelForCausalLM.from_pretrained(
    base_model,
    quantization_config=bnb_config,
    device_map={'':torch.cuda.current_device()}
)
model, tokenizer = setup_chat_format(model, tokenizer)
model = prepare_model_for_kbit_training(model)

Special tokens have been added in the vocabulary, make sure the associated word embeddings are fine-tuned or trained.

Loading checkpoint shards:   0%|          | 0/4 [00:00<?, ?it/s]

Dataset

mlabonne/orpo-dpo-mix-40k를 로드하고 apply_chat_template() 함수를 사용하여 "chosen" 및 "rejected" 열을 ChatML 형식으로 변환합니다. 실행 시간이 너무 오래 걸리므로 전체 데이터 세트가 아닌 1,000개의 샘플만 사용하겠습니다.

dataset_name = "mlabonne/orpo-dpo-mix-40k"
dataset = load_dataset(
    dataset_name, 
    split="all"
)
dataset = dataset.shuffle(seed=42).select(range(100))

def format_chat_template(row):
    row["chosen"] = tokenizer.apply_chat_template(
        row["chosen"], 
        tokenize=False
    )
    row["rejected"] = tokenizer.apply_chat_template(
        row["rejected"], 
        tokenize=False
    )
    return row

dataset = dataset.map(
    format_chat_template,
    num_proc= os.cpu_count(),
)
dataset = dataset.train_test_split(test_size=0.01)

Training

orpo_args = ORPOConfig(
    learning_rate=8e-6,
    beta=0.1,
    lr_scheduler_type="linear",
    max_length=1024,
    max_prompt_length=512,
    per_device_train_batch_size=2,
    per_device_eval_batch_size=2,
    gradient_accumulation_steps=4,
    optim="paged_adamw_8bit",
    num_train_epochs=1,
    evaluation_strategy="steps",
    eval_steps=0.2,
    logging_steps=1,
    warmup_steps=10,
    report_to="wandb",
    output_dir="./results/",
)

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

Map:   0%|          | 0/99 [00:00<?, ? examples/s]
Map:   0%|          | 0/1 [00:00<?, ? examples/s]

ㅈwwwand

Merge & Push Model

마지막으로 기존 모델과 fine-tune 모델을 merge하고 Huggingface hub에 push를 합니다.

# Flush memory
del trainer, model
gc.collect()
torch.cuda.empty_cache()

# Reload tokenizer and model
tokenizer = AutoTokenizer.from_pretrained(base_model)
model = AutoModelForCausalLM.from_pretrained(
    base_model,
    low_cpu_mem_usage=True,
    return_dict=True,
    torch_dtype=torch.float16,
    device_map="auto",
)
model, tokenizer = setup_chat_format(model, tokenizer)

# Merge adapter with base model
model = PeftModel.from_pretrained(model, new_model)
model = model.merge_and_unload()

model.push_to_hub(new_model, use_temp_dir=False)
tokenizer.push_to_hub(new_model, use_temp_dir=False)