大模型强化学习——PPO项目实战
【PPO算法介绍】
PPO（Proximal Policy Optimization）是一种强化学习算法，它的目标是找到一个策略，使得根据这个策略采取行动可以获得最大的累积奖励。
PPO的主要思想是在更新策略时，尽量让新策略不要偏离旧策略太远。这是通过在目标函数中添加一个额外的项来实现的，这个额外的项会惩罚新策略和旧策略之间的差异。这样可以避免在更新策略时出现过大的跳跃，从而提高学习的稳定性。
举一个通俗的例子，假设你正在玩一个游戏，你已经找到了一个还不错的策略，可以让你获得一定的分数。现在，你想要改进你的策略，以便获得更高的分数。但是，你不希望新的策略和旧的策略差距太大，因为这可能会让你的分数大幅度下降。所以，你会尽量让新的策略和旧的策略相近，这就是PPO的主要思想。
在实际操作中，PPO通过使用一种叫做“裁剪目标函数”的技术来实现这个思想。这个技术会限制新策略和旧策略之间的差异，从而避免更新过程中的大幅度跳跃。
针对大模型应用：
在强化学习中，"策略"是指模型在给定状态下选择动作的规则。对于一个对话生成模型来说，状态可能是当前的对话历史，动作则是生成下一个词或者句子。
"更新模型的生成策略"，就是通过学习过程改变模型选择下一个词或句子的规则。具体来说，模型会根据与环境的交互反馈（例如，奖励信号）来调整其参数，这样在相同的对话历史下，模型可能会选择生成不同的下一个词或句子。
例如，假设在一个对话生成任务中，模型在给定"天气"这个输入时，原本会生成"很好"这个输出。但是，如果在训练过程中，模型发现生成"很热"这个输出会得到更高的奖励，那么模型就会调整其参数，使得在下次给定"天气"这个输入时，模型会更倾向于生成"很热"这个输出。这就是一个更新生成策略的例子。
需要注意的是，这个"策略"并不是显式定义的规则，而是模型参数的函数。也就是说，模型的策略是通过调整模型的参数来改变的。
RL阶段实战，通过强化学习PPO算法对SFT模型进行优化，帮助读者深入理解ChatGPT模型在RL阶段的任务流程。
项目主要结构如下：
- data 存放数据的文件夹
- ppo_train.json 用于强化学习的文档数据
- rm_model RM阶段训练完成模型的文件路径
- config.json
- pytorch_model.bin
- vocab.txt
- sft_model SFT阶段训练完成模型的文件路径
- config.json
- pytorch_model.bin
- vocab.txt
- ppo_model PPO阶段训练完成模型的文件路径
- config.json
- pytorch_model.bin
- vocab.txt
- data_set.py 模型所需数据类文件
- model.py 模型文件
- train.py 模型训练文件
- predict.py 模型推理文件
注意：由于GitHub不方便放模型文件，因此sft_model文件夹、rm_model文件夹和ppo_model文件夹中的模型bin文件，请从百度云盘中下载。
| 文件名称 | 下载地址 | 提取码 |
|---|---|---|
| sft_model | 百度云 | iks4 |
| rm_model | 百度云 | 64wt |
| ppo_model | 百度云 | s8b7 |
环境配置
模型训练或推理所需环境，请参考requirements.txt文件。
模型训练
模型的训练流程如下所示：
模型训练需要运行train.py文件，会自动生成output_dir文件夹，存放每个epoch保存的模型文件。
命令如下：
python3 train.py --device 0 \
--train_file_path "data/ppo_train.json" \
--max_len 768 \
--query_len 64 \
--batch_size 16
注意：当服务器资源不同或读者更换数据等时，可以在模型训练时修改响应参数，详细参数说明见代码或阅读书9.3.3小节。
模型训练示例如下：
模型训练阶段损失值变化如下：
模型推理
模型训练需要运行predict.py文件，可以采用项目中以提供的模型，也可以采用自己训练后的模型。
命令如下：
python3 predict.py --device 0 --model_path ppo_model
注意：如果修改模型路径，请修改--model_path参数。
模型推理示例如下：
样例1：
输入的正文为：大莱龙铁路位于山东省北部环渤海地区，西起位于益羊铁路的潍坊大家洼车站，向东经海化、寿光、寒亭、昌邑、平度、莱州、招远、终到龙口，连接山东半岛羊角沟、潍坊、莱州、龙口四个港口，全长175公里，工程建设概算总投资11.42亿元。铁路西与德大铁路、黄大铁路在大家洼站接轨，东与龙烟铁路相连。
生成的第1个问题为：该项目建成后对于做什么？
生成的第2个问题为：该铁路线建成后会带动什么方面？
样例2：
输入的正文为：椰子猫（学名：'），又名椰子狸，为分布于南亚及东南亚的一种麝猫。椰子猫平均重3.2公斤，体长53厘米，尾巴长48厘米。它们的毛粗糙，一般呈灰色，脚、耳朵及吻都是黑色的。它们的身体上有三间黑色斑纹，面部的斑纹则像浣熊，尾巴没有斑纹。椰子猫是夜间活动及杂食性的。它们在亚洲的生态位与在北美洲的浣熊相近。牠们栖息在森林、有树木的公园及花园之内。它们的爪锋利，可以用来攀爬。椰子猫分布在印度南部、斯里兰卡、东南亚及中国南部。
生成的第1个问题为：椰子猫是什么族群？
生成的第2个问题为：椰子猫到底是什么？
注意
需要36GB内存才可以训练。我自己单机无法运行：
File "C:\Python311\Lib\site-packages\transformers\pytorch_utils.py", line 106, in forward
x = torch.addmm(self.bias, x.view(-1, x.size(-1)), self.weight)
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
RuntimeError: [enforce fail at ..\c10\core\impl\alloc_cpu.cpp:72] data. DefaultCPUAllocator: not enough memory: you tried to allocate 35684352 bytes.
【核心代码解读】
from utils import get_advantages_and_returns, actor_loss_function, critic_loss_function
from model import ActorModel, RewardModel, CriticModel
import argparse
from data_set import ExamplesSampler
import os
from transformers.models.bert import BertTokenizer
from torch.optim import Adam
import random
import numpy as np
import torch
try:
from torch.utils.tensorboard import SummaryWriter
except ImportError:
from tensorboard import SummaryWriter
def make_experience(args, actor_model, critic_model, ori_model, reward_model, input_ids, generate_kwargs):
"""获取经验数据"""
actor_model.eval()
critic_model.eval()
with torch.no_grad():
# 获取prompt内容长度
prompt_length = input_ids.shape[1]
# 使用动作模型通过已有提示生成指定内容，其中：seq_outputs为返回序列，包含prompt+生成的answer
seq_outputs, attention_mask = actor_model.generate(input_ids, **generate_kwargs)
# 通过动作模型和原始模型同时计算生成结果对应的log_probs
action_log_probs = actor_model(seq_outputs, attention_mask)
base_action_log_probs = ori_model(seq_outputs, attention_mask)
# 通过评判模型计算生成的answer的分值
value, _ = critic_model(seq_outputs, attention_mask, prompt_length)
value = value[:, :-1]
# 通过奖励模型计算生成奖励值，并对奖励值进行裁剪
_, reward_score = reward_model.forward(seq_outputs, attention_mask, prompt_length=prompt_length)
reward_clip = torch.clamp(reward_score, -args.reward_clip_eps, args.reward_clip_eps)
# reward_clip = reward_score
# 对动作模型和原始模型的log_probs进行kl散度计算，防止动作模型偏离原始模型
kl_divergence = -args.kl_coef * (action_log_probs - base_action_log_probs)
rewards = kl_divergence
start_ids = input_ids.shape[1] - 1
action_mask = attention_mask[:, 1:]
ends_ids = start_ids + action_mask[:, start_ids:].sum(1)
batch_size = action_log_probs.shape[0]
# 将奖励值加到生成的answer最后一个token上
for j in range(batch_size):
rewards[j, start_ids:ends_ids[j]][-1] += reward_clip[j]
# 通过奖励值计算优势函数
advantages, returns = get_advantages_and_returns(value, rewards, start_ids, args.gamma, args.lam)
experience = {"input_ids": input_ids, "seq_outputs": seq_outputs, "attention_mask": attention_mask,
"action_log_probs": action_log_probs, "value": value, "reward_score": reward_score,
"advantages": advantages, "returns": returns}
return experience
def update_model(args, experience_list, actor_model, actor_optimizer, critic_model, critic_optimizer, tb_write,
ppo_step):
"""模型更新"""
# 根据强化学习训练轮数，进行模型更新
for _ in range(args.ppo_epoch):
# 随机打乱经验池中的数据，并进行数据遍历
random.shuffle(experience_list)
for i_e, experience in enumerate(experience_list):
ppo_step += 1
start_ids = experience["input_ids"].size()[-1] - 1
# 获取actor模型的log_probs
action_log_probs = actor_model(experience["seq_outputs"], experience["attention_mask"])
action_mask = experience["attention_mask"][:, 1:]
# 计算actor模型损失值
actor_loss = actor_loss_function(action_log_probs[:, start_ids:],
experience["action_log_probs"][:, start_ids:], experience["advantages"],
action_mask[:, start_ids:], args.policy_clip_eps)
# actor模型梯度回传，梯度更新
actor_loss.backward()
tb_write.add_scalar("actor_loss", actor_loss.item(), ppo_step)
torch.nn.utils.clip_grad_norm_(actor_model.parameters(), args.max_grad_norm)
actor_optimizer.step()
actor_optimizer.zero_grad()
# 计算critic模型的value
value, _ = critic_model(experience["seq_outputs"], experience["attention_mask"],
experience["input_ids"].size()[-1])
value = value[:, :-1]
# 计算critic模型损失值
critic_loss = critic_loss_function(value[:, start_ids:], experience["value"][:, start_ids:],
experience["returns"], action_mask[:, start_ids:], args.value_clip_eps)
# actor模型梯度回传，梯度更新
critic_loss.backward()
tb_write.add_scalar("critic_loss", critic_loss.item(), ppo_step)
torch.nn.utils.clip_grad_norm_(critic_model.parameters(), args.max_grad_norm)
critic_optimizer.step()
critic_optimizer.zero_grad()
return ppo_step
def train(args, ori_model, actor_model, reward_model, critic_model, tokenizer, dataset, device, tb_write):
"""模型训练"""
# 根据actor模型和critic模型构建actor优化器和critic优化器
actor_optimizer = Adam(actor_model.parameters(), lr=args.learning_rate, eps=args.adam_epsilon)
critic_optimizer = Adam(critic_model.parameters(), lr=args.learning_rate, eps=args.adam_epsilon)
cnt_timesteps = 0
ppo_step = 0
experience_list = []
mean_reward = []
# 训练
for i in range(args.num_episodes):
for timestep in range(args.max_timesteps):
cnt_timesteps += 1
# 从数据集中随机抽取batch_size大小数据
prompt_list = dataset.sample(args.batch_size)
# 生成模型所需的input_ids
input_ids = tokenizer.batch_encode_plus(prompt_list, return_tensors="pt",
max_length=args.max_len - args.query_len - 3,
truncation=True, padding=True)["input_ids"]
input_ids = input_ids.to(device)
generate_kwargs = {
"min_length": -1,
"max_length": input_ids.shape[1] + args.query_len,
"top_p": args.top_p,
"repetition_penalty": args.repetition_penalty,
"do_sample": args.do_sample,
"pad_token_id": tokenizer.pad_token_id,
"eos_token_id": tokenizer.eos_token_id,
"num_return_sequences": args.num_return_sequences}
# 生成经验数据，并添加到经验池中
experience = make_experience(args, actor_model, critic_model, ori_model, reward_model, input_ids,
generate_kwargs)
experience_list.append(experience)
# 记录数据中的奖励值
mean_reward.extend(experience["reward_score"].detach().cpu().numpy().tolist())
# 当到达更新步数，进行模型更新
if (cnt_timesteps % args.update_timesteps == 0) and (cnt_timesteps != 0):
# 打印并记录平均奖励值
mr = np.mean(np.array(mean_reward))
tb_write.add_scalar("mean_reward", mr, cnt_timesteps)
print("mean_reward", mr)
actor_model.train()
critic_model.train()
# 模型更新
ppo_step = update_model(args, experience_list, actor_model, actor_optimizer, critic_model,
critic_optimizer, tb_write, ppo_step)
# 模型更新后，将经验池清空
experience_list = []
mean_reward = []
# 模型保存
actor_model.save_pretrained(os.path.join(args.output_dir, "checkpoint-{}".format(ppo_step)))
tokenizer.save_pretrained(os.path.join(args.output_dir, "checkpoint-{}".format(ppo_step)))
print("save model")
def set_args():
"""设置训练模型所需参数"""
parser = argparse.ArgumentParser()
parser.add_argument('--device', default='2', type=str, help='设置训练或测试时使用的显卡')
parser.add_argument('--train_file_path', default='data/ppo_train.json',
type=str, help='训练数据集')
parser.add_argument('--ori_model_path', default='sft_model/', type=str, help='SFT模型')
parser.add_argument('--reward_model_path', default='rm_model/', type=str, help='奖励模型路径')
parser.add_argument('--max_len', default=768, type=int, help='模型最大长度')
parser.add_argument('--query_len', default=64, type=int, help='生成问题的最大长度')
parser.add_argument('--batch_size', default=16, type=int, help='批次大小')
parser.add_argument('--num_episodes', default=3, type=int, help='循环次数')
parser.add_argument('--max_timesteps', default=80, type=int, help='单次训练最大步骤')
parser.add_argument('--update_timesteps', default=20, type=int, help='模型更新步数')
parser.add_argument('--kl_coef', default=0.02, type=float, help='kl散度概率')
parser.add_argument('--ppo_epoch', default=2, type=int, help='强化学习训练轮数')
parser.add_argument('--policy_clip_eps', default=0.2, type=float, help='策略裁剪')
parser.add_argument('--value_clip_eps', default=0.2, type=float, help='值裁剪')
parser.add_argument('--top_p', default=1.0, type=float, help='解码Top-p概率')
parser.add_argument('--repetition_penalty', default=1.4, type=float, help='重复惩罚率')
parser.add_argument('--do_sample', default=True, type=bool, help='随机解码')
parser.add_argument('--num_return_sequences', default=1, type=int, help='生成内容个数')
parser.add_argument('--max_grad_norm', default=1.0, type=float, help='')
parser.add_argument('--reward_clip_eps', default=5.0, type=float, help='奖励值裁剪')
parser.add_argument('--gamma', default=1.0, type=float, help='优势函数gamma值')
parser.add_argument('--lam', default=0.95, type=float, help='优势函数lambda值')
parser.add_argument('--learning_rate', default=1e-5, type=float, help='学习率')
parser.add_argument('--adam_epsilon', default=1e-5, type=float, help='Adam优化器的epsilon值')
parser.add_argument('--output_dir', default="output_dir", type=str, help='模型保存路径')
parser.add_argument('--seed', default=2048, type=int, help='')
return parser.parse_args()
def main():
# 设置模型训练参数
args = set_args()
# 设置显卡信息
os.environ["CUDA_DEVICE_ORDER"] = "PCI_BUS_ID"
os.environ["CUDA_VISIBLE_DEVICES"] = args.device
# 获取device信息，用于模型训练
device = torch.device("cuda" if torch.cuda.is_available() and int(args.device) >= 0 else "cpu")
# 设置随机种子，方便模型复现
if args.seed:
torch.manual_seed(args.seed)
random.seed(args.seed)
np.random.seed(args.seed)
if not os.path.exists(args.output_dir):
os.mkdir(args.output_dir)
tb_write = SummaryWriter()
# 实例化原始模型、Actor模型、Reward模型和Critic模型
ori_model = ActorModel(args.ori_model_path)
ori_model.to(device)
actor_model = ActorModel(args.ori_model_path)
actor_model.to(device)
reward_model = RewardModel.from_pretrained(args.reward_model_path)
reward_model.to(device)
critic_model = CriticModel.from_pretrained(args.reward_model_path)
critic_model.to(device)
# 实例化tokenizer
tokenizer = BertTokenizer.from_pretrained(args.ori_model_path, padding_side='left')
tokenizer.eos_token_id = tokenizer.sep_token_id
# 加载训练数据
dataset = ExamplesSampler(args.train_file_path)
print("数据量：{}".format(dataset.__len__()))
# 开始训练
train(args, ori_model, actor_model, reward_model, critic_model, tokenizer, dataset, device, tb_write)
if __name__ == '__main__':
main()
代码的主要功能和流程：
1. 导入所需的库和模块：这包括用于强化学习的工具函数，如get_advantages_and_returns，actor_loss_function和critic_loss_function，以及模型类，如ActorModel，RewardModel和CriticModel。
2. 定义函数make_experience：这个函数用于生成经验数据，它使用actor模型和critic模型来生成新的序列，并计算相关的奖励和优势。
3. 定义函数update_model：这个函数用于更新模型的参数。它使用PPO算法来更新actor模型和critic模型的参数。
4. 定义函数train：这个函数是训练循环的主体。它首先初始化模型和优化器，然后对每个训练步骤生成经验数据，并在适当的时候更新模型的参数。
5. 定义函数set_args：这个函数用于设置训练参数。
6. 定义函数main：这个函数是脚本的入口点。它首先设置训练参数，然后加载模型和数据，最后开始训练。
这个脚本的主要流程是：首先，它会加载模型和数据，然后开始训练循环。在每个训练步骤，它会生成新的经验数据，并在适当的时候更新模型的参数。训练完成后，它会保存模型的参数。
【PPO算法体现】
在这段代码中，PPO（Proximal Policy Optimization）算法主要体现在update_model函数中。这个函数负责更新模型的参数。
在update_model函数中，首先对经验池中的数据进行随机打乱，然后遍历每个经验数据，计算actor模型和critic模型的损失值，然后进行梯度回传和参数更新。
对于actor模型，损失函数是actor_loss_function，它计算的是策略的损失。这个损失函数中实现了PPO的核心思想，即限制新策略和旧策略之间的差异。
对于critic模型，损失函数是critic_loss_function，它计算的是值函数的损失。
【PPO算法步骤】
PPO（Proximal Policy Optimization）算法的实现步骤大致如下：
1. 初始化：初始化策略网络和价值网络。策略网络用于生成动作，价值网络用于估计每个状态的价值。
2. 采样：使用当前的策略网络进行多次采样，每次采样会生成一个轨迹，轨迹包含了状态、动作、奖励等信息。
3. 计算优势函数：对于每个轨迹中的每个状态，使用价值网络和奖励信息计算优势函数。优势函数表示采取某个动作比按照价值网络的预测采取动作好多少。
4. 更新策略网络：使用优势函数和PPO的目标函数来更新策略网络。PPO的目标函数会限制新策略和旧策略之间的差异，防止更新过程中出现大的跳跃。
5. 更新价值网络：使用轨迹中的奖励信息和价值网络的预测来更新价值网络。
6. 重复：重复上述步骤，直到满足停止条件，如达到最大迭代次数或策略网络的性能达到预设的阈值。
以上是PPO算法的基本步骤，实际的实现可能会有所不同，例如可能会添加一些额外的步骤来提高性能，如使用多个并行的环境进行采样，或使用更复杂的网络结构。
【策略网络】
策略网络（Policy Network）是强化学习中的一个重要概念，它是一个函数，用于在给定状态下生成动作。
在深度强化学习中，策略网络通常由神经网络来实现。这个神经网络的输入是状态，输出是在该状态下采取每个可能动作的概率。
策略网络有两种形式：
1. 确定性策略：在给定状态下，总是生成同一个动作。这种策略网络的输出是一个动作，而不是动作的概率。
2. 随机性策略：在给定状态下，生成每个可能动作的概率。这种策略网络的输出是一个概率分布。
在强化学习的训练过程中，我们通常会使用一种叫做策略梯度（Policy Gradient）的方法来更新策略网络，使得它能生成更好的动作。这个过程通常涉及到一种叫做优势函数（Advantage Function）的概念，它用于衡量在某个状态下，采取某个动作相比于按照当前策略采取动作的优势有多大。
【价值网络】
价值网络（Value Network）是强化学习中的一个重要概念，它是一个函数，用于估计在某个状态下，按照当前策略采取行动能获得的预期回报。
在深度强化学习中，价值网络通常由神经网络来实现。这个神经网络的输入是状态，输出是该状态的价值。
价值网络有两种形式：
1. 状态价值函数（V-function）：V(s)表示在状态s下，按照当前策略采取行动能获得的预期回报。
2. 动作价值函数（Q-function）：Q(s, a)表示在状态s下，采取动作a后能获得的预期回报。
在强化学习的训练过程中，我们通常会使用一种叫做值迭代（Value Iteration）的方法来更新价值网络，使得它能更准确地估计状态或动作的价值。这个过程通常涉及到一种叫做贝尔曼方程（Bellman Equation）的递归公式。
【优势函数】
势函数（Advantage Function）在强化学习中是一个非常重要的概念，它用于衡量在某个状态下，采取某个动作相比于按照当前策略采取动作的优势有多大。
优势函数A(s, a)的定义如下：
A(s, a) = Q(s, a) - V(s)
其中，Q(s, a)是动作价值函数，表示在状态s下采取动作a后能获得的预期回报；V(s)是状态价值函数，表示在状态s下按照当前策略π采取动作能获得的预期回报。
优势函数的值如果为正，表示采取动作a比按照当前策略的预期回报要高，反之则低。因此，优势函数可以用来指导策略的更新，即倾向于增大优势函数为正的动作的概率，减小优势函数为负的动作的概率。
【常用的优势函数】
在强化学习中，优势函数（Advantage Function）是一个重要的概念，它衡量在某个状态下，采取某个动作相比于按照当前策略采取动作的优势有多大。优势函数的计算方法有很多种，以下是一些常用的优势函数计算方法：
1. 基础优势函数：最基础的优势函数定义是 A(s, a) = Q(s, a) - V(s)，其中 Q(s, a) 是动作价值函数，V(s) 是状态价值函数。
2. 蒙特卡洛（Monte Carlo）估计：这种方法直接使用实际的回报减去状态价值函数来估计优势函数，即 A(s, a) = G_t - V(s)，其中 G_t 是从状态 s 开始的实际回报。
3. 时序差分（Temporal Difference）估计：这种方法使用一步的预期回报减去状态价值函数来估计优势函数，即 A(s, a) = r + γV(s') - V(s)，其中 r 是即时奖励，γ 是折扣因子，s' 是下一个状态。
4. 广义优势估计（Generalized Advantage Estimation，GAE）：这种方法是一种折衷的优势函数估计方法，它通过一个参数 λ 来控制蒙特卡洛估计和时序差分估计的权重，以达到方差和偏差之间的平衡。
PPO（Proximal Policy Optimization）算法通常使用的是广义优势估计（Generalized Advantage Estimation，GAE）作为优势函数。
GAE是一种平衡偏差和方差的优势函数估计方法。它通过引入一个衰减因子λ，结合了蒙特卡洛（Monte Carlo）方法和时序差分（Temporal Difference）方法的优点。λ的值在0和1之间，当λ为0时，GAE等同于时序差分估计；当λ为1时，GAE等同于蒙特卡洛估计。
GAE的计算公式如下：
A^GAE_t = δ_t + (γλ)δ_{t+1} + (γλ)^2 δ_{t+2} + ... + (γλ)^{T-t+1} δ_{T-1}
其中，δ_t = r_t + γV(s_{t+1}) - V(s_t)，r_t是即时奖励，γ是折扣因子，V(s)是状态价值函数。
在PPO算法中，使用GAE作为优势函数可以有效地减小估计的方差，提高学习的稳定性。
PPO算法使用示例：https://github.com/ericyangyu/PPO-for-Beginners
注意gym和pyglet的版本，最新的无法运行！
