在人工智能大模型迅猛發(fā)展的當(dāng)下，大模型參數(shù)和計(jì)算量呈指數(shù)級(jí)增長(zhǎng)，大規(guī)模深度學(xué)習(xí)模型的訓(xùn)練離不開(kāi)多硬件設(shè)備的分布式計(jì)算。在鯤鵬昇騰科教創(chuàng)新卓越中心的算力支持下，北京大學(xué)計(jì)算機(jī)學(xué)院崔斌教授團(tuán)隊(duì)創(chuàng)新研發(fā)了面向大模型的高效分布式訓(xùn)練框架，大幅提升了大規(guī)模分布式訓(xùn)練模型的效率。
針對(duì)模型訓(xùn)練任務(wù)的多樣性和復(fù)雜性所帶來(lái)的負(fù)載不均問(wèn)題，研究團(tuán)隊(duì)創(chuàng)新設(shè)計(jì)出了細(xì)粒度模型切分與并行策略搜索算法。此方法依托昇騰強(qiáng)大的計(jì)算資源管理能力，及算子優(yōu)化技術(shù)對(duì)訓(xùn)練策略的適配，完成了統(tǒng)一訓(xùn)練接口到智能切分策略的全過(guò)程。首先通過(guò)總結(jié)多種大模型訓(xùn)練的共同特點(diǎn)，設(shè)計(jì)出統(tǒng)一的接口來(lái)啟動(dòng)和管理不同任務(wù)負(fù)載的模型訓(xùn)練任務(wù)，對(duì)訓(xùn)練時(shí)所花費(fèi)的算力、內(nèi)存、網(wǎng)絡(luò)通信等進(jìn)行精確地計(jì)算。接著基于這些數(shù)據(jù)細(xì)致拆解龐大且復(fù)雜的大模型，根據(jù)不同模塊分配不同的訓(xùn)練策略以適應(yīng)各模塊的負(fù)載差異，實(shí)現(xiàn)訓(xùn)練任務(wù)的高效并行。目前，該方案已實(shí)現(xiàn)比分片數(shù)據(jù)并行、3D 并行等全局模版化并行方案提升15% 的訓(xùn)練效率。
除此之外，團(tuán)隊(duì)還解決了分布式計(jì)算所涉及的硬件間通信傳輸效率問(wèn)題。結(jié)合昇騰高速互聯(lián)總線技術(shù)的高帶寬低時(shí)延優(yōu)勢(shì)，系統(tǒng)會(huì)很根據(jù)不同的通信需求，對(duì)硬件設(shè)備進(jìn)行分組以優(yōu)化組隊(duì)時(shí)間，運(yùn)用計(jì)算通信重疊技術(shù)讓“計(jì)算”和“通信”同步進(jìn)行，提高訓(xùn)練流水線的效率，并在模型切分的決策時(shí)考慮計(jì)算通信重疊的性能影響，綜合多方面因素選取最適合的分布式運(yùn)行方案，最終實(shí)現(xiàn)數(shù)據(jù)傳輸效率和資源利用率的最大化。
該研究成果不僅為模型大規(guī)模訓(xùn)練提供了高效的解決方案，更展現(xiàn)了自主算力在分布式計(jì)算領(lǐng)域的巨大潛力。目前，研究成果已在國(guó)際頂級(jí)學(xué)術(shù)會(huì)議NeurIPS、ICLR、AAAI發(fā)表3篇論文，為國(guó)內(nèi)AI技術(shù)突破提供了理論支撐與實(shí)踐范式。
北京大學(xué) 鯤鵬昇騰科教創(chuàng)新卓越中心的這一突破彰顯了校企協(xié)同創(chuàng)新的顯著成效。未來(lái)，中心將持續(xù)加速AI前沿技術(shù)在自主計(jì)算平臺(tái)的深度落地，為我國(guó)人工智能產(chǎn)業(yè)的自主化突破提供強(qiáng)勁動(dòng)能。
