電子發(fā)燒友網(wǎng)報(bào)道（文/梁浩斌）中國(guó)超算的高光時(shí)刻，是2016年到2017年期間，神威·太湖之光連續(xù)四次登頂全球超級(jí)計(jì)算機(jī)TOP 500榜單。但在神威·太湖之光之后，國(guó)內(nèi)超級(jí)計(jì)算機(jī)一直頗為低調(diào)，其繼任者神威“Oceanlite”沒(méi)有公開(kāi)太多信息，也沒(méi)有再參與TOP500排名。

不過(guò)近期中國(guó)科學(xué)技術(shù)大學(xué)的一個(gè)研究團(tuán)隊(duì)，在神威“Oceanlite”上利用大規(guī)模神經(jīng)網(wǎng)絡(luò)，實(shí)現(xiàn)了量子化學(xué)模擬的突破，在真實(shí)分子尺度下完成復(fù)雜量子化學(xué)建模。該成果也在近期發(fā)布在高性能計(jì)算領(lǐng)域頂級(jí)期刊IEEE TPDS上。

量子化學(xué)是什么？

量子化學(xué)研究的其實(shí)是分子、原子層面的量子行為（如電子云分布、化學(xué)鍵作用），比如通過(guò)捕捉化學(xué)反應(yīng)中間的 “過(guò)渡態(tài)”（反應(yīng)最關(guān)鍵、最短暫的步驟），來(lái)解釋實(shí)驗(yàn)無(wú)法理解的反應(yīng)本質(zhì)。

然而這些過(guò)程根本無(wú)法通過(guò)傳統(tǒng)實(shí)驗(yàn)直接觀測(cè)或精準(zhǔn)控制，首先電子運(yùn)動(dòng)速度極快，且一旦實(shí)驗(yàn)設(shè)備對(duì)其進(jìn)行觀測(cè)就會(huì)干擾其狀態(tài)。

在微觀尺度下，分子的量子行為遵循薛定諤方程，但直接求解這一方程的計(jì)算量會(huì)隨著電子數(shù)量的增加呈指數(shù)級(jí)爆炸式增長(zhǎng)。例如，一個(gè)僅包含 10 個(gè)電子的分子，其可能的量子態(tài)就超過(guò) 10 萬(wàn)億種，而真實(shí)的藥物分子或催化劑往往包含數(shù)十個(gè)電子和上百個(gè)自旋軌道。

那么如何計(jì)算？傳統(tǒng)的近似方法如密度泛函理論（DFT）雖然能處理較大規(guī)模的系統(tǒng)，但在描述強(qiáng)電子關(guān)聯(lián)現(xiàn)象（如化學(xué)鍵的斷裂與形成）時(shí)會(huì)出現(xiàn)顯著誤差。量子計(jì)算機(jī)理論上能解決這一問(wèn)題，但其當(dāng)前仍處于含噪聲中等規(guī)模（NISQ）階段，無(wú)法支撐實(shí)際科研需求。

簡(jiǎn)單來(lái)說(shuō)就是“實(shí)驗(yàn)做不到，計(jì)算太復(fù)雜”。

如何在超算上實(shí)現(xiàn)量子化學(xué)模擬？

這次中國(guó)科學(xué)技術(shù)大學(xué)研究團(tuán)隊(duì)的核心是突破神經(jīng)網(wǎng)絡(luò)量子態(tài)（NNQS）在大規(guī)模電子結(jié)構(gòu)分析中的瓶頸，實(shí)現(xiàn)真實(shí)分子尺度的量子模擬，無(wú)需依賴(lài)未來(lái)的量子計(jì)算機(jī)。

團(tuán)隊(duì)利用神威“Oceanlite”超算，開(kāi)發(fā)了一個(gè)名為數(shù)據(jù)并行NNQS-Transformer框架，在這個(gè)框架中主要有三大核心創(chuàng)新。

首先是采用了混合架構(gòu)設(shè)計(jì)，研究團(tuán)隊(duì)將Transformer 的注意力機(jī)制引入量子波函數(shù)構(gòu)建，通過(guò)動(dòng)態(tài)捕捉電子間的長(zhǎng)程相互作用，解決了傳統(tǒng)神經(jīng)網(wǎng)絡(luò)量子態(tài)（NNQS）在強(qiáng)關(guān)聯(lián)系統(tǒng)中的精度不足問(wèn)題。

