NVIDIA 发布 Nemotron-Labs-3-Puzzle-75B-A9B：压缩混合 MoE 模型，服务器吞吐量提升 2.03 倍
NVIDIA Releases Nemotron-Labs-3-Puzzle-75B-A9B: A Compressed Hybrid MoE LLM Delivering 2.03x Server Throughput at Matched User Throughput
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 检查点。
把120B MoE压到75B后，同节点服务吞吐几乎翻倍，单张H100百万token并发从1涨到8。对长上下文RAG和AI编码助手这类对吞吐敏感的产品来说，这不是论文调优，是实打实的降本方案。
像 Nemotron-3-Super 这样的大型混合 MoE 模型虽然准确，但服务成本高昂。其激活参数、KV cache 和 Mamba 状态限制了单个节点在给定每用户 token 速率下所能承载的用户数量。NVIDIA AI 团队发布了 Nemotron-Labs-3-Puzzle-75B-A9B，这是 Nemotron-3-Super 的压缩变体。父模型拥有 120.7B 总参数和 12.8B 激活参数。压缩模型拥有 75.3B 总参数和 9.3B 激活参数。
部署目标在架构搜索开始之前就已确定。目标一是在每用户每秒 100 tokens 的速率下实现 2 倍服务器吞吐量。目标二是在单块 H100 上支持 8 个并发的 1M-token 请求。Hugging Face 上有三个检查点：BF16、FP8 和 NVFP4。
TL;DR
- 120.7B/12.8B 激活参数压缩至 75.3B/9.3B 激活参数，同时保留了 88 块的混合布局。
- 在匹配的 NVFP4 和匹配的用户吞吐量下，8xB200 总吞吐量相比 Super 提升 1.60 倍至 2.14 倍。
- 单块 H100 的 1M-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 | 比例 |
|---|---|---|---|
| 总参数量 | 120.7B | 75.3B | 62.4% |
| 激活参数量 | 12.8B | 9.3B | 73.1% |
| Mamba SSM 状态大小 | 128 | 96 | 75% |
| MoE 路由专家中间层维度 | 2688 | 1280-2688 | 均值 59.9% |
| 每个 token 激活的路由专家数 | 22 | 4-18 | 均值 50% |
| 激活路由专家容量（相对值） | 100% | 8.7%-62.3% | 均值 30.9% |
路由专家数量、共享专家规模和 MoE 隐层规模均保持不变。注意力层未做改动。该研究给出的理由是，Nemotron-3-Super 在 KV cache 方面已经非常高效。Mamba 层采用统一剪枝，因为推理框架不支持每层使用不同的 SSM 状态大小。
结果并不是一个被均匀缩小的教师模型。上图展示了各层的分配情况。Puzzle 在选定的中间层和靠后层保留了容量，而在其他层则大幅削减。
基准测试与性能
下表报告了在单个 8xB200 节点上、采用单步解码时的帕累托最优总吞吐量。
| 场景（输入/输出） | UT 下限 | Super（tok/s） | Puzzle-75B-A9B（tok/s） | 提升 |
|---|---|---|---|---|
| 50K / 2K | <= 100 | 5,128 | 8,210 | 1.60x |
| 50K / 2K | <= 125 | 3,784 | 6,412 | 1.69x |
| 50K / 2K | ≥ 150 | 2,532 | 4,523 | 1.79x |
| 8K / 64K | ≥ 100 | 20,939 | 42,601 | 2.03x |
| 8K / 64K | ≥ 125 | 13,074 | 27,918 | 2.14x |
| 8K / 64K | ≥ 150 | 8,522 | 18,047 | 2.12x |
两个模型均以匹配的 NVFP4 权重、FP8 KV cache 和 FP16 Mamba state 提供服务。因此，这一差距反映的是压缩，而非数值格式的变化。以 prefill 为主的 50K/2K 场景收益最小。以 decode 为主的 8K/64K 场景收益最大。
在单台 8xH100 节点上，UT = 100 时，收益较小。50K/2K 上为 1.91x，8K/64K 上为 1.82x。那里的两个模型均使用 FP8 权重、FP8 KV cache 和 FP32 Mamba state。
在单台 H100 上以 1M 上下文运行时，瓶颈从计算转为内存。Super 的 NVFP4 权重占据 80 GB HBM 预算中的约 70 GB。每个 1M token 请求增加约 4 GB 的 KV cache。因此有效并发数为 1。
Puzzle-75B-A9B 的 NVFP4 权重占据约 44.5 GB。注意力布局未变，因此每请求的 KV 开销未变。1M 上下文下的并发数升至 8。该并发数下的聚合 decode 吞吐量约为 Super 单请求吞吐量的 4x。990K token 提示词的 prefill 约快 1.2x。
迭代式 Puzzle 如何工作
Puzzle 是一个分解式神经架构搜索框架，在此实现为 Puzzletron。它定义了一个由备选层实现组成的离散搜索空间。每个备选方案获得一个质量分数。随后，一个混合整数规划在部署约束下为每一层选择一个备选方案。
三种剪枝技术构成搜索空间：
- 中间通道剪枝：每个路由专家内部的通道按其对专家输出的贡献进行排序。同一 MoE 层内的所有专家都被剪枝到统一大小，以保证内核兼容性。
- Top-k 缩减：一个 token 被路由到的专家数量因层而异，最多可达父模型的 k=22。
- Mamba SSM 剪枝：SSM 状态大小从 128 个通道降至 96 个通道。
对 SSM 结果进行了测量。将 128 个通道降至 96 个通道，使 SSM 内核在解码期间提速 1.2x 至 1.3x。这一效果在 8 到 512 的批大小范围内均成立。通道按其对 Mamba 层输出的估计贡献进行排序。该估计基于 67M token 的验证数据取平均。附录 A 表明，在激进剪枝下，这优于随机通道选择。
原始公式假设替换质量影响近似可加。每个候选块都在未修改的父模型内评分。这忽略了替换之间的高阶交互。
迭代式 Puzzle 将有界压缩与短时知识蒸馏恢复交替进行。它构建出一个序列 M0、M1、…… MR，而不是直接跳到目标。评分是针对当前已压缩的模型重新计算的，而非原始父模型。
共使用了三个阶段：
- MoE 权重降至教师容量的 75%，Mamba SSM 状态降至 75%。用 24B tokens 进行修复。
- MoE 权重降至教师容量的 60%。用 43.2B tokens 进行修复。
- 激活的路由专家预算降至 50%，采用异构分配。用 52.8B tokens 进行修复。
上表将这一方法与相同目标下的单步 Puzzle 基线进行了对比。三步流程在十个基准上平均为 69.05，而基线为 68.48。增益出现在 MMLU-Pro、GPQA、HLE、AA-LCR、LiveCodeBench、SciCode 和 RULER-256K 上。IFBench-Instruction 下降了 0.2 分，IFBench-Prompt 下降了 0.5 分。
恢复：知识蒸馏、RL 与冗长度
知识蒸馏在 Nemotron-3-Nano 的 30% 预训练数据和 70% SFT 数据上运行。在 Puzzle 阶段，KD 使用了 32K 序列长度。随后恢复阶段以 128K 训练，并扩展到 512K。预算最高为 100B tokens，全局批大小为 16M tokens，在 Megatron-LM 中进行。
RL 后训练采用了 Nemotron-3-Super RL 流程的第 2 阶段，聚焦于软件工程。阶段 2.1 进行了单步工具使用对比。阶段 2.2 转向端到端沙盒 RL，其中智能体最多运行 200 轮。两个阶段均使用 KL 惩罚为 0。团队对学习率进行了扫描，然后对所得权重取平均。
上图图 4 展示了每个阶段各自的贡献。短上下文知识蒸馏将大多数类别恢复到 Nemotron-3-Super 的 97% 以上。随后，长上下文知识蒸馏专门提升了长输入和长生成基准的表现。研究团队表示，在这些实验中，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+SR |
| KV cache | FP8 | FP8 | FP8 |
| Router | FP32 | FP32 | FP32 |
| Attention QKV/输出、MoE 潜在投影、LM head | BF16 | BF16 | BF16 |
两种方案均在 256 个后训练 SFT 样本上完成校准。NVFP4 使用的是最大校准，而非 Super 所用的 AutoQuantize 敏感度搜索。由此得到的检查点量化程度略为激进，但表现相近。
NVFP4 在 Hopper 上并不原生支持。它仍被用于 1M 上下文的 H100 目标，因为在那里 HBM 容量是瓶颈。
Puzzle-75B-A9B 从 Super 继承了一个共享的 MTP 头。参数在 MTP 各步骤之间共享，因此推理时一个头会递归应用。直接迁移 Super 训练好的头得到了相近的接受长度。
研究团队随后识别出一种训练-推理不匹配。教师强制式的 MTP 训练输入的是完整的移位隐藏状态序列。而自回归草稿则输入目标模型与 MTP 生成隐藏状态的混合。接受率在更深的草稿位置上下降。
对迁移后的头进行持续训练可以解决这一问题。在 SPEED-Bench 上，草稿长度为 7 时，平均接受长度从 3.45 升至 4.34。这大约是 25% 到 30%，集中在较后的草稿位置上。与 Super 不同，NVFP4 检查点几乎没有退化：4.31 对 4.34。
压缩在何处有益、在何处有害
| 基准测试（BF16） | Super | Puzzle-75B-A9B | 差值 |
|---|---|---|---|
| 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 checkpoint 得分为 93.2，高于 BF16 的 92.2。HLE 是 NVFP4 最明显的代价，从 16.5 降至 15.7。FP8 结果见附录 E，与 BF16 表现非常接近。FP8 checkpoint 未报告 SWE-Bench 成绩。
用例
- 单 GPU 上的超长上下文 RAG：一个 1M 上下文的文档分析服务从 1 个并发请求提升到 8 个。在该并发量下，聚合解码吞吐量约为 4 倍。
- 交互式编程助手：在 8K/64K 场景下，当 UT >= 100 tok/s 时，单个节点可服务的 token 量为 2.03 倍。按冗长度调整后，即每分钟完成的请求数为 2.16 倍。
- 预填充密集型文档流水线：50K/2K 场景仅获得 1.60 倍提升。当提示词处理主导计算时，压缩带来的帮助较小。
- 智能体 SWE 循环：请对照你的任务组合来检查 2.6 分的 SWE-Bench 差距。RL 恢复针对的正是这一能力，且仅部分恢复了它。
部署探索器
查看论文和模型权重。另外，欢迎在Twitter上关注我们，别忘了加入我们拥有15 万+ 成员的 ML SubReddit，并订阅我们的新闻通讯。等等！你在用 Telegram 吗？现在你也可以在 Telegram 上加入我们了。
来源：MarkTechPost（RSS） · marktechpost.com