GRPOTrainer原理解析

最新推荐文章于 2025-02-22 16:00:41 发布

asd8705

最新推荐文章于 2025-02-22 16:00:41 发布

阅读量1.4k

点赞数 21

文章标签： GRPO 强化学习

本文链接：https://blog.csdn.net/asd8705/article/details/145549097

版权

GRPOTrainer

User: GRPOTrainer
Kimi: GRPOTrainer 是基于 Hugging Face 的 transformers 库中的 Trainer 类实现的一个训练器，专门用于实现 Group Relative Policy Optimization（GRPO）算法[¹][²]。以下是关于 GRPOTrainer 的详细介绍：

定义

GRPOTrainer 是一个用于训练语言模型的强化学习工具，它通过优化策略网络，使得模型在面对不同输入时能够生成更有可能获得高奖励的输出[²]。它是 transformers 库中 Trainer 的子类，支持日志记录、检查点、分布式训练和参数高效微调（PEFT）等功能[⁴]。

用途

GRPOTrainer 主要用于优化语言模型的文本生成任务，尤其是在需要通过奖励函数来引导模型生成特定风格或内容的场景中。它通过比较同一组内不同样本的奖励值来计算相对优势，从而优化策略网络[²][³]。

核心功能

相对优势计算：GRPO 算法关注组内样本之间的相对优势，而非绝对奖励值。通过组内归一化计算每个样本的优势，减少奖励函数的偏差和方差[²][³]。
策略优化：通过策略梯度公式和 KL 正则化更新策略网络，使模型逐渐学习到最优的行为策略[²][³]。
灵活的奖励函数：支持多种类型的奖励函数，包括预训练模型、自定义函数等[¹]。
高效训练：支持梯度累积、分布式训练等技术，提高训练效率[¹][⁴]。

使用方法

以下是使用 GRPOTrainer 的基本步骤：

准备数据集：数据集需要包含 prompt 列，可以是标准文本或对话式结构[¹]。
定义奖励函数：可以使用预训练的奖励模型或自定义奖励函数[¹][³]。
初始化训练器：传入模型、奖励函数、训练数据集等参数[¹]。
开始训练：调用 trainer.train() 方法进行训练[¹]。

示例代码

from datasets import load_dataset
from trl import GRPOTrainer

# 加载数据集
dataset = load_dataset("trl-lib/tldr", split="train")

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

# 开始训练
trainer.train()

[¹]

GRPOTrainer 的出现为语言模型的强化学习训练提供了一个高效且灵活的工具，尤其适用于需要通过奖励机制优化模型输出的场景[²][³]。

GRPOTrainer是怎么计算损失是的

User: GRPOTrainer是怎么计算损失是的
Kimi: GRPOTrainer 计算损失的过程主要基于 Group Relative Policy Optimization（GRPO）算法的核心思想，即最大化生成文本的优势，同时确保模型策略接近参考策略。以下是损失计算的具体步骤[⁸][⁹]：

1. 生成完成

在每个训练步骤中，GRPOTrainer 会从数据集中采样一批提示（prompts），并为每个提示生成一组完成（completions），记为 $o_i$ 。

2. 计算优势

对于每个生成的序列 $o_i$ ，使用奖励模型计算奖励值 $r_i$ 。为了与奖励模型的比较性质保持一致，优势（advantage）会进行规范化处理：

$\hat{A}_{i,t} = \frac{r_i - \text{mean}(\mathbf{r})}{\text{std}(\mathbf{r})}$

其中， $\text{mean}(\mathbf{r})$ 和 $\text{std}(\mathbf{r})$ 分别是所有奖励值的均值和标准差[⁸]。

3. 估计 KL 散度

KL 散度用于衡量当前策略与参考策略之间的差异。其近似计算公式为：

$\mathbb{D}_{\text{KL}}\left[\pi_\theta \| \pi_{\text{ref}}\right] = \frac{\pi_{\text{ref}}(o_{i,t} \mid q, o_{i,<t})}{\pi_\theta(o_{i,t} \mid q, o_{i,<t})} - \log \frac{\pi_{\text{ref}}(o_{i,t} \mid q, o_{i,<t})}{\pi_\theta(o_{i,t} \mid q, o_{i,<t})} - 1$

其中， $\pi_\theta$ 是当前策略， $\pi_{\text{ref}}$ 是参考策略[⁸]。

