隨著人工智能技術的飛速發展，大模型參數量激增。MoE 稀疏大模型因能在提升容量的同時保持訓練計算量，成為了當下的研究熱點，但推理時的性能瓶頸仍制約著其進一步發展。在此背景下，中國科學技術大學人工智能與數據科學學院執行院長張燕詠率領團隊，在中國科學技術大學 鯤鵬昇騰科教創新卓越中心的算力支持下，成功研發出了基于專家選擇分布預測的專家負載均衡和all2all通信均衡框架。
該框架的研究致力于解決兩大難題——專家負載不均及較大的計算通信開銷，并通過“三步走”成功實現了推理性能的優化。為了讓模型在并行推理時能更合理地分配資源，團隊首先從專家選擇的規律入手，通過分析相似 token 在 MoE 層選擇專家的特征并歸類，成功預測出token的專家選擇路徑偏好。
接著便是解決推理過程中的專家負載不均衡問題。團隊依據token對專家選擇偏好的預測結果，精準判斷出專家動態訪問的頻率和重要性，通過復制高頻專家、替換非重要專家的方式實現了負載平衡。其中昇騰推理引擎MindIE提供的性能分析工具可實現性能數據的可視化，在實時監測專家負載、統計訪問頻率等方面提供數據參考，發揮了重要作用。
降低分布式計算的通信時延是第三步。團隊根據計算節點上的數據對專家選擇的預測結果，將專家提前部署到對應的節點上，并聯合需要緊密協作的專家放在同一計算單元，讓數據與所需專家在物理位置上更接近，同時結合昇騰384超節點采用的高速總線互聯技術，最終大幅降低了跨節點和跨計算單元的通信消耗。
在系列技術突破下，此次項目成果顯著。基于該創新框架，在多尺寸、多結構的 MoE 模型中，推理時間、all2all 時間、MoE 層時間及負載不均分數較現有主流方案提升超 30%；多卡復雜場景里，前三項指標提升 30%，推理時間提升 20%。
本次項目的成果，為開發者借助昇騰技術優化 MoE 稀疏大模型推理性能提供了可貴借鑒，將加速其在各領域的落地應用。未來，中國科學技術大學 鯤鵬昇騰科教創新卓越中心將持續深化產學研協同創新，依托昇騰軟硬件平臺，在前沿AI模型優化領域持續突破，為自主創新人工智能戰略與數字經濟發展貢獻力量。
