Mozart: Modularized and Efficient MoE Training on 3.5D Wafer-Scale Chiplet Architectures

この論文は、モジュール化された計算と疎性を活用して大規模言語モデルの訓練効率を向上させるため、脳のアナロジーに基づき、3.5D ウエハスケールチップレットアーキテクチャ向けにアルゴリズムとハードウェアを共設計した新しいフレームワーク「Mozart」を提案し、通信オーバーヘッドの削減やリソース利用率の向上を実現したことを述べています。

要約

モーツァルト：AI の「天才チーム」を効率よく動かす新しい仕組み

この論文は、最近話題の超大規模な AI（大規模言語モデル）を、より速く、より安く、より賢く動かすための新しい「設計図」を紹介しています。

タイトルにある**「モーツァルト（Mozart）」**は、この新しい仕組みの名前です。なぜモーツァルトかというと、人間の脳のように「専門家のチーム」を編成し、彼らが最高のパフォーマンスを発揮できるように、アルゴリズム（計算のルール）とハードウェア（物理的なチップ）を完璧に調和させるからです。

以下に、専門用語を使わずに、わかりやすい例え話で解説します。

---

1. 問題点：AI の「天才チーム」には大きな悩みがあった

最近の AI は**「Mixture-of-Experts（MoE）」という仕組みを使っています。これは、1 人の万能な天才ではなく、「何百人もの専門家のチーム」**を雇っているようなものです。

• 仕組み: 質問が来ると、AI は「この質問には A 博士と B 博士が適している」と判断し、その 2 人だけを選んで回答を作ります。
• メリット: 非常に賢いのに、計算コストは抑えられます（全員が動く必要がないから）。
• デメリット（今回の論文が解決したいこと）:
  • 移動の無駄: 専門家がバラバラの部屋（チップ）にいて、チームで話し合うために頻繁に移動（データ通信）する必要があります。
  • 待ち時間: 「誰が動くか」がその都度変わるため、部屋が空いたり、人が集まったりして、計算リソースがムダになります。
  • 混雑: 専門家の間をデータが飛び交うと、道路が渋滞して遅くなります。

2. 解決策：モーツァルトの 3 つのアイデア

モーツァルトは、この問題を解決するために「アルゴリズム（ルール）」と「ハードウェア（建物）」の両面からアプローチしました。

① 専門家同士を「近所」に住まわせる（配置の最適化）

【例え】
もし、料理の専門家と食材の専門家、そして味付けの専門家が、それぞれ遠くの島に住んでいて、毎日船で移動して会議をする必要があったらどうでしょう？とても非効率ですよね。
【モーツァルトのやり方】
過去のデータ（どんな質問が来るか）を分析し、「よく一緒に働く専門家ペア」を見つけ出します。そして、**「よく一緒に働く人同士を、同じ建物（チップ）の隣に住まわせる」**ように配置します。

• 効果: 移動距離が短くなり、通信の渋滞が解消されます。

② 流れ作業で「待ち時間」をゼロにする（細かなスケジュール管理）

【例え】
工場で製品を作る際、部品が届くのを待って機械が止まっている時間はありませんか？モーツァルトは、**「部品（データ）が来るのと同時に、次の作業が始まる」**ような流れ作業を実現します。
【モーツァルトのやり方】
AI が計算している最中に、次の計算に必要なデータを同時に読み込みます。「計算中」と「データ読み込み」を同時に並行して行うことで、待ち時間をなくします。

• 効果: チームが常に動き続け、全体のスピードが劇的に上がります。

③ 超巨大な「3.5D ワーファスケール・チップレット」を作る

【例え】
従来の AI チップは、大きな「一軒家」のようなものでした。しかし、AI が大きくなると、一軒家では部屋が足りなくなります。
モーツァルトは、**「巨大な敷地（ウェーハ）の上に、小さな家（チップレット）を何十棟も並べ、地下鉄（3D 接続）と高速道路（2.5D 接続）でつなげた巨大な都市」**を作りました。

• 特徴:
  • 3D 積み重ね: 計算する部屋と、データを置く倉庫を垂直に重ねて、移動距離を極限まで短くしました。
  • ツリー構造: 都市の中心に「交通整理員（スイッチ）」を置き、専門家のチームがスムーズに集まれるように設計しました。

3. 結果：どれくらい速くなった？

