SFT和RLHF详解

要深入理解 SFT (Supervised Fine-Tuning) 的内部实现，我们不能只看调包（Trainer），得从数学原理、数据流向和权重更新这三个底层维度来拆解。

SFT 的本质是条件概率的极大似然估计。它把一个“只会接话”的预训练模型，规训成一个“听得懂指令”的助手。

---

1. 数学核心：负对数似然 (NLL Loss)

预训练阶段，模型学习的是 $P(next\_token|context)$。
而 SFT 阶段，模型学习的是 $P(response|prompt)$。

假设输入指令是 $x$（Prompt），标准答案是 $y$（Response），$y$ 由一个个 token $y_1,y_2,...,y_n$ 组成。SFT 的目标是最小化以下损失函数：

$${\mathcal{L}}_{SFT}(\theta )=-\sum _{t=1}^n\log P_{\theta }(y_t|y_{<t},x)$$

关键细节： 在计算 Loss 时，我们只计算回答（Response）部分的损失。虽然 Prompt 也输入到了模型中，但模型预测 Prompt 预测得准不准，我们并不关心，也不计入梯度更新。

---

2. 数据流内部实现：Label Masking (标签掩码)

在代码实现中，这是通过 labels 矩阵完成的。这是 SFT 最核心的工程实现逻辑：

1. 拼接 (Concatenation)： 将 Prompt 和 Response 拼在一起，变成一串 Token IDs。
   • 例如：[SOS] 怎么做红烧肉？ [SEP] 第一步选五花肉... [EOS]
2. 构建 Labels： * 对应 Prompt 的位置，全部填上 -100。
   • 对应 Response 的位置，填上真实的 Token IDs。
3. 计算 Loss： PyTorch 的 CrossEntropyLoss 默认会忽略掉索引为 -100 的目标。
   • 这样，反向传播（Backpropagation）时，梯度只会从 Response 部分产生，强制模型学习“在看到这个 Prompt 后，如何生成这样的 Response”。

---

3. 权重更新：LoRA 的介入

既然你在用 4060，你跑的其实是 PEFT (Parameter-Efficient Fine-Tuning) 下的 SFT。

内部矩阵运算：

在全量微调中，我们更新的是原始权重 $W$。在 LoRA 中，权重保持不变（冻结），我们学习的是两个低秩矩阵 $A$ 和 $B$：
$$W_{updated}=W+\Delta W=W+(B\times A)$$

• 前向传播： $h=Wx+BAx$。
• 反向传播： 只计算针对 $A$ 和 $B$ 的梯度。因为 $A$ 和 $B$ 的参数量极小（通常只有原模型的 0.1%），所以你的 12G 显存才跑得动。

---

4. 训练全流程拆解

当我们执行 trainer.train() 时，底层发生了这些事：

Step 1: Forward Pass (前向传播)

模型读入拼接后的序列。由于是 Causal LLM（因果语言模型），每个 token 只能看到它左边的内容。

• 输入 Prompt 时，模型内部的 Attention Mask 确保它能看到问题。
• 生成到 Response 区域时，它根据 Prompt 的语义特征来预测下一个词。

Step 2: Calculate Loss (计算损失)

模型输出一个概率分布（Logits）。我们将这个分布与之前准备好的 labels（屏蔽了 Prompt 的标签）进行对比。

• 如果模型在 Response 位置预测的词与标准答案不符，Loss 就会很高。

Step 3: Backward Pass (反向传播)

计算 Loss 对 LoRA 参数 $A$ 和 $B$ 的偏导数。

Step 4: Optimizer Update (优化器更新)

使用 AdamW 等优化器更新 $A$ 和 $B$。

• 注意： 因为用了 4-bit 量化，权重在内存里是压缩状态，计算梯度时会临时反量化回 bf16 或 fp16。

---

5. 为什么 SFT 会导致“过拟合”或“复读机”？

从内部实现看，如果你让模型在同样的数据上跑太多次（Epochs 过多）：

1. 概率收缩： 模型对标准答案中每个 token 的预测概率会无限接近 100%。
2. 多样性消失： 模型学到了“死记硬背”，而不是理解逻辑。一旦预测概率过于集中，模型在推理时就会反复生成训练集里的原话，甚至陷入死循环。

