大模型强化学习——PPO项目实战

发布时间 2023-09-30 09:53:15作者: bonelee

【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.

 

【核心代码解读】  

# -*- coding:utf-8 -*-
# @project: ChatGPT
# @filename: train
# @author: 刘聪NLP
# @zhihu: https://www.zhihu.com/people/LiuCongNLP
# @contact: logcongcong@gmail.com
# @time: 2023/4/14 17:40
"""
    文件说明:
            
"""

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的版本,最新的无法运行!