Бэкенды Experts

Все реализации Mixture-of-Experts (MoE) выполняют одни и те же вычисления высокого уровня. Для каждого токена роутер выбирает k экспертов. Затем скрытое состояние токена проецируется через параметры выбранных экспертов и агрегируется с весами маршрутизации. Различие между бэкендами экспертов заключается в том, как выполняются матричные умножения этих экспертов.

ExpertsInterface предоставляет оптимизированные бэкенды для экспертов. Он отделяет реализацию экспертов от кода модели, упрощая эксперименты с различными функциями. Новые бэкенды добавляются через тот же интерфейс.

Бэкенды "batched_mm" и "grouped_mm" также выполняют квантованных экспертов в формате FP8 и FP4 (упакованных в int8) через ядро Triton finegrained-fp8, считывая масштабы в формате float32 или UE8M0. Они служат резервным вариантом для квантованных чекпоинтов, когда бэкенд "deepgemm" недоступен.

Переключение на этапе декодирования

На GPU модель, загруженная с параметром experts_implementation="grouped_mm", автоматически переключается на "batched_mm" для стадии декодирования при генерации, что значительно быстрее при малом количестве токенов. Исходный бэкенд восстанавливается после завершения генерации. На CPU grouped_mm остаётся активным на протяжении всей генерации, так как он эффективнее при любом размере входных данных.

Переключатель распространяется на слои MoE в модели верхнего уровня и в любом субконфигурационном backbone, например в text_config модели vision-language. Только записи grouped_mm переключаются на batched_mm. Эксперты, работающие на любом другом backend, сохраняют его.

Настройка backend для экспертов

Используйте аргумент experts_implementation в from_pretrained(), чтобы создать модель с конкретным бэкендом экспертов.

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(). Он возвращает dict с одним элементом для модели и одним элементом для каждой подконфигурации.

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

Бэкенд экспертов, специфичных для базовой модели

Мультимодальные модели могут иметь несколько подконфигураций (например, различные бэкбоны). Вы можете задать отдельный бэкенд экспертов для каждой подконфигурации, передав dict в experts_implementation при загрузке.

Ключи в отображении должны совпадать с именами подконфигураций.

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

Установите backend для экспертов глобально с пустым ключом.

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

DeepGEMM

Бэкенд "deepgemm" направляет умножения матриц экспертов через ядра DeepGEMM, распространяемые в репозитории kernels-community/deep-gemm. Он работает с неквантованными экспертами в формате bfloat16 и с экспертами, квантованными в FP8/FP4 и загруженными через Fine-grained FP8.

Бэкенд "deepgemm" требует:

• CUDA GPU с вычислительной способностью ≥ 9.0 (Hopper или новее).
• Полный набор CUDA toolkit с nvcc для JIT-компиляции ядра, версия 12.3 или новее для Hopper и 12.9 или новее для Blackwell. Установка только с рантаймом не содержит nvcc и отклоняется. Transformers находит toolkit через CUDA_HOME, CUDA_PATH, nvcc в PATH, затем /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" автоматически выбирает ядро сгруппированного GEMM для FP8 (или FP4 на Blackwell). 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" требуют масштабы экспертов UE8M0, являющиеся степенью двойки. Чекпоинт, квантованный с обычными масштабами float32 (scale_fmt="float"), вызывает ValueError при первом прямом проходе, а не молча искажает вывод. Загрузите чекпоинт, квантованный с scale_fmt="ue8m0", или переключитесь на grouped_mm или batched_mm, которые напрямую используют блочные масштабы float32. Hopper (SM90+) использует масштабы float32 на пути DeepGEMM без преобразования.

Основная причина передачи FineGrainedFP8Config для предварительно квантованного чекпоинта — это деквантование обратно в bfloat16, в этом случае эксперты работают в bfloat16, а не по пути FP8/FP4 DeepGEMM.

from transformers import AutoModelForCausalLM, FineGrainedFP8Config

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

Fused Mega MoE на Blackwell

На Blackwell (SM100+) установите experts_implementation="deepgemm_megamoe", чтобы запустить единое слиянное ядро, объединяющее диспетчеризацию с экспертной параллельностью, проекцию вверх, SwiGLU, проекцию вниз и объединение EP, перекрывая передачи по NVLink вычислениями на тензорных ядрах.

Этот бэкенд требует:

• GPU Blackwell (вычислительная способность ≥ 10.0) с CUDA toolkit (nvcc) версии 12.9 или новее.
• Экспертные веса в FP4-упаковке в паре с масштабами весов UE8M0 (предварительно квантованный чекпоинт обычно объявляет expert_dtype="fp4" и scale_fmt="ue8m0" в своём конфиге).
• Группа процессов torch.distributed для экспертно-параллельной группы, которая автоматически предоставляется обёрткой экспертной параллели, когда 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

Бэкенд "sonicmoe" объединяет прямой проход маршрутизируемого MoE (диспетчеризация, проекция gated up, активация, проекция down) в набор высокооптимизированных групповых GEMM-ядер CuteDSL, созданных на основе библиотеки quack и распространяемых через kernels-community/sonic-moe.

Бэкенд "sonicmoe" требует:

• CUDA GPU с вычислительной способностью ≥ 9.0 (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 в настоящее время поддерживает только bfloat16 при компиляции с torch.compile. Кроме того, он несовместим с CUDA-графиками, поэтому при компиляции необходимо использовать 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 и без него.