随着人工智能技术的飞速发展，大模型参数量激增。MoE 稀疏大模型因能在提升容量的同时保持训练计算量，成为了当下的研究热点，但推理时的性能瓶颈仍制约着其进一步发展。在此背景下，中国科学技术大学人工智能与数据科学学院执行院长张燕咏率领团队，在中国科学技术大学 鲲鹏昇腾科教创新卓越中心的算力支持下，成功研发出了基于专家选择分布预测的专家负载均衡和all2all通信均衡框架。
该框架的研究致力于解决两大难题——专家负载不均及较大的计算通信开销，并通过“三步走”成功实现了推理性能的优化。为了让模型在并行推理时能更合理地分配资源，团队首先从专家选择的规律入手，通过分析相似 token 在 MoE 层选择专家的特征并归类，成功预测出token的专家选择路径偏好。
接着便是解决推理过程中的专家负载不均衡问题。团队依据token对专家选择偏好的预测结果，精准判断出专家动态访问的频率和重要性，通过复制高频专家、替换非重要专家的方式实现了负载平衡。其中昇腾推理引擎MindIE提供的性能分析工具可实现性能数据的可视化，在实时监测专家负载、统计访问频率等方面提供数据参考，发挥了重要作用。
降低分布式计算的通信时延是第三步。团队根据计算节点上的数据对专家选择的预测结果，将专家提前部署到对应的节点上，并联合需要紧密协作的专家放在同一计算单元，让数据与所需专家在物理位置上更接近，同时结合昇腾384超节点采用的高速总线互联技术，最终大幅降低了跨节点和跨计算单元的通信消耗。
在系列技术突破下，此次项目成果显著。基于该创新框架，在多尺寸、多结构的 MoE 模型中，推理时间、all2all 时间、MoE 层时间及负载不均分数较现有主流方案提升超 30%；多卡复杂场景里，前三项指标提升 30%，推理时间提升 20%。
本次项目的成果，为开发者借助昇腾技术优化 MoE 稀疏大模型推理性能提供了可贵借鉴，将加速其在各领域的落地应用。未来，中国科学技术大学 鲲鹏昇腾科教创新卓越中心将持续深化产学研协同创新，依托昇腾软硬件平台，在前沿AI模型优化领域持续突破，为自主创新人工智能战略与数字经济发展贡献力量。
