华为GTS LocMoE+：高可扩展性亲和度 MoE 架构，低开销实现主动路由

AIxiv专栏是机器之心发布学术、技术内容的栏目。过去数年，机器之心AIxiv专栏接收报道了2000多篇内容，覆盖全球各大高校与企业的顶级实验室，有效促进了学术交流与传播。如果您有优秀的工作想要分享，欢迎投稿或者联系报道。投稿邮箱：liyazhou@jiqizhixin.com；zhaoyunfeng@jiqizhixin.com

本论文共同一作为李婧博士，孙志杰和林大超博士，主要成员来自GTS AI计算Lab，主要研究及落地领域包含LLM训推加速、AI训练保障和图计算。

MoE 因其在训推流程中低销高效的特点，近两年在大语言模型领域大放异彩。作为 MoE 的灵魂，专家如何能够发挥出最大的学习潜能，相关的研究与讨论层出不穷。此前，华为 GTS AI 计算 Lab 的研究团队提出了 LocMoE ，包括新颖的路由网络结构、辅助降低通信开销的本地性 loss 等，引发了广泛关注。

LocMoE 的上述设计有效缓解了部分 MoE 经典结构在训练中的瓶颈，如: 专家路由算法可能无法有效区分 token，以及通信同步效率受限于节点内和节点间的传输带宽差异等。此外，LocMoE 证明并求解出在能够成功处理判别性 token 的前提下专家容量的下限。该下限是根据 token 被动派发给专家的场景下，判别性 token 存在于 token 批次中的概率分布而得出。那么，如果专家同样具备择优 token 的能力，判别性 token 被处理的概率将大幅提高，专家容量下限将能进一步压缩。

在上述设想的基础上，该团队进一步提出了基于低开销主动路由的 MoE 架构，将其命名为 LocMoE+。LocMoE+ 继承了 LocMoE 高判别性专家以及本地性通信优势，进一步改造路由策略，定义 token 与专家之间的亲和度指标，并以此指标入手，更有效率完成 token 分派，进而提升训练效率。

论文链接：https://arxiv.org/pdf/2406.00023

论文简介

该论文的核心思路是通过结合传统被动路由与专家主动路由，提升一定容量下处理判别性 token 的概率，从而降低样本噪声和提升训练效率。该论文从 token 和其被分派的专家之间的关系入手，以一种低计算开销的方案量化和定义了专家与 token 之间的亲和性。据此，该论文实现了全局的自适应路由策略，并基于亲和性得分以专家维度对 token 进行重排和选择。同时，专家容量的下限被证明随着 token 特征分布的稳定而逐渐减小，训练开销得以降低。

该论文率先结合两种路由机制，根据学习性路由策略中 token 倾向于路由至与其夹角更小的专家的发现，打破了主动路由现存方案开销过大影响训练效率的障碍，并与被动路由的本质保持统一。

值得一提的是，作者选用了和 LocMoE 工作完全不同的硬件环境（服务器型号，NPU卡型号，集群组网方案）、训练框架和骨干模型，以证明该系列工作的高扩展性和易于移植性。

自适应双向路由分派机制

背景介绍

传统 MoE 有两种路由分派机制：

（1）hard router，直接把整个 token 特征进行分派；

（2）soft router，把 token 特征的加权组合进行分派。

本文延续考虑（1）因其较低的计算代价。而对于 hard router 场景，又可分为 1）Token Choice Router (TCR)， 也就是让每个 token 去选 top-k 专家；2）Expert Choice Router (ECR)，也就是让每个专家去选择 top-C 合适的 token。由于容量限制，每个专家接收的 token 数有上限 C， 因此在场景 1）中，会对每个专家接收的 token进行截断：

前人工作指出 MoE 训练分为两阶段：Phase 1. 路由训练，保障路由可以合理分派 token，也就是能把不同领域或者差异大的 token 进行区分和分派给不同专家；Phase 2. 由于 token 路由的作用，每个专家接收到相同领域或者性质相近的 token，每个专家在经历一定训练后就可以习得相关领域和性质的知识。总结来说，MoE 每步训练“成功”的关键在于 token 分派的正确和合理性。

该文贡献

（1）通过 softmax 激活函数推演，专家和 token 之间的余弦相似度能够较为准确地衡量亲和度：

（2）从理论建模角度，分析 TCR 和 ECR 两个常见场景下单次训练成功率：

基于理论，作者指出

在模型训练初期，路由分派 token 能力不足时，每次训练 TCR 要比 ECR 以更高概率成功训练，而且需要较大的专家容量保证选到合适的 token。

在模型训练后期，路由有一定能力正确分派 token 时，每次训练 ECR 要比 TCR 以更高概率成功训练，此时只需要较小的容量就可以选到合适的 token。

该理论也很符合直观，路由没有分派能力时，不如让 token 随机选择专家，而当路由有一定分派能力，也就是专家能选到合适 token 的时候，使用 ECR 更合适。因此作者推荐 TCR 向 ECR 方式的转变，提出了全局级别的自适应路由切换策略，同时基于专家容量的需求估计，在训练后期使用更小的专家容量。

实验结果

该论文的实验在昇腾 910B3 NPU 的自建集群上进行，得益于华为专有的高速缓存一致性协议高性能计算系统（HCCS）在多卡场景下实现的高性能设备间数据通信，以及专为昇腾处理器设计的华为集合通信库（HCCL）在 HCCS 等高速链路上实现高性能分布式训练。实验采用兼容昇腾 NPU 的 PyTorch for Ascend 框架和专为昇腾设备定制的加速库 AscendSpeed 和训练框架 ModelLink，关注 LLM 并行策略和通信遮掩优化。

训练效率

实验结果表明，在不影响模型训练收敛性或有效性的情况下，每个专家至少需要处理的token数量较基线可以减少60% 以上。结合通信优化，在 32 卡、64 卡和 256 卡的集群规模下，训练效率平均提高 5.4% 至 46.6%。

显存占用

LocMoE+ 对于显存占用也存在一定增益，尤其是集群规模较小计算较密集的场景。使用 Ascend Insight 工具对显存监控样本进行分析，可以看出，LocMoE+ 显存使用相比于基线下降了 4.57% 至 16.27%，相比于 LocMoE 下降了 2.86% 至 10.5%。随着集群规模的增大，显存占用的差距随之缩小。

效果评测

开源评测集 C-Eval 和 TeleQnA，以及自主构建的 ICT 领域评测集 GDAD 被用于评估 LocMoE+ 在通用知识和领域知识上的能力。其中，GDAD 共计覆盖了 47 个子项，包含 18060 条样本，考察模型在领域任务，领域能力认证考试和通用能力三大评估体系中的表现。

经过充分 SFT 后，LocMoE+ 相比于基线在领域任务能力的 16 个子能力上平均提升了约 20.1%，相比于 LocMoE 则提升了约3.5%。领域能力认证考试则分别提升了 16% 和 4.8%。在通用能力的 18 个子能力中，LocMoE+ 分别提升约 13.9% 和 4.8%。整体而言，LocMoE+ 在 GDAD、C-Eval 和 TeleQnA 上分别表现出 9.7% 到 14.1% 的效果提升。