Command Palette
Search for a command to run...
Nemotron-Labs-3-Puzzle-75B-A9B: Compressing Hybrid MoE LLMs
Nemotron-Labs-3-Puzzle-75B-A9B: Compressing Hybrid MoE LLMs
Abstract
We present Nemotron-Labs-3-Puzzle-75B-A9B, a compressed variant of Nemotron-3-Super optimized for interactive deployment. We designed the model to maximize server throughput under high user throughput constraints. In interactive serving workloads on a single 8×B200 node, Puzzle-75B-A9B achieves approximately 2× higher server throughput than Nemotron-3-Super at matched user throughput constraints. In ultra-long-context deployment on a single H100 GPU, the compressed model increases 1M-token concurrency from 1 request to 8 requests. Puzzle-75B-A9B is constructed using a multi-stage pipeline that combines the Iterative Puzzle compression framework with knowledge distillation, reinforcement learning, quantization, and Multi-Token Prediction head. The compression process jointly optimizes heterogeneous MoE pruning, active parameter budget, and Mamba pruning to improve inference efficiency while preserving model quality. We evaluate Puzzle-75B-A9B on a broad suite of reasoning, coding, multilingual, long-context, and agentic benchmarks. Despite substantial compression, the model retains strong downstream accuracy relative to the parent model across a wide range of tasks. These results demonstrate that large hybrid MoE models can be substantially optimized for deployment efficiency while maintaining strong downstream capability. Our model is publicly available on Hugging Face.
One-sentence Summary
Nemotron-Labs-3-Puzzle-75B-A9B, a compressed hybrid MoE LLM, achieves up to 2× higher server throughput on an 8×B200 node and 8× 1M-token concurrency on a single H100 GPU by employing a multi-stage compression pipeline that integrates Iterative Puzzle, knowledge distillation, reinforcement learning, quantization, and Multi-Token Prediction, jointly optimizing MoE pruning, active parameter budget, and Mamba pruning for efficient interactive deployment.
Key Contributions
- Nemotron-Labs-3-Puzzle-75B-A9B is a compressed variant of Nemotron-3-Super built through a multi-stage pipeline that combines the Iterative Puzzle framework with knowledge distillation, reinforcement learning, quantization, and a Multi-Token Prediction head, jointly optimizing heterogeneous MoE pruning, active parameter budget, and Mamba pruning to improve inference efficiency while preserving quality.
- The compressed model achieves approximately 2× higher server throughput than the parent model on a single 8×B200 node under matched user throughput constraints for interactive serving, and increases sustainable 1M-token concurrency from 1 request to 8 requests on a single H100 GPU for ultra-long-context deployment.
- Evaluations across reasoning, coding, multilingual, long-context, and agentic benchmarks show that Puzzle-75B-A9B retains strong downstream accuracy relative to the parent model, demonstrating that large hybrid MoE models can be substantially optimized for deployment efficiency while maintaining strong capability.
Introduction
The authors tackle the inference efficiency bottleneck in deploying large language models, focusing on hybrid Mixture-of-Experts (MoE) architectures like Nemotron-3-Super. Production settings impose tight constraints on throughput, latency, and device memory, yet existing post-training compression methods (pruning, quantization, distillation) often force a rigid tradeoff between efficiency and capability preservation, especially when optimizing multiple interdependent dimensions such as active parameters, KV-cache size, and Mamba state. The authors introduce Iterative Puzzle, a sequential compression framework that interleaves hardware-aware structural pruning with short knowledge distillation recovery phases, enabling non-uniform capacity allocation across network depth, routed experts, active experts, and SSM state size. The resulting model, Nemotron-Labs-3-Puzzle-75B-A9B, compresses the 120.7B/12.8B total/active parent to 75.3B/9.3B while delivering roughly 2× server throughput on interactive workloads and supporting 8 concurrent 1M-token requests on a single GPU, compared to the parent model’s limit of 1.
Method
The authors propose a two-stage pipeline for model compression and performance recovery, transforming the original Nemotron-3-Super model into a smaller, deployment-efficient architecture. The first stage focuses on structural compression using an iterative variant of the Puzzle framework, while the second stage addresses performance degradation through knowledge distillation and reinforcement learning.
In the structural compression stage, the authors leverage the Puzzle framework, a decomposed neural architecture search method. The search space primarily targets Mixture-of-Experts (MoE) pruning strategies and Mamba State Space Model (SSM) pruning. For MoE layers, they apply intermediate channel pruning and top-k reduction, allowing the framework to build a heterogeneous architecture that allocates different capacities to each layer based on its impact on accuracy. Mamba SSM state sizes are also pruned to accelerate the decode stage. To overcome the limitations of single-shot architecture search, which fails to model higher-order interactions between block replacements, the authors introduce Iterative Puzzle. This sequential procedure alternates between moderate structural compression and short knowledge distillation recovery phases. Let M0 denote the original pretrained model and T the teacher model. Iterative Puzzle constructs a sequence of models M0→M1→⋯→MR. At each iteration r, a Puzzle search space is constructed around Mr−1, and a compressed intermediate model M˘r is produced under an incremental pruning budget. This model is then healed via knowledge distillation from T:
Mr=Heal(M˘r,T)
This iterative approach allows replacement scores to be recomputed in the context of the current compressed model, making the search adaptive to changing internal representations.
Following compression, the authors apply a comprehensive recovery phase. Knowledge distillation is performed using a mixed dataset of pretraining and supervised fine-tuning data. In the iterative compression phase, distillation matches the logits of the original model at a 32K sequence length. In the final recovery phase, the context length is extended to 128K and then 512K tokens to restore long-context capabilities. Subsequently, reinforcement learning is applied to recover capabilities sensitive to compression, specifically targeting software engineering tasks. The RL pipeline involves single-step tool-use comparison and full end-to-end sandbox environments. To ensure training stability and high model quality, the authors employ a checkpoint averaging strategy across multiple runs with varying learning rates.
The implementation involves three specific compression and recovery stages. The first stage reduces MoE weights and Mamba SSM state size to 75% of the teacher capacity. The second stage further reduces MoE weights to 60%. The final stage constrains the activated routed-expert budget to 50% of the teacher, allowing heterogeneous allocation across layers. The resulting architecture preserves the parent hybrid block layout with 88 total blocks.
As shown in the figure below, the Puzzle framework selects a highly non-uniform allocation of active routed MoE capacity across the model depth.
Rather than uniformly scaling down the teacher, the final architecture aggressively prunes many layers while preserving substantially more routed capacity in selected middle and late layers, demonstrating that active MoE capacity is allocated according to measured layer sensitivity.
To further improve inference speed, the authors apply post-training quantization to produce FP8 and NVFP4 checkpoints tailored for Hopper-class and Blackwell-class GPUs, respectively. Calibration is performed on subsets of fine-tuning data, with specific precision choices made for different operators to balance memory overhead and accuracy.
Finally, the authors leverage native Multi-Token Prediction capabilities for speculative decoding. The model utilizes a shared MTP head formulation where parameters are shared across MTP steps, allowing a single prediction head to be applied recursively at inference time. To address the mismatch between teacher-forced MTP training and autoregressive speculative decoding, the authors continue training the transferred MTP head. This training scheme improves the stability of the internal drafter under autoregressive use, particularly at deeper draft positions where distribution mismatches are most pronounced, significantly boosting acceptance lengths without introducing additional runtime mechanisms.
Experiment
The evaluation benchmarks Nemotron-3-Puzzle-75B-A9B against its parent model, showing that iterative compression, knowledge distillation, and recovery training preserve most reasoning, coding, long-context, and multilingual capabilities, with only modest degradation even under aggressive quantization. Ablations confirm that iterative compression outperforms a single-shot approach and that short-context KD, long-context KD, and a small RL step each contribute to recovering accuracy and reducing verbosity. Performance analysis on an 8×B200 node demonstrates that the compressed model delivers 1.6× to 2.14× higher total throughput under interactive user-throughput constraints, and further gains are possible through prefill disaggregation and multi-token prediction, yielding a strong accuracy–efficiency tradeoff for production deployment.
The Puzzle-75B-A9B compressed model retains 62% of the parent's total parameters and 73% of active parameters by pruning the Mamba SSM state size to 75% and applying heterogeneous MoE routing, which reduces mean activated experts per token to half and active routed expert capacity to 31%. On an 8×B200 node, this architectural compression yields 1.6× to 2.14× higher total throughput under interactive serving constraints, while using the compressed model only for prefill alongside the full decode model avoids quality loss. The compression keeps the hybrid block layout but shrinks the Mamba state size to 75% and varies MoE dimensions across layers, bringing active parameters down to 73% of the parent. Activated routed experts per token drop to a mean of 50%, with per-layer active expert capacity ranging from 8.7% to 62.3% of the parent. Throughput on a single 8×B200 node improves by 1.6× to 2.14× at user-throughput thresholds of 100, 125, and 150 tokens per second. Using the compressed model for prefill and the full model for decode preserves benchmark scores, confirming that the gains stem from targeted compression of the compute-heavy prefill stage.
The FP8 checkpoint quantizes the heaviest compute operators—MoE and Mamba GEMMs plus the KV cache—to FP8, while keeping attention projections, embeddings, and the router in higher precision. The NVFP4 checkpoint further reduces the MoE GEMMs to 4-bit precision, making it slightly more aggressive but still matching accuracy, and is paired with FP16 Mamba state using stochastic rounding. The FP8 checkpoint applies FP8 only to the sparse and shared MoE GEMMs, Mamba linear layers, and KV cache, leaving all other components in BF16 or FP32. The NVFP4 checkpoint pushes the MoE GEMMs to NVFP4 while keeping Mamba GEMMs in FP8, resulting in a quantization that is slightly more aggressive than the AutoQuantize baseline but achieves similar accuracy.
The Puzzle-75B-A9B model retains most of the parent Super model's accuracy across general knowledge and reasoning benchmarks, with only minor drops. Under NVFP4 quantization, Puzzle-75B-A9B shows negligible degradation compared to its BF16 scores, and in some cases matches the quantized Super model exactly, demonstrating robustness to aggressive precision reduction. Puzzle-75B-A9B achieves 82.4 on MMLU-Pro in BF16, just 1.4 points below Super's 83.8, and the NVFP4 score of 82.2 is nearly identical. On AIME25, Puzzle-75B-A9B in NVFP4 matches Super's NVFP4 score of 89.9 exactly, while its BF16 score of 89.7 is 2.5 points behind Super's 92.2. Across all five reported benchmarks, the NVFP4 variant of Puzzle-75B-A9B never drops more than 1.3 points below the BF16 variant, and often the gap is 0.2 points or less. Reasoning benchmarks like HMMT and GPQA show Puzzle-75B-A9B trailing Super by only 0.8–1.6 points in BF16, and the NVFP4 gaps are similarly modest.
The compressed Puzzle-75B-A9B model nearly matches the accuracy of the original Nemotron-3-Super (70.74% vs. 71.93%) while delivering 2.03× higher server throughput in single-step mode. With multi-token prediction, throughput rises to 4.63× that of Super, and when adjusted for generation verbosity, it completes 4.91× more user requests per minute. The smaller Nano model achieves higher raw throughput but at the cost of significantly lower accuracy and more verbose outputs, leading to a lower effective request completion rate. Puzzle-75B-A9B with multi-token prediction achieves 4.63× the total throughput of the Super model at a fixed user-perceived rate of 100 tokens per second. In single-step inference, Puzzle-75B-A9B provides 2.03× the throughput of Super while accuracy drops by only 1.19 percentage points.
The Puzzle-75B-A9B model's MTP head achieves consistently higher average acceptance lengths than the original Super MTP across all evaluated categories, with notable gains in multilingual and coding tasks. Quantization to NVFP4 has only a marginal impact on acceptance lengths, demonstrating that the improved draft quality is robust to compression. Puzzle-75B-A9B MTP acceptance lengths exceed Super MTP in every category, with gains of more than 0.8 tokens on average. NVFP4 quantization reduces acceptance lengths by at most 0.08 tokens, so the draft quality remains virtually unchanged.
The evaluation examines a compressed Mamba-MoE hybrid model that prunes the Mamba state size and applies heterogeneous expert routing to reduce active parameters and activated expert counts. Throughput tests on an 8×B200 node demonstrate that using this compressed model solely for prefill while retaining the full model for decode yields up to 2.14× higher interactive serving throughput without any quality loss. Benchmark assessments show the compressed model nearly matches the parent’s accuracy and tolerates FP8 and NVFP4 quantization with negligible degradation, while its multi-token prediction head achieves consistently higher acceptance lengths, further amplifying effective throughput.