人类反馈强化学习奖励模型微调和RLHF微调

一只黑鸟

1358人浏览 · 2025-12-14 22:31:25

一只黑鸟 · 2025-12-14 22:31:25 发布

人类反馈强化学习奖励模型微调和RLHF微调

一、奖励模型训练核心流程

奖励模型训练的核心是对**优选回答（chosen）和拒绝回答（rejected）**进行成对排序，让模型学习给更优回答更高的分数。

1. 环境与依赖导入

# 导入LLaMA模型和分词器
from transformers import LlamaModel, LlamaTokenizer
# 导入PyTorch数据加载工具
from torch.utils.data import DataLoader, RandomSampler, SequentialSampler
# 导入自定义数据处理工具（来自dschat）
from dschat.utils.data.data_utils import create_prompt_dataset, DataCollatorReward
# 导入PyTorch核心模块
from torch import nn
import torch

2. 模型与分词器初始化

# 模型路径（以Chinese-Llama-2-7b为例）
args_model_name_or_path = '/mnt/chinese-llama-2-7b'
# 初始化分词器
tokenizer = LlamaTokenizer.from_pretrained(args_model_name_or_path)
# 初始化奖励模型的基础模型（LLaMA）
critic_model = LlamaModel.from_pretrained(args_model_name_or_path, device_map='auto')
# 设置模型为评估模式（后续训练时切换为train）
critic_model.eval()
# 输出文件夹
output_folder = '/tmp/'

3. 训练样本组织

# 训练阶段标记（2代表奖励模型训练）
train_phase = 2
# 创建prompt数据集，划分训练集和验证集
train_dataset, eval_dataset = create_prompt_dataset(
    args.local_rank,  # 分布式训练rank
    args.data_path,   # 数据路径
    args.data_split,  # 数据划分比例
    args.local_rank,
    args.data_output_path,
    train_phase,
    args.seed,
    tokenizer,
    args.max_seq_len  # 最大序列长度
)

4. 样本信息展示

用于验证数据加载是否正确，查看 chosen 和 rejected 样本的编码与内容：

# 遍历训练集，提取单条样本的chosen和rejected数据
for chosen_input_ids, chosen_attention_mask, reject_input_ids, reject_attention_mask in train_dataset:
    # 打印优选样本输入
    print("优选数据样本输入(chosen_input)-------------------")
    # 解码token_id为文本，跳过特殊token
    inputs = tokenizer.decode(chosen_input_ids[0], skip_special_tokens=True)
    print(f"原文: {inputs}")
    print(f"input_ids: {chosen_input_ids}")
    print(f"attention_mask: {chosen_attention_mask}")

    # 打印拒绝样本输入
    print("拒绝数据样本输入(reject_input)-------------------")
    print(f"reject_input_ids.shape: {reject_input_ids.shape}")
    inputs = tokenizer.decode(reject_input_ids[0], skip_special_tokens=True)
    print(f"原文: {inputs}")
    break  # 仅展示第一条样本

5. 奖励模型损失计算（forward功能模拟）

核心是计算成对排序损失，让优选样本的奖励值高于拒绝样本：

# 初始化数据整理器
data_collator = DataCollatorReward()
# 创建数据加载器
train_dataloader = DataLoader(
    train_dataset,
    collate_fn=data_collator,
    batch_size=args.per_device_train_batch_size  # 单卡批次大小
)

loss = 0
chosen_mean_scores = []  # 优选样本的平均奖励
rejected_mean_scores = []  # 拒绝样本的平均奖励

# 定义奖励头（将模型输出映射为标量奖励值）
v_head = nn.Linear(critic_model.config.hidden_size, 1, bias=False)

