RL后訓(xùn)練GPU用量壓縮到十分之一的工程拆解)
去年年底參加一場大模型后訓(xùn)練方向的閉門技術(shù)交流主講人一上來放了頁 PPT萬億參數(shù)模型做強(qiáng)化學(xué)習(xí)RL后訓(xùn)練GPU 用量壓到常規(guī)方案的十分之一。臺(tái)下第一反應(yīng)是懷疑第二反應(yīng)是拍屏保存。會(huì)后我把自己能調(diào)動(dòng)的小規(guī)模集群拉起來按這套思路復(fù)現(xiàn)了一遍關(guān)鍵環(huán)節(jié)又把公開資料里的訓(xùn)練賬單逐項(xiàng)核對過。結(jié)論是這不是標(biāo)題黨而是一套組合拳式的工程優(yōu)化。核心不是某個(gè)單點(diǎn)黑科技而是顯存賬、帶寬賬、調(diào)度賬三本賬一起算。這篇把整個(gè)拆解完整寫下來包括顯存計(jì)算過程、三層減法的原理、百卡規(guī)模的具體配置參數(shù)以及復(fù)現(xiàn)過程中踩過的幾個(gè)坑。先說明“10% GPU”的含義它不是說把單卡利用率做到 10% 反而跑通了而是把集群規(guī)模從千卡量級(jí)壓到百卡量級(jí)資源總量正好是常規(guī)方案的十分之一。下面從算賬開始講。1. 先把賬算清楚萬億參數(shù) RL 后訓(xùn)練的顯存黑洞1.1 四個(gè)模型同臺(tái)顯存不是翻倍而是翻四倍RL 后訓(xùn)練和普通 SFT 最大的區(qū)別在于訓(xùn)練過程中不只有一個(gè)模型在跑。以主流 PPO 框架為例顯存里同時(shí)要駐留四類角色策略模型參與訓(xùn)練的唯一主角權(quán)重要能更新所以除了權(quán)重本身還要有梯度和優(yōu)化器狀態(tài)。參考模型和策略模型同結(jié)構(gòu)、同規(guī)模但完全凍結(jié)用來計(jì)算新舊策略的 KL 散度防止策略一步跑飛。獎(jiǎng)勵(lì)模型給采樣結(jié)果打分模型本身可以比策略小一個(gè)量級(jí)但推理時(shí)同樣占顯存。采樣服務(wù)其實(shí)就是用當(dāng)前策略模型做推理生成 responseKV cache 和推理中間狀態(tài)都要顯存。很多人第一次做大規(guī)模 RL 時(shí)都會(huì)犯同一個(gè)錯(cuò)誤只按“一個(gè)模型訓(xùn)練”的賬準(zhǔn)備顯存等跑起來才發(fā)現(xiàn)參考模型和采樣服務(wù)把顯存吃光了。策略模型、參考模型、采樣服務(wù)三套本質(zhì)上都是同一個(gè)萬億參數(shù)級(jí)別模型顯存需求量不是翻一倍而是直接翻三四倍。這是先決條件也是后文所有優(yōu)化動(dòng)作的前提。1.2 常規(guī)賬單為什么動(dòng)輒上千卡把萬億參數(shù)的顯存賬攤開算。假設(shè)模型參數(shù)量 P 1T即 10^12 參數(shù)采用常見的 BF16 混合精度訓(xùn)練。只算模型狀態(tài)權(quán)重 梯度 優(yōu)化器狀態(tài)一個(gè)參數(shù)在標(biāo)準(zhǔn) AdamW 流程里要占 16 字節(jié)狀態(tài)精度每參數(shù)字節(jié)BF16 權(quán)重半精度2 字節(jié)BF16 梯度半精度2 字節(jié)FP32 優(yōu)化器主權(quán)重全精度4 字節(jié)一階動(dòng)量 m全精度4 字節(jié)二階動(dòng)量 v全精度4 字節(jié)合計(jì)16 字節(jié)1T 參數(shù)乘以 16 字節(jié)就是 16TB。一張 80GB 顯存的訓(xùn)練卡只放模型狀態(tài)也要 200 張才裝得下。這是策略模型本身的賬參考模型凍結(jié)權(quán)重再加 2TBBF16獎(jiǎng)勵(lì)模型另算采樣階段 KV cache 輕松破幾百 GB。疊加上激活值、通信緩沖、中間臨時(shí)張量常規(guī)配置沒有 500 到 1000 張卡萬億參數(shù) RL 后訓(xùn)練根本推不動(dòng)。這里有一層容易忽略的賬常規(guī)部署還會(huì)把訓(xùn)練池和采樣池完全分開。訓(xùn)練池 300 卡采樣池再 300 卡兩套集群各管各的資源總量直接翻倍。所以“千卡才能跑”的結(jié)論一點(diǎn)都不夸張。1.3 目標(biāo)賬單100 張卡如何夠用這套方案的總思路是把三筆最大開支分別處理優(yōu)化器狀態(tài)12TB直接不要了。梯度2TB邊算邊更算完即丟。權(quán)重2TB按層分片到每張卡配合 CPU 內(nèi)存動(dòng)態(tài)駐留。最終模型狀態(tài)賬單只剩 2TB 的 BF16 權(quán)重。攤到 100 張 80GB 卡上每卡本地切片 20GB剩下的 60GB 留給激活值、通信緩沖、KV cache 和臨時(shí)計(jì)算。參考模型量化后常駐 CPU 內(nèi)存獎(jiǎng)勵(lì)模型用小體量替代采樣和訓(xùn)練共用一批卡。這樣百卡規(guī)模就能閉環(huán)跑起來正好對應(yīng)標(biāo)題里“十分之一 GPU”這個(gè)數(shù)字。2. 三層減法為什么能把顯存砍一個(gè)數(shù)量級(jí)2.1 第一層不帶“記憶”的優(yōu)化器AdamW 每個(gè)參數(shù)要占 8 字節(jié)動(dòng)量 二階矩因?yàn)樗涀∶總€(gè)參數(shù)的歷史梯度一階矩和二階矩。在萬億參數(shù)規(guī)模下這份“記憶”本身就是最大的顯存黑洞。第一刀就砍在這里采用無狀態(tài)的極低顯存優(yōu)化思路不保存動(dòng)量和方差反向傳播計(jì)算到某一層時(shí)直接用當(dāng)前層梯度做一次參數(shù)更新。生活類比常規(guī)優(yōu)化器像記賬式理財(cái)每一筆開銷都要入賬、算歷史均值再?zèng)Q定下一步無狀態(tài)優(yōu)化器像現(xiàn)金管理手里有多少花多少花完即忘。它在小模型上收斂速度通常遜于 AdamW但在超大模型上省掉的 12TB 顯存可以換成更大 batch實(shí)際吞吐不降反升。代價(jià)是學(xué)習(xí)率必須給得更保守——實(shí)測下來無狀態(tài)優(yōu)化器的合理學(xué)習(xí)率約為 AdamW 的十分之一到二十分之一而且梯度裁剪幾乎不能省。2.2 第二層梯度即算即棄標(biāo)準(zhǔn)反向傳播會(huì)把整張計(jì)算圖所有層的梯度攢進(jìn)一個(gè)梯度張量等全部算完再統(tǒng)一更新。萬億模型這一攢就是 2TB。進(jìn)階做法是徹底取消“攢梯度”這一步反向傳播從最后一層往前推算到第 i 層梯度立刻更新第 i 層參數(shù)然后釋放該層梯度。整個(gè)過程只是多了幾次零散的參數(shù)更新操作但整卡再也不需要一塊 2TB 的全局梯度緩沖區(qū)。這里有個(gè)實(shí)現(xiàn)細(xì)節(jié)容易被忽略梯度不落地通信不能省。分片模式下反向階段需要再次 gather 權(quán)重才能算出梯度所以每輪訓(xùn)練至少要經(jīng)歷兩次參數(shù)聚合。要做異步流水讓通信時(shí)間藏進(jìn)矩陣計(jì)算里否則省下顯存卻跑出一個(gè)“通信饑餓”的訓(xùn)練過程??蚣軐用嫱ǔ?huì)有 prefetch 或 overlap 選項(xiàng)默認(rèn)不一定打開務(wù)必手工確認(rèn)。2.3 第三層零冗余分片 CPU 動(dòng)態(tài)駐留權(quán)重 2TB 攤到 100 張卡上每卡只有 20GB 本地切片。這里用零冗余參數(shù)分片同一份參數(shù)在整個(gè)集群只保留一份訓(xùn)練時(shí)按層 all-gather 聚合。注意是按層不是一次性全 gather。每層權(quán)重在萬億模型里通常只有幾個(gè) GB算完當(dāng)前層立刻丟棄再取下一層。這樣每卡的峰值占用是“本地 20GB 切片 當(dāng)前層 gather 結(jié)果”而不是恐怖的 2TB 全量。配合 CPU 動(dòng)態(tài)駐留冷參數(shù)當(dāng)前不參與計(jì)算的層放在 CPU 內(nèi)存里GPU 只保留正在計(jì)算的那段。參考模型更極端——它完全凍結(jié)可以量化到 INT8 甚至更低常駐 CPU推理時(shí)按需取。這一層有個(gè)很多人沒想透的點(diǎn)CPU 內(nèi)存不是問題CPU 帶寬才是。CPU 到 GPU 的互聯(lián)帶寬通常幾十 GB/s顯存帶寬是幾個(gè) TB/s差一個(gè)數(shù)量級(jí)。所以 CPU 卸載不是無腦把參數(shù)丟內(nèi)存而是要做預(yù)取隊(duì)列、按層流水、雙緩沖。更聰明的做法是只在 pipeline 冷卻階段卸載熱階段保持 GPU 駐留把每一寸帶寬都用在刀刃上。3. 實(shí)操過程100 張卡上把 RL 閉環(huán)跑起來3.1 訓(xùn)練與采樣同池化不單獨(dú)養(yǎng)推理艦隊(duì)RL 有三件活采樣推理、打分推理、訓(xùn)練前向 反向更新。常規(guī)方案把訓(xùn)練池和采樣池分開各用各的卡顯存需求直接翻倍。這套方案的核心動(dòng)作是把訓(xùn)練和采樣放到同一批卡上用交錯(cuò)調(diào)度錯(cuò)峰使用pipeline 并行的訓(xùn)練前向/反向之間存在天然氣泡某個(gè) stage 算完當(dāng)前 micro-batch 會(huì)空轉(zhuǎn)等待這些空檔全部拿去跑采樣推理。調(diào)度上我建議做一個(gè)顯式的雙環(huán)結(jié)構(gòu)訓(xùn)練環(huán)按 micro-batch 推進(jìn)采樣環(huán)從訓(xùn)練隊(duì)列尾部拿當(dāng)前策略參數(shù)快照去生成 response。兩個(gè)環(huán)共享同一批 worker靠隊(duì)列水位控制節(jié)奏。還要設(shè)一個(gè) min_available_mem 閾值——訓(xùn)練任務(wù)把顯存吃得太緊時(shí)采樣環(huán)自動(dòng)降速反之提速。實(shí)測下來采樣吞吐在訓(xùn)練滿載時(shí)會(huì)下降 20% 左右但省掉一整支推理集群這筆買賣非常劃算。3.2 策略模型訓(xùn)練配置一組可以直接用的參數(shù)復(fù)現(xiàn)時(shí)用的配置如下模型規(guī)模可以按比例縮放全局 batch2M tokens/stepmicro-batch 4096 tokens梯度累積 512 步。重計(jì)算開啟按層激活重計(jì)算顯存峰值約降一半代價(jià)是 1.5 倍左右計(jì)算量。學(xué)習(xí)率1e-6 起步warmup 200 步cosine 衰減到 1e-7。無狀態(tài)優(yōu)化器對學(xué)習(xí)率極其敏感寧小勿大。梯度裁剪max_grad_norm 0.3硬裁剪。PPO 側(cè)clip range 0.1KL penalty 系數(shù) 0.01 起步后續(xù)動(dòng)態(tài)調(diào)整。參考錨點(diǎn)每 200 步用當(dāng)前策略重算一次參考 logits中間步驟用緩存錨點(diǎn)做 KL 近似。這些參數(shù)來自小規(guī)模復(fù)現(xiàn)和公開資料的交叉對照換語言模型、換任務(wù)分布后建議重新掃一遍學(xué)習(xí)率。特別強(qiáng)調(diào)無狀態(tài)優(yōu)化器沒有二階動(dòng)量做“自動(dòng)剎車”調(diào)大學(xué)習(xí)率的后果不是平滑變快而是某一步直接飛掉。我在小模型上驗(yàn)證過學(xué)習(xí)率從 1e-6 調(diào)到 3e-6reward 曲線前期只快 10%后期直接塌方。大幅度更新會(huì)讓 rollout 分布瞬間偏離錨點(diǎn)KL 懲罰都救不回來。3.3 參考模型與獎(jiǎng)勵(lì)模型的輕量化技巧參考模型是“重量級(jí)配角”。完全凍結(jié)的特性讓它非常適合量化我用 INT8 權(quán)重量化 CPU 常駐推理時(shí)按層把量化權(quán)重拉到 GPU。更進(jìn)一步可以用 KL 錨點(diǎn)緩存——KL 散度逐 token 計(jì)算每隔若干步才用參考模型算一次 logits 并緩存中間步驟直接讀緩存。注意這里有個(gè)隱患策略在訓(xùn)練中不斷偏移錨點(diǎn)會(huì)過期。所以必須配 KL 漂移監(jiān)控當(dāng)策略輸出分布與錨點(diǎn)分布差異超過閾值常用 KL 均值 0.3時(shí)強(qiáng)制觸發(fā)一次錨點(diǎn)刷新。否則舊錨點(diǎn)會(huì)讓 KL 估計(jì)失真訓(xùn)練會(huì)跑出假象獎(jiǎng)勵(lì)。獎(jiǎng)勵(lì)模型相對好辦。傳統(tǒng)做法是單獨(dú)放大模型當(dāng)獎(jiǎng)勵(lì)模型這套方案建議“小獎(jiǎng)勵(lì)模型 規(guī)則打分器”混合數(shù)學(xué)、代碼類任務(wù)用規(guī)則驗(yàn)證器或執(zhí)行結(jié)果開放性任務(wù)用小一個(gè)量級(jí)的模型打分實(shí)在需要強(qiáng)語義打分的再上 FP8 推理。獎(jiǎng)勵(lì)模型精度直接影響 PPO 的信用分配不要為省顯存壓太狠——最低我只敢壓到原規(guī)模的三分之一再小獎(jiǎng)勵(lì)信號(hào)就開始失真。3.4 采樣階段的分頁 KV 與顯存復(fù)用采樣階段最大的顯存殺手是 KV cache。萬億模型的 attention 層數(shù)多、通道數(shù)大每個(gè) token 的 KV 要上千字節(jié)幾萬條并發(fā)序列隨便就是幾百 GB。這部分的處理方案分頁 KV cache把 KV 按固定 block 管理類似操作系統(tǒng)的虛擬內(nèi)存分頁解決碎片化并支持顯存不足時(shí)換出。低比特 KVK、V 分別降到 FP8顯存直接減半對 RL 采樣質(zhì)量影響可控。冷塊換出長時(shí)間不訪問的舊序列塊換到 CPU 內(nèi)存新序列用 GPU 塊。序列長度上限給所有 prompt 設(shè)統(tǒng)一 completion 上限避免極少數(shù)長尾序列吃穿整個(gè)顯存池。另一個(gè)復(fù)現(xiàn)時(shí)確認(rèn)過的細(xì)節(jié)采樣用動(dòng)態(tài)連續(xù)批處理而不是靜態(tài)按 batch 等待。靜態(tài)批處理在序列長度差異大時(shí)短序列會(huì)空等長序列浪費(fèi)的是顯存駐留時(shí)間連續(xù)批處理讓短序列先走空出來的 block 立刻給新序列吞吐能提升 30% 以上。這套機(jī)制配合訓(xùn)練采樣同池化是百卡閉環(huán)里最容易被忽略但回報(bào)率最高的模塊。4. 踩坑實(shí)錄這些問題網(wǎng)上基本查不到4.1 顯存沒爆時(shí)間爆了CPU 卸載的帶寬陷阱第一次跑通時(shí)顯存峰值很漂亮但訓(xùn)練速度慘不忍睹單 step 比常規(guī)方案慢 6 倍。查了半天發(fā)現(xiàn)罪魁禍?zhǔn)资?CPU 卸載沒做預(yù)取每一層都同步等權(quán)重從 CPU 搬進(jìn) GPU互聯(lián)帶寬成了真正的瓶頸。解決方式是把參數(shù)預(yù)取隊(duì)列加長到兩層以上開啟雙緩沖異步傳輸并且只在 pipeline 冷卻階段做卸載熱階段保持 GPU 駐留。調(diào)完之后單 step 只慢 1.8 倍顯存峰值穩(wěn)定在 50GB 以下。這個(gè)性價(jià)比可以接受。4.2 錨點(diǎn)緩存引發(fā)的“隱性獎(jiǎng)勵(lì)坍塌”這是一次很有意思的翻車。用 KL 錨點(diǎn)緩存跑了一個(gè)多小時(shí)reward 曲線單調(diào)上漲正當(dāng)我以為一切正常reward 突然斷崖式下跌并持續(xù)震蕩。定位后發(fā)現(xiàn)原因策略在更新中逐漸偏離早期錨點(diǎn)KL 近似誤差越來越大模型學(xué)會(huì)了用較大的 KL 漂移換取表面高分——本質(zhì)上是錨點(diǎn)欺騙。加上 KL 均值監(jiān)控和強(qiáng)制刷新之后reward 曲線才恢復(fù)健康?,F(xiàn)在的觀點(diǎn)是錨點(diǎn)緩存只適合中期提速寧可多花幾次前向計(jì)算也要保證 KL 估計(jì)可靠。4.3 同一批卡既訓(xùn)練又采樣差點(diǎn)死鎖訓(xùn)練環(huán)和采樣環(huán)共享 worker 的初版跑了半小時(shí)全部卡住采樣環(huán)等訓(xùn)練釋放顯存訓(xùn)練環(huán)等采樣環(huán)歸還通信原語互相等成了死鎖。解決方式是加隊(duì)列水位上限和超時(shí)看門狗訓(xùn)練隊(duì)列積壓超過閾值采樣環(huán)強(qiáng)制削峰任何通信原語等待超過 5 分鐘自動(dòng)丟棄當(dāng)前采樣任務(wù)并釋放資源。之后還發(fā)現(xiàn)一個(gè)更穩(wěn)妥的原則采樣任務(wù)優(yōu)先級(jí)必須低于訓(xùn)練任務(wù)。RL 訓(xùn)練中斷的代價(jià)遠(yuǎn)高于采樣中斷寧可丟幾個(gè)樣本也不要打斷訓(xùn)練 step。4.4 梯度范數(shù)“心電圖”與回滾策略無狀態(tài)優(yōu)化器沒有二階矩梯度范數(shù)在 RL 的嘈雜獎(jiǎng)勵(lì)下會(huì)跳得厲害。我親眼見過一個(gè) step 梯度范數(shù)從 3 跳到 300 再跳回 2。兩個(gè)對策一是 max_grad_norm 嚴(yán)格限制二是“心跳檢查”——每 N 步記錄各層梯度范數(shù)均值和方差若某層超過 5 倍歷史均值跳過這次更新并恢復(fù)上一個(gè) checkpoint。這套回滾機(jī)制讓整體訓(xùn)練穩(wěn)定很多代價(jià)只是偶爾跳過幾個(gè) step?,F(xiàn)象可能原因處理方式顯存低但 step 極慢CPU 卸載同步等待預(yù)取隊(duì)列 雙緩沖 熱階段駐留reward 先漲后崩KL 錨點(diǎn)過期KL 漂移監(jiān)控 強(qiáng)制刷新錨點(diǎn)集群整體卡死訓(xùn)練/采樣環(huán)互相等待隊(duì)列水位上限 超時(shí)看門狗梯度范數(shù)劇烈跳動(dòng)無狀態(tài)優(yōu)化器 嘈雜獎(jiǎng)勵(lì)梯度裁剪 心跳檢查 回滾5. 實(shí)測效果與這套方案的適用范圍5.1 復(fù)現(xiàn)實(shí)測的數(shù)據(jù)賬我復(fù)現(xiàn)的是機(jī)制本身用的是開源推理模型在小規(guī)模上驗(yàn)證目的不是跑 1T而是確認(rèn)每一層優(yōu)化是否成立。在百億級(jí)模型規(guī)模下常規(guī)方案需要訓(xùn)練池與采樣池分離加起來幾十張 80GB 卡這套方案只用了大約四分之一的數(shù)量同一批卡完成訓(xùn)練和采樣顯存峰值約 58GB吞吐約為常規(guī)方案的 7 成。reward 曲線和 KL 曲線形狀與常規(guī)方案一致穩(wěn)定性略差但完全可控。換算到萬億參數(shù)就是 80 到 128 張卡對比常規(guī)的 800 到 1000 張卡數(shù)量級(jí)完全對得上。當(dāng)然也要說實(shí)話百卡規(guī)模的單 step 時(shí)間一定比千卡慢因?yàn)?CPU 卸載和重計(jì)算都是有代價(jià)的。這套方案省的是“資源總量”和“采購成本”不是“單次訓(xùn)練墻鐘時(shí)間”。如果項(xiàng)目預(yù)算充足、時(shí)間緊張常規(guī)方案仍然更省時(shí)間如果卡緊張、預(yù)算有限這套方案幾乎是唯一解。5.2 適用的場景邊界這套方案的甜區(qū)是“顯存極度緊張 模型極大 獎(jiǎng)勵(lì)信號(hào)明確”的 RL 后訓(xùn)練比如代碼 RL、數(shù)理推理 RL。不適合的場景有三類需要精細(xì)優(yōu)化器的指令微調(diào)小模型用無狀態(tài)優(yōu)化器收益不大反而容易掉點(diǎn)。極長上下文 RLKV cache 絕對量太大低比特和分頁也救不回來。獎(jiǎng)勵(lì)模型本身必須很大才能區(qū)分質(zhì)量的場景獎(jiǎng)勵(lì)模型會(huì)變成新的顯存瓶頸這個(gè)賬就不好平衡了。最后分享一點(diǎn)個(gè)人體會(huì)。這套東西的核心不是變魔法而是把賬算清楚——顯存賬、帶寬賬、調(diào)度賬三本賬一起算才能把“十分之一 GPU”做實(shí)。如果只學(xué)一招我建議學(xué)它的參數(shù)分層駐留策略熱駐留、溫預(yù)取、冷卸載這是整套方案里回報(bào)率最高的一層工程優(yōu)化。用開源框架搭建之前一定確認(rèn)三件事梯度是否真的即算即棄、參數(shù)預(yù)取是否流水化、調(diào)度器有沒有防死鎖看門狗。三件事都確認(rèn)過了百卡萬億 RL 這件事才算真正落地。