この「モーツァルト」方式を使って、実際の AI モデル（Qwen3 や DeepSeek など）を動かした実験では、従来の方法に比べて約 2 倍（1.9 倍〜2.4 倍）も速く動作することが確認されました。

• 意味: 同じ計算をするのに、半分以下の時間で終わる、あるいは同じ時間でより多くのことを学べるようになります。

まとめ

この論文は、**「AI の天才チーム（MoE）が、物理的な制約（通信やメモリ）に邪魔されずに、最大限の力を発揮できるように、住み心地の良い都市（ハードウェア）と、効率的なルール（アルゴリズム）をセットで提案した」**というものです。

まるで、人間の脳のように、必要な専門家が素早く集まり、無駄な移動なく協働できる環境を作ったことで、次世代の AI をより現実的に、かつ高効率に動かす道を開いたと言えます。

技術概要

論文「Mozart: Modularized and Efficient MoE Training on 3.5D Wafer-Scale Chiplet Architectures」の技術的サマリー

この論文は、大規模言語モデル（LLM）における「Mixture-of-Experts（MoE）」アーキテクチャの学習効率を向上させるため、アルゴリズムとハードウェアを統合的に設計したフレームワーク「Mozart」を提案しています。特に、3.5D ウエハスケールのチップレットアーキテクチャ上で、モジュール化された MoE-LLM の効率的な学習（ポストトレーニング）を実現することを目的としています。

以下に、問題定義、手法、主要な貢献、結果、および意義について詳細を記述します。

1. 背景と問題定義

背景

• MoE の効率性: MoE アーキテクチャは、入力トークンに応じて専門的なサブネットワーク（エキスパート）を動的に選択・活性化することで、大規模なパラメータ数を持ちながら計算コストを抑制します。
• ハードウェアの課題: しかし、MoE の本質的なスパース性（疎性）は、従来の GPU や CPU などのハードウェアプラットフォームにおいて以下の課題を引き起こします。
  • メモリ局所性の欠如: 活性化されるエキスパートが不均一であり、メモリアクセスパターンが予測困難。
  • 通信オーバーヘッド: 並列ユニット間（エキスパート間）での「All-to-All」通信がボトルネックとなり、同期待ちが発生。
  • リソース利用効率の低さ: 動的で不均一な計算負荷により、計算リソースが遊休状態になることが多い。

既存アプローチの限界

• 従来の 2.5D/3D チップレット設計は、主に推論に焦点を当てており、学習時の動的な負荷分散や、MoE の微細なモジュール性を考慮したシステムレベルの最適化が不足しています。
• 既存のワエハスケール設計（例：FRED）は、粗粒度の静的なワークロード分割を採用しており、MoE のような疎なモジュールモデルには不向きで、チップレット間の通信過多やリソースの非効率な利用を招きます。

2. 提案手法：Mozart

Mozart は、アルゴリズム側とハードウェア側を協調設計（Co-design）するフレームワークです。

2.1 アルゴリズム側の最適化

1. エキスパートのクラスタリングと配置（Expert Clustering & Allocation）

   • 動機: トレーニングデータにおける「どのエキスパートが頻繁に同時に活性化されるか（共活性化パターン）」をプロファイリングします。
   • 手法: 共活性化頻度が高いエキスパート同士をクラスタリングし、それらを同じまたは隣接するチップレットに配置します。
   • 効果: これにより、異なるチップレット間での通信（All-to-All）を減らし、パッケージ内での通信を最大化します。また、負荷分散を目的として、クラスタをチップレットグループに均等に割り当てます（整数計画問題として定式化）。

2. 微細なスケジューリング（Fine-grained Scheduling）

   • ストリーミングトークンとエキスパート: 大規模なパラメータを DRAM に保存し、必要に応じてチップレットにロードする際、通信と計算のオーバーラップを最大化します。
   • 手法:
     • エキスパートのストリーミング: 負荷の重いエキスパートクラスタを優先的にロードします。
     • トークンのストリーミング: トークンバッチをマイクロバッチに分割し、DRAM からのアクティベーション保存（バックワードパス）とオンチップ計算を並行して実行します。
   • 効果: DRAM とチップレット間の通信遅延を計算時間と隠蔽（Overlap）し、スループットを向上させます。

2.2 ハードウェア側のアーキテクチャ