# 遍历训练批次
for step, batch in enumerate(train_dataloader):
    # 模型前向传播，获取输出
    transformer_outputs = critic_model(**batch)
    hidden_states = transformer_outputs[0]  # 获取模型最后一层隐藏状态
    rewards = v_head(hidden_states).squeeze(-1)  # 计算奖励值，压缩最后一维

    # 分离优选和拒绝样本的奖励
    chosen_reward = rewards[0]  # 优选样本奖励
    rejected_reward = rewards[1]  # 拒绝样本奖励

    # 定义PAD token id（以eos_token为例）
    PAD_ID = tokenizer.eos_token_id
    compute_fp32_loss = True
    seq_len = len(chosen_input_ids)

    # 找到填充标记的索引，确定有效序列长度
    c_inds = (chosen_input_ids == PAD_ID).nonzero()  # 找出PAD标记的位置
    c_ind = c_inds[0].item() if len(c_inds) > 0 else seq_len  # 第一个PAD的位置，无则取序列长度

    # -------- 关键：计算样本差异位置（divergence_ind） --------
    # 找到chosen和rejected序列第一次不同的位置
    check_divergence = (chosen_input_ids != reject_input_ids).nonzero()
    divergence_ind = check_divergence[0].item() if len(check_divergence) > 0 else seq_len

    # 截断奖励：只计算差异位置前的奖励，减少噪声
    end_ind = max(c_ind, divergence_ind)
    c_truncated_reward = chosen_reward[:end_ind]
    r_truncated_reward = rejected_reward[:end_ind]

    # 计算损失：使用logsigmoid，让chosen奖励 > rejected奖励
    loss += -torch.nn.functional.logsigmoid(c_truncated_reward - r_truncated_reward).mean()

    # 保存平均奖励值
    chosen_mean_scores.append(c_truncated_reward.mean().item())
    rejected_mean_scores.append(r_truncated_reward.mean().item())

# 打印损失和平均奖励
print(f"总损失: {loss.item()}")
print(f"优选样本平均奖励: {sum(chosen_mean_scores)/len(chosen_mean_scores)}")
print(f"拒绝样本平均奖励: {sum(rejected_mean_scores)/len(rejected_mean_scores)}")

注意：

v_head ：将LLaMA的隐藏状态映射为标量奖励值，是奖励模型的核心组件。
divergence_ind ：找到优选和拒绝样本的第一个差异位置，仅计算该位置前的奖励，减少无关噪声。
logsigmoid 损失：通过 -logsigmoid(c - r) 让模型学习让 c > r ，实现成对排序。

二、RLHF微调核心概念

RLHF微调阶段采用演员-评论家（Actor-Critic）算法，涉及4个核心模型：

actor_model：生成模型（演员），负责生成回答。
critic_model：评论模型（评论家），评估actor生成内容的好坏并给出奖励。
ref_model：参考模型，作为基准评估actor的改进效果。
reward_model：奖励模型，为critic提供奖励信号。

1. RLHF运行环境准备

import argparse

# 解析命令行参数
args = argparse.ArgumentParser()
args.add_argument('--local_rank', type=int, default=0)
args.add_argument('--data_path', type=str, default='/mnt/datasets/rm-static')
args.add_argument('--data_split', type=str, default='2,4,4')
args.add_argument('--data_output_path', type=str, default='/tmp/data_files')
args.add_argument('--max_seq_len', type=int, default=512)
args.add_argument('--seed', type=int, default=123)
args.add_argument('--num_train_epochs', type=int, default=1)
args.add_argument('--per_device_train_batch_size', type=int, default=1)
args = args.parse_args()

2. 关键训练参数

模型权重：使用 chinese-llama-2-1.3B （30GB CPU/16GB GPU即可运行）。
训练循环：包含训练轮次（train_epochs）、PPO轮次（ppo_epochs）、经验批次（generation_batches），参数均设为1（推荐配置）。
数据来源：采用 Duhas/rm-static 数据集，包含 prompt 、 response 、 chosen 、 rejected 字段。

三、关键函数与公式

1. logsigmoid函数：

用于损失计算，公式为：
$\text{logsigmoid}(x) = \log(\text{sigmoid}(x)) = \log\left(\frac{1}{1+e^{-x}}\right)$
作用是将奖励差值映射到对数概率空间，实现成对排序的损失优化。
2. 奖励截断逻辑：
通过 divergence_ind 找到样本差异位置，仅计算该位置前的奖励，减少无效计算和噪声干扰。

四、RLHF训练数据处理与Prompt解析

