DeepSpeed MoE 训练完全指南：稀疏专家混合层 API、并行拓扑组合与 ZeRO-Offload 实战

DeepSpeed MoE 训练完全指南：稀疏专家混合层 API、并行拓扑组合与 ZeRO-Offload 实战

【免费下载链接】DeepSpeedDeepSpeed is a deep learning optimization library that makes distributed training and inference easy, efficient, and effective.项目地址: https://gitcode.com/GitHub_Trending/de/DeepSpeed

导读：本文是 DeepSpeed 官方《Mixture of Experts (DeepSpeed MoE)》训练教程的深度解读与实践手册。MoE 属于"稀疏激活"模型——总参数量可高达万亿级，但每次前向只激活其中一小部分专家，从而以近似稠密小模型的计算成本换取大幅精度提升。本文围绕仓库内 docs/_tutorials/mixture-of-experts.md 展开，结合 deepspeed/moe 目录下的真实源码实现与 tests/unit/v1/moe/test_moe.py 测试用例，系统讲透deepspeed.moe.layer.MoE层 API、专家/数据/ZeRO 的多种并行组合方式、Pyramid-Residual MoE（PR-MoE）结构，以及用 ZeRO-Offload 在有限 GPU 上训练超大 MoE 模型的完整配置。读完本文，你将能在自己的模型上亲手接入 MoE 层、正确初始化专家进程组并给出可运行的 ZeRO stage-2 配置。

MoE 是什么：用稀疏激活换取"参数扩容不增算力"

DeepSpeed v0.5 开始引入了对混合专家（Mixture of Experts, MoE）模型训练的原生支持。MoE 模型是新兴的一类稀疏激活（sparsely activated）模型，其核心特性是：计算成本相对参数量呈次线性增长。一个典型的例子是 Switch Transformer——其参数量超过 1.6 万亿，但训练它所需的计算量约等于训练一个 100 亿参数的稠密模型。这意味着在"恒定计算预算"约束下，通过 MoE 扩大模型规模可以带来显著的精度收益。

理解这一点的关键在于 MoE 层内部的拓扑结构：

• 门控网络（Gate / Router）：一个线性层，负责把每个 token 路由到 top-k 个专家；
• 专家集合（Experts）：一系列同构的子网络（通常就是 MLP 或nn.Linear），由同一个 expert 模块deepcopy得到；
• 负载均衡损失（auxiliary load-balancing loss）：抑制"所有 token 都涌向少数几个专家"的坍缩现象。

从当前仓库源码看，上述组件对应三个文件：

1. deepspeed/moe/layer.py 定义了面向用户的核心 APIdeepspeed.moe.layer.MoE；
2. deepspeed/moe/experts.py 中的Experts容器负责将同一个 expert 模块复制num_local_experts份，并给每个专家参数打上allreduce = False与group_name标记（这是后续按"专家/共享"分组通信的关键）；
3. deepspeed/moe/sharded_moe.py 实现了TopKGate门控与MOELayer调度主体，其代码注记表明它改编自 fairscale 与 GShard 论文（arXiv:2006.16668）中的算法流程。

前置条件与两种接入路径

官方教程给出的最基本要求是DeepSpeed MoE 需要 PyTorch 1.8 及以上版本（这一约束也直接体现在 tests/unit/v1/moe/test_moe.py 的 skip 逻辑里：DeepSpeed MoE tests need torch 1.8 or higher to run correctly）。

本教程聚焦于显式 MoE 层 API的接入方式，即用户手动在模型里插入MoE层。如果你希望避免手改模型，仓库还提供了AutoEP（Automatic Expert Parallelism）方案：从 DeepSpeed 配置中自动检测并替换受支持的 Hugging Face MoE 层，其相关实现位于 deepspeed/module_inject（含 auto_ep、auto_tp 等模块），本文不展开。此外仓库在 deepspeed/moe 之外还维护了用于 v2 推理引擎的 MoE 模块（deepspeed/inference/v2/modules/interfaces/moe_base.py），属于推理侧实现，训练场景仍以上述三个文件为核心。

专家进程组与五种并行形式（E / D / Z / M 的自由组合）

DeepSpeed MoE 支持五种不同形态的并行，且能同时利用 GPU 与 CPU 内存。灵活的进程组设计允许用户混合叠加主流并行技术。下表是官方教程给出的完整并行配置矩阵，本文原样保留并加以说明：

