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