DPO(Direct Preference Optimization) 논문 심층 분석 — RLHF 없이 LLM 정렬하기

# RLHF (PPO) 학습 루프 의사 코드
import torch

# 4개 모델을 동시에 로드해야 함
policy_model = load_model("sft_checkpoint")
reference_model = load_model("sft_checkpoint")  # frozen
reward_model = load_model("reward_model_checkpoint")
value_model = load_model("value_head_checkpoint")

for batch in dataloader:
    prompts = batch["prompt"]

    # 1. 정책 모델로 응답 생성 (온라인 생성 필요)
    responses = policy_model.generate(prompts)

    # 2. 보상 모델로 점수 산출
    rewards = reward_model(prompts, responses)

    # 3. Value function으로 advantage 추정 (GAE)
    values = value_model(prompts, responses)
    advantages = compute_gae(rewards, values, gamma=0.99, lam=0.95)

    # 4. PPO 클리핑 목적 함수 최적화
    ratio = policy_model.log_prob(responses) - reference_model.log_prob(responses)
    clipped_ratio = torch.clamp(ratio, 1 - clip_eps, 1 + clip_eps)
    policy_loss = -torch.min(ratio * advantages, clipped_ratio * advantages).mean()

    # 5. KL 페널티 추가
    kl_penalty = beta * kl_divergence(policy_model, reference_model, responses)
    total_loss = policy_loss + kl_penalty

    total_loss.backward()
    optimizer.step()

# DPO 학습 루프 의사 코드
import torch
import torch.nn.functional as F

# 2개 모델만 필요
policy_model = load_model("sft_checkpoint")
reference_model = load_model("sft_checkpoint")  # frozen

for batch in dataloader:
    # 오프라인 데이터에서 직접 학습 (온라인 생성 불필요)
    prompts = batch["prompt"]
    chosen = batch["chosen"]
    rejected = batch["rejected"]

    # 1. 각 모델의 log probability 계산
    pi_chosen = policy_model.log_prob(chosen, prompts)
    pi_rejected = policy_model.log_prob(rejected, prompts)
    ref_chosen = reference_model.log_prob(chosen, prompts)
    ref_rejected = reference_model.log_prob(rejected, prompts)

    # 2. DPO 손실 함수 (단 3줄!)
    log_ratio_chosen = pi_chosen - ref_chosen
    log_ratio_rejected = pi_rejected - ref_rejected
    loss = -F.logsigmoid(beta * (log_ratio_chosen - log_ratio_rejected)).mean()

    loss.backward()
    optimizer.step()

import torch
import torch.nn.functional as F
from torch import nn


def dpo_loss(
    policy_chosen_logps: torch.Tensor,
    policy_rejected_logps: torch.Tensor,
    reference_chosen_logps: torch.Tensor,
    reference_rejected_logps: torch.Tensor,
    beta: float = 0.1,
    label_smoothing: float = 0.0,
) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
    """
    DPO 손실 함수 구현.

    Args:
        policy_chosen_logps: 정책 모델의 chosen 응답 log probability
        policy_rejected_logps: 정책 모델의 rejected 응답 log probability
        reference_chosen_logps: 참조 모델의 chosen 응답 log probability
        reference_rejected_logps: 참조 모델의 rejected 응답 log probability
        beta: KL 제약 강도 (temperature)
        label_smoothing: 레이블 스무딩 계수 (0이면 표준 DPO)

    Returns:
        loss: DPO 손실값
        chosen_rewards: chosen 응답의 암묵적 보상
        rejected_rewards: rejected 응답의 암묵적 보상
    """
    # 암묵적 보상 계산: r(x,y) = beta * log(pi/pi_ref)
    chosen_rewards = beta * (policy_chosen_logps - reference_chosen_logps)
    rejected_rewards = beta * (policy_rejected_logps - reference_rejected_logps)

    # 보상 마진 계산
    reward_margin = chosen_rewards - rejected_rewards

    # 레이블 스무딩 적용 (선택적)
    if label_smoothing > 0:
        loss = (
            -F.logsigmoid(reward_margin) * (1 - label_smoothing)
            - F.logsigmoid(-reward_margin) * label_smoothing
        )
    else:
        loss = -F.logsigmoid(reward_margin)

    return loss.mean(), chosen_rewards.mean(), rejected_rewards.mean()


# 사용 예시
batch_size = 4
seq_len = 128

# 모의 log probability 값
policy_chosen_logps = torch.randn(batch_size) * 0.1 - 2.0
policy_rejected_logps = torch.randn(batch_size) * 0.1 - 2.5
reference_chosen_logps = torch.randn(batch_size) * 0.1 - 2.0
reference_rejected_logps = torch.randn(batch_size) * 0.1 - 2.5

loss, chosen_r, rejected_r = dpo_loss(
    policy_chosen_logps,
    policy_rejected_logps,
    reference_chosen_logps,
    reference_rejected_logps,
    beta=0.1,
)
print(f"DPO Loss: {loss.item():.4f}")
print(f"Chosen Reward: {chosen_r.item():.4f}")
print(f"Rejected Reward: {rejected_r.item():.4f}")

import torch
import torch.nn.functional as F


