揭開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。