🤔 なぜ今MoEなのか?

LLMの世界はここ数年、「Dense(密)モデル」のスケーリング則に大きく依存してきました。パラメータ数とデータを増やせば性能が上がるというシンプルな方程式です。しかし、GPT-4、Llama 3、DeepSeek V3/R1などのモデルが数百億〜数千億パラメータに達した今、Denseモデルの限界が明確になっています。

  • 学習コストの高騰: 全パラメータを常に活性化するため、FLOPあたりの効率が悪い
  • 推論レイテンシの増大: モデルサイズが大きくなるほど、トークンあたりの処理時間が線形に増加
  • デプロイの困難: 単一GPUにモデルを載せることすら不可能に

Mixture of Experts (MoE) は、この問題に対する現実的な解決策です。すべてのパラメータを常に使う代わりに、入力トークンごとに必要なエキスパート(Expert)サブネットワークだけを選択的に活性化します。結果として、総パラメータ数は数百億でも、推論時に活性化されるパラメータは数十億に抑えられます。

例えば gpt-oss-20b モデルは総パラメータ210億ですが、トークンあたり4つのExpertだけを活性化し、実質的には 約36億の活性パラメータ で動作します。M3 Ultra Mac(メモリ帯域幅 ~800GB/s)でbfloat16の場合の推論速度を計算してみましょう:

推論速度 ≈ 800 / (3.6 × 2) ≈ 111 tokens/s

実測値は約115 tok/sで、理論値とほぼ一致します。つまり、210億パラメータモデルの性能を36億パラメータレベルの速度で使える ということです。

参考資料: 本記事の内容は、Hugging Face公式ブログ を基に再構成しています。

Diagram showing Mixture of Experts architecture with router and multiple expert sub-networks in a Transformer layer Algorithm Concept Visual

🔧 Transformers v5:MoEのための3大核⼼変更

Hugging Face transformers ライブラリは元々Denseモデルを前提に設計されていました。MoEモデルが急増(DeepSeek R1以降、爆発的に増加)したことで、ライブラリ内部の前提を完全に見直す必要が生じました。v5で導入された3つの核心的な変更を見ていきます。

1. Weight Loading Refactor:チェックポイントとランタイムの不一致を解消

MoEのチェックポイントを開くと、Expertが個別に保存されています。

// DeepSeek-V3 チェックポイントインデックス例
{
  "model.layers.3.mlp.experts.0.gate_proj.weight": ...,
  "model.layers.3.mlp.experts.1.gate_proj.weight": ...,
  // ... 全256個のExpert
  "model.layers.3.mlp.experts.255.gate_proj.weight": ...
}

しかしGPUで効率的に演算するには(Grouped GEMM、Fused MoEカーネル)、すべてのExpertの重みを1つの連続したテンソルにパッキングする必要があります。つまり、チェックポイントは256個の分離したテンソルですが、ランタイムは1個のパックされたテンソルを欲しているわけです。

この問題を解決するために導入されたのが WeightConverter 抽象化です。核心的な考え方は:

チェックポイントは単なるテンソルのシリアライズされたソースに過ぎない。ローディングとは、これらのテンソルを望みのランタイムレイアウトに変換するパイプラインである。

from transformers.modeling_utils import WeightConverter, MergeModulelist, Concatenate

# 例:Expertの重みを1つのテンソルにまとめる
converter = WeightConverter(
    source_patterns=[
        "block_sparse_moe.experts.*.w1.weight",
        "block_sparse_moe.experts.*.w3.weight",
    ],
    target_key="mlp.experts.gate_up_proj",
    operations=[
        MergeModulelist(dim=0),  # Expertリストを1つのテンソルに
        Concatenate(dim=1),      # gateとup projectionを結合
    ],
)

この変換は 遅延(lazy)方式 で実行されます。ローダーはチェックポイントのキーを一度スキャンし、コンバータのパターンとマッチするキーをグループ化し、各コンバータの依存関係が準備でき次第、スレッドプールで非同期に実行します。例えば MergeModulelist は、該当レイヤーの全Expertがロードされるまで待ってから一度に結合します。

ベンチマーク結果(Qwen/Qwen1.5-110B-Chat、A100 80GB 1台):

バージョン戦略ローディングモード時間
v4.57.6device_map="auto"Threadpool66.24s
v4.57.6device_map="auto"Sequential67.29s
v4.57.6TPOOM
v5device_map="auto"Async (default)20.71s
v5device_map="auto"Sync45.3s
v5TPAsync10.1s
v5TPSync19.28s

単に「スレッドを増やした」から速くなったわけではありません。シングルパスルーティング、非同期マテリアライズ、変換認識スケジューリング の組み合わせにより、メモリピークを回避しながらExpertパッキングとプロジェクションフュージョンをロードタイムに処理できるようになりました。

2. Expert Backend:プラグイン可能な実行アーキテクチャ

