ARTICLE · INTELLIGENCE

战地情报 · 详情页

来自尧图项目组的一线实战观察与深度解析

一个模型通吃五类化学体系:fairchem(ocp)的UMA机制、数据资产与选型指南

一个模型通吃五类化学体系:fairchem(ocp)的UMA机制、数据资产与选型指南

一个模型通吃五类化学体系:fairchem(ocp)的UMA机制、数据资产与选型指南

【免费下载链接】ocpFAIR Chemistry's library of machine learning methods for chemistry项目地址: https://gitcode.com/GitHub_Trending/oc/ocp

你手头有 800 万条 DFT(密度泛函理论,原子模拟的"高精度标尺")计算结果,但真正算一次新的反应路径还是要等数天——Open Catalyst Project(ocp,现名fairchem,FAIR Chemistry 团队维护的机器学习势仓库)要消掉的就是这道成本鸿沟。它把五大开放数据集、约5.2 亿条 DFT 计算喂进一个预训练模型UMA(Universal Model for Atoms,面向原子的通用模型),再借 ASE(Atomic Simulation Environment,原子模拟框架)计算器接口,把能量、力、应力的预测直接接进你现有的弛豫与分子动力学流程。这篇文章讲清三件事:UMA 为什么"一个模型、多套精度",六大开源数据集各自能干什么,以及从安装到批量筛选的完整上手路径。

上图是仓库自带基准(基准脚本):横轴是在 NVIDIA H200 上跑 1296 原子冰 Ih 体系一个 ASE NVE 步的平均耗时,纵轴是 OMol25 公开验证集(272 万条 ωB97M-V 结构)上的力平均绝对误差(force MAE)。这张图也是你判断"要不要从 DFT 切到机器学习势"的第一手依据。

UMA 的核心机制:一个模型如何同时服务五个化学领域

架构关键词 MoLE:大参数、快推理、能量守恒

UMA 是一个等变图神经网络(equivariant GNN:输入结构旋转/平移后,输出力、应力等矢量量按同样规则变换,天然满足物理对称性),核心创新是MoLE(Mixture of Linear Experts,线性专家混合)架构:模型总参数可以做到很大,但推理时通过一个动态路由的单一输出头只激活很小一部分参数——官方给出的两个版本,uma-s 激活 6.6M/总 290M,uma-m 激活 50M/总 1.4B,在材料、分子、催化等多个领域达到当时的 SOTA(state-of-the-art)水平。

路由需要四组输入,这四组输入恰好决定了模型"以什么口径说话":

  • task:指明要模拟的 DFT 数据集/理论水平(omol、omat、oc20、oc25、odac、omc、oc22);
  • 总电荷与自旋多重度(仅omol任务使用,默认 0/1);
  • 元素组成:对每个原子的嵌入做平均,得到整体组成嵌入。

为什么 task 选错,能量就没法直接比

每个数据集都是用不同 DFT 泛函、不同赝势算出来的,UMA 为每个 task 学了一个独立的嵌入,输出能量对应各自的原生理论(level of theory)。官方文档明确提醒:不同 task 的能量不能当作同一次参考计算的结果直接相减(快速上手文档)。一句话:task 选错,不是结果不准,而是口径错了。

UMA 版本与任务覆盖

版本时间任务数关键变化
UMA 1.02025-055首发oc20/omat/omol/odac/omc;MoLE + 等变 GNN,5 亿条 DFT 训练
UMA 1.120255修复能量尺寸外推(size extensivity)缺陷
UMA 1.2(含 1p2、1p2p1)2026-037新增oc25、oc22;纳入 OC22、OC25 与扩展 OMol25(含聚合物 OPoly26);官方称小模型 turbo 模式提速约 50%、Open Molecules 测试集精度提升约 40%

1.0 检查点已不受支持(加载会直接抛错);完整记录见 UMA 发布日志,当前库版本 fairchem-core 2.22.0 还进一步加快了 UMA 推理。

