第28章 长上下文 GRPO 训练（R1 推理模型）

Unsloth 下载 ☰ 下载 博客 长上下文 GRPO 2025年2月20日 • 作者：Daniel & Michael 2025年2月20日 • 作者：Daniel & Michael 现在，你只需 5GB VRAM 即可训练 Qwen2.5 (1.5B) 推理模型——相比两周前我们发布的 GRPO 版本所需的 7GB，内存占用大幅降低！ 我们强烈推荐阅读我们的指南，了解关于 GRPO + 奖励函数/验证器的所有内容。 目前，实现更长的上下文长度是 GRPO 面临的最大挑战之一。我们新推导的 Unsloth 高效 GRPO 算法，在内存占用上比所有其他 GRPO LoRA/QLoRA 实现（包括使用 Flash Attention 2 (FA2) 的）低 90%，同时支持 10 倍更长的上下文长度。 在使用 TRL + FA2 的 GRPO 配置中，Llama 3.1 (8B) 在 20K 上下文长度下训练需要 510.8GB VRAM。然而，Unsloth 的 90% 内存缩减将该需求降至仅 54.3GB（相同配置下）。 试用我们免费的进阶版 GRPO notebook，支持 10 倍更长的上下文：Llama 3.1 (3B) 在 Colab 上运行 查看我们其他模型（如 Phi-4）的 GRPO notebooks。 ❤️ 后记：如果你喜欢我们的工作，别忘了给我们点个 ⭐Star：github.com/unslothai/unsloth🦥 长上下文场景节省 90% 显存 用 Unsloth 做 GRPO 时，我们通过多种技巧，相比使用 Flash Attention 2 的标准实现，能智能地把显存占用降低 90% 以上。比如在 20K 上下文长度、每个 prompt 生成 8 次的设置下，Unsloth 跑 Llama 3.1 8B 只需 54.3GB 显存，而标准实现要 510.8GB（节省 90%）。 我们为 GRPO 新设计的高效线性内存算法把内存占用砍掉 8 倍以上——在 num_generations = 8、上下文长度 20K 的条件下省下 68.5GB 内存，而且借助 torch.compile 实际上还更快了。 我们还利用了之前发布的 Unsloth 智能梯度检查点算法：异步地把中间激活值卸载到系统内存，速度只慢 1%。由于需要 num_generations = 8，这一项就省下了高达 372GB 显存。通过梯度中间累积，还能进一步降低内存占用。 与其它框架的实现不同，Unsloth 和底层推理引擎（vLLM）共享同一块 GPU / CUDA 内存空间，这一项又省下 16GB 显存。

指标	🦥 Unsloth	TRL + FA2
训练内存开销（GB）	42GB	414GB
GRPO 内存开销（GB）	9.8GB	78.3GB
推理开销（GB）	0GB	16GB
20K 上下文的推理 KV Cache（GB）	2.5GB	2.5GB
总内存占用	54.3GB（节省 90%）	510.8GB

在典型的标准 GRPO 实现中，需要创建 2 个大小为 (8, 20K) 的 logits 张量来计算 GRPO loss，这要占用 2 * 2 字节 * 8（生成次数） * 20K（上下文长度） * 128256（词表大小） = 78.3GB 显存。

Unsloth 把长上下文 GRPO 的内存占用降低了 8 倍，20K 上下文长度下只需额外 9.8GB 显存！

