英伟达推出Nemotron-Labs-3-Puzzle-75B-A9B:压缩混合MoE大模型,服务器吞吐量提升2.03倍
阅读原文· marktechpost.com这是对 Nemotron-3-Super 的一次实用压缩,重点不在能力提升而在把 1M 上下文并发从 1 拉到 8,把吞吐量翻倍,是长上下文服务部署的工程优化参考。
英伟达发布Nemotron-Labs-3-Puzzle-75B-A9B,是Nemotron-3-Super的压缩变体。母模型总参120.7B、活跃12.8B;压缩后总参75.3B、活跃9.3B。在8×B200节点上,匹配NVFP4权重和用户吞吐量,总吞吐量提升1.60x至2.14x,8K/64K场景达2.03x。单张H100上1M token并发数从1升至8,权重从70GB降至44.5GB。采用迭代Puzzle压缩方法,平均分比单步Puzzle提升0.57。已发布BF16、FP8和NVFP4三种格式checkpoint。
像 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 激活的路由专家数 | 4 | 2-4 | 平均 50% 活跃路由专家容量(相对值) | 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 上处理 100 万上下文时,瓶颈约束从计算转向内存。Super 的 NVFP4 权重占用约 70 GB 的 80 GB HBM 预算。每个 100 万 token 的请求会增加约 4 GB 的 KV 缓存。因此,有效并发数为 1。
Puzzle-75B-A9B 的 NVFP4 权重占用约 44.5 GB。注意力布局不变,因此每个请求的 KV 开销不变。在 100 万上下文时的并发数提升至 8。该并发下的总解码吞吐量大约是 Super 单请求吞吐量的 4 倍。对 99 万 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 范围内,当 UT >= 100 tok/s 时,单个节点可提供 2.03 倍的 token 量。根据冗长度调整后,相当于每分钟完成的请求数提升了 2.16 倍。
预填充密集型文档处理管线:在 50K/2K 范围内仅获得 1.60 倍的提升。当提示词处理占据计算主导时,压缩的帮助较小。
智能体 SWE 循环:请对照你的任务组合,检查 SWE-Bench 上 2.6 个点的差距。RL 恢复技术针对这一能力进行了优化,但仅部分恢复了它。