4. 计算损失

损失函数的目标是最大化优势，同时通过 KL 散度惩罚偏离参考策略的行为。损失函数定义为：

$\mathcal{L}_{\text{GRPO}}(\theta) = -\frac{1}{G} \sum_{i=1}^G \frac{1}{\|o_i\|} \sum_{t=1}^{\|o_i\|} \left[ \frac{\pi_\theta(o_{i,t} \mid q, o_{i,<t})}{\left[\pi_\theta(o_{i,t} \mid q, o_{i,<t})\right]_{\text{no grad}}} \hat{A}_{i,t} - \beta \mathbb{D}_{\text{KL}}\left[\pi_\theta \| \pi_{\text{ref}}\right] \right]$

其中：

第一项是缩放后的优势。
第二项通过 KL 散度惩罚偏离参考策略的行为， $\beta$ 是一个超参数，用于平衡两项[⁹]。

在实际实现中，GRPOTrainer 会在每个训练步骤中计算上述损失，并通过反向传播更新模型参数[¹⁰]。

能够详细讲一下这个损失函数吗？

User: 能够详细讲一下这个损失函数吗？
Kimi: GRPOTrainer 的损失函数是基于 Group Relative Policy Optimization（GRPO）算法的核心思想设计的，其目的是通过最大化生成文本的优势（advantage），同时确保模型策略保持接近参考策略（reference policy）。以下是损失函数的详细解析[¹³][¹⁴]：

损失函数的定义

GRPO 的损失函数定义如下：

其中：

$G$ 是生成的样本组数。
$o_i$ 是第 $i$ 组生成的序列。
$\pi_\theta(o_{i,t} \mid q, o_{i,<t})$ 是当前策略下生成第 $t$ 个词元的概率。
$\hat{A}_{i,t}$ 是第 $t$ 个词元的规范化优势（normalized advantage）。
$\beta$ 是一个超参数，用于平衡优势项和 KL 散度项。
$\mathbb{D}_{\text{KL}}\left[\pi_\theta \| \pi_{\text{ref}}\right]$ 是当前策略与参考策略之间的 KL 散度。

损失函数的组成部分

优势项（Advantage Term）：
- 优势 $\hat{A}_{i,t}$ 是通过奖励函数计算得到的规范化奖励值，用于衡量生成的文本相对于其他候选文本的优劣。
- 优势项的目标是最大化生成文本的奖励值，即鼓励模型生成更有可能获得高奖励的文本。
KL 散度项（KL Divergence Term）：
- KL 散度 $\mathbb{D}_{\text{KL}}\left[\pi_\theta \| \pi_{\text{ref}}\right]$ 用于衡量当前策略与参考策略之间的差异。
- KL 散度项的作用是防止模型策略过度偏离参考策略，从而保持模型的稳定性和一致性。

损失函数的计算步骤

计算词元概率：
- 使用当前模型计算每个生成词元的概率 $\pi_\theta(o_{i,t} \mid q, o_{i,<t})$ 。
- 同时，计算参考策略下的词元概率 $\pi_{\text{ref}}(o_{i,t} \mid q, o_{i,<t})$ 。
计算 KL 散度：
- KL 散度的计算公式为：

$\mathbb{D}_{\text{KL}}\left[\pi_\theta \| \pi_{\text{ref}}\right] = \pi_{\text{ref}}(o_{i,t} \mid q, o_{i,<t}) \left( \frac{\pi_{\text{ref}}(o_{i,t} \mid q, o_{i,<t})}{\pi_\theta(o_{i,t} \mid q, o_{i,<t})} - 1 \right)$
3. 计算每个词元的损失：
- 每个词元的损失为：

$per_token_loss = − ( π θ ( o i , t ∣ q , o i , < t ) [ π θ ( o i , t ∣ q , o i , < t ) ] no grad A ^ i , t − β D KL [ π θ ∥ π ref ] ) \text{per\_token\_loss} = -\left( \frac{\pi_\theta(o_{i,t} \mid q, o_{i,<t})}{\left[\pi_\theta(o_{i,t} \mid q, o_{i,<t})\right]_{\text{no grad}}} \hat{A}_{i,t} - \beta \mathbb{D}_{\text{KL}}\left[\pi_\theta \| \pi_{\text{ref}}\right] \right)$
4. 计算总损失：
- 将所有词元的损失加权平均，得到最终的损失值：

