大模型分布式训练并行技术（八）-MOE并行

近年来，随着Transformer、MOE架构的提出，使得深度学习模型轻松突破上万亿规模参数，传统的单机单卡模式已经无法满足超大模型进行训练的要求。因此，我们需要基于单机多卡、甚至是多机多卡进行分布式大模型的训练。

而利用AI集群，使深度学习算法更好地从大量数据中高效地训练出性能优良的大模型是分布式机器学习的首要目标。为了实现该目标，一般需要根据硬件资源与数据/模型规模的匹配情况，考虑对计算任务、训练数据和模型进行划分，从而进行分布式存储和分布式训练。因此，分布式训练相关技术值得我们进行深入分析其背后的机理。

前面的文章中讲述了数据并行、流水线并行、张量并行、序列并行、自动并行等多种并行技术。但现在的模型越来越大，训练样本越来越多，每个样本都需要经过模型的全部计算，这就导致了训练成本的平方级增长。而当我们希望在牺牲极少的计算效率的情况下，把模型规模提升上百倍、千倍，通常就需要使用 MOE（Mixture-of-Experts）并行。因此，本文接下来给大家分享 MOE 并行。

MOE

---

通常来讲，模型规模的扩展会导致训练成本显著增加，计算资源的限制成为了大规模密集模型训练的瓶颈。为了解决这个问题，一种基于稀疏 MoE 层的深度学习模型架构被提出，即将大模型拆分成多个小模型(专家，expert)， 每轮迭代根据样本决定激活一部分专家用于计算，达到了节省计算资源的效果； 并引入可训练并确保稀疏性的门( gate )机制，以保证计算能力的优化。

与密集模型不同，MoE 将模型的某一层扩展为多个具有相同结构的专家网络( expert )，并由门( gate )网络决定激活哪些 expert 用于计算，从而实现超大规模稀疏模型的训练。

以下图为例，模型包含 3 个模型层，如(a)到(b)所示，将中间层扩展为具有 n 个 expert 的 MoE 结构，并引入 Gating network 和 Top_k 机制，MoE 细节如下图©所示。

计算过程如下述公式。

