HyperAIHyperAI

Command Palette

Search for a command to run...

NVIDIA Megatron で新興最適化器による LLM 訓練加速を推進

ニューラルネットワークの学習アルゴリズム「Shampoo」の高次最適化手法は長年研究されていますが、最近の超大規模言語モデル(LLM)のトレーニングにおいて特に顕著な成果を上げています。中でも、Kimi K2 や GLM-5 などの主要なオープンソースモデルに採用されている「Muon(Newton-Schulz 法による直交化モーメント)」は、その効率性で注目を集めています。NVIDIA は、この Muon を含む最先端の最適化手法を大規模モデルに実装するための包括的なサポート技術を発表しました。NVIDIA GB300 NVL72 システムを用いた実験では、256 個の GPU で Kimi K2 を、8 個の GPU で Qwen3 30B をトレーニングした結果、従来の AdamW 最適化手法との比較で学習スループットにほぼ差がないことが確認されました。むしろ、ニューロン・シュルツ反復計算の浮動小数点演算(FLOPs)を含めると、Muon を使用した方がモデルの FLOPs 利用率(MFU)が向上したとの報告です。これらの測定は、PyTorch ネイティブのライブラリである NeMo Framework 内の「NeMo Megatron Bridge 26.02」を利用して行われました。大規模な Muon トレーニングの普及には、正則化ステップにおける計算コストの増大、メモリ消費、混合精度学習における数値的不安定性、そして数千の GPU 間での通信ボトルネックといった課題がありました。これを克服するため、NVIDIA は「層ごとの分散最適化」を採用しました。従来の要素ごとの分散方式では全層の勾配を収集できない問題を解決し、各 GPU が全層のパラメータを持つことで、プレコンディショナーの計算を可能にしました。これにより、NeMo Megatron Core に完全に統合され、異なるサイズの層間通信にも対応しています。また、テンソル並列環境での「Newton-Schulz 反復」に対処する技術として、データ全取得モード、計算分散モード、そして通信不要のブロックワイズモードの 3 つの方式を提供しています。さらに、通信を非同期化する技術や、SYRK 演算を活用した浮動小数点演算の削減、並列化する全結合通信の統合など、複数の最適化技術が実装されています。NVIDIA は Muon の他、SOAP などの他の先進最適化手法の研究も支援しており、Megatron Core を通じて誰でも大規模 LLM のトレーニングにこれらの手法を適用できるようになりました。開発者は GitHub リポジトリの例に従うことで、すぐにトレーニングの再現や実装を開始できます。これらの技術革新により、高次最適化手法は、LLM トレーニングの効率性を大きく押し上げる不可欠な要素となりつつあります。

関連リンク