Experts バックエンド

すべての Mixture-of-Experts (MoE) 実装は、同じ高レベルの計算を実行します。各トークンに対して、ルーターが k 個のエキスパートを選択します。その後、トークンの隠れ状態は選択されたエキスパートのパラメータによって射影され、ルーティング重みで集約されます。エキスパートのバックエンド間の違いは、これらのエキスパート行列積が どのように 実行されるかという点です。

ExpertsInterface は、最適化されたエキスパートのバックエンドを提供します。エキスパートの実装をモデルコードから分離することで、異なる関数を用いた実験を簡素化します。同じインターフェースを通じて新しいバックエンドを追加できます。

"batched_mm" および "grouped_mm" バックエンドも、Triton の finegrained-fp8 カーネルを通じて、FP8 および FP4（int8 パック）量子化エキスパートを実行し、float32 または UE8M0 スケールを読み取ります。"deepgemm" バックエンドが利用できない場合の、量子化チェックポイントのフォールバックとして機能します。

デコード段階での切り替え

GPU 上で、experts_implementation="grouped_mm" を指定して読み込まれたモデルは、生成のデコード段階において、トークン数が少ない場合に大幅に高速な "batched_mm" に自動的に切り替わります。生成が完了すると、元のバックエンドに復元されます。CPU 上では、grouped_mm はすべての入力サイズで効率的であるため、生成全体を通じて有効なままになります。

スイッチは、トップレベルのモデルおよび、ビジョン・言語モデルのtext_configのようなサブ設定のバックボーンにある MoE レイヤーに到達します。grouped_mmエントリのみがbatched_mmに切り替わります。その他のバックエンドで実行されているエキスパートは、そのまま維持されます。

expertsバックエンドを設定する

from_pretrained() の experts_implementation 引数を使用して、特定の experts バックエンドでモデルをインスタンス化します。

from transformers import AutoModelForCausalLM

model = AutoModelForCausalLM.from_pretrained(
    "Qwen/Qwen1.5-MoE-A2.7B",
    dtype="bfloat16",
    experts_implementation="batched_mm",
)

モデルを再読み込みせずに、set_experts_implementation() を使用して実行時にエキスパートのバックエンドを切り替えます。

model.set_experts_implementation("eager")

現在実行されているバックエンドを get_experts_implementation() で読み取ります。この関数は、モデル用のエントリと、サブ設定ごとに1つずつのエントリを含む dict を返します。

model.get_experts_implementation()
# {"": "grouped_mm", "text_config": "grouped_mm", "vision_config": "eager"}

バックボーン固有のエクスパートのバックエンド

マルチモーダルモデルには複数のサブ設定（例えば、異なるバックボーン）を持つ場合があります。experts_implementation に dict を渡すことで、読み込み時にサブ設定ごとに異なるエキスパートのバックエンドを設定できます。

マッピングのキーは、サブ設定の名前と一致する必要があります。

from transformers import AutoModelForImageTextToText

experts_implementation_per_backbone = {
    "text_config": "grouped_mm",
    "vision_config": "eager",
}

model = AutoModelForImageTextToText.from_pretrained(
    "Qwen/Qwen3-VL-Moe",
    experts_implementation=experts_implementation_per_backbone,
)

空のキーで、experts バックエンドをグローバルに設定します。

model = AutoModelForCausalLM.from_pretrained(
    "Qwen/Qwen1.5-MoE-A2.7B",
    experts_implementation={"": "batched_mm"},
)

DeepGEMM

"deepgemm" バックエンドは、kernels-community/deep-gemm によって配布される DeepGEMM カーネルを介して、エキスパートの行列積をルーティングします。非量子化の bfloat16 エクスパート、および Fine-grained FP8 経由で読み込まれた FP8/FP4 量子化エキスパートに対応しています。

"deepgemm"バックエンドには、以下が必要です:

• 計算能力が 9.0 以上の CUDA GPU（Hopper 以降）。
• カーネルの JIT コンパイルに nvcc を備えた完全な CUDA ツールキットが必要です。Hopper ではバージョン 12.3 以降、Blackwell では 12.9 以降が対象です。ランタイムのみのインストールには nvcc が含まれていないため、使用できません。Transformers は CUDA_HOME、CUDA_PATH、PATH 上の nvcc、その後に /usr/local/cuda の順でツールキットを検索します。
• kernelsパッケージ。

from transformers import AutoModelForCausalLM

model = AutoModelForCausalLM.from_pretrained(
    "deepseek-ai/DeepSeek-V3",
    dtype="bfloat16",
    experts_implementation="deepgemm",
)

カーネルは最初のフォワードパス時に遅延ロードされます。

FP8およびFP4量子化されたエキスパート

DeepSeekスタイルのチェックポイントは通常、事前に量子化されており、独自の量子化設定を保持しているため、FineGrainedFP8Configを渡す必要はありません。"deepgemm"バックエンドは、FP8（BlackwellではFP4）のグループ化GEMMカーネルを自動的に選択します。DeepGEMMは動的な行ごとの活性化スケール（activation_scheme="dynamic"）を必要とし、静的（テンソルごと）な活性化量子化は拒否されます。

from transformers import AutoModelForCausalLM

model = AutoModelForCausalLM.from_pretrained(
    "deepseek-ai/DeepSeek-V3",
    experts_implementation="deepgemm",
)