上述第 1 个公式表示了包含 n 个专家的 MoE 层的计算过程。具体来讲，首先对样本 x 进行门控计算， W 表示权重矩阵；然后，由 Softmax 处理后获得样本 x 被分配到各个 expert 的权重； 然后，只取前 k (通常取 1 或者 2）个最大权重；最终，整个 MoE Layer 的计算结果就是选中的 k 个专家网络输出的加权和。

MOE 分布式并行策略

---

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

MOE + 数据并行

该策略是在数据并行模式下包含MOE架构，门网络(gate)和专家网络都被复制地放置在各个运算单元上。下图展示了一个有三个专家的两路数据并行MoE模型进行前向计算的方式。

该方式通常来说，对于现有的代码侵入性较小。但该方式唯一的问题是，专家的数量受到单个计算单元(如：GPU)的内存大小限制。

MOE + 模型并行

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

下图展示了一个有六个专家网络的模型被两路专家并行地训练。注意：专家1-3被放置在第一个计算单元上，而专家4-6被放置在第二个计算单元上。

该模式针对不同的模型和设备拓扑需要专门的并行策略，同时会引入额外的通信，因此，相较于数据并行+MOE策略，侵入性更强。

除了上述两种MOE并行方案之外，还可以MOE+数据并行+模型并行、MOE+ZeRO增强的数据并行等。

业界大模型的 MOE 并行方案

---

GShard

GShard 是第一个将 MoE 的思想拓展到 Transformer 上的工作。具体的做法就是把 Transformer 的 encoder 和 decoder 中每隔一个（every other）的FFN层，替换成 position-wise 的 MoE 层，使用的都是 Top-2 gating network。

此处之外，GShard还加入了很多其他设计：

• Expert capacity balancing：强制每个expert处理的tokens数量在一定范围内。
• Local group dispatching：通过把一个batch内所有的tokens分组，来实现并行化计算。
• Auxiliary loss：为了缓解“赢者通吃”问题，尽可能把token均分给各个专家。
• Random routing：在Top-2 gating的设计下，两个expert如何更高效地进行routing。

Switch-Transformer

Switch-Transformer 是在T5模型的基础上加入了 MoE 设计，并在C4数据集上预训练，得到了一个“又快又好”的预训练大模型。

Swith Transformer 简化了MoE的routing算法，从而大大提高了计算效率，具体如下图所示：

Swith Transformer 其设计的指导原则是以一种简单高效的实现方式尽可能地把Transformer模型的参数量做大。跟其他MoE模型的一个显著不同就是，Switch Transformer 的 gating network 每次只 route 到 1 个 expert，而其他的模型都是至少2个。这样就是最稀疏的MoE了，因此单单从MoE layer的计算效率上讲是最高的了。

GLaM

这是 Google 在2021年底推出的一个超大模型，完整的 GLaM 总共有 1.2T 参数，每个 MoE 包含 64 个专家，总共 32 个 MoE 层，但在推理期间，模型只会激活 97B 的参数，占总参数的 8%。

GLaM 的体系架构，每个输入 token 都被动态路由到从 64 个专家网络中选择的两个专家网络中进行预测，如下图所示。

GLaM比GPT-3大7倍，但是由于使用了Sparse MoE的设计，训练成本却只有GPT-3的1/3，并且推理过程中的计算量减少了约一半；同时，在29个NLP任务上超越了GPT-3。

AI 训练框架中的 MOE 并行训练

---

从 Google 发布的很多的论文和超大参数规模模型（千/万亿参数）可以看到，其基本都使用了 MOE 架构。除此之外，业界很多的AI训练框架中也继承了 MOE 并行，比如：PaddlePaddle、DeepSpeed、ColossalAI等。

PaddlePaddle 中的 MOE 并行

下面是一个在动态图模式下使用 PaddlePaddle 框架进行 MoE 架构的适配和训练示例。

1  # 导入需要的包
2  import paddle
3  from paddle.nn import Layer, LayerList, Linear, Dropout
4  from paddle.incubate.distributed.models.moe import MoELayer
5  from paddle.distributed.collective import Group
6  from paddle.distributed import fleet
7  import numpy as np
8
9  # 专家数
10 num_experts = 8
11
12 d_model = 512
13 d_hidden = 2048
14
15
16 # 封装专家层
17 class ExpertLayer(Layer):
18     def __init__(self, d_model, d_hidden, name=None):
19         super().__init__()
20         self.htoh4 = Linear(d_model, d_hidden)
21         self.h4toh = Linear(d_hidden, d_model)
22
23     def forward(self, x):
24         x = self.htoh4(x)
25         x = self.h4toh(x)
26         return x
27
28
29 # 初始化分布式环境，并构建 expert 通信组 moe_group
30 fleet.init(is_collective=True)
31 moe_group = paddle.distributed.new_group(list(range(fleet.worker_num())))
32 
33 
34 gate_config = {
35     "type": "gshard",
36     "top_k": 2,
37 }
38
39
40 experts_list = LayerList()
41 for expi in range(num_experts):
42    exp_layer = ExpertLayer(d_model, d_hidden)
43    experts_list.append(exp_layer)
44
45
46 # 调用 MoELayer API 封装并创建出 MoE 模型
47 class Model(Layer):
48 	def __init__(self, d_model, d_hidden, name=None):
49 	    super().__init__()
50	    self.linear1 = Linear(d_model, d_model)
51	    self.moe_layer = MoELayer(d_model = d_model,
52	                            experts=experts_list,
53	                            gate=gate_config,
54	                            moe_group=moe_group,
55	                            recompute_interval=0)
56
57	    self.linear2 = Linear(d_model, d_model)
58 	    self.dropout = Dropout(p=0.1)
59
60 	def forward(self, x):
61	    x = self.linear1(x)
62	    x = self.moe_layer(x)
63	    x = self.linear2(x)
64	    x = self.dropout(x)
65	    return x
66
67
68 model = Model(d_model, d_hidden)
69 optim = paddle.optimizer.SGD(parameters=model.parameters())
70
71 # 创建数据集，开始训练
72 for step in range(1, 100):
73    x = paddle.rand([4, 256, d_model])
74
75     y = model(x)
76     loss = y.mean()
77     loss.backward()
78     optim.step()
79
80    optim.clear_grad()
81
82     print("=== step : {}, loss : {}".format(step, loss.numpy()))

DeepSpeed 中的 MOE 并行

DeepSpeed中也提供了对 MOE 并行的支持。目前，DeepSpeed MoE 支持五种不同的并行形式，可以同时利用GPU和CPU内存，具体如下表所示。

下面是使用 ZeRO-Offload (stage 2) 和 DeepSpeed MOE组合的样例：

1
2  # MOE 模型架构
3  class Net(nn.Module):
4      def __init__(self):
5          super(Net, self).__init__()
6          self.conv1 = nn.Conv2d(3, 6, 5)
7          self.pool = nn.MaxPool2d(2, 2)
8          self.conv2 = nn.Conv2d(6, 16, 5)
9          self.fc1 = nn.Linear(16 * 5 * 5, 120)
10         self.fc2 = nn.Linear(120, 84)
11         if args.moe:
12             # MoE 层
13             fc3 = nn.Linear(84, 84)
14             self.moe_layer_list = []
15             for n_e in args.num_experts:
16                 # 基于专家数创建 MOE 层
17                 self.moe_layer_list.append(
18                     deepspeed.moe.layer.MoE(
19                         hidden_size=84,
20                         expert=fc3,
21                         num_experts=n_e,
22                         ep_size=args.ep_world_size,
23                         use_residual=args.mlp_type == 'residual',
24                         k=args.top_k,
25                         min_capacity=args.min_capacity,
26                         noisy_gate_policy=args.noisy_gate_policy))
27             self.moe_layer_list = nn.ModuleList(self.moe_layer_list)
28             self.fc4 = nn.Linear(84, 10)
29         else:
30             # 原始模型层
31             self.fc3 = nn.Linear(84, 10)
32
33     def forward(self, x):
34         x = self.pool(F.relu(self.conv1(x)))
35         x = self.pool(F.relu(self.conv2(x)))
36         x = x.view(-1, 16 * 5 * 5)
37         x = F.relu(self.fc1(x))
38         x = F.relu(self.fc2(x))
39         if args.moe:
40             # 将原始 FFN 层替换成 MoE 层
41             for layer in self.moe_layer_list:
42                 x, _, _ = layer(x)
43             x = self.fc4(x)
44         else:
45             x = self.fc3(x)
46         return x
47
48
49 net = Net()
50
51
52 # 组合 ZeRO-Offload (stage 2) 和 DeepSpeed MOE
53 def create_moe_param_groups(model):
54     from deepspeed.moe.utils import split_params_into_different_moe_groups_for_optimizer
55
56     parameters = {
57         'params': [p for p in model.parameters()],
58         'name': 'parameters'
59    }
60 
61     return split_params_into_different_moe_groups_for_optimizer(parameters)
62
63
64 parameters = filter(lambda p: p.requires_grad, net.parameters())
65 if args.moe_param_group:
66     parameters = create_moe_param_groups(net)
67
68
69 ds_config = {
70   "train_batch_size": 16,
71   "steps_per_print": 2000,
72   "optimizer": {
73     "type": "Adam",
74     "params": {
75       "lr": 0.001,
76       "betas": [
77         0.8,
78         0.999
79       ],
80       "eps": 1e-8,
81       "weight_decay": 3e-7
82     }
83   },
84   "scheduler": {
85     "type": "WarmupLR",
86     "params": {
87       "warmup_min_lr": 0,
88       "warmup_max_lr": 0.001,
89       "warmup_num_steps": 1000
90     }
91   },
92   "gradient_clipping": 1.0,
93   "prescale_gradients": False,
94   "bf16": {
95       "enabled": args.dtype == "bf16"
96   },
97   "fp16": {
98       "enabled": args.dtype == "fp16",
99       "fp16_master_weights_and_grads": False,
100      "loss_scale": 0,
101      "loss_scale_window": 500,
102      "hysteresis": 2,
103      "min_loss_scale": 1,
104      "initial_scale_power": 15
105  },
106  "wall_clock_breakdown": False,
107  "zero_optimization": {
108      "stage": args.stage,
109      "allgather_partitions": True,
110      "reduce_scatter": True,
111      "allgather_bucket_size": 50000000,
112      "reduce_bucket_size": 50000000,
113      "overlap_comm": True,
114      "contiguous_gradients": True,
115      "cpu_offload": True
116  }
117}
118
119# 初始化
120 model_engine, optimizer, trainloader, __ = deepspeed.initialize(
121    args=args, model=net, model_parameters=parameters, training_data=trainset, config=ds_config)
122 ...
123

总结

---

本文简要介绍了目前业界的一些 MOE 并行方案。如果说Transformer结构使得模型突破到上亿参数量，那么稀疏 MoE 结构可以在不显著增加计算成本的情况下，使模型参数量进一步突破，达到上千亿、万亿规模。虽然，1990年左右 MOE 的概念就已经出现了；但是可以预见，MOE 将在通往AGI的道路上扮演越来越重要的角色。

码字不易，如果觉得我的文章能够能够给您带来帮助，期待您的点赞收藏加关注~~

---

最后如果您也对AI大模型感兴趣想学习却苦于没有方向👀

小编给自己收藏整理好的学习资料分享出来给大家💖

👉获取方式：

😝有需要的小伙伴，可以保存图片到wx扫描二v码关注免费领取【保证100%免费】🆓

👉AI大模型学习路线汇总👈

大模型学习路线图，整体分为7个大的阶段：（全套教程文末领取哈）

第一阶段： 从大模型系统设计入手，讲解大模型的主要方法；

第二阶段： 在通过大模型提示词工程从Prompts角度入手更好发挥模型的作用；

第三阶段： 大模型平台应用开发借助阿里云PAI平台构建电商领域虚拟试衣系统；

第四阶段： 大模型知识库应用开发以LangChain框架为例，构建物流行业咨询智能问答系统；

第五阶段： 大模型微调开发借助以大健康、新零售、新媒体领域构建适合当前领域大模型；

第六阶段： 以SD多模态大模型为主，搭建了文生图小程序案例；

第七阶段： 以大模型平台应用与开发为主，通过星火大模型，文心大模型等成熟大模型构建大模型行业应用。

👉如何学习AI大模型？👈

作为一名热心肠的互联网老兵，我决定把宝贵的AI知识分享给大家。 至于能学习到多少就看你的学习毅力和能力了 。我已将重要的AI大模型资料包括AI大模型入门学习思维导图、精品AI大模型学习书籍手册、视频教程、实战学习等录播视频免费分享出来。

这份完整版的大模型 AI 学习资料已经上传CSDN，朋友们如果需要可以微信扫描下方CSDN官方认证二维码免费领取【保证100%免费】

一、全套AGI大模型学习路线

AI大模型时代的学习之旅：从基础到前沿，掌握人工智能的核心技能！

二、640套AI大模型报告合集

这套包含640份报告的合集，涵盖了AI大模型的理论研究、技术实现、行业应用等多个方面。无论您是科研人员、工程师，还是对AI大模型感兴趣的爱好者，这套报告合集都将为您提供宝贵的信息和启示。

三、AI大模型经典PDF籍

随着人工智能技术的飞速发展，AI大模型已经成为了当今科技领域的一大热点。这些大型预训练模型，如GPT-3、BERT、XLNet等，以其强大的语言理解和生成能力，正在改变我们对人工智能的认识。 那以下这些PDF籍就是非常不错的学习资源。

四、AI大模型商业化落地方案

作为普通人，入局大模型时代需要持续学习和实践，不断提高自己的技能和认知水平，同时也需要有责任感和伦理意识，为人工智能的健康发展贡献力量。