打開網(wǎng)易新聞 查看精彩圖片

選自oxen.ai

作者:Greg Schoeninger

編譯:陳陳、澤南

RTX 3080 移動(dòng)版能訓(xùn)練哪種大模型?本文為那些 GPU 資源有限時(shí)使用 GRPO 訓(xùn)練的開發(fā)者提供了寶貴的指導(dǎo)。

自 DeepSeek-R1 發(fā)布以來,群組相對(duì)策略優(yōu)化(GRPO)因其有效性和易于訓(xùn)練而成為大型語言模型強(qiáng)化學(xué)習(xí)的熱門話題。R1 論文展示了如何使用 GRPO 從遵循 LLM(DeepSeek-v3)的基本指令轉(zhuǎn)變?yōu)橥评砟P停―eepSeek-R1)。

GRPO 是一種在線學(xué)習(xí)算法(online learning algorithm),它通過使用訓(xùn)練過程中由訓(xùn)練模型自身生成的數(shù)據(jù)來進(jìn)行迭代改進(jìn)。GRPO 的目標(biāo)是最大化生成補(bǔ)全(completions)的優(yōu)勢(shì)函數(shù)(advantage),同時(shí)確保模型保持在參考策略(reference policy)附近。

打開網(wǎng)易新聞 查看精彩圖片

本文的目的是幫你節(jié)省一些時(shí)間,讓你根據(jù)硬件預(yù)算選擇合適的模型大小。在開始微調(diào)時(shí),你必須做出的重要決定是選擇模型大小,以及你是執(zhí)行完全微調(diào)還是參數(shù)高效微調(diào)(PEFT)。

文章作者來自 AI 公司 Oxen.ai 的 CEO Greg Schoeninger。

打開網(wǎng)易新聞 查看精彩圖片

原文鏈接:https://www.oxen.ai/blog/grpo-vram-requirements-for-the-gpu-poor

作者表示,他發(fā)現(xiàn) trl 庫中已經(jīng)有一個(gè)易于使用的 GRPO 實(shí)現(xiàn),便立刻開始了訓(xùn)練,使用的硬件是配備了 16GB 顯存的 Nvidia GeForce RTX 3080 的小型筆記本電腦。正如大家可能遇到的問題,作者發(fā)現(xiàn)示例代碼中的參數(shù)設(shè)置導(dǎo)致了一個(gè)巨大的顯存不足(OOM,out of memory )錯(cuò)誤。

  1. torch
  2. OutOfMemoryError
  3. CUDA
  4. out
  5. of memory
  6. Tried
  7. to allocate
  8. 1.90
  9. GiB
  10. GPU
  11. 0
  12. has a total capacity of
  13. GiB
  14. of which
  15. 1.28
  16. GiB
  17. is
  18. free
  19. Including
  20. non
  21. PyTorch
  22. memory
  23. this
  24. process has
  25. GiB
  26. memory
  27. in
  28. use
  29. Of
  30. the allocated memory
  31. GiB
  32. is
  33. allocated
  34. by
  35. PyTorch
  36. and
  37. 2.41
  38. GiB
  39. is
  40. reserved
  41. by
  42. PyTorch
  43. but unallocated
  44. If
  45. reserved but unallocated memory
  46. is
  47. large
  48. try
  49. setting PYTORCH_CUDA_ALLOC_CONF
  50. expandable_segments
  51. True
  52. to avoid fragmentation
  53. See
  54. documentation
  55. for
  56. Memory
  57. Management
  58. //pytorch.org/docs/stable/notes/cuda.html#environment-variables)

實(shí)際使用情況

作者表示,他們進(jìn)行了一系列實(shí)驗(yàn),以確定訓(xùn)練各種大小的模型所需的顯存(VRAM)要求。參數(shù)數(shù)量從 5 億到 140 億不等,他們比較了權(quán)重的完全微調(diào)與參數(shù)高效微調(diào)(使用 LoRA),所有訓(xùn)練運(yùn)行都在英偉達(dá) H100 上完成,因此這里的 OOM 意味著 >80GB 的 VRAM。

打開網(wǎng)易新聞 查看精彩圖片

