在 Colossal-AI 中为你的模型集成混合专家（MoE）：并行配置、MoE 层构建与训练实战

说明：本文章基于仓库文档 docs/source/en/advanced_tutorials/integrate_mixture_of_experts_into_your_model.md（另有中文版 docs/source/zh-Hans/advanced_tutorials/integrate_mixture_of_experts_into_your_model.md）编写，用于指导开发者把 MoE 架构接入自己的模型，并在 Colossal-AI 的模型并行与数据并行框架下训练。

导读

本文是一篇面向模型开发者的 MoE（Mixture of Experts，混合专家）接入实战指南。以 GShard、Switch Transformer 开启的稀疏化专家范式为背景，讲解如何在一个已有 Transformer 模型中用 Colossal-AI 提供的 MoE 并行能力完成三件事：在 config.py 中声明 MoE 模型并行组、用 Experts/Router/MoeLayer 等组件搭建稀疏 FFN 层、以及借助 colossalai.initialize 与 MoeGradientHandler/MoeLoss 正确完成反向传播与辅助损失训练。读完本文，你将掌握一套可在 Colossal-AI 项目中直接套用的 MoE 接入代码模板，并理解“MoE 模型并行组”与“MoE 数据并行组”在专家分布式存储下的划分逻辑。

需要提前说明的是：本教程对应的 API 属于 Colossal-AI 的早期（legacy）版本，在当前的仓库快照中，相关实现已归档于 colossalai/legacy 目录下；同时仓库中的 MoE 层也演化出了功能更完整的 SparseMLP 实现。下文会同时给出两者之间的对应关系，方便你在真实代码中定位。

为什么在 Colossal-AI 中集成 MoE：动机与适用边界

自从 Switch Transformer 问世以来，AI 社区发现 MoE 是一种在不显著增加每个 token 计算量的前提下扩大模型容量的有效手段：模型由一组并行的“专家”子网络构成，输入 token 通过一个可学习的 router 被路由到少数几个专家上执行，从而实现参数规模的稀疏激活。

Colossal-AI 为此提供了面向 MoE 模型的并行支持（早期访问版本），其最突出的卖点是便利性：目标就是让用户能轻易地把 MoE 与已有的模型并行（model parallelism）和数据并行（data parallelism）结合起来使用。在早期版本中，该实现存在两个已知短板（文档原文明确声明）：

1. 在大批量（large batch size）与长序列（long sequence length）训练下效率不佳；
2. 与张量并行（tensor parallelism）暂时不兼容。

文档同时说明团队正在通过系统优化解决训练效率问题，而张量并行兼容性需要更多适配工作。从当前仓库源码看，MoE 层实现已进一步演进——例如 colossalai/legacy/moe/layer/layers.py 中的 SparseMLP 已把并行模式扩展为支持 "EP"、"TP" 与 None（见其 forward 中对 parallel 的分支处理），并引入了 kernel 优化、通信重叠等开关，可作为你评估当前能力边界时的参考。

第一步：搭建 MoE 运行环境（并行配置）

教程的第一个核心操作是在项目根目录创建 config.py，用于声明训练时希望启用的特性。要启用 MoE，必须添加一个名为 parallel 的 dict，并在其中设置 moe 键的值；moe.size 即专家的模型并行大小，代表一个并行组内用来并行化训练的专家数量。

MOE_MODEL_PARALLEL_SIZE = ...
parallel = dict(
    moe=dict(size=MOE_MODEL_PARALLEL_SIZE)
)

对 size 的直观理解（沿用教程原文）：假设 size = 4，那么 4 个进程会被分配到 4 块连续 GPU 上，这 4 个进程组成一个 MoE 模型并行组；组内每块 GPU 只持有全部专家中的一部分（局部专家）。调大模型并行大小的代价权衡是：通信开销下降，但单卡计算量上升、激活值显存占用增加。

另外两点需要记住的默认行为：

• size 越小，说明单组内专家分得越散；示例中若 MOE_MODEL_PARALLEL_SIZE = E 且专家总数也设为 E，则一个 Transformer Encoder 在单个模型并行组内做前向时，token 先经 router 打分、再被分发给组内各 GPU 上的专家子网处理、最后聚合回输出，这正是 GShard 论文所描述的“MoE Transformer”数据流形态（原文配图即引自 GShard 论文）。
• 总的数据并行大小会被自动检测，默认等于 GPU 总数，无需手动声明。

