我要投稿

「古董」GPU也能跑DeepSeek同款GRPO！显存只需1/10，上下文爆涨10倍

发布日期：2025-03-11 04:33:54 浏览次数： 2168

作者：新智元

微信搜一搜，关注“新智元”

开源微调神器Unsloth带着黑科技又来了：上次更新把GRPO需要的内存见到了7GB，这次只需要5GB的VRAM，就能训练自己的推理模型Qwen2.5（1.5B），比上次要少2GB。

这次彻底把推理模型训练显存打下来了！

这次把GRPO训练推理模型的上下文变长10倍，同时需要的显存少了90%。

使用最新的Unsloth，只要5GB显存就能训练自己的推理模型，而且Qwen2.5-1.5B不会损失准确率。

5GB显存什么概念呢？

16年开始发售的GPU比如GTX 1060的显存都有8GB。16年GTX 1060放到现在，堪称电子古董！

目前，实现更长的上下文是GRPO面临的最大挑战之一。

与其他GRPO LoRA/QLoRA实现相比，即使是基于Flash Attention 2（FA2）的实现，Unsloth新推出的高效GRPO算法上下文长度增加了10倍，同时使用的VRAM只要10%。

在配备TRL+FA2的GRPO设置中，Llama 3.1（8B）在20K上下文长度下，训练需要510.8GB的VRAM。

而Unsloth将VRAM减少了90%，降至仅54.3GB。

减少长上下文90%VRAM

和使用Flash Attention 2的标准实现相比，Unsloth使用多种技巧，巧妙地把GRPO的VRAM使用量减少了90%多！

在20K的上下文长度下，每个提示生成8次，Unsloth在Llama-3.1-8B模型上仅使用54.3GB的VRAM，而标准实现需要510.8GB（Unsloth减少了90%）。这一切得益于下列3项突破：

全新设计的内存高效线性算法：将GRPO的内存使用量削减了8倍以上，节省了68.5GB的内存。借助torch.compile，在num_generations=8和20K上下文长度下，实际上还更快。
利用了Unsloth已发布的智能梯度checkpoint算法：将中间激活值异步卸载到系统RAM中，速度仅慢了1%。由于需要num_generations=8，这节省了高达372GB的VRAM。通过中间梯度累积，甚至可以进一步减少内存使用。
与底层推理引擎（vLLM）共享相同的GPU/CUDA内存空间，不像其他包中的实现那样。这又节省了16GB的VRAM。

Unsloth和基于Flash Attention 2（FA2）的标准实现内存比较

在典型的GRPO标准实现中，需要创建两个大小为(8，20K)的logits来计算GRPO损失。这需要2*2字节*8（生成次数）*20K（上下文长度）*128256（词汇表大小）=78.3GB的VRAM。

Unsloth将长上下文GRPO的内存使用量削减了8倍，因此对于20K的上下文长度，只需要额外的9.8GBVRAM！

还需要以16位格式存储KV缓存。Llama3.18B有32层，K和V的大小均为1024。因此，对于20K的上下文长度，内存使用量=2*2字节*32层*20K上下文长度*1024=每个批次2.5GB。

可以将vLLM的批次大小设置为8，但为了节省VRAM，在计算中将其保持为1。否则，需要20GB来存储KV缓存。

数学原理

分组相对策略优化（Group Relative Policy Optimization，GRPO），出自DeepSeek去年发表的论文。

如果一生只能读一篇DeepSeek的论文，网友建议选择首次提出GRPO的DeepSeekMath论文。

论文链接：https://arxiv.org/abs/2402.03300

随后在DeepSeek的论文中，利用GRPO算法创建了DeepSeek-R1。

发现的问题

在这里利用了Hugging Face的TRL GRPO实现。

注意到，TRL实现的公式如下：

其中使用的是反向KL散度（而不是正向KL散度）。β是一个设为0.04的缩放因子，A是考虑所有奖励函数后得到的优势值。q是新训练的模型，P是原始参考模型。

然后注意到，该实现将反向KL散度计算为：

但这真的是正确的吗？

首先尝试推导并整理类似项：

这意味着什么？实现中可能缺少一个与q（新分布项）的乘法吗？

但这似乎是正确的，和DeepSeek-Math论文第14页首次引入GRPO时一样。

DeepSeek-Math论文第14页：在损失函数中添加KL散度，正则化GRPO算法

同样，John Schulman的博客也提到，反向KL项的无偏估计，实际上并不需要额外的q项。

链接地址：http://joschu.net/blog/kl-approx.html

在博客中看到：

还发现了一个有趣的现象：

torch.exp(q-q.detach()) * advantages.unsqueeze(1)

这应该等于1，对吗？

Hugging Face的TRL GRPO实现

实际上，发现这是必要的——似乎自动梯度autograd引擎可能无法正确传播梯度。

因此，进行了4个实验：

使用参考实现的常规GRPO（红线）
移除detach代码（蓝线）
按照之前讨论的完整反向KL，添加额外项（黄线）
使用正向KL散度代替（绿线）

总体来说，移除detach显然会破坏训练，所以必须保留它——这很可能需要进一步调查。其他实现似乎也类似？可能需要运行模型更长时间，以观察不同的效果。

在所有实现中，还利用了logsumexp技巧：

Unsloth高效GRPO算法

但没想到华人工程师Horace He的线性交叉熵实现，带给unsloth灵感并成功应用于GRPO！

Horace He，在Meta从事PyTorch相关工作

实际上，unsloth发现了一些令人惊讶的要点：

1 GRPO参考实现使用的是反向KL散度，而不是正向KL散度。

2 如果不正确处理，在float16混合精度（以及float8）上直接实现线性交叉熵，并使用自动混合精度缩放机制，会导致崩溃。

3 发现了GRPO损失实现中的其他一些奇怪之处，主要是在反向KL散度的公式表述方面。

线性交叉商链接：https://gist.github.com/Chillee/22cd93e11b887db1f596ab754d60a899

其他功能

GRPO的完整日志记录

之前，unsloth只显示总聚合奖励函数本身，新版本为所有奖励函数提供完整的日志记录详情！

也不再需要调用函数来给GRPO打补丁了！也就是说，新版本会自动处理，可以删除下列代码：

from unsloth import PatchFastRLPatchFastRL("GRPO", FastLanguageModel)

vLLM推理选项

现在在vLLM中还能使用FP8 KV缓存，这可以在较新的GPU（RTX 3090、A100及更新型号）上将KV缓存空间使用量减少2倍。

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, fast_inference = True, max_lora_rank = lora_rank,gpu_memory_utilization = 0.6, float8_kv_cache = True, )

如果想在vLLM中使用min_p=0.1或其他采样参数，也支持传递vLLM的SamplingParams参数中的任何内容！

max_prompt_length = 256from trl import GRPOConfig, GRPOTrainerfrom unsloth import vLLMSamplingParamsvllm_sampling_params = vLLMSamplingParams(min_p = 0.1,seed = 3407,...)training_args = GRPOConfig(...vllm_sampling_params = vllm_sampling_params,temperature = 1.5,)