Backends de Experts

Todas as implementações de Mixture-of-Experts (MoE) executam o mesmo cálculo de alto nível. Para cada token, um roteador seleciona k especialistas. O estado oculto do token é então projetado pelos parâmetros dos especialistas selecionados e agregado com os pesos de roteamento. A diferença entre os backends de especialistas é como essas multiplicações de matriz dos especialistas são executadas.

A ExpertsInterface fornece backends de experts otimizados. Ela desacopla a implementação dos experts do código do modelo para simplificar a experimentação com diferentes funções. Adicione novos backends por meio da mesma interface.

Os backends "batched_mm" e "grouped_mm" também executam especialistas quantizados em FP8 e FP4 (empacotados em int8) por meio do kernel Triton finegrained-fp8, lendo escalas em float32 ou UE8M0. Eles atuam como fallback para checkpoints quantizados quando o backend "deepgemm" não está disponível.

Comutação da etapa de decodificação

Na GPU, um modelo carregado com experts_implementation="grouped_mm" alterna automaticamente para "batched_mm" na etapa de decodificação da geração, o que é significativamente mais rápido para contagens menores de tokens. O backend original é restaurado assim que a geração termina. Na CPU, grouped_mm permanece ativo durante toda a geração, pois é mais eficiente em todos os tamanhos de entrada.

O switch atinge as camadas MoE no modelo de nível superior e em qualquer backbone de subconfiguração, como o text_config de um modelo de visão-linguagem. Apenas as entradas grouped_mm são alteradas para batched_mm. Os especialistas que executam qualquer outro backend mantêm o valor atual.

Defina um backend de especialistas

Use o argumento experts_implementation em from_pretrained() para instanciar um modelo com um backend de experts específico.

from transformers import AutoModelForCausalLM

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

Alterne entre backends de especialistas em tempo de execução sem recarregar o modelo usando set_experts_implementation().

model.set_experts_implementation("eager")

Leia o backend que está sendo executado atualmente com get_experts_implementation(). Ele retorna um dict com uma entrada para o modelo e uma entrada para cada sub-configuração.

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

Backend de especialistas específicos do backbone

Modelos multimodais podem ter múltiplas subconfigurações (por exemplo, backbones diferentes). Você pode definir um backend de especialistas diferente para cada subconfiguração passando um dict para experts_implementation no momento do carregamento.

As chaves no mapeamento devem corresponder aos nomes das subconfigurações.

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,
)

Defina o backend de especialistas globalmente com uma chave vazia.

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

DeepGEMM

O backend "deepgemm" roteia os matmuls de especialistas por meio dos kernels DeepGEMM distribuídos por kernels-community/deep-gemm. Ele funciona com especialistas bfloat16 sem quantização e com especialistas quantizados em FP8/FP4 carregados por meio de Fine-grained FP8.

O backend "deepgemm" requer:

• GPU CUDA com capacidade de computação ≥ 9.0 (Hopper ou mais recente).
• Um toolkit CUDA completo com nvcc para a compilação JIT do kernel, versão 12.3 ou posterior em Hopper e 12.9 ou posterior em Blackwell. Uma instalação apenas de runtime não possui nvcc e é rejeitada. O Transformers localiza o toolkit por meio de CUDA_HOME, CUDA_PATH, nvcc no PATH e, em seguida, /usr/local/cuda.
• O pacote kernels.

from transformers import AutoModelForCausalLM

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

O kernel é carregado de forma preguiçosa na primeira chamada de forward.

Especialistas quantizados em FP8 e FP4

Checkpoints no estilo DeepSeek geralmente são pré-quantizados e possuem sua própria configuração de quantização, portanto, não é necessário passar um FineGrainedFP8Config. O backend "deepgemm" seleciona automaticamente o kernel de grouped-GEMM em FP8 (ou FP4 no Blackwell). O DeepGEMM requer escalas de ativação dinâmicas por linha (activation_scheme="dynamic") e rejeita a quantização de ativação estática (por tensor).