RLHF微调采用 Duhas/rm-static 数据集，包含 prompt 、 response 、 chosen 、 rejected 字段，核心是对prompt进行编码和解码，提取模型训练所需的输入特征。

1. 数据加载与Prompt编码展示

from dschat.utils.data.data_utils import create_prompt_dataset
from transformers import LlamaTokenizer

# 初始化分词器
tokenizer = LlamaTokenizer.from_pretrained(args.model_name_or_path)
# 创建prompt数据集（仅使用prompt字段）
train_dataset, _ = create_prompt_dataset(
    args.local_rank,
    args.data_path,
    args.data_split,
    args.local_rank,
    args.data_output_path,
    train_phase=3,  # RLHF微调阶段标记
    args.seed,
    tokenizer,
    args.max_seq_len
)

# 遍历数据集，展示prompt的ids和attention_mask
for step, batch_prompt in enumerate(train_dataset):
    print(f"批次prompt: {batch_prompt}")
    print(f"prompt_attention_mask: {batch_prompt['prompt_attention_mask']}")
    break  # 仅展示第一条样本

其中：

train_phase=3 ：标记为RLHF微调阶段，数据集仅加载 prompt 相关字段。
prompt_attention_mask ：用于标识prompt中有效token的位置，避免模型关注padding部分。

2. Prompt填充与序列生成

RLHF中需将prompt与演员模型生成的回答拼接，填充符号分布在两侧，核心代码实现序列生成与标记反转：

import torch
from dschat.utils.data import DataCollatorRLHF
from torch.utils.data import DataLoader

# 初始化数据整理器
data_collator = DataCollatorRLHF(args.max_prompt_seq_len)
# 创建训练数据加载器
prompt_train_dataloader = DataLoader(
    train_dataset,
    collate_fn=data_collator,
    batch_size=args.per_device_generation_batch_size
)

# 生成填充的提示语句，反转标记顺序
pad_token_id = tokenizer.pad_token_id
for step, batch_prompt in enumerate(prompt_train_dataloader):
    prompts = batch_prompt['prompt']
    mask = batch_prompt['prompt_at_mask']
    prompt_length = _generate_sequence(prompts, mask, step)  # 生成序列函数
    out_seq = torch.cat([prompts, torch.tensor([[pad_token_id]] * len(prompts))], dim=1)  # 拼接填充token
    print(f"生成序列长度: {prompt_length}")
    print(f"拼接后序列: {out_seq.shape}")
    break

# 序列生成辅助函数
def _generate_sequence(prompts, mask, step):
    """根据prompt和mask生成序列长度"""
    prompt_len = prompts.shape[1]
    mask_len = mask.shape[1]
    return max(prompt_len, mask_len) + step  # 动态调整序列长度

注意：

DataCollatorRLHF ：专为RLHF设计的数据整理器，处理prompt的padding和mask。
_generate_sequence ：根据prompt和mask的长度动态生成序列，适配模型输入要求。

五、对数概率计算（Actor/Ref模型）

RLHF中需要计算**演员模型（actor_model）和参考模型（ref_model）**的对数概率，用于后续损失计算与策略优化。

1. 模型初始化

from transformers import LlamaForCausalLM, RewardModel

# 初始化演员模型（生成模型）
actor_model = LlamaForCausalLM.from_pretrained(args.actor_model_name_or_path, device_map='auto')
# 初始化参考模型（基准模型）
ref_model = LlamaForCausalLM.from_pretrained(args.ref_model_name_or_path, device_map='auto')
# 初始化奖励模型（用于评估生成结果）
reward_model = RewardModel.from_pretrained(args.reward_model_name_or_path, device_map='auto')

# 设置模型模式
actor_model.train()
ref_model.eval()  # 参考模型仅用于推理，不训练

2. 对数概率计算核心代码

import torch.nn.functional as F

def compute_log_probs(logits, labels, pad_token_id):
    """计算序列的对数概率"""
    # 移除最后一个token的logits（无后续标签）
    logits = logits[:, :-1, :]
    # 移除第一个token的labels（无前置logits）
    labels = labels[:, 1:]
    
    # 将logits转为对数概率
    log_probs = F.log_softmax(logits, dim=-1)
    # 根据labels提取对应位置的对数概率
    gather_log_probs = log_probs.gather(dim=-1, index=labels.unsqueeze(-1)).squeeze(-1)
    
    # 屏蔽padding部分的对数概率
    mask = (labels != pad_token_id).long()
    log_probs = gather_log_probs * mask
    
    return log_probs