1. 3.5D ウエハスケールチップレットアーキテクチャ
   • 3D ロジックオンメモリスタック: 各計算チップレットに論理回路（Logic）と SRAM を垂直に積層（ハイブリッドボンディング）し、中間データ（アクティベーション）への高速アクセスと局所性を確保します。
   • 2.5D NoP-Tree トポロジー:
     • Attention チップレット: 木の中心（Dispatch ノード）に配置され、高い DRAM バンド幅を持ちます。
     • エキスパートチップレット: 木の葉（Leaf ノード）に配置され、計算リソースに特化しています。
     • スイッチ: 木構造のノードに配置され、ネットワーク内での MoE 出力の集約（In-network aggregation）を行い、チップレット間通信を削減します。
   • 階層化メモリ構造: モデル重みはオフチップ DRAM に、頻繁にアクセスされるアクティベーションはローカル SRAM に配置する 2 段階のメモリ階層を採用しています。

3. 主要な貢献

1. アルゴリズム・ハードウェア協調設計フレームワークの提案:
   • MoE-LLM のポストトレーニングを 3.5D チップレット上で効率的に行うための包括的なフレームワーク「Mozart」を提案しました。
   • 専門家配置戦略、通信効率化、微細スケジューリングという 3 つの最適化手法を統合しています。
2. モジュール化された MoE 構造に適応したアーキテクチャ設計:
   • 2.5D NoP-Tree トポロジーと 3D 積層メモリを備えた 3.5D チップレットアーキテクチャを設計し、疎な計算と通信のオーバーラップをハードウェアレベルでサポートします。
3. 実証的な性能向上:
   • 3 つの主要なオープンソース MoE モデル（Qwen3-30B, OLMoE-1B-7B, DeepSeek-MoE-16B）を用いた評価により、ベースラインと比較して最大 2.37 倍の高速化（レイテンシ削減）を達成しました。

4. 実験結果

• 評価対象: Qwen3-30B-A3B, OLMoE-1B-7B-0924, DeepSeek-MoE-16B-Base の 3 モデル。
• 比較対象: 最適化なしのベースライン、および各最適化手法を段階的に適用した変種（Mozart-A, B, C）。
• 結果:
  • 全体性能: Mozart-C（全最適化適用）は、Qwen3 で 1.92 倍、OLMoE で 2.37 倍、DeepSeek-MoE で 2.17 倍の高速化を達成。
  • 通信オーバーヘッドの削減: 提案されたエキスパート配置により、All-to-All 通信データ量が削減され、レイテンシと相関して減少しました。
  • シーケンス長への耐性: 長いシーケンス長（512 トークンなど）においても、通信と計算のオーバーラップにより、ベースラインに比べてレイテンシの増加が抑制されました。
  • メモリ帯域幅の影響: HBM2（高帯域幅メモリ）を使用した場合、最適化の効果が顕著に現れました。一方、SSD などの低速メモリでは、メモリ転送自体がボトルネックとなるため、最適化の恩恵は限定的でした。
• ボトルネック分析: 現在の設計では、計算リソースの並列化は成功していますが、重みのストリーミング（ロード）が逐次的に行われるため、システム全体はメモリバウンドであることが判明しました。

5. 意義と将来展望

• 意義:
  • MoE-LLM の大規模展開における「通信コスト」と「リソース利用効率」という根本的な課題に対し、アルゴリズムとハードウェアの両面から解決策を提示しました。
  • 人間の脳のモジュール性を模倣した設計思想が、次世代の AI ハードウェアにおいて有効であることを実証しました。
  • 3.5D チップレット技術が、単なる推論だけでなく、大規模モデルの学習プロセスにおいても有効活用できる可能性を示唆しています。
• 限界と将来の課題:
  • 現在の設計では Attention モジュールが単一のチップレットに割り当てられており、リソース制約によりサブオプティマルな遅延が発生する可能性があります（データ並列化やテンソル並列化での解決が期待されます）。
  • 高通信負荷下ではスイッチがボトルネックとなる可能性があります。将来的にはスイッチ領域の拡大や帯域幅の向上が検討されるべきです。

総じて、Mozart は、次世代の LLM 学習を可能にするための、モジュール化されたハードウェアアーキテクチャとそれを駆動するアルゴリズムの重要な一歩を示す研究です。