进程组创建：从手动初始化到按层自动创建

为了支持不同形态的并行，DeepSpeed 在内部创建了多种进程组。早期的帮助函数位于deepspeed/utils/groups.py（即仓库中的 deepspeed/utils/groups.py）：

deepspeed.utils.groups.initialize(ep_size="desired expert-parallel world size")

注意：该函数在当前仓库中已被正式废弃，训练代码不再需要手动调用它。打开 deepspeed/utils/groups.py 可以看到，initialize现在会直接抛出DeprecatedException，并给出明确指引：

Please do not use the groups.initialize() API as it is deprecated. Instead, pass the desired ep_size to deepspeed.moe.layer.MoE(..,ep_size,..)

也就是说，现在把ep_size作为参数直接传给每一个 MoE 层即可。这一新 API 允许用户构建"不同 MoE 层拥有不同专家数量、不同专家并行度"的模型。

进程组的实际创建逻辑在 MoE 层被挂到引擎后触发。查看 deepspeed/moe/layer.py 的set_deepspeed_parallelism/_create_process_groups：当某个ep_size的进程组尚不存在时，会根据是否启用了张量并行（groups.mpu）选择调用：

• groups._create_expert_and_data_parallel(ep_size, ...)（纯 E+D 场景）；
• groups._create_expert_data_and_model_parallel(ep_size, mpu, ...)（结合 Tensor/Pipeline/Model 并行的 E+D+M 场景）。

每个进程组以ep_size_{ep_size}命名并在多个进程组字典（如_EXPERT_PARALLEL_GROUP、_EXPERT_DATA_PARALLEL_GROUP）中登记，同名的组只创建一次。这里可以引出一个重要的通信拓扑事实（注释于 deepspeed/utils/groups.py 的 legacy E+D 示例中）：

• expert parallel group：负责专家间的 All-to-All 分发，不做梯度 all-reduce；
• expert data parallel group：只在持有同一批专家的副本之间做专家参数的梯度 all-reduce；
• data parallel group：对非专家参数做梯度 all-reduce。

MoE 层的语义约定：同维度专家与维度适配

MoE层有一个需要特别留意的约定：hidden_size既是该层的输入维度也是输出维度（expert 模块输入输出同维）。这会给部分视觉/卷积模型带来模型定义的改动——因为它们的输入输出维度在特定位置并不一致。以 CIFAR-10 示例为例，官方教程把第三个全连接层替换为 MoE 层，并为此额外增加一个维度桥接的全连接层（其输入维度等于 MoE 层输出维度）。

原始模型配置：

self.fc3 = nn.Linear(84, 10)

替换为 MoE 层之后：

self.fc3 = nn.Linear(84, 84) self.fc3 = deepspeed.moe.layer.MoE(hidden_size=84, expert=self.fc3, num_experts=args.num_experts, ep_size=<desired expert-parallel world size> ...) self.fc4 = nn.Linear(84, 10)

对照 deepspeed/moe/layer.py 的构造签名，MoE 层会把你的 expert 模块包装为每卡num_experts / ep_size个本地专家。为什么fc3需要从nn.Linear(84, 10)改为nn.Linear(84, 84)？因为专家内部要做84→84的变换以维持hidden_size恒定，而最终的类别映射（84→10）则下沉到新增的fc4。

MoE 层 API 全参数解析

deepspeed.moe.layer.MoE的完整签名可从 deepspeed/moe/layer.py 的 docstring 与构造函数中确认，参数及其默认值如下：

从源码可以看到两个硬性校验：deepspeed/moe/layer.py 断言num_experts % ep_size == 0（专家总数必须能被专家并行度整除）；deepspeed/moe/layer.py 对noisy_gate_policy的合法性做了白名单检查。

• 关于use_residual：当它为 True 时，层内会保留一个不经门控的旁路self.mlp = expert，并新增一个nn.Linear(hidden_size, 2)的coefficient层，对 MoE 输出与 MLP 旁路输出做 softmax 加权融合（deepspeed/moe/layer.py）。
• 关于容量（capacity）：_capacity的计算在 deepspeed/moe/sharded_moe.py：capacity = ceil((num_tokens / num_experts) * capacity_factor)，且不小于min_capacity。它决定每个专家在一个 batch 内最多处理多少个 token，是稀疏调度正确性与吞吐平衡的枢纽。