# 推理阶段计算logits
inputs = batch_prompt['input_ids'].to(device)
attention_mask = batch_prompt['attention_mask'].to(device)

# 演员模型前向传播
actor_logits = actor_model(inputs, attention_mask=attention_mask).logits
# 参考模型前向传播
ref_logits = ref_model(inputs, attention_mask=attention_mask).logits

# 计算对数概率
actor_log_probs = compute_log_probs(actor_logits, inputs, pad_token_id)
ref_log_probs = compute_log_probs(ref_logits, inputs, pad_token_id)

print(f"演员模型对数概率形状: {actor_log_probs.shape}")
print(f"参考模型对数概率形状: {ref_log_probs.shape}")

值得注意的是：

F.log_softmax ：将logits归一化为对数概率分布，避免数值溢出。
torch.gather ：根据标签提取对应token的对数概率，仅保留有效token的概率值。
mask ：屏蔽padding部分的对数概率，确保损失计算仅针对有效序列。

六、期望奖励计算

奖励模型通过 forward_value 函数输出标量奖励，用于评估演员模型生成序列的质量，核心是将隐藏状态映射为奖励值。

奖励模型前向传播

def compute_reward_score(reward_model, inputs, attention_mask, pad_token_id):
    """计算奖励分数"""
    # 奖励模型前向传播，获取隐藏状态
    outputs = reward_model(inputs, attention_mask=attention_mask)
    hidden_states = outputs.hidden_states[-1]  # 取最后一层隐藏状态
    
    # 奖励头：将隐藏状态映射为标量奖励
    v_head = torch.nn.Linear(reward_model.config.hidden_size, 1, bias=False).to(device)
    rewards = v_head(hidden_states).squeeze(-1)
    
    # 截断padding部分的奖励
    mask = (inputs != pad_token_id).long()
    truncated_rewards = rewards * mask
    
    # 计算平均奖励（有效token）
    reward_score = truncated_rewards.sum(dim=1) / mask.sum(dim=1)
    return reward_score

# 计算演员模型生成序列的奖励
reward_score = compute_reward_score(reward_model, inputs, attention_mask, pad_token_id)
print(f"序列奖励分数: {reward_score}")

七、KL散度计算

KL散度用于衡量演员模型与参考模型的概率分布差异，避免演员模型生成的内容偏离基准过多，公式为：
$DKL(P∣∣Q)=∑iP(i)log⁡(P(i)Q(i))D_{KL}(P||Q) = \sum_{i} P(i)\log\left(\frac{P(i)}{Q(i)}\right)$
其中P为演员模型的概率分布，Q为参考模型的概率分布。

1. KL散度计算代码

def compute_kl_divergence(actor_log_probs, ref_log_probs, mask):
    """计算KL散度（仅针对有效token）"""
    # 计算KL散度：P*log(P/Q) = actor_log_probs - ref_log_probs
    kl_div = (actor_log_probs - ref_log_probs) * mask
    # 计算平均KL散度
    kl_div_mean = kl_div.sum(dim=1) / mask.sum(dim=1)
    return kl_div_mean

# 计算KL散度
kl_div = compute_kl_divergence(actor_log_probs, ref_log_probs, mask)
print(f"KL散度: {kl_div}")

注意喵：

KL散度仅计算有效token的差异，避免padding部分干扰结果。
若KL散度过大，说明演员模型与参考模型差异过大，需通过损失函数约束。

八、RLHF损失函数整合

RLHF的最终损失为奖励损失与KL散度损失的加权和
其中\beta为KL散度的权重系数，用于平衡奖励优化与策略稳定性。

beta = 0.01  # KL散度权重
total_loss = -reward_score.mean() + beta * kl_div.mean()
print(f"RLHF总损失: {total_loss.item()}")

