PPO算法实践：手把手教会你PPO算法的工程应用

PPO 模型的训练

我们需要的模型

• 1. 基准模型：一般是SFT后的模型作为基准，新训练的模型不能和这个模型的概率分布相差太大。
• 2. 训练模型： 他的结构和基准模型是一样的。
• 3. reward模型：对一个问答序列进行打分，输出是一个分数。
• 4. 状态价值模型：对每个状态进行评估，对截止到目前的序列预测到序列生成结束后这个序列的期望回报是多少，对每个token都输出分数，输出是一个分数。

我们可以使用LoRA技术，只使用一个大模型，多个LoRA层，来完成这个任务。减少训练时对显存的占用。训练模型和状态价值模型可以共用一个loRA层，不同的头来实现。

实现流程伪代码

for batch_prompt in prompt_dataset:
    batch_response = active_model.generate(batch_prompt)# 策略模型的响应
    batch_data = concat(batch_prompt, batch_response)# 连接问题和响应
    batch_scores = reward_model(batch_data)# 计算得分

    batch_all_probs, batch_probs, batch_all_values = active_model.forward_pass(batch_data)# 对批次数据进行前向传播，得到所有可能动作的概率（`batch_all_probs`）、选择动作的概（`batch_probs`）和所有可能动作的价值（`batch_all_values`）
    ref_all_probs, ref_probs, ref_all_values = ref_model.forward_pass(batch_data)# 计算基础模型的所有可能动作的概率（`ref_all_probs`）、选择动作的概率（`ref_probs`）和所有可能动作的价值（`ref_all_values`）
    kls = compute_KL(batch_all_probs, ref_all_probs)# 计算KL散度
    rewards = compute_rewards(batch_scores, kls)# 根据得分和KL散度计算奖励。
    advantages = compute_advantages(batch_all_values, rewards)# 计算优势函数，即奖励与价值函数估计之间的差异。
    returns = advantages + batch_all_values# 计算回报，即优势函数与价值函数估计的和。
 
   for i in range(epoch):
       active_all_probs, active_probs, active_all_values = active_model.forward_pass(batch_data)

       loss_state_value = torch.mean((returns - active_all_values) ** 2)
       ratio = active_probs / batch_probs
       loss_ppo = torch.mean(-advantages * ratio)
       loss = loss_ppo + value_loss_rate * loss_state_value
       loss.backward()
       optimizer.step()
       optimizer.zero_grad()

$$Los{s}_{PPO}=-\frac{1}{N}\sum _{n=1}^N\sum _{t=1}^{T_n}{A}_{{\theta }^{\prime \prime }}^{GAE}(s_n^t,a_n^t)\frac{P_{\theta }(a_n^t|s_n^t)}{P_{{\theta }^{\prime \prime }}(a_n^t|s_n^t)}$$
代码片段中，PPO（Proximal Policy Optimization）算法的损失函数体现在以下部分：

ratio = active_probs / batch_probs
loss_ppo = torch.mean(-advantages * ratio)

让我们详细解释这些代码行是如何与PPO算法的损失函数公式相对应的：

代码中的公式解释

PPO算法的损失函数公式为：

$$Los{s}_{PPO}=-\frac{1}{N}\sum _{n=1}^N\sum _{t=1}^{T_n}{A}_{{\theta }^{\prime \prime }}^{GAE}(s_n^t,a_n^t)\frac{P_{\theta }(a_n^t|s_n^t)}{P_{{\theta }^{\prime \prime }}(a_n^t|s_n^t)}$$
其中：

• $N$ 是批次大小。
• $T_n$ 是每个样本的时间步数。
• $A_{{\theta }^{\prime \prime }}^{GAE}(s_n^t,a_n^t)$ 是优势函数，使用GAE（Generalized Advantage Estimation）计算。
• $P_{\theta }(a_n^t|s_n^t)$ 是新策略在状态 $s_n^t$ 下选择动作 $a_n^t$ 的概率。
• $P_{{\theta }^{\prime \prime }}(a_n^t|s_n^t)$ 是旧策略在状态 $s_n^t$ 下选择动作 $a_n^t$ 的概率。

代码解释

1. 计算概率比率：

   ratio = active_probs / batch_probs
   

   • active_probs 是新策略在给定状态下选择动作的概率。
   • batch_probs 是基准模型在给定状态下选择动作的概率。
   • 这对应于公式中的 $\frac{P_{\theta }(a_n^t|s_n^t)}{P_{{\theta }^{\prime \prime }}(a_n^t|s_n^t)}$。

