古董GPU也能跑DeepSeek同款GRPO!显存只需1/10,上下文爆涨10倍
ztj100 2025-04-26 22:45 40 浏览 0 评论
编辑:KingHZ
【新智元导读】开源微调神器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项。
链接地址: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 PatchFastRL
PatchFastRL("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 = 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,
)
相关推荐
- SpringBoot整合SpringSecurity+JWT
-
作者|Sans_https://juejin.im/post/5da82f066fb9a04e2a73daec一.说明SpringSecurity是一个用于Java企业级应用程序的安全框架,主要包含...
- 「计算机毕设」一个精美的JAVA博客系统源码分享
-
前言大家好,我是程序员it分享师,今天给大家带来一个精美的博客系统源码!可以自己买一个便宜的云服务器,当自己的博客网站,记录一下自己学习的心得。开发技术博客系统源码基于SpringBoot,shiro...
- springboot教务管理系统+微信小程序云开发附带源码
-
今天给大家分享的程序是基于springboot的管理,前端是小程序,系统非常的nice,不管是学习还是毕设都非常的靠谱。本系统主要分为pc端后台管理和微信小程序端,pc端有三个角色:管理员、学生、教师...
- SpringBoot+LayUI后台管理系统开发脚手架
-
源码获取方式:关注,转发之后私信回复【源码】即可免费获取到!项目简介本项目本着避免重复造轮子的原则,建立一套快速开发JavaWEB项目(springboot-mini),能满足大部分后台管理系统基础开...
- Spring Boot的Security安全控制——认识SpringSecurity!
-
SpringBoot的Security安全控制在Web项目开发中,安全控制是非常重要的,不同的人配置不同的权限,这样的系统才安全。最常见的权限框架有Shiro和SpringSecurity。Shi...
- 前同事2024年接私活已入百万,都是用这几个开源的SpringBoot项目
-
前言不得不佩服SpringBoot的生态如此强大,今天给大家推荐几款优秀的后台管理系统,小伙伴们再也不用从头到尾撸一个项目了。SmartAdmin...
- 值得学习的15 个优秀开源的 Spring Boot 学习项目
-
SpringBoot算是目前Java领域最火的技术栈了,除了书呢?当然就是开源项目了,今天整理15个开源领域非常不错的SpringBoot项目供大家学习,参考。高富帅的路上只能帮你到这里了,...
- 开发企业官网就用这个基于SpringBoot的CMS系统,真香
-
前言推荐这个项目是因为使用手册部署手册非常...
- 2021年超详细的java学习路线总结—纯干货分享
-
本文整理了java开发的学习路线和相关的学习资源,非常适合零基础入门java的同学,希望大家在学习的时候,能够节省时间。纯干货,良心推荐!第一阶段:Java基础...
- jeecg-boot学习总结及使用心得(jeecgboot简单吗)
-
jeecg-boot学习总结及使用心得1.jeecg-boot是一个真正前后端分离的模版项目,便于二次开发,使用的都是较流行的新技术,后端技术主要有spring-boot2.x、shiro、Myb...
- 后勤集团原料管理系统springboot+Layui+MybatisPlus+Shiro源代码
-
本项目为前几天收费帮学妹做的一个项目,JavaEEJSP项目,在工作环境中基本使用不到,但是很多学校把这个当作编程入门的项目来做,故分享出本项目供初学者参考。一、项目描述后勤集团原料管理系统spr...
- 白卷开源SpringBoot+Vue的前后端分离入门项目
-
简介白卷是一个简单的前后端分离项目,主要采用Vue.js+SpringBoot技术栈开发。除了用作入门练习,作者还希望该项目可以作为一些常见Web项目的脚手架,帮助大家简化搭建网站的流程。...
- Spring Security 自动踢掉前一个登录用户,一个配置搞定
-
登录成功后,自动踢掉前一个登录用户,松哥第一次见到这个功能,就是在扣扣里边见到的,当时觉得挺好玩的。自己做开发后,也遇到过一模一样的需求,正好最近的SpringSecurity系列正在连载,就结...
- 收藏起来!这款开源在线考试系统,我爱了
-
大家好,我是为广大程序员兄弟操碎了心的小编,每天推荐一个小工具/源码,装满你的收藏夹,每天分享一个小技巧,让你轻松节省开发效率,实现不加班不熬夜不掉头发,是我的目标!今天小编推荐一款基于Spr...
- Shiro框架:认证和授权原理(shiro权限认证流程)
-
优质文章,及时送达前言Shiro作为解决权限问题的常用框架,常用于解决认证、授权、加密、会话管理等场景。本文将对Shiro的认证和授权原理进行介绍:Shiro可以做什么?、Shiro是由什么组成的?举...
你 发表评论:
欢迎- 一周热门
- 最近发表
-
- SpringBoot整合SpringSecurity+JWT
- 「计算机毕设」一个精美的JAVA博客系统源码分享
- springboot教务管理系统+微信小程序云开发附带源码
- SpringBoot+LayUI后台管理系统开发脚手架
- Spring Boot的Security安全控制——认识SpringSecurity!
- 前同事2024年接私活已入百万,都是用这几个开源的SpringBoot项目
- 值得学习的15 个优秀开源的 Spring Boot 学习项目
- 开发企业官网就用这个基于SpringBoot的CMS系统,真香
- 2021年超详细的java学习路线总结—纯干货分享
- jeecg-boot学习总结及使用心得(jeecgboot简单吗)
- 标签列表
-
- idea eval reset (50)
- vue dispatch (70)
- update canceled (42)
- order by asc (53)
- spring gateway (67)
- 简单代码编程 贪吃蛇 (40)
- transforms.resize (33)
- redisson trylock (35)
- 卸载node (35)
- np.reshape (33)
- torch.arange (34)
- npm 源 (35)
- vue3 deep (35)
- win10 ssh (35)
- vue foreach (34)
- idea设置编码为utf8 (35)
- vue 数组添加元素 (34)
- std find (34)
- tablefield注解用途 (35)
- python str转json (34)
- java websocket客户端 (34)
- tensor.view (34)
- java jackson (34)
- vmware17pro最新密钥 (34)
- mysql单表最大数据量 (35)