🤖 machine learning

DMuon: Efficient Distributed Muon Training with Near-Adam Overhead

本論文は、既存の学習パイプラインにシームレスに統合され、オプティマイザーステップのレイテンシを劇的に削減しつつAdamWに近い効率性を実現することで、大規模な基盤モデルにおける行列直交化に基づく最適化のスケール可能な利用を可能にする、Muonオプティマイザのオープンソースの分散実装であるDMuonを導入するものである。

原著者： Vincent Chen, Starrick Liu, Regis Cheng, Dance Yang, Shalfun Li, Ryan Yu, Lucy Liang, Hang Su, Roy Gan, Hao Wang, Qian Wang

公開日 2026-06-26

📖 1 分で読めます☕ さくっと読める

CC BY 4.0

原著者： Vincent Chen, Starrick Liu, Regis Cheng, Dance Yang, Shalfun Li, Ryan Yu, Lucy Liang, Hang Su, Roy Gan, Hao Wang, Qian Wang

原論文は CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/) でライセンスされています。 ✨ これは以下の論文のAI生成解説です。著者が執筆または承認したものではありません。技術的な正確性については原論文を参照してください。免責事項の全文を読む

巨大なロボットの脳（大規模言語モデルや、身体性AI）を訓練している場面を想像してください。それを教えるためには、「オプティマイザー（最適化アルゴリズム）」、つまりロボットのミスを見て、その内部の接続（重み）を調整し、より賢くするための「コーチ」が必要です。

長年、標準的なコーチは AdamW でした。これは、何千人もの小さな作業員が、一度に一つの小さなネジを直していくようなものです。高速で効率的ですが、それぞれのネジを独立して扱います。

最近、Muon という、よりスマートな新しいコーチが発見されました。Muonは、ネジを一つずつ直す代わりに、ロボットの「パネル全体（データ行列）」を見て、それらをまとめて調整します。この「グループセラピー」的なアプローチにより、ロボットはより速く学習し、より高いパフォーマンスに到達できます。しかし、一つ問題があります。Muonは、分散システム（多くのコンピュータが連携して働くクラスター）上で実行すると、非常に動作が遅くなるのです。

問題：「全員参加」のミーティングによるボトルネック
分散型トレーニングの設定では、ロボットの脳は多くのコンピュータ（GPU）に分割されています。

AdamW は、リモートチームのように機能します。コンピュータAは自分のネジを直し、コンピュータBは自分のネジを直します。彼らは互いに多くを話す必要はありません。
Muon は、委員会のように機能します。パネルを修正するために、すべてのコンピュータがまずパネルの「全体」を見る必要があります。

素朴な実装におけるMuonでは、すべてのコンピュータが一旦作業を停止し、全員から全データをダウンロードし、複雑な計算を行い、その結果を送り戻さなければなりません。これは、膨大な1,000ページのレポートを全員が読み終えるまで、決定を下すことができない大規模な会議を開催するようなものです。これでは時間がかかりすぎ、ミーティング（オプティマイザー・ステップ）が実際の作業（学習ステップ）よりも長くなってしまいます。実際、論文では、これが学習自体の2倍以上の時間をコストとして支払わせると指摘されています！

解決策：DMuon (Distributed Muon)
X Square Robot Team は、MuonをAdamW並みに高速に動作させるための新しいシステム、DMuon を構築しました。彼らは、コーチの数学的な仕組みを変えるのではなく、「チームの働き方」を変えることでこれを解決しました。

彼らが以下の3つのトリックを用いて、このボトルネックをどのように解決したのかを説明します。

1. 「指定されたエキスパート」戦略（オーナー中心の実行）

全員がすべてのパネルを直そうとするのではなく、DMuonは各パネルに対して特定のコンピュータを「オーナー（所有者）」として割り当てます。

以前： 全員がデータを集め、全員が計算を行い、全員が時間を無駄にしていた。
現在： 「オーナー」となるコンピュータだけがデータを集め、重い計算を行います。他のコンピュータは、待機するか、あるいは別の作業を行います。
たとえ： 建設現場を想像してください。すべての作業員が家全体を建てようとするのではなく、屋根の専門家、配管の専門家というように役割が決まっています。他の作業員は屋根を作ろうとはせず、専門家に必要な道具を手渡すだけです。これにより、同じ作業を二度繰り返す無駄を防いでいます。

