Transformer 中的 Mixture of Experts:深度技术解析
核心概念:MoE 在保持 Transformer 主干的同时,将某些密集前馈层替换为一组专家。每个 token 由路由器选择少量专家处理,实现"稀疏激活"。
为什么选择 MoE?
- 更好的计算效率:给定固定训练 FLOP 预算,MoE 通常优于密集模型
- 自然的并行化轴:专家提供计算图的结构边界
- 行业采用:Qwen3.5, MiniMax M2, GLM-5, Kimi K2.5, DeepSeek R1 等
关键架构组件
1. 专家路由机制
不同 token 根据其隐藏表示激活不同专家。模型容量取决于总参数,但推理速度取决于活跃参数。
2. Transformers v5 重量加载优化
| 版本 | 策略 | 加载模式 | 时间 |
|---|---|---|---|
| v4.57.6 | device_map="auto" | Threadpool | 66.24s |
| v5 | device_map="auto" | Async | 20.71s |
| v5 | TP | Async | 10.1s |
速度提升来自:单次路由、异步实例化、转换感知调度
3. 专家后端系统
- eager: 循环遍历选中的专家 - 用于正确性参考和调试
- batched_mm: 使用 torch.bmm API - 适合小 batch、GPU 重载工作负载
- grouped_mm: 使用 torch._grouped_mm API - 大 batch 或内存受限设置下表现出色
4. 专家并行 (Expert Parallelism)
将专家分布到多个设备,每个设备只加载分配给自己的专家子集。
from transformers import AutoModelForCausalLM, AutoTokenizer
from transformers.distributed.configuration_utils import DistributedConfig
distributed_config = DistributedConfig(enable_expert_parallel=True)
model = AutoModelForCausalLM.from_pretrained(
"openai/gpt-oss-120b",
dtype="auto",
distributed_config=distributed_config,
)
训练优化:与 Unsloth 合作
- ~12× 更快 MoE 训练
- >35% VRAM 减少
- ~6× 更长上下文
- 12-30× 整体加速(相比 v4)
主流 MoE 模型时间线
- 2023.12: Mixtral-8x7B
- 2024: DeepSeek V2
- 2025.01: DeepSeek R1(明确拐点)
- 2026: Qwen3.5, MiniMax M2, GLM-5, Kimi K2.5