GRPO 算法损失函数原理与代码实现

论文：《DeepSeek-R1: Incentivizing Reasoning Capability in LLMs via Reinforcement Learning》

一、GRPO 损失函数

二、GRPO 算法核心组成部分

GRPO 算法可分解为四个关键部分：

1. 策略损失（policy loss）：模型在有适配器和没有适配器情况下的词元概率分布比率。
2. 优势值（advantages）：从奖励函数中计算得出。
3. 比率裁剪（clip）：确保在任何单独步骤中都没有大的损失值。
4. KL 散度：确保训练过程中，模型不会偏离基准模型太多。

1. 模型加载与初始化

首先加载所需的模型和分词器，并打印模型的网络结构和生成文本的效果。

from transformers import AutoModelForCausalLM, AutoTokenizer
import torch

# 初始化 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)

prompt = "The quick brown fox jumped over the "
input_ids = tokenizer(prompt, return_tensors="pt")
print(input_ids)

# Generate next 2 tokens with torch.no_grad()
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
generated_text = tokenizer.decode(outputs[0], skip_special_tokens=True)
generated_portion = generated_text[len(prompt):]
print(f"Generated text: {prompt}{generated_portion}")

在 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_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 的对数概率。

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
    completion_mask = torch.zeros(total_length, dtype=torch.float32)
    completion_mask[prompt_length:] = 1.0
    
    return input_ids, attention_mask, completion_mask

import torch.nn.functional as F

def compute_log_probs(model, input_ids, attention_mask):
    outputs = model(input_ids, attention_mask=attention_mask)
    # outputs.logits 是神经网络输出中未经过归一化的概率
    log_probs = F.log_softmax(outputs.logits, dim=-1)
    return log_probs.gather(
        dim=-1, index=input_ids.unsqueeze(-1)
    ).squeeze(-1)

策略损失函数的含义是策略模型对每个 token 的对数概率，与参考模型对每个 token 的对数概率的比率。这个比率（ratio）越大，说明策略模型对生成的 token 的信心越大。

advantages 是通过奖励函数转换来的优势值，强化学习的目标就是获得最大的优势值。

在 GRPO 算法中根据优势值缩放比率，来判断策略模型信心大的生成 token 的方向是不是和优势值方向一致。

代码中的 completion_mask 是由模型生成的 token 的位置掩码，在计算策略损失时，只考虑由模型生成的 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)
    
    # 根据优势值缩放比率
    policy_loss = ratio * advantage
    
    # We want to maximize reward, so we make the loss negative
    per_token_loss = -policy_loss
    
    # 只考虑输出 tokens 的损失
    loss = (per_token_loss * completion_mask).sum() / completion_mask.sum()
    return loss

当参考模型和策略模型相同，比率是 1，policy_loss 等于 advantages。这也说明为什么奖励函数生成的奖励分数需要多样性，如果生成的奖励分数是固定的，那么得到优势值也都是 0，导致损失为 0，阻止训练过程。

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 = torch.exp(token_log_probs - ref_token_log_probs)
    unclipped = ratio * advantage
    
    # 裁剪比率：将比率控制在一个范围，防止比率过大或过小
    clipped = torch.clamp(ratio, 1-epsilon, 1+epsilon) * advantage
    policy_loss = torch.min(unclipped, clipped)
    
    per_token_loss = -policy_loss
    loss = (per_token_loss * completion_mask).sum() / completion_mask.sum()
    return loss

加入比率裁剪之后，损失值比之前的更靠近 -2.0，说明损失稳定很多。

4. 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 = torch.exp(token_log_probs - ref_token_log_probs)
    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 散度是惩罚值，越小越好
    per_token_loss = -(policy_loss - beta * per_token_kl)
    
    loss = (per_token_loss * completion_mask).sum() / completion_mask.sum()
    return loss

KL 散度计算公式中 delta 与 KL 损失值之间的关系如下：

1. 当 delta=0 时，意味着策略模型和参考模型对输出分配了相同的概率。
2. 当 delta>0 时，意味着策略模型相比于参考模型对生成的词元有更多的信心，在这种情况下 KL 散度是一个非常小的惩罚项。
3. 当 delta<0 时，意味着策略模型相比于参考模型对生成的词元缺乏信心，KL 散度会非常迅速地增加，告诉模型它偏离了参考模型。