中國(guó)團(tuán)隊(duì)刷新紀(jì)錄：超算“神威·Oceanlite”上實(shí)現(xiàn)AI驅(qū)動(dòng)的最大規(guī)模量子化學(xué)模擬
2016到2017年，中國(guó)超級(jí)計(jì)算機(jī)“神威·太湖之光”曾四度登頂全球超級(jí)計(jì)算機(jī)TOP 500榜單，將中國(guó)超算推上世界之巔。此后，國(guó)產(chǎn)超算變得相對(duì)低調(diào)，繼任者“神威·Oceanlite”相關(guān)信息鮮少公開，也未再現(xiàn)身TOP 500排名。然而，國(guó)內(nèi)科研團(tuán)隊(duì)并未止步，反而在幕后默默發(fā)力。
近日，中國(guó)科學(xué)技術(shù)大學(xué)團(tuán)隊(duì)在“神威·Oceanlite”上，利用大規(guī)模神經(jīng)網(wǎng)絡(luò)技術(shù)取得了突破性成果：首次在真實(shí)分子尺度下完成了復(fù)雜的量子化學(xué)建模。這一成果發(fā)表于高性能計(jì)算領(lǐng)域頂級(jí)期刊《IEEE TPDS》，為超算在量子化學(xué)領(lǐng)域的新應(yīng)用樹立了標(biāo)桿。
量子化學(xué)：用計(jì)算洞見分子世界
量子化學(xué)研究的是分子、原子以及電子等微觀粒子的量子行為。比如化學(xué)反應(yīng)中關(guān)鍵但難以捕捉的“過渡態(tài)”，是理解反應(yīng)本質(zhì)的關(guān)鍵環(huán)節(jié)。受限于實(shí)驗(yàn)手段，很難直接觀測(cè)這些瞬時(shí)過程，因?yàn)樵谖⒂^世界里，電子運(yùn)動(dòng)極快，一旦試圖測(cè)量就會(huì)干擾系統(tǒng)本身。
理論上，分子的量子行為遵循薛定諤方程。但這個(gè)方程的復(fù)雜性隨系統(tǒng)中電子數(shù)增長(zhǎng)而呈指數(shù)暴漲。例如，僅10個(gè)電子的分子，就可能存在超過10萬億種量子態(tài)。藥物分子、催化劑等實(shí)際分子則往往擁有幾十電子、上百自旋軌道，計(jì)算量更是天文數(shù)字。
傳統(tǒng)上，科研人員通過密度泛函理論（DFT）等近似方法來求解這些復(fù)雜問題，雖然能處理較大體系，但對(duì)于強(qiáng)電子關(guān)聯(lián)現(xiàn)象（如化學(xué)鍵斷裂、形成）容易產(chǎn)生較大誤差。理論上量子計(jì)算機(jī)能夠解決這一難題，但以目前技術(shù)，量子計(jì)算機(jī)還遠(yuǎn)未到實(shí)用階段。
AI+超算：破解“計(jì)算太復(fù)雜”的難題
這一次，中科大團(tuán)隊(duì)的突破點(diǎn)，在于結(jié)合大規(guī)模神經(jīng)網(wǎng)絡(luò)與超級(jí)計(jì)算機(jī)硬件，讓經(jīng)典計(jì)算機(jī)也能逼近量子計(jì)算機(jī)才能勝任的分子模擬難題。團(tuán)隊(duì)開發(fā)的數(shù)據(jù)并行型NNQS-Transformer框架，是一套專為超算“神威·Oceanlite”打造的神經(jīng)網(wǎng)絡(luò)量子態(tài)解決方案，有三大核心創(chuàng)新：
混合架構(gòu)創(chuàng)新
框架結(jié)合了Transformer注意力機(jī)制，用于動(dòng)態(tài)捕捉電子間復(fù)雜的長(zhǎng)程相互作用。這樣，能夠顯著提升強(qiáng)關(guān)聯(lián)系統(tǒng)中的計(jì)算精度，克服傳統(tǒng)神經(jīng)網(wǎng)絡(luò)量子態(tài)（NNQS）方法面臨的精度瓶頸。
高效并行設(shè)計(jì)
針對(duì)“神威·Oceanlite”采用的國(guó)產(chǎn)SW26010-Pro CPU（單芯384核，整機(jī)4100萬核），團(tuán)隊(duì)構(gòu)建了MPE（管理核心）與CPE（計(jì)算核心）的分工體系。MPE負(fù)責(zé)跨節(jié)點(diǎn)數(shù)據(jù)通信，CPE配備512位向量引擎，專注本地高效量子計(jì)算。如此分工配合，使得在3700萬個(gè)CPE核心上進(jìn)行大規(guī)模模擬時(shí)，任務(wù)分配可以實(shí)時(shí)調(diào)整，計(jì)算資源利用率達(dá)到92%的強(qiáng)擴(kuò)展性、98%的弱擴(kuò)展性，極大提升了超算效率。
軟硬件適配與底層優(yōu)化
“神威·Oceanlite”CPU架構(gòu)采用本地內(nèi)存設(shè)計(jì)，原本更適合AI訓(xùn)練，但量子化學(xué)模擬面對(duì)隨機(jī)采樣、不規(guī)則計(jì)算等特殊難題。團(tuán)隊(duì)針對(duì)性地用Julia語言開發(fā)代碼，結(jié)合底層硬件深度優(yōu)化，實(shí)現(xiàn)了任務(wù)與硬件之間的高效對(duì)接。
最終，研究團(tuán)隊(duì)成功實(shí)現(xiàn)了包含120個(gè)自旋軌道的分子模擬，這是目前經(jīng)典超算下規(guī)模最大的AI驅(qū)動(dòng)量子化學(xué)計(jì)算。
推動(dòng)新材料、藥物發(fā)現(xiàn)模式變革
用超級(jí)計(jì)算機(jī)模擬量子化學(xué)的意義，首先在于理解化學(xué)反應(yīng)的本質(zhì)過程。通過虛擬模擬，可以重現(xiàn)如光合作用等關(guān)鍵“過渡態(tài)”，為人工制氫等原創(chuàng)技術(shù)提供理論依據(jù)。
更進(jìn)一步，在新材料、藥物研發(fā)領(lǐng)域，量子級(jí)模擬能預(yù)測(cè)藥物分子與蛋白、病毒的結(jié)合能力，幫助提前篩選新藥。材料領(lǐng)域則可以借此模擬分子結(jié)構(gòu)，預(yù)測(cè)材料的物理化學(xué)性質(zhì)，反向設(shè)計(jì)催化劑、電池材料等，實(shí)現(xiàn)從原理到產(chǎn)業(yè)的創(chuàng)新飛躍。
結(jié)語
這項(xiàng)研究表明，超級(jí)計(jì)算機(jī)不僅僅是硬件算力的較量，更是軟硬件協(xié)同與創(chuàng)新生態(tài)建設(shè)的重要舞臺(tái)。當(dāng)中國(guó)超算榜單排名不再是關(guān)注重點(diǎn)，平臺(tái)化科研和自主生態(tài)正成為新的發(fā)展重心。同樣令人振奮的是，在量子計(jì)算機(jī)尚處于“嬰兒期”時(shí)，經(jīng)典超算與AI結(jié)合，依舊擁有強(qiáng)大的科研創(chuàng)新潛力——計(jì)算邊界仍在持續(xù)拓展。