from transformers import AutoModelForCausalLM

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

Para pesos de especialistas empacotados em FP4 (estilo DeepSeek V4), a GPU deve ser SM100+ (Blackwell). A configuração do checkpoint geralmente define expert_dtype="fp4" e scale_fmt="ue8m0".

No Blackwell (SM100+), os kernels de experts "deepgemm" e "deepgemm_megamoe" exigem escalas de experts UE8M0 que sejam potências de dois. Um checkpoint quantizado com escalas float32 comuns (scale_fmt="float") levanta um ValueError no primeiro forward, em vez de corromper silenciosamente a saída. Carregue um checkpoint quantizado com scale_fmt="ue8m0", ou alterne para grouped_mm ou batched_mm, que consomem diretamente escalas de bloco float32. O Hopper (SM90+) consome escalas float32 no caminho DeepGEMM sem conversão.

O principal motivo para passar um FineGrainedFP8Config para um checkpoint pré-quantizado é desquantizá-lo de volta para bfloat16, caso em que os especialistas são executados em bfloat16 em vez do caminho DeepGEMM FP8/FP4.

from transformers import AutoModelForCausalLM, FineGrainedFP8Config

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

Fused Mega MoE no Blackwell

No Blackwell (SM100+), defina experts_implementation="deepgemm_megamoe" para executar um único kernel fundido que combina o despacho paralelo de especialistas, a projeção para cima, o SwiGLU, a projeção para baixo e a combinação de EP, sobrepondo as transferências NVLink ao cálculo nos tensor cores.

Este backend requer:

• Uma GPU Blackwell (capacidade de computação ≥ 10.0) com um toolkit CUDA (nvcc) 12.9 ou posterior.
• Pesos de especialistas empacotados em FP4 combinados com escalas de pesos UE8M0 (o checkpoint pré-quantizado normalmente declara expert_dtype="fp4" e scale_fmt="ue8m0" em sua configuração).
• Um grupo de processos torch.distributed para o grupo de paralelismo de especialistas, fornecido automaticamente pelo encapsulamento de paralelismo de especialistas quando ep_size > 1.

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

O backend "sonicmoe" funde o forward do MoE roteado (dispatch, projeção gated up, ativação, projeção down) em um conjunto de kernels CuteDSL grouped-GEMM altamente otimizados, construídos sobre a biblioteca quack e distribuídos por kernels-community/sonic-moe.

O backend "sonicmoe" requer:

• GPU CUDA com capacidade de computação ≥ 9.0 (Hopper ou mais recente).
• O pacote kernels e o pacote nvidia-cutlass-dsl.
• Especialistas com ativação com gate (silu, gelu ou relu, mapeados para SwiGLU/GeGLU/ReGLU).

from transformers import AutoModelForCausalLM

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

Se os requisitos não forem atendidos, o forward lança ImportError e você deve escolher uma experts_implementation diferente.

torch.compile

Os backends "eager", "batched_mm" e "grouped_mm" são compatíveis com torch.compile em diferentes graus. A tabela a seguir resume essa compatibilidade. Os backends "deepgemm", "deepgemm_megamoe" e "sonicmoe" usam kernels CUDA externos e não estão cobertos por esta tabela.

Observações:

• O backend de especialistas grouped_mm atualmente suporta apenas bfloat16 quando compilado com torch.compile. Além disso, ele não é compatível com CUDA graphs, portanto, você deve usar mode=None ou mode="max-autotune-no-cudagraphs" ao compilar.
• O backend de especialistas eager usa uma operação dependente dos dados para encontrar quais especialistas são usados em uma passagem de frente. Essa operação não é compatível com a compilação de grafo completo (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")

Benchmarks

Este benchmark compara diferentes tamanhos de entrada e implementações de especialistas com e sem torch.compile.