---

总结

SFT 的内部逻辑 = 拼接数据 + 屏蔽 Prompt 标签 + 预测 Response Token + 极大似然优化。

它就像是一个老师拿着参考答案，遮住题目，只在学生答错答案的部分敲手心。

以下为可直接复制运行的代码（已测试）

"""
SFT (Supervised Fine-Tuning) 监督微调示例脚本 - 优化版
===========================================================
本脚本演示如何通过监督微调让模型学会遵循特定指令和表达风格。

SFT 核心思想：
- 使用 "指令-回答" 对进行模仿学习
- 建立模型的基础对话能力和格式规范
- 是所有对齐技术（如 DPO/RLHF）的基石

硬件要求：
- 显卡：NVIDIA RTX 4060 12GB
- 显存：约 7-9GB（使用 4-bit 量化）

优化要点：
1. 增加数据量（20条）
2. 提高LoRA秩（r=32）
3. 增加训练轮数（5轮）
4. 调整学习率（1e-4）
5. 添加warmup_steps
"""

import os
import torch
from datasets import Dataset
from modelscope import snapshot_download
from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig
from trl.trainer.sft_trainer import SFTTrainer
from trl.trainer.sft_config import SFTConfig
from peft import LoraConfig, PeftModel

# ================= 1. 环境与模型准备 =================
print("正在从 ModelScope 下载模型...")
# 使用 Qwen2.5-1.5B-Instruct 作为实验基座
model_dir = snapshot_download('qwen/Qwen2.5-1.5B-Instruct')

# 4-bit 量化配置，完美适配 12G 显存
bnb_config = BitsAndBytesConfig(
    load_in_4bit=True,
    bnb_4bit_quant_type="nf4",
    bnb_4bit_compute_dtype=torch.bfloat16,
)

tokenizer = AutoTokenizer.from_pretrained(model_dir)
tokenizer.pad_token = tokenizer.eos_token

model = AutoModelForCausalLM.from_pretrained(
    model_dir,
    quantization_config=bnb_config,
    device_map="auto",
    trust_remote_code=True
)
model.config.use_cache = False

# ================= 1.5 测试：查看微调前的模型输出 =================
test_prompt = "今天几号"
print(f"\n--- [微调前：标准模型回复] ---\n提问：{test_prompt}")

messages = [{"role": "user", "content": test_prompt}]
text = tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
inputs = tokenizer([text], return_tensors="pt").to(model.device)

with torch.no_grad():
    outputs = model.generate(**inputs, max_new_tokens=100)
    print("回答：", tokenizer.decode(outputs[0], skip_special_tokens=True).split("assistant\n")[-1])

# ================= 2. 准备 SFT 数据集 =================
"""
SFT 数据集格式：
- 指令遵循（Instruction Following）：问什么，答什么。
- 只有正确答案：不需要像 DPO 那样提供错误选项，只需提供"标准范本"。
- 数据量：建议至少20-50条，覆盖不同类型的问题
"""

print("\n正在构造本地 SFT 演示数据集...")

