昨天，分享了Kimi-K2-Thinking的测试体验。
其中，有提到，关于Kimi-K2-Thinking的量化感知训练（QAT）部分。
跟知乎@刘少伟 少伟聊过之后，特此分享，主要是他的心路历程。
知乎：https://www.zhihu.com/question/1969558404759544488/answer/1970539327902679960
开头一句话总结，
个人认为，在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会越来越多，我们的量化方案当然也会随之不断演进。
最后
真正想总结的已经写在文章的最开始了。希望这次K2-Thinking的发布能向大家传递出一个信息：在LLM时代，“量化”是一个可以跟“SOTA”和“Frontier”放在一起，甚至加快模型达到“Frontier”速度的概念。
PS：都看到这里，来个点赞 、在看 、关注 吧。 您的支持是我坚持的最大动力！
欢迎多多关注公众号「刘聪NLP」，加入交流群，交个朋友吧，一起学习，一起进步！