代码落地参考：parallel 配置最终会被解析并驱动专家在各 GPU 上的划分。当前仓库中这一逻辑位于 colossalai/legacy/moe/manager.py，其中 MoeManager（模块内以单例 MOE_MANAGER 形式暴露）提供 get_info(num_experts, ...) 来计算每个进程上的局部专家数并推导出相应的数据并行组与专家并行组，同时还负责汇总 router 产生的辅助损失（add_loss/get_loss，见 manager.py 中 router_aux_loss、router_z_loss 相关代码）。

进程组的再划分：MoE 模型并行组 vs. MoE 数据并行组

由于一个 MoE 模型并行组内，所有专家被分配到多张 GPU 上、每张卡只拥有部分专家，因此原有数据并行组在反向传播阶段对专家参数做梯度归约时不再正确。为此 Colossal-AI 引入了一种新并行组——MoE 数据并行组（moe data parallel group）。

教程给出了区分示例：WORLD_SIZE=4、MOE_MODEL_PARALLEL_SIZE=2 时，常规数据并行组把 4 个进程当作同一份模型参数的 4 份拷贝；而 MoE 下的划分会形成两个相互正交的组：每 2 个进程组成一个 MoE 模型并行组（共同负责完整的一组专家），不同 MoE 模型并行组之间则构成用于数据并行复制与专家梯度同步的 MoE 数据并行组关系（原文配图“MoE process group”即为此示意，此处以文字复述其含义）。

梯度处理与运行环境变量

就梯度处理而言，Colossal-AI 提供了 MoeGradientHandler 来对模型每一个参数做 all-reduce。使用方式分两种：

• 若你通过 colossalai.initialize 函数创建训练引擎，MoE 梯度处理器会被自动挂载到引擎上，无需手动处理；
• 否则，梯度归约需要你自己负责。

当前仓库中该 Handler 的实现位于 colossalai/legacy/engine/gradient_handler/_moe_gradient_handler.py（class MoeGradientHandler(BaseGradientHandler)），其文档字符串明确说明：它在数据并行组与 MoE 模型并行之间执行 all-reduce 集体通信来归约梯度。与之配套的单元测试在 tests/test_legacy/test_moe/test_grad_handler.py。

MoE 运行环境的所有参数都存放在 colossalai.global_variables.moe_env 中，训练前可以读取其中的配置参数来核对你的设置是否正确：

from colossalai.global_variables import moe_env

第二步：创建 MoE 层（Experts + Router + MoeLayer）

MoE 层的构件可以从 colossalai.nn.moe（在教程写作时的 API 命名空间）获取；在当前仓库快照中这些模块已归档到 colossalai/legacy/moe 目录（详见 colossalai/legacy/moe/layer/init.py，其从 experts.py、layers.py、routers.py 统一导出），噪声生成器等相关辅助类则集中在 colossalai/legacy/moe/utils.py。阅读源码时可按下表定位对应实现：

教程中的角色	当前仓库中的实现位置	说明
专家集合 Experts	colossalai/legacy/moe/layer/experts.py	各 GPU 上局部专家参数的容器与计算
专家 MLP VanillaFFN 类	colossalai/shardformer/layer/moe 下的 MLPExperts 等	单专家子网络权重 wi/wo（可选 gated）
路由器 Top1Router/Top2Router	MoE 路由相关代码	top-1 / top-2 专家选择与容量控制
噪声器 NormalNoiseGenerator 等	colossalai/legacy/moe/utils.py	NormalNoiseGenerator、UniformNoiseGenerator 及 get_noise_generator 工厂
整体封装 MoeLayer	已演进为 colossalai/legacy/moe/layer/layers.py 的 SparseMLP	门控 + 路由 + 分发/聚合 + 专家计算

先为所有进程设置随机种子

在创建 MoE 层之前，应当像下面这样为所有进程统一设置随机种子：

from colossalai.context.random import moe_set_seed
from model_zoo.moe.models import Widenet

moe_set_seed(42)
model = Widenet(num_experts=4, capacity_factor=1.2)

moe_set_seed 的作用是：为 MoE 模型并行组内的不同进程设置不同的种子，从而帮助专家参数在初始化时就彼此不同（这正是“每个专家都应该长得不一样”的稀疏模型初始化需求）。该辅助函数在当前仓库中实现于 colossalai/legacy/context/random/_helper.py。这里 Widenet（即“Go Wider Instead of Deeper”论文提出的基于 MoE 宽化模型的示例）来自教程配套的外部示例工程，你完全可以用自己的 Transformer 模型替换——MoE 层通常被用来替换标准 Transformer 中的 FFN 子层。

组装专家实例与路由器实例

教程给出了在模型 zoo 中组装 MoE 层的示例代码，模式是“先建 router，再建 experts，最后包一层 MoeLayer”：

from colossalai.nn.layer.moe import Experts, MoeLayer, Top2Router, NormalNoiseGenerator


