MiniMind GRPO 训练源码深度解析

GRPO算法流程图 图 1：GRPO 训练流程示意图

GRPO 算法与 MiniMind 对齐实战导读#

在 LLM 的训练流程中，预训练提供知识，SFT 提供指令对齐，而 RLHF 则让模型真正学会“偏好”的权衡。GRPO 是一种更轻量的偏好优化方案，它不依赖显式 Critic，而是对同一 prompt 的多条回复在组内做相对标准化，用优势信号直接更新策略。

本文将以 train_grpo.py 为主线，重点拆解以下关键模块：

组内相对优势 (Group-Relative Advantage)：同一 prompt 生成多条回复后，在组内标准化得到优势信号。
无 Critic 的直接策略优化：直接对 token 级对数概率加权，并加入 KL 惩罚，限制策略偏移。
混合奖励工程：Reward Model 打分 + 规则奖励共同构成最终奖励，适配推理模型的结构化输出。

GRPO 算法核心原理简述#

1. 组内采样与奖励#

对每个 prompt 生成 $K$ 个候选回复：

\{y_1, y_2, \dots, y_K\}

计算每个回复的奖励：

r_i = R(x, y_i)

在组内做标准化，得到优势：

\mu = \frac{1}{K}\sum_{i=1}^K r_i,\quad \sigma = \sqrt{\frac{1}{K}\sum_{i=1}^K (r_i - \mu)^2}

A_i = \frac{r_i - \mu}{\sigma + \epsilon}

2. 策略更新目标#

GRPO 直接对策略的 token 级对数概率做加权，并加入 KL 惩罚项：

\mathcal{L} = -\mathbb{E}\left[A_i \cdot \log \pi_\theta(y_i|x)\right] + \beta\,\mathrm{KL}(\pi_\theta \parallel \pi_{\text{ref}})

其中：

$\pi_\theta$ 为当前策略模型。
$\pi_{\text{ref}}$ 为参考模型（通常是 SFT 权重）。
$\beta$ 控制 KL 惩罚强度。

实现中常用 token 级 KL 近似：

\Delta = \log \pi_{\text{ref}} - \log \pi_\theta,\quad \mathrm{KL}_{\text{token}} \approx \exp(\Delta) - \Delta - 1

3. 训练流程概览#

取一批 prompts。
每个 prompt 生成 $K$ 条回复。
计算奖励并组内标准化得到优势。
计算策略与参考模型的 per-token logp。
构造损失（优势项 + KL 惩罚）。
反向传播并更新参数。

全局引用与环境初始化 (Imports & Setup)#

这一段完成运行时环境准备：修正包路径、引入训练依赖、设置分布式与告警过滤，为后续函数与主入口打好基础。

代码：全局引用与环境初始化

1
import os
2
import sys
3

4
__package__ = "trainer"
5
sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), '..')))
6

7
import argparse
8
import re
9
import gc
10
import warnings
11
import torch
12
import torch.distributed as dist
13
from transformers import AutoTokenizer
14
from contextlib import nullcontext
15
from torch import optim
16
from torch.nn.parallel import DistributedDataParallel
17
from torch.utils.data import DataLoader, DistributedSampler
18
from torch.optim.lr_scheduler import CosineAnnealingLR
19
from transformers import AutoModel
20
from model.model_minimind import MiniMindConfig, MiniMindForCausalLM
21
from dataset.lm_dataset import RLAIFDataset
22
from trainer.trainer_utils import Logger, is_main_process, lm_checkpoint, init_distributed_mode, setup_seed, SkipBatchSampler, init_model
23

24
warnings.filterwarnings('ignore')

奖励计算逻辑 (calculate_rewards)#

该函数负责整合规则奖励与 Reward Model 评分，生成每条回复的总奖励，为后续优势估计与策略更新提供信号。

代码：calculate_rewards

1
def calculate_rewards(prompts, responses, reward_model, reward_tokenizer):
2
    """整合所有奖励函数计算总奖励"""
3
    def reasoning_model_reward(rewards):
4
        # 正则化匹配response整体格式
5
        pattern = r"^<think>\n.*?\n</think>\n<answer>\n.*?\n</answer>$"
6
        pattern2 = r"^<think>\n.*?\n</think>\n\n<answer>\n.*?\n</answer>$"