在表格中,你可以找到 GSM8K 數(shù)據(jù)集上訓(xùn)練的前 100 步中的峰值內(nèi)存使用情況。用于實(shí)驗(yàn)的模型是:

打開網(wǎng)易新聞 查看精彩圖片

所有實(shí)驗(yàn)均使用 Shadeform 的 GPU 市場(chǎng)完成,因此每次實(shí)驗(yàn)只需要花費(fèi)幾美元 H100。

實(shí)驗(yàn)結(jié)果表明,內(nèi)存需求隨著模型大小和訓(xùn)練方式的不同而顯著變化。例如,全參數(shù)微調(diào)比 PEFT 需要更多的內(nèi)存。

為什么 GRPO 對(duì)內(nèi)存需求較高

這要從 GRPO 的原理說起,這是它的流程圖。

打開網(wǎng)易新聞 查看精彩圖片

GRPO 對(duì)內(nèi)存需求較高的原因在于,其內(nèi)部涉及多個(gè)模型,并且在訓(xùn)練數(shù)據(jù)中每個(gè)查詢會(huì)產(chǎn)生多個(gè)輸出。上圖中的策略模型、參考模型和獎(jiǎng)勵(lì)模型各自都是一個(gè)需要進(jìn)行推理的 LLM。(盡管從技術(shù)上講,獎(jiǎng)勵(lì)模型可能不需要參數(shù)化,可以只是一個(gè) Python 函數(shù)或正則表達(dá)式,但不影響 GRPO 對(duì)內(nèi)存的高需求。)

為什么 8-Bit 優(yōu)化和梯度檢查點(diǎn)有助于減少內(nèi)存占用?

通常來講,訓(xùn)練一個(gè)大型語言模型需要在內(nèi)存中存儲(chǔ)三種主要類型的信息:模型參數(shù)、模型學(xué)習(xí)所需的梯度、優(yōu)化器的跟蹤數(shù)據(jù)。

對(duì)上述內(nèi)容我們可以這樣理解:如果模型的參數(shù)占用了 X 的空間,那么梯度也會(huì)占用大約相同的空間。然后,像 AdamW 這樣的優(yōu)化器需要更多的空間,因?yàn)樗鼈兙拖褚粋€(gè)記錄員,跟蹤最近的更新歷史,以便更好地決定未來的優(yōu)化。

為了減輕這種內(nèi)存負(fù)擔(dān),通常采用兩種技術(shù):

  • 首先,可以使用像 AdamW 這樣的 8-bit 優(yōu)化器版本,它們能更高效地存儲(chǔ)跟蹤數(shù)據(jù),同時(shí)仍保持良好的性能 —— 類似于壓縮照片可以節(jié)省空間,同時(shí)保留大部分圖像質(zhì)量;
  • 其次,使用梯度檢查點(diǎn)技術(shù),這就像在訓(xùn)練過程中拍攝快照,而不是記錄所有內(nèi)容。雖然這會(huì)使訓(xùn)練速度減慢約 20-30%,但它顯著減少了內(nèi)存使用。

結(jié)合這些技術(shù),即使對(duì) GPU 資源有限的人來說,也能夠訓(xùn)練更大的模型。

代碼示例

像 trl 這樣的庫已經(jīng)開始支持 GRPO,使得微調(diào)由 transformers 構(gòu)成的 LLM 變得非常簡(jiǎn)單。代碼也非常簡(jiǎn)潔,只需將訓(xùn)練器替換為 GRPOTrainer 并定義一些獎(jiǎng)勵(lì)即可。GRPO 的最小代碼量大約只有 99 行,如果你使用的是像 meta-llama/Llama-3.2-1B-Instruct 這樣的小型模型和像 openai/GSM8K 這樣的數(shù)據(jù)集,可以非??焖俚貑?dòng)。