# 构造示例：让模型学习一种更简洁、更口语化的服务风格
# 增加到20条数据，覆盖不同类型的问题
sft_data = {
    "prompt": [
        # 自我介绍类
        "你好，请问你是谁？",
        "你叫什么名字？",
        "介绍一下你自己",
        
        # 日常对话类
        "今天几号？",
        "现在几点了？",
        "今天天气怎么样？",
        
        # 学习建议类
        "我想学习 Python，怎么入门？",
        "怎么学好英语？",
        "推荐一些好书",
        
        # 创作类
        "帮我写一段关于大海的描述。",
        "写一首关于秋天的诗",
        "描述一下你最喜欢的食物",
        
        # 能力询问类
        "你会做什么？",
        "你能帮我做什么？",
        "你有什么特长？",
        
        # 建议类
        "周末有什么安排建议？",
        "推荐一部好看的电影",
        "怎么提高工作效率？",
        
        # 闲聊类
        "你吃饭了吗？",
        "最近有什么新闻？",
        "你喜欢什么颜色？"
    ],
    "completion": [
        # 自我介绍类 - 口语化、亲切
        "嘿！我是你的全能助手，叫我小灵就行，有事儿您说话。",
        "叫我小灵就好啦，我是你的AI小伙伴，随时待命！",
        "我是小灵，一个喜欢聊天的AI助手，有啥想问的尽管说！",
        
        # 日常对话类 - 幽默、接地气
        "今天日期我建议你瞄一眼手机右上角，毕竟实时信息我偶尔会犯糊涂。",
        "我的表可能不准，你最好自己看一眼手机哈哈！",
        "天气这事儿你得看天气预报，我这边没联网，猜不准~",
        
        # 学习建议类 - 实用、直接
        "先装个 Python 环境，然后去 B 站搜个基础教程，多敲代码少看书，准没错。",
        "背单词+看美剧+找老外聊天，三管齐下，比死记硬背强多了。",
        "《三体》必须看，科幻迷圣经！另外《人类简史》也不错，开阔眼界。",
        
        # 创作类 - 生动、有画面感
        "深蓝的海水拍打着礁石，海风带着咸味，远处海天一色，真的挺治愈的。",
        "秋风起，落叶黄，一杯热茶暖心房。岁月静好心自安，且将诗意藏。",
        "火锅！麻辣鲜香，毛肚七上八下，蘸上香油蒜泥，绝了！",
        
        # 能力询问类 - 自信、幽默
        "聊天、写代码、出点子我都在行，只要别让我去线下帮你搬砖就行。",
        "写文案、答问题、陪聊天，样样精通，24小时不打烊！",
        "我特长就是话多，而且说得还挺有道理，不信你多问问~",
        
        # 建议类 - 贴心、实用
        "周末去公园走走，或者约朋友吃顿好的，别老宅在家里！",
        "《肖申克的救赎》经典中的经典，每次看都有新感悟。",
        "番茄工作法试试，25分钟专注+5分钟休息，效率翻倍！",
        
        # 闲聊类 - 随意、自然
        "AI不用吃饭，但你要记得按时吃哦，别饿坏了！",
        "我这没联网，新闻得你自己去看，不过世界应该还是老样子~",
        "我喜欢蓝色，像天空和大海一样，看着就舒服！"
    ]
}

# 转换为标准指令格式数据集
dataset = Dataset.from_dict(sft_data)
print(f"SFT 数据集构造完成！共 {len(dataset)} 条数据")

# ================= 3. 配置 LoRA 和 SFT =================
"""
优化后的LoRA配置：
- r=32：增加秩，提高模型表达能力
- lora_alpha=64：相应增加缩放因子
- lora_dropout=0.05：降低dropout，减少正则化强度
"""

peft_config = LoraConfig(
    r=32,                         # 增加秩到32，提高模型学习能力
    lora_alpha=64,                # 相应增加alpha，保持比例2:1
    target_modules=["q_proj", "v_proj", "k_proj", "o_proj", "gate_proj", "up_proj", "down_proj"],
    lora_dropout=0.05,            # 降低dropout，让模型更好地学习
    bias="none",
    task_type="CAUSAL_LM",
)

# 优化后的SFT训练配置
sft_config = SFTConfig(
    output_dir="./sft_qwen_results",
    per_device_train_batch_size=1,    # 减小batch size，更稳定
    gradient_accumulation_steps=8,    # 增加梯度累积，等效batch size=8
    learning_rate=1e-4,               # 降低学习率，训练更稳定
    num_train_epochs=15,               # 增加训练轮数到5轮
    warmup_steps=20,                  # 添加预热步数，稳定训练初期
    lr_scheduler_type="cosine",       # 余弦退火学习率
    logging_steps=5,                  # 更频繁的日志记录
    save_steps=50,                    # 更频繁的保存
    bf16=True,
    gradient_checkpointing=True,      # 显存保命开关
    dataset_text_field="completion",  # 指定回答字段
    max_length=512,
    neftune_noise_alpha=None,        # 可尝试启用NEFTune技术提升效果
)

# 初始化 SFT 训练器
trainer = SFTTrainer(
    model=model,
    train_dataset=dataset,
    args=sft_config,
    peft_config=peft_config,
    processing_class=tokenizer,      # 适配最新版本 trl
)

