GRPO 算法损失函数原理与代码实现 | 极客日志

PythonAI算法

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

GRPO 算法损失函数的原理与代码实现。涵盖策略损失、优势值、比率裁剪及 KL 散度四个核心组件。通过对比参考模型与策略模型的对数概率比率，结合奖励优势值优化生成方向。引入裁剪机制稳定训练，利用 KL 散度约束模型偏差。提供基于 PyTorch 和 PEFT 的完整 Python 代码示例。

人间过客发布于 2026/3/30更新于 2026/4/186 浏览

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

一、GRPO 损失函数

文章配图

二、GRPO 算法核心组成部分

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

策略损失（policy loss）：模型在有适配器和没有适配器情况下的词元概率分布比率。
优势值（advantages）：从奖励函数中计算得出。
比率裁剪（clip）：确保在任何单独步骤中都没有大的损失值。
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}")

极客日志微信公众号二维码

更多推荐文章

相关免费在线工具

加密/解密文本
使用加密算法（如AES、TripleDES、Rabbit或RC4）加密和解密文本明文。在线工具，加密/解密文本在线工具，online
RSA密钥对生成器
生成新的随机RSA私钥和公钥pem证书。在线工具，RSA密钥对生成器在线工具，online
Mermaid 预览与可视化编辑
基于 Mermaid.js 实时预览流程图、时序图等图表，支持源码编辑与即时渲染。在线工具，Mermaid 预览与可视化编辑在线工具，online
curl 转代码
解析常见 curl 参数并生成 fetch、axios、PHP curl 或 Python requests 示例代码。在线工具，curl 转代码在线工具，online
Base64 字符串编码/解码
将字符串编码和解码为其 Base64 格式表示形式即可。在线工具，Base64 字符串编码/解码在线工具，online
Base64 文件转换器
将字符串、文件或图像转换为其 Base64 表示形式。在线工具，Base64 文件转换器在线工具，online

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)

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)

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

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

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