noisy_func = NormalNoiseGenerator(num_experts)
shared_router = Top2Router(capacity_factor,
                           noisy_func=noisy_func)
shared_experts = Experts(expert=VanillaFFN,
                         num_experts=num_experts,
                         **moe_mlp_args(
                             d_model=d_model,
                             d_ff=d_ff,
                             drop_rate=drop_rate
                         ))
ffn = MoeLayer(dim_model=d_model, num_experts=num_experts,
               router=shared_router, experts=shared_experts)

对这段代码做逐块拆解：

• NormalNoiseGenerator(num_experts)：为 router 的 logits 叠加一个按专家数归一化的高斯噪声。Switch/ViT-MoE 一类的做法是在训练阶段给门控打分加噪声，以促进专家负载均衡、避免路由坍缩到少数专家。当前仓库的 get_noise_generator 工厂（colossalai/legacy/moe/utils.py）会依据 noise_type 返回 UniformNoiseGenerator 或 NormalNoiseGenerator(num_experts) 实例，对应 Jitter 与 Gaussian 两种策略。
• Top2Router(capacity_factor, noisy_func=...)：top-2 路由器，即每个 token 被分给得分最高的两个专家。capacity_factor 用于控制每个专家最多能接收的 token 数量（容量），超过上限的 token 会被丢弃；它直接决定了批次内计算的稀疏度上限。
• Experts(expert=VanillaFFN, num_experts=..., **expert_init_args)：专家容器。每个 GPU 上应持有的局部专家数会在 Experts 初始化内部被自动计算，你只需要指定每个专家所属的类（如 VanillaFFN）以及用于构造该专家类的参数（例如隐藏维 d_model、FFN 中间维 d_ff、dropout drop_rate）。自动切分局部专家数目的逻辑可追溯至 MOE_MANAGER.get_info(...)（colossalai/legacy/moe/manager.py），它依据你配置的 parallel.moe.size 推导出局部专家数。
• MoeLayer(...)：对外的完整稀疏层。创建完 experts 与 router 之后，MoeLayer 内部真正需要初始化的只剩 gate（门控）模块，即那张 num_experts × hidden_size 的门控权重矩阵（在现代 SparseMLP 中对应 self.gate_weight，使用 std=sqrt(0.1/hidden_size) 的高斯初始化，见 layers.py）。

从当前实现看，一个 MoE 稀疏层的完整前向数据流大致是（SparseMLP.forward）：

1. 把输入 (batch, seq, hidden) reshape 成扁平 token (batch*seq, hidden)；
2. 门控线性层产出 gate_logits，并在送入 router 前转成 fp32 以保证路由精度；
3. router 依据 logits（含噪声与 top-k 选择、容量因子限制）输出每个 token 的 dispatch 掩码与 combine 权重；
4. 通过 MoeDispatch（或掩码矩阵乘）把 token 打包成 (num_experts, capacity, hidden)；
5. 在 EP 模式下经 AllToAll / 分层 HierarchicalAllToAll 把数据送到持有对应专家的 GPU，专家 MLP 计算后 AllToAll 送回；TP 模式则走 AllGather + ReduceScatter（参见 _ep_process/_tp_process，通信算子实现在 colossalai/moe/_operation.py）；
6. 用 combine 权重把各专家输出加权聚合回 (batch, seq, hidden)。

注意：教程示例中的 Top2Router、VanillaFFN、moe_mlp_args 等名称来自其配套示例工程，正式使用前请先确认当前所安装版本对应的导入路径。若使用仓库中现成的 SparseMLP（位于 colossalai/legacy/moe/layer/layers.py），则更完整的可配置参数包括：num_experts、hidden_size、intermediate_size、router_top_k、parallel（"EP"/"TP"/None）、router_loss、router_norm、训练/评估期的容量因子 router_capacity_factor_train（默认 1.25）/router_capacity_factor_eval（默认 2.0）、router_min_capacity（默认 4）、router_noisy_policy、router_drop_tks、mlp_gated，以及 enable_load_balance（配套 colossalai/legacy/moe/load_balance.py 中的 LoadBalancer 实现专家负载均衡）、enable_kernel、enable_comm_overlap 等高级开关，可作为你 DIY 稀疏层时对照文档语义的权威参考。

第三步：训练你的 MoE 模型

用 colossalai.initialize 挂载 MoE 梯度处理器

