做材料計(jì)算和AI交叉方向的開(kāi)發(fā)者對(duì)下面的場(chǎng)景應(yīng)該不陌生傳統(tǒng)材料研發(fā)靠試錯(cuò)一個(gè)配方從設(shè)計(jì)到驗(yàn)證動(dòng)不動(dòng)按年算高通量篩選又受限于算力和數(shù)據(jù)標(biāo)注成本。這兩年AI for Science喊得響但真把模型跑起來(lái)、把推理塞進(jìn)生產(chǎn)流程坑比想象中多得多。我最近在做的這套“智能煉金術(shù)”系統(tǒng)就是把AI材料生成模型和CANN異構(gòu)計(jì)算加速結(jié)合起來(lái)走通了一條從候選材料生成、性質(zhì)預(yù)測(cè)到實(shí)驗(yàn)驗(yàn)證的閉環(huán)。今天這篇就完整拆一下設(shè)計(jì)思路和落地過(guò)程中踩過(guò)的坑給同樣做材料AI落地的朋友一個(gè)參考。這套系統(tǒng)解決的核心問(wèn)題很直接材料發(fā)現(xiàn)的起點(diǎn)不再是實(shí)驗(yàn)試錯(cuò)而是讓生成模型在元素周期表的約束空間里自動(dòng)產(chǎn)出候選結(jié)構(gòu)再用圖神經(jīng)網(wǎng)絡(luò)快速預(yù)測(cè)帶隙、形成能、力學(xué)模量等關(guān)鍵性質(zhì)最后疊加一個(gè)基于CANN優(yōu)化的推理加速層把單條候選的評(píng)估延遲從秒級(jí)壓到毫秒級(jí)。適合正在做材料基因工程、虛擬篩選、生成式分子設(shè)計(jì)的開(kāi)發(fā)者也適合想了解國(guó)產(chǎn)異構(gòu)計(jì)算棧在科學(xué)計(jì)算場(chǎng)景怎么落地的同學(xué)。1. 項(xiàng)目整體設(shè)計(jì)思路拆解1.1 為什么材料發(fā)現(xiàn)需要AI生成模型而不是傳統(tǒng)篩選老路子做新材料核心是“枚舉加過(guò)濾”。先在已知結(jié)構(gòu)數(shù)據(jù)庫(kù)里做替換、摻雜、插層再用第一性原理計(jì)算去驗(yàn)證最后挑性能達(dá)標(biāo)的進(jìn)實(shí)驗(yàn)。這套思路的問(wèn)題在于搜索空間形狀不友好——元素周期表組合起來(lái)是天文數(shù)字窮舉不現(xiàn)實(shí)隨機(jī)采樣大概率落在沒(méi)有熱力學(xué)穩(wěn)定性的區(qū)域。生成模型解決的是“從無(wú)到有”的問(wèn)題。擴(kuò)散模型和流匹配模型能在學(xué)習(xí)過(guò)的晶體結(jié)構(gòu)分布上采樣新結(jié)構(gòu)自編碼器能把結(jié)構(gòu)編碼成隱向量再在隱空間做插值和方向性編輯。我們自己對(duì)比過(guò)同樣生成1000個(gè)候選結(jié)構(gòu)純隨機(jī)采樣只有不到3%能通過(guò)基礎(chǔ)電負(fù)性和原子間距過(guò)濾而生成模型采樣的通過(guò)率能到30%以上這就是“煉金”和“搖獎(jiǎng)”的區(qū)別。但生成模型有一個(gè)代價(jià)很大的副作用計(jì)算開(kāi)銷(xiāo)翻倍往上走。一個(gè)中等規(guī)模的晶體擴(kuò)散模型單次反向采樣在CPU上要跑幾十秒在GPU上也要兩三秒。放到材料篩選中一輪要跑幾萬(wàn)個(gè)候選這個(gè)時(shí)間賬根本算不過(guò)來(lái)。所以整套系統(tǒng)的關(guān)鍵瓶頸根本不在于模型本身的精度而在于能不能把推理吞吐提上去。這也是我們引入CANN加速的直接動(dòng)因。1.2 系統(tǒng)整體架構(gòu)四層各司其職整套系統(tǒng)按數(shù)據(jù)、模型、加速、應(yīng)用四層劃分各層之間只通過(guò)標(biāo)準(zhǔn)接口通信這樣后續(xù)替換任何一層都不會(huì)牽連其他模塊。數(shù)據(jù)層負(fù)責(zé)把所有材料數(shù)據(jù)統(tǒng)一成圖結(jié)構(gòu)原子是節(jié)點(diǎn)鍵是邊再把元素屬性、電荷、半徑、電負(fù)性這些特征拼成節(jié)點(diǎn)特征向量。模型層包含兩個(gè)核心模塊一個(gè)是變分自編碼器負(fù)責(zé)結(jié)構(gòu)生成一個(gè)是圖神經(jīng)網(wǎng)絡(luò)負(fù)責(zé)性質(zhì)預(yù)測(cè)。加速層是整個(gè)系統(tǒng)的主心骨所有模型的推理、張量算子、內(nèi)存調(diào)度全部經(jīng)由CANN工具鏈做統(tǒng)一編譯和優(yōu)化。應(yīng)用層則是面向材料研究人員的篩選工作臺(tái)輸入目標(biāo)帶隙或硬度范圍輸出候選列表和解釋性特征。這種拆法最大的收益是模塊化。數(shù)據(jù)層和模型層完全對(duì)齊學(xué)術(shù)前沿讀論文就能復(fù)現(xiàn)加速層可以隨時(shí)切換設(shè)備而不改上游代碼應(yīng)用層獨(dú)立出來(lái)之后做Web服務(wù)也好、做批量任務(wù)也好都不會(huì)干擾核心計(jì)算邏輯。1.3 為什么選CANN而不是普通深度學(xué)習(xí)框架很多人會(huì)問(wèn)PyTorch加CUDA已經(jīng)很好用為什么還要引入CANN這一套最直接的原因是部署環(huán)境里沒(méi)有傳統(tǒng)GPU可用只有國(guó)產(chǎn)AI加速卡。CANN是面向這類(lèi)異構(gòu)AI處理器的軟件棧功能上和CUDA對(duì)位但生態(tài)、算子庫(kù)、圖優(yōu)化策略都有自己的一套邏輯。更重要的原因是模型部署階段的實(shí)際收益。材料模型相比大語(yǔ)言模型規(guī)模小但算子種類(lèi)很雜圖卷積、注意力、自定義損失、大量 elementwise 操作混在一起。這種模型在傳統(tǒng)框架上直接推理kernel launch 開(kāi)銷(xiāo)占比很高小算子和數(shù)據(jù)搬運(yùn)經(jīng)常吃掉一半時(shí)間。CANN的圖融合功能能把相鄰小算子熔接到一起大幅減少kernel調(diào)度次數(shù)這一點(diǎn)我們?cè)趯?shí)測(cè)里感受非常明顯。另外一個(gè)考慮是國(guó)產(chǎn)化替代的趨勢(shì)。做科研項(xiàng)目除了性能和效果還需要考慮供應(yīng)鏈和生態(tài)的長(zhǎng)期穩(wěn)定性。把模型遷移到CANN上意味著從芯片到軟件棧完全自主可控后續(xù)做橫向項(xiàng)目或者商業(yè)化交付時(shí)不會(huì)被單一陣營(yíng)的生態(tài)綁定。2. 核心模塊與關(guān)鍵技術(shù)點(diǎn)2.1 材料數(shù)據(jù)集的構(gòu)造與特征工程模型效果的上限由數(shù)據(jù)決定。我們用了公開(kāi)材料數(shù)據(jù)庫(kù)的數(shù)據(jù)初始規(guī)模在十幾萬(wàn)量級(jí)但原始數(shù)據(jù)并不能直接喂給模型需要做幾層預(yù)處理。第一步是結(jié)構(gòu)清洗。去掉實(shí)驗(yàn)數(shù)據(jù)里明顯不合理的結(jié)構(gòu)比如原子間距離小于共價(jià)半徑之和80%的電荷不平衡超過(guò)閾值的以及重復(fù)結(jié)構(gòu)同一空間群同一晶格參數(shù)。這一步過(guò)濾掉的數(shù)據(jù)量比想象中多大約有8%的原始條目不合格。第二步是特征工程。每個(gè)原子節(jié)點(diǎn)我們拼接了12維特征包括族數(shù)、周期數(shù)、電負(fù)性、原子半徑、價(jià)電子數(shù)、第一電離能等這些均為數(shù)值型不需要?dú)w一化到同一量綱因?yàn)閳D神經(jīng)網(wǎng)絡(luò)本身能處理不同尺度。邊的特征我們用了鍵長(zhǎng)、鍵級(jí)以及兩個(gè)原子的電負(fù)性差值后者對(duì)離子鍵傾向的預(yù)測(cè)很重要。第三步是數(shù)據(jù)增強(qiáng)。材料數(shù)據(jù)天然存在類(lèi)別不平衡高帶隙絕緣體遠(yuǎn)遠(yuǎn)多于金屬如果直接訓(xùn)練模型對(duì)金屬體系完全無(wú)感。我們做了組合增強(qiáng)對(duì)同一結(jié)構(gòu)做晶胞擴(kuò)胞和超胞平移在不改變物理性質(zhì)的前提下增加樣本多樣性和原子排列之間的等變性擾動(dòng)。實(shí)際跑下來(lái)增強(qiáng)之后的模型在同一驗(yàn)證集上的帶隙預(yù)測(cè)MAE降低了約22%。2.2 生成模型選型變分自編碼器的空間變換結(jié)構(gòu)生成我們最終選了變分自編碼器VAE而不是擴(kuò)散模型。原因很實(shí)際VAE的隱空間是連續(xù)的可以直接做插值和屬性條件的向量運(yùn)算這對(duì)于目標(biāo)導(dǎo)向設(shè)計(jì)非常重要。比如想生成帶隙處于1.0到1.5電子伏特之間的半導(dǎo)體材料可以在隱空間里訓(xùn)練一個(gè)條件映射器把屬性向量映射成隱向量的偏移量再解碼為新結(jié)構(gòu)。解碼器輸出的是晶格矩陣和原子坐標(biāo)這個(gè)輸出必須滿足物理約束。我們額外加了三項(xiàng)正則一是原子間最小距離約束防止解碼出物理上不可能的重疊結(jié)構(gòu)二是晶格體積范圍約束把生成的胞體積限制在合理區(qū)間三是電負(fù)性匹配約束避免生成極端活潑的組合。這三項(xiàng)直接寫(xiě)進(jìn)自定義損失函數(shù)里相對(duì)權(quán)重靠超參搜索確定。VAE的訓(xùn)練損失包含重建損失、KL散度和物理約束損失。重建損失不是簡(jiǎn)單的坐標(biāo)MSE而是經(jīng)過(guò)一個(gè)可微的最近鄰匹配后計(jì)算結(jié)構(gòu)相似度。這個(gè)細(xì)節(jié)很關(guān)鍵因?yàn)橥粋€(gè)結(jié)構(gòu)可以有多種原子排列索引方式直接比較坐標(biāo)會(huì)讓模型學(xué)不到正確的幾何關(guān)系。訓(xùn)練過(guò)程用了三個(gè)模型并行一個(gè)編碼器、一個(gè)解碼器、一個(gè)輔助的屬性回歸器。屬性回歸器的作用是把隱向量和性質(zhì)對(duì)齊這樣后續(xù)做條件生成時(shí)才有依據(jù)。整個(gè)訓(xùn)練在CANN環(huán)境上跑了約300個(gè)epoch混合精度開(kāi)啟了之后單卡吞吐比FP32模式提升了2.6倍左右顯存占用也降了一半。2.3 性質(zhì)預(yù)測(cè)模型圖神經(jīng)網(wǎng)絡(luò)的等變關(guān)鍵生成模型負(fù)責(zé)廣撒網(wǎng)性質(zhì)預(yù)測(cè)模型負(fù)責(zé)準(zhǔn)確過(guò)濾。我們用的是信息傳遞風(fēng)格的圖神經(jīng)網(wǎng)絡(luò)和材料領(lǐng)域常用的CGCNN結(jié)構(gòu)類(lèi)似但針對(duì)等變性質(zhì)做了修改。關(guān)鍵點(diǎn)是向量特征的處理。普通圖神經(jīng)網(wǎng)絡(luò)把節(jié)點(diǎn)特征放在標(biāo)量空間里也就是每個(gè)原子只有一個(gè)特征向量所有鄰居消息加權(quán)求和。但對(duì)于力、偶極矩、彈性張量這些方向性物理量標(biāo)量表達(dá)會(huì)丟失信息。我們?cè)谙鬟f階段加入了球諧函數(shù)展開(kāi)的向量特征讓模型天然具備旋轉(zhuǎn)等變性。這個(gè)改動(dòng)不是錦上添花而是質(zhì)變?cè)谕獠框?yàn)證集上標(biāo)量模型對(duì)六方晶系材料的彈性模量預(yù)測(cè)誤差顯著高于等變模型因?yàn)榉较蛐孕畔⒈粯?biāo)量化之后完全丟失了。還有一個(gè)容易忽略的點(diǎn)是圖的構(gòu)建邊界。晶體結(jié)構(gòu)是周期性體系不能只連近鄰原子還要考慮周期性鏡像。我們把cutoff半徑設(shè)成8埃同時(shí)把周期性平移后的鄰居也納入圖中。這個(gè)半徑的選擇有講究太小會(huì)丟失長(zhǎng)程相互作用太大圖會(huì)變得過(guò)于稠密內(nèi)存和計(jì)算開(kāi)銷(xiāo)指數(shù)上升。8埃是在物理精度和計(jì)算效率之間折中的結(jié)果。模型結(jié)構(gòu)進(jìn)一步疊加了注意力機(jī)制讓中心原子對(duì)不同鄰居的貢獻(xiàn)做自適應(yīng)加權(quán)。這一步解決了傳統(tǒng)GNN中“平均化”導(dǎo)致的選擇性不足問(wèn)題特別是對(duì)不同配位環(huán)境下的同種元素注意力機(jī)制能區(qū)分它們的實(shí)際結(jié)合狀態(tài)。2.4 CANN加速層算子融合與混合精度策略這是整套系統(tǒng)里最“工程”的部分也是收益最明顯的一部分。把PyTorch模型遷到CANN上不是換一個(gè)后端那么簡(jiǎn)單要做四個(gè)層面的適配。第一層是算子映射檢查。模型里99%的標(biāo)準(zhǔn)算子卷積、全連接、歸一化、激活、Softmax都能直接映射到CANN的算子庫(kù)但自定義的物理約束損失函數(shù)里有一些自定義操作比如球諧展開(kāi)和周期性鄰居索引這些算子原生庫(kù)沒(méi)有需要通過(guò)自定義算子開(kāi)發(fā)接口補(bǔ)充。我們第一次跑算子適配檢查的時(shí)候發(fā)現(xiàn)總共有6個(gè)自定義算子無(wú)法映射后來(lái)重寫(xiě)了其中4個(gè)的純張量版用已有的基礎(chǔ)算子組合實(shí)現(xiàn)剩下2個(gè)寫(xiě)了自定義TBE算子。這一步建議前期就做靜態(tài)檢查不要等到訓(xùn)練中途才報(bào)錯(cuò)否則定位成本非常高。第二層是圖融合。CANN的編譯器會(huì)對(duì)計(jì)算圖做融合優(yōu)化但融合策略受限于算子本身的屬性。我們實(shí)測(cè)下來(lái)融合對(duì)elementwise類(lèi)小算子的收益最大。舉個(gè)例子一個(gè)后面跟著ReLU和Dropout的BatchNorm三層單獨(dú)跑是3次kernel launch融合后就變成1次。在節(jié)點(diǎn)數(shù)量很大的圖神經(jīng)網(wǎng)絡(luò)里這種小算子出現(xiàn)的頻率極高積少成多整體推理耗時(shí)能下降40%以上。第三層是混合精度。模型遷移時(shí)我們直接開(kāi)了混合精度模式FP16計(jì)算、FP32存儲(chǔ)。在CANN工具鏈上開(kāi)啟混合精度比想象中簡(jiǎn)單只要在模型轉(zhuǎn)換時(shí)指定精度檔位即可但需要注意loss scale的設(shè)定。材料性質(zhì)預(yù)測(cè)模型和圖像模型不同輸出范圍差異較大帶隙可能只有0.1電子伏特彈性模量則是幾十吉帕。如果loss scale不夠大梯度下溢會(huì)直接導(dǎo)致訓(xùn)練發(fā)散。我們最終把loss scale設(shè)到1024同時(shí)打開(kāi)動(dòng)態(tài)調(diào)整跑完整個(gè)訓(xùn)練過(guò)程沒(méi)有出現(xiàn)nan。第四層是內(nèi)存復(fù)用。材料模型的輸入圖規(guī)模不固定導(dǎo)致內(nèi)存調(diào)度容易碎片化。CANN的靜態(tài)內(nèi)存規(guī)劃模式可以在編譯期就把張量的生命周期分析好復(fù)用同一塊內(nèi)存給生命周期不重疊的算子使用。這一步對(duì)長(zhǎng)序列采樣特別有作用生成模型的反向采樣中會(huì)反復(fù)創(chuàng)建中間張量原先每秒只能采樣4個(gè)結(jié)構(gòu)內(nèi)存復(fù)用后提升到了每秒10個(gè)。3. 實(shí)操落地的完整流程與性能指標(biāo)3.1 環(huán)境配置與算子適配檢查先給出一份可以直接抄的配置清單。操作系統(tǒng)是64位Linux服務(wù)器內(nèi)核5.10以上AI加速卡驅(qū)動(dòng)已安裝版本和CANN工具包版本嚴(yán)格配套Python環(huán)境用CANN自帶的依賴(lài)包管理避免自己配PyTorch和CANN的版本錯(cuò)位。算子適配檢查用自帶的靜態(tài)分析工具跑了一遍這一步強(qiáng)烈建議所有遷移項(xiàng)目都做工具會(huì)解析模型結(jié)構(gòu)輸出一張算子映射表標(biāo)明每個(gè)算子是原生支持的、需要原型實(shí)現(xiàn)的、還是完全不支持的。我們第一次跑完發(fā)現(xiàn)原生算子覆蓋率為82%另外11%可以通過(guò)已有算子組合實(shí)現(xiàn)剩下7%必須手寫(xiě)。又花了一周把自定義算子補(bǔ)上最終覆蓋率打到了100%。# 算子適配檢查命令路徑按實(shí)際環(huán)境修改 /usr/local/Ascend/ascend-toolkit/latest/tools/msopstool/msopstool --check ./model_origin.onnx --output ./op_report.json3.2 用ATC工具完成模型轉(zhuǎn)換的三步經(jīng)驗(yàn)?zāi)Ｐ娃D(zhuǎn)換是CANN路線里最關(guān)鍵的一步相當(dāng)于把訓(xùn)練好的PyTorch模型翻譯成加速卡能高效執(zhí)行的中間表示。我們實(shí)際操作下來(lái)轉(zhuǎn)換步驟本身不復(fù)雜但要注意的問(wèn)題不少。第一步先把PyTorch模型導(dǎo)出成ONNX格式。這里有一個(gè)很隱蔽的坑模型帶自定義算子時(shí)導(dǎo)出過(guò)程可能會(huì)把自定義算子折疊成標(biāo)準(zhǔn)算子導(dǎo)致推理結(jié)果和原模型不一致。解決辦法是在導(dǎo)出前把自定義算子注冊(cè)到onnx的符號(hào)表里并在轉(zhuǎn)換后做輸出比對(duì)。我們用一組隨機(jī)輸入分別跑原模型和轉(zhuǎn)換后模型輸出誤差控制在1e-4量級(jí)才算通過(guò)。第二步是用ATC工具做模型轉(zhuǎn)換。核心參數(shù)是輸入節(jié)點(diǎn)的shape、輸出節(jié)點(diǎn)的名稱(chēng)、以及精度檔位。輸入shape建議直接設(shè)成固定尺寸因?yàn)橛袆?dòng)態(tài)shape的模型在轉(zhuǎn)換時(shí)要做多檔適配生成的模型文件和推理性能都會(huì)有折損。材料模型的輸入是圖結(jié)構(gòu)動(dòng)態(tài)性主要在鄰居數(shù)上我們通過(guò)padding統(tǒng)一到了固定節(jié)點(diǎn)數(shù)。第三步是保存推理配置。轉(zhuǎn)換產(chǎn)物是一個(gè)模型文件加一份配置json里面包含了輸入輸出格式、動(dòng)態(tài)批次、精度類(lèi)型等信息。這一步很多人會(huì)忽略直接拿著模型文件就去推理結(jié)果輸入輸出對(duì)齊錯(cuò)誤。配置json建議保留后續(xù)做推理服務(wù)的時(shí)候能直接引用。3.3 推理性能實(shí)測(cè)結(jié)果對(duì)比這是項(xiàng)目組最滿意的部分。我們?cè)谕慌蜻x材料上對(duì)比了純PyTorch CPU推理、PyTorch CUDA模擬推理同型號(hào)消費(fèi)級(jí)GPU以及CANN加速后的性能?；鶞?zhǔn)是1000條候選材料結(jié)構(gòu)每條結(jié)構(gòu)節(jié)點(diǎn)數(shù)平均64邊數(shù)平均320。表格里的數(shù)據(jù)是實(shí)測(cè)均值推理方式單條平均延遲ms1000條總耗時(shí)s相對(duì)CPU加速比PyTorch CPU458458.01xPyTorch GPU18.518.524.8xCANN混合精度7.87.858.7x單條7.8毫秒是什么概念原來(lái)CPU跑完一輪篩選要等一個(gè)小時(shí)現(xiàn)在不到一分鐘就能看到候選列表。留給材料研究者的等待時(shí)間從“可以去泡杯咖啡”變成了“刷新頁(yè)面就出結(jié)果”。這個(gè)量級(jí)的加速才能讓交互式材料設(shè)計(jì)真正落地。注意延遲還可以更低。我們測(cè)了batch推理模式把多條候選打包成一個(gè)大圖同時(shí)推理單條延遲進(jìn)一步降到4.2毫秒。但代價(jià)是顯存占用線性上漲batch設(shè)到256時(shí)峰值顯存已經(jīng)超過(guò)16GB實(shí)際部署時(shí)需要在吞吐和內(nèi)存之間找平衡。3.4 端到端工作流從生成到驗(yàn)證的閉環(huán)整套系統(tǒng)跑通后的流程是這樣的用條件VAE在目標(biāo)帶隙區(qū)間生成候選結(jié)構(gòu)經(jīng)過(guò)物理過(guò)濾器初篩密度、電負(fù)性差、晶胞體積剩下約35%進(jìn)入圖神經(jīng)網(wǎng)絡(luò)做性質(zhì)預(yù)測(cè)。預(yù)測(cè)結(jié)果按多目標(biāo)打分排序取Top 50進(jìn)入第一性原理驗(yàn)證。驗(yàn)證結(jié)果回傳給數(shù)據(jù)庫(kù)再作為訓(xùn)練數(shù)據(jù)進(jìn)入下一輪迭代。閉環(huán)設(shè)計(jì)帶來(lái)的最大收益是模型越用越準(zhǔn)。每一輪實(shí)驗(yàn)驗(yàn)證都在給訓(xùn)練集補(bǔ)充“硬樣本”也就是那些生成模型打分很高但計(jì)算驗(yàn)證不通過(guò)的反例。把這些反例加進(jìn)訓(xùn)練集下一輪模型就知道哪些組合是陷阱候選質(zhì)量的通過(guò)率會(huì)逐步上升。我們記錄了三輪迭代的數(shù)據(jù)初始通過(guò)率27%第二輪34%第三輪41%。4. 實(shí)際踩過(guò)的坑與排查技巧4.1 自定義算子在轉(zhuǎn)換時(shí)被錯(cuò)誤折疊這是整個(gè)項(xiàng)目里最耗時(shí)的一個(gè)坑前前后后花了一周才定位。癥狀是模型在PyTorch里推理完全正常轉(zhuǎn)換到CANN上之后輸出變成噪聲。排查過(guò)程極其痛苦先是懷疑混合精度引入誤差把精度檔位降到FP32還是不行然后懷疑數(shù)據(jù)預(yù)處理不一致對(duì)照了很多遍也沒(méi)問(wèn)題最后用逐層輸出比對(duì)工具定位到某一層——自定義的球諧展開(kāi)算子被折疊成了普通全連接。原因在于導(dǎo)出ONNX時(shí)球諧展開(kāi)的計(jì)算圖被優(yōu)化成了它認(rèn)為等效的矩陣乘法但數(shù)學(xué)上兩者在非線性激活區(qū)域并不嚴(yán)格一致。解決辦法有兩個(gè)一是導(dǎo)出前顯式關(guān)閉ONNX的優(yōu)化pass二是寫(xiě)一個(gè)自定義算子模塊并在ONNX中注冊(cè)強(qiáng)制它作為整體導(dǎo)出。我們用了第二種方案徹底解決了問(wèn)題。4.2 混合精度訓(xùn)練在損失scale上的教訓(xùn)第一次開(kāi)混合精度訓(xùn)練跑了20個(gè)epoch直接發(fā)散loss飆到NaN。排查之后發(fā)現(xiàn)問(wèn)題出在loss scale設(shè)得太低。材料性質(zhì)數(shù)值范圍跨度大一個(gè)batch里既有0.1量級(jí)的帶隙又有幾百GPa的模量梯度幅值差距能到三個(gè)數(shù)量級(jí)。固定loss scale為128時(shí)小梯度直接下溢成0調(diào)成1024之后又擔(dān)心溢出還好動(dòng)態(tài)loss scale機(jī)制自動(dòng)調(diào)整在訓(xùn)練早期快速爬升到了1024并穩(wěn)定下來(lái)。這個(gè)教訓(xùn)給團(tuán)隊(duì)的啟示是不要照搬圖像任務(wù)里的默認(rèn)參數(shù)。混合精度的優(yōu)勢(shì)毋庸置疑但loss scale和輸出范圍強(qiáng)相關(guān)需要針對(duì)任務(wù)重新標(biāo)定。4.3 內(nèi)存墻問(wèn)題比算力墻更早出現(xiàn)CANN的算力優(yōu)勢(shì)很明顯但吞吐上來(lái)了之后內(nèi)存帶寬立刻變成瓶頸。具體現(xiàn)象是推理延遲已經(jīng)降到很低了再壓batch size上去延遲沒(méi)有繼續(xù)下降反而開(kāi)始上升同時(shí)CPU側(cè)的數(shù)據(jù)預(yù)處理來(lái)不及供數(shù)GPU利用率掉到60%以下。排查之后發(fā)現(xiàn)瓶頸有三個(gè)一是數(shù)據(jù)加載進(jìn)程用了Python多進(jìn)程進(jìn)程切換開(kāi)銷(xiāo)大二是圖數(shù)據(jù)格式是Python對(duì)象列表序列化反序列化代價(jià)高三是在Host側(cè)做鄰居采樣然后在設(shè)備側(cè)重新建圖這部分完全串行化了。改動(dòng)方案是把圖數(shù)據(jù)轉(zhuǎn)成統(tǒng)一張量格式預(yù)處理階段離線完成鄰居索引構(gòu)建推理時(shí)直接加載Device可讀格式。修改之后端到端吞吐又漲了約35%。4.4 數(shù)據(jù)不平衡導(dǎo)致生成模型偏科材料數(shù)據(jù)的類(lèi)別分布天然不均衡帶隙預(yù)測(cè)任務(wù)里高帶隙樣本占比接近七成。如果不做處理模型對(duì)金屬材料預(yù)測(cè)精度會(huì)退化到“完全沒(méi)學(xué)過(guò)”的程度。我們用了混合策略一是在特征工程階段做平衡采樣保證每個(gè)batch里高低帶隙樣本比例接近二是在訓(xùn)練損失里給少數(shù)類(lèi)樣本加權(quán)重三是用Focal Loss替代普通MSE降低易分樣本的損失貢獻(xiàn)。三個(gè)措施疊加之后金屬類(lèi)別的預(yù)測(cè)MAE最終和絕緣體類(lèi)別在同一水平線上。這個(gè)問(wèn)題的隱蔽性在于整體MAE指標(biāo)看起來(lái)依然不錯(cuò)不細(xì)分類(lèi)別根本發(fā)現(xiàn)不了。做材料AI一定不要只盯全局指標(biāo)要按材料類(lèi)別、空間群、元素主族分別做分層評(píng)估才能真正暴露模型的短板。4.5 模型更新與緩存失效的一致性系統(tǒng)上線運(yùn)行后訓(xùn)練數(shù)據(jù)不斷增補(bǔ)模型版本也在迭代。但推理服務(wù)如果一直掛著舊模型前端看到的候選質(zhì)量和新模型的差距會(huì)越拉越大。我們最初的做法是人肉定時(shí)重啟服務(wù)加載新模型效率低還容易漏。后來(lái)改成自動(dòng)模型版本管理新模型訓(xùn)練完自動(dòng)評(píng)估指標(biāo)優(yōu)于當(dāng)前線上版本則發(fā)布推理服務(wù)通過(guò)監(jiān)聽(tīng)文件變動(dòng)實(shí)現(xiàn)熱加載。模型加載期間的服務(wù)調(diào)用會(huì)暫時(shí)路由到舊版本加載完成后再切換整個(gè)過(guò)程完全無(wú)縫。這個(gè)機(jī)制保證了系統(tǒng)在持續(xù)學(xué)習(xí)的同時(shí)線上服務(wù)不會(huì)出現(xiàn)中斷或新舊模型混用的問(wèn)題。5. 這套系統(tǒng)后續(xù)還能怎么擴(kuò)展目前單點(diǎn)性能已經(jīng)達(dá)標(biāo)了但整套系統(tǒng)的天花板還很遠(yuǎn)。從材料研究的角度下一步可以做跨尺度建模不但預(yù)測(cè)塊體材料的性質(zhì)還要把表面、界面、缺陷、非晶態(tài)這些復(fù)雜體系納入進(jìn)來(lái)。從AI的角度可以引入主動(dòng)學(xué)習(xí)策略讓模型主動(dòng)挑選“最不確定”的候選去做實(shí)驗(yàn)驗(yàn)證用最少的實(shí)驗(yàn)代價(jià)換取最大的模型提升。從工程的角度CANN加速層目前只服務(wù)了推理訓(xùn)練階段還有優(yōu)化空間。圖神經(jīng)網(wǎng)絡(luò)在訓(xùn)練時(shí)動(dòng)態(tài)圖變化頻繁CANN的圖優(yōu)化每次都要重算這部分性能比推理差不少。如果能把訓(xùn)練側(cè)的圖編譯策略做起來(lái)整套系統(tǒng)的迭代速度還會(huì)上一個(gè)臺(tái)階。我個(gè)人在實(shí)際操作中最深刻的體會(huì)是AI for Science這類(lèi)項(xiàng)目模型精度再高算力跟不上就不可能落地。材料發(fā)現(xiàn)的價(jià)值鏈非常長(zhǎng)光有好的生成模型不夠光有快的推理引擎也不夠只有把模型、數(shù)據(jù)、加速、驗(yàn)證閉環(huán)串起來(lái)AI才能真正變成材料實(shí)驗(yàn)室里的一臺(tái)不知疲倦的煉金爐。這套系統(tǒng)的構(gòu)建過(guò)程不算輕松但路線已經(jīng)趟通了希望這篇總結(jié)能幫你少走幾步彎路。