门控与调度：top-1 / top-2 的底层流程

在 deepspeed/moe/sharded_moe.py 中：

• TopKGate.forward先做 fp32 化的线性打分（支持Jitter/RSample噪声门控），然后按k分派到top1gating、top2gating或通用topkgating；
• 门控输出l_aux（负载均衡损失）、capacity、路由索引与 gate 权重；
• MOELayer.forward依据路由把 token 按 (expert, capacity slot) 排列，ep_size > 1时在专家并行组内执行两次_AllToAll（分发 token 与回收专家输出），最后按 gate 权重加权还原为原始序列形状（deepspeed/moe/sharded_moe.py）。

Pyramid-Residual MoE（PR-MoE）：金字塔 + 残差专家

官方在近期工作中提出了Pyramid-Residual MoE（PR-MoE）架构（arXiv:2201.05596），目标是显著提升参数效率。要构建 PR-MoE 模型，只需在原 API 上多做两件事：

1. 构造金字塔结构：把num_experts传成一个 list，例如[4, 8]，列表长度与模型中 MoE 层的数目一致，不同层承载不同数量的专家；
2. 启用残差专家：设置use_residual=True，把该层标记为 Residual MoE 层。

代码形态：

self.experts = deepspeed.moe.layer.MoE(hidden_size=input_dim, expert=ExpertModule(), num_experts=[..], ep_size=ep_size, use_residual=True)

注意：PR-MoE 混合使用金字塔专家数与残差融合必须配合使用——list 形式的num_experts是为了跨层形成金字塔，而use_residual=True保证每层在门控专家之外还有一个 MLP 残差通路做加权融合。仓库单元测试 tests/unit/v1/moe/test_moe.py 使用SimplePRMoEModel对use_residual参数在 zero stage 0/1/2 与不同ep_size组合下进行了覆盖验证，说明 PR-MoE 与 ZeRO 各阶段均可正常协同训练。

一个可推理的端到端场景：8 专家切 4 卡

下面用官方教程中的场景把前面的概念串起来。假设：

WORLD_SIZE = 4 # 总 GPU 数 EP_WORLD_SIZE = 2 # 每个专家并行组内的 GPU 数 EXPERTS = [8] # 该层专家总数

模型代码只需这样写：

self.experts = deepspeed.moe.layer.MoE(hidden_size=input_dim, expert=ExpertModule(), num_experts=EXPERTS, ep_size=EP_WORLD_SIZE)

运行起来后，DeepSpeed 运行时将训练一个共 8 个专家、分布到 4 张 GPU、每卡 4 个专家的 MoE 模型——也就是表格中的E + D（Expert + Data）模式：2 个 rank 组成一个专家并行组共享同一批 8 个专家，WORLD_SIZE / EP_WORLD_SIZE = 2组之间构成专家数据并行。

教程还给出了等价的逐层改造代码骨架：

import torch import deepspeed import deepspeed.utils.groups as groups from deepspeed.moe.layer import MoE WORLD_SIZE = 4 EP_WORLD_SIZE = 2 EXPERTS = 8 fc3 = torch.nn.Linear(84, 84) fc3 = MoE(hidden_size=84, expert=fc3, num_experts=EXPERTS, ep_size=EP_WORLD_SIZE, k=1) fc4 = torch.nn.Linear(84, 10)

这里k=1表示每个 token 只路由到 1 个专家（top-1 gating）。一个可以直接本地阅读的实现事实是：num_local_experts = num_experts // ep_size（deepspeed/moe/layer.py），而Experts容器会通过copy.deepcopy在每卡复制出num_local_experts个专家实例，并为每个专家参数设置param.allreduce = False与param.group_name = ep_size_{ep_size}（deepspeed/moe/experts.py）。allreduce = False正是 MoE 参数与非专家参数在梯度通信上"分道扬镳"的判据。

Random Token Selection：默认开启的收敛优化

传统的 MoE 训练存在"有偏选择（biased selection）"问题——超容量专家会按固定规则丢弃一部分 token，导致梯度估计偏差、影响收敛。DeepSpeed 为此设计了一套Random Token Selection（随机 token 选择）技术，用于显著改善收敛性。

