从RLHF走到DPO：AI直接偏好优化笔记

# 早期尝试的奖励模型训练代码
def train_reward_model(train_data, val_data):
    model = AutoModelForSequenceClassification.from_pretrained(
        "gpt2", num_labels=1
    )

    # 训练了太多轮
    for epoch in range(50):  # 这里就错了
        train_loss = train_one_epoch(model, train_data)
        val_acc = evaluate(model, val_data)
        print(f"Epoch {epoch}: val_acc={val_acc}")

# PPO 配置简直是玄学
ppo_config = {
    "learning_rate": 1.41e-5,  # 这个数字是怎么来的？
    "batch_size": 128,         # 显存刚好够用
    "mini_batch_size": 16,     # 折中方案
    "gradient_accumulation_steps": 8,  # 再大显存就不够了
    "clip_range": 0.2,         # 听说大家都是0.2
    "kl_coef": 0.1,            # 这个要调吗？
    # ... 还有一堆参数
}

# 偏好数据格式
preference_data = [
    {
        "prompt": "用户：什么是量子计算？",
        "chosen": "量子计算利用量子力学原理进行计算，与传统计算机使用二进制位不同，量子计算机使用量子位（qubits）可以同时处于多种状态，从而在特定问题上实现指数级加速。",
        "rejected": "量子计算很复杂，它涉及薛定谔方程、量子隧穿、量子纠缠等等好多概念，基本上就是利用微观粒子的奇怪特性来计算..."
    },
    # ... 更多数据
]

from transformers import AutoModelForCausalLM, AutoTokenizer

# 使用经过 SFT 的模型作为起点
model_name = "local-sft-model"  # 已经过监督微调的模型
model = AutoModelForCausalLM.from_pretrained(
    model_name,
    torch_dtype=torch.bfloat16,
    device_map="auto"
)
tokenizer = AutoTokenizer.from_pretrained(model_name)

from trl import DPOTrainer, DPOConfig
from datasets import Dataset

# 转换数据格式
train_dataset = Dataset.from_list(preference_data)

# DPO 配置
dpo_config = DPOConfig(
    output_dir="./dpo_model",
    per_device_train_batch_size=4,
    gradient_accumulation_steps=8,
    learning_rate=5e-7,      # 比 RLHF 小很多
    logging_steps=10,
    save_steps=100,
    num_train_epochs=3,      # 通常 2-3 轮就够了
    beta=0.1,                 # KL 散度系数，控制偏离程度
    max_length=512,
    max_prompt_length=256,
)

# 初始化训练器
dpo_trainer = DPOTrainer(
    model=model,
    ref_model=None,          # DPO 会自动创建参考模型的副本
    args=dpo_config,
    train_dataset=train_dataset,
    tokenizer= tokenizer,
    beta=0.1,
)

# 开始训练
dpo_trainer.train()

# 训练日志示例
{
    "epoch": 1,
    "train_loss": 0.456,        # DPO 损失
    "rewards/chosen": 0.789,    # 选中回答的隐含奖励
    "rewards/rejected": -0.456, # 拒绝回答的隐含奖励
    "rewards/acc": 0.85,        # 偏好预测准确率
    "rewards/margins": 1.245,   # 奖励边际（越大越好）
}

# 使用梯度检查点
dpo_config = DPOConfig(
    # ... 其他参数
    gradient_checkpointing=True,  # 减少显存
    optim="adamw_torch_fused",     # 融合 AdamW
    bf16=True,                     # 使用 BF16
)

# 使用 DeepSpeed
dpo_config = DPOConfig(
    # ... 其他参数
    deepspeed="ds_config.json",    # DeepSpeed 配置
)

# 降低学习率
dpo_config = DPOConfig(
    learning_rate=1e-7,  # 从 5e-7 降低
    warmup_ratio=0.1,    # 添加 warmup
    lr_scheduler_type="cosine",  # 使用余弦调度
)

# 调整 beta
dpo_config = DPOConfig(
    beta=0.2,  # 从 0.1 增加，限制偏离程度
)

# 清洗数据
def clean_preference_data(data):
    cleaned = []
    for item in data:
        # 确保偏好差异明显
        chosen_score = evaluate_response_quality(item['chosen'])
        rejected_score = evaluate_response_quality(item['rejected'])

        if chosen_score - rejected_score > 0.5:
            cleaned.append(item)
    return cleaned

# 1. 减少训练轮数
dpo_config = DPOConfig(
    num_train_epochs=2,  # 从 3 减少到 2
)

# 2. 增加 beta 值
dpo_config = DPOConfig(
    beta=0.15,  # 从 0.1 增加到 0.15
)

# 3. 添加早停机制
from transformers import EarlyStoppingCallback

dpo_trainer = DPOTrainer(
    # ... 其他参数
    callbacks=[EarlyStoppingCallback(early_stopping_patience=3)]
)

# 1. 统一偏好数据的风格
def normalize_response_style(response):
    # 确保所有 chosen 回答风格一致
    return clean_and_normalize(response)

# 2. 使用更大的 beta 值
dpo_config = DPOConfig(
    beta=0.2,  # 限制风格漂移
)

# 3. 在训练时监控风格指标
def monitor_style_during_training(model, test_prompts):
    results = []
    for prompt in test_prompts:
        response = generate(model, prompt)
        avg_length = len(response.split())
        results.append(avg_length)
    return sum(results) / len(results)

# 成本计算（以 A100 80G 云服务器为例）

rlhf_cost = {
    "服务器时间": "7天 × 24小时",
    "人工标注": "200小时",
    "调试时间": "100小时",
    "总成本": "约15000元"
}

dpo_cost = {
    "服务器时间": "1.5天 × 24小时",
    "人工标注": "100小时",  # 数据需求量更少
    "调试时间": "20小时",   # 调试简单很多
    "总成本": "约3500元"
}

# 节省约 77% 的成本