圖片

在大模型訓練領域，顯存一直是一個讓研究者和開發者頭疼的問題。特別是在進行長文本上下文訓練時，動輒需要幾百GB的顯存需求，這讓很多研究者望而卻步。不過最近，AI基礎設施優化團隊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：顯存優化的突破性進展

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

訓練內存成本：從414GB降至42GB

GRPO內存成本：從78.3GB降至9.8GB

推理內存開銷：從16GB降至0GB

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

圖片

2、技術創新：三重優化方案

Unsloth團隊采用了三個關鍵的技術創新來實現這一突破：

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

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

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

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

參考 GRPO 實現使用反向 KL 散度，而不是正向 KL 散度。

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

我們發現 GRPO 損失函數實現中存在其他問題——主要是在反向 KL 散度的公式表達上。

Unsloth團隊進行了 4 個實驗：

通過參考實現（紅線）進行常規 GRPO

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

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

前向 KL 散度（綠色線）

圖片

一般來說，移除 detach 確實會破壞所有訓練，所以我們必須保留它——這很可能需要更多的調查。看起來其他所有實現似乎都很相似？我們可能需要運行模型更長時間以看到不同的效果。

在所有實現中，Unsloth團隊還利用了 logsumexp 技巧

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

這項技術突破的意義遠不止于數字的優化。它意味著：

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

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

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

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

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