NVIDIA 发布 Nemotron-Labs-3-Puzzle-75B-A9B:压缩混合 MoE 模型,服务器吞吐量提升 2.03 倍
阅读原文· marktechpost.com把120B MoE压到75B后,同节点服务吞吐几乎翻倍,单张H100百万token并发从1涨到8。对长上下文RAG和AI编码助手这类对吞吐敏感的产品来说,这不是论文调优,是实打实的降本方案。
NVIDIA 发布 Nemotron-3-Super 的压缩变体 Nemotron-Labs-3-Puzzle-75B-A9B,总参数从 120.7B 降至 75.3B,活跃参数从 12.8B 降至 9.3B,保持 88 块混合布局(40 Mamba、40 MoE、8 注意力)。在 8×B200 节点上,8K/64K 场景匹配用户吞吐量≥100 tok/s 时,服务器吞吐量提升 2.03 倍。单 H100 上 1M-token 并发从 1 增至 8,权重占用从 70 GB 降至 44.5 GB。迭代式 Puzzle 方法平均得分比单步高 0.57。代价:Arena-Hard-V2 降 4.2 分、SWE-Bench 降 2.6 分。Hugging Face 提供 BF16、FP8、NVFP4 检查点。
像 Nemotron-3-Super 这样的大型混合 MoE 模型虽然准确,但部署成本高昂。其活跃参数、KV 缓存和 Mamba 状态限制了单个节点在给定每用户 token 速率下所能承载的用户数量。NVIDIA AI 团队发布了 Nemotron-3-Super 的压缩版本——Nemotron-Labs-3-Puzzle-75B-A9B。原始模型总参数量为 1207 亿,活跃参数量为 128 亿。压缩后的模型总参数量为 753 亿,活跃参数量为 93 亿。
部署目标在架构搜索开始前就已确定。目标一:在每用户每秒 100 个 token 的速率下,实现 2 倍的服务器吞吐量。目标二:在单块 H100 上,支持 8 个并发的 100 万 token 请求。Hugging Face 上提供了三个检查点:BF16、FP8 和 NVFP4。
1207 亿/128 亿活跃参数压缩至 753 亿/93 亿活跃参数,同时保留了原有的 88 块混合布局。
在匹配的 NVFP4 精度和匹配的用户吞吐量下,8 块 B200 的总吞吐量相比 Super 提升了 1.60 倍至 2.14 倍。
单块 H100 上 100 万 token 的并发数从 1 提升至 8,这得益于权重从 70 GB 降至 44.5 GB。
在相同压缩目标下,迭代式 Puzzle 方法比单步式 Puzzle 方法平均高出 0.57 分。
Arena-Hard-V2(-4.2)和 SWE-Bench(-2.6)是实际付出的代价;RULER 和 AA-LCR 则几乎没有变化。
Nemotron-Labs-3-Puzzle-75B-A9B
Nemotron-3-Super 是一个混合了 Mamba-Transformer 架构的 MoE 模型。Puzzle-75B-A9B 完全保留了原始模型的块布局。它共有 88 个块:40 个 Mamba 块、40 个 MoE 块和 8 个注意力块。
发生变化的是这些块内部的容量:
数量 | Super | Puzzle-75B-A9B | 比例 总参数 | 1207 亿 | 753 亿 | 62.4% 活跃参数 | 128 亿 | 93 亿 | 73.1% Mamba SSM 状态大小 | 128 | 96 | 75% MoE 路由专家中间层大小 | 2688 | 1280-2688 | 均值 59.9% 每个 token 激活的路由专家数 | 2 | 2-4 | 均值 18% 活跃路由专家容量(相对值) | 100% | 8.7%-62.3% | 均值 30.9%
路由专家数量、共享专家大小以及 MoE 潜在空间大小保持不变。注意力层未做改动。该研究提出的理由是,Nemotron-3-Super 本身在 KV 缓存方面已经非常高效。Mamba 层被统一剪枝,因为推理框架不支持每层使用不同的 SSM 状态大小。
其结果并非一个均匀缩小的教师模型。上图展示了在深度维度上的分配情况。Puzzle 在选定的中间层和后层保留了容量,而在其他层则大幅削减。
基准测试与性能表现
下表报告了在单个 8×B200 节点上,采用单步解码时的帕累托最优总吞吐量。
场景(输入/输出)UT 下限Super(tok/s)Puzzle-75B-A9B(tok/s)提升倍数50K / 2K>= 1005,1288,2101.60x50K / 2K>= 1253,7846,4121.69x50K / 2K>= 1502,5324,5231.79x8K / 64K>= 10020,93942,6012.03x8K / 64K>= 12513,07427,9182.14x8K / 64K>= 1508,52218,0472.12x
两个模型均在匹配的 NVFP4 权重、FP8 KV 缓存和 FP16 Mamba 状态下提供服务。因此,性能差距反映的是压缩效果,而非数值格式的变化。以预填充为主的 50K/2K 场景提升最小。以解码为主的 8K/64K 场景提升最大。
在单个 8×H100 节点上,UT = 100 时,提升幅度较小。在 50K/2K 场景下为 1.91 倍,在 8K/64K 场景下为 1.82 倍。这两个模型均使用 FP8 权重、FP8 KV 缓存和 FP32 Mamba 状态。
在单个 H100 上处理 1M 上下文时,瓶颈约束从计算转向内存。Super 的 NVFP4 权重占用约 70 GB 的 80 GB HBM 预算。每个 1M token 的请求会增加约 4 GB 的 KV 缓存。因此,有效并发数为 1。
Puzzle-75B-A9B 的 NVFP4 权重占用约 44.5 GB。注意力机制布局不变,因此每个请求的 KV 缓存开销不变。在 1M 上下文下的并发数提升至 8。该并发下的聚合解码吞吐量大约是 Super 单请求吞吐量的 4 倍。对 990K token 提示词的预填充速度大约快 1.2 倍。
迭代式 Puzzle 的工作原理
Puzzle 是一个分解式神经架构搜索框架,在此以 Puzzletron 实现。它定义了一个由备选层实现组成的离散搜索空间。每个备选方案都会获得一个质量评分。然后,一个混合整数规划会在部署约束下为每一层选择一个备选方案。
三种剪枝技术构成了该搜索空间:
中间通道剪枝:每个路由专家内部的通道根据其对专家输出的贡献进行排序。一个 MoE 层内的所有专家都被剪枝到统一大小,以保证内核兼容性。
Top-k 缩减:每个 token 被路由到的专家数量因层而异,上限为父模型的 k=22。
Mamba SSM 剪枝:SSM 状态大小从 128 通道降至 96 通道。
对 SSM 结果进行了测量。将 128 通道降至 96 通道,在解码阶段使 SSM 内核速度提升 1.2 倍至 1.3 倍。这一效果在 8 到 512 的批量大小范围内均成立。通道根据其对 Mamba 层输出的预估贡献进行排序。该预估基于 6700 万 token 的验证数据取平均值。附录 A 显示,在激进剪枝条件下,此方法优于随机通道选择。
原始公式假设替换质量的影响近似可加。每个候选模块在未经修改的父模型内部进行评分。这忽略了替换之间的高阶交互作用。
迭代式 Puzzle 将有限压缩与短程知识蒸馏恢复交替进行。它构建一个序列 M0, M1, … MR,而不是直接跳转到目标模型。评分针对当前压缩后的模型重新计算,而非原始父模型。
共使用了三个阶段:
MoE 权重降至教师模型容量的 75%,Mamba SSM 状态降至 75%。使用 240 亿 token 进行恢复训练。
MoE 权重降至教师模型容量的 60%。使用 432 亿 token 进行恢复训练。
激活的路由专家预算降至 50%,并采用异构分配。使用 528 亿 token 进行恢复训练。
上表将本方法与相同目标下的单步 Puzzle 基线进行了对比。三步流程在十个基准测试上的平均得分为 69.05,而基线为 68.48。在 MMLU-Pro、GPQA、HLE、AA-LCR、LiveCodeBench、SciCode 和 RULER-256K 上均取得了提升。IFBench-Instruction 下降了 0.2 分,IFBench-Prompt 下降了 0.5 分。
恢复:蒸馏、强化学习与冗长度控制
知识蒸馏使用了来自 Nemotron-3-Nano 的 30% 预训练数据和 70% SFT 数据。在 Puzzle 阶段,KD 使用 32K 序列长度。恢复阶段则在 128K 长度下训练,并扩展至 512K。预算上限为 1000 亿 token,全局批次大小为 1600 万 token,在 Megatron-LM 框架中运行。
RL 后训练采用了 Nemotron-3-Super RL 流水线的第二阶段,重点聚焦软件工程。阶段 2.1 进行了单步工具使用对比。阶段 2.2 转向端到端沙箱 RL,智能体在此阶段最多可运行 200 轮。两个阶段均使用 KL 惩罚系数为 0。团队对学习率进行了扫描,然后对得到的权重进行了平均。
上述图 4 展示了每个阶段的贡献。短上下文 KD 将大多数类别恢复至 Nemotron-3-Super 的 97% 以上。长上下文 KD 则专门提升了长输入和长生成基准的性能。研究团队表示,RL 在这些实验中的影响较小。
冗长性是一个不易察觉的细节。在最后一次 Puzzle 迭代之后,模型生成的 token 数量是 Super 的 132%。经过完整的恢复流水线后,这一比例降至 99%。
部署:量化与多 token 预测
生成了两种训练后量化方案:FP8 W8A8 面向 Hopper,NVFP4 W4A4 面向 Blackwell。
组件 | BF16 基线 | FP8 检查点 | NVFP4 检查点 稀疏与共享 MoE GEMM | BF16 | FP8 | NVFP4 Mamba GEMM | BF16 | FP8 | FP8 Mamba SSM 缓存 | FP32 | FP32 | FP16+SRK KV 缓存 | FP8 | FP8 | FP8 路由器 | FP32 | FP32 | FP32 注意力 QKV/输出、MoE 潜在投影、LM 头 | BF16 | BF16 | BF16
两种方案均在 256 个训练后 SFT 样本上进行了校准。NVFP4 使用了最大校准,而非 Super 所使用的 AutoQuantize 敏感性搜索。由此产生的检查点量化程度略高,但性能表现相似。
NVFP4 在 Hopper 上并非原生支持。它仍被用于 100 万上下文窗口的 H100 目标,因为该场景受限于 HBM 容量。
Puzzle-75B-A9B 继承了 Super 的共享 MTP 头。参数在 MTP 步骤之间共享,因此一个头在推理时递归应用。直接迁移 Super 训练好的头,得到了相似的接受长度。
研究团队随后发现了一个训练-推理不匹配问题。教师强制 MTP 训练输入的是完整的移位隐藏状态序列。而自回归草稿生成输入的则是目标模型和 MTP 生成隐藏状态的混合体。在较深的草稿位置上,接受率会下降。
对迁移后的头部进行持续训练解决了这一问题。在 draft 长度为 7 的 SPEED-Bench 上,平均接受长度从 3.45 提升至 4.34。这大约提升了 25% 到 30%,主要集中在较后的 draft 位置。与 Super 不同,NVFP4 检查点几乎没有退化:4.31 对比 4.34。
压缩在何处有益,在何处有害
基准测试 (BF16) | Super | Puzzle-75B-A9B | Delta MMLU-Pro | 83.8 | 82.4 | -1.4 AIME25 (无工具) | 92.2 | 89.7 | -2.5 GPQA (无工具) | 80.5 | 78.6 | -1.9 LiveCodeBench | 82.1 | 81.1 | -1.0 SciCode (子任务) | 42.3 | 40.6 | -1.7 SWE-Bench (OpenHands) | 59.5 | 56.9 | -2.6 Arena-Hard-V2 | 72.8 | 68.6 | -4.2 AA-LCR | 56.8 | 56.9 | +0.1 RULER 1M | 93.9 | 92.2 | -1.7 MMLU-ProX | 79.5 | 77.5 | -2.0
该研究论文自身的总结是,指令遵循和智能体类评测的损失最大。Arena-Hard-V2 是最差的情况,下降了 4.2 个点。RULER 在 256K、512K 和 1M 长度下,分数波动大致保持在 1 到 2 个点以内。
有三项 BF16 结果没有出现退化。AA-LCR 提升了 0.1,Scale AI Multi-Challenge 持平于 56.6,TauBench Telecom 提升了 0.4。
NVFP4 在压缩基础上带来的额外代价很小。在 RULER 1M 上,NVFP4 检查点得分为 93.2,高于 BF16 的 92.2。HLE 是 NVFP4 代价最明显的例子,分数从 16.5 降至 15.7。FP8 的结果见附录 E,与 BF16 的表现非常接近。FP8 检查点未报告 SWE-Bench 结果。
使用场景
单 GPU 上的超长上下文 RAG:一个处理 1M 上下文的文档分析服务,从 1 个并发请求提升至 8 个。在该并发度下,聚合解码吞吐量大约提升 4 倍。
交互式编程助手:在 8K/64K 范围内,当用户吞吐量 >= 100 tok/s 时,单个节点可提供 2.03 倍的 token 量。根据冗长度调整后,相当于每分钟完成的请求数提升了 2.16 倍。
预填充密集型文档处理管线:在 50K/2K 范围内仅获得 1.60 倍的提升。当提示词处理占据计算主导时,压缩的帮助较小。
智能体 SWE 循环:请对照你的任务组合,检查 SWE-Bench 上 2.6 个点的差距。RL 恢复技术针对这一能力进行了优化,但仅部分恢复了它。
在匹配的 NVFP4 和匹配的用户吞吐量下,总吞吐量相比 Super 提升了 1.60 倍到 2.14 倍
单张 H100 上的 1M token 并发数从 1 个请求提升至 8 个
在 draft 长度为 7 的 SPEED-Bench 上,MTP 接受长度从 3.45 提升至 4.34
在 RULER 的 256K、512K 和 1M 长度下,长上下文准确率保持在 1 到 2 个点以内
生成冗余在达到 Super 的 99% 时终止,因此 token 增益在请求层面得以保留。
发布了三个检查点:BF16、FP8 和 NVFP4。
Arena-Hard-V2 下降了 4.2 个点,SWE-Bench 下降了 2.6 个点。
RL 恢复带来的影响很小,论文中已直接说明。
Mamba 剪枝是均匀的,因为框架无法按层改变 SSM 状态大小。
潜在维度剪枝被放弃:NVFP4 MoE 内核需要潜在维度是 512 的倍数。
在这篇 v2 预印本中,正文和表格在多个吞吐量倍数上存在不一致。
分离式预填充的增益仅为 5% 到 7%,并且增加了服务复杂性。