7
        matches_pattern = [re.match(pattern, response, re.S) for response in responses]
8
        matches_pattern2 = [re.match(pattern2, response, re.S) for response in responses]
9

10
        format_rewards = []
11
        for match_pattern, match_pattern2 in zip(matches_pattern, matches_pattern2):
12
            # 如果符合pattern1和pattern2人任意格式，就加0.5
13
            if match_pattern or match_pattern2:
14
                format_rewards.append(0.5)
15
            else:
16
                format_rewards.append(0.0)
17
        # 把格式奖励逐元素
18
        rewards += torch.tensor(format_rewards, device=args.device)
19

20
        def mark_num(text):
21
            reward = 0
22
            # 独立的格式打分，不检查顺序
23
            if text.count("<think>") == 1: reward += 0.25
24
            if text.count("</think>") == 1: reward += 0.25
25
            if text.count("<answer>") == 1: reward += 0.25
26
            if text.count("</answer>") == 1: reward += 0.25
27
            return reward
28
        # 奖励逐个加到
29
        mark_rewards = [mark_num(response) for response in responses]
30
        rewards += torch.tensor(mark_rewards, device=args.device)
31
        return rewards
32

33
    rewards = torch.zeros(len(responses), device=args.device)
34
    if args.reasoning == 1:
35
        rewards = reasoning_model_reward(rewards)
36
    # 不进行梯度计算
37
    with torch.no_grad():
38
        reward_model_scores = []
39
        # 得到Batch_Size
40
        batch_size = len(prompts)
41
        # 用于分数裁剪
42
        scale = 3.0
43

44
        for i in range(batch_size):
45
            for j in range(args.num_generations):
46
                # 二维展开到一维索引
47
                response_idx = i * args.num_generations + j
48
                response = responses[response_idx]
49
                prompt = prompts[i]
50

51
                pattern = r"<\|im_start\|>(system|user|assistant)\s+(.*?)<\|im_end\|>"
52
                matches = re.findall(pattern, prompt, re.DOTALL)
53
                # 将Prompts中形如pattern的段落解析成message列表
54
                messages = [{"role": role, "content": content.strip()} for role, content in matches]
55
                # 进行message拼接
56
                tmp_chat = messages + [{"role": "assistant", "content": response}]
57
                # score是标量, 裁剪到[-3, 3]
58
                score = reward_model.get_score(reward_tokenizer, tmp_chat)
59
                score = max(min(score, scale), -scale)
60

61
                if args.reasoning == 1:
62
                    answer_match = re.search(r'<answer>(.*?)</answer>', response, re.DOTALL)
63
                    if answer_match:
64
                        # 从<answer></answer>中的内容弄给单独打分, 并与全文评分融合
65
                        answer_content = answer_match.group(1).strip()
66
                        tmp_chat = messages + [{"role": "assistant", "content": answer_content}]
67
                        answer_score = reward_model.get_score(reward_tokenizer, tmp_chat)
68
                        answer_score = max(min(answer_score, scale), -scale)
69
                        # 40%来自全文, 60%来自答案
70
                        score = score * 0.4 + answer_score * 0.6
71

72
                reward_model_scores.append(score)
73

74
        reward_model_scores = torch.tensor(reward_model_scores, device=args.device)
75
        rewards += reward_model_scores
76

77
    return rewards

GRPO 单轮训练 (grpo_train_epoch)#

这一段涵盖生成采样、对数概率计算、优势归一化、KL 约束以及梯度更新，是 GRPO 训练的核心循环。

代码：grpo_train_epoch

1
def grpo_train_epoch(epoch, loader, iters, ref_model, reward_model, reward_tokenizer, start_step=0, wandb=None):
2
    for step, batch in enumerate(loader, start=start_step + 1):
3
        prompts = batch['prompt']  # list[str], length B
4
        # 左侧padding(序列左边补pad), 不自动添加BOS/EOS等特殊Token
5
        prompt_inputs = tokenizer(prompts, return_tensors="pt", padding=True, return_token_type_ids=False,
6
                                  padding_side="left", add_special_tokens=False).to(args.device)  # input_ids: [B, P], attention_mask: [B, P] P是本batch中最长的prompt token数
7
        if args.max_seq_len:
