HAI-LLM MoE：一键实现大模型的混合专家并行训练

混合专家系统（Mixture of Experts, 简写为 MoE）是在神经网络领域发展起来的一种集成学习技术。传统的深度学习模型在训练时，整个网络都会参与计算。随着模型越来越大，训练使用的样本数据越来越多，训练所需的计算量越来越难以承受。而 MoE 可以动态激活部分神经网络，从而实现在不增加计算量的前提下大幅度增加模型的参数量。

之前介绍 GPT-4 的文章中提到，GPT-4 采用 MoE 技术，可以使 LLM 模型轻松突破上万亿规模的参数，通过巧妙的 MoE 并行训练设计，可以让训练这种参数规模的模型成为可能。幻方x深度求索在萤火智算集群上进行大量的实践，研发了一套轻量的 MoE 并行训练接口，整合入了 HAI-LLM 大模型训练框架中。

本文接下来将为大家详细介绍。

关注下方“深度求索”官方公众号立即体验大模型

↓

MoE

混合专家系统（MoE）会将任务分解为若干子任务，每个子任务上训练一个专家模型（Expert Model），并开发一个门控模型（Gating Model）。门控模型依据输入样本来学习使用哪个专家模型，并组合输出结果。尽管该技术最初是在神经网络上用专家和门控模型来描述的，但这种概念可以推广到任何类型的模型。

子任务和专家

一些建模任务非常复杂，因此需要尽可能的把它们划分为子任务来处理。这是解决问题的一种分治方法，可以应用到更广泛的复杂场景问题的解决。例如，对于 LLM 模型可以根据输入问题所在的领域，划分若干个子空间，然后就每个子空间上训练模型，这样就可以获得特定子问题的专家模型。最终，模型整体会根据输入，学习调用哪个专家模型。

混合专家系统

混合专家系统，简称 MoE，是一种集成学习技术。在神经网络中，一些研究者提出了分解输入空间的 MoE 方法，以便每个专家网络根据不同的子空间训练，再由门控网络组合各种专家网络，输出最终结果。下图描述了 MoE 的基础架构。

在 MoE 架构中，一组专家和一个门控相互合作，通过将输入空间划分为一组嵌套的区域来解决非线性监督学习问题。总的来说，MoE 的过程是如下四个步骤：

将任务划分为子任务；
为每个子任务开发专家模型；
使用门控模型来决定使用哪个专家模型；
池化预测和门控模型输出以进行预测。

MoE 分布式并行策略

上面讲述了 MoE 整体结构，下面来讲述含 MoE 架构的模型的分布式并行策略。

上图展示了不同的并行策略的示意图。可以看到，MoE 分布式并行策略是基于数据并行和模型并行的进一步改进。

MoE + 数据并行

该策略是在数据并行策略下包含 MoE 架构，门网络和专家网络都被复制地放置在各个 GPU 上。下图展示了一个案例，其包含三个专家网络，模型进行两路数据并行进行前向计算。

该方式与现有的数据并行方式容易结合，代码改造量小，但专家网络的数量会受到单个 GPU 内存大小限制。

MoE + 模型并行

与数据并行的方式不同，该策略门网络依然是复制地被放在每个 GPU 上，但是专家网络被独立地分别放置在各个GPU 上。因此，需引入额外的通信操作。该策略可以允许更多的专家网络们同时被训练，而其数量限制与 GPU 的数量是正相关的。

下图展示了一个案例，其包含六个专家网络，模型被两路模型并行进行前向计算。这里，专家1-3被放置在第一个 GPU 上，而专家4-6被放置在第二个 GPU 上。

该策略可以支持更多的专家网络训练，但针对不同的模型和设备拓扑，需要定制专门的并行策略，因此，相较于 MoE + 数据并行的方式，代码改造量会更大。

除了上述两种 MoE 并行方案之外，还可以 MoE+数据并行+模型并行、MoE+ZeRO 增强的数据并行等。原理类似，可以参考之前的文章《HAI-LLM：高效且轻量的大模型训练工具》。

HAI-LLM 上的使用

HAI-LLM 提供了一个简单易用的接口，用户只需在训练配置（cfg）中加入 MoE 结构的设定，即可构建具有 MoE 并行训练能力的模型。配置案例如下：

from hai_llm.model.moe.llama import MoeLLaMABlockfrom hai_llm.model.moe.gate import GShardGate, SwitchGate, HashGate
gpt = dict(    ...    moe=dict(        n_expert=16,        capacity_factor=(1.2, 1.5),  # (train, val)        alpha=0.01,        top_k=2,        gate=GShardGate,        log_auxloss=True,    ),)
parallel = dict(    pp_size=2,    tp_size=2,    activation_checkpoint=False,)
...
model_builder = dict(    ...    moe_block=MoeLLaMABlock,    ...)

可以看到，HAI LLM 中提供了 MoeLLaMABlock 与三种门结构的实现，通过在原来的模型定义 gpt 和模型构建 model_builder 中加入对应的组件即可。这里要注意的是，因为不同门结构的特性，GShardGate 支持 top_k == 2 而 SwitchGate 和 HashGate 只支持 top_k == 1 。

MoeLLaMABlock 的前向传播过程主要包含 self attention 和 moe mlp 两个环节：

class MoeLLaMABlock(nn.Module):    ...        def forward(self, x, **kwargs):        residual = None        if isinstance(x, tuple):            x, residual = x            ### attention block        #   ┌────────────────────────────────────────────────────────────┐        #   │                                                            ▼        #   x ──► norm1 ──► (f ──► Wqkv) -> attn -> (out_proj ──► g) ──► + ──► y        y, residual = self.norm1(x, residual)        y = self.self_attn(y)            ### mlp block        #   ┌─────────────────────────────────┐        #   │                                 ▼        #   y ──► norm2 ──► gate ──► MoEs ──► + ──► z        z, residual = self.norm2(y, residual)        z = self.moe_mlp(z, **kwargs)            return (z, residual) if self.return_residual else (z + residual)

总结

本文简要介绍了 MoE 并行方案和基于 HAI-LLM 的训练实践。如果说 Transformer 结构使得模型可以突破到上亿的参数量，那么与 MoE 结合则可以在不显著增加计算成本的情况下，使模型参数量进一步突破到上千亿、万亿的规模。虽然，MoE 的概念很早就出现了，但与大模型的结合，相信 MoE 会扮演越来越重要的角色。

END

High-Flyer AI

如果你想参与构建下一代通用人工智能（AGI），那么请不要犹豫，加入 DeepSeek，与我们共同在 AGI 征程上“深度求索”吧！点击下方「 阅读原文 」