Nemotron Labs 3 Puzzle 75B A9Bの紹介：2.03倍のサーバースループットを実現する圧縮ハイブリッドMoE LLM

大規模ハイブリッドMoEモデル（Nemotron-3-Superなど）は高精度ですが、サービス提供にはコストがかかります。アクティブパラメータ、KVキャッシュ、Mamba状態によって、ノードが特定のユーザーあたりトークンレートで収容できるユーザー数が制限されます。NVIDIA AIチームは、Nemotron-3-Superの圧縮版であるNemotron-Labs-3-Puzzle-75B-A9Bをリリースしました。元モデルの総パラメータは120.7B、アクティブパラメータは12.8B。圧縮後はそれぞれ75.3B、9.3Bです。

デプロイメント目標はアーキテクチャ探索開始前に固定されていました。目標1：ユーザーあたり毎秒100トークンでサーバースループット2倍。目標2：単一H100上で1Mトークンリクエストを8つ同時処理。Hugging FaceではBF16、FP8、NVFP4の3つのチェックポイントが提供されています。

クイックサマリー：120.7B/12.8Bアクティブパラメータが75.3B/9.3Bアクティブに圧縮され、88ブロックのハイブリッドレイアウトを維持。8xB200の総スループットは、NVFP4とユーザースループットを一致させた条件でSuper比1.60倍～2.14倍。単一H100での1Mトークン同時実行数は、重量が70GBから44.5GBに減少したことにより1から8に。反復型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ブロック（Mamba 40、MoE 40、Attention 8）を保持。変更点はブロック内部の容量です：総パラメータ62.4%、アクティブパラメータ73.1%、Mamba SSM状態サイズ75%、MoEルーテッドエキスパート中間サイズ平均59.9%、トークンあたりのアクティブルーテッドエキスパート数平均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ウェイトは80GB HBMの約70GBを占有し、各1Mトークンリクエストは約4GBのKVキャッシュを追加するため、実効同時実行数は1。Puzzle-75B-A9BのNVFP4ウェイトは約44.5GB、アテンションレイアウトは不変のためリクエストあたりKVコストは不変、1Mコンテキストでの同時実行数は8に向上。この同時実行数での総デコードスループットはSuperの単一リクエストスループットの約4倍。990Kトークンプロンプトのプリフィルは約1.2倍高速。

反復型Puzzleの仕組み：Puzzleは分解型ニューラルアーキテクチャ探索フレームワークで、ここではPuzzletronとして実装。代替層実装の離散探索空間を定義し、各代替案に品質スコアを割り当て、混合整数計画法でデプロイ制約下で各層に1つの代替案を選択。3つのプルーニング技術が探索空間を構成：中間チャネルプルーニング（各ルーテッドエキスパート内のチャネルをその出力への寄与でランク付け、同一MoE層内の全エキスパートをカーネル互換性のため均一サイズにプルーニング）、Top-k削減（トークンがルーティングされるエキスパート数を層ごとに変化、最大で親モデルのk=22まで）、Mamba SSMプルーニング（SSM状態サイズを128から96チャネルに削減）。SSM結果は測定済み：128チャネルから96チャネルへの削減で、デコード時にSSMカーネルが1.2倍～1.3倍高速化（バッチサイズ8～512）。チャネルはMamba層出力への推定寄与でランク付けされ、6700万トークンの検証データの平均に基づく。付録Aは、積極的プルーニング下でランダムチャネル選択よりも優れていることを示す。

元の定式化では、置換品質への影響は近似的に加法的と仮定し、各候補ブロックを未修正の親モデル内でスコアリング。置換間の高次相互作用を無視。反復型Puzzleは、限定された圧縮と短い知識蒸留回復を交互に行い、目標に直接ジャンプする代わりにシーケンスM0, M1, … MRを構築。スコアは元の親モデルではなく、現在の圧縮モデルに対して再計算。

3段階が使用されました：段階1：MoEウェイトを教師容量の75%、Mamba SSM状態を75%に削減、240億トークンで回復。段階2：MoEウェイトを教師容量の60%に削減、432億トークンで回復。段階3：アクティブルーテッドエキスパート予算を50%に削減（不均一割り当て）、528億トークンで回復。上表は同じ圧縮目標での単一段階Puzzleベースラインとの比較：3段階法は10ベンチマーク平均69.05、単段階は68.48。MMLU-Pro、GPQA、HLE、AA-LCR、LiveCodeBench、SciCode、RULER-256Kで向上、IFBench-Instructionは0.2ポイント低下、IFBench-Promptは0.5低下。

回復：蒸留、強化学習、冗長性。知識蒸留はNemotron-3-Nanoからのプリトレーニングデータ30%とSFTデータ70%を使用。Puzzleフェーズでは32Kシーケンス長でKD、回復フェーズでは128Kでトレーニングし、512Kにスケール。予算は最大100Bトークン、グローバルバッチ16Mトークン、Megatron-LMを使用。RL後トレーニングはNemotron-3-Super RLパイプラインのステージ2を採用し、ソフトウェアエンジニアリングに焦点：フェーズ2.1で単一ステップツール使用比較、フェーズ2.2でエンドツーエンドサンドボックスRL（エージェント最大200ターン）。両フェーズともKLペナルティ0。チームは学習率をスイープし、得られた重みを平均。上図は各ステージの寄与を示す：ショートコンテキストKDはほとんどのカテゴリをNemotron-3-Superの97%超に回復、ロングコンテキストKDは特に長入力・長生成ベンチマークを向上。研究チームはRLの影響は今回の実験では小さかったと述べる。

冗長性は静かな詳細：最後のPuzzle反復後、モデルはSuperのトークン数の132%を生成。完全な回復パイプライン後は99%に低下。

デプロイメント：量子化とマルチトークン予測。2つのトレーニング後量子化レシピ：FP8 W8A8（Hopper向け）、NVFP4 W4A4（Blackwell向け）。コンポーネントは、スパース＆共有MoE GEMM、Mamba GEMM、Mamba SSMキャッシュ、KVキャッシュ、ルーターなどが異なる精度。両レシピとも256個のトレーニング後SFTサンプルでキャリブレーション。NVFP4はSuperで使用されたAutoQuantize感度検索ではなく最大キャリブレーションを使用。その結果、チェックポイントはやや積極的に量子化されているが、パフォーマンスは同程度。NVFP4はHopperではネイティブサポート外だが、ここではHBM容量がボトルネックとなる1MコンテキストH100ターゲットに使用。

Puzzle-75B-A9BはSuperの共有MTPヘッドを継承。パラメータはMTPステップ間で共有されるため、1つのヘッドが推論時に再帰的に適用。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ポイントの差。3つの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で1ノードが2.03倍のトークン、冗長性調整後は1分あたり完了リクエスト数2.16倍。プリフィル負荷の高いドキュメントパイプライン：50K/2K領域は1.60倍のみ、プロンプト処理が計算を支配する場合圧縮の効果は薄い。エージェントSWEループ：2.6ポイントのSWE-Benchギャップをタスクミックスに対して評価、RL回復はこの能力をターゲットにしたが部分的にしか回復できず。

デプロイメントエクスプローラー：論文とモデルウェイトを参照。Twitterでフォローし、150k+ ML SubRedditに参加、ニュースレターを購読。Telegramでも参加可能。GitHub Repo、Hugging Faceページ、製品リリース、ウェビナーなどのプロモーションでパートナーシップが必要な場合はご連絡ください。本記事はMarkTechPostに初出。