$per_token_loss \text{loss} = \frac{1}{G} \sum_{i=1}^G \frac{1}{\|o_i\|} \sum_{t=1}^{\|o_i\|} \text{per\_token\_loss}$

实现细节

在 GRPOTrainer 的实现中，compute_loss 方法负责计算上述损失函数[¹⁴]。具体代码如下：

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

总结

GRPOTrainer 的损失函数通过结合优势项和 KL 散度项，实现了在强化学习框架下对语言模型的优化。优势项鼓励模型生成高奖励的文本，而 KL 散度项则确保模型策略的稳定性[¹³][¹⁴]。

prompt_ids和completion_ids分别指的是？

User: prompt_ids和completion_ids分别指的是？
Kimi: 在 GRPOTrainer 的上下文中，prompt_ids 和 completion_ids 是用于文本生成任务的两个关键输入部分，它们分别代表：

1. `prompt_ids`

prompt_ids 是提示文本（Prompt）的 Token ID 表示。提示文本是用户输入的初始文本，用于引导模型生成后续内容。它为模型提供了上下文信息，帮助模型理解生成任务的背景和目标[¹⁷]。

例如：

如果用户希望模型生成一段关于“人工智能”的介绍，那么提示文本可能是：“人工智能是一种……”。
在实际应用中，prompt_ids 是将提示文本经过 Tokenizer 处理后得到的 Token ID 序列，用于输入到模型中[¹⁷]。

2. `completion_ids`

completion_ids 是模型根据提示文本生成的完成文本（Completion）的 Token ID 表示。完成文本是模型根据提示文本生成的后续内容[¹⁷]。

例如：

如果提示文本是：“人工智能是一种……”，那么模型生成的完成文本可能是：“……能够模拟人类智能行为的技术，广泛应用于……”。
同样，completion_ids 是将生成的完成文本经过 Tokenizer 处理后得到的 Token ID 序列[¹⁷]。

在 `GRPOTrainer` 中的作用

prompt_ids：作为输入，提供给模型初始的上下文信息，帮助模型理解生成任务的背景。
completion_ids：作为目标输出，用于计算奖励值和损失函数。模型的目标是生成高质量的完成文本，以获得更高的奖励[¹⁷]。

示例

假设我们有一个提示文本：“今天天气很好，适合……”，经过 Tokenizer 处理后得到的 prompt_ids 可能是 [101, 123, 456, 789]。模型根据这个提示生成的完成文本是：“适合出去郊游”，对应的 completion_ids 可能是 [234, 567, 890][¹⁷]。

在 GRPOTrainer 的训练过程中，prompt_ids 和 completion_ids 一起用于计算生成文本的奖励值和损失函数，从而优化模型的生成能力。

KL损失函数详细计算过程

User: KL损失函数详细计算过程
Kimi: 在 GRPOTrainer 中，KL 散度损失函数用于衡量当前策略（模型生成的概率分布）与参考策略（参考模型的概率分布）之间的差异。以下是 KL 散度损失函数的详细计算过程：

KL 散度的定义

KL 散度（Kullback-Leibler Divergence）用于衡量两个概率分布 $P$ 和 $Q$ 之间的差异，其定义为：

$D_{\text{KL}}(P \| Q) = \sum_{i} P(i) \log \frac{P(i)}{Q(i)}$

其中， $P (i)$ 和 $Q (i)$ 分别是两个概率分布中事件 $i$ 的概率[²⁶]。

KL 散度损失函数的计算步骤

准备概率分布：
- $P$ 是参考策略的概率分布（通常是一个预训练模型或目标分布）。
- $Q$ 是当前策略的概率分布（模型生成的分布）。
计算每个事件的概率比：
- 对于每个事件 $i$ ，计算 $P (i)$ 和 $Q (i)$ 的比值 $\frac{P(i)}{Q(i)}$ 。
计算对数概率比：
- 对每个事件 $i$ ，计算 $\log \frac{P(i)}{Q(i)}$ 。
计算加权和：
- 将每个事件的对数概率比乘以其在参考分布 $P$ 中的概率 $P (i)$ ，并求和：

$D_{\text{KL}}(P \| Q) = \sum_{i} P(i) \log \frac{P(i)}{Q(i)}$
5. 处理数值稳定性：
- 在实际计算中，需要确保 $Q (i) > 0$ ，以避免对数计算中的数值问题。如果 $Q (i) = 0$ ，可以通过添加一个小的常数（如 $10^{-8}$ ）来避免数值不稳定[²⁸]。