def compute_log_probs(
    logits: torch.Tensor,
    labels: torch.Tensor,
    attention_mask: torch.Tensor,
) -> torch.Tensor:
    """
    토큰 단위 log probability를 계산하고 시퀀스 단위로 합산.

    Args:
        logits: 모델 출력 로짓 (batch_size, seq_len, vocab_size)
        labels: 정답 토큰 ID (batch_size, seq_len)
        attention_mask: 어텐션 마스크 (batch_size, seq_len)

    Returns:
        시퀀스 단위 log probability (batch_size,)
    """
    # 로짓을 1 토큰 앞으로 이동 (next token prediction)
    shift_logits = logits[:, :-1, :]
    shift_labels = labels[:, 1:]
    shift_mask = attention_mask[:, 1:]

    # 토큰 단위 log probability
    log_probs = F.log_softmax(shift_logits, dim=-1)
    per_token_logps = torch.gather(
        log_probs, dim=2, index=shift_labels.unsqueeze(2)
    ).squeeze(2)

    # 마스킹된 토큰은 0으로 처리하고 시퀀스 합산
    per_token_logps = per_token_logps * shift_mask
    return per_token_logps.sum(dim=-1)

from datasets import load_dataset
from transformers import AutoModelForCausalLM, AutoTokenizer
from trl import DPOConfig, DPOTrainer
from peft import LoraConfig

# 1. 모델과 토크나이저 로드
model_name = "Qwen/Qwen2.5-7B-Instruct"
model = AutoModelForCausalLM.from_pretrained(
    model_name,
    torch_dtype="bfloat16",
    attn_implementation="flash_attention_2",
)
tokenizer = AutoTokenizer.from_pretrained(model_name)
if tokenizer.pad_token is None:
    tokenizer.pad_token = tokenizer.eos_token

# 2. LoRA 설정 (메모리 절약)
peft_config = LoraConfig(
    r=16,
    lora_alpha=32,
    lora_dropout=0.05,
    target_modules=["q_proj", "k_proj", "v_proj", "o_proj"],
    task_type="CAUSAL_LM",
)

# 3. 선호 데이터 로드 (prompt, chosen, rejected 구조)
dataset = load_dataset("trl-lib/ultrafeedback_binarized", split="train")

# 4. DPO 학습 설정
training_args = DPOConfig(
    output_dir="./dpo-qwen2.5-7b",
    per_device_train_batch_size=2,
    gradient_accumulation_steps=8,
    learning_rate=5e-6,         # SFT 대비 10~100배 낮게
    beta=0.1,                   # KL 제약 강도
    max_length=1024,
    max_prompt_length=512,
    num_train_epochs=1,
    bf16=True,
    logging_steps=10,
    save_strategy="steps",
    save_steps=500,
    warmup_ratio=0.1,
    gradient_checkpointing=True,
    remove_unused_columns=False,
)

# 5. DPO Trainer 생성 및 학습
trainer = DPOTrainer(
    model=model,
    args=training_args,
    train_dataset=dataset,
    processing_class=tokenizer,
    peft_config=peft_config,
)
trainer.train()
trainer.save_model("./dpo-qwen2.5-7b-final")

# 필수 패키지 설치
pip install trl>=0.12.0 transformers>=4.46.0 peft>=0.13.0 \
    datasets accelerate bitsandbytes flash-attn

# 멀티 GPU 학습 실행 (DeepSpeed ZeRO-3)
accelerate launch --config_file deepspeed_zero3.yaml \
    --num_processes 4 \
    train_dpo.py

# 학습 모니터링
tensorboard --logdir ./dpo-qwen2.5-7b/runs

# 합성 선호 데이터 생성 파이프라인 예시
from openai import OpenAI

client = OpenAI()

def generate_preference_pair(prompt: str, model_response_a: str, model_response_b: str) -> dict:
    """GPT-4를 사용하여 선호 판정을 수행"""
    judge_prompt = f"""다음 두 응답을 비교하여 어느 것이 더 도움이 되고,
정확하며, 안전한지 판단하세요.

프롬프트: {prompt}

응답 A: {model_response_a}

응답 B: {model_response_b}

더 나은 응답의 번호만 출력하세요 (A 또는 B):"""

    result = client.chat.completions.create(
        model="gpt-4",
        messages=[{"role": "user", "content": judge_prompt}],
        max_tokens=1,
    )
    choice = result.choices[0].message.content.strip()

    if choice == "A":
        return {"prompt": prompt, "chosen": model_response_a, "rejected": model_response_b}
    else:
        return {"prompt": prompt, "chosen": model_response_b, "rejected": model_response_a}

# DPO 학습 중 모니터링해야 할 핵심 지표
def log_dpo_metrics(chosen_rewards, rejected_rewards, loss):
    """학습 중 핵심 지표를 로깅하는 함수"""
    reward_margin = (chosen_rewards - rejected_rewards).mean()
    accuracy = (chosen_rewards > rejected_rewards).float().mean()

    metrics = {
        "dpo/loss": loss.item(),
        "dpo/reward_margin": reward_margin.item(),
        "dpo/accuracy": accuracy.item(),
        "dpo/chosen_reward_mean": chosen_rewards.mean().item(),
        "dpo/rejected_reward_mean": rejected_rewards.mean().item(),
    }
    # reward_margin이 지속적으로 증가하면 과적합 징후
    # accuracy가 1.0에 도달하면 학습 중단 고려
    return metrics