2023 年 12 月，首个开源 MoE 大模型 Mixtral 8×7B 发布，在多种基准测试中，其表现近乎超越了 GPT-3.5 和 LLaMA 2 70B，而推理开支仅相当于 12B 左右的稠密模型。为进一步提升模型功能，稠密 LLM 常由于其参数规模急剧扩张而面临严峻的训练成本。MoE 的广泛应用，使得在算计成本相对不变的条件下，模型容量可以或许得到显著扩展。此特性无疑使得 MoE 成为推动 LLM 发展的关键技术。
MoE 设计的初衷，是使模型的学习更加 “术业有专攻”，其有效性已得到业界肯定。然而现有 MoE 架构训练中的弊端也逐渐凸显，主要包括：大师负载失衡、大师内样本混杂而大师间同质化现象严重、额外的通讯开支等等。
为了缓解现有 MoE 普遍存在的训练效率与功能瓶颈，专精于高功能算计、LLM 训练加速的华为 GTS AI 算计 Lab的研究团队提出了名为 LocMoE 的全新 MoE 架构，从路由体制角度出发，以期降低稀疏 LLM 训练成本的同时，提升其功能。
论文链接：https://arxiv.org/abs/2401.13920
论文简介
在这项工作中，作家发现之前的 MoE 路由体制往往会导致路由至同一大师收集的 token 差异较大，干扰大师收集的收敛；而路由至不同大师的 token 相似性较高，造成大师间同质化现象严重，最终影响模型语义理解与生成的能力。作家通过理论阐明了大师路由与输入数据特征之间的关联，并在 NLP 规模首次证明大师收集存在容量下限。在此理论基础上，大师路由的门控权重经正交化处理后，明显增强了大师收集间的区分度，处理远小于原先规模的 token，可以或许在规模评测中达到相近的效果。同时该研究针对 MoE 架构中固有的 All-To-All 通讯瓶颈，结合负载 / 通讯优化，提出高效高能的 MoE 架构。
具体来说，作家提出了一种名为 LocMoE 的新颖 MoE 架构，将其嵌入到盘古大模型的骨干中以增强其能力。LocMoE 旨在增强路由体制的可解释性，同时降低额外通讯与算计开支。首先，作家发现 token 总倾向于路由至与该 token 夹角最小的大师，当大师间门控权重向量满足正交时，大师收集间处理的 token 能尽可能避免同质化。
因此，本文采用 GAP 层提取 token 特征，将其作为路由的依据。GAP 层特性上满足门控权重正交的条件，算计量相比 Dense 层也得到大幅下降。基于上述结构，作家通过理论求解出在不影响模型 loss 前提下，大师处理的 token 规模的下限，以降低大师收集的负载。此外，作家结合 auxiliary loss，提出了 locality loss 对路由进行软束缚，降低跨机 All-To-All 通讯开支。最后，采用通讯遮掩等工程优化，进一步提升稀疏大模型整体的训练功能。
作家将 LocMoE 架构嵌入到盘古-Σ 38B 模型中，采用语义相似度较高的 ICT 规模数据进行训练，检验其规模知识的学习能力。在十项下游任务中，LocMoE 的准确性普遍高于原生盘古-Σ，训练功能每步提升 10%~20%。该 MoE 架构还具有较强的通用性和易于移植性，可以或许低成本嵌入到其他硬件规格和其他 MoE 架构的 LLM 骨干中。
当前，LocMoE 已部署至华为 ICT 服务规模专业知识问答应用 AskO3 上，AskO3 已上线华为 O3 知识社区，获得数万工程师用户群体广泛好评。
创新点剖析
路由与数据特征的关系
针对现有 MoE 路由体制普遍缺乏可解释性的问题，作家分析了 token 路由的本质，并设计了可以或许将 token 有效区分的结构。对于某个 token，学习性的路由策略普遍选择门控权重与该 token 乘积更大的大师进行分配：
那么，如果大师的门控权重满足正交，可以或许使得大师具备更高的判别性。同时，可以或许得出 token 倾向于被路由至与其夹角更小的大师：
作家最终选取 GAP 作为提取 token 特征的结构，其门控权重可以或许满足正交的条件：
上述路由体制的实质描述了路由判别能力与大师 token 间最小夹角之间的关联，如图所示。
图：LocMoE 路由体制示意图
大师容量下界
在上述理论的基础上，作家发现，大师容量存在下界，即，在输入数据确定的情况下，大师处理远小于经验值规模的 token，也能达到相当的功能。该问题可以转化为，求解使得至少一个具有类别判别性 token 被路由至某个大师的最小 token 规模，作为所有大师容量拉齐时的下界。同时可以或许得出，合理的大师容量与 token 和门控权重间的最小夹角呈负相关，并随着夹角的减小呈指数级增长。经实验证实，大师容量设为该下界时，未对模型收敛性和 loss 曲线产生影响。
本地性束缚
LocMoE 在 MoE 层的 loss 包含两部分：auxiliary loss 和 locality loss。auxiliary loss 首次在 sparsely-gated MoE 中提出，同时应用于 SwitchTransformer，用以提升大师负载均衡性：
然而，跨机 All-To-All 带来的额外通讯开支仍无法避免。因此，作家添加了本地性束缚，使得在大师负载均衡的前提下，token 更倾向于被分派给本地设备的大师，最终达到束缚平衡。locality loss 采用当前 token 分布与完全本地化分布之间的差异即 KL 散度来量化，从而将部分机间通讯转为机内通讯，充分利用机内互联高带宽。
实验结果
作家分别在包含 64 张、128 张和 256 张昇腾 910A NPU 的集群上进行了实验，主要与两款经典的 MoE 结构：Hash (来自 Facebook) 和 Switch (来自 Google) 进行比较。
训练效率
作家记录了各实验组训练过程中算计、通讯、遮掩以及闲置的耗时。其中，在 64N 和 128N 的情况下，LocMoE 的算计开支和通讯开支都是最低的。尽管 256N 时 LocMoE 的算计开支仍然最低，但部分设备不包含大师使得本地性通讯转换失效，说明了 LocMoE 在算计及通讯方面同时存在显著增益的适用条件是大师数至少大于等于节点数。
图：多种 MoE 结构在不同集群配置下的训练效率
大师负载
为了验证束缚项对于大师负载的影响，作家分析了路由至每个大师的 token 分布情况。为了达到负载均衡，通过 RRE 模块实现的 Hash 路由采用静态路由表的硬束缚确保分配的均衡性，LocMoE 和 Switch 则考虑到 token 的具体特征而进行路由。作为学习型路由，在 auxiliary 和本地性束缚项的作用下，LocMoE 大师的均衡性明显优于 Switch，表现出稳定且较高的资源利用率。
图：多种 MoE 结构下大师负载情况
分配给大师的样本相似性
对于支撑 LocMoE 提出的关键理论，作家采用实验对其进行了验证。左图表明路由到同一大师的 token 相似性更高，接近于 1。右图则表明 token 与其路由至的大师对应的门控权重相似度分布相较其他大师更靠右，从而证实了 token 倾向于路由至与其夹角最小的大师的理论前提，并标记出大师容量下限求解的关键参数 δ。
图：路由至同一 / 不同大师 token 相似性（左）；token 与其路由至的大师的相似性（右）
NLP 规模下游任务
盘古-Σ 已经使用了来自金融、健康等超过 40 个规模的语料进行预训练，证明其从多规模文本数据中学习知识的能力。在本项工作中，作家使用 ICT 服务的规模数据，包含无线收集、光、运营商 IT 等产品线的技术报告和工具手册等，评估 LocMoE 在专业规模知识的学习表现。根据概念间逻辑复杂程度分为 L1 至 L3，梳理出十类 NLP 规模任务的评测数据集。与原生盘古-Σ 相比，LocMoE 使得模型语义理解和表达能力都有一定程度的提高。
图：与原生盘古-Σ 相比，NLP 规模下游任务表现