圖片

在大模型訓(xùn)練領(lǐng)域，顯存一直是一個(gè)讓研究者和開發(fā)者頭疼的問題。特別是在進(jìn)行長文本上下文訓(xùn)練時(shí)，動輒需要幾百GB的顯存需求，這讓很多研究者望而卻步。不過最近，AI基礎(chǔ)設(shè)施優(yōu)化團(tuán)隊(duì)Unsloth帶來了一個(gè)重大突破 - 他們推出的新算法可以讓GRPO訓(xùn)練所需顯存減少高達(dá)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)化的突破性進(jìn)展

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

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

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

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

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

圖片

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

Unsloth團(tuán)隊(duì)采用了三個(gè)關(guān)鍵的技術(shù)創(chuàng)新來實(shí)現(xiàn)這一突破：

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

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

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

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

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

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

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

Unsloth團(tuán)隊(duì)進(jìn)行了 4 個(gè)實(shí)驗(yàn)：

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

移除斷開代碼（藍(lán)色線條）

完整反向 KL，如前所述增加一個(gè)額外項(xiàng)（黃色線）

前向 KL 散度（綠色線）

圖片

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

在所有實(shí)現(xiàn)中，Unsloth團(tuán)隊(duì)還利用了 logsumexp 技巧

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

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

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

看完這篇文章，是不是對AI訓(xùn)練的未來更有信心了？如果你也對大模型訓(xùn)練感興趣，不妨關(guān)注Unsloth團(tuán)隊(duì)的GitHub項(xiàng)目，開啟你的AI訓(xùn)練之旅！

文章標(biāo)題：Long-context GRPO 長上下文 GRPO

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

本文轉(zhuǎn)載自 ??AI帝國??，作者： 無影寺