8
            # 截断后, Tensor[B, L], 因为取的是序列末尾的Token, 所以要倒着截断
9
            prompt_inputs["input_ids"] = prompt_inputs["input_ids"][:, -args.max_seq_len:]
10
            prompt_inputs["attention_mask"] = prompt_inputs["attention_mask"][:, -args.max_seq_len:]
11

12
        with torch.no_grad():
13
            # DDP 模型需要使用 .module 访问 generate 方法
14
            model_for_gen = model.module if isinstance(model, DistributedDataParallel) else model
15
            # 进行采样, num_return_sequences即G, 生成G条回答
16
            outputs = model_for_gen.generate(
17
                **prompt_inputs, max_new_tokens=args.max_gen_len, do_sample=True, temperature=0.8,
18
                num_return_sequences=args.num_generations, pad_token_id=tokenizer.pad_token_id)  # [B*num_gen, P+R], 每条序列包含原Prompt+生成Token
19
        # prompt_inputs["input_ids"].size(1)就是P, outputs[:, P:]取每条序列从第P位开始的部分
20
        completion_ids = outputs[:, prompt_inputs["input_ids"].size(1):]  # [B*num_gen, R]
21

22
        def get_per_token_logps(mdl, input_ids, n_keep):
23
            # N=B*G, L=P+R, 形状不变:Tensor[N, L]
24
            input_ids = input_ids.detach().clone() if input_ids.is_inference() else input_ids
25
            # Tensor[N, n_kepp+1, V], V是词表大小, n_keep是生成的response token数R
26
            # 拿到最后n_keep+1个词的分数, 再去掉最后一个
27
            logits = mdl(input_ids, logits_to_keep=n_keep + 1).logits[:, :-1, :]
28
            per_token_logps = []
29
            # logits形状是[N, n_keep, V], input_ids[:, -n_keep:]形状是[N, n_keep]
30
            for logits_row, ids_row in zip(logits, input_ids[:, -n_keep:]):
31
                # logits_row是某一条样本的[n_keep, V], ids_row是同一条样本的[n_keep]
32
                ids_row = ids_row.detach().clone() if ids_row.is_inference() else ids_row
33
                # unsqueeze(1)后变成[n_keepp, 1], 在词表维度上按照真实token id取值, 得到每个位置对应真实token的log概率, 形状为[n_keep, 1], squeeze(1)去掉多余维度, 得到[n_keep]
34
                per_token_logps.append(torch.gather(logits_row.log_softmax(dim=-1), 1, ids_row.unsqueeze(1)).squeeze(1))
35
            # 所有样本拼接成张量, 最终形状[N, n_keep]
36
            return torch.stack(per_token_logps)
37
        # 进入自动混合精读上下文
38
        with autocast_ctx:
39
            # 计算最后R个token的逐token log概率, 返回形状[B*G, R]
40
            per_token_logps = get_per_token_logps(model, outputs, completion_ids.size(1))  # [B*num_gen, R]
41
            # 如果是MoE, 再跑一次拿到MoE的辅助输出
42
            res = model(outputs) if lm_config.use_moe else None
43
            aux_loss = res.aux_loss if res is not None else torch.tensor(0.0, device=args.device)
44

45
        with torch.no_grad():
46
            # 用参考模型计算同样输出序列的逐token log概率, 形状[B*G, R]
47
            ref_per_token_logps = get_per_token_logps(ref_model, outputs, completion_ids.size(1))  # [B*num_gen, R]
48
        # 解码成文本
49
        completions = tokenizer.batch_decode(completion_ids, skip_special_tokens=True)
50
        # 用奖励模型计算每条response的标量奖励
51
        rewards = calculate_rewards(prompts, completions, reward_model, reward_tokenizer).to(args.device)  # [B*num_gen]
52
        # 按照prompt分组, 计算优势, [B, G]
53
        grouped_rewards = rewards.view(-1, args.num_generations)  # [B, num_gen]
54
        # 求均值/标准差并广播回[B*G]
55
        mean_r = grouped_rewards.mean(dim=1).repeat_interleave(args.num_generations)  # [B*num_gen]
56
        std_r = grouped_rewards.std(dim=1).repeat_interleave(args.num_generations)  # [B*num_gen]
57
        # 计算标准化优势
