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

41 0 0

文章摘要

开源微调工具Unsloth近期发布了重大更新，进一步优化了GRPO（Group Relative Policy Optimization）训练算法，显著提升了性能并大幅降低了显存需求。此次更新使得上下文长度增加了10倍，同时所需的显存仅为原来的10%。具体而言，使用Unsloth工具，仅需5GB的显存即可训练推理模型Qwen2.5（1.5B），且不会损失准确率。这一改进使得即使是较老的GPU（如GTX 1060）也能胜任此类任务，极大地降低了硬件门槛。

在技术细节上，Unsloth通过三项关键突破实现了显存的大幅优化。首先，全新设计的内存高效线性算法将GRPO的内存使用量削减了8倍以上，节省了68.5GB的内存。其次，智能梯度checkpoint算法通过将中间激活值异步卸载到系统RAM中，节省了高达372GB的VRAM，且速度仅慢了1%。最后，与底层推理引擎（vLLM）共享相同的GPU/CUDA内存空间，进一步节省了16GB的VRAM。这些优化使得在20K的上下文长度下，Llama 3.1（8B）模型的训练显存需求从510.8GB降至54.3GB。

此外，Unsloth还对GRPO算法本身进行了深入分析和改进。研究发现，GRPO参考实现使用的是反向KL散度而非正向KL散度，且在float16混合精度下直接实现线性交叉熵会导致崩溃。通过实验，团队验证了移除detach代码会破坏训练，因此必须保留这一设计。同时，Unsloth还提供了完整的日志记录功能，并支持在vLLM中使用FP8 KV缓存，进一步优化了显存使用。

在应用层面，Unsloth的更新不仅降低了显存需求，还提升了训练效率。例如，在配备TRL+FA2的GRPO设置中，Llama 3.1（8B）在20K上下文长度下的训练显存需求减少了90%。这一改进使得更多开发者和研究者能够在资源有限的情况下进行高效的模型训练。此外，Unsloth还支持传递vLLM的SamplingParams参数，进一步增强了其灵活性和适用性。

总体而言，Unsloth的这次更新在显存优化、算法改进和功能扩展方面均取得了显著进展。通过降低硬件门槛和提升训练效率，Unsloth为开源社区和AI研究领域提供了强有力的支持，推动了GRPO算法及相关技术的广泛应用。