KV 缓存也需要以 16bit 精度存储。Llama 3.1 8B 拥有 32 层，K 和 V 的维度均为 1024。因此，20K 上下文长度的显存占用为：2 × 2 字节 × 32 层 × 20K 上下文长度 × 1024 = 每个批次 2.5GB。我们将 vLLM 的批处理大小设置为 8，但在计算时将其保持为 1 以节省显存，否则 KV 缓存将需要 20GB 显存。🦥 Unsloth 高效 GRPO 算法 我们受到 Horace He 的线性交叉熵实现的启发，并成功使其适用于 GRPO！在这个过程中，我们发现了一些令人意外的点： 参考的 GRPO 实现使用的是反向 KL 散度，而非前向 KL 散度。 如果未妥善处理，在 float16 混合精度（以及 float8）下通过自动混合精度缩放机制来实现线性交叉熵，将会导致失败。 我们还在 GRPO 损失的实现中发现了一些其他怪癖——主要涉及反向 KL 散度的表述方式。 💡 GRPO 的数学原理及发现的问题 GRPO 算法最早在 2024 年 2 月至 4 月间由 DeepSeek 在其数学论文中提出。随后，DeepSeek 在其论文中提到的 DeepSeek R1 中应用了 GRPO 算法。 此处我们采用 Hugging Face 的 TRL GRPO 实现。我们注意到 TRL 执行的是以下公式： L = \frac{1}{n}\sum{\beta D_{\text{KL}}}\big( q \,\|\, p \big) + A 这里使用的是反向 KL 散度（而非前向 KL 散度）。Beta 是缩放因子，设为 0.04；A 是考虑所有奖励函数后获得的优势值。 Q 是训练后的新模型，P 是原始参考模型。 值得注意的是，我们发现该实现中反向 KL 散度的计算公式为： $$ \begin{align} p &= \sigma (f(x)) \\ q &= \sigma (f'(x)) \\ D_{\text{KL}}\big( q \,\|\, p \big)_i &= \exp(\log(p)-\log(q))-(\log(p)-\log(q)) - 1 \\ &= \exp\bigg(log\bigg(\frac{p}{q}\bigg)\bigg)-log\bigg(\frac{p}{q}\bigg) - 1 \\ &= \frac{p}{q} - log\bigg(\frac{p}{q}\bigg) - 1 \end{align} $$ 但这在数学上真的正确吗？我们先尝试推导，合并同类项： $$ \begin{align} D_{\text{KL}}\big( q \,\|\, p \big) &= \sum q \bigg[ \frac{p}{q} - \log{\bigg(\frac{p}{q}\bigg)} - 1 \bigg] \\ &= \sum q \frac{p}{q} - \sum q \log{\bigg(\frac{p}{q}\bigg)} - \sum q \\ &= \sum p - \sum q \log{\bigg(\frac{p}{q}\bigg)} - 1 \\ &= 1 - \sum q \log{\bigg(\frac{p}{q}\bigg)} - 1 \\ &= - \sum q \log{\bigg(\frac{p}{q}\bigg)} \\ D_{\text{KL}}\big( q \,\|\, p \big) &= \sum q \log{\bigg(\frac{q}{p}\bigg)} \\ \end{align} $$ 这意味着，该实现可能缺少了 $q$（新分布项）的乘积？但 DeepSeek Math 论文在第 14 页首次引入 GRPO 时似乎认为这是正确的。同样，John Schulman 的博客也指出，反向 KL 项的无偏估计量实际上并不需要额外的 $q$ 项。博客中给出的公式如下： $$ \begin{align} r &= \frac{p(x)}{q(x)} \\ \text{KL}[q, p] &= (r-1)-\log{r} \\ &= \frac{p}{q} - 1 - \log{\frac{p}{q}} \end{align} $$ 我们还发现了一个有趣的代码片段：`torch.exp(q - q.detach()) * advantages.unsqueeze(1)`
这个应该求值后等于 1，对吧？
我们实际发现这一步是必要的——似乎 autograd 引擎在梯度传播上有问题。

于是我们做了 4 组实验：用参考实现跑常规 GRPO（红线）；去掉 detach 代码（蓝线）；如前所述，在完整 reverse KL 上加一个额外项（黄线）；改用 forward KL 散度（绿线）。总体来看，去掉 detach 会彻底破坏训练，所以必须保留它——这一点大概率还需要进一步排查。其他几种实现看起来效果相近？也许需要把模型跑得更久一些才能看出差异。

在所有实现中，我们还会使用 logsumexp 技巧： log σ(x) = log { exp(x) / ∑ exp(x) } = x − log ∑ exp(x) = x − logsumexp(x) \begin{align} \log\sigma(x) = \log{\frac{\exp(x)}{\sum{\exp(x)}}} &= x - \log\sum{\exp(x)} \\ &= x - \text{logsumexp}(x) \end{align} 📈 GRPO 的完整日志记录 我们为所有奖励函数提供了完整的日志细节！此前，我们只展示聚合后的总奖励函数。 现在，你也不需要调用函数来给 GRPO 打补丁了！也就是说，删除顶部的这段代码（我们会自动处理）：from unsloth import PatchFastRL

PatchFastRL("GRPO", FastLanguageModel)🖥️ vLLM 推理选项 我们还支持在 vLLM 中使用 FP8 KV 缓存，这能让更新的 GPU（RTX 3090、A100 及更新型号）将 KV 缓存空间占用降低 50%。
model, tokenizer = FastLanguageModel.from_pretrained(
model_name = "meta-llama/meta-Llama-3.1-8B-Instruct",
max_seq_length = max_seq_length,
load_in_4bit = True, # LoRA 16bit 请设为 False
fast_inference = True, # 启用 vLLM 高速推理
max_lora_rank = lora_rank,
gpu_memory_utilization = 0.6, # 若内存不足请调低此值
float8_kv_cache = True, # 启用 float8 KV 缓存
) 如果你想使用 min_p = 0.1 或其他 vLLM 采样参数，我们也支持直接传入 vLLM 的 SamplingParams 中任意参数！
max_prompt_length = 256
from trl import GRPOConfig, GRPOTrainer
from unsloth import vLLMSamplingParams
vllm_sampling_params = vLLMSamplingParams(
min_p = 0.1,
seed = 3407,
...
)
training_args = GRPOConfig(
...
vllm_sampling_params = vllm_sampling_params,
temperature = 1.5,
)✨ 其他更新 🦥 在 vLLM 中直接运行 Unsloth 动态 4-bit 量化 现在可以直接在 vLLM 中运行并执行我们动态量化的推理。这得益于我们为 vLLM 仓库提交并被接受的 PR。点击这里查看示例和基准测试，了解我们的动态量化如何大幅提升精度，远超标准 4-bit 量化。🚀 运行 Perplexity 的 R1-1776 现在你可以下载我们为 Perplexity AI 新推出的 R1-1776 模型制作的动态 GGUF 文件。该模型是对 DeepSeek-R1 的微调版本，在保留推理能力的同时去除了所有审查限制。你可以在自己的设备上本地运行！🐱 GitHub Universe 访谈 今年十月，在 GitHub 2024 Universe 大会上，我们接受了一段精彩的访谈，视频现已发布！我们聊了来自澳大利亚的背景、Unsloth 的构建过程，以及大家有多么了不起等等。请在 YouTube 上观看 💕 致谢 感谢 Eyera、Edd 和 Keith 再次协助本次版本发布。非常感谢每一位使用并分享 Unsloth 的用户，我们深表感激。 🙏

一如既往，欢迎加入我们的 Reddit 页面和 Discord 服务器，获取帮助或仅仅表达支持！你也可以在 Twitter 和通讯中关注我们。感谢阅读！Daniel & Michael Han 🦥
2025 年 2 月 20 日 全面了解 GRPO 了解更多