Các backend chuyên gia

Tất cả các triển khai Mixture-of-Experts (MoE) đều thực hiện cùng một phép tính ở cấp độ cao. Đối với mỗi token, một router sẽ chọn k chuyên gia. Trạng thái ẩn của token sau đó được chiếu qua các tham số của các chuyên gia đã chọn và được tổng hợp với các trọng số định tuyến. Sự khác biệt giữa các backend chuyên gia nằm ở cách các phép nhân ma trận chuyên gia đó được thực thi.

ExpertsInterface cung cấp các backend chuyên gia được tối ưu hóa. Nó tách biệt việc triển khai chuyên gia khỏi mã mô hình để đơn giản hóa việc thử nghiệm với các hàm khác nhau. Thêm các backend mới thông qua cùng một giao diện.

Các backend "batched_mm" và "grouped_mm" cũng chạy các chuyên gia lượng tử hóa FP8 và FP4 (đóng gói int8) thông qua kernel Triton finegrained-fp8, đọc các scale dạng float32 hoặc UE8M0. Chúng đóng vai trò là phương án dự phòng cho các checkpoint đã lượng tử hóa khi backend "deepgemm" không khả dụng.

Chuyển đổi giai đoạn giải mã

Trên GPU, một model được nạp với experts_implementation="grouped_mm" sẽ tự động chuyển sang "batched_mm" cho giai đoạn decode của quá trình sinh, giúp nhanh hơn đáng kể khi số lượng token thấp. Backend gốc sẽ được khôi phục ngay khi quá trình sinh kết thúc. Trên CPU, grouped_mm vẫn hoạt động xuyên suốt quá trình sinh vì nó hiệu quả hơn ở mọi kích thước đầu vào.

Công tắc này áp dụng cho các lớp MoE trong mô hình cấp cao nhất và trong bất kỳ backbone cấu hình con nào, chẳng hạn như text_config của một mô hình thị giác-ngôn ngữ. Chỉ các mục grouped_mm mới chuyển sang batched_mm. Các chuyên gia chạy trên bất kỳ backend nào khác sẽ giữ nguyên.

Đặt backend chuyên gia

Sử dụng đối số experts_implementation trong from_pretrained() để khởi tạo một mô hình với backend experts cụ thể.

from transformers import AutoModelForCausalLM

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

Chuyển đổi giữa các backend chuyên gia tại thời điểm chạy mà không cần tải lại mô hình bằng cách sử dụng set_experts_implementation().

model.set_experts_implementation("eager")

Đọc backend đang chạy hiện tại bằng get_experts_implementation(). Hàm này trả về một dict với một mục cho model và một mục cho mỗi sub-config.

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

Backend chuyên gia cụ thể cho backbone

Các mô hình đa phương thức có thể có nhiều cấu hình con (ví dụ: các backbone khác nhau). Bạn có thể đặt một backend chuyên gia khác nhau cho từng cấu hình con bằng cách truyền một dict vào experts_implementation khi tải.

Các khóa trong ánh xạ phải khớp với tên cấu hình con.

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

Đặt backend chuyên gia toàn cục với một khóa rỗng.

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

DeepGEMM

Backend "deepgemm" định tuyến các phép nhân ma trận chuyên gia thông qua các kernel DeepGEMM được phân phối bởi kernels-community/deep-gemm. Nó hoạt động với các chuyên gia bfloat16 chưa lượng tử hóa và với các chuyên gia được lượng tử hóa FP8/FP4 được nạp thông qua Fine-grained FP8.

Backend "deepgemm" yêu cầu:

• GPU CUDA với khả năng tính toán ≥ 9.0 (Hopper hoặc mới hơn).
• Bộ công cụ CUDA đầy đủ với nvcc để biên dịch JIT cho kernel, phiên bản 12.3 trở lên trên Hopper và 12.9 trở lên trên Blackwell. Cài đặt chỉ runtime không có nvcc và sẽ bị từ chối. Transformers xác định vị trí bộ công cụ thông qua CUDA_HOME, CUDA_PATH, nvcc trên PATH, sau đó là /usr/local/cuda.
• Gói kernels.

from transformers import AutoModelForCausalLM

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

Kernel được nạp một cách lười biếng vào lần forward đầu tiên.

Các chuyên gia lượng tử hóa FP8 và FP4

