NVIDIA MegatronによるLLM学習高速化へ向けた新世代オプティマイザの高度化

シュリンプ（Shampoo）などの高次最適化アルゴリズム（higher-order optimization algorithms）は、少なくとも過去10年にわたりニューラルネットワークの学習に効果的に適用されてきました。これらの手法は、近年では主要な大規模言語モデル（LLM: Large Language Model）に適用されることで顕著な成功を収めています。特に、ニュートン-シュルツ法による直交化モーメント（Muon: MomentUm Orthogonalized by Newton-Schulz）は、Kimi K2やGLM-5など、現在公開されている最高水準のオープンソースモデルの一部の学習に使用されています。 本記事では、NVIDIAがMuonやその他の最先端の新興オプティマイザ（optimizer）に対して包括的なサポートを提供する方法、および大規模モデルの学習を可能にする技術について解説します。 NVIDIA GB300 NVL72におけるMuonの学習パフォーマンス 表1は、NVIDIA GB300 NVL72システム上でMuonおよびAdamWオプティマイザ（optimizer）を用いたKimi K2モデルとQwen3 30Bモデルの学習スループット（throughput）をまとめたものです。次のセクションで紹介する技術を用いることで、AdamWと比較してMuonオプティマイザ使用時の学習パフォーマンスの低下は極めて小さいことが示されています。ニュートン-シュルツ反復（Newton-Schulz iteration）における行列乗算のFLOPsをカウントした場合、Muonの方がモデルFLOPs利用率（MFU: Model FLOps Utilization）が高くなります。 これらの計測は、NeMo Framework内のPyTorchネイティブライブラリであるNVIDIA NeMo Megatron Bridge 26.02を使用して達成されました。このライブラリは、主要なLLMおよびVLMモデルに対して事前学習（pretraining）、SFT（Supervised Fine-Tuning）、LoRAを提供します。NVIDIAチームは、Kimi K2の学習にPP4DP64EP64構成で256個のNVIDIA GB300 GPUを、Qwen3 30B-A3Bの学習にDP8EP8構成で8個のNVIDIA GB300 GPUを使用しました。 NVIDIA GB300 MXFP8（TFLOPs/s/GPU） モデル AdamW Muon Kimi K2 1,051 1,080（モデル1,029 + Muon 51） Qwen3 30B-A3B 713 721（モデル686 + Muon 35） 表1. NVIDIA Megatron-Bridge 26.02で計測された学習スループット 詳細な実験設定および計測の再現手順については、以下のクイックスタート（Quickstart）セクションを参照してください。 大規模なMuon学習を可能にする技術 Muonなどの高次オプティマイザはLLMの学習効率向上が期待されますが、大規模な実装にはいくつかの課題があります。ニュートン-シュルツ反復や固有値分解を用いた前処理（preconditioning）ステップは、計算コストとメモリ消費を大幅に増加させます。さらに、混合精度学習（mixed-precision training）や勾配蓄積（gradient accumulation）の過程で数値不安定性が発生する可能性があります。同期化された直交化更新を数千個のGPUに効率的に分散配置することは、通信ボトルネックを引き起こし、全体的な効率向上のメリットを損なう可能性があります。 このセクションでは、これらの課題を克服し数千個のGPUにスケールするための技術を紹介しています。選択された技術は、汎用性、スループット、実装の複雑さの間でバランスが取れています。単純なMuonに特化した最適化を優先するのではなく、Muon、SOAP（ShampoO with Adam in the Preconditioner’s eigenbasis）、その他の複雑なオプティマイザに一般化可能な技術を優先しています。 レイヤーごとの分散オプティマイザ（layer-wise distributed optimizer） 従来の要素ごとの分散オプティマイザ（element-wise distributed optimizer）（AdamWなどの要素ごとのオプティマイザに大規模適用されるもの）は、以下の機能を提供します： パーティショニング状態（Partitioning states）：すべてのGPUが全データを保持するのではなく、オプティマイザの状態（optimizer states）は利用可能なすべてのGPUに均等に分割されます。 リデュース・スキャッター勾配（Reduce-scatter gradient）：すべての勾配に対してリデュース・スキャッター演算が行われ、各GPUは自身が「所有する」パラメータに対応する勾配の一部を受け取ります。 ローカル更新（Local updates）：各GPUは自身が「所有する」モデルパラメータの特定の部分のみを更新します。 アールギャザーパラメータ（AllGather parameters）：更新後、GPUは通信を行い、次の順伝播（forward pass）のために全デバイスがモデル全体の更新された重みを持つことを確認します。 より高度なバージョンでは、演算を細分化し、通信と計算処理をオーバーラップ（重畳）させます。 要素ごとの分散オプティマイザ（Element-wise distributed optimizers）はAdamオプティマイザとよく機能してきました。しかし、Muonを含む多くの新興オプティマイザは、その層の重み更新を計算するために層全体の勾配を必要とします。重みとオプティマイザの状態がデータ並列（data parallel, DP）ランク間で均等に分散されている場合、各GPUで利用可能なデータに基づいて更新を計算することはできません。完全な更新を計算するためのデータを収集するために、追加の通信が必要になります。 このタイプのオプティマイザをサポートするために、層ごとの分散オプティマイザ（layer-wise distributed optimizer）が導入されました。異なる層のパラメータは、異なるDPランクに分散されます。各GPUにはプリコンディショナ（preconditioner）を計算できるだけの完全な層分のパラメータが用意されています。 図1. DPランク間における要素ごとのパラメータ分散と層ごとのパラメータ分散の比較。層ごとの分散では、各ランクに完全な層が割り当てられ、全層のプリコンディショニングを可能にします。 層ごとの分散に伴う変更の一つは、可変サイズ通信（variable-size communication）です。完全な層のサイズが異なるため、各GPUはall_gathervを通じて、異なるGPUから異なるサイズの重み更新を収集する必要があります。 この層ごとの分散オプティマイザは、高度な並列処理、混合精度（mixed precision）、最適化されたGPUカーネルを用いて大規模モデルの構築とトレーニングを行うためのオープンソースライブラリであるNVIDIA Megatron Coreに完全に統合されています。実装はlayer_wise_optimizer.pyで確認できます。 分散ニュートン・シュルツ（Distributed Newton-Schulz） 大規模なLLMトレーニングでは、他の並列処理手法と組み合わせてテンソル並列（tensor parallelism, TP）がよく使用されますが、これにより独自の課題が生じます。TPは特定の次元に沿って複数のGPUに個別の重み行列を分割するため、単一のデバイスが完全なパラメータテンソルを保持することはありません。Muonオプティマイザの重要な直交化ステップでは、計算のためにモーメンタムバッファ行列全体へのアクセスが必要です。このモーメンタムバッファがデバイス間でシャードされている場合、シャードされたモーメンタムと重みを処理するために追加の通信が必要になります。 TensorParallelMuonでは、このようなシナリオを処理するために2つの手法が提供されています：複製モード（duplicated mode）と分散モード（distributed mode）です。 複製モード（Duplicated mode） まず、TP（テンソル並列）ドメイン内でモーメンタムが all-gather（全結合）され、各 GPU がニュートン-シュルツ反復（Newton-Schulz iteration）に必要な全データを持つようにします。その後、各 GPU は完全なニュートン-シュルツ反復を実行し、対応する直交化されたモーメンタムによって自身が管理する重みを更新します。このモードは、各 GPU が同じ NS 反復（NS interactions）を実行する代償として、ネットワークレイテンシを最適化します。各更新において、NS 反復の回数（通常は5回）に関わらず、事前に必要となる通信は all-gather が1回だけです。 分散モード 名が示す通り、NS 反復の計算は TP ドメイン内の GPU に分散されます。各 NS 反復では3回の行列乗算（matrix multiplication）が行われ、各反復の最初の行列乗算後に all-reduce（全縮小）が必要です。このモードは、より頻繁な通信を行う代償として、計算効率を最適化します。 さらに、ブロックワイズモード（blockwise mode）もサポートされており、これは各 GPU が所有するモーメンタムと重みのみを用いて直交化（orthogonalization）と更新を行います。他の2つのモードと比較して計算コストが低く、通信も不要です。ただし、これは全モーメンタム行列をまとめて直交化することとは等価ではありません。数学的には概念的にブロック直交化（block orthogonalization）であり、したがってこの名前が付けられており、『Scalable Second Order Optimization for Deep Learning』で最初に紹介されたブロッキング手法に類似しています。 追加の最適化 その他の最適化手法により、トレーニングスループット（training throughput）をさらに向上させることができます。このセクションでは重要なものいくつかを簡潔に解説し、これらは NVIDIA Emerging Optimizers 研究プロジェクトを通じて段階的に導入されていきます。 通信の非同期化（Communication hiding） 現在の NVIDIA Emerging Optimizers プロジェクトのリリースでは、パラメータはオプティマイザステップ（optimizer step）直後に収集され完全に公開されますが、これは特定の状況下でボトルネックとなる可能性があります。要素ごとの分散オプティマイザ（element-wise distributed optimizers）で用いられるのと同じ手法を適用でき、パラメータの収集を次のバッチのフォワードステップ（forward step）に遅延させ、計算と重畳させることができます。 負荷分散（Load balancing） アテンションおよび MLP 重み行列（attention and MLP weight matrices）のサイズが様々であるため、プリコンディショニング（preconditioning）に起因するオプティマイザの計算コストも層によって異なります。理論上でも完全な負荷分散は達成できません。現在、層は同じ DP ドメイン（DP domain）内の GPU にラウンドロビン方式で分散されています。計算コストと通信コストをある程度見積もることで、より良い負荷分散を実現する方法もあります。 SYRK と融合された all-reduce NS 反復の3回の行列乗算のうち最初の2回を SYRK（対称ランクK更新）にマッピングできるため、浮動小数点演算（floating point operations）の約半分を節約できます。対角に沿ったタイルを完全に計算する必要があるため、節約幅は半分よりわずかに少なくなります。現在のリリースでは Triton カーネル（Triton kernels）が提供されています。 SYRKは圧縮形式の上位（または下位）三角行列ではなく、完全な行列を出力します。分散モードでは、all-reduce（全結合）依然として半分ではなく完全な行列を通信する必要があります。より高度なバージョンでは、通信をSYRKカーネルに融合でき、帯域幅の節約だけでなく、より細かな粒度（つまりタイルレベル）で通信を隠蔽することも可能です。CuTe DSLの実装は将来のリリースで予定されています。 研究のためにNVIDIAがサポートする他のオプティマイザは何ですか？ Muonに加え、NVIDIAは研究コミュニティが探索できる多くの他のオプティマイザもサポートしています。以下を含む： 直交化オプティマイザの究極形態であるMOP（偏微分解由による... -> 極分解によるモーメンタム直交化） REKLSにおける、固有値分解とKL補正を用いてステップごとに固有基底を更新する高度なSOAPバリアント Muonの学習結果の再現 このセクションでは、表1のスループット結果を再現する方法と、MuonがNVIDIAソフトウェアスタックにどのように統合されているかを説明します。 クイックスタート NVIDIA-NeMo/Megatron-Bridge GitHubリポジトリのパフォーマンスレシピに記載の手順に従ってください。 例として、256個のNVIDIA GB300 GPUでKimiを実行する場合、Megatron Bridgeリポジトリのルートから開始します： CONTAINER="nvcr.io/nvidia/nemo:26.04"python scripts/performance/setup_experiment.py --account -i ${CONTAINER} --partition -m kimi -mr kimi_k2 --log_dir --num_gpus 256 --gpus_per_node 4 -t "00:15:00" -g gb300 -c fp8_mx -hf 統合 MuonはMegatron Coreに統合されており、Muonクラスとレイヤーごとの分散ラッパーが含まれています。 Muonの場合、`use_distributed_optimizer`は要素ごとの対応物ではなく、レイヤーごとの分散オプティマイザに自動的にディスパッチすることに注意してください。 LLM学習のための新興オプティマイザの開始 Muonのような高次オプティマイザは、LLM学習効率の限界を押し広げる上で不可欠であることが証明されています。レイヤーごとの分散最適化、カスタマイズされた分散ニュートン-シュルツ反復（Newton-Schulz iterations）、そして通信隠蔽（communication hiding）および融合SYRK/all-reduceカーネルに関する継続的な作業を通じて、これらの強力な手法を大規模に効果的にデプロイできるようになりました。Megatron Coreを使用すれば、今すぐこれらの手法の使用を開始できます。 Muonおよびその他の新興オプティマイザはMegatron Coreに統合されています。 MOPやREKLSなどの研究用オプティマイザは、NVIDIA-NeMo/Emerging-Optimizers GitHubリポジトリで実験用に利用可能です。 設定を選択する際は、以下の考慮事項を念頭に置いてください： 分散NS（ニュートン-シュルツ）モード：ネットワークレイテンシが支配的な場合は重複モードを使用します。計算がボトルネックとなる場合は分散モードを使用します。ブロックごとのモードは通信を完全に回避しますが、完全な直交化を近似します。 スケーラビリティ：レイヤーごとの分散オプティマイザは、最小限のオーバーヘッドで大きなDP（データ並列）スケールにおいてMuonを可能にします。表1は、NVIDIA GB300におけるAdamWのスループットとほぼ同等であることを示しています。 今後の改善課題：通信の隠蔽（Communication hiding）、より高度な負荷分散（Load balancing）、そしてSYRK/All-reduceの融合カーネルにより、今後のリリース版で残るスループット（Throughput）の差をさらに縮小していきます。 早速始めたい方は、Megatron Bridgeのパフォーマンスレシピをご覧ください。

NVIDIA MegatronによるLLM学習高速化へ向けた新世代オプティマイザの高度化

背景や根拠まで確認しますか？

調べる

選ぶ

サイト