58
        advantages = torch.clamp((rewards - mean_r) / (std_r + 1e-4), -10, 10)
59
        advantages = (advantages - advantages.mean()) / (advantages.std() + 1e-8)  # [B*num_gen]
60
        # 计算每条样本第一个EOS的位置, 没有就默认R
61
        is_eos = completion_ids == tokenizer.eos_token_id  # [B*num_gen, R]
62
        eos_idx = torch.full((is_eos.size(0),), is_eos.size(1), dtype=torch.long, device=args.device)
63
        eos_idx[is_eos.any(dim=1)] = is_eos.int().argmax(dim=1)[is_eos.any(dim=1)]
64
        completion_mask = (torch.arange(is_eos.size(1), device=args.device).expand(is_eos.size(0), -1) <= eos_idx.unsqueeze(1)).int()  # [B*num_gen, R]
65
        # 相当于 log p_ref - log p_pi
66
        kl_div = ref_per_token_logps - per_token_logps
67
        # KL的一种平滑形式, 来自exp(x)-x-1, 数值稳定且对偏离有惩罚
68
        per_token_kl = torch.exp(kl_div) - kl_div - 1  # [B*num_gen, R]
69
        # 前向值为exp(0)=1, 但反向梯度等价于 ∇log π_θ, 再乘以优势
70
        per_token_loss = -(torch.exp(per_token_logps - per_token_logps.detach()) * advantages.unsqueeze(1) - args.beta * per_token_kl)  # [B*num_gen, R]
71
        # 做反向传播
72
        policy_loss = ((per_token_loss * completion_mask).sum(dim=1) / completion_mask.sum(dim=1)).mean()
73
        loss = (policy_loss + aux_loss) / args.accumulation_steps  # scalar
74
        loss.backward()
75
        # 参数更新
76
        if (step + 1) % args.accumulation_steps == 0:
77
            if args.grad_clip > 0:
78
                torch.nn.utils.clip_grad_norm_(model.parameters(), args.grad_clip)
79
            optimizer.step()
80
            scheduler.step()
81
            optimizer.zero_grad()
82
        # 日志打印, wandb记录
83
        if step % args.log_interval == 0 or step == iters:
84
            policy_loss_val = loss.item() * args.accumulation_steps
85
            current_aux_loss = aux_loss.item()
86
            avg_reward_val = rewards.mean().item()
87
            avg_len_val = completion_mask.sum(dim=1).float().mean().item()
88
            current_lr = optimizer.param_groups[0]['lr']
89

90
            Logger(f'Epoch:[{epoch + 1}/{args.epochs}]({step}/{iters}), '
91
                   f'Actor Loss: {policy_loss_val:.4f}, Aux Loss: {current_aux_loss:.4f}, Reward: {avg_reward_val:.4f}, '
92
                   f'Avg Response Len: {avg_len_val:.2f}, Learning Rate: {current_lr:.8f}')
93

94
            if wandb and is_main_process():
95
                wandb.log({
96
                    "policy_loss": policy_loss_val,
97
                    "aux_loss": current_aux_loss,
98
                    "reward": avg_reward_val,
99
                    "avg_response_len": avg_len_val,
100
                    "advantages_mean": advantages.mean().item(),
101
                    "learning_rate": current_lr
102
                })
103

104
        if (step % args.save_interval == 0 or step == iters - 1) and is_main_process():
105
            model.eval()
106
            moe_suffix = '_moe' if lm_config.use_moe else ''
107
            ckp = f'{args.save_dir}/{args.save_weight}_{lm_config.hidden_size}{moe_suffix}.pth'
108
            raw_model = model.module if isinstance(model, DistributedDataParallel) else model
109
            raw_model = getattr(raw_model, '_orig_mod', raw_model)
110
            state_dict = raw_model.state_dict()
111
            torch.save({k: v.half().cpu() for k, v in state_dict.items()}, ckp)
112
            lm_checkpoint(lm_config, weight=args.save_weight, model=model, optimizer=optimizer,
113
                         epoch=epoch, step=step, wandb=wandb, save_dir='../checkpoints', scheduler=scheduler)
114
            model.train()
115
            del state_dict
116

117
        del prompt_inputs, outputs, completion_ids, per_token_logps, ref_per_token_logps
