人工智能技术的迅猛进步带来了大模型参数量的急剧增长,MoE稀疏大模型因其独特的优势——在提升模型容量的同时保持训练计算量,成为学术界和工业界的热门话题。然而,推理性能上的瓶颈一直是制约其广泛应用的关键因素。面对这一挑战,中国科学技术大学人工智能与数据科学学院的张燕咏院长领导的研究团队,在鲲鹏昇腾科教创新卓越中心的算力支撑下,取得了一项重要突破。
该研究团队聚焦于MoE稀疏大模型推理时的两大难题:专家负载不均衡和计算通信开销过大。他们提出了一种创新的解决方案,即基于专家选择分布预测的专家负载均衡和all2all通信均衡框架。这一框架通过细致的分析和预测,实现了推理性能的优化。
首先,团队深入研究了专家选择的规律,通过分析相似token在MoE层选择专家的特征,成功预测了token的专家选择路径偏好。这一步骤为后续的负载均衡打下了坚实的基础。
紧接着,团队着手解决专家负载不均衡的问题。他们利用预测结果,精准地判断出专家动态访问的频率和重要性,通过复制高频专家和替换非重要专家的策略,实现了负载的均衡分布。在这个过程中,昇腾推理引擎MindIE的性能分析工具发挥了关键作用,它提供了实时的专家负载监测和访问频率统计数据,为负载平衡策略的制定提供了有力支持。
最后,团队致力于降低分布式计算的通信时延。他们根据计算节点上的数据对专家选择的预测结果,将专家提前部署到对应的节点上,并将需要紧密协作的专家放在同一计算单元内。这一策略使得数据与所需专家在物理位置上更加接近,从而减少了通信开销。同时,结合昇腾384超节点的高速总线互联技术,团队进一步降低了跨节点和跨计算单元的通信消耗。
这一系列的技术突破带来了显著的成果。在多种尺寸和结构的MoE模型中,该创新框架使得推理时间、all2all时间、MoE层时间及负载不均分数相较于现有主流方案提升了超过30%。在多卡复杂场景中,前三项指标同样提升了30%,而推理时间更是提升了20%。
此次项目的成功不仅为开发者提供了宝贵的经验借鉴,也加速了MoE稀疏大模型在各领域的落地应用。中国科学技术大学鲲鹏昇腾科教创新卓越中心将继续深化产学研协同创新,依托昇腾软硬件平台,在前沿AI模型优化领域不断探索,为自主创新人工智能战略和数字经济的发展贡献力量。