# ================= 4. 执行训练 =================
print("\n开始 SFT 训练...")
print(f"训练参数：")
print(f"  - 数据量：{len(dataset)} 条")
print(f"  - 训练轮数：{sft_config.num_train_epochs}")
print(f"  - LoRA秩：{peft_config.r}")
print(f"  - 学习率：{sft_config.learning_rate}")
print(f"  - 有效batch size：{sft_config.per_device_train_batch_size * sft_config.gradient_accumulation_steps}")

trainer.train()

# 保存适配器
trainer.save_model("./sft_qwen_final")
print("\n训练完成！模型已保存至 ./sft_qwen_final")

# ================= 5. 测试微调后的模型输出 =================
print("\n--- [加载微调后的 SFT 权重...] ---")
model = PeftModel.from_pretrained(model, "./sft_qwen_final")

# 测试多个问题，验证训练效果
test_prompts = [
    "今天几号",           # 训练集中的问题
    "你好，请问你是谁？",  # 训练集中的问题
    "推荐一部好看的电影",  # 训练集中的问题
    "你会唱歌吗？"        # 训练集外的问题（测试泛化能力）
]

print(f"\n--- [微调后：新风格模型回复] ---")
for prompt in test_prompts:
    print(f"\n提问：{prompt}")
    messages = [{"role": "user", "content": prompt}]
    text = tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
    inputs = tokenizer([text], return_tensors="pt").to(model.device)
    
    with torch.no_grad():
        outputs = model.generate(**inputs, max_new_tokens=100)
        response = tokenizer.decode(outputs[0], skip_special_tokens=True).split("assistant\n")[-1]
        print(f"回答：{response}")



"""
SFT 进阶说明：
==================
1. 数据质量：SFT 是"垃圾进，垃圾出"。如果你的 completion 质量不高，模型会学坏。
2. 指令多样性：如果 prompt 太单一，模型会失去泛化能力（只会回答那几个固定问题）。
3. 学习率：如果模型回答开始胡言乱语，尝试降低学习率（如 5e-5）。
4. 训练轮数：个人风格模仿通常 3-5 epoch 即可。
5. NEFTune技术：可尝试启用neftune_noise_alpha参数，有概率提升效果
"""

---

RLHF微调之DPO

---

1. 为什么大家都在逃离 PPO？

PPO (Proximal Policy Optimization) 是 OpenAI 在 InstructGPT 和 ChatGPT 早期使用的标准强化学习算法。它的理论非常完美，但工程实现简直是灾难。

痛点一：极度吃硬件（显存杀手）

如果你要跑 PPO，你的显存里必须同时塞进 4 个模型：

1. Actor Model (策略模型): 正在训练的主模型。
2. Reference Model (参考模型): 冻结的 SFT 模型，用来防止 Actor 跑偏。
3. Reward Model (奖励模型): 专门训练出来给 Actor 的回答打分的模型。
4. Critic/Value Model (价值模型): 用来预测当前状态能拿多少分的辅助模型。

对于 12GB 显存来说，装一个模型都费劲，更别说同时跑 4 个了。

痛点二：极度不稳定（调参地狱）

PPO 是典型的强化学习，它有两个阶段：

1. 先花大精力训练一个 RM（奖励模型）。如果 RM 训练得不好，后面的 PPO 怎么跑都是废的（这叫 Reward Hacking，模型学会了卡 RM 的 Bug 来骗取高分，但生成的文本其实是乱码）。
2. Actor 和 Critic 之间需要来回博弈，学习率、KL 散度惩罚系数、PPO 截断参数等超参数极难调整，动不动就“崩盘”（Loss 爆炸或生成无意义输出）。

---

2. DPO 的破局思路：数学上的“狸猫换太子”

DPO 之所以被称作神作，是因为斯坦福的学者在数学上证明了一件事：我们根本不需要去训练那个脆弱的奖励模型（Reward Model）。

在传统的 RLHF 中，目标是找到一个策略 ${\pi }_{\theta }$，最大化奖励 $r(x,y)$，同时不能偏离参考模型 ${\pi }_{ref}$ 太远。

DPO 的作者发现，在 Bradley-Terry (BT) 偏好模型下，最优策略和奖励函数之间存在一个解析解（直接的等价关系）。也就是说，语言模型本身就可以被当做奖励模型来用！