Checkpoint theo phong cách DeepSeek thường đã được lượng tử hóa trước và đi kèm cấu hình lượng tử hóa riêng, vì vậy bạn không cần truyền FineGrainedFP8Config. Backend "deepgemm" sẽ tự động chọn kernel grouped-GEMM FP8 (hoặc FP4 trên Blackwell). DeepGEMM yêu cầu thang kích hoạt động theo từng hàng (activation_scheme="dynamic") và từ chối lượng tử hóa kích hoạt tĩnh (theo tensor).

from transformers import AutoModelForCausalLM

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

Đối với trọng số chuyên gia đóng gói FP4 (kiểu DeepSeek V4), GPU phải là SM100+ (Blackwell). Cấu hình checkpoint thường đặt expert_dtype="fp4" và scale_fmt="ue8m0".

Trên Blackwell (SM100+), các kernel chuyên gia "deepgemm" và "deepgemm_megamoe" yêu cầu các thang đo chuyên gia UE8M0 là lũy thừa của hai. Một checkpoint được lượng tử hóa với thang đo float32 thông thường (scale_fmt="float") sẽ gây ra ValueError ở bước forward đầu tiên thay vì âm thầm làm hỏng đầu ra. Hãy tải một checkpoint được lượng tử hóa với scale_fmt="ue8m0", hoặc chuyển sang grouped_mm hoặc batched_mm, vốn sử dụng trực tiếp các thang đo khối float32. Hopper (SM90+) sử dụng các thang đo float32 trên đường dẫn DeepGEMM mà không cần chuyển đổi.

Lý do chính để truyền một FineGrainedFP8Config cho một checkpoint đã lượng tử hóa trước là để khử lượng tử hóa nó trở lại bfloat16, trong trường hợp đó các chuyên gia sẽ chạy ở bfloat16 thay vì trên đường dẫn 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 trên Blackwell

Trên Blackwell (SM100+), đặt experts_implementation="deepgemm_megamoe" để chạy một kernel hợp nhất duy nhất, kết hợp việc dispatch song song chuyên gia, phép chiếu lên, SwiGLU, phép chiếu xuống và việc combine EP, đồng thời xen kẽ các thao tác truyền NVLink với tính toán trên tensor-core.

Backend này yêu cầu:

• GPU Blackwell (khả năng tính toán ≥ 10.0) với bộ công cụ CUDA (nvcc) 12.9 hoặc mới hơn.
• Trọng số chuyên gia đóng gói FP4 kết hợp với thang đo trọng số UE8M0 (checkpoint đã lượng tử hóa trước thường khai báo expert_dtype="fp4" và scale_fmt="ue8m0" trong cấu hình của nó).
• Một nhóm tiến trình torch.distributed cho nhóm song song chuyên gia, được bọc song song chuyên gia cung cấp tự động khi 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

Backend "sonicmoe" hợp nhất quá trình forward của MoE có định tuyến (dispatch, gated up projection, activation, down projection) thành một tập hợp các kernel grouped-GEMM CuteDSL được tối ưu hóa cao, được xây dựng trên thư viện quack và phân phối bởi kernels-community/sonic-moe.

Backend "sonicmoe" yêu cầu:

• GPU CUDA với khả năng tính toán ≥ 9.0 (Hopper hoặc mới hơn).
• Gói kernels và gói nvidia-cutlass-dsl.
• Các chuyên gia có kích hoạt có cổng (silu, gelu, hoặc relu, được ánh xạ sang SwiGLU/GeGLU/ReGLU).

from transformers import AutoModelForCausalLM

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

Nếu các yêu cầu không được đáp ứng, hàm forward sẽ gây ra ImportError và bạn nên chọn một experts_implementation khác.

torch.compile

Các backend "eager", "batched_mm" và "grouped_mm" tương thích với torch.compile ở các mức độ khác nhau. Bảng dưới đây tóm tắt mức độ tương thích của chúng. Các backend "deepgemm", "deepgemm_megamoe" và "sonicmoe" được định tuyến thông qua các CUDA kernel bên ngoài và không được đề cập trong bảng này.

Ghi chú:

• Backend chuyên gia grouped_mm hiện chỉ hỗ trợ bfloat16 khi được biên dịch bằng torch.compile. Ngoài ra, nó không tương thích với CUDA graphs, vì vậy bạn phải sử dụng mode=None hoặc mode="max-autotune-no-cudagraphs" khi biên dịch.
• Backend chuyên gia eager sử dụng một thao tác phụ thuộc vào dữ liệu để tìm ra các chuyên gia được sử dụng trong một bước truyền thẳng. Thao tác này không tương thích với biên dịch đồ thị đầy đủ (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

Benchmark này so sánh các kích thước đầu vào khác nhau và các triển khai chuyên gia có và không có torch.compile.