这个特性已经内置在 DeepSpeed 运行时中并默认开启，用户无需设置任何 config flag 或命令行参数即可享受收益。从源码层面印证：MoE构造参数use_rts: bool = True（deepspeed/moe/layer.py）一路传入TopKGate与top1gating；在 deepspeed/moe/sharded_moe.py 中，use_rts=True时会给 one-hot 路由掩码乘上一个 [0,1) 均匀分布的随机数，再沿 token 维度做torch.topk挑选哪些 token 真正进入专家容量槽——相当于让"谁被丢弃"的决策随机化，从而消除系统性偏差。若想关闭可显式传use_rts=False。

ZeRO-Offload 组合 MoE：在有限 GPU 上训练超大模型

MoE 的参数量虽大，但绝大多数参数集中在专家上，而专家参数天然被ep_size切分到多卡，本身不占单卡重复显存。真正占显存的瓶颈是非专家（底座）参数与优化器状态。因此官方教程给出的超大模型方案是：MoE（切专家）+ ZeRO-Offload（切底座参数并把优化器状态卸载到 CPU 内存）。

第一步：构建 MoE 参数组

ZeRO 优化器依赖两个不同的参数分组（专家参数组 + 共享/非专家参数组）来差异化地做通信与 offload。官方推荐的参数组创建函数如下：

def create_moe_param_groups(model): from deepspeed.moe.utils import split_params_into_different_moe_groups_for_optimizer parameters = {'params': [p for p in model.parameters()], 'name': 'parameters'} return split_params_into_different_moe_groups_for_optimizer(parameters)

其底层实现在 deepspeed/moe/utils.py：split_params_into_different_moe_groups_for_optimizer依据is_moe_param(p)（即p.allreduce == False，见 deepspeed/moe/utils.py）把参数拆分，并进一步按param.group_name（每个不同的ep_size一个名字）细分成多个'moe': True的组；非专家参数留在普通组中。此外它还按max_group_size（默认178956971，接近 int32 上限）对超大专家组做再切分，避免单组参数规模溢出。

随后把这两个参数组喂给 ZeRO stage-2 优化器即可：

net = Net() parameters = create_moe_param_groups(net) model_engine, optimizer, trainloader, __ = deepspeed.initialize( args=args, model=net, model_parameters=parameters, training_data=trainset)

关于"自动化"，需要特别说明一个演进事实：当前仓库已经自动完成了上述参数组拆分。在 deepspeed/runtime/engine.py 的deepspeed.initialize路径中会调用configure_moe_param_groups(model_parameters)；该函数（deepspeed/moe/utils.py）能自动识别输入是裸Parameter列表还是已带分组的 dict 列表，缺失 MoE 分组时自动补建。因此现代代码即使不手动调create_moe_param_groups，引擎也能跑通——tests/unit/v1/moe/test_moe.py 的TestSimpleMoE直接deepspeed.initialize(config=config_dict, model=model)并注释"should automatically create moe param groups in deepspeed backend"即为实证。不过手动构造分组仍然有效，且在需要精细控制分组命名/数量时是推荐做法。

第二步：ZeRO-Offload 的 ds_config

要在 cifar 类示例上以ZeRO-Offload（stage 2）+ MoE方式运行，请在ds_config中设置：

"zero_optimization": { "stage": 2, "allgather_partitions": true, "reduce_scatter": true, "allgather_bucket_size": 50000000, "reduce_bucket_size": 50000000, "overlap_comm": true, "contiguous_gradients": true, "cpu_offload": true }

各字段作用与约束如下：

• stage: 2：启用 ZeRO stage-2（梯度分片），这是与cpu_offload配合的经典配置；当前实现同时支持 stage 0/1/2 下运行 MoE（见测试参数化 tests/unit/v1/moe/test_moe.py）。
• allgather_partitions / reduce_scatter：用 reduce-scatter 替代逐参数 all-reduce，是 ZeRO-2 减小通信量的关键，默认建议true。
• allgather_bucket_size / reduce_bucket_size（单位：元素个数，示例取 50,000,000）：梯度归约与参数聚合的桶大小，影响通信粒度与峰值内存的平衡，可按模型 hidden size 与 batch 适当调大。
• overlap_comm：梯度通信与计算重叠，掩盖通信延迟。
• contiguous_gradients：把梯度搬进连续缓冲区，减少内存碎片与通信次数。
• cpu_offload: true：核心开关，把 stage-2 的优化器状态与梯度卸载到 CPU 内存，实现"GPU+CPU 混合显存"。