# 反向传播与优化
optimizer = torch.optim.AdamW(actor_model.parameters(), lr=1e-5)
optimizer.zero_grad()
total_loss.backward()
optimizer.step()

九、优势函数（Advantage Function）

优势函数是强化学习中衡量“动作价值与状态价值差异”的核心指标，公式为： $A (s, a) = Q (s, a) - V (s)$
其中Q(s,a)为动作价值函数，V(s)为状态价值函数。在RLHF中，优势函数用于评估演员模型生成动作的优劣，指导策略优化。

1. GAE（广义优势估计）计算

为平衡偏差与方差，RLHF采用**GAE（Generalized Advantage Estimation）**计算优势，核心代码如下：

import torch

def compute_gae(rewards, values, gamma=0.9, lam=0.95):
    """
    计算广义优势估计（GAE）
    :param rewards: 奖励序列 [seq_len]
    :param values: 评论家模型预测的价值序列 [seq_len]
    :param gamma: 折扣因子
    :param lam: 衰减因子
    :return: advantages: 优势序列, returns: 回报序列
    """
    advantages = []
    last_gae = 0
    # 从后向前计算优势
    for t in reversed(range(len(rewards))):
        # 计算TD误差：r_t + γ*V(t+1) - V(t)
        delta = rewards[t] + gamma * (values[t+1] if t+1 < len(values) else 0) - values[t]
        # 累积优势：delta + γ*λ*last_gae
        last_gae = delta + gamma * lam * last_gae
        advantages.insert(0, last_gae)
    
    # 计算回报（Return）：优势 + 价值
    returns = torch.tensor(advantages) + torch.tensor(values[:-1])
    # 标准化优势（可选，提升训练稳定性）
    advantages = (torch.tensor(advantages) - torch.mean(torch.tensor(advantages))) / (torch.std(torch.tensor(advantages)) + 1e-8)
    
    return advantages, returns

# 示例调用
rewards = [1.2, 0.8, -0.3]  # 奖励序列
values = [0.5, 0.7, 0.9, 0.6]  # 价值序列（长度比奖励多1）
advantages, returns = compute_gae(rewards, values)
print(f"优势序列: {advantages}")
print(f"回报序列: {returns}")

注意：

gamma ：折扣因子，控制未来奖励的权重，通常取0.9~0.99。
lam ：衰减因子，平衡偏差与方差，通常取0.95。
反向遍历计算优势：从最后一步向前累积，确保每一步的优势包含后续奖励的影响。

十、演员模型（Actor）损失计算

演员模型的损失核心是策略梯度损失，结合KL散度约束与奖励截断，避免模型更新幅度过大。

1. 演员损失核心代码

def actor_loss_fn(logprobs, old_logprobs, advantages, mask, cliprange=0.2):
    """
    计算演员模型损失（PPO裁剪）
    :param logprobs: 当前策略的对数概率 [batch, seq_len]
    :param old_logprobs: 旧策略的对数概率 [batch, seq_len]
    :param advantages: 优势序列 [batch, seq_len]
    :param mask: 有效token掩码 [batch, seq_len]
    :param cliprange: PPO裁剪范围
    :return: actor_loss: 演员损失
    """
    # 计算概率比：exp(new_logprob - old_logprob)
    ratio = torch.exp(logprobs - old_logprobs)
    # 未裁剪的损失：优势 * 概率比
    pg_loss1 = advantages * ratio
    # 裁剪后的损失：优势 * clip(ratio, 1-cliprange, 1+cliprange)
    pg_loss2 = advantages * torch.clamp(ratio, 1 - cliprange, 1 + cliprange)
    # 取两者最小值（PPO裁剪），避免梯度爆炸
    pg_loss = -torch.min(pg_loss1, pg_loss2) * mask
    # 求平均损失（仅计算有效token）
    pg_loss = pg_loss.sum() / mask.sum()
    
    return pg_loss

# 示例调用
cliprange = 0.2
actor_loss = actor_loss_fn(actor_log_probs, old_actor_log_probs, advantages, mask, cliprange)
print(f"演员模型损失: {actor_loss.item()}")

其中：

