打开网易新闻 查看精彩图片

模型压缩是一项需要在精度与性能之间反复权衡的工作。NVIDIA在开发Nemotron 3.5 Lightning NVFP4检查点时,选择了量化感知蒸馏(QAD)这一方案,将模型从66 GB压缩至22 GB,同时实现最高4倍的吞吐量提升,并保持接近原始精度。

背景:为什么PTQ不够用

训练后量化(PTQ)是最常见的模型压缩手段,适用于大多数场景。但当目标是在更严格的内存约束下实现更高吞吐量时,PTQ的局限就显现了——激进量化带来的精度损失难以挽回。QAD则不同,它允许模型在训练过程中主动适应量化噪声,从而在大幅压缩的同时将精度损失降到最低。

QAD的工作原理

QAD采用"教师-学生"框架。完整精度的BF16模型充当冻结的教师,量化后的模型作为学生。整个过程分两个阶段:

第一阶段,对BF16基础模型运行PTQ,将其权重量化至W4A16,生成量化学生检查点。此阶段有意将量化推得更激进——中位精度恢复率目标定在95%至99%之间,而非PTQ单独使用时的99%以上。这种适度的精度下降是刻意为之,目的是为QAD留出足够的恢复空间。

第二阶段,学生模型的每次前向传播都经过模拟量化,使其在推理时能提前感知量化噪声。同时,通过KL散度损失函数,学生被训练成与教师的logit输出对齐,而非仅仅预测下一个Token。两个信号共同作用,使QAD能在激进量化的同时维持高质量。

PTQ配方的选择

开发团队测试了多种PTQ配方,核心差异在于权重校准方式和Mamba投影层、KV缓存的量化激进程度。最终选用的是four_over_six配方,在W4A16 Mamba线性层上取得了最佳的精度-性能权衡。

所有配方共享若干基本设置:lm_head量化为W4A16(称为"忠实lm_head"),注意力投影层保留BF16,校准使用1000个样本,在单台NVIDIA DGX B300上完成。校准序列长度方面,实验证明32K序列长度在four_over_six配方上效果最佳。

量化尺度处理策略

QAD训练中有两种尺度处理策略,选择取决于PTQ配方类型:

动态尺度QAD:从最大值校准的PTQ检查点出发,权重和激活尺度在训练中实时重算,两者随模型学习同步调整。

冻结尺度QAD:从基于均方误差(MSE)校准的PTQ检查点出发,由于MSE的尺度搜索代价过高,无法每步重算,因此PTQ阶段校准好的尺度在整个训练过程中保持固定,只有权重被更新。

训练配置

在中间检查点上的实验验证了QAD的有效性。以检查点A为例,PTQ单独使用时中位精度恢复率为96.33%,AIME 2025下降3.70分;经过200次QAD迭代后,恢复率提升至99.72%,AIME 2025与BF16基线的差距缩小至0.57分。检查点B的实验同样印证了这一规律:PTQ达到95.84%,QAD达到98.53%,AA v4.1 Index从20.03升至23.48(BF16基线为24.81)。

序列长度对训练效果影响显著,尤其是长上下文基准测试。最终QAD运行采用522K序列长度,数据集使用NVIDIA开源的Nemotron-Post-Training v1和v2。蒸馏以5e-6的恒定学习率进行,关闭dropout,梯度裁剪设为1.0,在两节点×8 GPU(TP=2, EP=4)上运行。

最终检查点评估

最终NVFP4检查点采用了更保守的量化策略以优化精度,PTQ中位精度恢复率已达99.24%,QAD为98.97%,两者都已非常接近BF16基线。QAD的收益主要体现在智能体和编程基准上:Terminal-Bench v2.1高出PTQ 3.79分,SWE-Bench多语言高出1.07分,HLE高出0.65分,BrowseComp、SWE-Bench Verified和PinchBench也有小幅提升。

端到端复现

完整QAD配方已集成至NVIDIA Model Optimizer,通过单一启动文件megatron_lm_qad.yaml即可端到端运行整个流水线,同时支持Nemotron 3 Nano和Nemotron 3.5 Lightning。同样的流程也可通过Megatron-Bridge运行,该库是NeMo框架中的PyTorch原生库,提供Hugging Face与Megatron Core之间的双向检查点转换。

最终NVFP4检查点已发布在Hugging Face上。完整工具链和文档可在NVIDIA Model Optimizer GitHub仓库中获取。

Q&A

Q1:量化感知蒸馏(QAD)和训练后量化(PTQ)有什么区别?

A:PTQ是在训练完成后直接对模型权重量化,简单高效但精度损失难以挽回,激进量化时尤为明显。QAD则在量化后继续训练,让学生模型在模拟量化环境中学习,同时通过KL散度损失对齐教师模型的输出,从而主动适应量化噪声并恢复精度。简单说,PTQ接受量化误差,QAD让模型适应并克服它。

Q2:Nemotron 3.5 Lightning NVFP4检查点的内存占用和推理速度有多大提升?

A:Nemotron 3.5 Lightning NVFP4检查点从BF16的66 GB压缩至22 GB,体积缩减至原来的约三分之一,同时实现最高4倍的推理吞吐量提升。这一压缩效果来自将大量权重量化为4位精度(W4A16 NVFP4格式),而精度损失通过QAD恢复至接近BF16基线水平,最终中位精度恢复率达到98.97%至99.24%。

Q3:如何在自己的模型上复现QAD流程?

A:可以通过NVIDIA Model Optimizer复现完整QAD流程。首先使用mtq.quantize对BF16模型运行PTQ,生成量化学生检查点;然后通过megatron_lm_qad.yaml启动文件运行蒸馏阶段,将学生模型对齐冻结的教师模型。数据方面推荐使用NVIDIA开源的Nemotron-Post-Training v1/v2数据集。最终导出步骤会生成可部署的Hugging Face NVFP4检查点。详细代码示例和配置文件可在NVIDIA Model Optimizer GitHub仓库中找到。