2. 「コンベアベルト」システム（通信のオーバーラップ）

「オーナー」が計算を行っている間、他のコンピュータはただ座って待っているわけではありません。DMuonはパイプラインを作成します。

フォワードパス（順伝播）： ロボットが「考えている（データを処理している）」間に、システムはバックグラウンドで次の道具のセットをエキスパートへ静かに送ります。
バックワードパス（逆伝播）： ロボットが「間違いから学んでいる」間に、システムはすでにエキスパートのための次のデータ収集を開始しています。
たとえ： レストランの厨房のようなものです。シェフ（オーナー）が野菜を切っています。シェイトが野菜を切り終えるのを待ってからウェイターが次の注文を持ってくるのではなく、シェフがまだ切っている間に、ウェイターは次の注文を持ってきます。待ち時間は「切る作業」の中に隠されます。

3. 「スマート組み立てライン」（バッチ処理とチューニング）

Muonが行う計算は複雑です。時には、パーツが巨大な壁のように大きいこともあれば、小さなタイルくらいに小さいこともあります。

問題： 小さなタイルを巨大な機械で処理しようとすると、機械が遊んでしまいます。逆に、巨大な壁を処理しようとすると、非常に時間がかかります。
解決策： DMuonは、多くの小さなタイルを一つの「バッチ」としてまとめ、機械が常に稼働するようにします。また、特定の形状のタイルに対して最も速い処理方法を自動的に見つける「チューナー」も使用しています。
たとえ： 封筒をたった一通送るために配送トラックを出すのではなく、DMuonはトラックがいっぱいになるまで待ち、まとめて発送します。また、目的地が都市なのか農村なのかに応じて、トラックのルートを変更します。

結果

論文では、ロボットの訓練モデル（Wall-OSS）や言語モデル（Qwen2.5）を含む、現実世界のモデルを用いてDMuonをテストしました。

スピード： DMuonは、従来の素朴な方法と比較して、「オプティマイザー・ステップ（ミーティング）」を6倍から163倍高速化しました。
効率性： 1ステップの訓練にかかる総時間は、標準的なAdamWの手法よりもわずか2%遅いだけです。
結論： DMuonを使用すれば、莫大な時間のペナルティを支払うことなく、Muonの優れた「グループセラピー」的な学習スタイルを利用できます。これにより、複雑で重苦しいプロセスをスムーズで効率的なものへと変え、巨大なAIモデルでの実用的な運用を可能にしました。

要約すると、DMuon は、最も賢いコーチ（Muon）が、コンピュータの高速道路で渋滞を起こすことなく走れるようにするための「交通管制システム」であり、次世代のAIを訓練することを初めて現実的なものにしたのです。

技術要約: DMuon – AdamWに近いオーバーヘッドで効率的な分散型Muonトレーニングを実現

1. 問題提起

行列を意識したオプティマイザ、特に Muon は、重み行列のモーメンタム集約された勾配に対してニュートン・シュルツ（NS）反復を適用することで、AdamWのような要素ごとのオプティマイザと比較して優れた収束特性を示すことが実証されています。このアプローチは、更新の特異値を一様なスペクトルへと導き、大規模モデル（LLMやエンボディド・ファウンデーション・モデルなど）の事前学習における計算効率の大幅な向上をもたらします。

しかし、現代の分散トレーニング環境にMuonをデプロイする場合、重大なシステム上のボトルネックに直面します。現代の分散トレーニングスタック（ZeRO、FSDPなど）は、更新がパラメータシャードごとに独立して計算されるという要素ごとの最適化の暗黙的な契約の下で動作しています。Muonはこの契約に違反します。なぜなら、その核となるNS反復には、必要な行列積（具体的には $XX^\top$ ）を計算するために、単一のデバイス上に完全な重み行列が存在する必要があるからです。