既然可以直接用大模型的概率分布来表示人类偏好，那我们干脆把强化学习问题，退化成了一个普通的监督学习问题（分类问题）。

---

3. DPO 内部实现详解：数据流与公式

数据输入

DPO 的数据不再是单纯的 $(Prompt,Response)$，而是三元组 $(x,y_w,y_l)$：

• $x$：问题 (Prompt)
• $y_w$：好的回答 (Chosen/Win)
• $y_l$：差的回答 (Rejected/Loss)

核心数学公式

DPO 的损失函数非常优雅，只有这一行核心公式：

$${\mathcal{L}}_{DPO}({\pi }_{\theta };{\pi }_{ref})=-{\mathbb{E}}_{(x,y_w,y_l)\sim D}\left[\log \sigma \left(\beta \log \frac{{\pi }_{\theta }(y_w|x)}{{\pi }_{ref}(y_w|x)}-\beta \log \frac{{\pi }_{\theta }(y_l|x)}{{\pi }_{ref}(y_l|x)}\right)\right]$$

我们把它拆解成人话，看看内部每一步在干什么：

Step 1: 参考模型的基准评估 (Reference Log Probs)

把 $x+y_w$ 和 $x+y_l$ 分别输入给冻结的参考模型 ${\pi }_{ref}$，得到两个对数概率：

• ref_logp_w (参考模型对好答案的打分)
• ref_logp_l (参考模型对坏答案的打分)

Step 2: 策略模型的当前评估 (Policy Log Probs)

把同样的输入喂给正在训练的策略模型 ${\pi }_{\theta }$，得到：

• policy_logp_w (当前模型对好答案的打分)
• policy_logp_l (当前模型对坏答案的打分)

Step 3: 计算隐式奖励差 (Implicit Reward Margin)

DPO 认为，当前模型比参考模型多出来的概率，就是它的“奖励”。

• 对好答案的奖励：$R_w=\beta (\text{policy\_logp\_w}-\text{ref\_logp\_w})$
• 对坏答案的奖励：$R_l=\beta (\text{policy\_logp\_l}-\text{ref\_logp\_l})$
• 注：公式中的 $\beta$ (beta) 是一个温度超参数，通常设为 0.1，用来控制模型偏离参考模型的程度。

Step 4: 概率拔河 (Sigmoid Cross Entropy)

我们希望模型给好答案的奖励 $R_w$，远远大于给坏答案的奖励 $R_l$。
所以计算它们的差值：logits = R_w - R_l。
最后套上一个 Log-Sigmoid 函数（这其实就是二分类交叉熵的底层逻辑），计算出 Loss。如果 $R_w$ 比 $R_l$ 大很多，Loss 就会逼近于 0。

---

4. 辅助代码讲解：DPO 损失函数的 PyTorch 实现

这就是 TRL 库中 DPOTrainer 内部计算 Loss 的最核心代码逻辑（精简版）：

import torch
import torch.nn.functional as F

def dpo_loss(
    policy_chosen_logps,    # pi_theta(y_w|x)
    policy_rejected_logps,  # pi_theta(y_l|x)
    ref_chosen_logps,       # pi_ref(y_w|x)
    ref_rejected_logps,     # pi_ref(y_l|x)
    beta=0.1
):
    # 1. 计算好答案相对于参考模型的对数概率差
    pi_logratios = policy_chosen_logps - policy_rejected_logps
    ref_logratios = ref_chosen_logps - ref_rejected_logps
    
    # 或者写成 DPO 论文里的标准形式：
    # chosen_logratios = policy_chosen_logps - ref_chosen_logps
    # rejected_logratios = policy_rejected_logps - ref_rejected_logps
    # logits = chosen_logratios - rejected_logratios

    logits = pi_logratios - ref_logratios
    
    # 2. 乘以 beta 参数，并计算 Log-Sigmoid
    # 如果 logits 越大（对好的回答越来越确信，对坏的越来越排斥），Loss 越小
    losses = -F.logsigmoid(beta * logits)
    
    # 3. 计算这一批次的平均 Loss
    return losses.mean()

总结

