聊聊DeepSeek训练——基于强化学习训练的底层逻辑

在接触DeepSeek时，就有疑问：是怎么基于RL算法训练的？就一句GRPO算法的介绍吗？GRPOTrainer[1]类是HuggingFace官方基于GRPO算法的实现，DeepSeek官方paper地址：GRPO[2]。Group Relative Policy Optimization（GRPO）是一种用于强化学习和自然语言处理的优化算法，旨在改进策略优化过程，尤其适用于处理文本生成等任务。

健忘的派大星

909人浏览 · 2025-02-13 10:14:16

健忘的派大星 · 2025-02-13 10:14:16 发布

❝

聊聊从GRPO RL算法实现来看大模型的训练逻辑

概述

在接触DeepSeek时，就有疑问：是怎么基于RL算法训练的？就一句GRPO算法的介绍吗？

GRPOTrainer[1]类是HuggingFace官方基于GRPO算法的实现，DeepSeek官方paper地址：GRPO[2]。

Group Relative Policy Optimization（GRPO）是一种用于强化学习和自然语言处理的优化算法，旨在改进策略优化过程，尤其适用于处理文本生成等任务。

其算法原理如下

基于相对优势： GRPO 算法关注的是组内样本之间的相对优势，而非绝对的奖励值。在一个批次的样本中，它通过比较不同样本的奖励来确定每个样本的相对优劣，以此作为优化策略的依据。这种相对优势的计算可以减少奖励函数的偏差和方差，使训练更加稳定。
策略优化： 其核心目标是优化策略网络，使得智能体在面对不同的输入时，能够选择更有可能获得高相对奖励的动作。通过不断调整策略网络的参数，让智能体逐渐学习到最优的行为策略。

GRPO算法公式如下：

计算流程图如下：

源码

GRPOTrainer继承自Trainer类，在Trainer类中封装了很多的训练逻辑，GRPOTrainer只需要实现该算法必要的步骤即可。

示例调用：

from datasets import load_dataset
  from trl import GRPOTrainer

  dataset = load_dataset("trl-lib/tldr", split="train")

  trainer = GRPOTrainer(
      model="Qwen/Qwen2-0.5B-Instruct",
      reward_funcs="weqweasdas/RM-Gemma-2B",
      train_dataset=dataset,
  )

  trainer.train()

代码结构：

总的来说，prediction_step、compute_loss、_prepare_inputs、_get_per_token_logps四个函数是GRPO算法的关键实现，同时结合父类Trainer，即可实现神经网络模型的前向迭代、反向传播，损失计算、梯度更新等训练流程。

prediction_step

该方法在父类Trainer中调用，而GRPOTrainer重写了该方法，最终会返回计算后的loss。

def prediction_step(self, model, inputs, prediction_loss_only, ignore_keys: Optional[list[str]] = None):
    # 预处理输入的input
    inputs = self._prepare_inputs(inputs)
    with torch.no_grad():
        with self.compute_loss_context_manager():
            # 计算loss
            loss = self.compute_loss(model, inputs)
        loss = loss.mean().detach()
    # 返回loss值
    return loss, None, None

_prepare_inputs

预处理输入的input，在这里会计算completion（即model的返回值）、reward值、GRPO算法的advantage值等，并最终都返回；其中GRPO算法的advantage计算公式为：

：平均值

：标准差

看看最后返回的结构与注释：

{
  "prompt_ids": prompt_ids, # prompt对应的数字表示，也可以理解为词元token
  "prompt_mask": prompt_mask, # mask掩码
  "completion_ids": completion_ids, # model返回的预测值的数字表示，即对prompt的补全
  "completion_mask": completion_mask, # mask掩码
  "ref_per_token_logps": ref_per_token_logps, # 引用模型的token概率
  "advantages": advantages, # GRPO算法的advantage，在Loss计算需要用到
}

关键代码片段及注释：