Expertが1つのテンソルにパックされた後は、どのように効率的にルーティングするか が問題です。MoEモデルは各トークンを異なるExpertに振り分けるため、ディスパッチ → 演算 → 収集・再整列のプロセスが必要です。

transformers v5では、@use_experts_implementation デコレータパターンを導入し、Expertの計算をモデル実装から分離しました。現在3つのバックエンドを提供しています:

  • eager: 選択されたExpertをループで回しながら1つずつ演算。正確性の参照やデバッグ用。
  • batched_mm: torch.bmm APIを使用。選択されたExpertの重みをトークンごとに複製し、単一のバッチGEMMを実行。小バッチ、GPUメモリが潤沢な場合に適している。
  • grouped_mm: torch._grouped_mm APIを使用。トークンをExpert IDでソート・グループ化し、単一のGrouped GEMMを実行。大バッチやメモリ制約のある環境で威力を発揮する。

3. Expert Parallelism:数百億パラメータモデルを複数GPUに分散

MoEモデルは単一GPUに全く収まらないほど巨大になることがあります。Expert Parallelism (EP) はExpertを複数デバイスに分散し、各デバイスが自身に割り当てられたExpertだけをロード・計算し、結果を集約します。

import torch
from transformers import AutoModelForCausalLM
from transformers.distributed.configuration_utils import DistributedConfig

# Expert Parallelism を有効化
distributed_config = DistributedConfig(enable_expert_parallel=True)

model = AutoModelForCausalLM.from_pretrained(
    "openai/gpt-oss-120b",
    dtype="auto",
    distributed_config=distributed_config,
)
# 実行:Expert数がGPU数で割り切れる必要あり
torchrun --nproc-per-node=8 script.py

EPの核心コンポーネント:

  • GroupedGemmParallel: Expert次元(dim=0)に沿って重みを分割。各デバイスは num_experts / num_devices 個のExpertだけをロード。
  • RouterParallel: グローバルExpertインデックスをローカルインデックスに再マッピングし、現在のランクに割り当てられていないExpertはマスキング。all-reduceで部分出力を結合。

Comparison chart of dense vs sparse model parameter count and inference speed with GPU memory bandwidth Programming Illustration

⚠️ MoEを使う際の注意点

MoEは万能薬ではありません。実際のプロジェクトに導入する前に、以下の点を必ず確認してください。

1. 推論バッチサイズによる性能逆転

batched_mm バックエンドはバッチが小さいときに高速ですが、バッチが大きくなるとメモリ使用量が爆発します。逆に grouped_mm はバッチが大きいときに効率的ですが、ソートのオーバーヘッドがあります。サービスの想定同時リクエスト数に応じてバックエンドを選択 する必要があります。

2. Expert間の不均衡(Load Balancing)

ルーターが特定のExpertだけを偏愛すると、そのExpertの重みだけが更新され、残りは死んだExpertになります。DeepSeek R1はauxiliary lossを使ってこの問題を緩和していますが、完全に解決されたわけではありません。学習時のモニタリングが必須です。

3. 量子化(Quantization)の複雑さ

Expertごとに量子化を適用するには、Expertが予測可能なパッキングレイアウトで存在する必要があります。v5のWeightConverterがこのパイプラインを可能にしましたが、まだすべての量子化手法(AWQ、GPTQなど)がMoEに完全最適化されているわけではありません。

4. 国内クラウド環境での適用コンテキスト

日本のクラウド環境(AWS Japan、GCP Tokyo、さくらインターネット等)では、まだNVIDIA H100/H200クラスタの普及率が限定的です。A100 80GB環境が主流であるため、Expert数が8を超えるモデル(例:DeepSeek-V3の256 Expert)はExpert Parallelismを十分に活用できない可能性があります。まずは8個以下のExpertを持つモデル(Mixtral 8x7Bなど)でPoCを進める ことをお勧めします。

Developer configuring expert parallelism settings in a Python script for distributed MoE model training IT Technology Image

🚀 次のステップ:実際に試してみる

理論は十分にカバーしました。実際にMoEモデルをファインチューニングしてみませんか?

合わせて読みたい記事

MoEアーキテクチャは、今後のLLMエコシステムの標準になる可能性が高いです。特にDeepSeek R1以降、オープンソースコミュニティでのMoE採用が加速しており、transformers ライブラリもそれに合わせて進化しています。単に「モデルが大きい」という理由でMoEを選ぶのではなく、自身のサービス特性(レイテンシ、スループット、コスト)に合ったExpert数と活性化戦略 を実験してみることをお勧めします。

ご質問やご意見があれば、コメントでお知らせください。一緒に議論できれば幸いです!

本コンテンツは、信頼性の高い情報源をもとにAIツールを活用して作成され、編集者によるレビューを経て公開されています。専門家によるアドバイスの代替となるものではありません。