隨著人工智能技術(shù)的飛速發(fā)展，大模型參數(shù)量激增。MoE 稀疏大模型因能在提升容量的同時(shí)保持訓(xùn)練計(jì)算量，成為了當(dāng)下的研究熱點(diǎn)，但推理時(shí)的性能瓶頸仍制約著其進(jìn)一步發(fā)展。在此背景下，中國(guó)科學(xué)技術(shù)大學(xué)人工智能與數(shù)據(jù)科學(xué)學(xué)院執(zhí)行院長(zhǎng)張燕詠率領(lǐng)團(tuán)隊(duì)，在中國(guó)科學(xué)技術(shù)大學(xué) 鯤鵬昇騰科教創(chuàng)新卓越中心的算力支持下，成功研發(fā)出了基于專(zhuān)家選擇分布預(yù)測(cè)的專(zhuān)家負(fù)載均衡和all2all通信均衡框架。

該框架的研究致力于解決兩大難題——專(zhuān)家負(fù)載不均及較大的計(jì)算通信開(kāi)銷(xiāo)，并通過(guò)“三步走”成功實(shí)現(xiàn)了推理性能的優(yōu)化。為了讓模型在并行推理時(shí)能更合理地分配資源，團(tuán)隊(duì)首先從專(zhuān)家選擇的規(guī)律入手，通過(guò)分析相似 token 在 MoE 層選擇專(zhuān)家的特征并歸類(lèi)，成功預(yù)測(cè)出token的專(zhuān)家選擇路徑偏好。

接著便是解決推理過(guò)程中的專(zhuān)家負(fù)載不均衡問(wèn)題。團(tuán)隊(duì)依據(jù)token對(duì)專(zhuān)家選擇偏好的預(yù)測(cè)結(jié)果，精準(zhǔn)判斷出專(zhuān)家動(dòng)態(tài)訪問(wèn)的頻率和重要性，通過(guò)復(fù)制高頻專(zhuān)家、替換非重要專(zhuān)家的方式實(shí)現(xiàn)了負(fù)載平衡。其中昇騰推理引擎MindIE提供的性能分析工具可實(shí)現(xiàn)性能數(shù)據(jù)的可視化，在實(shí)時(shí)監(jiān)測(cè)專(zhuān)家負(fù)載、統(tǒng)計(jì)訪問(wèn)頻率等方面提供數(shù)據(jù)參考，發(fā)揮了重要作用。

降低分布式計(jì)算的通信時(shí)延是第三步。團(tuán)隊(duì)根據(jù)計(jì)算節(jié)點(diǎn)上的數(shù)據(jù)對(duì)專(zhuān)家選擇的預(yù)測(cè)結(jié)果，將專(zhuān)家提前部署到對(duì)應(yīng)的節(jié)點(diǎn)上，并聯(lián)合需要緊密協(xié)作的專(zhuān)家放在同一計(jì)算單元，讓數(shù)據(jù)與所需專(zhuān)家在物理位置上更接近，同時(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)目的成果，為開(kāi)發(fā)者借助昇騰技術(shù)優(yōu)化 MoE 稀疏大模型推理性能提供了可貴借鑒，將加速其在各領(lǐng)域的落地應(yīng)用。未來(lái)，中國(guó)科學(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)力量。