素朴な分散実装（「Gather-then-Compute」方式）では、すべてのランクが以下の手順を踏む必要があります：

全てのシャードから、削減されたフル勾配行列をすべてオールギャザー（All-gather）する。
ローカルでNS反復を実行する。
自身が所有していないシャードに対する結果を破棄する。

このアプローチは、主に2つのコストを発生させます：

通信オーバーヘッド: 全行列を再構成することは、モデルサイズと分散幅に比例するコレクティブ通信を導入します。
冗長な計算: すべてのランクが同じ高価な直交化を実行するため、計算コストがランク数（ $D$ ）倍に増大します。

実際、これらのコストにより、分散環境におけるMuonのオプティマイザーステップ時間は、フォワードおよびバックワードパスの合計ウォールクロック時間を上回るか、それに匹敵するようになり、Muonのアルゴリズム上の利点を損ないます。このペナルティは、標準的なLLMの事前学習と比較して、フォワード・バックワードパスが相対的に軽いエンボディド・モデルの学習（例：VLA）において特に深刻です。

2. メソドロジー: DMuon システム設計

DMuon は、MuonとAdamWの間の性能差を埋めるために設計された、オープンソースの分散実装です。これは、基礎となるトレーニングフレームワーク（PyTorch/FSDP2など）を変更することなく、通信の最適化、ロードバランシング、およびカーネルレベルの効率性を統合した、協調設計されたシステムアーキテクチャを通じて実現されます。

2.1 オーナー中心の実行戦略

すべてのランクで計算を複製する代わりに、DMuonは各行列パラメータを単一のオーナーランクに割り当てます。

オーナーの役割: オーナーは権威あるオプティマイザ状態を保持し、行列ごとに正確に一度だけMuon更新（ニュートン・シュルツ反復）を計算します。
非オーナーの役割: 他のランクは、フォワード/バックワード計算に必要な場合にのみパラメータを実体化し、直交化は行いません。
通信フロー:
- フォワード: オーナーは、レイヤー実行とパイプライン化された形で、更新されたパラメータを非オーナーに非同期的に公開します。
- バックワード: 勾配はパックされたバッファに蓄積され、オーナーへとリデュースされます。
- オーバーラップ: DMuonは、通信と計算をオーバーラップさせるために、細粒度のウェイトレイアウトを用いた2段階の階層構造（イントラノードおよびインターノード）を採用しており、コレクティブの競合を防ぎます。

2.2 効率的なグラム・ニュートン・シュルツ反復

オーナー側での更新コストを削減するために、DMuonはグラム・ニュートン・シュルツ定式化を採用しています。

グラム空間における漸化式: 元の $m \times n$ 空間で操作するのではなく、反復を $m \times m$ のグラム空間（ $G = XX^\top$ ）で行います。これにより、 $m < n$ の場合、計算量は $O(m^2n)$ から $O(m^3)$ へと減少します。
対称性の活用: $G$ は対称行列であるため、DMuonはSYRK形式の実行パスを採用し、支配的な行列積の計算量のみを半分に抑えます。
バッチ処理: 小規模な行列によるGPU占有率の問題に対処するため、DMuonは類似した形状の行列をバッチ化し、グラムNS漸化式を共同で実行することで、カーネルディスパッチのオーバーヘッドを償却します。

2.3 計算量に基づいたロードバランシング

行列の形状によってMuonの更新コストは大きく異なるため、素朴な割り当て（例：ラウンドロビン）はストラグラー（遅延ランク）を生じさせます。DMu를は、オーナーの割り当てを**混合整数線形計画法（MILP）**として定式化しています。

プロファイリング: 初期化中に、システムは各異なる行列形状に対する様々なバッチサイズとカーネルスケジュールの実行時間をベンチマークします。
最適化: MILPソルバーは、予測されるメイクスパン（最大負荷）を最小化するように機能し、最も重い行列が負荷を分散するように配置します。
フォールバック: 極めて大きな探索空間に対しては、初期化時間を制限するために貪欲法（Greedy heuristic）が使用されます。

2.4 形状適応型カーネル・オートチューニング