不要忘记使用 colossalai 提供的 colossalai.initialize 函数——只有通过它创建训练引擎，梯度处理器才会被自动添加。在 colossalai.initialize 内部会自动创建一个 MoeGradientHandler 对象来统一处理 MoE 模型的反向传播，即它会遍历模型所有参数执行正确的 all-reduce 归约（其中既包含普通张量并行/数据并行语义下的参数，也包含跨 MoE 模型并行组的专家参数）。前文已给出其当前实现文件 colossalai/legacy/engine/gradient_handler/_moe_gradient_handler.py。若你绕开 initialize 自行管理训练循环，则梯度处理必须自己负责，否则专家参数在并行组间的梯度会不完整。

用 MoeLoss 包装损失函数以叠加辅助损失

为了把 router 的辅助损失（auxiliary loss）纳入整体训练目标，损失函数应当用 MoeLoss 包装。教程给出如下示例：

criterion = MoeLoss(
    aux_weight=0.01,
    loss_fn=nn.CrossEntropyLoss,
    label_smoothing=0.1
)

参数含义：

• aux_weight：路由辅助损失在总损失中的权重。router 训练若不加以约束，容易陷入“赢者通吃”的负载不均衡；给每个专家分配负载的均匀性损失加上一个小的权重（示例取 0.01）可以抑制这种坍缩。
• loss_fn：底层的主任务损失类，这里直接复用 PyTorch 的 nn.CrossEntropyLoss（传入类而非实例，由 MoeLoss 在内部实例化）。
• label_smoothing：交叉熵的标签平滑系数（示例取 0.1），可缓解过度自信并提升泛化。

关于辅助损失的收集与存放，当前仓库的机制同样在 colossalai/legacy/moe/manager.py 中：MoeManager 维护 router_aux_loss 与 router_z_loss 两个累加列表，训练过程中 router 产生的辅助损失通过 add_loss 登记、经 get_loss 取出。更完整的工程化示例可参考仓库内的开源 MoE 模型训练配置 colossalai/legacy/moe/openmoe/README.md（其中 router_aux_loss_factor、router_z_loss_factor、label_smoothing 等命令行参数与本教程的 aux_weight/label_smoothing 语义一致，且额外支持 use_kernel、use_layernorm_kernel 等 kernel 优化开关）。

启动训练

最后，直接使用 colossalai 提供的 trainer 或 engine 进行训练即可——前提依然是经由 colossalai.initialize 完成初始化；否则梯度需要自行处理。教程源文件末尾记录了它的文档级测试命令，单卡环境下可直接用 PyTorch 原生启动器运行一份完整的 MoE 接入脚本：

torchrun --standalone --nproc_per_node=1 integrate_mixture_of_experts_into_your_model.py

多卡时按你的并行规模调整 --nproc_per_node，并结合 config.py 中的 parallel.moe.size 一起规划进程数与专家数的匹配关系（例如将专家总数设为 MOE_MODEL_PARALLEL_SIZE 的整数倍，使每个 MoE 模型并行组内的专家能够被均匀切分到各 GPU）。

仓库内的延伸阅读与验证材料

• 中文对照版教程：docs/source/zh-Hans/advanced_tutorials/integrate_mixture_of_experts_into_your_model.md
• MoE 组件实现：colossalai/legacy/moe/layer/layers.py、colossalai/legacy/moe/layer/experts.py、colossalai/legacy/moe/layer/routers.py
• 并行组与局部专家管理、辅助损失收集：colossalai/legacy/moe/manager.py
• 负载均衡器：colossalai/legacy/moe/load_balance.py
• MoE 通信算子（AllToAll/AllGather/MoeDispatch 等）：colossalai/moe/_operation.py
• 专家张量的并行信息标记（EP/DP 组查询）：colossalai/tensor/moe_tensor/api.py
• MoE 相关单元测试：tests/test_legacy/test_moe（含梯度处理器测试 test_grad_handler.py、负载均衡测试 test_moe_load_balance.py 等）

结语与注意事项

把 MoE 接入自有模型并不需要你手写 all-to-all 与梯度归约：配置好 parallel.moe.size，用 Experts/Router 组装稀疏 FFN 并替换模型中的稠密 FFN，再经由 colossalai.initialize + MoeLoss 训练即可。最后提醒三点：其一，教程 API 属于早期命名空间，实际导入路径请以你所安装的 Colossal-AI 版本为准，当前仓库归档位置在 colossalai/legacy/moe；其二，路由辅助损失权重、容量因子、专家数与 moe.size 的匹配关系是决定收敛质量与显存/通信开销的关键超参，建议从教程示例值（aux_weight=0.01、capacity_factor≈1.2~1.25）出发做小规模实验；其三，早期版本在大 batch、长序列下的效率短板及与张量并行的兼容性限制已在教程中明确声明，部署到大规模训练前应结合当前版本的 SparseMLP 能力与官方变更记录做评估。