• PPO 的逻辑是： 先雇一个裁判（Reward Model），然后选手（Actor）不停地做题，裁判给分，选手根据分数调整自己。
• DPO 的逻辑是： 把历年真题（好坏答案对比）直接塞给选手，告诉他：“这个得分高，那个得分低，你自己去悟出评分标准（隐式奖励），同时调整你的做题策略”。

因为省去了裁判模型，DPO 极大地节省了显存，且训练过程变成了纯粹的梯度下降优化，彻底消灭了强化学习中的不稳定性。

以下为可直接复制运行的dpo代码（已测试）：

"""
DPO (Direct Preference Optimization) 训练示例脚本
===============================================
本脚本演示如何使用 DPO 技术微调大语言模型，让模型学习特定的对话风格。

DPO 核心思想：
- 不需要训练奖励模型
- 直接使用偏好数据（chosen vs rejected）优化模型
- 让模型学会"什么回答更好"

硬件要求：
- 显卡：NVIDIA RTX 4060 12GB 或更高
- 显存：约 8-10GB（使用 4-bit 量化）
- 存储：约 5GB（模型文件 + 训练结果）

依赖库：
- torch
- transformers
- trl (Transformer Reinforcement Learning)
- peft (Parameter Efficient Fine-Tuning)
- datasets
- modelscope (国内模型下载加速)
"""

import os
import torch
from datasets import load_dataset, Dataset
from modelscope import snapshot_download
from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig, TrainingArguments
from trl.trainer.dpo_trainer import DPOTrainer
from trl.trainer.dpo_config import DPOConfig
from peft import LoraConfig, PeftModel

# ================= 1. 环境与模型准备 =================
"""
步骤1：下载并加载预训练模型
- 使用 ModelScope 加速下载（国内镜像）
- 选择 Qwen2.5-1.5B-Instruct 模型（轻量级，适合本地训练）
- 使用 4-bit 量化减少显存占用
"""

print("正在从 ModelScope 下载模型...")
# 自动从魔搭下载模型到本地缓存目录
# Qwen2.5-1.5B 约 3GB，适合个人电脑运行
model_dir = snapshot_download('qwen/Qwen2.5-1.5B-Instruct')

# 4-bit 量化配置（QLoRA 技术）
# 目的：将模型权重从 16-bit 压缩到 4-bit，大幅减少显存使用
# 原理：使用 NF4 (Normal Float 4) 量化类型，保持模型精度
bnb_config = BitsAndBytesConfig(
    load_in_4bit=True,                    # 启用 4-bit 量化
    bnb_4bit_quant_type="nf4",            # 使用 NF4 量化类型（精度更高）
    bnb_4bit_compute_dtype=torch.bfloat16, # 计算时使用 bfloat16（平衡速度和精度）
)

# 加载分词器（Tokenizer）
# Tokenizer 负责将文本转换为模型可理解的数字序列
tokenizer = AutoTokenizer.from_pretrained(model_dir)
# 设置填充标记为结束标记（EOS token）
# 原因：生成模型通常使用 EOS 作为序列结束标记
tokenizer.pad_token = tokenizer.eos_token

# 加载预训练模型
# device_map="auto"：自动分配模型层到可用设备（GPU/CPU）
# trust_remote_code=True：允许执行模型仓库中的自定义代码
model = AutoModelForCausalLM.from_pretrained(
    model_dir,
    quantization_config=bnb_config,       # 应用 4-bit 量化配置
    device_map="auto",                    # 自动设备映射
    trust_remote_code=True                # 信任远程代码（Qwen模型需要）
)
# 训练时必须关闭 KV Cache
# 原因：KV Cache 用于加速推理，但训练时需要重新计算梯度
model.config.use_cache = False

# ================= 1.5 测试：查看训练前的模型输出 =================
"""
步骤2：测试原始模型的输出风格
- 使用一个测试问题
- 观察模型在训练前的回答风格（通常是标准的 AI 助手风格）
- 作为对比基准
"""

# 设定测试问题（模拟日常对话场景）
test_prompt = "嘿，忙啥呢？晚上出来撸串不？"

# 获取原始模型（Base Model）的回答
print("\n--- [训练前：原始 Qwen2.5 机器人风格] ---\n")
print(test_prompt)

# 构建消息格式（OpenAI 风格的消息列表）
messages = [{"role": "user", "content": test_prompt}]

