ThAME: 3D Memory-Enabled Heterogeneous Accelerator for LLM Mixture of Experts

インターネットが、日々賢くなっていく本が並ぶ巨大な図書館であると想像してみてください。これらの「スマートな本」は大規模言語モデル（LLM）と呼ばれ、物語を書いたり、数学の問題を解いたり、人間のようにチャットしたりすることができます。しかし、これらのモデルが真に賢くなるためには、数十億の事実をメモリに保持するほど巨大である必要があります。最近、科学者たちは「Mixture of Experts（MoE：混合エキスパート）」と呼ばれる巧妙なトリックを発見しました。一つの巨大な脳ですべてを記憶しようとする代わりに、MoEは小さな専門家チームを構築します。質問を受けると、システムはチーム全員を呼び起こすのではなく、その答えを知っている数人のエキスパートだけを呼び出します。これは、司書に本を持ってきてもらうよう頼むようなものです。図書館のスタッフ全員を起こす代わりに、歴史セクションのベルを鳴らすだけなのです。

しかし、問題があります。現実の世界では、これらの専門家を呼び出すことは非常に煩雑です。エキスパートたちはあちこちに散らばっており、司書は正しい本を持ってくるためにあちこちを走り回らなければならず、それから答えを持ってくるために戻ってこなければなりません。この「あちこちを走り回ること」が交通渋滞を引き起こします。コンピュータは、実際に思考する時間よりも、データが到着するのを待っている時間の方が長くなってしまいます。これが「メモリーウォール（メモリの壁）」であり、最も速いスーパーコンピュータでさえも速度を低下させるボトルネックとなります。問いはこうです。どのようにすれば、読者のすぐ隣に本があり、ランナーたちが交通渋滞を避けるための完璧な高速道路システムを備えた図書館を構築できるのでしょうか？

---

問題点：脳内での交通渋滞

この論文は、まさにこの問題を解決するために、ThAME（3D Memory-Enabled Heterogeneous Accelerator：3Dメモリ対応ヘテロジニアス・アクセラレータ）と呼ばれる新しいアーキテクチャを紹介しています。著者らは、現在のコンピュータは、「思考」の部分（CPUまたはGPU）が「メモリ」の部分（データが存在する場所）から遠く離れた、忙しい都市のようなものであると説明しています。大規模言語モデルがMixture of Expertsのトリックを使用すると、混沌とした状況が生じます。

学校の食堂を想像してみてください。生徒たち（トークン）が、食べ物を得るために異なるランチカウンター（エキスパート）へ行く必要があります。通常の食堂では、全員が同じ列に並びます。しかし、MoEシステムでは、食べたいものに応じて、生徒はランダムに異なるカウンターへと送られます。あるカウンターには大行列ができ、別のカウンターには誰も来ないこともあります。生徒たちは食事を得るために食堂を駆け巡り、それから食事を混ぜ合わせるために中央のテーブルへと走って戻ってこなければなりません。これは「スキャッター・ギャザー（分散・収集）」と呼ばれるトラフィックパターンを生み出します。つまり、生徒たちがあらゆる方向に走り回り、衝突や長い待ち時間が発生する混沌とした混雑です。論文は、標準的なコンピュータチップはこの種の予測不能で無秩序なトラフィックには作られていないと主張しています。チップはグリッドロック（交通膠着）状態に陥り、システム全体が停滞してしまうのです。

解決策：特別な高速道路を持つ3D都市

これを解決するために、著者らは、これらのAIモデルのために全く新しいハイテク都市を建設することに相当する「ThAME」を提案しています。彼らは、機能させるために3つの主要なアイデアを用いています。

1. 適切な仕事には適切な道具を（ヘテロジニアス・メモリ）：
   論文は、すべてのメモリが同じではないことを指摘しています。AIの特定の部分は非常に頻繁に書き込まれる必要がありますが（ホワイトボードのように）、他の部分は巨大な本棚から読み取られるだけで済みます。

   • 「ホワイトボード」の部分（アテンション・メカニズム）には、高速で書き換えが容易ですが、容量を多く消費するDRAMを使用します。
   • 「巨大な本棚」の部分（エキスパートの重み）には、極めて高密度で、頻繁に書き換える必要のない3Dメモリの一種であるFeFET-NANDを使用します。
   • これは、参照用書籍が巨大で高密度のタワー（FeFET）に積み上げられている一方で、アクティブなワークステーションには異なる種類の高速なデスクが使われている図書館のようなものです。これらを垂直方向に（3Dとして）積み重ねることで、プロセッサのすぐ隣に膨大な量のデータを配置でき、「ランナー」が移動する距離を短縮できます。

