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