trl 項(xiàng)目地址:https://github.com/huggingface/trl?ref=ghost.oxen.ai

  1. import
  2. torch
  3. from
  4. datasets
  5. import
  6. load_dataset
  7. Dataset
  8. from
  9. transformers
  10. import
  11. AutoTokenizer
  12. AutoModelForCausalLM
  13. from
  14. trl
  15. import
  16. GRPOConfig
  17. GRPOTrainer
  18. import
  19. re
  20. SYSTEM_PROMPT
  21. Respond in the following format:
  22. def
  23. extract_hash_answer
  24. text
  25. str
  26. str
  27. None
  28. if
  29. "####"
  30. not
  31. in
  32. text
  33. return
  34. None
  35. return
  36. text
  37. split
  38. "####"
  39. 1
  40. strip
  41. def
  42. get_gsm8k_questions
  43. split
  44. "train"
  45. Dataset
  46. data
  47. load_dataset
  48. 'openai/gsm8k'
  49. 'main'
  50. split
  51. data
  52. data
  53. map
  54. lambda
  55. 'prompt'
  56. 'role'
  57. 'system'
  58. 'content'
  59. SYSTEM_PROMPT
  60. },
  61. 'role'
  62. 'user'
  63. 'content'
  64. 'question'
  65. ],
  66. 'answer'
  67. extract_hash_answer
  68. 'answer'
  69. return
  70. data
  71. def
  72. extract_xml_answer
  73. text
  74. str
  75. str
  76. answer
  77. text
  78. split
  79. 1
  80. answer
  81. answer
  82. split
  83. ""
  84. 0
  85. return
  86. answer
  87. strip
  88. def
  89. format_reward_func
  90. completions
  91. kwargs
  92. list
  93. float
  94. """Reward function that checks if the completion has a specific format."""
  95. pattern
  96. r
  97. "^\n\n$"
  98. \n.*?\n
  99. \n.*?\n
  100. responses
  101. completion
  102. 0
  103. "content"
  104. for
  105. completion
  106. in
  107. completions
  108. matches
  109. re
  110. match
  111. pattern
  112. r
  113. for
  114. r
  115. in
  116. responses
  117. return
  118. 0.5
  119. if
  120. match
  121. else
  122. 0.0
  123. for
  124. match
  125. in
  126. matches
  127. def
  128. accuracy_reward_func
  129. prompts
  130. completions
  131. answer
  132. kwargs
  133. list
  134. float
  135. """Reward function that extracts the answer from the xml tags and compares it to the correct answer."""
  136. responses
  137. completion
  138. 0
  139. 'content'
  140. for
  141. completion
  142. in
  143. completions
  144. extracted_responses
  145. extract_xml_answer
  146. r
  147. for
  148. r
  149. in
  150. responses
  151. return
  152. 2.0
  153. if
  154. r
  155. a
  156. else
  157. 0.0
  158. for
  159. r
  160. a
  161. in
  162. zip
  163. extracted_responses
  164. answer
  165. def
  166. main
  167. dataset
  168. get_gsm8k_questions
  169. model_name
  170. "meta-llama/Llama-3.2-1B-Instruct"
  171. model
  172. AutoModelForCausalLM
  173. from_pretrained
  174. model_name
  175. torch_dtype
  176. torch
  177. bfloat16
  178. attn_implementation
  179. "flash_attention_2"
  180. device_map
  181. None
  182. to
  183. "cuda"
  184. tokenizer
  185. AutoTokenizer
  186. from_pretrained
  187. model_name
  188. tokenizer
  189. pad_token
  190. tokenizer
  191. eos_token
  192. training_args
  193. GRPOConfig
  194. output_dir
  195. "output"
  196. learning_rate
  197. 5e-6
  198. adam_beta1
  199. 0.9
  200. adam_beta2
  201. 0.99
  202. weight_decay
  203. 0.1
  204. warmup_ratio
  205. 0.1
  206. lr_scheduler_type
  207. 'cosine'
  208. logging_steps
  209. 1
  210. bf16
  211. True
  212. per_device_train_batch_size
  213. 1
  214. gradient_accumulation_steps
  215. 4
  216. num_generations
  217. 4
  218. max_prompt_length
  219. 256
  220. max_completion_length
  221. 786
  222. num_train_epochs
  223. 1
  224. save_steps
  225. 100
  226. save_total_limit
  227. 1
  228. max_grad_norm
  229. 0.1
  230. log_on_each_node
  231. False
  232. trainer
  233. GRPOTrainer
  234. model
  235. model
  236. processing_class
  237. tokenizer
  238. reward_funcs
  239. format_reward_func
  240. accuracy_reward_func
  241. ],
  242. args
  243. training_args
  244. train_dataset
  245. dataset
  246. trainer
  247. train
  248. if
  249. __name__
  250. "__main__"
  251. main