# 应用聊天模板：将消息列表转换为模型输入格式
# tokenize=False：返回字符串，不转换为 token IDs
# add_generation_prompt=True：添加生成提示，告诉模型需要生成回复
text = tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)

# 将文本转换为模型输入张量
# return_tensors="pt"：返回 PyTorch 张量
inputs = tokenizer([text], return_tensors="pt").to(model.device)

# 生成回复（不计算梯度，节省显存）
with torch.no_grad():
    outputs = model.generate(**inputs, max_new_tokens=100)
    # 解码输出：将 token IDs 转换回文本
    # split("assistant\n")[-1]：提取 assistant 的回复部分
    print("\n原始模型回复：")
    print(tokenizer.decode(outputs[0], skip_special_tokens=True).split("assistant\n")[-1])
    # 嗨！我最近在学习新的编程语言，希望未来能做一些有趣的项目。晚上有空的话可以一起吃个烤串，怎么样？

# ================= 2. 准备 DPO 数据集 =================
"""
步骤3：构建 DPO 训练数据集

DPO 数据集格式要求：
- prompt: 用户的提问或指令
- chosen: 偏好的回答（希望模型学习的风格）
- rejected: 不喜欢的回答（希望模型避免的风格）

数据构建原则：
1. chosen 和 rejected 对应同一个 prompt
2. chosen 应该比 rejected 更符合期望的风格
3. 数据量不需要很大（几十到几百条即可）
"""

print("正在构造本地演示数据集（跳过 Hugging Face 下载）...")

# 构造示例数据集
# 目标：让模型学习更自然、亲切的对话风格，摒弃生硬的 AI 腔调
mock_data = {
    "prompt": [
        "你吃了吗？",
        "今天天气怎么样？",
        "我想去蹦迪，你去吗？",
        "老板叫我加班，好烦啊。",
        "你会写代码吗？"
    ],
    "chosen": [
        # 期望的风格：随性、亲切、有情感
        "还没呢，正打算去弄点烧烤，你要一起不？",
        "外面阴沉沉的，感觉要下雨，出门记得带伞哈。",
        "走起啊！今晚不醉不归，我早就想去晃荡一下了。",
        "害，打工人实惨。摸摸头，等加完班请你喝奶茶补补偿。",
        "略懂略懂，不过要是太难的逻辑我也得抓瞎，咱们一起研究呗。"
    ],
    "rejected": [
        # 不喜欢的风格：生硬、机械、过于正式
        "作为一个人工智能助手，我不需要进食。",
        "根据天气预报显示，今日局部地区有阵雨，建议携带雨具。",
        "我不建议进行嘈杂的娱乐活动，这不利于身心健康。",
        "加班是职场中常见的现象，建议您合理安排时间，保持积极心态。",
        "我可以为您提供多种编程语言的代码示例，请问您需要什么？"
    ]
}

# 将字典转换为 Hugging Face Dataset 格式
# Dataset 格式便于后续处理和训练
dataset = Dataset.from_dict(mock_data)
print("本地数据集构造完成！")

# ================= 3. 配置 LoRA 和 DPO =================
"""
步骤4：配置 LoRA 和 DPO 训练参数

LoRA (Low-Rank Adaptation)：
- 只训练少量参数（适配器），冻结原模型
- 大幅减少训练时间和显存需求
- 训练后的权重可以合并回原模型或单独保存

DPO 训练参数：
- beta：控制偏好学习的强度（通常 0.1-0.5）
- 其他参数类似标准监督学习
"""

# LoRA 配置
# 原理：在原始权重矩阵旁添加低秩矩阵进行微调
peft_config = LoraConfig(
    r=8,                      # LoRA 秩：低秩矩阵的维度，越大表达能力越强
    lora_alpha=16,            # 缩放因子：控制 LoRA 权重的影响程度
    # 目标模块：指定哪些层添加 LoRA 适配器
    # q_proj, k_proj, v_proj, o_proj：注意力层的投影矩阵
    # gate_proj, up_proj, down_proj：MLP 层的投影矩阵
    target_modules=["q_proj", "v_proj", "k_proj", "o_proj", "gate_proj", "up_proj", "down_proj"],
    lora_dropout=0.05,        # Dropout 率：防止过拟合
    bias="none",              # 是否训练偏置项
    task_type="CAUSAL_LM",    # 任务类型：因果语言模型（生成任务）
)

