GRPO 算法（损失函数）——原理讲解与代码讲解
视频讲解链接：8.calculating-loss-in-grpo.zh_en_哔哩哔哩_bilibili
论文：《DeepSeek-R1: Incentivizing Reasoning Capability in LLMs via Reinforcement Learning》
一、GRPO 损失函数
二、GRPO 算法可以分解为四个关键组成部分：
（1）策略损失（policy loss）：模型在有适配器和没有适配器情况下的词元概率分布比率
（2）从奖励函数中计算出来的优势值（advantages）
（3）比率裁剪（clip）：用于确保在任何单独步骤中都没有大的损失值
（4）KL散度：用于确保在训练过程中，我们正在训练的模型不会偏离它已经知道的基准模型太多
下面我们对每个部分逐一介绍。
1. 首先加载所需的模型和分词器，并打印模型的网络结构和生成文本的效果。
from transformers import AutoModelForCausalLM, AutoTokenizer # 初始化 model 和 tokenizer model_str = "babylm/babyllama-100m-2024" base_model = AutoModelForCausalLM.from_pretrained(model_str) tokenizer = AutoTokenizer.from_pretrained(model_str) # pad on the left so we can append new tokenizer on the right tokenizer.padding_side = "left" tokenizer.truncation_side = "left" print(base_model) import torch prompt = "The quick brown fox jumped over the " # Tokenizer the prompt input_ids = tokenizer(prompt, return_tensors="pt") print(input_ids) # 将文本进行分词，得到 token，在尾部添加了一个终止符 # Generate next 2 tokens with torch.no_grad(): outputs = base_model.generate( **input_ids, max_new_tokens=2, pad_token_id=tokenizer.pad_token_id ) # Decode the generated text 将 token 重新转换为文本 generated_text = tokenizer.decode( outputs[0], skip_special_tokens=True ) generated_portion = generated_text[len(prompt):] print(f"Generated text: {prompt}\033[94m{generated_portion}\033[0m")在 GRPO 中，使用两个不同的模型来指导学习过程，一个是参考模型（没有 LoRA 的基础模型，在整个训练过程中保持冻结），另一个是策略模型（要训练的模型，使用一组在整个学习过程中不断更新的 LORA 权重）。下面我们来分别定义参考模型（ref_model）和策略模型（model）
import copy from peft import LoraConfig, get_peft_model # Create a copy of the base model to use as the reference model（参考模型） ref_model = copy.deepcopy(base_model) # 初始化 LoRA 配置文件：用于我们想要添加到模型中的 LoRA 权重 lora_config = LoraConfig( r=8, # 秩 lora_alpha=32, target_modules=["q_proj", "v_proj"], # 插入权重的目标模块 init_lora_weights=False, bias="none", task_type="CAUSAL_LM" ) # Apply LoRA to model model = get_peft_model(base_model, lora_config) print(model)2. 有了参考模型和策略模型之后，我们可以实现策略函数（policy_loss）部分。
两个辅助函数：
（1）prepare_inputs：将文本 prompt 和 completion 转换为 tokenizer，并进行合并，方便作为模型的提示词输入。
（2）compute_log_probs：计算通过模型生成的每个 token 的对数概率。概率越大说明模型对生
成的 token 的准确信心越大。
def prepare_inputs(prompt, completion): # Tokenization prompt_tokens = tokenizer(prompt, return_tensors="pt") completion_tokens = tokenizer(completion, return_tensors="pt") # Combined input input_ids = torch.cat( [ prompt_tokens["input_ids"], completion_tokens["input_ids"] ], dim = 1 ) # 注意力掩码 attention_mask = torch.cat( [ prompt_tokens["attention_mask"], completion_tokens["attention_mask"] ], dim = 1 ) prompt_length = prompt_tokens["input_ids"].shape[1] completion_length = completion_tokens["input_ids"].shape[1] total_length = prompt_length + completion_length # 补全掩码：Create a mask to identify the tokens that were generated by the model in the full sequence completion_mask = torch.zeros(total_length, dtype=torch.float32) completion_mask[prompt_length:] = 1.0 return input_ids, attention_mask, completion_maskimport torch.nn.functional as F def compute_log_probs(model, input_ids, attention_mask): outputs = model(input_ids, attention_masks=attention_mask) # outputs.logits 是神经网络输出中未经过归一化的概率，下一步通常是 softmax log_probs = F.log_softmax(outputs.logits, dim=-1) return log_probs.gather( dim=-1, index=input_ids.unsqueeze(-1) ).squeeze(-1)策略损失函数的含义是策略模型对每个 token 的对数概率，与参考模型对每个 token 的对数概率的比率。这个比率（retio）越大，说明策略模型对生成的 token 的信心越大，说明策略模型生成的结果越好。
advantages 是通过奖励函数转换来的优势值，强化学习的目标就是获得最大的优势值。
在 GRPO 算法中根据优势值缩放比率，来判断策略模型信心大的生成 token 的方向是不是和优势值方向一致（举个例子：对于优势值为负的策略，尽管策略模型对它生成的 token 信心很大，但这只能说明策略模型生成文本的功能很好，但不是我们想要的方向）。
代码中的 completion_mask 是由模型生成的 token 的位置掩码（它是一个前面全是0，后面全是1的掩码），在计算策略损失时，我们只考虑由模型生成的 token 的损失。
def grpo_loss(model, ref_model, prompt, completion, advantage): input_ids, attention_mask, completion_mask = prepare_inputs(prompt, completion) # 策略模型对数概率 token_log_probs = compute_log_probs( model, input_ids, attention_mask ) # 参考模型对数概率 with torch.no_grad(): ref_token_log_probs = compute_log_probs( ref_model, input_ids, attention_mask ) # 这个比率（ratio）表示策略模型生成的 token 相比于参考模型，是具有更高的概率还是更低的概率 ratio = torch.exp(token_log_probs - ref_token_log_probs) # 根据优势值缩放比率，生成的 token 是否对模型学习方向有正向作用 policy_loss = ratio * advantage # We want to maximize reward, so we make the loss negative # because optimizers minimize loss per_token_loss = -policy_loss # 只考虑输出 tokens 的损失 loss = (per_token_loss * completion_mask).sum() / completion_mask.sum() return loss下面两个输出结果反应出：当参考模型和策略模型相同，比率是 1 ，policy_loss 等于 advantages。这也说明为什么奖励函数生成的奖励分数需要多样性，如果生成的奖励分数是固定的，那么得到优势值也都是0，导致损失为0，阻止训练过程。
grpo_loss(model, ref_model, prompt, "fence and", advantage=2.0) # tensor(-7.5770, grad_fn=<DivBackward0>) grpo_loss(ref_model, ref_model, prompt, "fence and", advantage=2.0) # tensor(-2., grad_fn=<DivBackward0>)3. 比率裁剪（clip）
policy_loss 是计算生成的每个 token 在策略模型和参考模型中的比率，但有时候某个 token 的比率会明显大于其他值，这是不可避免的，但这样会造成强化学习过程非常不稳定。所以我们需要一种方法来防止在任何单步训练步骤中产生过大的损失值。
实现这一点方法是比率裁剪（clip）：将比率控制在一个范围内，防止比率过大或过小，代码如下。
def grpo_loss_with_clip(model, ref_model, prompt, completion, advantage, epsilon = 0.2): input_ids, attention_mask, completion_mask = prepare_inputs(prompt, completion) # 策略模型对数概率 token_log_probs = compute_log_probs( model, input_ids, attention_mask ) # 参考模型对数概率 with torch.no_grad(): ref_token_log_probs = compute_log_probs( ref_model, input_ids, attention_mask ) # 这个比率（ratio）表示策略模型生成的 token 相比于参考模型，是具有更高的概率还是更低的概率 ratio = torch.exp(token_log_probs - ref_token_log_probs) # 根据优势值缩放比率，生成的 token 是否对模型学习方向有正向作用 unclipped = ratio * advantage # 裁剪比率：将比率控制在一个范围，防止比率过大或过小 clipped = torch.clamp(ratio, 1-epsilon, 1+epsilon) * advantage policy_loss = torch.min(unclipped, clipped) # We want to maximize reward, so we make the loss negative # because optimizers minimize loss per_token_loss = -policy_loss # 只考虑输出 tokens 的损失 loss = (per_token_loss * completion_mask).sum() / completion_mask.sum() return loss加入比率裁剪之后，损失值比之前的更靠近 -2.0 （ -2.0 表示策略模型与参考模型对生成 token 信心相同，这是我们希望的最终结果），说明损失稳定很多（由 -7.5770 → -2.4000）。
grpo_loss_with_clip(model, ref_model, prompt, "fence and", advantage=2.0, epsilon = 0.2) # tensor(-2.4000, grad_fn=<DivBackward0>)4. KL 散度（惩罚项）
强化学习中的一个常见问题是：随着更新策略模型，它在生成的词元上开始与参考模型有所偏差。因此我们通常引入一个称为 KL 散度的惩罚项，以防止策略模型偏离参考模型太远。KL 散度是一种衡量策略模型与参考模型之间的分布偏差程度的方法。
def grpo_loss_with_kl(model, ref_model, prompt, completion, advantage, epsilon = 0.2, beta = 0.1): input_ids, attention_mask, completion_mask = prepare_inputs(prompt, completion) # 策略模型对数概率 token_log_probs = compute_log_probs( model, input_ids, attention_mask ) # 参考模型对数概率 with torch.no_grad(): ref_token_log_probs = compute_log_probs( ref_model, input_ids, attention_mask ) # 这个比率（ratio）表示策略模型生成的 token 相比于参考模型，是具有更高的概率还是更低的概率 ratio = torch.exp(token_log_probs - ref_token_log_probs) # 根据优势值缩放比率，生成的 token 是否对模型学习方向有正向作用 unclipped = ratio * advantage # 裁剪比率：将比率控制在一个范围，防止比率过大或过小 clipped = torch.clamp(ratio, 1-epsilon, 1+epsilon) * advantage policy_loss = torch.min(unclipped, clipped) # 当 delta 为正值时，意味着策略模型相比于参考模型对生成的词元更有信心 delta = token_log_probs - ref_token_log_probs per_token_kl = torch.exp(-delta) - (-delta) -1 # policy_loss 是优势值，越大越好 # Kl 散度是惩罚值，越小越好 # We want to maximize reward, so we make the loss negative because optimizers minimize loss per_token_loss = -(policy_loss - beta * per_token_kl) # 只考虑输出 tokens 的损失 loss = (per_token_loss * completion_mask).sum() / completion_mask.sum() return lossKL 散度计算公式中 delta 与 KL 损失值之间的关系如下。
（1）当 delta=0 时，意味着策略模型和参考模型对输出分配了相同的概率。所以他们对生成的词元具有相同的信心。
（2）当 delta>0 时，意味着策略模型相比于参考模型对生成的词元有更多的信心，在这种情况下 KL 散度是一个非常小的惩罚项，它告诉策略模型生成的词元很好，同时防止策略模型与参考模型偏差太多。
（3）当 delta<0 时，意味着策略模型相比于参考模型对生成的词元缺乏信心，KL 散度会非常迅速地增加，很快告诉模型它偏离了参考模型，需要向更接近参考模型的位置来修正方向。
