在当今人工智能迅速发展的背景下，大规模模型（如MoE，即Mixture of Experts）因其能在扩展模型容量的同时降低训练计算成本而成为研究的热点。这类模型在推理阶段常常面临性能瓶颈，严重制约了其在实际应用中的广泛落地。对此，中国科学技术大学人工智能与数据科学学院的张燕咏教授及其团队在近日取得了显著的技术突破，成功研发出了一种基于专家选择分布预测的专家负载均衡和all2all通信均衡框架，为MoE稀疏大模型的推理性能提升提供了新的解决方案。
一、技术背景与挑战
随着人工智能应用场景日益复杂，对模型的性能和处理能力也提出了更高要求。MoE稀疏大模型因具备优良的性能与可扩展性，广泛用于自然语言处理、计算机视觉等领域。推理过程中的计算通信开销及专家负载不均始终是其面临的两大核心问题。如何在保持模型性能的基础上，优化推理时的资源分配与通信效率，成为了各研究团队亟待解决的难题。
二、专家选择与负载均衡的创新方法
在张教授的团队中，研究人员关注到了“专家选择”的规律，通过对相似token在MoE层选择专家的特征进行深入分析，构建了专家选择路径偏好的预测模型。这一创新的思路使团队能够在并行推理时更加精准地分配计算资源，避免了因负载不均而造成的性能下降。
在解决负载不均的问题上，团队通过对token选择偏好的预测结果，能够即时了解各专家的动态访问频率和重要性。信息，团队实施了高频专家的复制与非重要专家的替换，有效实现了负载的均衡。值得一提的是，昇腾推理引擎MindIE提供的性能分析工具，在实时监测专家的负载与访问频率方面发挥了不可或缺的支持。
三、降低通信延迟的有效策略
为了进一步提升推理性能，团队将重心转向了分布式计算中的通信时延问题。他们根据计算节点上的数据，结合专家选择的预测结果，提前将专家部署至对应的计算节点。这一措施让需要紧密合作的专家得以放置在同一计算单元中，大幅提升了数据与专家的物理接近度。同时，依托昇腾384超节点的高速总线互联技术，团队有效降低了跨节点及跨计算单元间的通信开销，进一步推动了推理速度的提升。
四、技术成果与未来展望
经过一系列的技术突破，团队的研究成果展现了显著的优势。基于这一创新框架，在不同尺寸和结构的MoE模型中，推理时间、all2all时间和MoE层时间的提升率均超过了30%；在多卡复杂场景下，推理时间也实现了20%的提升。这一系列改进的有效实施不仅为现有主流方案提供了更具竞争力的选择，也为后续开发者在使用昇腾技术优化MoE稀疏大模型推理性能方面提供了宝贵的经验和借鉴。
展望未来，中国科学技术大学鲲鹏昇腾科教创新卓越中心将继续深化产学研的协同创新，依托昇腾软硬件，在AI模型优化及相关应用领域持续推进技术研发与突破，以期为我国的自主创新人工智能战略和数字经济发展贡献更多力量。
中国科学技术大学团队的成功研发不仅标志着MoE稀疏大模型推理性能的重大进展，也为推动该技术的广泛应用铺平了道路。随着这一技术的不断成熟，我们有理由相信未来会在更广泛的领域中看到人工智能的应用，从而加速数字经济和智能社会的发展进程。