NVIDIAはLLama Nemotronシリーズモデルを発表しました。これは推論モードと通常のチャットモードの間で動的に切り替えることができ、さまざまなタスク要件に適応します。
Llama-Nemotronの主要技術解説
ニューラルアーキテクチャ探索（NAS）：
- ブロックごとの局所蒸留：
Llama 3命令モデルを起点として、Puzzleフレームワークは各代替サブブロックを独立して並列に訓練し、親ブロックの機能を維持しつつ、レイテンシの削減、メモリ使用量の低減、スループットの向上などの計算性能を高めます。
例えば、LN-Ultraモデルの場合、Llama 3.1-405B-Instructモデルから始め、計算性能を向上させつつ親ブロックの機能を維持するために各代替サブブロックを訓練します。
訓練プロセスでは、各代替ブロックは親ブロックの機能に近似できるよう訓練され、同時に計算性能の向上が実現されます。
例えば、一部のブロックはアテンションメカニズムを削除して計算とKVキャッシュメモリの消費を削減する一方で、他のブロックはフィードフォワードネットワーク（FFN）の中間サイズを調整することで異なる程度の圧縮を実現する場合があります。
- 混合整数計画法（MIP）：
代替ブロックのライブラリを構築した後、Puzzleフレームワークは混合整数計画ソルバーを利用して、与えられた制約に基づいて各レイヤーの最適なブロックを選択します。
例えば、LN-Superモデルの場合、制約には単一のNVIDIA H100 GPUで少なくとも5倍のスループット向上を実現し、FP8精度で約300Kのキャッシュトークンをサポートすることが含まれます。
MIPソルバーを使用して、与えられた制約（ハードウェア互換性、最大許容レイテンシ、総メモリ予算、または期待される推論スループットなど）に基づいてブロックライブラリから各レイヤーの最適なブロックを選択し、完全なモデルを構築します。
MIPソルバーは目的関数を最適化することで、制約条件を満たしながら、ブロックライブラリから最適なブロックの組み合わせを選択し、完全なモデルを構築します。
例えば、LN-Ultraモデルの場合、最終モデルは8つのH100 GPUで少なくとも1.5倍のレイテンシ削減を達成し、FP8精度で最大3Mのキャッシュトークンをサポートしました。
- FFN融合：
LN-Ultraモデルでは、FFN融合技術が導入されています。Puzzleが一部のアテンションレイヤーを削除した後、モデルには連続したFFNブロックが頻繁に現れます。
例えば、モデルに2つの連続したFFNブロックがある場合、FFN融合技術はそれらをより広範なFFNレイヤーに置き換えます。これは並列実行可能であり、シーケンシャルなステップを減らし、計算利用率を向上させます。
FFN融合により、LN-Ultraモデルは推論レイテンシにおいて著しい改善を達成し、最終的に1.71倍のレイテンシ向上が実現されました。
知識蒸留と継続的事前学習：
- 知識蒸留：
LN-Superモデルは、Distillation Mixデータセットで40Bトークンの知識蒸留訓練を行います。
例えば、LN-Superモデルの出力を教師モデルの出力と比較することで、LN-Superモデルのパラメータを調整し、教師モデルの振る舞いをよりよく近似できるようにします。
LN-UltraモデルはまずDistillation Mixデータセットで65Bトークンの知識蒸留訓練を行い、その後Nemotron-Hステージ4事前学習データセットで継続的な事前学習を88Bトークン行います。
例えば、知識蒸留フェーズでは、LN-Ultraモデルは教師モデルの出力を学習することで、段階的に自身の性能を向上させます。
継続的な事前学習フェーズでは、モデルは知識範囲をさらに拡大し、最終的に主要なベンチマークで参照モデルであるLlama 3.1-405B-Instructを上回りました。
- 継続的事前学習：
LN-Ultraは知識蒸留後、Nemotron-Hステージ4事前学習データセットで継続的に事前学習を行い、さらに性能を向上させます。
例えば、継続的な事前学習フェーズで、LN-Ultraモデルは大量のラベルなしデータを学習することで、その語彙と言語パターンを拡張し、推論タスクでより優れたパフォーマンスを発揮します。
教師ありファインチューニング（SFT）：
- データ準備：
推論データと非推論データの両方を含む混合データセットを構築します。
例えば、推論データでは、各プロンプトに「detailed thinking on」という指示が含まれており、モデルは詳細な推論プロセスを出力する必要があります。
非推論データでは、各プロンプトに「detailed thinking off」という指示が含まれており、モデルは簡潔な応答を出力する必要があります。
推論データについては、さらに数学、コーディング、科学、および一般的な領域に細分化されます。
例えば、数学領域では、Art of Problem Solving（AoPS）コミュニティフォーラムから数学の問題を収集し、DeepSeek-R1やQwen2.5-Math-7BInstructなどのモデルを使用して推論と非推論の解法を生成し、その後フィルタリングと検証ステップを経て、データの品質と正確性を確保します。
- 訓練プロセス：
初期段階は推論データの訓練に焦点を当てます
中間段階で非推論データを導入します
最終段階はチャット、指示追従、ツール呼び出しの混合データに焦点を当てます。
すべてのモデルは、指示調整データに基づくトークンごとの交差エントロピー損失を使用して訓練されます。
例えば、訓練プロセスでは、モデルの出力が目標出力と比較され、交差エントロピー損失を計算することでモデルのパラメータが調整されます。
モデルの規模と要件に応じて、異なる学習率、シーケンス長、訓練エポックが使用されます。
例えば、LN-Nanoモデルは3段階のSFTプロセスを採用しています。
LN-Superモデルは全データセットで1エポック訓練されます。
LN-Ultraモデルは、線形ウォームアップやコサイン減衰学習率調整を含むより複雑な訓練戦略を採用し、訓練の安定性と収束性を確保します。
大規模強化学習：
- 訓練アルゴリズム：
LN-Ultraでは、科学的推論能力を向上させるために、Group Relative Policy Optimization（GRPO）アルゴリズムを使用した強化学習を採用しています。
- データ処理：
応答を個別に生成し、合格率を計算することで、合格率の低いプロンプトをフィルタリングし、訓練データの難易度を高めます。
同時に、事前計算された合格率を難易度指標として使用するカリキュラム訓練戦略を採用し、各バッチの難易度分布を動的に調整することで、モデルが簡単なタスクから複雑なタスクへと段階的に学習できるようにします。
- 報酬メカニズム：
正確性報酬とフォーマット報酬を使用してモデルの学習を誘導します。正確性報酬は、モデルが生成した応答が正解と一致するかどうかを判断して与えられ、モデルが正確な質問応答ができるようにします。フォーマット報酬は、モデルが推論モードで正しく思考プロセスを出力しているか、非推論モードで思考タグの出力を避けているかを確認することで、モデルが正しい出力フォーマットに従うように誘導します。
- 推論モード切り替え
軽量なシステムプロンプト「detailed thinking on/off」を介して、推論モードと通常のチャットモード間の動的切り替えを実現します。
https://arxiv.org/pdf/2505.00949
https://huggingface.co/collections/nvidia/llama-nemotron-67d92346030a2691293f200b