118
        del completions, rewards, grouped_rewards, mean_r, std_r, advantages, completion_mask

主入口与训练流程 (Main Entry)#

主入口负责参数解析、分布式初始化、模型/数据/优化器构建、断点恢复、DDP 包装以及整体训练调度。

代码：主入口

1
if __name__ == "__main__":
2
    parser = argparse.ArgumentParser(description="MiniMind GRPO (Group Relative Policy Optimization)")
3
    parser.add_argument("--save_dir", type=str, default="../out", help="模型保存目录")
4
    parser.add_argument('--save_weight', default='grpo', type=str, help="保存权重的前缀名")
5
    parser.add_argument("--epochs", type=int, default=1, help="训练轮数")
6
    parser.add_argument("--batch_size", type=int, default=2, help="batch size")
7
    parser.add_argument("--learning_rate", type=float, default=8e-8, help="初始学习率")
8
    parser.add_argument("--device", type=str, default="cuda:0" if torch.cuda.is_available() else "cpu", help="训练设备")
9
    parser.add_argument("--dtype", type=str, default="bfloat16", help="混合精度类型")
10
    parser.add_argument("--num_workers", type=int, default=8, help="数据加载线程数")
11
    parser.add_argument("--accumulation_steps", type=int, default=1, help="梯度累积步数")
12
    parser.add_argument("--grad_clip", type=float, default=1.0, help="梯度裁剪阈值")
13
    parser.add_argument("--log_interval", type=int, default=1, help="日志打印间隔")
14
    parser.add_argument("--save_interval", type=int, default=10, help="模型保存间隔")
15
    parser.add_argument('--hidden_size', default=512, type=int, help="隐藏层维度")
16
    parser.add_argument('--num_hidden_layers', default=8, type=int, help="隐藏层数量")
17
    parser.add_argument('--use_moe', default=0, type=int, choices=[0, 1], help="是否使用MoE架构（0=否，1=是）")
18
    parser.add_argument('--max_seq_len', default=66, type=int, help="Prompt最大长度")
19
    parser.add_argument("--max_gen_len", type=int, default=1536, help="生成的最大长度")
20
    parser.add_argument("--data_path", type=str, default="../dataset/rlaif-mini.jsonl", help="RLAIF数据路径")
21
    parser.add_argument("--num_generations", type=int, default=8, help="每个prompt生成的样本数")
22
    parser.add_argument("--beta", type=float, default=0.02, help="KL惩罚系数")
23
    parser.add_argument("--reasoning", type=int, default=1, choices=[0, 1], help='推理模型类型（0=普通模型，1=推理模型）')
24
    parser.add_argument("--reward_model_path", type=str, default="../../internlm2-1_8b-reward", help="Reward模型路径")
25
    parser.add_argument('--from_resume', default=0, type=int, choices=[0, 1], help="是否自动检测&续训（0=否，1=是）")
26
    parser.add_argument("--use_wandb", action="store_true", help="是否使用wandb")
27
    parser.add_argument("--wandb_project", type=str, default="MiniMind-GRPO", help="wandb项目名")
28
    parser.add_argument("--use_compile", default=0, type=int, choices=[0, 1], help="是否使用torch.compile加速（0=否，1=是）")
29
    args = parser.parse_args()
30

31
    # ========== 1. 初始化环境和随机种子 ==========
32
    local_rank = init_distributed_mode()
33
    if dist.is_initialized(): args.device = f"cuda:{local_rank}"
34
    setup_seed(42 + (dist.get_rank() if dist.is_initialized() else 0))
35

36
    # ========== 2. 配置目录、模型参数、检查ckp ==========
37
    os.makedirs(args.save_dir, exist_ok=True)
38
    lm_config = MiniMindConfig(hidden_size=args.hidden_size, num_hidden_layers=args.num_hidden_layers,
39
                               max_seq_len=args.max_seq_len + args.max_gen_len, use_moe=bool(args.use_moe))
40
    ckp_data = lm_checkpoint(lm_config, weight=args.save_weight, save_dir='../checkpoints') if args.from_resume==1 else None
41

42
    # ========== 3. 设置混合精度 ==========
43
    device_type = "cuda" if "cuda" in args.device else "cpu"
