TRACE：用 rollout 引導的量化感知訓練，實現 MoE 語言模型的 FP4 強化學習

論文要解決什麼問題

強化學習（RL）後訓練已經成為提升大語言模型推理、程式碼和長程任務能力的主流手段。但它有一個很重的成本：每一輪訓練都要先用當前策略做大量 rollout（取樣生成），再用這些樣本更新模型。論文指出，rollout 階段帶來了顯著的計算和視訊記憶體開銷，這促使業界嘗試用低精度做 rollout。FP4 是目前硬體支援的最激進的浮點精度之一，權重、啟用和 KV 快取都能壓到 4 位。

問題在於，現有的 FP4 RL 方法有一個關鍵缺陷。按照論文的描述，它們主要是分別最佳化訓練路徑和 rollout 路徑各自的量化精度，而沒有直接縮小兩條量化執行路徑之間的差異。RL 對這種差異非常敏感：rollout 用 FP4 取樣出的軌跡，是在訓練側以另一套數值行為計算梯度的。兩邊的數值不一致，就會變成離策略偏差，訓練變得不穩定，最終效能下降。

TRACE 的核心思路

TRACE 的全稱是 Train-Rollout Quantization Alignment via Compact GuidancE，即藉助緊湊指引實現訓練與 rollout 的量化對齊。它面向混合專家（MoE）語言模型，包含兩個核心設計。

第一個設計是 rollout 引導的量化感知訓練（QAT）。傳統 QAT 在訓練側自行決定每個數值如何舍入到 FP4 網格。TRACE 改為讓 rollout 側的量化結果去指導訓練側的 FP4 舍入決策。換句話說，訓練側不再各自為政，而是儘量復現 rollout 側實際發生的舍入，從而直接降低 train-rollout 差異。這是論文區別於以往方法的主要創新點：最佳化目標從“每條路徑各自量化得準”變成“兩條路徑彼此對得齊”。

第二個設計是高效的量化資訊快取方案。要讓訓練側知道 rollout 側的舍入結果，就得儲存並傳遞相應的量化資訊，這會帶來儲存和通訊開銷。TRACE 選擇性地只保留來自較深層的尾數（mantissa）和縮放因子（scale）資訊，以此壓縮這部分額外成本。摘要沒有給出具體保留哪些層、佔用多少額外視訊記憶體，這些細節需要讀正文確認。

工作原理的工程理解

FP4 通常採用分塊縮放：一小組數值共享一個縮放因子，組內每個數值只用 4 位表示。舍入方向（向上還是向下取最近的網格點）在每個數值上都是一個離散決定。兩條路徑如果對同一個權重或啟用做出不同的舍入決定，輸出就會出現微小但系統性的偏差。偏差經過多層 MoE 網路累積，再疊加路由器選擇專家的離散性，就可能被放大。

TRACE 的做法可以這樣理解：rollout 引擎先完成 FP4 量化並記錄結果，訓練側在前向計算時參考這份記錄來決定舍入，使兩邊的計算圖在數值上儘量一致。因為深層的誤差對最終輸出影響更直接，快取深層資訊是在精度與開銷之間的折中。以上是對摘要描述的機制解讀，具體演算法細節以論文正文為準。

實驗結果與效能

論文在四個大規模 MoE 語言模型上做了評估，覆蓋推理、程式碼和長程 RL 任務。根據摘要，主要結論有三點。 其一，TRACE 實現了 FP4 權重與啟用、以及 FP4 KV 快取的聯合 rollout，同時 RL 效能與 BF16 rollout 相當。這意味著低精度並沒有以犧牲最終效果為代價。 其二，rollout 速度最高提升 5.4 倍。注意“最高”二字：這是最好情形下的數字，實際加速取決於模型、序列長度、批次大小和硬體，不能直接當作平均值。 其三，與“對 BF16 訓練出的策略做事後 FP4 量化”相比，TRACE 的最終 FP4 效能更強。這個對比很關鍵，它說明讓模型在訓練中就適應 FP4，比訓練完再壓縮更好。

摘要沒有給出具體的基準分數、視訊記憶體節省比例和模型名稱，本文不做推測。

對開發者與企業的影響

對做 RL 後訓練的團隊，rollout 往往佔據訓練牆鍾時間的很大一部分，尤其在長思維鏈和長程任務裡。如果 FP4 rollout 在不損失效果的前提下帶來數倍加速，同樣的 GPU 預算能跑更多迭代，或者縮短實驗週期。對 MoE 模型而言，權重體積大、KV 快取壓力大，FP4 的收益更明顯。

在生態層面，這項工作依賴支援 FP4 的硬體和推理核心。採用門檻包括：推理引擎要能匯出量化結果，訓練框架要能消費這些結果，兩者之間的資料通路需要工程改造。它也提示了一個更一般的設計原則：在 RL 系統裡，訓練與推理的數值一致性應當作為一等目標，而不是副產品。

侷限與未來方向

有幾點需要保持清醒。第一，摘要中的 5.4 倍是上限數字。第二，實驗覆蓋四個 MoE 模型，是否能推廣到稠密模型或更小規模，需要獨立驗證。第三，rollout 引導增加了系統耦合：訓練側依賴 rollout 側的輸出，會帶來額外的同步、儲存與通訊，即使已經做了快取壓縮。第四，這是一篇剛釋出的預印本，尚未經過同行評審，結論應等待復現。

未來的方向包括：把對齊思想擴充套件到更多精度格式、與非同步 RL 系統結合、研究快取策略在更長序列下的擴充套件性，以及評估在安全性和穩健性等維度上是否同樣不受影響。

小結

TRACE 抓住了 FP4 RL 的真正瓶頸：不是單條路徑量化得不夠準，而是訓練與 rollout 兩條路徑對不齊。

透過 rollout 引導的 QAT 與緊湊的量化資訊快取，它在摘要所述的實驗裡做到了與 BF16 rollout 相當的 RL 效能，並獲得最高 5.4 倍的 rollout 加速。是否值得采用，要看你的硬體、框架和模型規模，建議在小規模任務上先復現。