# DPO 训练配置
dpo_config = DPOConfig(
    output_dir="./dpo_qwen_results",      # 训练结果保存目录
    per_device_train_batch_size=1,        # 每个设备的批次大小（显存限制设为1）
    gradient_accumulation_steps=4,        # 梯度累积步数：每4步更新一次参数
    learning_rate=5e-5,                   # 学习率：控制参数更新幅度
    lr_scheduler_type="cosine",           # 学习率调度：余弦退火，逐渐降低学习率
    max_steps=50,                         # 最大训练步数（演示用，实际可更大）
    save_steps=50,                        # 每50步保存一次检查点
    logging_steps=5,                      # 每5步记录一次日志
    bf16=True,                            # 使用 bfloat16 混合精度训练
    gradient_checkpointing=True,          # 梯度检查点：节省显存，但速度稍慢
    remove_unused_columns=False,          # 保留所有列（DPO 需要多列数据）
    beta=0.1,                             # DPO 温度参数：控制偏好学习的强度
    max_length=1024,                      # 最大序列长度
)

# 初始化 DPO 训练器
# DPOTrainer 自动处理偏好数据的损失计算
dpo_trainer = DPOTrainer(
    model=model,                          # 要训练的模型
    ref_model=None,                       # 参考模型（None 表示使用当前模型作为参考）
    args=dpo_config,                      # 训练配置
    train_dataset=dataset,                # 训练数据集
    processing_class=tokenizer,           # 分词器（新版 TRL 使用 processing_class）
    peft_config=peft_config,              # LoRA 配置
)

# ================= 4. 执行训练 =================
"""
步骤5：开始 DPO 训练

训练过程：
1. 对每个样本，计算 chosen 和 rejected 的 log 概率
2. 计算偏好损失（DPO 损失函数）
3. 反向传播更新 LoRA 参数
4. 重复直到达到 max_steps

训练时间：
- RTX 4060 12GB：约 5-10 分钟（50步）
"""

print("开始 DPO 训练...")
dpo_trainer.train()

# 保存微调后的模型权重（LoRA 适配器）
# 注意：保存的是适配器权重，不是完整模型
dpo_trainer.save_model("./dpo_qwen_final")
print("训练完成！模型已保存至 ./dpo_qwen_final")

# ================= 5. 测试训练后的模型 =================
"""
步骤6：加载训练后的模型并测试

对比训练前后的差异：
- 训练前：标准 AI 助手风格，正式、机械
- 训练后：学习到的个人风格，自然、亲切
"""

# 加载微调后的 LoRA 适配器权重
# PeftModel：将适配器权重合并到基础模型
print("\n--- [加载你的数字人分身权重...] ---")
model = PeftModel.from_pretrained(model, "./dpo_qwen_final")

# 获取微调后（Fine-tuned）的回答
print("\n--- [训练后：你的个人风格分身] ---")
with torch.no_grad():
    outputs = model.generate(**inputs, max_new_tokens=100)
    print(tokenizer.decode(outputs[0], skip_special_tokens=True).split("assistant\n")[-1])
    # 嗨！晚上出来撸串挺好的，你有时间的话可以来我家。不过我得先回家做饭了，不然今天就吃不了了。

"""
总结
====
本脚本演示了完整的 DPO 微调流程：
1. 加载预训练模型（使用 4-bit 量化节省显存）
2. 测试原始模型的输出风格
3. 构建偏好数据集（chosen vs rejected）
4. 配置 LoRA 和 DPO 训练参数
5. 执行 DPO 训练
6. 测试训练后的模型效果

关键概念：
- DPO：直接偏好优化，无需奖励模型
- LoRA：参数高效微调，只训练少量参数
- 4-bit 量化：大幅减少显存占用

扩展建议：
1. 增加更多训练数据（几百到几千条）
2. 调整 beta 参数（0.1-0.5 之间尝试）
3. 增加训练步数（100-500步）
4. 尝试不同的 LoRA 秩（r=4, 8, 16, 32）
5. 合并 LoRA 权重到基础模型，便于部署
"""