摘要：开源微调神器Unsloth带着黑科技又来了：短短两周后，再次优化DeepSeek-R1同款GRPO训练算法，上下文变长10倍，而显存只需原来的1/10！
【导读】 开源微调神器Unsloth带着黑科技又来了：短短两周后，再次优化DeepSeek-R1同款GRPO训练算法，上下文变长10倍，而显存只需原来的1/10！
开源微调神器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项。 
在博客中看到： 
还发现了一个有趣的现象： 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，在Meta从事PyTorch相关工作 实际上，unsloth发现了一些令人惊讶的要点： 1 GRPO参考实现使用的是反向KL散度，而不是正向KL散度。 2 如果不正确处理，在float16混合精度（以及float8）上直接实现线性交叉熵，并使用自动混合精度缩放机制，会导致崩溃。 3 发现了GRPO损失实现中的其他一些奇怪之处，主要是在反向KL散度的公式表述方面。 
线性交叉商链接：https://gist.github.com/Chillee/22cd93e11b887db1f596ab754d60a899 其他功能
之前，unsloth只显示总聚合奖励函数本身，新版本为所有奖励函数提供完整的日志记录详情！ 
也不再需要调用函数来给GRPO打补丁了！也就是说，新版本会自动处理，可以删除下列代码： from unsloth import PatchFastRL现在在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,...)) 来源：东窗史谈
免责声明：本站系转载，并不代表本网赞同其观点和对其真实性负责。如涉及作品内容、版权和其它问题，请在30日内与本站联系，我们将在第一时间删除内容!
