You can now train your own Reasoning model locally with just 5GB VRAM!
Hey guys! Thanks so much for the support on our GRPO release 2 weeks ago! Today, we're excited to announce that you can now train your own reasoning model with just 5GB VRAM for Qwen2.5 (1.5B) - down from 7GB in the previous Unsloth release! GRPO is the algorithm behind DeepSeek-R1 and how it was trained.
The best part about GRPO is it doesn't matter if you train a small model compared to a larger model as you can fit in more faster training time compared to a larger model so the end result will be very similar! You can also leave GRPO training running in the background of your PC while you do other things!
- This is thanks to our newly derived Efficient GRPO algorithm which enables 10x longer context lengths while using 90% less VRAM vs. all other GRPO LoRA/QLoRA implementations, even those utilizing Flash Attention 2 (FA2).
- With a GRPO setup using TRL + FA2, Llama 3.1 (8B) training at 20K context length demands 510.8GB of VRAM. However, Unsloth’s 90% VRAM reduction brings the requirement down to just 54.3GB in the same setup.
- We leverage our gradient checkpointing algorithm which we released a while ago. It smartly offloads intermediate activations to system RAM asynchronously whilst being only 1% slower. This shaves a whopping 372GB VRAM since we need num_generations = 8. We can reduce this memory usage even further through intermediate gradient accumulation.
- Try our free GRPO notebook with 10x longer context: Llama 3.1 (8B) on Colab
Blog for more details on the algorithm, the Maths behind GRPO, issues we found and more: https://unsloth.ai/blog/grpo
GRPO VRAM Breakdown:
| Metric | 🦥 Unsloth | TRL + FA2 |
|---|---|---|
| Training Memory Cost (GB) | 42GB | 414GB |
| GRPO Memory Cost (GB) | 9.8GB | 78.3GB |
| Inference Cost (GB) | 0GB | 16GB |
| Inference KV Cache for 20K context (GB) | 2.5GB | 2.5GB |
| Total Memory Usage | 54.3GB (90% less) | 510.8GB |
- We also now provide full logging details for all reward functions now! Previously we only showed the total aggregated reward function itself.
- You can now run and do inference with our 4-bit dynamic quants directly in vLLM.
- Also we spent a lot of time on our Guide for everything on GRPO + reward functions/verifiers so would highly recommend you guys to read it: docs.unsloth.ai/basics/reasoning
Thank you guys once again for all the support it truly means so much to us! We also have a major release coming within the next few weeks which I know you guys have been waiting for - and we're also excited for it. 🦥