最近和本文的作者合作進(jìn)行神威上的相關(guān)工作，學(xué)習(xí)一下老師的上一篇工作。

TensorKMC: Kinetic Monte Carlo Simulation of 50 Trillion Atoms Driven by Deep Learning on a New Generation of Sunway Supercomputer

????????文章提出了一種三編碼算法和空置緩存機(jī)制，高效地將第一性原理勢(shì)神經(jīng)網(wǎng)絡(luò)和原子的動(dòng)力學(xué)蒙特卡洛模擬（AKMC）結(jié)合。通過(guò)實(shí)現(xiàn)了快速的特征運(yùn)算符、NNP的大型算子融合和內(nèi)存優(yōu)化，TensorKMC算法在SW26010-pro上模擬了54trillions原子，并最高可以擴(kuò)展至27456000個(gè)核心。

INTRODUCTION

????????動(dòng)力學(xué)蒙特卡洛（KMC）是一種中尺度的模擬方法，用于研究材料的宏觀失效，例如斷裂、磨損、腐蝕，其中，原子的動(dòng)力學(xué)蒙特卡洛方法（AKMC）將跳躍事件對(duì)應(yīng)至晶格點(diǎn)狀態(tài)的改變。AKMC與第一性原理的方法DFT結(jié)合，提高了精確度，但受限于較高的計(jì)算開(kāi)銷(xiāo)。分子動(dòng)力學(xué)模擬（MD）領(lǐng)域通過(guò)結(jié)合DFT與神經(jīng)網(wǎng)絡(luò)取得了較大進(jìn)展，但由于較小的時(shí)間步，使得MD無(wú)法模擬較大的時(shí)間尺度，不能替代KMC方法。

????????OpenKMC是一種AKMC的實(shí)現(xiàn)，用于研究錒系合金中空位缺陷的動(dòng)態(tài)演化。文章結(jié)合OpenKMC和TensorAlloy，在超級(jí)計(jì)算機(jī)上模擬核反應(yīng)壓力容器上合金的微結(jié)構(gòu)演化。主要的挑戰(zhàn)包括：將適用于MD的TensorAlloy，高效地與AKMC結(jié)合；內(nèi)存占用過(guò)大；如何針對(duì)新神威架構(gòu)優(yōu)化。

BACKGROUND

????????AKMC模型的基本假設(shè)是，空位擴(kuò)散通過(guò)一系列的最近鄰居跳躍完成。AMKC首先計(jì)算所有可能發(fā)生的空位跳躍的概率，概率取決于不同狀態(tài)的能量變化，隨后按照概率隨機(jī)發(fā)生空位跳躍，最后計(jì)算時(shí)間間隔。

????????并行AKMC框架使用MPI進(jìn)行數(shù)據(jù)并行，由于晶體中原子時(shí)刻處于晶格點(diǎn)上，這勢(shì)必面臨跳躍沖突的問(wèn)題，即不同的MPI進(jìn)程中的原子跳躍到相同的位置。解決方案是通過(guò)對(duì)每個(gè)MPI進(jìn)程劃分子晶格分別計(jì)算來(lái)避免這個(gè)問(wèn)題。

????????超級(jí)計(jì)算機(jī)使用了搭載新眾核CPU SW26010-pro的新一代神威超級(jí)計(jì)算機(jī)。

????????AKMC與MD有著許多不同，主要體現(xiàn)在：MD中原子位置不固定，但AKMC中原子始終處于晶格點(diǎn)上；MD中通過(guò)原子能量計(jì)算受力，AKMC中通過(guò)原子能量驅(qū)動(dòng)演化；AKMC若使用傳統(tǒng)的數(shù)據(jù)并行算法會(huì)發(fā)生跳躍沖突的問(wèn)題；MD的每個(gè)時(shí)間步需要計(jì)算所有原子，AKMC中只有空位跳躍附近的原子會(huì)被影響。而根據(jù)MD的原則構(gòu)建的OpenKMC，遇到了許多挑戰(zhàn)：存儲(chǔ)了所有原子的性質(zhì)，導(dǎo)致內(nèi)存占用過(guò)大；對(duì)所有原子計(jì)算能量，導(dǎo)致冗余的計(jì)算；計(jì)算所需的內(nèi)存超過(guò)LDM。

INNOVATION AND OPTIMIZATION

Triple-Encoding

????????每個(gè)空位跳躍只由空位以及周?chē)淖罱従訁⑴c，影響范圍被限制在截?cái)嗑嚯x之內(nèi)。又由于晶體中空位濃度很低，每次計(jì)算只需要考慮很少部分的原子，這樣通過(guò)對(duì)每個(gè)空位及周邊原子進(jìn)行編碼，可以節(jié)省內(nèi)存并便于CPE參與計(jì)算。文章提出了Triple-Encoding方法，使用三個(gè)張量：相對(duì)坐標(biāo)編碼表（CET）、近鄰列表編碼表（NET）、空位編碼表（VET），來(lái)描述一個(gè)空位系統(tǒng)。

Vacancy Cache Mechanism

????????TensorKMC中不再將所有原子信息緩存到內(nèi)存之中，而是只緩存空位系統(tǒng)的編碼表，并在每次計(jì)算完成后更新。

Memory Usage

????????TensorKMC通過(guò)以算代存，將根據(jù)空間坐標(biāo)訪問(wèn)晶格點(diǎn)性質(zhì)的過(guò)程由2次內(nèi)存訪問(wèn)減少為1次，同時(shí)減少了內(nèi)存占用。

Fast Feature Operator

????????TensorKMC使用的特征函數(shù)如下，其中p、q為超參數(shù)。又由于晶體中距離r是離散分布的，因此可以對(duì)函數(shù)的指數(shù)部分預(yù)先建立關(guān)于r、p、q的表，便于使用for循環(huán)計(jì)算，并利用CPE并行加速。

Big-Fusion Operator

卷積神經(jīng)網(wǎng)絡(luò)是內(nèi)存密集的任務(wù)，通過(guò)改進(jìn)算法，隱藏內(nèi)存訪問(wèn)和數(shù)據(jù)交換時(shí)間，將這個(gè)過(guò)程變?yōu)橛?jì)算密集的。

a. 將1*1卷積計(jì)算轉(zhuǎn)化為矩陣乘，ci指輸入通道，co指輸出通道。

b. 將矩陣乘、偏置和ReLU合并為一個(gè)kernel。

c. 將卷積神經(jīng)網(wǎng)絡(luò)的所有l(wèi)ayers合并為一個(gè)kernel，減少內(nèi)存訪問(wèn)次數(shù)。

d. 模型分布式存儲(chǔ)在CPEs上，通過(guò)RMA共享參數(shù)。

e. 使用2個(gè)buffer，流水線式計(jì)算所有state的能量。

f. CPEs內(nèi)部的流水線計(jì)算。

EVALUATION

這部分就略過(guò)了……

參考文獻(xiàn)

[1] SHANG H, CHEN X, GAO X, 等. TensorKMC: kinetic Monte Carlo simulation of 50 trillion atoms driven by deep learning on a new generation of Sunway supercomputer[C/OL]//Proceedings of the International Conference for High Performance Computing, Networking, Storage and Analysis. St. Louis Missouri: ACM, 2021: 1-14[2023-02-26]. https://dl.acm.org/doi/10.1145/3458817.3476174.