全球數(shù)值天氣預(yù)報(bào)正邁入千米級(jí)非靜力模擬的新時(shí)代。然而,高分辨率非靜力模式面臨三大核心挑戰(zhàn):第一,模式在不連續(xù)區(qū)域中容易出現(xiàn)非物理振蕩,影響模擬的真實(shí)性;第二,能精細(xì)刻畫垂直運(yùn)動(dòng)的非靜力模式,其時(shí)間積分步長(zhǎng)受限于垂直CFL條件,導(dǎo)致計(jì)算效率受制明顯;第三,隨著GPU為代表的異構(gòu)算力的爆發(fā)式增長(zhǎng),異構(gòu)計(jì)算已成為高性能計(jì)算(HPC)的核心發(fā)展方向,然而氣象模式對(duì)異構(gòu)算力的適配不足。這些瓶頸制約著未來全球公里級(jí)模式的發(fā)展。
針對(duì)上述難題,中國(guó)科學(xué)院大氣物理研究所地球系統(tǒng)數(shù)值模擬與應(yīng)用全國(guó)重點(diǎn)實(shí)驗(yàn)室陳曦研究員團(tuán)隊(duì)自主研發(fā)了高性能動(dòng)力框架 LMARSpy。該動(dòng)力框架采用A網(wǎng)格來儲(chǔ)存變量,使用專為大氣定制的低馬赫數(shù)黎曼求解器(LMARS)求解可壓縮歐拉方程組,兼具高精度與高穩(wěn)定性。該框架從算法層面專為GPU優(yōu)化,是一個(gè)面向異構(gòu)計(jì)算架構(gòu)的非靜力動(dòng)力核心,集成了梯度保持的單調(diào)性限制器與具有守恒特性的垂直隱式求解器。
為抑制不連續(xù)區(qū)域的數(shù)值振蕩,該研究開發(fā)了能在不連續(xù)區(qū)域保持高梯度的單調(diào)性限制器(圖1)。另外,為突破非靜力模式時(shí)間步長(zhǎng)瓶頸,該研究開發(fā)了解除垂直CFL限制的垂直隱式求解器。在垂直網(wǎng)格距遠(yuǎn)小于水平網(wǎng)格距的情況下,使用該求解器可以獲得數(shù)量級(jí)的計(jì)算速度提升(圖2)。
為增強(qiáng)異構(gòu)算力的可移植性,LMARSpy采用Python語(yǔ)言開發(fā),模式和并行框架代碼與高性能計(jì)算后端分離,可靈活切換不同計(jì)算后端(如PyTorch等)。該設(shè)計(jì)使其同時(shí)支持跨節(jié)點(diǎn)MPI多CPU計(jì)算與跨節(jié)點(diǎn)NCCL多GPU高性能計(jì)算,兼顧通用性與擴(kuò)展性。該研究還通過即時(shí)編譯技術(shù)(JIT)讓代碼的運(yùn)行速度接近機(jī)器碼水平,實(shí)現(xiàn)異構(gòu)計(jì)算機(jī)的高性能計(jì)算。另外,LMARSpy具有很強(qiáng)的擴(kuò)展性,在跨節(jié)點(diǎn)并行測(cè)試中的擴(kuò)展性一直保持在90%以上,能在大規(guī)模GPU集群上高效并行計(jì)算(圖3)。

圖1:Robert不連續(xù)熱泡實(shí)驗(yàn)。上方為未使用限制器的結(jié)果,下方為使用限制器的結(jié)果。第1~2列為不同時(shí)刻的模擬結(jié)果,第3列為沿垂直中線的截面。

圖2:Robert連續(xù)熱泡實(shí)驗(yàn)。上方為顯式方案結(jié)果,下方為使用垂直隱式求解器的結(jié)果,第1~3列為不同時(shí)刻的模擬結(jié)果。

圖3:左圖為將二維的Robert連續(xù)熱泡實(shí)驗(yàn)拓展到三維的模擬結(jié)果,以此來進(jìn)行大規(guī)模網(wǎng)格的異構(gòu)計(jì)算擴(kuò)展性測(cè)試。右圖為在不同GPU數(shù)量下,保持每GPU工作負(fù)載不變時(shí)的單位工作負(fù)載墻時(shí)間,虛線為256個(gè)CPU核心在相同任務(wù)規(guī)模下的計(jì)算時(shí)間基準(zhǔn)。
該研究為全球公里級(jí)非靜力模式中遇到的關(guān)鍵問題提供了適配GPU的算法解決方案,并為構(gòu)建新一代大氣環(huán)流模式(GCM)以及區(qū)域模式奠定了堅(jiān)實(shí)基礎(chǔ)。同時(shí)借助其基于機(jī)器學(xué)習(xí)框架的可擴(kuò)展特性,LMARSpy還能與AI大模型深度融合,為推動(dòng)數(shù)值天氣預(yù)報(bào)邁向智能化新時(shí)代提供強(qiáng)大引擎。
相關(guān)成果發(fā)表于《Journal of Advances in Modeling Earth Systems (JAMES)》,中國(guó)科學(xué)院大氣物理研究所研究生張偉康為第一作者,中國(guó)科學(xué)院大氣物理研究所陳曦研究員為通訊作者。研究得到國(guó)家自然科學(xué)基金(42275174,42288101)和國(guó)家重點(diǎn)研發(fā)計(jì)劃(2022YFF0802000,2022YFF0802001,2024YFB4204800)等項(xiàng)目的共同資助,以及國(guó)家重大科技基礎(chǔ)設(shè)施“地球系統(tǒng)數(shù)值模擬裝置”(EarthLab)等單位的支持。
論文信息:Zhang,?W.,&?Chen,?X.?(2025).?LMARSpy: A GPU-ready nonhydrostatic?dynamical core with a sharpness-preserving monotonicity limiter and a conservative vertical implicit solver.?Journal of Advances in Modeling Earth Systems,?17,e2025MS005056.?https://doi.org/10.1029/2025MS005056