致敬Kimi K2：基于slime的全流程INT4量化感知RL訓練

受 Kimi K2 團隊啟發，SGLang RL 團隊成功落地了 INT4 量化感知訓練（QAT） 流程方案。通過 “訓練端偽量化 + 推理端真實量化（W4A16）” 的方案組合，我們實現了媲美 BF16 全精度訓練的穩定性與訓推一致性，同時 INT4 極致壓縮也將 1TB 級超大模型的采樣任務容納于單機 H200 (141G) 顯存內，消除了跨機通信瓶頸，顯著提高了 Rollout 效率，為社區提供了兼顧高性能與低成本的開源參考。

近期，SGLang RL 團隊在強化學習的訓練穩定性，訓練效率與適用場景方面取得了重要進展，具體包括：

• Unified multi-turn VLM/LLM 多輪采樣范式：我們提供了 VLM 多輪采樣范式的實現blog，開發者只需編寫一套定制化的 rollout 函數，即可像訓練 LLM 一樣，輕松開啟 VLM 的多輪強化學習。
• 穩定性提升：我們實現了Rollout Router Replay機制，顯著提升了 MoE 模型在 RL 訓練過程中的穩定性。
• 低精度訓練：我們在 RL 場景中成功實現了全流程 FP8 訓練與采樣，進一步釋放了硬件性能。
• 投機采樣：我們在 RL 場景中成功實踐了投機采樣，實現了大規模訓練的無損加速。

在此基礎上，我們更進一步，在 slime 框架上成功復現并落地了INT4 量化感知訓練（QAT）全流程方案。該方案深受 Kimi 團隊 K2-Thinking 技術報告中關于W4A16 QAT (Quantization-Aware Training)實踐的啟發。為了致敬先行者并回饋社區，本文將詳細剖析我們在開源生態中打通全流程的技術細節，旨在為社區提供一份兼顧穩定性與性能的可落地參考。

核心收益概覽：

