Transformer 中的专家混合（MoEs）

Hugging Face 已彻底改造了 transformers 库，以提供对专家混合（MoE）架构的一流支持。此重新设计解决了 MoE 模型在检查点存储方式与硬件执行方式之间的根本不匹配，实现了显著更快的权重加载、更高效的通过可插拔后端进行推理，以及通过专家并行实现可扩展的分布。

MoE 基础：容量 vs. 活动参数

专家混合模型用一组称为专家的可学习子网络取代密集前馈层。路由器为每个 token 选择少量专家进行处理，从而将模型的总体容量与推理成本解耦。

  • 总参数: 决定模型的整体容量和质量。
  • 活动参数: 决定每个 token 的推理速度和计算成本。

例如，gpt-oss-20b 模型拥有 21B 总参数，但每个 token 只激活约 3.6B 参数（使用 32 个专家中的 4 个）。这使得模型在保持 21B 参数系统质量的同时，以 3.6B 参数模型的速度运行，在 M3 Ultra Mac 上实现约每秒 115 个 token。

权重加载重构与 WeightConverter

transformers 中的传统权重加载假设检查点张量与运行时参数之间是一对一映射。MoE 检查点通常会独立序列化专家（例如，256 个独立张量），但优化后的运行时内核需要将专家打包成单个连续张量，以进行分组 GEMM 操作。

为了解决此问题，Hugging Face 通过 WeightConverter 抽象引入了一个 转换管线。这将加载过程从简单的逐键复制转变为动态转换：

  • MergeModulelist: 将多个专家张量堆叠为单个连续张量。
  • SplitModulelist: 将打包的张量拆分回单个专家。
  • 惰性实例化: 加载器仅在依赖准备好时，通过线程池实例化张量，且只扫描一次键，降低内存峰值并避免重复扫描。

权重加载基准

使用单个 A100（80GB）上的 Qwen/Qwen1.5-110B-Chat 进行基准测试，展示了相较于 v4，v5 管线的显著加速：

| 版本 | 策略 | 加载模式 | 时间 | | :--- | :--- | :--- | | v4.57.6 | device_map="auto" | 线程池 | 66.24s | | v5 | device_map="auto" | 异步（默认） | 20.71s | | v5 | TP | 异步 | 10.1s |

可插拔专家后端

为将专家计算与模型实现解耦，Hugging Face 使用 @use_experts_implementation 装饰器引入了专家后端系统。这使得模型能够在运行时在三种执行策略之间切换：

  1. eager: 对选定的专家进行循环；主要用于调试和正确性验证。
  2. batched_mm: 使用 torch.bmm 为每个 token 复制专家权重；针对小批量和 GPU 密集型工作负载进行优化。
  3. grouped_mm: 使用 torch._grouped_mm 按专家 ID 对 token 排序并执行单次分组 GEMM；针对大批量和内存受限环境进行优化。

专家并行（EP）

专家并行使得拥有数千亿参数的模型能够通过在多个 GPU 上分配专家来实现扩展。不同于标准张量并行，每个设备仅加载其分配的专家子集（num_experts / num_devices）。

这通过 DistributedConfig(enable_expert_parallel=True) 实现，并依赖两个核心组件：

  • GroupedGemmParallel: 处理沿专家维度（dim=0）的专家权重分片。
  • RouterParallel: 将全局专家索引重新映射为本地索引，并使用 all-reduce 操作合并跨设备的部分输出。

使用 Unsloth 优化的 MoE 训练

通过与 Unsloth 的合作，Hugging Face 利用专家后端抽象和 PyTorch 的 torch._grouped_mm API，并结合自定义 Triton 分组 GEMM 与 LoRA 内核，实现了更快的 MoE 训练。这些优化提供了：

  • MoE 训练速度提升最高可达 12 倍。
  • VRAM 使用量降低超过 35%。
  • 上下文窗口长度约提升 6 倍。
  • 相较于 v4 transformers 实现，总体加速达 12‑30 倍。

SUMMARY: Hugging Face 重新设计了 transformers 库，通过全新的权重加载重构、可插拔的专家后端以及原生的专家并行，使专家混合（MoEs）成为一等公民。

TITLE: Transformer 中的专家混合（MoEs）

Sources