实例 GRPO

import torch
from torch.nn.functional import kl_div, log_softmax, softmax
from transformers import AutoModelForCausalLM, AutoTokenizer
from datasets import load_dataset
import re
from copy import deepcopy

# 1. 加载数据集
dataset = load_dataset("openai/gsm8k", split="train[:5%]")

# 2. 奖励函数
# 组件：奖励函数，用于评估每一个响应的好坏
def reward_function(completions, answers, **kwargs):
    rewards = []
    pattern = r"<answer>(.*?)</answer>"
    for completion, correct_answer in zip(completions, answers):
        try:
            match = re.search(pattern, completion)
            reward = 1.0 if match and match.group(1).strip() == str(correct_answer) else 0.0
        except:
            reward = 0.0
        rewards.append(reward)
    return rewards

# 3. 数据预处理
def format_dataset(example):
    system_prompt = (
        "Solve the math problem step-by-step, providing reasoning in <think> tags "
        "and the final answer in <answer> tags."
    )
    return {
        "prompt": [
            {"role": "system", "content": system_prompt},
            {"role": "user", "content": example["question"]}
        ],
        "answer": example["answer"]
    }

dataset = dataset.map(format_dataset)

# 4. 初始化模型和分词器
model_name = "Qwen/Qwen2-0.5B-Instruct"
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModelForCausalLM.from_pretrained(model_name, torch_dtype=torch.float16)
model_old = deepcopy(model)  # 参考模型（旧策略），保持不变
model.to("cuda" if torch.cuda.is_available() else "cpu")
model_old.to("cuda" if torch.cuda.is_available() else "cpu")

# 5. 自定义训练循环
def train_with_kl(dataset, num_epochs=3, batch_size=4, num_generation=4, kl_weight=0.1):
    # Adam 优化器的目标是最小化
    # 这里默认是更新所有参数 0.5B 
    optimizer = torch.optim.AdamW(model.parameters(), lr=1e-5)
    for epoch in range(num_epochs):
        for i in range(0, len(dataset), batch_size):
            batch = dataset[i:i + batch_size]
            prompts = batch["prompt"]
            answers = batch["answer"]

            # 生成多个候选输出
            completions = []
            new_logits = []
            old_logits = []
            for prompt in prompts:
                inputs = tokenizer([f"{prompt[0]['content']}\n{prompt[1]['content']}"], 
                                    return_tensors="pt", 
                                    padding=True).to(model.device)
                # 组件分组采样：每一个prompt 生成4个响应，形成一个响应组
                for _ in range(num_generation):
                    outputs = model.generate(**inputs, 
                                             max_new_tokens=50, 
                                             do_sample=True, 
                                             return_dict_in_generate=True, 
                                             output_scores=True)
                    # 生成响应，evla值
                    completion = tokenizer.decode(outputs.sequences[0], skip_special_tokens=True)
                    completions.append(completion)
                    # 获取 logits , 最后一层的 logits
                    logits = torch.stack(outputs.scores, dim=1)[:, -1, :]
                    new_logits.append(logits)
                # 旧策略 logits，通过no_grad() 确保旧policy不参与梯度计算
                with torch.no_grad():
                    ref_outputs = model_old(**inputs, return_dict=True)
                    old_logits.append(ref_outputs.logits[:, -1, :])

            # 计算奖励
            # 评估响应，响应值与真值比较
            rewards = reward_function(completions, answers)

            # 计算 KL 散度
            kl_loss = 0
            for new_logit, old_logit in zip(new_logits, old_logits):
                kl_loss += kl_div(
                    log_softmax(new_logit, dim=-1),
                    softmax(old_logit, dim=-1),
                    reduction="batchmean"
                )
            kl_loss /= num_generation

            # 计算奖励loss
            reward_loss = -torch.tensor(rewards, device=model.device).mean()
            # 总loss
            total_loss = reward_loss + kl_weight * kl_loss
            # 优化
            optimizer.zero_grad()
            total_loss.backward()
            optimizer.step()

            print(f"Epoch {epoch+1}, Batch {i//batch_size+1}, Total Loss: {total_loss.item()}, KL Loss: {kl_loss.item()}")

# 6. 开始训练
train_with_kl(dataset, num_epochs=3, batch_size=4, num_generation=4, kl_weight=0.1)

kl_loss = 0
for new_logit, old_logit in zip(new_logits, old_logits):
    kl_loss += kl_div(
        log_softmax(new_logit, dim=-1),
        softmax(old_logit, dim=-1),
        reduction="batchmean"
    )
kl_loss /= num_generation

new_logits = [
    torch.tensor([[2.1, 0.5, ..., -1.2], ..., [1.8, 0.3, ..., -0.9]]),  # 第一个响应
    torch.tensor([[1.9, 0.7, ..., -1.0], ..., [2.0, 0.4, ..., -1.1]]),  # 第二个响应
    torch.tensor([[...]]),
    torch.tensor([[...]])
]
old_logits = [
    torch.tensor([[1.5, 0.6, ..., -0.8], ..., [1.7, 0.5, ..., -0.7]]),  # 老模型输出
    torch.tensor([[...]]),
    torch.tensor([[...]]),
    torch.tensor([[...]])
]