• 突破顯存瓶頸：通過權重壓縮與低比特量化，使 1TB 級別的 K2 類模型能縮容至單機 H200 (141G) 顯存內，避免了跨機通信瓶頸。
• 訓推一致：訓練端利用 QAT 確保權重符合 INT4 分布，推理端執行 W4A16 (Weights INT4, activations BF16 ) 計算；二者均通過 BF16 Tensor Core 進行運算，實現了媲美 BF16 全精度的訓推一致性。
• 單機效率倍增：在超大模型場景下，INT4 策略大幅降低了顯存與帶寬壓力，Rollout 效率顯著超越 W8A8 (Weights FP8 , Activations FP8）。

本項目由 SGLang RL 團隊、 InfiXAI 團隊、螞蟻集團 Asystem & 阿福 Infra 團隊， slime 團隊與 RadixArk Miles 團隊聯合完成。相關功能與 recipe 已經同步到了slime與Miles社區，歡迎大家試用與貢獻。我們也在更進一步向 MXFP8 與 NVFP4 發起挑戰。同時，由衷感謝Verda Cloud為本工作提供的計算資源。

1. 技術方案概覽

1.1 總體流程

我們實現了從訓練到推理的完整 QAT INT4 閉環的方案，如下圖所示：

圖1 QAT INT4 全流程

在QAT 訓練階段，訓練側在維護 BF16 主權重（Master Weights）的基礎上，前向傳播通過偽量化（Fake Quantization）引入量化噪聲。所謂 “偽”，是指該步驟并未真正將 BF16 數據類型轉換為低精度的 INT4，而是保持浮點計算路徑不變，通過插入量化再反量化（Quant-Dequant）操作來模擬低精度的計算。

具體而言，高精度權重在經過 “離散化映射到 INT4” 后被立即還原，雖然其物理存儲格式仍為浮點，但數值精度已實質性降低。這種原值與還原值之間的差異引入了量化誤差，在數學上等效于向網絡注入了噪聲，迫使模型在訓練階段就通過梯度更新去適應這種精度損失。

反向傳播則利用STE (Straight-Through Estimator)技術跳過了量化算子的不可導特性。量化過程的核心操作是 “取整（Rounding）”，其數學形態為階梯函數，導數在幾乎所有位置均為 0。這意味著在標準反向傳播過程中，梯度信號傳導至此處會因“梯度消失”而徹底中斷，導致底層的主權重無法獲得更新。

對此，STE 采用了 “梯度透傳” 策略：在反向傳播計算時，將取整函數的導數定義為 1（即視為恒等映射）。這一機制相當于在不可導的 “斷崖” 上架設了一座橋梁，讓梯度能夠越過取整層，有效回傳至高精度的浮點權重，確保 QAT 訓練鏈路的閉環。

在權重轉換階段，我們將訓練收斂的 BF16 權重導出并執行真實量化（Real Quantization），將其轉換為推理引擎適配的 INT4 格式（如 Marlin）。

進入RL Rollout階段，由 SGLang 加載 INT4 Weights 并執行高效的 W4A16（INT4 權重 x BF16 激活）推理，生成的經驗數據（Experience）將回流至第一階段用于下一輪 RL 訓練，從而構成一個自洽的迭代閉環。

1.2 核心策略選擇

在量化格式上，我們參考Kimi-K2-Thinking選用了INT4 (W4A16)方案。這主要考慮到相比 FP4，INT4 在現有硬件（Pre-Blackwell 架構）上的支持更加廣泛，并且業界已有成熟高效的 Marlin Kernel 實現。實驗表明，在 1×32 量化 Scale 粒度下，INT4 動態范圍充足、精度穩定，其性能與生態鏈路均已高度優化。作為工業界 “足夠好（Good Enough）” 的量化標準，INT4 在性能、風險與維護成本間實現理性平衡。當然，我們后續也計劃在 NVIDIA Blackwell 系列硬件上進一步展開 FP4 RL 的探索。

在訓練方法方面，我們采用了Fake Quantization 配合 STE的經典組合。通過維護 BF16 主權重，在前向計算中模擬量化噪聲，并在反向傳播時直通梯度，這種方式最大程度地保證了低精度訓練的收斂性與穩定性。

2. 訓練側：Megatron-LM 的偽量化改造

2.1 Fake Quantization 與 STE 實現

圖2

這一階段的核心目標是在訓練過程中實時模擬量化誤差，迫使模型 “學會” 適應低精度表示。為此，我們采用了Fake Quantization機制：盡管權重在存儲和更新時仍保持高精度的 BF16 格式，但在前向傳播的實際計算中，會被暫時映射到 INT4 的精度范圍參與運算。

具體實現上，我們在 megatron/core/extensions/transformer_engine.py 中的 _FakeInt4QuantizationSTE 類構建了核心邏輯。基于分組最大絕對值進行動態量化（Dynamic Quantization），模擬 INT4 的 [-7, 7] 數值范圍及截斷操作，但在計算時仍使用 BF16 類型，僅引入量化誤差。

而在關鍵的反向傳播環節，我們引入了STE機制，確保梯度能夠直接穿透量化層，不經修改地回傳以更新主權重，從而保證訓練的連續性。

2.2 Fake Quantization 對比實驗

為了驗證 QAT 方案的必要性，并探究訓練與推理精度不匹配帶來的具體影響，我們設計了一組消融實驗，分別在 “開啟 QAT INT4 訓練，BF16 Rollout” 和 “關閉 QAT 訓練，直接進行 INT4 Rollout” 兩種非對稱場景下進行了測試，并以對數概率絕對差值（Logprob Abs Diff）作為訓推不一致的觀測指標。

圖3 Rollout 側 BF16，訓練側對比 QAT INT4 效果

圖3展示了 “開啟 QAT INT4 訓練，BF16 Rollout” 的場景（即紅線部分）。可以看到，即使我們使用了高精度的 BF16 進行推理，誤差依然顯著偏高。這是因為在 QAT 過程中，模型權重已經針對 INT4 的量化噪聲進行了 “適應性調整” 或補償；推理時若移除量化步驟，這種補償反而成為擾動，導致特性分布偏移（Distribution Shift）。

圖4 Rollout 側 INT4 Weight Only，訓練側對比 QAT INT4 效果

圖4則展示了 “關閉 QAT 訓練，直接進行 INT4 Rollout” 的場景（即紅線部分）。這對應了傳統的訓練后量化（PTQ）模式。由于模型在訓練階段從未接觸過量化噪聲，直接將權重壓縮至 INT4 不僅造成信息的劇烈丟失，更導致推理時的特征分布與訓練時產生偏移，致使誤差隨著訓練步數呈現震蕩上升的趨勢。

結論：實驗有力地證明，訓練端的 Fake Quantization 與推理端的 Real Quantization 必須協同開啟。只有當訓練時的模擬噪聲與推理時的真實量化精度嚴格對齊，才能有效抑制訓推不一致，避免分布偏移，將誤差控制在接近基線的水平，從而真正打通低精度 RL 訓練的全流程。

3. 權重更新階段

3.1 權重流轉與動態格式適配

圖5

為了復用 SGLang 在推理端已有的優化，我們直接采用了其內置的Marlin INT4作為 INT4 的推理方案。然而，這在工程落地時我們遇到了顯著的 “格式鴻溝”：QAT 訓練產出的是類似 Hugging face 上的標準格式權重，而 SGLang 推理引擎的 Marlin Kernel 則強制要求權重必須經過特定的打包（Pack）與重排（Permute）處理，方能被 Kernel 高效讀取。

面對 RL 訓練中頻繁的權重更新需求，首先需要解決格式兼容性問題。為此，我們設計了一套逆向的 `restore_weights_before_loading`保護機制。該機制利用緩存的 `_original_shapes` 元數據，能夠在權重更新動作發生前，強制將當前內存中的 Marlin 權重格式還原（Resize）回原始形狀。這一設計有效防止了因維度不匹配導致的運行時錯誤，確保模型能夠在標準權重格式與 Marlin 權重格式之間平滑切換。此外，我們還在系統層面新增了 `post_process_weights` API，允許控制平面根據訓練節奏顯式觸發這一流程。

而針對權重加載完成后的格式適配挑戰，我們在 `compressed_tensors_moe.py` 中實現了一套動態權重管理機制。在模型權重加載結束階段，系統會自動觸發 `process_weights_after_loading` 流程，底層調用 `gptq_marlin_moe_repack` 與 `marlin_moe_permute_scales` 等算子，在內存中即時將標準權重轉換為高度優化的 Marlin 權重格式，從而最大化推理時的訪存與計算效率。

3.2 權重更新時的量化

圖6

進入核心的Real Quantization環節。不同于訓練時的 Fake Quantization，這一步通過代碼中的 `int4_block_quantize` 函數執行不可逆的精度壓縮操作：基于設定的 Group Size，計算每組權重的縮放因子（Scale），并將高精度浮點數映射到 `[-7, 7]` 的 INT4 整數域。

為了最大化顯存利用率，接著執行位寬打包（Packing）操作。由于 PyTorch 缺乏原生的 INT4 數據類型，我們通過 `pack_int4_to_int32` 函數利用位運算技巧，將 8 個 INT4 數值緊湊地 “壓縮” 進 1 個 INT32 整數中（即 `8 × 4 bits = 32 bits`）。最終，這些經過壓縮的 Packed Weights 連同 Scale 因子被傳輸至推理引擎，完成了從 “訓練格式” 到 “推理格式” 的轉換。

4. 推理階段

圖7

極簡打包與零開銷解包

在 RL 訓練的 Rollout 階段，我們直接復用了 SGLang 優化成熟的 W4A16 量化方案。SGLang 使用緊湊的 INT4 格式，將兩個 4-bit 權重打包進一個字節，相比 BF16 節省了 75% 的內存。在推理時，Triton kernel 通過高效的位移和掩碼操作（>> 4 和 & 0xF）快速解包，得益于計算與 IO 的并行覆蓋，該過程幾乎實現了零額外延遲。

MoE 算子深度融合

• 顯存優化：SGLang 引入動態的 moe_align_block_size，根據當前 Token 數量和 Expert 分布自動選擇 block_size ，將同一 Expert 的 Token 聚集并對齊，提升顯存帶寬利用率。
• 計算融合：SGLang 引擎除集成了高效的Marlin INT4實現、還將 gating 部分 fuse 成一個高性能的 kernel，避免了反復啟動 kernel 和讀寫中間結果。同時，該 INT4 推理方案兼容 GPTQ 和 AWQ 等主流量化格式，以及支持對稱與非對稱兩種模式。

5. INT4 QAT RL 效果

5.1 訓練效果

• 訓練側

圖8 Qwen3-235B-A22B Raw-Reward對比

圖9 Kimi-K2-Thinking Raw-Reward對比

上圖展示了基于 slime 框架，Qwen3-235B-A22B 與 Kimi-K2-Thinking 模型在 dapo-math-17k 數據集上的訓練表現。通過對比實驗發現，相較于 “BF16 訓 - BF16 推” 及 “BF16 訓 - FP8 推”，“BF16 訓 - INT4 推” 配置下的 Raw-Reward 仍能保持穩健增長，且其增長趨勢與前兩者基本一致，證明了該方案在訓練過程中的有效性。

• 評估側

圖10 Qwen3-235B-A22B AIME數據集評估對比

圖11 Kimi-K2-Thinking AIME數據集評估對比

為了更加嚴謹地評估模型能力的演進，我們每隔 10 個訓練步長就在 aime-2024 基準測試集上進行一次評估。上圖給出了 Qwen3-235B-A22B 與 Kimi-K2-Thinking 在不同 RL 訓練配置下的模型評分增長軌跡。

實驗表明：“BF16 訓 - INT4 推” 方案不僅在評估分數上呈現出穩健的上升態勢，且其性能提升的斜率與最終達到的峰值，均與 “BF16 訓 - BF16 推” 和 “BF16 訓 - FP8 推” 方案保持了較高的重合度。這種高度的一致性有力地證明了模型在經過低比特量化后，其核心表示能力并未受損，保證了在大幅降低計算開銷的同時，依然能夠實現與全精度推理相媲美甚至完全看齊的泛化表現。

5.2 訓推差異

圖12

圖13

為了直觀評估方案效果，我們在 Qwen3-30B 與 Qwen3-235B 模型上進行了的 QAT RL 訓練驗證。圖中 Y 軸反映了訓練側與推理側輸出的 Logprob 絕對差值，數值越低意味一致性越強。實驗結果顯示，INT4（綠色虛線）與 BF16 基準（紅色實線）呈現出驚人的重合度，且顯著低于表現出較高誤差水平的 FP8（藍色虛線）。這證實了 INT4 QAT 策略能有效規避 “BF16 訓 - FP8 推” 模式下的精度損失，實現與全精度無異的訓推表現。

這種一致性背后的原因我們推測為兩點：

• 截斷誤差抑制：訓練側的 Fake Quantization 將權重限制在 INT4 值域內。這種數值范圍的約束，有效降低了矩陣乘法中 Accumulator 累加時因并行計算順序不確定性引發的浮點舍入誤差（Floating-point Rounding Error），即改善了所謂的“大數加小數”精度丟失問題。
• 高精度計算：推理側采用 W4A16 模式，其核心計算全程基于BF16 Tensor Core進行，確保了運算精度與訓練階段的高度對齊。

5.3 Rollout 加速

圖14 Qwen3-235B-A22B Rollout 性能對比

從 Qwen3-235B 的 Rollout 性能對比圖中可以直觀看到，雖然 INT4（綠色點劃線）與 FP8（藍色虛線）均較 BF16 基線（紅色實線）實現了顯著加速，但兩者彼此之間并未拉開巨大的性能鴻溝。這一現象主要受限于當前的硬件特性：由于 NVIDIA H 系列 GPU 沒有原生的 INT4 Tensor Core， W4A16 方案本質上利用的還是 BF16 Tensor Core 進行計算，雖然大幅降低了顯存帶寬壓力，但在吞吐上無法像 W8A8 一樣利用原生 FP8 Tensor Core 進行加速從而獲得計算增益。因此，在單步推理耗時上，INT4 僅表現出微弱的優勢，與 FP8 基本處于同一性能梯隊。

圖15 Kimi-K2-Thinking Rollout 性能對比

對于 Kimi-K2-Thinking Rollout 性能的對比。首先觀察雙節點場景下的通信瓶頸：圖中 FP8（紅線）與 INT4（藍線）呈現出相似的水平。因為 H 系列 GPU 缺乏原生的 INT4 計算單元，INT4 無法在計算層面提供加速，因此整體性能依然受限于跨節點的通信帶寬。

然而，綠線所代表的單節點表現揭示了 INT4 的真正價值 —— 顯存壓縮。通過將模型體積減半，我們成功將 1TB 級別的超大模型完整加載至單機顯存中。這直接消除了昂貴的跨機通信開銷，將 Rollout 耗時大幅縮減。這有力地證明，在當前硬件環境下，INT4 QAT 的核心收益在于通過壓縮顯存，解鎖了高效的單機部署 Rollout 方案。

6. 總結與未來工作

slime 的這項工作不僅證明了在開源生態中復現工業界前沿方案的可行性，也為超大規模模型的低成本訓練探索了新的路徑。我們期望這套方案助力更多開發者深入理解 QAT 技術，并推動其在 RL 場景下的實際落地與廣泛應用。

通過在開源框架上的復現，我們驗證了 Kimi 團隊所提出的 INT4 QAT 方案的有效性：

• 精度復現：在 slime 的復現實驗中，我們同樣觀察到了 INT4 QAT 的精度優勢，實現了與 BF16 基線一致的效果。
• 效率提升：RL Rollout 階段的吞吐提升顯著，驗證了低比特量化在 RL 場景下的巨大價值。

未來工作：

• 訓練端效率優化：目前，由于在訓練過程中引入了 QAT Fake Quantization 計算，帶來了較大的額外性能開銷，導致訓練速度明顯低于 BF16 模式。這在一定程度上折損了 Rollout 階段帶來的端到端性能收益。我們后續計劃提出一套全新的優化方案，旨在解決這一訓練側的效率瓶頸，實現全鏈路的加速。
• 推理側 FP4：隨著 NVIDIA Blackwell 架構的逐步普及，我們將積極探索 FP4 精度在 RL 訓練與推理中的應用可行性，以期進一步挖掘硬件潛力。

slime 在 QAT INT4 的嘗試不僅證明了在開源生態中復現工業界前沿方案的可行性，也為超大規模模型的低成本訓練探索了新的路徑。我們期望這套方案助力更多開發者深入理解 QAT 技術，并推動其在 RL 場景下的實際落地與廣泛應用。

致謝

SGLang RL Team: Ji Li, Yefei Chen, Xi Chen, BBuf

InfiXAI Team: Mingfa Feng, Congkai Xie, Shuo Cai

螞蟻集團 Asystem & 阿福 Infra 團隊：Yanan Gao, Zhiling Ye, Yuan Wang, Xingliang Shi

RadixArk Miles Team: Chenyang Zhao, Yueming Yuan, Jiajun Li, Yusheng Su, Mao Cheng, Tom, Banghua Zhu

slime Team: Zilin Zhu, Chengxing Xie, Lei Li, Haisha Zhao