今天给大家来带好友知乎@ybq关于OpenRLHF的学习笔记，主要介绍其中的各种loss内容。
从这篇文章开始，我会不定期分享利用 OpenRLHF 学习 RLHF 的一些心得。我平常读代码喜欢开门见山，直接去看 loss 函数是什么形式，再去理解代码的其他环节，所以就从 loss 开始分享吧。
代码详见：https://github.com/OpenRLHF/OpenRLHF/blob/main/openrlhf/models/loss.py
基础
在研究 loss 函数前，建议把下面几个公式和图先焊死在脑子中。
SFT 家族
GPTLMLoss
class GPTLMLoss(nn.Module):
"""
GPT Language Model Loss
"""
def __init__(self):
super().__init__()
self.IGNORE_INDEX = -100
self.loss = nn.CrossEntropyLoss(ignore_index=self.IGNORE_INDEX)
def forward(self, logits: torch.Tensor, labels: torch.Tensor) -> torch.Tensor:
shift_logits = logits[..., :-1, :].contiguous()
shift_labels = labels[..., 1:].contiguous()
# Flatten the tokens
return self.loss(shift_logits.view(-1, shift_logits.size(-1)), shift_labels.view(-1))
没啥多说的，最常见的 gpt loss 函数，也就是 pretrain / sft 的 loss 函数，通过 self.IGNORE_INDEX 来实现 prompt 的 loss_mask 。
KDLoss
# Adapted from https://github.com/microsoft/LMOps/blob/main/minillm/finetune.py#L166
class KDLoss(nn.Module):
"""
Language Model Knowledge Distillation Loss
"""
def __init__(self):
super().__init__()
self.IGNORE_INDEX = -100
def forward(self, logits: torch.Tensor, teacher_logits: torch.Tensor, label: torch.Tensor) -> torch.Tensor:
teacher_probs = F.softmax(teacher_logits, dim=-1, dtype=torch.float32)
inf_mask = torch.isinf(logits)
logprobs = F.log_softmax(logits, dim=-1, dtype=torch.float32)
prod_probs = torch.masked_fill(teacher_probs * logprobs, inf_mask, 0)
x = torch.sum(prod_probs, dim=-1).view(-1)
mask = (label != self.IGNORE_INDEX).int()
distil_loss = -torch.sum(x * mask.view(-1), dim=0) / torch.sum(mask.view(-1), dim=0)
return distil_loss
第二种 sft 的 loss 函数：知识蒸馏的 loss 函数。需要在同源 tokenizer 的情况下，利用一个大模型的 logits 分布结果，来让小模型学习软标签。当然，embedding 毕竟只是一个线性层，可以考虑再给模型外挂一个线性层，把 model_A 的 tokenizer 映射到 model_B 的 tokenizer，进而实现利用 qwen 蒸馏 llama 的美好愿景，不知道有没有大佬做过类似的尝试。
言归正传，我们都知道知识蒸馏是用 KL 散度作为 loss 函数的，但代码里也没看见 KL 散度公式啊，不妨一起简单推导下。
其中， 是教师模型的概率分布， 是学生模型的概率分布。在实际的优化过程中，KL 散度中的第一项 是关于教师模型的熵，对于学生模型的参数优化是一个常数，可丢弃。因此，我们通常只需要最小化第二项，即交叉熵损失：
知识蒸馏本身没啥痛点，只要能解决 seq_len * vocab_size 大小的 logits 通讯问题，这就是个简单纯粹有效的优化小模型的极佳方案。不过传统的 KL 往往是 soft_label 和 hard_label 的加权组合，这在 OpenRLHF 的代码中没有体现出来，大家有需要的话可以自行实践：
lm_loss = F.cross_entropy(
logits.view(-1, logits.size(-1)),
label.view(-1),
ignore_index=self.IGNORE_INDEX
)
total_loss = alpha * lm_loss + beta * distil_loss
DPO 家族
DPOLoss
class DPOLoss(nn.Module):
"""
DPO Loss
"""
def __init__(self, beta: float, label_smoothing: float = 0.0, ipo: bool = False) -> None:
super().__init__()
self.beta = beta
self.label_smoothing = label_smoothing
self.ipo = ipo
def forward(
self,
policy_chosen_logps: torch.Tensor,
policy_rejected_logps: torch.Tensor,
reference_chosen_logps: torch.Tensor,
reference_rejected_logps: torch.Tensor,
) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
pi_logratios = policy_chosen_logps - policy_rejected_logps
ref_logratios = reference_chosen_logps - reference_rejected_logps
logits = pi_logratios - ref_logratios
if self.ipo:
losses = (logits - 1 / (2 * self.beta)) ** 2 # Eq. 17 of https://arxiv.org/pdf/2310.12036v2.pdf
else:
# Eq. 3 https://ericmitchell.ai/cdpo.pdf; label_smoothing=0 gives original DPO (Eq. 7 of https://arxiv.org/pdf/2305.18290.pdf)
losses = (
-F.logsigmoid(self.beta * logits) * (1 - self.label_smoothing)
- F.logsigmoid(-self.beta * logits) * self.label_smoothing
)
loss = losses.mean()
chosen_rewards = self.beta * (policy_chosen_logps - reference_chosen_logps).detach()
rejected_rewards = self.beta * (policy_rejected_logps - reference_rejected_logps).detach()
return loss, chosen_rewards, rejected_rewards
我们熟悉的 dpo 的 loss 函数，看上去没提供任何 trick，实践中如果 chosen_rewards 和 rejected_rewards 都下降，可以考虑给正例 / 负例再加一个系数。
除了原始的 loss 函数，OpenRLHF 为我们提供了额外两个选项：
IPO：论文中的 loss 表达式
我没有实践过这个算法就不多评价了，似乎重点是加了一个正则项。
CDPO：大概就是给 DPO 加了个 label_smoothing。
标签平滑是一种正则化方法，它通过将硬标签转换为软标签来防止模型过度自信。具体来说，对于二分类问题：原本的样本是正例就是正例，是负例就是负例，平滑后变成了：(1 - self.label_smoothing) 的概率是正例，self.label_smoothing 的概率是负例。
具体在 DPO 算法中的含义，一个 pair 对，以 (1 - self.label_smoothing) 的概率认为 good_sentence 比 bad_sentence 质量高，以 self.label_smoothing 的概率认为 bad_sentence 比 good_sentence 质量高。从而避免了模型对训练数据的过度拟合和过度自信。
理解这个平滑代码实现的关键点在于下面这两个公式，相信负例的 loss 可以动手笔划一下：
相信正例的 loss： 相信负例的 loss：
KTOLoss
# Adapted from https://github.com/ContextualAI/HALOs/blob/ca9b7e3eeea220c0944ad8095d641da33f907a7e/trainers.py#L770
class KTOLoss(nn.Module):
"""
KTO loss for uneven sampling
"""
def __init__(
self, beta: float, desirable_weight: float, undesirable_weight: float, world_size: int, device: torch.device
) -> None:
super().__init__()
self.beta = beta
self.world_size = world_size
self.device = device
self.desirable_weight = desirable_weight
self.undesirable_weight = undesirable_weight
def forward(
self,
policy_chosen_logps: torch.FloatTensor,
policy_rejected_logps: torch.FloatTensor,
policy_KL_logps: torch.FloatTensor,
reference_chosen_logps: torch.FloatTensor,
reference_rejected_logps: torch.FloatTensor,
reference_KL_logps: torch.FloatTensor,
) -> Tuple[torch.FloatTensor, torch.FloatTensor, torch.FloatTensor]:
KL = (policy_KL_logps - reference_KL_logps).mean().detach()
# all_reduce sums up the KL estimates across all devices (gradient will also be scaled by world size)
dist.all_reduce(KL, op=dist.ReduceOp.SUM)
# take average (will also scale gradients appropriately)
KL = (KL / self.world_size).clamp(min=0)
if policy_chosen_logps.shape[0] != 0:
chosen_logratios = policy_chosen_logps - reference_chosen_logps
chosen_losses = 1 - F.sigmoid(self.beta * (chosen_logratios - KL))
chosen_rewards = self.beta * chosen_logratios.detach()
else:
# important to cast to policy_dtype; otherwise error will occur during all_gather
chosen_losses = torch.Tensor([]).to(policy_rejected_logps.dtype).to(self.device)
chosen_rewards = torch.Tensor([]).to(policy_rejected_logps.dtype).to(self.device)
if policy_rejected_logps.shape[0] != 0:
rejected_logratios = policy_rejected_logps - reference_rejected_logps
rejected_losses = 1 - F.sigmoid(self.beta * (KL - rejected_logratios))
rejected_rewards = self.beta * rejected_logratios.detach()
else:
# important to cast to policy_dtype; otherwise error will occur during all_gather
rejected_losses = torch.Tensor([]).to(policy_chosen_logps.dtype).to(self.device)
rejected_rewards = torch.Tensor([]).to(policy_chosen_logps.dtype).to(self.device)
losses = torch.cat(
(self.desirable_weight * chosen_losses, self.undesirable_weight * rejected_losses), 0
).mean()
return losses, chosen_rewards, rejected_rewards, KL
kto 的 loss 函数，对标 dpo 的一个工作。这个算法我之前也没有实践过，但是“kto 不需要偏好 pair 对”这一优点吸引了我，所以也简单研究了一下它的原理和实现代码。
kto 的算法思想说是借鉴了“前景理论”，这个概念对我一个程序员来说太高深了，还是直接去讨论一下算法怎么实现的吧。
kto 的训练数据是 prompt + response + label，这个 label 就是 1 或者 -1，代表着 response 的质量是否被认可。label 是 1 的被称为正例，label 是 -1 的被称为负例。我们看到 loss 函数中要做一个判断 if policy_chosen_logps.shape[0] != 0 的操作，这是因为如果该条训练数据为负例，那么 policy_chosen_logps 这个变量就是一个空 tensor，反之亦然。和 dpo 相比最大的区别是：dpo 的每一条 prompt 需要同时具有正例和负例，kto 的每一条 prompt 则只需要有正例或负例中的一个即可。
kto 正例和负例的 loss 函数分别如下所示：
1 - sigmoid 是一个单调递减函数，这说明：kto 的 loss 函数在正例中鼓励策略模型尽量大于参考点 KL，在负例中则鼓励模型尽量小于参考点 KL，也是一个比较明显的学习正例打压负例的损失函数。self.desirable_weight 和 self.undesirable_weight 则是正向和负向样本各自的权重损失，调参用的。
kto 代码的理解难点是，这个 KL 并不是一条训练样本的 KL，而是一批样本的平均 KL（代码中的 dist.all_reduce），并且为了训练稳定这个 KL 也是不进行反向传播的（代码中的 detach），只是拿来控制损失的饱和度，并且做了 clamp(min=0) 处理。至于这么设计的原因，反正原论文就这么写的，我没具体看公式是怎么推的，不敢瞎分析，感兴趣的可以自己推推公式。
VanillaKTOLoss
# Adapted from https://github.com/ContextualAI/HALOs/blob/ca9b7e3eeea220c0944ad8095d641da33f907a7e/trainers.py#L742
class VanillaKTOLoss(nn.Module):
"""
KTO loss for even sampling
"""
def __init__(self, beta: float) -> None:
super().__init__()
self.beta = beta
def forward(
self,
policy_chosen_logps: torch.FloatTensor,
policy_rejected_logps: torch.FloatTensor,
reference_chosen_logps: torch.FloatTensor,
reference_rejected_logps: torch.FloatTensor,
) -> Tuple[torch.FloatTensor, torch.FloatTensor, torch.FloatTensor]:
chosen_KL = (policy_chosen_logps - reference_chosen_logps).mean().clamp(min=0)
rejected_KL = (policy_rejected_logps - reference_rejected_logps).mean().clamp(min=0)
chosen_logratios = policy_chosen_logps - reference_chosen_logps
rejected_logratios = policy_rejected_logps - reference_rejected_logps
losses = torch.cat(
(
1 - F.sigmoid(self.beta * (chosen_logratios - rejected_KL)),
1 - F.sigmoid(self.beta * (chosen_KL - rejected_logratios)),
),
0,
).mean()
chosen_rewards = self.beta * (policy_chosen_logps - reference_chosen_logps).detach()
rejected_rewards = self.beta * (policy_rejected_logps - reference_rejected_logps).detach()
return losses, chosen_rewards, rejected_rewards
kto 的变种，看代码实现的话，主要 diff 应该是去掉了参考点 KL ，并且正负样本要 1:1 均衡（OpenRLHF 代码库没写，但是注释里的 https://github.com/ContextualAI/HALOs 这个代码库写了）
乍一看，这个均匀采样 kto 的 loss 函数和 dpo 已经很相似了，但其实还是有本质区别的。dpo 的重点是 margin，也就是正例和负例的 loss 是要做减法的，均匀采样 kto 用的是 torch.cat()，也就是说正例和负例的 loss 相互之间毫无影响，各自朝着各自的 label 去优化。
需要留意的细节是，chosen_KL 和 rejected_KL 也做了 clamp(min=0) 的操作。这里给出我对 RLHF 代码的一条学习心得：不要放过任何一个 clamp / clip 操作背后的原因。
RLHF 家族
PolicyLoss
class PolicyLoss(nn.Module):
"""
Policy Loss for PPO
"""
def __init__(self, clip_eps: float = 0.2) -> None:
super().__init__()
self.clip_eps = clip_eps
def forward(
self,
log_probs: torch.Tensor,
old_log_probs: torch.Tensor,
advantages: torch.Tensor,
action_mask: Optional[torch.Tensor] = None,
) -> torch.Tensor:
ratio = (log_probs - old_log_probs).exp()
surr1 = ratio * advantages
surr2 = ratio.clamp(1 - self.clip_eps, 1 + self.clip_eps) * advantages
loss = -torch.min(surr1, surr2)
loss = masked_mean(loss, action_mask, dim=-1).mean()
return loss
rlhf 中，actor_model （也就是被优化的模型）的 loss 函数，大概是三个步骤：
计算新旧策略的概率比例
利用优势函数指导更新方向； 限制策略更新幅度 -torch.min(surr1, surr2)，选择未剪辑和剪辑后损失项的最小值。取负号是因为在最小化损失函数，但 PPO 的目标是最大化期望收益。
代码写的很清晰简洁，和 ppo 论文完全吻合，上面的两个公式也都是 ppo 论文的原始公式。对这里的代码实现有疑惑的，可以结合 ppo 论文一起读。
ValueLoss
class ValueLoss(nn.Module):
"""
Value Loss for PPO
"""
def __init__(self, clip_eps: float = None) -> None:
super().__init__()
self.clip_eps = clip_eps
def forward(
self,
values: torch.Tensor,
old_values: torch.Tensor,
returns: torch.Tensor,
action_mask: Optional[torch.Tensor] = None,
) -> torch.Tensor:
if self.clip_eps is not None:
values_clipped = old_values + (values - old_values).clamp(-self.clip_eps, self.clip_eps)
surr1 = (values_clipped - returns) ** 2
surr2 = (values - returns) ** 2
loss = torch.max(surr1, surr2)
else:
loss = (values - returns) ** 2
loss = masked_mean(loss, action_mask, dim=-1).mean()
return 0.5 * loss
rlhf 中，critic_model 的 loss 函数。
如果不剪辑损失函数： 如果要剪辑损失函数，便需要对价值函数的更新进行剪辑，防止价值估计发生过大的变化，和上面的 policy model 的剪辑是一个道理： 剪辑的损失函数： 。也就是说，如果有剪辑，通过 torch.max(surr1, surr2) 选择让 loss 最大化的更新策略； loss 最终乘以 0.5，也算是平方误差损失函数中常见的缩放因子了。
这里我之前被一个地方绊住过，可以分享一下我曾经的疑惑点：clamp 的意义既然是防止模型进行较大的参数更新，那为什么 value function 的 loss 还要选 torch.max() 呢，不应该是 torch.min() 更合理吗？
我目前的观点：在策略函数中，模型更新幅度的大与小，和 loss 的大小并无直接关系。新的 values 距离 old_values 越近，代表着价值估计的目标更新幅度越小。显然，old_values + (values - old_values).clamp(-self.clip_eps, self.clip_eps) 对应的 values_clipped ，是一定比原始的 values 更接近 old_values 的。
surr1，surr2 分别代表使用 values 和 values_clipped 进行模型更新的 loss：
surr1 > surr2：说明 clip 的过分了，导致 loss 变小可能会更新不动，那就放弃 clip，选择 values 来更新； surr1 < surr2：说明用更保守的更新策略 values_clipped，得到了更大的 loss。模型期望的更新幅度小，训练动力还大，没有比这更好的事情了。
PairWiseLoss
class PairWiseLoss(nn.Module):
"""
Pairwise Loss for Reward Model
"""
def forward(
self, chosen_reward: torch.Tensor, reject_reward: torch.Tensor, margin: torch.Tensor = None
) -> torch.Tensor:
if margin is not None:
loss = -F.logsigmoid(chosen_reward - reject_reward - margin)
else:
loss = -F.logsigmoid(chosen_reward - reject_reward)
return loss.mean()
主角登场，reward_model 的 loss 函数，以最简单的形式干最多的活！
这里有意思的点是 OpenRLHF 提供了一个 margin 的选项。还记得文章开头给大家画出来的 求导后的曲线吗？x 越小，梯度的绝对值越大，就越能避免梯度消失。原本 positive_reward - negative_reward = 2 的时候，已经没梯度训不动了，现在 positive_reward - negative_reward = 2 + margin 的时候才会训不动。
这个 margin 和 dpo 的 refernece_model 非常类似，都是常量。我曾经疑惑过这种常量是不是没啥大用，后来动手求了求导就明白了：这些被 logsigmoid() 包裹起来的常量，会影响梯度的大小，决定梯度在什么情况下趋近于零，进而也会影响模型训练的动力。
LogExpLoss
class LogExpLoss(nn.Module):
"""
Pairwise Loss for Reward Model
Details: https://arxiv.org/abs/2204.05862
"""
def forward(
self, chosen_reward: torch.Tensor, reject_reward: torch.Tensor, margin: torch.Tensor = None
) -> torch.Tensor:
loss = torch.log(1 + torch.exp(reject_reward - chosen_reward)).mean()
return loss
想不到吧，reward_model 也有变种。
看上去是用 取代了 ，这妥妥的就是一个等价变化啊。
PRMLoss
class PRMLoss(nn.Module):
"""
Process Reward Model Loss
"""
def __init__(self, placeholder_token_id: int, reward_token_ids: Optional[list[int]] = None):
super().__init__()
self.IGNORE_INDEX = -100
self.loss = nn.CrossEntropyLoss(ignore_index=self.IGNORE_INDEX)
self.placeholder_token_id = placeholder_token_id
self.reward_token_ids = reward_token_ids
def forward(self, inputs: torch.Tensor, logits: torch.Tensor, labels: torch.Tensor, *, return_acc: bool = False):
placeholder_mask = inputs == self.placeholder_token_id
logits = logits[placeholder_mask]
labels = labels[placeholder_mask]
if labels.dtype == torch.float:
# soft label
assert len(self.reward_token_ids) == 2, "reward_token_ids should have 2 tokens for soft labels"
logits = logits[..., self.reward_token_ids]
positive_labels = labels.to(logits.dtype)
negative_labels = 1 - positive_labels
negative_labels[positive_labels != -100] = 1 - positive_labels[positive_labels != -100]
labels = torch.stack([positive_labels, negative_labels], dim=-1)
elif self.reward_token_ids is not None:
# hard label with reward_token_ids set. (otherwise the whole vocab will be trained together.)
logits = logits[..., self.reward_token_ids]
# this is slow....
for i, token in enumerate(self.reward_token_ids):
labels = torch.where(labels == token, i, labels)
loss = self.loss(logits, labels)
if not return_acc:
return loss
if labels.dtype == logits.dtype:
labels = labels.argmax(dim=-1)
acc = (logits.argmax(dim=-1) == labels).float().mean()
return loss, acc
当红炸子鸡，学界认为的通往 o1 的钥匙，process_reward_model 的 loss 函数。出于对 o1 的尊重，我逐行解读一下这个 loss 。
首先看下 PRM 训练集合的样子：
{
"inputs": "Janet pays $40/hour for 3 hours per week of clarinet lessons and $28/hour for 5 hours a week of piano lessons. How much more does she spend on piano lessons than clarinet lessons in a year? Step 1: Janet spends 3 hours + 5 hours = <<3+5=8>>8 hours per week on music lessons. ки Step 2: She spends 40 * 3 = <<40*3=120>>120 on clarinet lessons per week. ки Step 3: She spends 28 * 5 = <<28*5=140>>140 on piano lessons per week. ки Step 4: Janet spends 120 + 140 = <<120+140=260>>260 on music lessons per week. ки Step 5: She spends 260 * 52 = <<260*52=13520>>13520 on music lessons in a year. The answer is: 13520 ки",
"labels": "Janet pays $40/hour for 3 hours per week of clarinet lessons and $28/hour for 5 hours a week of piano lessons. How much more does she spend on piano lessons than clarinet lessons in a year? Step 1: Janet spends 3 hours + 5 hours = <<3+5=8>>8 hours per week on music lessons. + Step 2: She spends 40 * 3 = <<40*3=120>>120 on clarinet lessons per week. + Step 3: She spends 28 * 5 = <<28*5=140>>140 on piano lessons per week. + Step 4: Janet spends 120 + 140 = <<120+140=260>>260 on music lessons per week. + Step 5: She spends 260 * 52 = <<260*52=13520>>13520 on music lessons in a year. The answer is: 13520 -",
"values": [ "+", "+", "+", "+", "-" ]
}
在 inputs 中，每个 step 后面会有一个 special_token：ки 在 labels 中，每个 step 后面会有一个 label_token：+ / - （代表着当前 step 的推理是否正确）
placeholder_mask = inputs == self.placeholder_token_id
logits = logits[placeholder_mask]
labels = labels[placeholder_mask]
logits 就是整个 inputs 过了一遍 llm 后得到的输出，形状为 seq_len * vocab_size （不考虑 batch_size），self.placeholder_token_id 就是 “ки” 对应的 id。使用这几行代码，上面的 case 中，logits 会变成 5 * vocab_size， label 会变成 5 * 1
logits = logits[..., self.reward_token_ids]
for i, token in enumerate(self.reward_token_ids):
labels = torch.where(labels == token, i, labels)
紧接着，先理解常规的 hard label，self.reward_token_ids 就是["+"对应的 id, "-"对应的 id]，labels 就是["+"对应的 id, "+"对应的 id, "+"对应的 id, "+"对应的 id , "-"对应的 id]。这几行代码成功提取出了每个 step 下，两个 label 各自对应的 logits，以及每个 step 的 label 是什么。
if labels.dtype == torch.float:
logits = logits[..., self.reward_token_ids]
assert len(self.reward_token_ids) == 2, "reward_token_ids should have 2 tokens for soft labels"
positive_labels = labels.to(logits.dtype)
negative_labels = 1 - positive_labels
negative_labels[positive_labels != -100] = 1 - positive_labels[positive_labels != -100]
labels = torch.stack([positive_labels, negative_labels], dim=-1)
再理解非常规的 soft_label，此时 labels 不再是 id，而是 float 类型，比如 labels = [0.8, 0.85, 0.9, 0.78, 0.1]，代表着每个 step 正确的概率（ assert len(self.reward_token_ids) == 2 是为了确保可以通过减法算出 step 错误的概率）。
理解了 soft_label 和 hard_label 分别是如何获得的之后，后面的 loss 计算和 acc 计算就没什么好多说的了。哦对，如果用 qwen 去跑这份代码，会遇到一个 tokenizer 的 bug，请教了下代码作者朱小霖大佬，大概是说不想为了 math-shepherd 加太多冗余逻辑，后续会给出一个优化版的代码（期待ing）
