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卡型号，集群组网方案）、训练框架和骨干模型，以证明该系列工作的高扩展性和易于移植性。（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% 的效果提升。
