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

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

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

MoE因其在训推流程中低销高效的特点，近两年在大语言模型领域大放异彩。作为MoE的灵魂，专家如何能够发挥出最大的学习潜能，相关的研究与讨论层出不穷。此前，华为GTSAI计算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）hardrouter，直接把整个token特征进行分派；

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

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

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

该文贡献

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

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

基于理论，作者指出

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

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

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

实验结果

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

训练效率

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

显存占用

LocMoE 对于显存占用也存在一定增益，尤其是集群规模较小计算较密集的场景。使用AscendInsight工具对显存监控样本进行分析，可以看出，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%的效果提升。