Nemotron-Labs-3-Puzzle-75B-A9B: Compressing Hybrid MoE LLMs

Background Large hybrid Mixture-of-Experts (MoE) models that combine attention with Mamba state-space layers offer strong accuracy, but their size and memory footprint limit interactive deployment throughput. The parent model here, Nemotron-3-Super, spans 120.7B total / 12.8B active parameters over a hybrid stack of 40 Mamba blocks, 40 MoE blocks, and 8 attention blocks. Which resource binds depends on the workload: long-input requests are compute-bound, long-output requests are IO-bound, and ultra-long-context requests are capped by device memory through the KV-cache and Mamba-state footprint.

배경 Attention과 Mamba state-space layer를 결합한 대형 hybrid Mixture-of-Experts(MoE) 모델은 높은 정확도를 제공하지만, 그 크기와 메모리 사용량이 interactive 배포 처리량을 제한한다. 여기서의 부모 모델 Nemotron-3-Super는 Mamba 블록 40개, MoE 블록 40개, attention 블록 8개의 hybrid 스택에 걸쳐 총 1207억 / 활성 128억 파라미터 규모다. 어떤 자원이 병목이 되는지는 워크로드에 따라 갈린다: 긴 입력 요청은 compute-bound, 긴 출력 요청은 IO-bound이며, 초장문맥 요청은 KV-cache와 Mamba-state 메모리 때문에 디바이스 메모리가 한계로 작용한다.

Problem Deploying a very large hybrid MoE model like Nemotron-3-Super under strict per-user latency constraints limits how many concurrent users a server can serve, especially at very long context lengths. Concretely, a single H100 can hold only 1 concurrent 1M-token request of the parent model once parameters, KV cache, and Mamba state are all accounted for. It is like a restaurant where every table must be served within minutes — the real question is not how much food the kitchen can cook in total, but how many tables it can seat at once without any of them waiting. The paper fixes the per-user floor at 100 tokens/s and asks how many users fit under that condition.

문제 Nemotron-3-Super 같은 매우 큰 hybrid MoE 모델을 엄격한 사용자별 지연 제약 하에서 배포하면, 특히 매우 긴 문맥 길이에서 서버가 동시에 처리할 수 있는 사용자 수가 제한된다. 실제로 H100 한 장에서는 파라미터·KV cache·Mamba state를 모두 합산하면 부모 모델의 1M-token 요청을 동시에 1개밖에 수용하지 못한다. 모든 테이블에 몇 분 안에 음식이 나가야 하는 레스토랑과 같다 — 진짜 질문은 주방이 총 몇 인분을 만들 수 있느냐가 아니라, 어느 테이블도 기다리게 하지 않으면서 동시에 몇 테이블을 받을 수 있느냐다. 논문은 사용자당 100 tokens/s를 하한으로 고정하고 그 조건에서의 동시 수용량을 묻는다.

Prior limitations Existing post-training compression — pruning, quantization, distillation — faces an acute efficiency-versus-capability tradeoff on hybrid MoE architectures, whose cost depends jointly on total parameters, active parameters, KV-cache size, and Mamba-state size. In particular, the original Puzzle framework is single-shot: each candidate layer replacement is scored by swapping it alone into the parent model, assuming quality impacts are approximately additive. That misses higher-order interactions — pruning one MoE layer changes the representations later layers see, making some blocks more important and others more redundant than their replace-one scores suggest.

기존 한계 Pruning·quantization·distillation 같은 기존 post-training 압축은 효율과 능력 보존 사이의 까다로운 tradeoff를 안고 있으며, 비용이 총 파라미터·활성 파라미터·KV-cache·Mamba-state 크기에 동시에 걸려 있는 hybrid MoE 구조에서 이 문제가 특히 심하다. 특히 원래의 Puzzle 프레임워크는 single-shot이다: 각 후보 layer 교체를 부모 모델에 하나씩만 끼워 넣어 점수를 매기고, 품질 영향이 근사적으로 가산적(additive)이라고 가정한다. 이는 고차 상호작용을 놓친다 — 한 MoE layer를 pruning하면 뒤 layer들이 보는 표현이 바뀌어, 어떤 블록은 replace-one 점수가 시사하는 것보다 더 중요해지고 어떤 블록은 더 불필요해진다.

Goal Produce a substantially smaller hybrid MoE model that maximizes server throughput under interactive deployment constraints while retaining as much of the parent model's downstream accuracy as possible. The paper states its deployment targets up front: roughly 2x server throughput over the parent at an operating point of at least 100 tokens/s per user on one 8xB200 node, and 8 concurrent 1M-token requests on a single H100. The optimization is anchored to two representative regimes: prefill-heavy 50K/2K (long-context analysis, RAG) and decode-heavy 8K/64K (reasoning-style generation).

