显存狂降80%!Unsloth黑科技优化GRPO流程,让人人都能训自己的Deepseek R1
显存狂降80%!Unsloth黑科技优化GRPO流程,让人人都能训自己的Deepseek R1 by 歸藏的AI工具箱
我们知道 Deepseek R1 核心的贡献是揭示了一个“aha”时刻,在 R1-Zero 中通过使用 GRPO (Group Relative Policy Optimization)在没有人类反馈的情况下自主学会了分配更多的思考时间。
开源社区也在其他模型上复现了类似的表现,不过成本很高,比如为Qwen2.5(1.5B)实现推理也需要 160G 显存,根本不是个人可以承受的。
但是今天 Unsloth AI 优化了优化了整个 GRPO 流程,让这个流程的显存占用节省了 80% ,极限情况下你甚至可以用 8G 显卡搞定整个过程训练自己的推理模型。
注意不是对 R1 蒸馏模型微调,而是将标准模型转化为完整的推理模型!
核心要点:
凭借 15GB 的显存,Unsloth 能够将任何参数高达 150 亿的模型,如 Llama 3.1(8B)、Phi-4(14B)、Mistral(7B)或 Qwen2.5(7B),转化为推理模型。极限情况下只需要 7G 显存。
此前,GRPO 仅支持全量微调,Unsloth AI 使其能够与 QLoRA 和 LoRA 协同工作
他们还提供了Colab笔记本,直接运行就行,这个可以直接将 Llama 8B 变成推理模型:https://colab.research.google.com/github/unslothai/notebooks/blob/main/nb/Llama3.1_(8B)-GRPO.ipynb
GRPO + “aha”时刻
Deepseek 的研究人员在用纯强化学习(RL)训练 R1-Zero 时观察到了一个“顿悟时刻”。该模型在没有人类指导或预设指令的情况下,通过重新评估其初始方法,学会了延长其思考时间。
在一个测试示例中,他们仅使用 GRPO 对 Phi-4 进行了 100 步的训练,结果已显而易见。未采用 GRPO 的模型不具备思维标记,而经过 GRPO 训练的模型不仅拥有该标记,还给出了正确答案。
这种魔法可以通过 GRPO 重现,这是一种无需依赖价值函数即可高效优化响应的强化学习算法,与依赖价值函数的近端策略优化(PPO)不同。Unsloth AI 用 GRPO 训练模型,旨在使其自主发展出自我验证和搜索能力——创造一个小小的“顿悟时刻”。
工作原理:
该模型生成多组响应。
每个响应的评分基于正确性或由某种设定的奖励函数创建的其他指标,而非基于LLM奖励模型。
该组的平均分数被计算出来。
每个响应的得分与组内平均值进行比较。
该模型经过强化,以优先选择得分更高的响应。
最初,人们需要收集大量数据来填充推理过程或思维链条。但 GRPO(DeepSeek 使用的算法)或其他强化学习算法能够引导模型自动展现推理能力并生成推理轨迹。相反,他们需要创建良好的奖励函数或验证器。
例如,如果模型得出正确答案,就给它 1 分;如果某些单词拼写错误,就扣 0.1 分,以此类推!他们可以提供许多函数来奖励这一过程。
Unsloth 中的 GRPO 如何使用
如果在本地使用 GRPO 与 Unsloth,请同时“pip install diffusers”,因为它是必需的依赖项。
请至少等待 300 步,奖励才会实际增加,并请使用最新版本的 vLLM。请记住,在 Colab 上的示例仅训练了一小时,因此结果并不理想。为了获得良好的结果,需要至少训练 12 小时(这是 GRPO 的工作原理),但请记住,这并不是强制性的,可以随时停止。
建议将 GRPO 应用于参数至少为 1.5B 的模型,以确保正确生成思维标记,因为较小的模型可能无法做到。如果使用的是基础模型,请确保拥有聊天模板。GRPO 的训练损失跟踪现已直接集成到 Unsloth 中,无需再依赖如 wandb 等外部工具。
Unsloth x vLLM:吞吐量提升 20 倍,VRAM 节省 50%
现在,你可以直接在微调堆栈中使用 vLLM,这大大提高了吞吐量,并允许您同时进行模型微调和推理!在 1 块 A100 40GB 显卡上,使用 Unsloth 对 Llama 3.2 3B Instruct 的动态 4 位量化,预计可达每秒 4000 个Token左右。而在 16GB 的 Tesla T4(免费 Colab GPU)上,您也能获得每秒 300 个Token的速度。
他们还神奇地消除了同时加载 vLLM 和 Unsloth 时的双倍内存占用,使得 Llama 3.1 8B 节省了约 5GB,Llama 3.2 3B 节省了 3GB(灵感来源于 Boris)。Unsloth 最初能够在 1 块 48GB GPU 上微调 Llama 3.3 70B Instruct,其中 Llama 3.3 70B 权重占用 40GB 显存。如果不消除双倍内存占用,同时加载 Unsloth 和 vLLM 时将需要>=80GB 的显存。
但有了 Unsloth,你仍能在不到 48GB 的显存中微调并享受快速推理的优势!要使用快速推理,首先安装 vllm,并通过 fast_inference 实例化 Unsloth:
pip install unsloth vllm
from unsloth import FastLanguageModel
model, tokenizer = FastLanguageModel.from_pretrained(
model_name = "unsloth/Llama-3.2-3B-Instruct",
fast_inference = True,
)
model.fast_generate(["Hello!"])
Unsloth 在 vLLM 中的发现
vLLM 现在可以加载 Unsloth 动态 4 位量化模型。正如他们展示的 1.58 位动态 R1 GGUF 一样,动态地将某些层量化为 4 位,而其他层保持 16 位,可以在保持模型小巧的同时显著提升准确性。
他们自动选择多个参数以考虑 RAM、VRAM 效率和最大吞吐量(如分块预填充Token数、最大序列数等)。他们在 vLLM 中默认启用-O3 优化并开启前缀缓存。他们发现,在旧 GPU 上使用 Flashinfer 实际上会慢 10%。FP8 KV 缓存会使速度降低 10%,但吞吐量潜力翻倍。
他们通过解析状态字典而非从磁盘加载的方式,在 vLLM 中实现了 LoRA 的加载——这可以使您的 GRPO 训练速度提升 1.5 倍。当前一个活跃的研究领域是探索如何在 vLLM 中直接编辑 LoRA 适配器(具体方法尚不明确)。若能实现,将大幅提升速度,因为目前他们正在进行不必要的 GPU 数据传输。
vLLM 在批处理生成时会出现奇怪的 VRAM 随机峰值。他们添加了一个批处理生成函数以减少内存峰值。
整理和分析不易,觉得有帮助的话可以点个赞或者再看,感谢🙏
来源:https://unsloth.ai/blog/r1-reasoning
原文链接