2. 特別な高速道路（NoC）：
   これがこの論文の最大の革新です。メモリが近くにあっても、「ランナー」（データ）は依然として異なるエキスパート・カウンター間を移動する必要があります。著者らは、標準的な道路ネットワーク（単純なグリッドやリング状のもの）は、トラフィックが予測不能な場合には機能しないことに気づきました。

   • 彼らは**階層型Network-on-Chip（NoC）**を設計しました。これは、地元の近隣地域が、ローカルなトラフィックを迅速に処理するための小さなプライベートな道路（クロスバー）を持ち、その後、これらの近隣地域が、異なるエリア間の大きなトラフィックの流れを処理できるスマートなツリー状の高速道路システムによって接続されている都市を想像してください。
   • このシステムは、生徒がどのように分散しているか（どのエキスパートが忙しいか）に関わらず、交通渋滞が最小限になるように、複雑な数学的手法を用いて最適化されました。それは、事故が発生する前に自動的に車を迂回させる交通管制システムのようなものです。

3. スマートな配車係（ディスパッチャー）：
   システムには、スマートな交通警察官のように機能するスケジューラーが含まれています。これは混雑状況を見て、全員がほぼ同時に終了するように、各エキスパート・カウンターに正確に何人の「ランナー」（コンピュータコア）を割り当てるかを決定します。これにより、一つの遅いエキスパートがグループ全体を停滞させることを防ぎます。

研究結果

著者らは単に紙の上でこれを作ったわけではありません。どのように動作するかを確認するために、詳細なシミュレーションを実行しました。彼らはThAMEを、既存の最高水準のシステム（StratumやH3D-Tなど）および標準的なGPUと比較しました。

• 速度： シミュレーションにおいて、ThAMEは驚異的な速さを示しました。テキスト生成において、既存の最高のハードウェアよりも最大15.7倍高速でした。これは、標準的なコンピュータが段落を書くのに10秒かかる場合、ThAMEなら1秒未満で完了できることを意味します。
• エネルギー： 同様のタスクに対して、最大9.8倍少ないエネルギーしか消費せず、非常に効率的でした。これは、これらの巨大なモデルを動かすには通常、大量の電力を消費するため、非常に重要なことです。
• 堅牢性： 論文は、メモリ技術が完璧でない場合（例えば、読み取り速度が予想よりも遅い場合）でも、ThAMEは十分な帯域幅が組み込まれているため、良好に動作することを示しています。

行われなかったこと

この論文が主張していないことも明記しておく必要があります。著者らは、物理的なチップを工場で製造し、実際のサーバーファームでテストしたわけではありません。彼らのすべての結果は、ハードウェアがどのように振る舞うかを模倣した非常に詳細なコンピュータモデルであるサイクル正確なシミュレーションに基づいています。また、彼らはすべてのAIの問題を解決したとも主張していません。彼らは、現在のコンピュータにとって最も困難な部分である「デコード（テキストを逐次生成する工程）」フェーズに特化して取り組んでいます。彼らは、「プリフィル（初期プロンプトの読み込み）」フェーズは標準的な強力なプロセッサ（TPU）によって処理されることを認めており、彼らのイノベーションは、その後に続く、メモリ負荷の高い混沌とした部分に特化したものです。

結論

この論文は、異なる種類のメモリを組み合わせ、チップ内にトラフィックに賢いカスタム高速道路を構築することで、ようやくMixture of Expertsモデルを高速かつ効率的に実行できることを示唆しています。これは現在はシミュレーション段階ですが、その結果は、この「スマートな高速道路を持つ3D都市」というアプローチが、電力網を使い果たすことなく、次世代の超知能AIを解き放つ鍵となる可能性を示唆するほど有望なものです。