목표 Interactive 배포 제약 하에서 서버 처리량을 최대화하면서, 부모 모델의 다운스트림 정확도를 최대한 유지하는 훨씬 작은 hybrid MoE 모델을 만드는 것. 논문은 배포 목표를 처음부터 명시한다: 8xB200 노드 1대에서 사용자당 최소 100 tokens/s 운영점 기준 부모 대비 약 2배 서버 처리량, 그리고 H100 한 장에서 1M-token 요청 8개 동시 처리. 최적화는 prefill-heavy 50K/2K(장문맥 분석·RAG)와 decode-heavy 8K/64K(reasoning형 생성)라는 두 대표 regime에 맞춰진다.

Method The Iterative Puzzle framework alternates moderate hardware-aware pruning with short knowledge-distillation (KD) healing phases, then finishes with long-context KD, reinforcement-learning (RL) recovery, and a Multi-Token Prediction (MTP) head. Three staged rounds each target a different axis: round 1 trims MoE weights and the Mamba SSM state to 75% of the teacher (healed on 24B tokens), round 2 cuts MoE weights to 60% (43.2B tokens), and round 3 constrains the activated routed-expert budget to 50%, letting a mixed-integer program distribute it non-uniformly across layers (52.8B tokens). Recovery KD distills the teacher's logits on a 30% pretraining / 70% SFT mix, growing sequence length from 32K to 128K and then 512K, while RL focuses on software engineering — the capability that degraded most under compression. The rhythm is that of pruning a bonsai: trim a little, let it recover, re-judge which branches matter, then trim again — rather than hacking out the final shape in one cut.

방법 Iterative Puzzle 프레임워크는 완만한 hardware-aware pruning과 짧은 knowledge-distillation(KD) 회복 단계를 번갈아 수행한 뒤, 장문맥 KD·강화학습(RL) 복구·Multi-Token Prediction(MTP) head로 마무리한다. 3개 라운드가 각기 다른 축을 겨냥한다: 1라운드는 MoE 가중치와 Mamba SSM state를 teacher의 75%로 줄이고(24B token으로 heal), 2라운드는 MoE 가중치를 60%로(43.2B token), 3라운드는 활성 routed-expert 예산을 50%로 제한하되 mixed-integer program이 layer별로 불균일하게 배분한다(52.8B token). 복구 KD는 pretraining 30% + SFT 70% 데이터에서 teacher logit을 증류하며 시퀀스 길이를 32K→128K→512K로 늘리고, RL은 압축으로 가장 크게 손상된 능력인 software engineering에 집중한다. 리듬은 분재 가지치기와 같다: 조금 자르고, 회복시키고, 어느 가지가 중요한지 다시 판단한 뒤 또 자른다 — 최종 형태를 한 번에 잘라내지 않는다.

Key idea Compression is treated as re-balancing rather than shrinking: a mixed-integer program allocates MoE capacity non-uniformly across depth from measured layer sensitivity, while the Mamba SSM state is pruned uniformly from 128 to 96 channels (an inference-framework constraint) and attention layers are left untouched. The final model keeps between 8.7% and 62.3% (mean 30.9%) of the teacher's per-layer active routed-expert capacity — like remixing a track by adjusting each instrument's fader individually instead of dragging down the master volume, layers that carry the signal keep their capacity while redundant ones are cut hard. Re-scoring candidates after each healing phase lets later pruning decisions account for how earlier cuts changed the network's internal representations — something a one-shot search cannot do.

핵심 아이디어 압축을 축소가 아니라 재배분으로 다룬다: mixed-integer program이 측정된 layer 민감도에 따라 MoE 용량을 깊이 방향으로 불균일하게 배분하고, Mamba SSM state는 (inference framework 제약으로) 전 layer 균일하게 128→96 채널로 pruning하며, attention layer는 손대지 않는다. 최종 모델은 layer별로 teacher 활성 routed-expert 용량의 8.7%~62.3%(평균 30.9%)만 남긴다 — master 볼륨을 통째로 내리는 대신 악기별 페이더를 따로 조절해 트랙을 리믹스하듯, 신호를 싣는 layer는 용량을 지키고 불필요한 layer는 과감히 깎는다. 매 회복 단계 후 후보 점수를 다시 계산하므로, 앞선 절단이 내부 표현을 어떻게 바꿨는지를 이후의 pruning 결정이 반영할 수 있다 — one-shot 탐색으로는 불가능한 부분이다.