第三步：为超大规模训练省出 fp16 主权重

在"极少数 GPU 上训练极大模型"的场景下，官方还引入了一个额外的显存优化手段，需要给 fp16 优化器添加如下配置：

"fp16": { "enabled": true, "fp16_master_weights_and_grads": true, }

其含义是：开启 fp16 混合精度训练，并让优化器不再单独维护一份 fp32 master weights/grads（fp16_master_weights_and_grads: true），从而为超大 MoE 模型省出可观的内存。需要留意的是，该开关牺牲一部分数值精度以换取容量，适合显存极度受限、且对数值稳定性有容忍度的训练场景。

进阶应用：NLG 大模型的 MoE / PR-MoE / MoS

官方教程为 MoE 应用于 NLG（自然语言生成，GPT 类）模型单独维护了一份更完整的实战文档，即本仓库中的 docs/_tutorials/mixture-of-experts-nlg.md，它基于 Megatron-LM 框架的 GPT 风格模型展开。相关核心结论与超参可以快速速览（详细内容请跟随上方仓库内链接阅读原教程）：

• 标准 MoE 新增超参：--num-experts（每层专家数，实验常用 128，更多专家收敛更好但边际递减）、--moe-expert-parallel-size（专家并行度，单卡专家数 =num-experts / moe-expert-parallel-size，其值不得超过 GPU 总数与num-experts）、--moe-loss-coeff（MoE 辅助损失系数，实验常用 0.01）、--moe-train-capacity-factor/--moe-eval-capacity-factor/--moe-min-capacity（单个专家的 token 容量，越大越利于收敛但负载更不均、训练更慢）、--disable-moe-token-dropping（完全移除专家容量限制，因负载不均问题仅推荐推理/评估阶段使用）。
• PR-MoE 差异超参：--num-experts需传与 MoE 层数等长的 list（建议靠近输出层的后段放更多专家）；--mlp-type可选standard/residual，设为residual即启用 Residual-MoE。实践上 NLG+MoE 相比稠密底座还宜降低学习率并拉长学习率衰减期。
• MoS（Mixture-of-Students）模型压缩：通过分阶段知识蒸馏压缩大 MoE 模型，需指定--mos、--load-teacher（教师模型 checkpoint 路径，必备）、--num-layers-teacher/--hidden-size-teacher/--num-experts-teacher（教师模型架构参数）；官方实验采用分阶段蒸馏，即在训练到某一阶段后停掉蒸馏 loss、仅用标准语言建模 loss 继续优化，避免教师模型全程约束损害学生精度。

小结：一条可落地的接入路线

把以上内容收敛为工程上的四个步骤：

1. 选型：确定每个 MoE 层的num_experts与ep_size（保证num_experts % ep_size == 0），按"每卡专家数 = num_experts / ep_size"反推显存与吞吐；
2. 改模型：把目标层替换为deepspeed.moe.layer.MoE(hidden_size=..., expert=..., num_experts=..., ep_size=..., k=1)，注意输入输出同维约束，必要时新增桥接层；构造金字塔（num_experts 传 list）与残差（use_residual=True）即得 PR-MoE；
3. 交引擎：把模型交给deepspeed.initialize（当前引擎会自动划分专家/共享参数组），并按需在ds_config配置 ZeRO stage 与cpu_offload、fp16_master_weights_and_grads；
4. 验证：先跑通仓库测试同款的SimpleMoEModel/SimplePRMoEModel分布式用例（tests/unit/v1/moe/test_moe.py），再迁移到真实模型，观察 gate loss（l_aux）与各专家 token 分布是否均衡。

这样，无论是为已有稠密模型做"稀疏化扩容"，还是在少量 GPU 上冲击超大 MoE 模型，你都有了从 API 到并行拓扑再到显存优化的完整方案。

【免费下载链接】DeepSpeedDeepSpeed is a deep learning optimization library that makes distributed training and inference easy, efficient, and effective.项目地址: https://gitcode.com/GitHub_Trending/de/DeepSpeed