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 MoEPR-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 都涌向少数几个专家的坍缩现象。从当前仓库源码看上述组件对应三个文件deepspeed/moe/layer.py 定义了面向用户的核心 APIdeepspeed.moe.layer.MoEdeepspeed/moe/experts.py 中的Experts容器负责将同一个 expert 模块复制num_local_experts份并给每个专家参数打上allreduce False与group_name标记这是后续按专家/共享分组通信的关键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层。如果你希望避免手改模型仓库还提供了AutoEPAutomatic 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 内存。灵活的进程组设计允许用户混合叠加主流并行技术。下表是官方教程给出的完整并行配置矩阵本文原样保留并加以说明缩写灵活并行配置收益EExpert通过增加专家数量扩展模型规模E DExpert Data扩展到多个数据并行组加速训练吞吐E ZExpert ZeRO 数据并行切分非专家参数支撑更大的底座模型E D MExpert Data Model支撑超大 hidden size底座模型比 EZ 更大E D ZExpert Data ZeRO 数据并行支撑超大 hidden size底座模型比 EDM 更大E Z-Off MExpert ZeRO-Offload Model在有限 GPU 上同时利用 GPU 与 CPU 内存训练超大 MoE 模型进程组创建从手动初始化到按层自动创建为了支持不同形态的并行DeepSpeed 在内部创建了多种进程组。早期的帮助函数位于deepspeed/utils/groups.py即仓库中的 deepspeed/utils/groups.pydeepspeed.utils.groups.initialize(ep_sizedesired 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, ...)纯 ED 场景groups._create_expert_data_and_model_parallel(ep_size, mpu, ...)结合 Tensor/Pipeline/Model 并行的 EDM 场景。每个进程组以ep_size_{ep_size}命名并在多个进程组字典如_EXPERT_PARALLEL_GROUP、_EXPERT_DATA_PARALLEL_GROUP中登记同名的组只创建一次。这里可以引出一个重要的通信拓扑事实注释于 deepspeed/utils/groups.py 的 legacy ED 示例中expert parallel group负责专家间的 All-to-All 分发不做梯度 all-reduceexpert data parallel group只在持有同一批专家的副本之间做专家参数的梯度 all-reducedata 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_size84, expertself.fc3, num_expertsargs.num_experts, ep_sizedesired 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 与构造函数中确认参数及其默认值如下参数默认值说明hidden_size必填模型隐藏维度同时是 MoE 层的输入与输出维度expert必填定义专家子网络的nn.Module如 MLP、nn.Linearnum_experts1该层专家总数传 list 则启用 Pyramid-MoE见下文ep_size1专家并行组大小即参与切分这批专家的 rank 数量k1top-k 门控值仅支持 k1 或 k2高版本也扩展了 topkgatingcapacity_factor1.0训练时每个专家的容量系数eval_capacity_factor1.0推理/评估时每个专家的容量系数min_capacity4忽略 capacity_factor 时每个专家的最小容量use_residualFalse设为 True 则该层变为 Residual MoEPR-MoE层noisy_gate_policyNone噪声门控策略合法取值Jitter、RSample或Nonedrop_tokensTrue是否丢弃超容量 token设为 False 等价于无限容量use_rtsTrue是否启用 Random Token Selection随机 token 选择use_tutelFalse是否使用 Tutel 优化需已安装 tutel 库enable_expert_tensor_parallelismFalse专家内部是否采用张量并行top2_2nd_expert_samplingTrue是否对第 2 个专家做采样Gumbel-max 技巧仅 k2 生效从源码可以看到两个硬性校验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.pycapacity 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 MoEPR-MoE金字塔 残差专家官方在近期工作中提出了Pyramid-Residual MoEPR-MoE架构arXiv:2201.05596目标是显著提升参数效率。要构建 PR-MoE 模型只需在原 API 上多做两件事构造金字塔结构把num_experts传成一个 list例如[4, 8]列表长度与模型中 MoE 层的数目一致不同层承载不同数量的专家启用残差专家设置use_residualTrue把该层标记为 Residual MoE 层。代码形态self.experts deepspeed.moe.layer.MoE(hidden_sizeinput_dim, expertExpertModule(), num_experts[..], ep_sizeep_size, use_residualTrue)注意PR-MoE 混合使用金字塔专家数与残差融合必须配合使用——list 形式的num_experts是为了跨层形成金字塔而use_residualTrue保证每层在门控专家之外还有一个 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_sizeinput_dim, expertExpertModule(), num_expertsEXPERTS, ep_sizeEP_WORLD_SIZE)运行起来后DeepSpeed 运行时将训练一个共 8 个专家、分布到 4 张 GPU、每卡 4 个专家的 MoE 模型——也就是表格中的E DExpert 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_size84, expertfc3, num_expertsEXPERTS, ep_sizeEP_WORLD_SIZE, k1) fc4 torch.nn.Linear(84, 10)这里k1表示每个 token 只路由到 1 个专家top-1 gating。一个可以直接本地阅读的实现事实是num_local_experts num_experts // ep_sizedeepspeed/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 Truedeepspeed/moe/layer.py一路传入TopKGate与top1gating在 deepspeed/moe/sharded_moe.py 中use_rtsTrue时会给 one-hot 路由掩码乘上一个 [0,1) 均匀分布的随机数再沿 token 维度做torch.topk挑选哪些 token 真正进入专家容量槽——相当于让谁被丢弃的决策随机化从而消除系统性偏差。若想关闭可显式传use_rtsFalse。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.pysplit_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( argsargs, modelnet, model_parametersparameters, training_datatrainset)关于自动化需要特别说明一个演进事实当前仓库已经自动完成了上述参数组拆分。在 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(configconfig_dict, modelmodel)并注释should automatically create moe param groups in deepspeed backend即为实证。不过手动构造分组仍然有效且在需要精细控制分组命名/数量时是推荐做法。第二步ZeRO-Offload 的 ds_config要在 cifar 类示例上以ZeRO-Offloadstage 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 内存实现GPUCPU 混合显存。第三步为超大规模训练省出 fp16 主权重在极少数 GPU 上训练极大模型的场景下官方还引入了一个额外的显存优化手段需要给 fp16 优化器添加如下配置fp16: { enabled: true, fp16_master_weights_and_grads: true, }其含义是开启 fp16 混合精度训练并让优化器不再单独维护一份 fp32 master weights/gradsfp16_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-coeffMoE 辅助损失系数实验常用 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。实践上 NLGMoE 相比稠密底座还宜降低学习率并拉长学习率衰减期。MoSMixture-of-Students模型压缩通过分阶段知识蒸馏压缩大 MoE 模型需指定--mos、--load-teacher教师模型 checkpoint 路径必备、--num-layers-teacher/--hidden-size-teacher/--num-experts-teacher教师模型架构参数官方实验采用分阶段蒸馏即在训练到某一阶段后停掉蒸馏 loss、仅用标准语言建模 loss 继续优化避免教师模型全程约束损害学生精度。小结一条可落地的接入路线把以上内容收敛为工程上的四个步骤选型确定每个 MoE 层的num_experts与ep_size保证num_experts % ep_size 0按每卡专家数 num_experts / ep_size反推显存与吞吐改模型把目标层替换为deepspeed.moe.layer.MoE(hidden_size..., expert..., num_experts..., ep_size..., k1)注意输入输出同维约束必要时新增桥接层构造金字塔num_experts 传 list与残差use_residualTrue即得 PR-MoE交引擎把模型交给deepspeed.initialize当前引擎会自动划分专家/共享参数组并按需在ds_config配置 ZeRO stage 与cpu_offload、fp16_master_weights_and_grads验证先跑通仓库测试同款的SimpleMoEModel/SimplePRMoEModel分布式用例tests/unit/v1/moe/test_moe.py再迁移到真实模型观察 gate lossl_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创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考