K2-Thinking背后的原生INT4量化的思考

NLP工作站

發布于 2025-11-11 07:24

瀏覽

0收藏

個人認為，在LLM的語境下，量化并非傳統的犧牲精度換速度，原生的低比特量化在param-scaling + test-time-scaling的持續演進下，會成為大模型訓練的一個標準范式。

量化一直是LLM領域一個非常火熱的研究方向，篇幅有限這一次就不嘗試對相關的基礎知識進行詳細展開了。

如果文章中出現你不太了解的概念，可以嘗試問問Kimi:)

這里簡單列一下后續主要會涉及到的一些縮寫和基礎解釋：

PTQ：訓練后量化，將訓練完成的模型權重通過特定的量化方法量化到對應的精度。
QAT：訓練感知量化，在訓練階段就讓模型提前適應某部分的權重/激活被量化到更低比特數下產生的精度損失。通常來說相比PTQ流程更加復雜但精度損失更小。
W(n)A(m)：模型權重(weight)使用n-bit，激活(activation)使用m-bit。

本文主要聊與這次K2-Thinking相關的W4A16 QAT相關的問題，對W8A8，W4A4等在訓練&Prefill場景實現加速的量化方案留待之后討論。

低比特量化對低延遲推理的重要性

當前，大模型推理根據優化目標的不同，存在兩個不同的trade-off方向（也就是大家看到的kimi turbo和非turbo）

高吞吐（面向成本）：以推理集群的吞吐量為主要目標，盡可能提高gpu計算資源的利用率。通常使用較大的并發數（batch size）使得推理過程中的GEMM（矩陣乘）達到compute bound。
低延遲（面向用戶體驗）：以用戶側的輸出延遲（TPOT）為主要目標，盡可能降低單次推理的延遲。通常會使用相對小的并發度以及小規模的單實例節點數。

Kimi-K2的MoE部分稀疏度達到的1/48，在我們當前的硬件環境下，decoding 階段 MoE 算子幾乎必然 memory-bound，權重的大小就決定了MoE的計算效率。

與此同時，模型權重占用的顯存越小，單個實例需要的最小GPU數量越少，多卡通信延遲越低，在GPU高速互聯（如nvlink）能夠全部覆蓋時還會進一步陡降。

而K2原始FP8權重的1TB大小，剛好是很多GPU單機高速互聯”裝不下“的邊界。因此，在decoding階段，W4A16量化的推理延遲是要顯著優于W8A8的。

為什么需要QAT而不是使用PTQ

事實上，在Kimi-K2發表后，開源社區中誕生了很多基于K2的low-bit PTQ模型，我們也了解到一些K2推理服務的供應商同樣采用了FP4等量化方式來加速推理。

我們內部也驗證了在4-bit這個精度下，PTQ可以做到在我們能觀測到的所有benchmark上近乎無損。

然而，當進入K2-Thinking的研發階段，這個結論被推翻了：隨著模型的生成長度變得越來越長，我們原本的block FP8推理精度和INT4 PTQ的結果呈現出了統計意義上的明顯差別。

一個可能的原因是隨著decoding計算次數的增加，量化產生的誤差被不斷累積了。

此外，我們還觀察到了INT4 PTQ的另一個劣勢：依賴校準集。

我們測試了一些在訓練集中出現過，但未在PTQ校準集中出現的case，發現FP8模型可以很好地背誦下這些訓練數據，而量化后的模型則會換一種表述方式甚至遺忘相關的內容。

關于這個問題目前的大致猜測是當moe非常稀疏時，盡管我們已經用了較大規模的校準數據，仍然會有部分專家只被路由到了少量token，進而導致這些專家的量化結果產生明顯的“失真”。

綜上，我們認為K2-Thinking如果想要實現低延遲的推理，使用量化損失更小的QAT方案勢在必行。

K2-Thinking的QAT方案

很遺憾，K2-Thinking的INT4 QAT并不存在一些社交媒體上期待的“天頂星科技”，因為我們發現相對基礎的QAT方案就可以比較輕松的做到相對baseline的無損，這也是為什么我們并沒有在技術博客花太多筆墨來詳細介紹相關的技術方案。

