揭开Nemotron Labs 3 Puzzle 75B A9B的面纱：压缩混合MoE大语言模型实现2.03倍服务器吞吐量

大型混合MoE模型（如Nemotron-3-Super）虽然精度高，但服务成本高昂。其活跃参数、KV缓存和Mamba状态限制了单个节点在给定每用户token速率下所能承载的用户数量。NVIDIA AI团队发布了Nemotron-Labs-3-Puzzle-75B-A9B，这是Nemotron-3-Super的压缩变体。原模型总参数120.7B，活跃参数12.8B；压缩后总参数75.3B，活跃参数9.3B。

部署目标在架构搜索开始前即已确定。目标一：在每用户每秒100 token的条件下，实现2倍服务器吞吐量。目标二：在单块H100上实现8个并发1M token请求。Hugging Face上提供了三个检查点：BF16、FP8和NVFP4。

快速概览：120.7B/12.8B活跃参数压缩至75.3B/9.3B活跃参数，保留了88块混合布局；8xB200总吞吐量在匹配NVFP4和匹配用户吞吐量的条件下，相比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-3-Super是一种混合型Mamba-Transformer MoE模型。Puzzle-75B-A9B完全保留了原模型的块布局，共88块：40个Mamba块、40个MoE块和8个注意力块。变化在于块内部的容量：总参数降至62.4%，活跃参数降至73.1%，Mamba SSM状态大小降至75%，MoE路由专家中间大小均值降至59.9%，每token激活的路由专家数降至50%（均值），活跃路由专家容量相对降至30.9%（均值）。路由专家数量、共享专家大小和MoE潜在大小保持不变。注意力层未改动，因为Nemotron-3-Super的KV缓存效率已经很高。Mamba层被均匀剪枝，因为推理框架不支持每层不同的SSM状态大小。

结果是并非均匀缩小的教师模型。上图显示了深度上的分配：Puzzle保留了中间和后期的部分层容量，而大幅削减了其他层。

基准测试与性能：下表报告了在单个8xB200节点上，单步解码的帕累托最优总吞吐量。在50K输入/2K输出场景下，用户吞吐量（UT）≥100时，Super为5,128 tok/s，Puzzle-75B-A9B为8,210 tok/s，提升1.60倍；UT≥125时提升1.69倍；UT≥150时提升1.79倍。在8K输入/64K输出场景下，UT≥100时提升2.03倍，UT≥125时提升2.14倍，UT≥150时提升2.12倍。两种模型均在匹配的NVFP4权重、FP8 KV缓存和FP16 Mamba状态下服务，因此差距反映的是压缩效果而非数值格式变化。预填充密集的50K/2K场景增益最小，而解码密集的8K/64K场景增益最大。

在单个8xH100节点上，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分。

恢复：蒸馏、强化学习与冗长度。知识蒸馏使用了30%的预训练数据和70%的来自Nemotron-3-Nano的SFT数据。Puzzle阶段使用32K序列长度进行KD，恢复阶段则训练在128K长度，并扩展至512K。预算最高达1000亿token，全局批大小为1600万token，使用Megatron-LM。RL后训练采用Nemotron-3-Super RL管线的第二阶段，专注于软件工程：第2.1阶段进行单步工具使用比较；第2.2阶段转向端到端沙盒RL，智能体最多运行200轮。两个阶段均使用KL惩罚为0。团队对学习率进行了扫描，然后平均了所得权重。上图显示了各阶段的贡献：短上下文KD将大多数类别恢复至Nemotron-3-Super的97%以上；长上下文KD则专门提升了长输入和长生成基准。研究团队表示，RL在此次实验中的影响较小。冗长度是一个低调的细节：最后一次Puzzle迭代后，模型生成的token数为Super的132%；经过完整恢复管线后，降至99%。

部署：量化与多token预测。两种训练后量化方案：FP8 W8A8针对Hopper，NVFP4 W4A4针对Blackwell。组件包括稀疏和共享MoE GEMM、Mamba GEMM、Mamba SSM缓存、KV缓存、路由器等，各有不同精度。两种方案均在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）对比显示，MMLU-Pro下降1.4分，AIME25下降2.5分，GPQA下降1.9分，LiveCodeBench下降1.0分，SciCode下降1.7分，SWE-Bench下降2.6分，Arena-Hard-V2下降4.2分，AA-LCR提升0.1分，RULER 1M下降1.7分，MMLU-ProX下降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紧密跟踪。SWE-Bench未报告FP8检查点结果。

用例：单GPU超长上下文RAG：文档分析服务在1M上下文下从1个并发请求增至8个，聚合解码吞吐量约为4倍。交互式编程助手：在8K/64K场景且UT≥100时，一个节点服务的token量为2.03倍，考虑冗长度调整后，每分钟完成的请求数为2.16倍。预填充密集的文档管线：50K/2K场景仅提升1.60倍，压缩在提示处理主导计算时帮助较小。智能体SWE循环：需根据任务混合情况评估2.6分的SWE-Bench差距；RL恢复曾针对此能力，但仅部分恢复。

部署探索：请查阅论文和模型权重。欢迎关注Twitter，加入15万+ML SubReddit和订阅Newsletter。也可通过Telegram加入。如需合作推广GitHub Repo、Hugging Face页面、产品发布或网络研讨会，请联系我们。本文首发于MarkTechPost。