RTX 3080 移動版能訓(xùn)練哪種大模型？本文為那些 GPU 資源有限時使用 GRPO 訓(xùn)練的開發(fā)者提供了寶貴的指導(dǎo)。
自 DeepSeek-R1 發(fā)布以來，群組相對策略優(yōu)化（GRPO）因其有效性和易于訓(xùn)練而成為大型語言模型強化學(xué)習(xí)的熱門話題。R1 論文展示了如何使用 GRPO 從遵循 LLM（DeepSeek-v3）的基本指令轉(zhuǎn)變?yōu)橥评砟Ｐ停―eepSeek-R1）。
GRPO 是一種在線學(xué)習(xí)算法（online learning algorithm），它通過使用訓(xùn)練過程中由訓(xùn)練模型自身生成的數(shù)據(jù)來進行迭代改進。GRPO 的目標(biāo)是最大化生成補全（completions）的優(yōu)勢函數(shù)（advantage），同時確保模型保持在參考策略（reference policy）附近。
本文的目的是幫你節(jié)省一些時間，讓你根據(jù)硬件預(yù)算選擇合適的模型大小。在開始微調(diào)時，你必須做出的重要決定是選擇模型大小，以及你是執(zhí)行完全微調(diào)還是參數(shù)高效微調(diào)（PEFT）。
文章作者來自 AI 公司 Oxen.ai 的 CEO Greg Schoeninger。
原文鏈接：https://www.oxen.ai/blog/grpo-vram-requirements-for-the-gpu-poor
作者表示，他發(fā)現(xiàn) trl 庫中已經(jīng)有一個易于使用的 GRPO 實現(xiàn)，便立刻開始了訓(xùn)練，使用的硬件是配備了 16GB 顯存的 Nvidia GeForce RTX 3080 的小型筆記本電腦。正如大家可能遇到的問題，作者發(fā)現(xiàn)示例代碼中的參數(shù)設(shè)置導(dǎo)致了一個巨大的顯存不足（OOM，out of memory ）錯誤。
torch
.
OutOfMemoryError
:
CUDA
out
of memory
.
Tried
to allocate
1.90
GiB
.
GPU
0
has a total capacity of
15.73
GiB
of which
1.28
GiB
is
free
.
Including
non
-
PyTorch
memory
,
this
process has
14.43
GiB
memory
in
use
.
Of
the allocated memory
11.82
GiB
is
allocated
by
PyTorch
,
and
2.41
GiB
is
reserved
by
PyTorch
but unallocated
.
If
reserved but unallocated memory
is
large
try
setting PYTORCH_CUDA_ALLOC_CONF
=
expandable_segments
:
True
to avoid fragmentation
.
See
documentation
for
Memory
Management
(
https
:
//pytorch.org/docs/stable/notes/cuda.html#environment-variables)
實際使用情況
作者表示，他們進行了一系列實驗，以確定訓(xùn)練各種大小的模型所需的顯存（VRAM）要求。參數(shù)數(shù)量從 5 億到 140 億不等，他們比較了權(quán)重的完全微調(diào)與參數(shù)高效微調(diào)（使用 LoRA），所有訓(xùn)練運行都在英偉達 H100 上完成，因此這里的 OOM 意味著 >80GB 的 VRAM。
在表格中，你可以找到 GSM8K 數(shù)據(jù)集上訓(xùn)練的前 100 步中的峰值內(nèi)存使用情況。用于實驗的模型是：
所有實驗均使用 Shadeform 的 GPU 市場完成，因此每次實驗只需要花費幾美元 H100。
實驗結(jié)果表明，內(nèi)存需求隨著模型大小和訓(xùn)練方式的不同而顯著變化。例如，全參數(shù)微調(diào)比 PEFT 需要更多的內(nèi)存。
為什么 GRPO 對內(nèi)存需求較高
這要從 GRPO 的原理說起，這是它的流程圖。
GRPO 對內(nèi)存需求較高的原因在于，其內(nèi)部涉及多個模型，并且在訓(xùn)練數(shù)據(jù)中每個查詢會產(chǎn)生多個輸出。上圖中的策略模型、參考模型和獎勵模型各自都是一個需要進行推理的 LLM。（盡管從技術(shù)上講，獎勵模型可能不需要參數(shù)化，可以只是一個 Python 函數(shù)或正則表達式，但不影響 GRPO 對內(nèi)存的高需求。）
為什么 8-Bit 優(yōu)化和梯度檢查點有助于減少內(nèi)存占用？
通常來講，訓(xùn)練一個大型語言模型需要在內(nèi)存中存儲三種主要類型的信息：模型參數(shù)、模型學(xué)習(xí)所需的梯度、優(yōu)化器的跟蹤數(shù)據(jù)。
對上述內(nèi)容我們可以這樣理解：如果模型的參數(shù)占用了 X 的空間，那么梯度也會占用大約相同的空間。然后，像 AdamW 這樣的優(yōu)化器需要更多的空間，因為它們就像一個記錄員，跟蹤最近的更新歷史，以便更好地決定未來的優(yōu)化。
為了減輕這種內(nèi)存負擔(dān)，通常采用兩種技術(shù)：
首先，可以使用像 AdamW 這樣的 8-bit 優(yōu)化器版本，它們能更高效地存儲跟蹤數(shù)據(jù)，同時仍保持良好的性能 —— 類似于壓縮照片可以節(jié)省空間，同時保留大部分圖像質(zhì)量；
其次，使用梯度檢查點技術(shù)，這就像在訓(xùn)練過程中拍攝快照，而不是記錄所有內(nèi)容。雖然這會使訓(xùn)練速度減慢約 20-30%，但它顯著減少了內(nèi)存使用。
結(jié)合這些技術(shù)，即使對 GPU 資源有限的人來說，也能夠訓(xùn)練更大的模型。
代碼示例
像 trl 這樣的庫已經(jīng)開始支持 GRPO，使得微調(diào)由 transformers 構(gòu)成的 LLM 變得非常簡單。代碼也非常簡潔，只需將訓(xùn)練器替換為 GRPOTrainer 并定義一些獎勵即可。GRPO 的最小代碼量大約只有 99 行，如果你使用的是像 meta-llama/Llama-3.2-1B-Instruct 這樣的小型模型和像 openai/GSM8K 這樣的數(shù)據(jù)集，可以非常快速地啟動。