本次我們weight-only QAT采用的就是常見的fake-quantization+STE（直通估計器）方案。

保存原始的bf16權重，通過量化-反量化得到模擬精度損失后的權重，進行矩陣乘，反向傳播時再將梯度直接更新到原始的bf16中。

相比方法上的創新，Infra team在完整支持這個feature過程中的迭代效率可能更讓我自豪。

由于RL訓練才是Thinking模型的重頭戲，所以我們在發現PTQ不work之后，需要快速完整打通: (1) QAT訓練邏輯 (2) int4的完整推理流程 (3) QAT得到的模型到int4推理模型的轉換 (4) RL rollout部署方案 (5) colocate更新訓推框架中int4權重。

而以上這些在短短幾天的級別就基本正確運轉起來，不得不感嘆一下我司的人才密度和跨團隊協作的高效率。

經過算法同學的多方驗證，以上這套QAT方案可以在整個post-training階段不改變任何訓練配方，不增加任何訓練token數的前提下實現近乎無損。具體反映在loss有非常非常微小的gap，而benchmark分數保持一致。

不少人忽略的一點：INT4 QAT在RL階段的好處

這里算是我寫這篇文章的最大動機。

目前外界對于我們Native INT4的分析其實都挺正確的，但似乎比較少人想過它對我們RL訓練本身的價值。

熟悉RL的同學應該都知道，當前主流RL rollout環節有一個繞不過去的痛點就是長尾問題（大部分rollout已經結束，只有少量請求依然在持續生成的階段，此時推理效率非常低下）。

而此時的負載情況和我們前面說的低延遲推理是高度一致的。

因此，當我們的模型在rollout階段已經是INT4的精度，長尾階段的效率可以得到顯著提升。

當然，INT4 QAT由于需要進行QDQ（量化-反量化），會輕微增加訓練耗時，但比起rollout效率提升，這個增量要小得多。

目前從端到端來看，相比baseline方案，int4精度下一個完整的RL iteration耗時可以有大約10%-20%的減少，視不同場景的負載情況變化。

此外，Quantized RL本身也是最近備受關注的topic。QeRL表明在rollout階段引入量化噪聲可以讓learning-rate更加魯棒，以及增大trace的多樣性。

而我們也觀測到，在INT4精度下，推理框架和訓練框架forward結果的差異性（Discrepancy)會明顯比bf16更小。

目前的猜測是由于INT4表示的空間更小，不同gemm kernel下accumulator累加順序不一致導致的“大數加小數”的截斷誤差得到了緩解。這些變化和特性具體會給RL帶來哪些新的可能性目前還并沒有研究透徹，我們之后還會進一步探究。

一些額外的思考

事實上，K2-Thinking并不是第一個原生low-bit量化的主流模型。gpt-oss此前就已經在post-training階段使用了mxfp4的QAT，它的文檔中只提到這是為了將120b模型裝到單卡中。

對此，我有一些猜測（純屬個人臆測，不對準確性負責）：gpt-oss不是openai的主力模型，他們沒必要為了“單卡裝得下”這一個feature專門設計一個QAT流程；更有可能的是，這套QAT流程已經廣泛應用到了他們的主力模型中，在開源的gpt-oss上沿用這套方案反而是更自然的操作。而他們這么做的主要動機，就是比我們更早看到了前面提到的各種low-bit QAT的潛在收益。

至于我們這一次為什么選擇INT4而非MXFP4/NVFP4等更“先進”的格式，也確實如大家所說，是為了更好地支持非Blackwell架構硬件。

在我們的quant scale粒度（1x32）下，INT4的表示能力是不輸給以上兩種FP4格式的（4bit總共就16個數，scale的靈活度也非常重要）。而且基于INT4的W4A16 MoE，業界已經有了相當高效的Marlin kernel實現。因此，單純看W4A16這個場景，無論是從適配性、精度還是效率上，INT4都可以認為是一個足夠優秀的選擇了。

不過顯然，W4A16并不是我們對量化探索的終點，我相信W4A8甚至W4A4都是足夠具備可行性的方案。

隨著國外&國產芯片的不斷演進，對MXFP4等新一代低精度算子提供硬件支持的GPU會越來越多，我們的量化方案當然也會隨之不斷演進。