Num Generations 有什么用

Num Generations 是一個(gè)超參數(shù),它決定了我們將在訓(xùn)練數(shù)據(jù)中對(duì)每個(gè)查詢采樣多少個(gè)補(bǔ)全。然而,這會(huì)顯著增加 VRAM 的消耗。

打開網(wǎng)易新聞 查看精彩圖片

目前有一個(gè)開放的 GitHub 問題,可能會(huì)幫助解決內(nèi)存瓶頸問題,可以參考如下鏈接

地址:https://github.com/huggingface/trl/issues/2709?ref=ghost.oxen.ai

對(duì)于 num_completions=8,16,64 (DeepSeekMath 論文使用的 64),作者表示,不用再次計(jì)算上述所有值,而是使用了 1B 參數(shù)模型進(jìn)行了測(cè)試,以顯示內(nèi)存增長(zhǎng)。不過,作者還是建議大家在內(nèi)存瓶頸得到修復(fù)之前使用 num_generations=4,也能獲得不錯(cuò)的性能。

打開網(wǎng)易新聞 查看精彩圖片

影響 VRAM 的一些因素

要對(duì)所有影響顯存(VRAM)使用的因素進(jìn)行全面的超參數(shù)驗(yàn)證,需要進(jìn)行大量的實(shí)驗(yàn)。簡(jiǎn)單起見,這里只指出了需要注意的設(shè)置,以及實(shí)驗(yàn)中使用的具體數(shù)值。

  • batch_size=1,由于 GRPO 為每個(gè)查詢生成多個(gè)響應(yīng),batch size 會(huì)迅速失控。
  • gradient_accumulation_steps=4,優(yōu)化器是另一個(gè)占用大量 VRAM 的地方。此參數(shù)決定了我們將存儲(chǔ)的梯度以幫助優(yōu)化器進(jìn)行其「爬山」過程。
  • num_completions=4,DeepSeekMath 論文中使用了 64。這完全超出了有些人的計(jì)算預(yù)算。
  • max_prompt_length=256,如果你想訓(xùn)練模型擁有更大上下文的推理能力,將不得不增加 VRAM。GSM8K 的提示相對(duì)較小,適合此測(cè)試。
  • max_completion_length=786,同樣,由于計(jì)算注意力的內(nèi)存有限,推理鏈在這里受到限制。上下文或生成的 token 越多,需要的內(nèi)存就越大。
  • LoRA target_modules=["q_proj", "k_proj", "o_proj", "up_proj", "down_proj"] 在這方面可以嘗試幾種不同的迭代。target_modules="all-linear" 是一種流行的方式,可以從你的 LoRA 中擠出最多的性能(就準(zhǔn)確性而言)。

對(duì) VRAM 使用的粗略估算

如果你正在使用 FP16 精度進(jìn)行訓(xùn)練,以下是一些簡(jiǎn)單的估算方法,可以幫助你了解內(nèi)存主要用在了哪些地方:

  • 模型參數(shù):每個(gè)參數(shù)占用 2 字節(jié)。
  • 參考模型參數(shù):每個(gè)參數(shù)占用 2 字節(jié)
  • 梯度:每個(gè)參數(shù)占用 2 字節(jié)。
  • 優(yōu)化器狀態(tài):每個(gè)參數(shù)占用 8 字節(jié)。
  • 8 位優(yōu)化器:每個(gè)參數(shù)占用 4 字節(jié)。
  • PEFT:有助于減少梯度的顯存占用。

最后是關(guān)于準(zhǔn)確率的。作者完成了一個(gè) 10 億參數(shù)的 Llama 3.2 模型的完整訓(xùn)練。在應(yīng)用 GRPO 之前,該模型在保留測(cè)試集上達(dá)到了約 19% 的準(zhǔn)確率,而在經(jīng)過一個(gè)訓(xùn)練周期后,模型的準(zhǔn)確率飆升至約 40.5%。雖然這離 SOTA 水平還差得很遠(yuǎn),但這展示了 GRPO 的強(qiáng)大潛力。