Validation Evaluated on a suite spanning general knowledge (MMLU-Pro), reasoning and coding (AIME25, HMMT, GPQA, LiveCodeBench, SciCode, HLE), agentic tasks (Terminal-Bench, SWE-Bench, TauBench V2), instruction following (IFBench, Arena-Hard-V2), long context (AA-LCR, RULER at 256K/512K/1M tokens), and multilingual benchmarks (MMLU-ProX, WMT24++), in both BF16 and NVFP4 precision. Throughput uses a Pareto methodology on a single 8xB200 node: sweep tensor parallelism over {1,2,4,8}, expert parallelism on/off, and batch size, then report the best total throughput at fixed user-throughput thresholds of 100/125/150 tokens/s. Ultra-long-context concurrency is measured on a single H100, and a dedicated ablation isolates iterative versus single-step Puzzle at the same compression target.

검증 General knowledge(MMLU-Pro), 추론·코딩(AIME25, HMMT, GPQA, LiveCodeBench, SciCode, HLE), agentic(Terminal-Bench, SWE-Bench, TauBench V2), instruction following(IFBench, Arena-Hard-V2), 장문맥(AA-LCR, RULER 256K/512K/1M token), 다국어(MMLU-ProX, WMT24++)를 아우르는 벤치마크 스위트를 BF16과 NVFP4 정밀도 모두에서 평가했다. 처리량은 8xB200 노드 1대에서 Pareto 방법론으로 측정한다: tensor parallelism {1,2,4,8}, expert parallelism on/off, batch size를 sweep한 뒤 사용자 처리량 100/125/150 tokens/s 임계값에서 최고 총 처리량을 보고한다. 초장문맥 동시성은 H100 한 장에서 측정했고, 동일 압축 목표에서 iterative vs single-step Puzzle을 가르는 전용 ablation도 수행했다.

Results Total parameters drop from 120.7B to 75.3B (62.4% of the original) and active parameters from 12.8B to 9.3B (73.1%), while single-step server throughput reaches 1.60-2.14x the parent across six operating points (about 2x in the decode-heavy 8K/64K regime), and 1M-token concurrency on one H100 rises from 1 to 8 requests. Accuracy stays close to the parent, e.g. AIME25 89.7 vs. 92.2, RULER-1M 92.2 vs. 93.9, MMLU-Pro 82.4 vs. 83.8 (BF16). In the ablation, Iterative Puzzle beats single-step Puzzle by +0.57 average points over ten benchmarks at the same compression target. Within the recovery pipeline, the first short-context KD stage alone restores most benchmark categories to over 97% of the teacher, while the measured impact of RL was small.

결과 총 파라미터는 1207억에서 753억(원본의 62.4%)으로, 활성 파라미터는 128억에서 93억(73.1%)으로 줄었고, single-step 서버 처리량은 6개 운영점에 걸쳐 부모의 1.60~2.14배(decode-heavy 8K/64K regime에서 약 2배)에 도달하며, H100 한 장의 1M-token 동시 처리는 1개에서 8개 요청으로 늘어난다. 정확도는 부모에 근접한다: AIME25 89.7 vs 92.2, RULER-1M 92.2 vs 93.9, MMLU-Pro 82.4 vs 83.8(BF16). Ablation에서 Iterative Puzzle은 single-step Puzzle 대비 동일 압축 목표 기준 10개 벤치마크 평균 +0.57점의 이득을 보인다. 복구 파이프라인에서는 첫 short-context KD 단계만으로 대부분의 벤치마크 카테고리가 teacher의 97% 이상으로 회복되고, RL의 측정된 효과는 작았다.

Comparison Compared directly against the parent Nemotron-3-Super at matched BF16 and NVFP4 precision, and positioned against the much smaller Nemotron-3-Nano (30B total / 3.5B active) on the accuracy-throughput plane. Continued MTP training raises mean acceptance length from 3.45 to 4.34 in BF16 (3.31 to 4.31 in NVFP4), a roughly 25-30% gain concentrated at deeper draft positions where the teacher-forced/autoregressive mismatch hurts most. Against Nano, raw token throughput favors the smaller model, but after adjusting for Nano's 1.6x generation verbosity, Puzzle-75B-A9B+MTP completes a comparable number of user requests per minute at near-parent accuracy.

비교 부모 모델 Nemotron-3-Super와 BF16·NVFP4 동일 정밀도에서 직접 비교하고, 훨씬 작은 Nemotron-3-Nano(총 30B/활성 3.5B)와는 정확도-처리량 평면에서 포지셔닝한다. 지속 MTP 학습은 평균 수락 길이(acceptance length)를 BF16에서 3.45→4.34(NVFP4에서는 3.31→4.31)로 올리는데, 약 25~30% 개선이며 teacher-forced 학습과 autoregressive 추론의 불일치가 가장 큰 깊은 draft 위치에서 이득이 집중된다. Nano 대비 raw token 처리량은 밀리지만, Nano의 1.6배 생성 verbosity를 보정하면 Puzzle-75B-A9B+MTP는 부모급 정확도로 분당 비슷한 수의 사용자 요청을 완료한다.