2. 计算PPO损失：

   loss_ppo = torch.mean(-advantages * ratio)
   

   • advantages 是优势函数的估计值，对应于公式中的 $A_{{\theta }^{\prime \prime }}^{GAE}(s_n^t,a_n^t)$。
   • -advantages * ratio 计算了PPO损失的一部分，对应于公式中的 $-A_{{\theta }^{\prime \prime }}^{GAE}(s_n^t,a_n^t)\frac{P_{\theta }(a_n^t|s_n^t)}{P_{{\theta }^{\prime \prime }}(a_n^t|s_n^t)}$。
   • torch.mean() 计算了所有样本的平均值，对应于公式中的 $\frac{1}{N}{\sum }_{n=1}^N{\sum }_{t=1}^{T_n}$。

3. 计算总损失：

   loss = loss_ppo + value_loss_rate * loss_state_value
   

   • loss_state_value 是状态价值损失，用于衡量价值函数的估计值与实际回报之间的差异。
   • value_loss_rate 是状态价值损失的权重。
   • 这对应于公式中的 $\frac{1}{N}{\sum }_{n=1}^N{\sum }_{t=1}^{T_n}$ 部分，但公式中没有直接体现状态价值损失。

总结来说，代码中的 loss_ppo 计算部分直接体现了PPO算法损失函数的核心思想，即通过计算新旧策略概率比率与优势函数的乘积来优化策略。而状态价值损失部分则是为了提高价值函数的估计精度。
batch训练为外循环训练，训练epoch为内循环训练。每次用当前训练的模型作为重要性采样的模型计算advantage，训练epoch次模型

数据准备阶段

1. 遍历数据集：

   • for batch_prompt in prompt_dataset:：遍历数据集中的每个批次的提示（prompt）。

2. 生成响应：

   • batch_response = active_model.generate(batch_prompt)：使用当前的策略模型（active_model）根据提示生成响应。

3. 合并数据：

   • batch_data = concat(batch_prompt, batch_response)：将提示和响应合并成一个批次的数据。

4. 计算奖励：

   • batch_scores = reward_model(batch_data)：使用奖励模型（reward_model）计算批次数据的得分，这些得分将用于计算奖励。

5. 前向传播：

   • batch_all_probs, batch_probs, batch_all_values = active_model.forward_pass(batch_data)：对批次数据进行前向传播，得到所有可能动作的概率（batch_all_probs）、选择动作的概率（batch_probs）和所有可能动作的价值（batch_all_values）。
   • ref_all_probs, ref_probs, ref_all_values = ref_model.forward_pass(batch_data)：对批次数据进行前向传播，得到参考模型（ref_model）的所有可能动作的概率（ref_all_probs）、选择动作的概率（ref_probs）和所有可能动作的价值（ref_all_values）。

6. 计算KL散度：

   • kls = compute_KL(batch_all_probs, ref_all_probs)：计算当前策略模型和参考模型之间的KL散度，用于衡量两个概率分布的差异。

7. 计算奖励：

   • rewards = compute_rewards(batch_scores, kls)：根据得分和KL散度计算奖励。

8. 计算优势：

   • advantages = compute_advantages(batch_all_values, rewards)：计算优势函数，即奖励与价值函数估计之间的差异。

9. 计算回报：

   • returns = advantages + batch_all_values：计算回报，即优势函数与价值函数估计的和。

训练阶段

1. 遍历训练周期：

   • for i in range(epoch):：遍历每个训练周期。

2. 前向传播：

   • active_all_probs, active_probs, active_all_values = active_model.forward_pass(batch_data)：再次对批次数据进行前向传播，得到当前策略模型的概率和价值。

3. 计算状态价值损失：

   • loss_state_value = torch.mean((returns - active_all_values) ** 2)：计算状态价值损失，即回报与价值函数估计之间的均方误差。

4. 计算概率比率：

   • ratio = active_probs / batch_probs：计算新旧策略选择动作的概率比率。

5. 计算PPO损失：

   • loss_ppo = torch.mean(-advantages * ratio)：计算PPO损失，即优势函数与概率比率的乘积的负均值。

6. 计算总损失：

   • loss = loss_ppo + value_loss_rate * loss_state_value：计算总损失，即PPO损失与状态价值损失的加权和。

7. 反向传播：

   • loss.backward()：对总损失进行反向传播，计算梯度。

8. 更新模型参数：

   • optimizer.step()：使用优化器更新模型参数。

9. 清零梯度：

   • optimizer.zero_grad()：清零梯度，为下一次迭代做准备。

实现代码