示例代码

以下是一个使用 PyTorch 计算 KL 散度的示例代码：

import torch
import torch.nn.functional as F

# 定义两个概率分布
p = torch.tensor([0.1, 0.4, 0.5], dtype=torch.float32)
q = torch.tensor([0.2, 0.3, 0.5], dtype=torch.float32)

# 计算 KL 散度
kl_div = F.kl_div(q.log(), p, reduction='sum')
print(f"KL Divergence: {kl_div.item()}")

在这个示例中，q.log() 是对 $Q$ 的每个元素取对数，p 是参考分布 $P$ 。reduction='sum' 表示对所有事件的 KL 散度求和[²⁸]。

在 `GRPOTrainer` 中的应用

在 GRPOTrainer 中，KL 散度损失函数用于惩罚当前策略与参考策略之间的差异。具体来说：

参考策略：通常是预训练模型的输出分布，表示为 $\pi_{\text{ref}}$ 。
当前策略：是模型生成的分布，表示为 $\pi_\theta$ 。
KL 散度损失：用于确保模型生成的分布接近参考分布，从而保持模型的稳定性和一致性[²⁶]。

通过最小化 KL 散度损失，模型可以学习到更接近参考策略的行为，同时避免过度偏离目标分布[³⁰]。

GRPOTrainer 当前策略是怎么样的

User: GRPOTrainer 当前策略是怎么样的
Kimi: 在 GRPOTrainer 中，当前策略（Current Policy） 是指模型在当前训练阶段所采用的生成策略，即模型根据输入的提示（prompt）生成文本的决策过程。当前策略是通过模型的参数（ $\theta$ ）定义的，这些参数决定了模型如何生成文本。

当前策略的具体实现

模型选择：
- 当前策略由一个预训练的语言模型实现，该模型可以是任何支持因果语言建模（Causal Language Modeling）的模型。
- 在 GRPOTrainer 中，可以通过指定 model 参数来加载预训练模型，例如：

model = "Qwen/Qwen2.5-3B-Instruct"

生成过程：
- 在每个训练步骤中，当前策略会根据输入的提示（prompt_ids）生成一组完成文本（completion_ids）。
- 生成的完成文本用于计算奖励值（rewards）和优势（advantages），进而优化模型的策略[³²]。
策略更新：
- 当前策略会通过优化目标函数进行更新，该目标函数结合了奖励值和 KL 散度项。
- KL 散度项用于确保当前策略与参考策略（Reference Policy）保持接近，从而维持策略的稳定性。

当前策略与参考策略的关系

参考策略：通常是一个固定或缓慢更新的策略，用于提供稳定的基线。
当前策略：在训练过程中不断更新，以最大化奖励值并保持与参考策略的相似性。

示例

假设我们使用 GRPOTrainer 训练一个模型，当前策略的更新过程可以总结为以下步骤：

生成文本：使用当前策略生成一组完成文本。
计算奖励：通过奖励函数计算每个生成文本的奖励值。
计算优势：通过规范化奖励值计算优势。
更新策略：根据优势和 KL 散度更新当前策略[³²]。

通过这种方式，GRPOTrainer 能够在训练过程中逐步优化当前策略，使其生成的文本更符合奖励函数的期望[³²]。

当前策略的概率分布怎么计算用python代码实现

User: 当前策略的概率分布怎么计算用python代码实现
Kimi: 在 GRPOTrainer 中，当前策略的概率分布是通过模型的输出 logits 计算得到的。具体来说，模型的输出 logits 表示每个词元的未归一化对数概率，通过 softmax 函数可以将其转换为概率分布。

以下是用 Python 和 PyTorch 实现当前策略概率分布计算的代码示例：

1. 加载模型和 Tokenizer

首先，加载预训练的语言模型和对应的 Tokenizer。这里以 Hugging Face 的 transformers 库为例：

from transformers import AutoModelForCausalLM, AutoTokenizer

# 加载预训练模型和 Tokenizer
model_name = "Qwen/Qwen2.5-3B-Instruct"
model = AutoModelForCausalLM.from_pretrained(model_name)
tokenizer = AutoTokenizer.from_pretrained(model_name)