其次是使用了分層通信模型，針對(duì)神威“Oceanlite”超算采用的 SW26010-Pro CPU，單芯片384核心，整機(jī)4100萬(wàn)核心集群的獨(dú)特架構(gòu)，設(shè)計(jì)了管理核心（MPE）與計(jì)算核心（CPE）的分工體系。管理核心負(fù)責(zé)節(jié)點(diǎn)間通信，而配備 512 位向量引擎的 CPE 專(zhuān)注于本地量子計(jì)算，包括隨機(jī)采樣等，顯著提升了數(shù)據(jù)訪問(wèn)效率。

SW26010 Pro是SW26010的升級(jí)版，CPU核心架構(gòu)從上一代4個(gè)核簇升級(jí)為6個(gè)，每個(gè)核簇由一個(gè)MPE和64個(gè)CPE組成，MPE擁有32KB L1指令緩存、32KB L1數(shù)據(jù)緩存以及512KB L2緩存。每個(gè)CPE擁有32KB的L1指令緩存和256KB軟件管理的高速臨時(shí)存儲(chǔ)。

最后，團(tuán)隊(duì)在3700萬(wàn)個(gè)CPE核心上進(jìn)行量子化學(xué)模擬計(jì)算，在龐大規(guī)模的核心數(shù)量下，通過(guò)實(shí)時(shí)調(diào)整各個(gè)核心的任務(wù)分配，實(shí)現(xiàn)了92%的強(qiáng)擴(kuò)展性和98%的弱擴(kuò)展性，即無(wú)論是提高問(wèn)題規(guī)模還是增加超算核心規(guī)模，都能保證系統(tǒng)性能，確保了超大規(guī)模計(jì)算中的資源高效利用。

神威超算的SW26010-ProCPU采用本地內(nèi)存而非緩存的設(shè)計(jì)，原本更適合規(guī)則化的 AI 訓(xùn)練任務(wù)，但量子化學(xué)模擬需處理隨機(jī)采樣、不規(guī)則局域能量計(jì)算，所以還需要針對(duì)性改造算法。為適配量子化學(xué)的不規(guī)則計(jì)算特性，團(tuán)隊(duì)采用Julia語(yǔ)言進(jìn)行代碼開(kāi)發(fā)，利用其可移植性和高性能特性，結(jié)合底層優(yōu)化實(shí)現(xiàn)了與硬件的無(wú)縫對(duì)接。

最終，該團(tuán)隊(duì)成功模擬了包含120 個(gè)自旋軌道的分子系統(tǒng)，這是目前經(jīng)典超算上規(guī)模最大的 AI 驅(qū)動(dòng)量子化學(xué)計(jì)算。

新材料、藥物研發(fā)領(lǐng)域前景廣泛

計(jì)算機(jī)模擬量子化學(xué)的基本目的是理解自然規(guī)律，通過(guò)計(jì)算機(jī)模擬，可以捕捉化學(xué)反應(yīng)中間的 “過(guò)渡態(tài)”，解釋實(shí)驗(yàn)無(wú)法理解反應(yīng)本質(zhì)。在實(shí)際應(yīng)用中，通過(guò)模擬還原一些自然環(huán)境中產(chǎn)生的化學(xué)反應(yīng)，比如模擬光合作用，電子在葉綠素中的傳遞路徑、過(guò)渡態(tài)的能量變化，為人工光合制氫提供理論指導(dǎo)。

而在這個(gè)過(guò)程中，對(duì)微觀現(xiàn)象的模擬對(duì)于創(chuàng)造新的材料更是突破性的新模式。在藥物研發(fā)領(lǐng)域，通過(guò)計(jì)算機(jī)模擬藥物分子與疾病靶點(diǎn)，如蛋白質(zhì)、病毒等的量子相互作用，預(yù)測(cè) 哪種分子能精準(zhǔn)結(jié)合靶點(diǎn)、阻斷疾病通路。

在材料領(lǐng)域，通過(guò)模擬分子、原子的量子結(jié)構(gòu)，預(yù)測(cè)材料的物理化學(xué)性質(zhì)，如導(dǎo)電性、催化活性、強(qiáng)度等，再反向設(shè)計(jì)出符合需求的材料。比如用于研發(fā)更高效的燃料電池催化劑、固態(tài)電池的電解質(zhì)材料等。

寫(xiě)在最后

從該項(xiàng)研究成果可以看出，超算硬件算力基礎(chǔ)，與軟件生態(tài)系統(tǒng)能力的高度配合同樣重要。目前中國(guó)神威“Oceanlite”超算不再參與算力排名，或許也表示國(guó)內(nèi)超算的指標(biāo)，從單純追求硬件參數(shù)，轉(zhuǎn)向了追求科研平臺(tái)落地，構(gòu)建自主的超算生態(tài)。

同時(shí)也證明了量子計(jì)算機(jī)實(shí)用化之前，經(jīng)典超算與AI結(jié)合，仍有非常大的潛力可以挖掘。