圖片

在大模型訓練領域，顯存一直是一個讓研究者和開發(fā)者頭疼的問題。特別是在進行長文本上下文訓練時，動輒需要幾百GB的顯存需求，這讓很多研究者望而卻步。不過最近，AI基礎設施優(yōu)化團隊Unsloth帶來了一個重大突破 - 他們推出的新算法可以讓GRPO訓練所需顯存減少高達90%！文章公布了Llama3.1(8B) GRPO在Colab上notebook，見：https://colab.research.google.com/github/unslothai/notebooks/blob/main/nb/Llama3.1_(8B)-GRPO.ipynb

1、從510GB到54GB：顯存優(yōu)化的突破性進展

在傳統(tǒng)的GRPO訓練方案中，要訓練一個支持20K上下文長度的Llama 3.1（8B）模型，需要高達510.8GB的顯存。這個量級的顯存需求，即便是頂配的訓練服務器也難以滿足。而Unsloth團隊通過其創(chuàng)新的算法優(yōu)化，將這一需求降低到了驚人的54.3GB，這意味著：

訓練內(nèi)存成本：從414GB降至42GB

GRPO內(nèi)存成本：從78.3GB降至9.8GB

推理內(nèi)存開銷：從16GB降至0GB

20K上下文的推理KV緩存：保持在2.5GB

圖片

2、技術創(chuàng)新：三重優(yōu)化方案

Unsloth團隊采用了三個關鍵的技術創(chuàng)新來實現(xiàn)這一突破：

全新的線性算法：團隊為GRPO開發(fā)了一個全新的內(nèi)存高效線性算法，這個優(yōu)化alone就減少了68.5GB的內(nèi)存使用。更令人驚喜的是，通過torch.compile的協(xié)助，這個算法在性能上還實現(xiàn)了提速。

智能梯度檢查點：通過將中間激活值異步卸載到系統(tǒng)RAM，在僅損失1%性能的情況下節(jié)省了驚人的372GB顯存。這個優(yōu)化特別適用于需要多次生成的場景。    

共享內(nèi)存空間：與其他實現(xiàn)不同，Unsloth可以與底層推理引擎(vLLM)共享GPU/CUDA內(nèi)存空間，這又節(jié)省了16GB顯存。

Unsloth團隊從 Horace 的線性交叉熵實現(xiàn)中獲得了靈感，并成功使其適用于 GRPO！實際上，我們發(fā)現(xiàn)了一些令人驚訝的點：

參考 GRPO 實現(xiàn)使用反向 KL 散度，而不是正向 KL 散度。

天真地實現(xiàn)浮點 16 混合精度（以及浮點 8）上的線性交叉熵，如果沒有正確處理，將因自動混合精度縮放機制而崩潰。

我們發(fā)現(xiàn) GRPO 損失函數(shù)實現(xiàn)中存在其他問題——主要是在反向 KL 散度的公式表達上。

Unsloth團隊進行了 4 個實驗：

通過參考實現(xiàn)（紅線）進行常規(guī) GRPO

移除斷開代碼（藍色線條）

完整反向 KL，如前所述增加一個額外項（黃色線）

前向 KL 散度（綠色線）

圖片

一般來說，移除 detach 確實會破壞所有訓練，所以我們必須保留它——這很可能需要更多的調(diào)查?？雌饋砥渌袑崿F(xiàn)似乎都很相似？我們可能需要運行模型更長時間以看到不同的效果。

在所有實現(xiàn)中，Unsloth團隊還利用了 logsumexp 技巧

3、實踐意義：讓更多開發(fā)者參與AI訓練

這項技術突破的意義遠不止于數(shù)字的優(yōu)化。它意味著：

• 降低硬件門檻：原本需要多卡集群才能完成的訓練任務，現(xiàn)在用單卡就能搞定。比如Qwen2.5 (1.5B)的訓練現(xiàn)在只需要5GB顯存！
• 提升研究效率：研究人員可以更快速地進行實驗驗證，加快模型迭代速度。
• 擴大應用場景：更多的小團隊和個人開發(fā)者現(xiàn)在也能嘗試大模型訓練，這將極大促進AI技術的普及和創(chuàng)新。

看完這篇文章，是不是對AI訓練的未來更有信心了？如果你也對大模型訓練感興趣，不妨關注Unsloth團隊的GitHub項目，開啟你的AI訓練之旅！

文章標題：Long-context GRPO 長上下文 GRPO

文章鏈接：?https://unsloth.ai/blog/grpo

本文轉載自 ??AI帝國??，作者： 無影寺