2. 准备输入文本

将输入文本转换为模型所需的 Token ID 序列：

# 输入文本
prompt = "今天天气很好，适合出去郊游。"

# 使用 Tokenizer 将文本转换为 Token ID 序列
inputs = tokenizer(prompt, return_tensors="pt")
input_ids = inputs["input_ids"]
attention_mask = inputs["attention_mask"]

3. 生成输出

使用模型生成输出 logits，并计算每个词元的概率分布：

import torch

# 将输入传递给模型，获取输出 logits
with torch.no_grad():  # 禁用梯度计算，节省内存
    outputs = model(input_ids=input_ids, attention_mask=attention_mask)
    logits = outputs.logits

# 获取最后一个词元的 logits（假设我们只关注最后一个词元的生成）
last_logits = logits[:, -1, :]

# 使用 softmax 函数将 logits 转换为概率分布
probabilities = torch.nn.functional.softmax(last_logits, dim=-1)

# 打印概率分布
print("概率分布：", probabilities)

4. 解释代码

model(input_ids=input_ids, attention_mask=attention_mask)：将输入的 Token ID 序列传递给模型，获取输出 logits。
logits[:, -1, :]：获取最后一个词元的 logits。logits 的形状通常是 (batch_size, sequence_length, vocab_size)，其中 vocab_size 是词汇表的大小。
torch.nn.functional.softmax(last_logits, dim=-1)：将 logits 转换为概率分布。softmax 函数将 logits 转换为归一化的概率分布，使得每个词元的概率之和为 1。

5. 完整代码示例

将上述步骤整合到一个完整的代码示例中：

from transformers import AutoModelForCausalLM, AutoTokenizer
import torch

# 加载预训练模型和 Tokenizer
model_name = "Qwen/Qwen2.5-3B-Instruct"
model = AutoModelForCausalLM.from_pretrained(model_name)
tokenizer = AutoTokenizer.from_pretrained(model_name)

# 输入文本
prompt = "今天天气很好，适合出去郊游。"

# 使用 Tokenizer 将文本转换为 Token ID 序列
inputs = tokenizer(prompt, return_tensors="pt")
input_ids = inputs["input_ids"]
attention_mask = inputs["attention_mask"]

# 将输入传递给模型，获取输出 logits
with torch.no_grad():  # 禁用梯度计算，节省内存
    outputs = model(input_ids=input_ids, attention_mask=attention_mask)
    logits = outputs.logits

# 获取最后一个词元的 logits
last_logits = logits[:, -1, :]

# 使用 softmax 函数将 logits 转换为概率分布
probabilities = torch.nn.functional.softmax(last_logits, dim=-1)

# 打印概率分布
print("概率分布：", probabilities)

6. 输出结果

运行上述代码后，probabilities 将是一个形状为 (1, vocab_size) 的张量，表示模型对下一个词元的概率预测。例如：

概率分布： tensor([[0.0012, 0.0023, 0.0045, ..., 0.0123, 0.0056, 0.0078]])

7. 进一步应用

在 GRPOTrainer 中，当前策略的概率分布用于计算 KL 散度和优势，从而优化模型的生成策略。通过这种方式，模型可以逐步学习生成更符合奖励函数期望的文本。

GRPOTrainer 中当前策略的概率分布的源码是怎么样的

User: GRPOTrainer 中当前策略的概率分布的源码是怎么样的
Kimi: 在 GRPOTrainer 中，当前策略的概率分布是通过模型的输出 logits 计算得到的。具体来说，模型的输出 logits 表示每个词元的未归一化对数概率，通过 softmax 函数可以将其转换为概率分布[³⁵]。

概率分布计算的源码实现

在 GRPOTrainer 的实现中，当前策略的概率分布计算通常在 compute_loss 方法中完成。以下是相关代码的简化版本，展示了如何计算当前策略的概率分布[³⁵]：

import torch
import torch.nn.functional as F

