トランスフォーマーにおけるエキスパートの混合 (MoEs)

Hugging Face は transformers ライブラリを大幅に改良し、Mixture of Experts (MoE) アーキテクチャのファーストクラスサポートを提供します。この再設計は、MoE モデルがチェックポイントに保存される方法とハードウェア上で実行される方法の根本的な不一致に対処し、重みのロードを大幅に高速化し、プラガブルなバックエンドによる効率的な推論、エキスパート並列化によるスケーラブルな分散を実現します。

MoE の基本概念：容量 vs. アクティブパラメータ

Mixture of Experts モデルは、密なフィードフォワード層を「エキスパート」と呼ばれる学習可能なサブネットワークの集合に置き換えます。ルーターは各トークンを処理するためにこれらのエキスパートのうち小さなサブセットを選択し、モデル全体の容量と推論コストを切り離します。

  • Total Parameters: モデル全体の容量と品質を決定します。
  • Active Parameters: 推論速度とトークンあたりの計算コストを決定します。

例えば、gpt-oss-20b モデルは総パラメータ数が 21B ですが、トークンあたり約 3.6B のパラメータ（32 エキスパートのうち 4 つ）しかアクティブにしません。これにより、21B パラメータシステムの品質を保ちつつ、3.6B パラメータモデルと同等の速度で動作し、M3 Ultra Mac 上でおおよそ 115 トークン/秒を達成します。

重みロードのリファクタリングと WeightConverter

従来の transformers における重みロードは、チェックポイントテンソルとランタイムパラメータが 1 対 1 に対応していることを前提としていました。MoE のチェックポイントは通常、エキスパートを個別にシリアライズします（例：256 個の別々のテンソル）。しかし、最適化されたランタイムカーネルは、グループ化された GEMM 演算のためにエキスパートを単一の連続テンソルにパックすることを要求します。

この課題を解決するために、Hugging Face は WeightConverter 抽象化を通じて 変換パイプライン を導入しました。これにより、ロードプロセスは単純なキーごとのコピーから動的な変換へとシフトします。

  • MergeModulelist: 複数のエキスパートテンソルを単一の連続テンソルにスタックします。
  • SplitModulelist: パックされたテンソルを個々のエキスパートに分割します。
  • Lazy Materialization（遅延実体化）: ローダーはキーを一度だけスキャンし、依存関係が整った時点でスレッドプールを介してテンソルを実体化するため、メモリピークを削減し、再スキャンを回避します。

重みロードベンチマーク

Qwen/Qwen1.5-110B-Chat を単一の A100（80GB）で使用したベンチマークは、v4 に比べて v5 パイプラインで大幅な高速化が示されています。

| バージョン | 戦略 | ロードモード | 時間 | | :--- | :--- | :--- | | v4.57.6 | device_map="auto" | スレッドプール | 66.24s | | v5 | device_map="auto" | 非同期（デフォルト） | 20.71s | | v5 | TP | 非同期 | 10.1s |

プラガブルなエキスパートバックエンド

エキスパート計算をモデル実装から切り離すために、Hugging Face は @use_experts_implementation デコレータを使用した Experts Backend システムを導入しました。これにより、モデルは実行時に 3 つの実行戦略間で切り替えることができます。

  1. eager: 選択されたエキスパートをループ処理します。主にデバッグと正確性の確認に使用されます。
  2. batched_mm: torch.bmm を使用してトークンごとにエキスパートの重みを複製します。小バッチや GPU 集中型ワークロードに最適化されています。
  3. grouped_mm: torch._grouped_mm を使用してトークンをエキスパート ID でソートし、単一のグループ化された GEMM を実行します。大バッチやメモリ制約のある環境に最適化されています。

エキスパート並列化 (EP)

エキスパート並列化は、数千億パラメータ規模のモデルがエキスパートをデバイス間で分散させることで複数の GPU にスケールできるようにします。標準的なテンソル並列化とは異なり、各デバイスは割り当てられたエキスパートのサブセット（num_experts / num_devices）のみをロードします。

これは DistributedConfig(enable_expert_parallel=True) によって実装され、以下の 2 つのコアコンポーネントに依存します。

  • GroupedGemmParallel: エキスパート次元（dim=0）に沿ったエキスパート重みのシャーディングを処理します。
  • RouterParallel: グローバルエキスパートインデックスをローカルインデックスに再マップし、all-reduce 操作を使用してデバイス間で部分出力を結合します。

Unsloth を用いた最適化された MoE トレーニング

Unsloth との協業により、Hugging Face はエキスパートバックエンド抽象化と PyTorch の torch._grouped_mm API、さらにカスタム Triton グループ化 GEMM と LoRA カーネルを組み合わせて、より高速な MoE トレーニングを実現しました。これらの最適化により、以下が提供されます。

  • MoE トレーニングが最大 12 倍高速化。
  • VRAM 使用量が 35% 以上削減。
  • コンテキストウィンドウが約 6 倍長くなります。
  • v4 の transformers 実装と比較して、全体で 12〜30 倍の速度向上。

Sources