44
    dtype = torch.bfloat16 if args.dtype == "bfloat16" else torch.float16
45
    autocast_ctx = nullcontext() if device_type == "cpu" else torch.cuda.amp.autocast(dtype=dtype)
46

47
    # ========== 4. 配wandb ==========
48
    wandb = None
49
    if args.use_wandb and is_main_process():
50
        import swanlab as wandb
51
        wandb_id = ckp_data.get('wandb_id') if ckp_data else None
52
        resume = 'must' if wandb_id else None
53
        wandb_run_name = f"MiniMind-GRPO-Epoch-{args.epochs}-BS-{args.batch_size}-LR-{args.learning_rate}"
54
        wandb.init(project=args.wandb_project, name=wandb_run_name, id=wandb_id, resume=resume)
55

56
    # ========== 5. 初始化模型和数据 ==========
57
    base_weight = "reason" if args.reasoning == 1 else "full_sft"
58
    # Policy模型
59
    model, tokenizer = init_model(lm_config, base_weight, device=args.device)
60
    if args.use_compile == 1:
61
        model = torch.compile(model)
62
        Logger('torch.compile enabled')
63
    # Reference模型
64
    ref_model, _ = init_model(lm_config, base_weight, device=args.device)
65
    ref_model = ref_model.eval().requires_grad_(False)
66
    # Reward模型
67
    reward_model = AutoModel.from_pretrained(
68
        args.reward_model_path, torch_dtype=torch.float16, trust_remote_code=True
69
    )
70
    reward_model = reward_model.to(args.device).eval().requires_grad_(False)
71
    reward_tokenizer = AutoTokenizer.from_pretrained(args.reward_model_path, trust_remote_code=True)
72
    # 数据和优化器
73
    train_ds = RLAIFDataset(args.data_path, tokenizer, max_length=lm_config.max_seq_len)
74
    train_sampler = DistributedSampler(train_ds) if dist.is_initialized() else None
75
    optimizer = optim.AdamW(model.parameters(), lr=args.learning_rate)
76
    loader_for_count = DataLoader(train_ds, batch_size=args.batch_size, sampler=train_sampler)
77
    iters = len(loader_for_count)
78
    total_optimizer_steps = (iters // args.accumulation_steps) * args.epochs
79
    scheduler = CosineAnnealingLR(optimizer, T_max=total_optimizer_steps, eta_min=args.learning_rate / 10)
80

81
    # ========== 6. 从ckp恢复状态 ==========
82
    start_epoch, start_step = 0, 0
83
    if ckp_data:
84
        model.load_state_dict(ckp_data['model'])
85
        optimizer.load_state_dict(ckp_data['optimizer'])
86
        scheduler.load_state_dict(ckp_data['scheduler'])
87
        start_epoch = ckp_data['epoch']
88
        start_step = ckp_data.get('step', 0)
89

90
    # ========== 7. DDP包模型 ==========
91
    if dist.is_initialized():
92
        model._ddp_params_and_buffers_to_ignore = {"freqs_cos", "freqs_sin"}
93
        model = DistributedDataParallel(model, device_ids=[local_rank])
94

95
    # ========== 8. 开始训练 ==========
96
    for epoch in range(start_epoch, args.epochs):
97
        train_sampler and train_sampler.set_epoch(epoch)
98
        setup_seed(42 + epoch); indices = torch.randperm(len(train_ds)).tolist()
99
        skip = start_step if (epoch == start_epoch and start_step > 0) else 0
100
        batch_sampler = SkipBatchSampler(train_sampler or indices, args.batch_size, skip)
101
        loader = DataLoader(train_ds, batch_sampler=batch_sampler, num_workers=args.num_workers, pin_memory=True)
102
        if skip > 0:
103
            Logger(f'Epoch [{epoch + 1}/{args.epochs}]: 跳过前{start_step}个step，从step {start_step + 1}开始')
104
            grpo_train_epoch(epoch, loader, len(loader) + skip, ref_model, reward_model, reward_tokenizer, start_step, wandb)
105
        else:
106
            grpo_train_epoch(epoch, loader, len(loader), ref_model, reward_model, reward_tokenizer, 0, wandb)
107

108
    # ========== 9. 清理分布进程 ==========
109
    if dist.is_initialized(): dist.destroy_process_group()