「古董」GPU也能跑DeepSeek同款GRPO！顯存只需1/10，上下文爆漲10倍
開源微調(diào)神器Unsloth帶著黑科技又來了：上次更新把GRPO需要的內(nèi)存見到了7GB，這次只需要5GB的VRAM，就能訓練自己的推理模型Qwen2.5（1.5B），比上次要少2GB。
這次徹底把推理模型訓練顯存打下來了！
這次把GRPO訓練推理模型的上下文變長10倍，同時需要的顯存少了90%。
使用最新的Unsloth，只要5GB顯存就能訓練自己的推理模型，而且Qwen2.5-1.5B不會損失準確率。
5GB顯存什么概念呢？
16年開始發(fā)售的GPU比如GTX 1060的顯存都有8GB。16年GTX 1060放到現(xiàn)在，堪稱電子古董！
目前，實現(xiàn)更長的上下文是GRPO面臨的最大挑戰(zhàn)之一。
與其他GRPO LoRA/QLoRA實現(xiàn)相比，即使是基于Flash Attention 2（FA2）的實現(xiàn)，Unsloth新推出的高效GRPO算法上下文長度增加了10倍，同時使用的VRAM只要10%。
在配備TRL+FA2的GRPO設(shè)置中，Llama 3.1（8B）在20K上下文長度下，訓練需要510.8GB的VRAM。
而Unsloth將VRAM減少了90%，降至僅54.3GB。
減少長上下文90%VRAM
和使用Flash Attention 2的標準實現(xiàn)相比，Unsloth使用多種技巧，巧妙地把GRPO的VRAM使用量減少了90%多！
在20K的上下文長度下，每個提示生成8次，Unsloth在Llama-3.1-8B模型上僅使用54.3GB的VRAM，而標準實現(xiàn)需要510.8GB（Unsloth減少了90%）。這一切得益于下列3項突破：
- 全新設(shè)計的內(nèi)存高效線性算法：將GRPO的內(nèi)存使用量削減了8倍以上，節(jié)省了68.5GB的內(nèi)存。借助torch.compile，在num_generatinotallow=8和20K上下文長度下，實際上還更快。
- 利用了Unsloth已發(fā)布的智能梯度checkpoint算法：將中間激活值異步卸載到系統(tǒng)RAM中，速度僅慢了1%。由于需要num_generatinotallow=8，這節(jié)省了高達372GB的VRAM。通過中間梯度累積，甚至可以進一步減少內(nèi)存使用。
- 與底層推理引擎（vLLM）共享相同的GPU/CUDA內(nèi)存空間，不像其他包中的實現(xiàn)那樣。這又節(jié)省了16GB的VRAM。
Unsloth和基于Flash Attention 2（FA2）的標準實現(xiàn)內(nèi)存比較
在典型的GRPO標準實現(xiàn)中，需要創(chuàng)建兩個大小為(8，20K)的logits來計算GRPO損失。這需要2*2字節(jié)*8（生成次數(shù)）*20K（上下文長度）*128256（詞匯表大?。?78.3GB的VRAM。
Unsloth將長上下文GRPO的內(nèi)存使用量削減了8倍，因此對于20K的上下文長度，只需要額外的9.8GBVRAM！
還需要以16位格式存儲KV緩存。Llama3.18B有32層，K和V的大小均為1024。因此，對于20K的上下文長度，內(nèi)存使用量=2*2字節(jié)*32層*20K上下文長度*1024=每個批次2.5GB。
可以將vLLM的批次大小設(shè)置為8，但為了節(jié)省VRAM，在計算中將其保持為1。否則，需要20GB來存儲KV緩存。
數(shù)學原理
分組相對策略優(yōu)化（Group Relative Policy Optimization，GRPO），出自DeepSeek去年發(fā)表的論文。
如果一生只能讀一篇DeepSeek的論文，網(wǎng)友建議選擇首次提出GRPO的DeepSeekMath論文。
論文鏈接：https://arxiv.org/abs/2402.03300
隨后在DeepSeek的論文中，利用GRPO算法創(chuàng)建了DeepSeek-R1。
發(fā)現(xiàn)的問題
在這里利用了Hugging Face的TRL GRPO實現(xiàn)。
注意到，TRL實現(xiàn)的公式如下：
其中使用的是反向KL散度（而不是正向KL散度）。β是一個設(shè)為0.04的縮放因子，A是考慮所有獎勵函數(shù)后得到的優(yōu)勢值。q是新訓練的模型，P是原始參考模型。
然后注意到，該實現(xiàn)將反向KL散度計算為：
但這真的是正確的嗎？
首先嘗試推導并整理類似項：
這意味著什么？實現(xiàn)中可能缺少一個與q（新分布項）的乘法嗎？
但這似乎是正確的，和DeepSeek-Math論文第14頁首次引入GRPO時一樣。
DeepSeek-Math論文第14頁：在損失函數(shù)中添加KL散度，正則化GRPO算法
同樣，John Schulman的博客也提到，反向KL項的無偏估計，實際上并不需要額外的q項。
鏈接地址：http://joschu.net/blog/kl-approx.html
在博客中看到：
還發(fā)現(xiàn)了一個有趣的現(xiàn)象：
torch.exp(q-q.detach()) * advantages.unsqueeze(1)這應(yīng)該等于1，對嗎？
Hugging Face的TRL GRPO實現(xiàn)
實際上，發(fā)現(xiàn)這是必要的——似乎自動梯度autograd引擎可能無法正確傳播梯度。
因此，進行了4個實驗：
- 使用參考實現(xiàn)的常規(guī)GRPO（紅線）
- 移除detach代碼（藍線）
- 按照之前討論的完整反向KL，添加額外項（黃線）
- 使用正向KL散度代替（綠線）
總體來說，移除detach顯然會破壞訓練，所以必須保留它——這很可能需要進一步調(diào)查。其他實現(xiàn)似乎也類似？可能需要運行模型更長時間，以觀察不同的效果。
在所有實現(xiàn)中，還利用了logsumexp技巧：
Unsloth高效GRPO算法
但沒想到華人工程師Horace He的線性交叉熵實現(xiàn)，帶給unsloth靈感并成功應(yīng)用于GRPO！
Horace He，在Meta從事PyTorch相關(guān)工作
實際上，unsloth發(fā)現(xiàn)了一些令人驚訝的要點：
1 GRPO參考實現(xiàn)使用的是反向KL散度，而不是正向KL散度。
2 如果不正確處理，在float16混合精度（以及float8）上直接實現(xiàn)線性交叉熵，并使用自動混合精度縮放機制，會導致崩潰。
3 發(fā)現(xiàn)了GRPO損失實現(xiàn)中的其他一些奇怪之處，主要是在反向KL散度的公式表述方面。
線性交叉商鏈接：https://gist.github.com/Chillee/22cd93e11b887db1f596ab754d60a899
其他功能
GRPO的完整日志記錄
之前，unsloth只顯示總聚合獎勵函數(shù)本身，新版本為所有獎勵函數(shù)提供完整的日志記錄詳情！
也不再需要調(diào)用函數(shù)來給GRPO打補丁了！也就是說，新版本會自動處理，可以刪除下列代碼：
from unsloth import PatchFastRL
PatchFastRL("GRPO", FastLanguageModel)vLLM推理選項
現(xiàn)在在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或其他采樣參數(shù)，也支持傳遞vLLM的SamplingParams參數(shù)中的任何內(nèi)容！
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,
)