Significance It demonstrates that a production-grade hybrid MoE+Mamba model can shed more than a third of its total parameters while preserving most downstream capability, by jointly pruning MoE experts, the active-parameter budget, and the Mamba state rather than any single axis. The evaluation framing is itself a contribution: reporting Pareto-optimal total throughput at fixed per-user throughput answers the actual deployment question ("how much can a server push through under this interactivity constraint?") rather than a kernel microbenchmark. The verbosity-adjusted request-completion comparison also warns that raw tokens/s overstates lighter models — a chattier model must generate more tokens per answer.

의의 실서비스급 hybrid MoE+Mamba 모델이 MoE expert·활성 파라미터 예산·Mamba state를 (단일 축이 아니라) 결합해 pruning하면 총 파라미터의 3분의 1 이상을 덜어내고도 대부분의 다운스트림 능력을 지킬 수 있음을 보여준다. 평가 프레이밍 자체도 기여다: 사용자별 처리량을 고정한 Pareto-optimal 총 처리량 보고는 kernel microbenchmark가 아니라 실제 배포 질문("이 interactivity 제약에서 서버가 얼마나 밀어낼 수 있나")에 답한다. Verbosity를 보정한 요청 완료율 비교는 raw tokens/s가 가벼운 모델을 과대평가한다는 경고이기도 하다 — 말이 많은 모델은 답 하나에 더 많은 token을 생성해야 하기 때문이다.

Limitations The paper states the largest accuracy gaps appear on instruction-following and agentic evaluations — Arena-Hard-V2 (68.6 vs. 72.8 in BF16) and certain TauBench domains — which it describes as more sensitive to aggressive compression. The RL recovery stage had only a small measured effect in the reported experiments. Structurally, the heterogeneous search was applied only to the MoE layers: Mamba layers were pruned uniformly due to inference-framework limitations and attention layers were left untouched, so two of the three layer types never benefit from per-layer allocation (reviewer judgment: likely remaining headroom).

한계 논문은 가장 큰 정확도 격차가 instruction-following·agentic 평가, 구체적으로 Arena-Hard-V2(BF16 68.6 vs 72.8)와 특정 TauBench 도메인에서 나타나며, 이들이 공격적 압축에 더 민감하다고 명시한다. RL 복구 단계의 측정된 효과는 보고된 실험에서 작았다. 구조적으로 heterogeneous 탐색은 MoE layer에만 적용됐다: Mamba layer는 inference framework 제약으로 균일하게 pruning했고 attention layer는 손대지 않았으므로, 세 layer 유형 중 둘은 layer별 배분의 혜택을 받지 못한다(리뷰어 판단: 남은 headroom일 가능성).

Future work The conclusion explicitly plans to apply Iterative Puzzle more aggressively and to expand the pruning techniques toward "extreme LLM compression" — much higher compression rates while preserving strong accuracy. Since the measured RL impact was small, the authors also state they plan to research more RL environments and additional post-training techniques. (Reviewer judgment: extending heterogeneous allocation to the Mamba and attention axes, once framework support lands, is a natural adjacent direction.)

향후 과제 결론에서 명시적으로, Iterative Puzzle을 더 공격적으로 적용하고 pruning 기법을 확장해 정확도를 지키면서 훨씬 높은 압축률의 "extreme LLM compression"에 도달하겠다는 계획을 밝힌다. RL의 측정 효과가 작았던 만큼, 더 많은 RL environment와 추가 post-training 기법 연구도 future work로 명시한다. (리뷰어 판단: framework 지원이 갖춰지면 Mamba·attention 축으로 heterogeneous 배분을 확장하는 것이 자연스러운 인접 방향이다.)

Released resources The models are publicly released on Hugging Face in three checkpoints — BF16, NVFP4 (Blackwell-targeted), and FP8 (Hopper-targeted). The compression stack runs on the Puzzletron implementation inside NVIDIA's open-source Model-Optimizer repository. MTP quality was evaluated with SPEED-Bench, and the parent Nemotron-3-Super checkpoints referenced in the paper are likewise public on Hugging Face.

자원 공개 모델은 Hugging Face에 BF16, NVFP4(Blackwell 대상), FP8(Hopper 대상) 세 체크포인트로 공개되어 있다. 압축 스택은 NVIDIA의 오픈소스 Model-Optimizer 저장소 안 Puzzletron 구현 위에서 동작한다. MTP 품질은 SPEED-Bench로 평가했으며, 논문이 참조하는 부모 Nemotron-3-Super 체크포인트들도 Hugging Face에 공개되어 있다.