「古董」GPU也能跑DeepSeek同款GRPO！顯存只需1/10，上下文爆漲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_generatinotallow=8和20K上下文長度下，實際上還更快。
- 利用了Unsloth已發布的智能梯度checkpoint算法：將中間激活值異步卸載到系統RAM中，速度僅慢了1%。由于需要num_generatinotallow=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,
)