FP4パックされたエキスパート重み（DeepSeek V4スタイル）の場合、GPUはSM100以上（Blackwell）である必要があります。チェックポイントの設定では、通常expert_dtype="fp4"とscale_fmt="ue8m0"が設定されています。

Blackwell (SM100+) では、"deepgemm" および "deepgemm_megamoe" の experts カーネルは、2 の累乗である UE8M0 の expert scales を必要とします。通常の float32 scales (scale_fmt="float") で量子化されたチェックポイントをロードすると、出力が静かに破損するのではなく、最初の forward で ValueError が発生します。scale_fmt="ue8m0" で量子化されたチェックポイントをロードするか、float32 の block scales を直接使用できる grouped_mm または batched_mm に切り替えてください。Hopper (SM90+) では、DeepGEMM パス上で変換なしで float32 scales を使用します。

事前量子化済みチェックポイントに対してFineGrainedFP8Configを渡す主な理由は、bfloat16へデクォンタイズすることであり、その場合、エキスパートはFP8/FP4 DeepGEMMパスではなくbfloat16で実行されます。

from transformers import AutoModelForCausalLM, FineGrainedFP8Config

model = AutoModelForCausalLM.from_pretrained(
    "deepseek-ai/DeepSeek-V3",
    quantization_config=FineGrainedFP8Config(dequantize=True),
    experts_implementation="deepgemm",
)

Blackwell での Fused Mega MoE

Blackwell (SM100+) では、experts_implementation="deepgemm_megamoe" を設定することで、エキスパート並列ディスパッチ、アップ投影、SwiGLU、ダウン投影、および EP コンバインを単一の融合カーネルとして実行できます。これにより、NVLink 転送とテンソルコアの計算をオーバーラップできます。

このバックエンドには以下が必要です:

• CUDA ツールキット（nvcc）12.9 以降を備えた Blackwell GPU（compute capability ≥ 10.0）。
• UE8M0の重みスケールと組み合わせたFP4パックされたエキスパート重み（事前量子化済みチェックポイントは、通常そのconfigでexpert_dtype="fp4"およびscale_fmt="ue8m0"を宣言しています）。
• ep_size > 1 の場合、expert-parallel ラッピングが自動的に提供する、expert-parallel グループ用の torch.distributed プロセスグループです。

import os

from transformers import AutoModelForCausalLM, DistributedConfig

distributed_config = DistributedConfig(
    tp_size=int(os.environ["WORLD_SIZE"]),
    ep_size=int(os.environ["WORLD_SIZE"]),
)
model = AutoModelForCausalLM.from_pretrained(
    "deepseek-ai/DeepSeek-V4",
    experts_implementation="deepgemm_megamoe",
    distributed_config=distributed_config,
)

SonicMoE

"sonicmoe" バックエンドは、ルーティングされた MoE のフォワード（ディスパッチ、ゲート付きアップ投影、活性化、ダウン投影）を、quack ライブラリを基盤として構築され、kernels-community/sonic-moe によって配布される、高度に最適化された CuteDSL グループ GEMM カーネルのセットに融合します。

"sonicmoe"バックエンドには、以下が必要です:

• 計算能力が 9.0 以上の CUDA GPU（Hopper 以降）。
• kernelsパッケージとnvidia-cutlass-dslパッケージ。
• ゲート付き活性化（silu、gelu、または relu、それぞれ SwiGLU/GeGLU/ReGLU にマッピングされる）を持つエキスパート。

from transformers import AutoModelForCausalLM

model = AutoModelForCausalLM.from_pretrained(
    "Qwen/Qwen1.5-MoE-A2.7B",
    dtype="bfloat16",
    experts_implementation="sonicmoe",
)

要件が満たされていない場合、forward は ImportError を送出します。その場合は、別の experts_implementation を選択してください。

torch.compile

"eager"、"batched_mm"、"grouped_mm" バックエンドは、torch.compile とある程度互換性があります。以下の表にその互換性をまとめます。"deepgemm"、"deepgemm_megamoe"、"sonicmoe" バックエンドは外部の CUDA カーネルを介してルーティングされるため、この表の対象外です。

注:

• grouped_mmエキスパートバックエンドは、torch.compileでコンパイルする場合、現在bfloat16のみをサポートしています。また、CUDA graphsと互換性がありません。そのため、コンパイル時にはmode=Noneまたはmode="max-autotune-no-cudagraphs"を使用する必要があります。
• eagerエキスパートバックエンドは、フォワードパスでどのエキスパートが使用されるかを特定するために、データ依存の操作を使用します。この操作は、グラフ全体のコンパイル（fullgraph=True）と互換性がありません。

import torch
from transformers import AutoModelForCausalLM

model = AutoModelForCausalLM.from_pretrained(
    "Qwen/Qwen1.5-MoE-A2.7B",
    dtype="bfloat16",
    experts_implementation="grouped_mm",
    device_map="auto",
).eval()

# Works for grouped_mm (no CUDA graphs)
model.forward = torch.compile(model.forward, mode="max-autotune-no-cudagraphs")

ベンチマーク

このベンチマークでは、torch.compileを使用する場合と使用しない場合の、異なる入力サイズとエキスパート実装を比較しています。