数据与模型资产:六大开源数据集一张表看懂

ocp 仓库的fairchem-data-*子包(数据工具源码)负责生成各数据集的输入构型,标签(能量/力/应力)以 LMDB(内存映射数据库,支持随机高效读取)形式分发。六个核心资产速览:

数据集领域规模理论水平(DFT)存储形式
OC20异相催化(82 种吸附质、约 1.15 万块体)1.33 亿+ 帧 / 约 46 万条吸附-表面弛豫VASP RPBES2EF 轨迹 + 预计算 LMDB;提供 200K/2M/20M/全量四档训练子集
OC22氧化物电催化剂约 6.2 万氧化物体系VASP PBE+U(自旋极化)预计算 LMDB
OC25固-液界面(显式溶剂)约 800 万条 DFT、150 万溶剂环境、平均 144 原子、88 元素VASP RPBE+D3ASE LMDB(*.aselmdb)
OMol25有机/无机分子、配合物、电解质1 亿+ 单点与弛豫ORCA6 wB97M-V/def2-TZVPDASE LMDB
OMat24无机块体材料训练 1.07M + 验证 1.02M 结构VASP PBE/PBE+UASE LMDB
ODAC23/25MOF 直接空气捕集CO₂/H₂O 吸附VASP PBE+D3ASE LMDB

催化方向的任务命名约定值得记住:S2EF(Structure to Energy and Forces,给定结构预测能量和力)、IS2RE(初始结构→弛豫能量)、IS2RS(初始结构→弛豫结构),验证集按 id / ood_ads(未见吸附质)/ ood_cat(未见催化剂)/ ood_both 四种子集切分,专门量化外推能力。

两个容易踩的坑:其一,UMa/OMat24 系的模型用 DFT+U 总能量标签训练,与 Materials Project 的计算设置不兼容(README 有专门提示),做形成能、凸包能这类差值计算时要小心;其二,OC20 训练集标签是吸附能,OC22 起改用总能量——跨数据集比较能量前先看标签口径。

从零到第一个结果:安装 fairchem 并跑通第一次弛豫

安装分三步(安装文档):

pip install fairchem-core # V2 已去掉 torch-geometric 等依赖 huggingface-cli login # UMA 检查点为门控仓库,需先申请访问 pip install "fairchem-core[extras]" # 可选:Ray 批量/多卡、Parquet 轨迹等

之后所有计算遵循同一个四步模式:构建 ASE Atoms 对象 → 加载 UMA → 选 task 挂 FAIRChemCalculator → 调 ASE 拿能量/力或起动力学。以 CO 在 Cu(100) 上弛豫为例:

from ase.build import fcc100, add_adsorbate, molecule from ase.optimize import LBFGS from fairchem.core import pretrained_mlip, FAIRChemCalculator predictor = pretrained_mlip.get_predict_unit("uma-s-1p2p1", device="cuda") calc = FAIRChemCalculator(predictor, task_name="oc20") slab = fcc100("Cu", (3, 3, 3), vacuum=8, periodic=True) add_adsorbate(slab, molecule("CO"), 2.0, "bridge") slab.calc = calc LBFGS(slab).run(fmax=0.05, steps=100)

换成分子体系时,记得把电荷和自旋写进atoms.info({"charge": 0, "spin": 1}),这是omol任务的硬性输入;周期性材料则选omat,它额外预测应力,因此 ASE 的FrechetCellFilter可以同时弛豫原子位置与晶胞(Hello World 教程)。

三条进阶路线:跑分子动力学、做高通量、搜过渡态

分子动力学:ASE 直连或 Hydra 配置,都支持 NVT/NPT

