訓練從數周壓到數小時：他們把通用原子勢訓練帶入Exascale時代

作者 | 論文團隊

編輯丨ScienceAI

在材料發現、催化設計、能源體系模擬、藥物研發等方向，原子尺度模擬一直都是底層工具，但傳統 DFT 計算成本極高，很難支撐大規模篩選。過去幾年，機器學習原子間勢雖然發展很快，但要真正做到「一個模型覆蓋材料、分子、催化、MOF 等多個領域」，并不容易。

原因很直接：這類模型訓練不僅要預測能量，還要通過自動微分嚴格求出原子力和應力，訓練中需要二階導數；同時，為了保證分子動力學的穩定性，又往往必須堅持 FP32 精度。二階訓練、高精度要求、超大原子圖，這三重約束疊加在一起，讓十億級通用原子勢的訓練長期停留在「理論上可擴、工程上難訓」的狀態。

最近，來自中國科學院計算技術研究所的研究團隊把通用機器學習原子間勢（uMLIP）的訓練規模推到了一個新量級，提出了十億級通用原子勢模型 MatRIS-MoE，以及配套的分布式訓練框架 Janus，在兩臺 Exascale 超算上實現了 1.2/1.0 EFLOPS 的單精度峰值性能，并把原本需要數周的訓練過程壓縮到數小時。

研究論文：Breaking the Training Barrier of Billion-Parameter Universal Machine Learning Interatomic Potentials

論文地址：https://arxiv.org/pdf/2604.15821v1

兼顧物理先驗與擴展性的模型架構

MatRIS-MoE 建立在 MatRIS 的不變架構之上，將原子體系表示為圖結構，模型對原子、成對距離和三體角度進行嵌入外，還加入任務嵌入、charge/spin 嵌入和全局特征嵌入，將不同數據集、DFT 泛函和體系類型對齊到統一表示空間中。相比原始 MatRIS，它不再只是單任務勢函數，而是面向多域統一建模的 universal MLIP。

更關鍵的是，作者沒有簡單堆大參數，而是引入了 MoE。在 MatRIS-MoE 中，MoE 層被插入到注意力前后，分別處理消息構造和特征更新。其一個重要設計是：路由不依賴瞬時坐標，而是按元素類型進行 Top-K 專家激活。這樣既能讓不同專家學習不同元素和化學環境下的規律，提升跨域表達能力，也能保持專家激活穩定，從而有助于維持連續、平滑的勢能面。

訓練上，這項工作依然堅持「保守式」路線：模型先預測總能量，再通過自動微分得到力和應力，而不是單獨開力頭直接擬合，這對保證物理一致性很重要。同時，作者還設計了 multi-task robust loss，在每個任務內部統計 batch loss 的均值和方差，并對離群樣本做平滑降權，以減輕異構任務之間的干擾。

兼顧 MoE 稀疏執行與二階反傳的訓練框架

如果說模型解決的是「怎么表達」，那么 Janus 解決的就是「怎么訓練」。

Janus 的核心是一個叫 FS-3D 的統一執行單元，把 FSDP、FSGP 和 FSEP 三種并行機制揉到了一起。簡單說，FSDP 負責切參數，降低模型參數、梯度和優化器狀態的靜態顯存占用；FSGP 負責切圖，把一個大原子圖劃分到多張卡上，降低激活開銷；FSEP (以 LAER-MoE 為基礎) 則專門面向 MoE，把專家參數分布到不同設備上，只在需要時恢復活躍專家。對于既有超大圖、又有 MoE、還要做 double backward 的 uMLIP 來說，這種三維統一分片的設計，是這篇工作的關鍵工程創新之一。

為了進一步解決 MoE 的低效問題，作者還提出了 JIT planning。傳統 MoE 訓練往往會保留很多當前 step 根本用不到的專家；作者利用體系靜態屬性決定專家激活的特點，在每個訓練 step 開始前，先對所有 MoE 層做批量路由，統計各專家 token 負載，再基于負載做局部規劃和全局合并，只恢復當前真正活躍的專家，并按負載盡量均衡地放到不同 rank 上。

更難的是，uMLIP 訓練不是普通的一次前向一次反向，而是包含 前向、一次反向、二次反向三個階段。Janus 為此專門實現了一個「二階訓練大模型」的運行時系統：參數按需恢復，前向階段記錄執行順序，后續階段復用順序做 prefetch 和 overlap，梯度同步則延遲到最終反向再進行。

結果與展望

結果也確實夠亮眼。論文使用了覆蓋分子、材料、催化表面、分子晶體和 MOF 的 4.73 億個原子構型，對應約 3.6 萬億條邊；最大模型規模達到 11.5B 參數，2.89B active parameters，支持處理多個領域的科學任務。最終，在兩臺 Exascale 平臺上，系統實現了超過 90% 的弱擴展并行效率，峰值達到 1.2 EFLOPS，把十億級通用原子勢訓練第一次真正推到了超算級可用的階段。

從這個角度看，這篇工作真正重要的地方，不只是「把一個模型訓到了 11.5B」，也不只是「跑到了 1.2 EFLOPS」，而是它證明了一件事：通用原子勢也可以像大模型一樣，被系統化地擴展、訓練和部署。

對 AI for Science 來說，這可能比單點精度提升更重要。因為從這里開始，通用原子勢不再只是論文里的 benchmark 選手，而是在向真正的科學基礎設施靠近。

作者介紹

中國科學院計算技術研究所及中國科學院大學周遠昌、王宏宇、杜奕明、汪焱、李明真、胡思宇為本文共同第一作者；中國科學院計算技術研究所賈偉樂研究員為本文通訊作者。賈偉樂 2020 年獲國際高性能計算應用領域最高獎戈登貝爾獎，2022 年再次入圍該獎項；相關成果入選兩院院士評選國內十大科技進展新聞，并獲得中國超算年度最佳應用獎及中國計算機學會高性能計算青年科技人才獎等多項榮譽。