def compute_loss(self, model, inputs, return_outputs=False):
    # 提取输入数据
    prompt_ids = inputs["prompt_ids"]
    completion_ids = inputs["completion_ids"]
    attention_mask = torch.cat([inputs["prompt_mask"], inputs["completion_mask"]], dim=1)

    # 将 prompt 和 completion 拼接起来作为模型的输入
    input_ids = torch.cat([prompt_ids, completion_ids], dim=1)

    # 获取模型的输出 logits
    outputs = model(input_ids=input_ids, attention_mask=attention_mask)
    logits = outputs.logits

    # 计算当前策略的概率分布
    # 我们只关注 completion 部分的 logits
    logits_to_keep = completion_ids.size(1)
    relevant_logits = logits[:, -logits_to_keep:, :]

    # 使用 softmax 函数将 logits 转换为概率分布
    current_policy_probs = F.softmax(relevant_logits, dim=-1)

    # 计算其他部分（如 KL 散度和优势）...
    # 这里省略了其他部分的代码，只展示了当前策略概率分布的计算

    return current_policy_probs

代码解析

输入数据准备：
- prompt_ids 和 completion_ids 分别是提示文本和生成文本的 Token ID 序列。
- attention_mask 用于指示模型哪些位置是有效的输入。
模型输出：
- 将 prompt_ids 和 completion_ids 拼接起来作为模型的输入。
- 模型输出的 logits 表示每个词元的未归一化对数概率。
概率分布计算：
- 提取与 completion_ids 对应的 logits 部分（relevant_logits）。
- 使用 torch.nn.functional.softmax 将 logits 转换为概率分布。

注意事项

数值稳定性：在实际实现中，可能需要对 logits 进行一些处理（如减去最大值）以提高数值稳定性。
KL 散度计算：计算 KL 散度时，需要参考策略的概率分布。参考策略的概率分布通常通过一个固定的模型或预训练模型计算得到[³⁵]。

通过上述代码，GRPOTrainer 可以计算当前策略的概率分布，进而用于计算 KL 散度和优化模型的生成策略。

GRPOTrainer 中参考策略的概率分布的源码是怎么样的

User: GRPOTrainer 中参考策略的概率分布的源码是怎么样的
Kimi: 在 GRPOTrainer 中，参考策略的概率分布是通过一个固定的参考模型计算得到的。参考模型通常是一个预训练的语言模型，其参数在训练过程中保持不变。参考策略的概率分布用于计算 KL 散度，以确保当前策略不会偏离参考策略太远[³⁸]。

以下是 GRPOTrainer 中计算参考策略概率分布的源码实现：

源码解析

在 GRPOTrainer 的 compute_loss 方法中，参考策略的概率分布是通过输入中的 ref_per_token_logps 获取的。ref_per_token_logps 是预先计算好的，表示参考策略下每个词元的对数概率[³⁸]。

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)

    # 使用当前模型计算每个输入 token 的对数概率
    per_token_logps = self._get_per_token_logps(model, input_ids, attention_mask, logits_to_keep)

    # 获取参考策略的对数概率
    ref_per_token_logps = inputs["ref_per_token_logps"]

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

    # 计算其他部分（如优势和损失）...
    # 省略了其他部分的代码，只展示了参考策略概率分布的使用

计算参考策略概率分布

参考策略的概率分布通常是通过一个固定的参考模型计算得到的。以下是一个简化的代码示例，展示如何计算参考策略的概率分布：

def compute_ref_per_token_logps(self, ref_model, input_ids, attention_mask, logits_to_keep):
    # 使用参考模型计算每个输入 token 的 logits
    with torch.no_grad():  # 禁用梯度计算，节省内存
        ref_logits = ref_model(input_ids=input_ids, attention_mask=attention_mask).logits

    # 提取与 completion 部分对应的 logits
    ref_logits = ref_logits[:, -logits_to_keep:, :]

    # 使用 log_softmax 函数将 logits 转换为对数概率分布
    ref_per_token_logps = torch.nn.functional.log_softmax(ref_logits, dim=-1)

    # 获取每个词元的实际 ID，并计算对应的对数概率
    ref_per_token_logps = torch.gather(ref_per_token_logps, dim=-1, index=completion_ids.unsqueeze(-1)).squeeze(-1)

    return ref_per_token_logps

代码解析

参考模型的输出：
- 使用参考模型（ref_model）计算输入的 logits。ref_model 是一个预训练的语言模型，其参数在训练过程中保持不变[³⁸]。
- 提取与 completion_ids 对应的 logits 部分，因为这部分是我们关注的生成文本[³⁸]。
对数概率分布：
- 使用 torch.nn.functional.log_softmax 将 logits 转换为对数概率分布[³⁸]。
- 通过 torch.gather 获取每个词元的实际对数概率[³⁸]。