def _prepare_inputs(self, inputs: dict[str, Union[torch.Tensor, Any]]) -> dict[str, Union[torch.Tensor, Any]]:
    # 处理输入的query，即prompt
    prompts = [x["prompt"] for x in inputs]
    prompts_text = [maybe_apply_chat_template(example, self.processing_class)["prompt"] for example in inputs]
    # 返回prompt的数字表示
    prompt_inputs = self.processing_class(
        prompts_text, return_tensors="pt", padding=True, padding_side="left", add_special_tokens=False
    )
    prompt_inputs = super()._prepare_inputs(prompt_inputs)
    prompt_ids, prompt_mask = prompt_inputs["input_ids"], prompt_inputs["attention_mask"]

    # .....其它code.....
    # ......
    # ......
    
    # Generate completions using vLLM: gather all prompts and use them in a single call in the main process
    # 使用vllm生成prompt的后续补全信息，即model的输出
    all_prompts_text = gather_object(prompts_text)
    if self.accelerator.is_main_process:
        # 调用模型生成补全
        outputs = self.llm.generate(all_prompts_text, sampling_params=self.sampling_params, use_tqdm=False)
        completion_ids = [out.token_ids for completions in outputs for out in completions.outputs]
    else:
        completion_ids = [None] * len(all_prompts_text) * self.num_generations

    # 解码输出的ids，将数字表示转换为文本表示
    completions = self.processing_class.batch_decode(completion_ids, skip_special_tokens=True)
    if is_conversational(inputs[0]):
        completions = [[{"role": "assistant", "content": completion}] for completion in completions]

    # 计算reward
    # 将prompt与model的completion作为参数，传入奖励函数中；具体的奖励逻辑由外部传入的reward_func参数决定
    output_reward_func = reward_func(prompts=prompts, completions=completions, **reward_kwargs)
    rewards_per_func[:, i] = torch.tensor(output_reward_func, dtype=torch.float32, device=device)
    rewards = rewards_per_func.sum(dim=1)

    # 根据GRPO算法的advantages
    mean_grouped_rewards = rewards.view(-1, self.num_generations).mean(dim=1)
    std_grouped_rewards = rewards.view(-1, self.num_generations).std(dim=1)

    mean_grouped_rewards = mean_grouped_rewards.repeat_interleave(self.num_generations, dim=0)
    std_grouped_rewards = std_grouped_rewards.repeat_interleave(self.num_generations, dim=0)
    # 这里...
    advantages = (rewards - mean_grouped_rewards) / (std_grouped_rewards + 1e-4)

    # 最终返回...
    return {
        "prompt_ids": prompt_ids,
        "prompt_mask": prompt_mask,
        "completion_ids": completion_ids,
        "completion_mask": completion_mask,
        "ref_per_token_logps": ref_per_token_logps,
        "advantages": advantages,
    }

总的来说，GRPO算法的advantage就是在这里计算的，最后返回。

compute_loss

loss计算，在前面做的处理会往inputs中放入一些必要的参数，而在loss计算时，将其作为参数。在这里会计算KL散度，结合之前计算的advantage值，基于GRPO算法的目标函数公式，最终计算出loss。该loss值会在父类Tainer中反向传播应用，并计算梯度。

 def compute_loss(self, model, inputs, return_outputs=False, num_items_in_batch=None):

    prompt_ids, prompt_mask = inputs["prompt_ids"], inputs["prompt_mask"]
    completion_ids, completion_mask = inputs["completion_ids"], inputs["completion_mask"]
    input_ids = torch.cat([prompt_ids, completion_ids], dim=1)
    attention_mask = torch.cat([prompt_mask, completion_mask], dim=1)
    logits_to_keep = completion_ids.size(1)  # we only need to compute the logits for the completion tokens

    # 获取词元概率 （model的输出与ref model的输出）
    per_token_logps = self._get_per_token_logps(model, input_ids, attention_mask, logits_to_keep)

    # 计算KL散度
    ref_per_token_logps = inputs["ref_per_token_logps"]
    per_token_kl = torch.exp(ref_per_token_logps - per_token_logps) - (ref_per_token_logps - per_token_logps) - 1

    # 计算loss：根据advantage，kl散度，基于GRPO算法的目标函数
    advantages = inputs["advantages"]
    per_token_loss = torch.exp(per_token_logps - per_token_logps.detach()) * advantages.unsqueeze(1)
    per_token_loss = -(per_token_loss - self.beta * per_token_kl)
    loss = ((per_token_loss * completion_mask).sum(dim=1) / completion_mask.sum(dim=1)).mean()

    return loss