摘要:在異構計算架構CANN的助力下,AI預測效能達到現有產品的1.5+倍,可預測規模較傳統方法提升10000+倍,為光伏材料、新能源電池、半導體材料研究帶來巨大的商業應用價值。
本文分享自華為雲社群《效能達1.5+倍!昇騰AI助力分子動力學模擬研究》,作者:昇騰CANN。
研究背景
分子動力學是一套分子模擬方法,主要依靠牛頓力學來模擬分子體系運動,是研究微觀世界的有效手段。但傳統研究手段的時間複雜度較高,僅限於研究數千量級原子的小型系統,在應用中存在普遍侷限性。而深勢科技提出的DeePMD-kit可將AI技術成功應用於分子動力學模擬,並實現了上億原子體系的模擬,目前該項技術已經發展成為端到端的開源解決方案,以科學智慧(AI for Science)為分子動力學的行業應用帶來了更多可能。
而昇騰AI憑藉超強算力及完整軟體棧配套,為科學智慧提供了天然土壤,透過昇騰AI硬體及異構計算架構CANN,藉助軟硬體協同能力充分釋放硬體算力,為DeePMD的大規模分子動力學模擬提供了高效能解決方案。
基於CANN自定義DeePMD高效能運算元
對諸如DeePMD之類的神經網路模型進行加速的關鍵之一是對網路中的運算元進行深度最佳化,從而發揮出硬體的澎湃算力,而作為專門面向AI場景的異構計算架構,CANN正是搭起了上層深度學習框架和底層AI硬體的橋樑。
基於CANN的運算元程式設計介面,開發人員在DeePMD網路中開發相關的自定義運算元,其中涉及到資料排序、讀取與儲存等離線操作,也涉及矩陣、標量的計算,這些自定義運算元不但對精度有較高的要求,而且也很大程度上決定了模型整體的效能。
CANN能夠高效協同昇騰不同異構單元自定義高效能運算元,充分釋放AI Core、AI CPU和Vector Core的異構算力。比如將離散的距離計算和排序部署在AI CPU上,將可以並行的矩陣、標量的計算部分部署在AI Core上,以發揮出每個計算單元的能力,充分發揮硬體計算效能。
高效協同昇騰不同異構單元自定義高效能運算元
基於CANN的DeePMD融合規則及網路最佳化
運算元融合也是網路效能最佳化重點之一,也是業界常用手段,在DeePMD原生實現中,運算元融合就作為效能工程的主要手段之一。
而在DeePMD網路效能最佳化中,運算元開發人員進一步結合網路特點,藉助CANN設計新的融合規則,包括PAD運算元支援動靜合一、Mul支援NZ+ND和ND+NZ、MatMul + Add + TanhGrad支援Buffer融合等,這些融合在整網效能提升中起到了關鍵作用。同時藉助CANN的智慧調優工具AOE,自動化完成子圖調優、運算元調優,以及TransData消除等一系列最佳化,極大提升了模型調優效率。
DeePMD網路融合示例
此外專案最佳化了整網混合精度推理流程,將半精度和單精度混合使用,有效減少記憶體佔用,充分釋放硬體算力,進一步加速模型執行,在確保分子模擬精度的前提下帶來最大效能收益。在不同場景對分子模擬的物理性質與模擬結果對比測試中,以銅拉伸應力應變模擬實驗為例可以看到,昇騰AI基礎軟硬體平臺的實驗結果已經非常近似實際資料。
昇騰AI基礎軟硬體平臺上銅的拉伸應力應變模擬結果
昇騰AI的整體最佳化解決方案使DeePMD-kit工具在分子動力學模擬計算上取得1.5+倍現有產品的效能提升成果,助力分子動力學服務成功商用。未來,昇騰AI將繼續以超強算力和軟硬協同能力助力科學計算進入科學智慧新階段,攜手更多夥伴凝心聚力,共同向上發展,構築科學智慧領先格局,全面邁進數智時代。