DMuonは、高パフォーマンスなカーネルを生成するために、DSL駆動のオートチューニングワークフロー（TileLangおよびCUTE DSLを使用）を利用しています。

ターゲットハードウェアおよび行列形状に特化したタイル形状、ブロックサイズ、およびソフトウェアパイプライン構成を探索します。
最適な構成はキャッシュされ、再利用されるため、トランスフォーマーの学習中に遭遇する特定の次元問題に対して、カーネルスケジュールが常に最適に近い状態を維持できます。

3. 主な貢献

細粒度な通信最適化: 階層的スケジューリングと細粒度のウェイトレイアウトを通じて通信オーバーヘッドを最小化し、正確なオプティマイザのセマンティクスを維持する、オーナー中心のランタイム。
形状適応型実行スタック: バッチ化されたグラム・ニュートン・シュルツ、対称性を考慮したカーネル、およびオートチューニングを組み合わせ、不均一な行列ワークロードを効率的に処理する高性能エンジン。
計算量に基づいたロードバランシング: 理論的なFLOP数ではなく、実際の実行時プロファイリングに基づいて行列オーナーを割り当てる、計測コストベースの最適化手法（MILP）。
ドロップイン統合: フレームワークレベルの変更をほとんど必要とせず、既存のトレーニングパイプライン（FSDP2など）に容易に統合できるモジュール。
プロダクション規模の検証: ロボティクス（エンボディド）および言語モデルの学習ワークロードの両方における、成功したデプロイメント。

4. 実験結果

著者らは、8から256個のA800 GPUにスケールする4つのプロダクションワークロード（Wall-OSS-0.5、Pi0、Wall-WM、Qwen2.5-7B）においてDMuonを評価しました。

オプティマイザーステップ時間: DMuonは、素朴な「Gather-then-Compute」方式の分散型Muonと比較して、オプティマイザーステップ時間において6.85倍から163.00倍の高速化を達成しました。
エンドツーエンドのステップ時間: DMuonは、総ステップ時間において1.48倍から3.01倍の高速化を実現しました。
AdamWに対するオーバーヘッド: DMuonのステップあたりのレイテンシはAdamWに近いレベルまで抑えられており、評価された全ワークロードにおいて平均**+2%**のオーバーヘッドのみとなっています。
スケーリング: Wall-OSSにおいて、DMuonはGPU数が増加してもAdamWに近いスループットを維持しますが、素朴なMuonは深刻な通信ボトルネックに苦しむことになります。

高速化の内訳:

対称グラムカーネル: 高速化の48%（支配的な積における算術演算を半減）。
オーナー・スケジューリングとロードバランシング: 32%（ $D$ 倍の冗長な計算を排除）。
オートチューニングとバッチ処理: 16%（小規模行列のカーネル起動オーバーヘッドを削減）。

5. 意義と主張

本論文は、DMuonが、行列を意識したオプティマイザと分散トレーニング・インフラストラクチャとの間のシステム的なミスマッチを解決したと主張しています。冗長な計算を排除し、通信を最適化することで、DMuonはMuonの収束の利点（高い計算効率や優れた最適化幾何学など）を、実用的なウォールクロック時間の高速化へと変換することを可能にします。

著者らは、DMuonを新しいアルゴリズムとしてではなく、Muonを大規模なプロダクション・トレーニングにおいて実行可能な、AdamWのドロップイン・リプレイスメントとするための、不可欠なシステム・イネーブラーとして位置付けています。この成果は、「行列を意識した最適化のペナルティ」を無視できるレベルまで低減できることを示しており、大規模言語モデル（LLM）とエンボディド・ファウンデーションモデル（VLA）の両方における効率的なスケーリングを可能にします。

なお、DMuonは更新ルールや収束特性を変更するものではなく、厳密にはシステムのオーバーヘッドのみに対処するものであることを明記しています。したがって、その恩恵は分散設定において最も顕著であり、シングルGPUトレーニングにおいては、排除すべき分散された冗長性が存在しないため、より緩やかな改善（オプティマイザ速度の約2倍の向上）に留まります。