ratio ：当前策略与旧策略的概率比，衡量策略更新幅度。
torch.clamp ：裁剪概率比，限制策略更新的最大幅度（通常设为0.2）。
损失取最小值：确保策略更新不会过于激进，提升训练稳定性。

十一、评论家模型（Critic）损失计算

评论家模型的损失是价值损失，通过预测状态价值与实际回报的均方误差衡量，同时加入裁剪避免价值估计波动过大。

1. 评论家损失核心代码

def critic_loss_fn(values, old_values, returns, mask, cliprange_value=0.2):
    """
    计算评论家模型损失（价值损失+裁剪）
    :param values: 当前价值预测 [batch, seq_len]
    :param old_values: 旧价值预测 [batch, seq_len]
    :param returns: 实际回报 [batch, seq_len]
    :param mask: 有效token掩码 [batch, seq_len]
    :param cliprange_value: 价值裁剪范围
    :return: critic_loss: 评论家损失
    """
    # 裁剪价值预测，限制更新幅度
    values_clipped = torch.clamp(values, old_values - cliprange_value, old_values + cliprange_value)
    # 未裁剪的价值损失：(value - return)^2
    vf_loss1 = (values - returns) ** 2
    # 裁剪后的价值损失：(clipped_value - return)^2
    vf_loss2 = (values_clipped - returns) ** 2
    # 取两者最大值，确保价值估计不会偏离实际回报过远
    vf_loss = 0.5 * torch.max(vf_loss1, vf_loss2) * mask
    # 求平均损失（仅计算有效token）
    vf_loss = vf_loss.sum() / mask.sum()
    
    return vf_loss

# 示例调用
cliprange_value = 0.2
critic_loss = critic_loss_fn(critic_values, old_critic_values, returns, mask, cliprange_value)
print(f"评论家模型损失: {critic_loss.item()}")

十二、RLHF总损失与训练整合

RLHF的总损失为演员损失、评论家损失与KL散度损失的加权和，核心代码如下：

def compute_total_loss(actor_loss, critic_loss, kl_div, kl_coeff=0.01):
    """
    计算RLHF总损失
    :param actor_loss: 演员损失
    :param critic_loss: 评论家损失
    :param kl_div: KL散度（演员与参考模型的差异）
    :param kl_coeff: KL散度权重
    :return: total_loss: 总损失
    """
    total_loss = actor_loss + 0.5 * critic_loss + kl_coeff * kl_div
    return total_loss

# 示例调用
kl_div = compute_kl_divergence(actor_log_probs, ref_log_probs, mask)  # 前文KL散度计算函数
total_loss = compute_total_loss(actor_loss, critic_loss, kl_div)
print(f"RLHF总损失: {total_loss.item()}")

# 反向传播与优化
optimizer = torch.optim.AdamW(list(actor_model.parameters()) + list(critic_model.parameters()), lr=1e-5)
optimizer.zero_grad()
total_loss.backward()
optimizer.step()

其中：

kl_coeff ：KL散度权重，平衡策略探索与稳定性，通常取0.01~0.1。
联合优化：同时更新演员和评论家模型的参数，实现策略与价值的协同学习。

十三、奖励截断与KL散度加权优化

为避免奖励值波动过大，需对奖励进行截断；同时将KL散度估计值加入奖励，约束演员模型与参考模型的差异：

def compute_reward_with_kl(reward_score, kl_div, clip_reward_value=5, kl_coeff=0.01):
    """
    计算带KL散度加权的截断奖励
    :param reward_score: 原始奖励分数
    :param kl_div: KL散度
    :param clip_reward_value: 奖励截断范围
    :param kl_coeff: KL散度权重
    :return: truncated_reward: 截断并加权后的奖励
    """
    # 奖励截断：限制在[-clip_reward_value, clip_reward_value]
    truncated_reward = torch.clamp(reward_score, -clip_reward_value, clip_reward_value)
    # 加入KL散度惩罚：奖励 = 原始奖励 - KL散度*权重
    truncated_reward = truncated_reward - kl_coeff * kl_div
    return truncated_reward

# 示例调用
truncated_reward = compute_reward_with_kl(reward_score, kl_div)
print(f"截断并加权后的奖励: {truncated_reward}")