FAIRChemCalculator实现了标准 ASE 计算器接口,所以 Langevin、VelocityVerlet 等 ASE 动力学类可以直接驱动,仓库同时提供 Hydra 配置(configs/uma/md/,nvt.yaml为 Langevin 300K、npt.yaml为 Berendsen 300K/1bar),一行fairchem -c configs/uma/md/nvt.yaml即可跑起来,支持断点续跑(checkpoint + resume 配置)、Parquet 轨迹落盘和 SLURM 提交。NVT/NPT/NVE/Bussi 等系综适配器一览见 MD 文档。注意turbo推理档(编译快速路径 + TF32)需要 Ampere 及以后的 GPU,旧卡回退inference_settings: default。

大规模与批量:workers 参数背后的 Ray 框架

多卡/多机并行不需要你手配 MPI:给get_predict_unit加workers=N,底层 Ray 自动分片。官方称该路线已验证到百卡级、百万原子体系,批量小体系弛豫在单卡 H100 上最高有4 倍相对串行的加速(InferenceBatcher 文档);更大规模的既有 MD 工作流还可以接 LAMMPS 集成。

过渡态搜索:CatTSunami 框架与 NEB 参考数据

对催化反应路径,仓库提供 CatTSunami 框架:用在 OC20 上训练的 GNN 做 nudged elastic band(NEB," nudged 弹性带",过渡态搜索的常用离散化方法)帧集的高通量枚举,并附932 条收敛 DFT NEB参考(解离、脱附、迁移三类反应,覆盖 2827 次尝试)用于验证模型。同类应用包还有 AdsorbML(最低吸附能)、fastcsp(晶体结构预测)、ocx(实验-计算衔接,应用源码)。

怎么选模型、task 与验证集:一张表加三条红线

模型选型(UMA 指南):

模型激活/总参数定位
uma-s-1p2p1(UMA 1.2.1 小模型)6.6M / 290M默认选择:多数基准 SOTA,速度与显存最优
uma-m-1p150M / 1.4B全指标最高精度,代价是更慢、更吃显存

task 按科学领域选,而不是按"哪种读得进我的结构"选:异相催化oc20、氧化物电催化oc22、(电)催化显式溶剂oc25、无机材料omat、分子/聚合物omol、分子晶体omc、MOF 捕集odac。注意oc25/oc22仅 UMA 1.2 起可用。

三条红线:

  • 跨 task 能量不可直接比较——泛函、赝势、色散校正都不同;
  • 除omol外,各 task 默认期望中性、闭壳层输入(电荷/自旋恒为 0/1),开放壳层体系要格外小心;omat不能自选磁性态;
  • omol训练数据全部是非周期体系,周期性结构慎用;UMa 加速的只是"探索",关键结论仍要用参考数据或更高精度方法复核(官方文档原话强调)。

最后按你的角色给行动建议:

入门者——装好fairchem-core,申请 UMA 访问后跑 Hello World 里的 CH₂ 自旋隙和 Fe 晶胞弛豫两个例子,先建立"结构→预测"的完整手感,再碰动力学。

催化/电催化研究者——按体系锁定 task(氧化物选oc22、界面带溶剂选oc25、通用金属表面选oc20),用四种子验证集(id/ood_ads/ood_cat/ood_both)评估泛化,需要反应能垒时直接读 CatTSunami 的教程 notebook 和 932 条 DFT 参考。

工程化开发者——从turbo推理档起步,用workers=N+ Ray 或 LAMMPS 扩展体系规模,批量小体系上 InferenceBatcher;要复现官方数字,以 benchmarks/omol 里的环境与脚本为准(Python 3.12 + PyTorch 2.13 + ASE 3.29 + fairchem-core 2.22.0,单卡 H200)。

【免费下载链接】ocpFAIR Chemistry's library of machine learning methods for chemistry项目地址: https://gitcode.com/GitHub_Trending/oc/ocp

创作声明:本文部分内容由AI辅助生成(AIGC),仅供参考

RELATED READING

延伸阅读

更多一线实战笔记与深度复盘,助您持续精进