# 导入必要的库
import torch  # PyTorch 深度学习框架
from peft import LoraConfig, TaskType  # PEFT 库，用于微调模型
from transformers import AutoTokenizer, BitsAndBytesConfig  # transformers 库，用于加载模型和分词器
from trl import AutoModelForCausalLMWithValueHead, PPOConfig, PPOTrainer  # trl 库，用于强化学习训练
from datasets import Dataset  # datasets 库，用于处理数据集
import json  # 用于处理 JSON 格式的数据

# 模型路径
model_path = r'D:\work\models\Meta-Llama-3.1-8B-Instruct'  # 模型文件的路径

# 加载分词器
tokenizer = AutoTokenizer.from_pretrained(model_path, use_fast=False)  # 加载预训练分词器
tokenizer.padding_side = "right"  # 设置填充方向为右侧
tokenizer.pad_token = tokenizer.eos_token  # 将填充符设置为结束符

# 配置 4 位量化参数
bnb_config = BitsAndBytesConfig(
    load_in_4bit=True,  # 启用 4 位量化
    bnb_4bit_use_double_quant=True,  # 使用双量化
    bnb_4bit_quant_type="nf4",  # 使用 nf4 量化类型
    bnb_4bit_compute_dtype=torch.bfloat16  # 计算数据类型为 bfloat16
)

# 配置 LoRA 微调参数
peft_config = LoraConfig(
    r=8,  # LoRA 的秩
    target_modules=["q_proj", "v_proj", "k_proj", "o_proj", "gate_proj", "down_proj", "up_proj"],  # 需要微调的模块
    task_type=TaskType.CAUSAL_LM,  # 任务类型为因果语言模型
    lora_alpha=16,  # LoRA 的 alpha 参数
    lora_dropout=0.05  # LoRA 的 dropout 概率
)

# 加载模型
model = AutoModelForCausalLMWithValueHead.from_pretrained(
    model_path,  # 模型路径
    reward_adapter="./reward_model",  # 奖励模型路径
    peft_config=peft_config,  # LoRA 配置
    quantization_config=bnb_config  # 量化配置
)
model.to("cuda")  # 将模型移动到 GPU 设备

# 加载训练数据
items = []
with open("./data/queries.json", "r", encoding="utf8") as f:  # 打开训练数据文件
    for line in f:
        items.append(json.loads(line))  # 解析 JSON 数据并添加到列表
queries_dataset = Dataset.from_list(items)  # 将数据转换为 Dataset 对象

# 定义数据收集器函数
def collator(data):
    queries = []
    for item in data:
        queries.append(tokenizer(item["query"], return_tensors="pt")["input_ids"].squeeze().to("cuda"))
        # 对每个 query 进行分词、转换为张量，并移动到 GPU 设备
    return queries

# 配置 PPO 训练参数
ppo_config = PPOConfig(kl_penalty="full", ppo_epochs=3, batch_size=2, mini_batch_size=1)
ppo_trainer = PPOTrainer(
    config=ppo_config,  # PPO 配置
    model=model,  # 微调模型
    ref_model=None,  # 参考模型（这里没有使用）
    tokenizer=tokenizer,  # 分词器
    dataset=queries_dataset,  # 训练数据集
    data_collator=collator  # 数据收集器
)

# 定义生成参数
generation_kwargs = {
    "min_length": -1,  # 最小生成长度
    "top_k": 0.0,  # top-k 采样参数
    "top_p": 1.0,  # top-p 采样参数
    "do_sample": True,  # 是否进行采样
    "pad_token_id": tokenizer.pad_token_id,  # 填充符 ID
    "max_new_tokens": 32,  # 最大新生成的 token 数量
}

# 开始 PPO 训练循环
for batch in ppo_trainer.dataloader:  # 遍历数据加载器
    query_tensors = batch  # 获取当前批次的 query 张量

    # 生成响应
    response_tensors = ppo_trainer.generate(query_tensors, return_prompt=False, **generation_kwargs)

    # 计算奖励分数
    scores = []
    for query, response in zip(query_tensors, response_tensors):
        input_ids = torch.concat([query, response], dim=0)  # 将 query 和 response 拼接
        input_ids = torch.unsqueeze(input_ids, dim=0)  # 增加一个维度
        score = ppo_trainer.model.compute_reward_score(input_ids=input_ids)[0, -1, 0]  # 计算奖励分数
        scores.append(score)

    # 执行 PPO 训练步骤
    stats = ppo_trainer.step(query_tensors, response_tensors, scores)

# 保存训练后的模型
ppo_trainer.save_pretrained("./rl_model")