人工智能技术的日新月异,正引领着大模型参数规模的不断攀升。MoE(混合专家网络)稀疏大模型以其独特的优势,即在增强模型容量的同时有效控制训练计算量,成为了科研界的新宠。然而,推理性能上的瓶颈,如同一道无形的枷锁,束缚着其进一步的发展步伐。在此关键时刻,中国科学技术大学人工智能与数据科学学院执行院长张燕咏带领的团队,携手中国科学技术大学鲲鹏昇腾科教创新卓越中心,借助强大的算力支持,成功研发出一项创新框架,该框架集成了专家选择分布预测、专家负载均衡以及all2all通信均衡技术。
这一创新框架直击两大核心难题:专家负载不均衡以及高昂的计算通信成本。团队首先深入剖析了专家选择的内在规律,通过对MoE层中相似token选择专家特征的细致分类与归纳,精准预测了token的专家选择路径倾向,为后续资源分配奠定了坚实基础。
针对专家负载不均衡的顽疾,团队利用预测结果,对专家动态访问频率和重要性进行了精确评估。通过复制高频专家、优化非关键专家配置的策略,实现了负载的有效均衡。在这一过程中,昇腾推理引擎MindIE的性能分析工具发挥了不可或缺的作用,其提供的可视化数据支持,为实时监测专家负载、统计访问频率等关键指标提供了有力依据。
为了进一步削减分布式计算中的通信时延,团队巧妙地将预测结果与数据布局相结合,提前将专家部署到对应节点,并将需紧密协作的专家置于同一计算单元,从而大幅缩短了数据与所需专家之间的物理距离。加之昇腾384超节点采用的高速总线互联技术,跨节点和跨计算单元的通信开销得到了显著降低。
实践是检验真理的唯一标准。该创新框架在实际应用中取得了显著成效。在多尺寸、多结构的MoE模型中,推理时间、all2all时间、MoE层时间及负载不均分数相较于现有主流方案,均有超过30%的提升。在更为复杂的多卡场景下,前三项指标同样提升了30%,而推理时间更是实现了20%的飞跃。
此次项目的成功,不仅为开发者利用昇腾技术优化MoE稀疏大模型推理性能提供了宝贵经验,更为加速其在各领域的应用落地开辟了新路径。中国科学技术大学鲲鹏昇腾科教创新卓越中心将继续深化产学研合作,依托昇腾软硬件平台,在AI模型优化领域不断探索,为推动我国人工智能战略与数字经济的蓬勃发展贡献力量。