FAIRChem 上手指南:用 UMA 机器学习势把催化剂 DFT 筛选提速千倍

FAIRChem 上手指南:用 UMA 机器学习势把催化剂 DFT 筛选提速千倍

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

FAIRChem 是 Meta FAIR Chemistry 团队的化学机器学习库,把开源数据集、预训练原子模型(UMA)与 ASE/LAMMPS 工作流整合在一个仓库里。做催化、材料、分子动力学的研究者可以用几行 Python 替代大部分 DFT 计算;本文给出安装、首个计算、一个完整场景走查与可复现的性能数字。

DFT 的瓶颈与 MLIP 的量化收益

DFT 算得准,但贵:一次 NEB 过渡态搜索可能烧掉数天 GPU,上千候选材料的筛选则以月计。FAIRChem 仓库里的 CatTSunami 应用给出了一个可引用的对照:在 932 条 DFT NEB 验证集(覆盖 98 个反应)上,"全 ML" 管线(图神经网络预测初/末态 + 弛豫 + NEB)相对全 DFT 流程取得2200 倍加速,70% 的预测过渡态与 DFT 一致,且没有针对该任务做额外训练。官方催化教程中,uma-s-1p2p1在催化工作流上约1000 倍于 DFT的速度,精度满足初筛级要求。

图1:932 条 RPBE DFT NEB 验证集上,不同 ML/DFT 混合策略的加速比与成功率;全 ML 策略为 2200 倍加速、70% 成功率

这个策略的分工是:ML 模型负责"看很多"(初筛、构型搜索、长轨迹),DFT 只负责"看重点"(对少数候选做最终确认)。下面的内容会说明 FAIRChem 里哪些部分承担这两件事。

能力地图:一个 task 参数切换 7 个化学领域

UMA 是一个模型、多个"任务档位":每个 task 对应一套 DFT 训练数据与理论层级,推理时用task_name切换即可。常用档位如下(来源:docs/core/uma.md):

两个使用注意:不同 task 的能量来自不同 DFT 层级,跨 task 的能量不能直接相减比较;omol还需要在输入里给出总电荷与自旋重数。

安装配置:两条命令跑通第一个计算

FAIRChem V2 已剥离 torch-geometric 等历史依赖,pip 一条即可装好;多 GPU 并行才需要额外的 Ray 包:

pip install fairchem-core # 多GPU并行可选:pip install fairchem-core[ray] huggingface-cli login # UMA 是 gated 模型,先在 Hugging Face 申请 UMA 仓库访问

若需二次开发,可克隆本仓库后以编辑模式安装:git clone https://gitcode.com/GitHub_Trending/oc/ocp,然后pip install -e packages/fairchem-core[dev]。注意 V2 与 V1 预训练权重不兼容,老代码请固定fairchem-core==1.10。

所有计算遵循同一个四步模式:构造 ASEAtoms→ 加载 UMA → 选 task 挂上FAIRChemCalculator→ 取能量/力或交给 ASE 优化器。加载模型只需两行:

from fairchem.core import pretrained_mlip, FAIRChemCalculator predictor = pretrained_mlip.get_predict_unit("uma-s-1p2p1", device="cuda")

uma-s-1p2p1是当前默认推荐档(6.6M 激活 / 290M 总参数);追求精度可换uma-m-1p1(50M / 1.4B),代价是更慢、更吃显存。

场景走查:CO 在 Cu(100) 表面的吸附弛豫

这是催化初筛里最典型的单元任务:给定表面 + 吸附质,求弛豫后的几何与吸附能。输入是一个 3×3×3 的 Cu(100) slab(真空层 8 Å),CO 放在 2.0 Å 高的 bridge 位;关键参数是task_name="oc20"(对应 RPBE/VASP 层级)和 LBFGS 的收敛阈值fmax=0.05 eV/Å、上限 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") # 2.0 Å 高度, bridge 位 slab.calc = calc LBFGS(slab).run(fmax=0.05, steps=100) print(slab.get_potential_energy()) # 弛豫后总能量 (eV)

怎么读结果:slab.get_potential_energy()给出的是 slab+CO 总能量;吸附能需要三个量相减 E_ads = E(slab+CO) − E(slab) − E(CO 孤立),用同一 predictor 分别对三个碎片调用即可。弛豫收敛时slab.get_forces()的模长应接近 0.05 eV/Å 量级;若远大于该值,说明步数不够或初构型离目标位点太远。

更多任务族:MD、自旋、MOF 与实验数据闭环

• 分子动力学:omoltask 挂 ASELangevin(0.1 fs 步长、400 K)即可跑轨迹;MD 场景建议加inference_settings="turbo",并设随机种子区分不同轨迹。
• 自旋分裂能:对单重/三重态 CH₂ 分别设atoms.info["spin"],能量差即 spin gap(教程见 docs/core/quickstart.md)。
• MOF 吸附:odactask 面向 CO₂/H₂O 在金属有机框架中的吸附。ODAC 数据侧的力场–DFT 对照显示,大部分吸附构型的相互作用能误差集中在 ±0.25 eV 内:

图2:ODAC 数据集中 CO₂(红)与 H₂O(蓝)吸附相互作用能的 FF−DFT 误差分布,峰值位于 ±0.25 eV 内

• 催化剂发现闭环:仓库内置 OCx24 应用(src/fairchem/applications/ocx/),把 685M 条吸附表面构型、692,764 个候选表面与实验库对接,用计算特征预测实验活性;CatTSunami 的位点枚举逻辑(5 个位点 × 5 = 25 种产物组合)也在这里:

图3:OCx24 数据集:左侧为计算数据(685M 吸附构型)与实验流程,右侧为"预测模型 + 推断驱动发现"的使用方式

图4:CatTSunami 的三步位点筛选:每步选 5 个低能量位点,共枚举 25 种产物组合

底层设计:5.2 亿 DFT 样本训练的 MoLE 等变图网络

UMA 是一套等变图神经网络,结构上的关键选择是MoLE(Mixture of Linear Experts):单个输出头加动态路由,让总参数量很大、每个 token 只激活一小部分权重,从而在"容量"和"推理速度"之间取得平衡。官方口径:约5.2 亿个 DFT 计算(UMA 1.2,早期文档写作 500M)覆盖 7 个数据集;模型为每个 task 学一个嵌入,推理时按 task + 电荷 + 自旋 + 元素组成路由到对应参数。

UMA 1.2 相比 1.1:速度约快 50%,Open Molecules 测试集上精度约高 40%,并新增 OC22/OC25 两个催化档位。精度与速度同时摆在一张图上(来源:benchmarks/omol/,2,720,091 个 ωB97M-V 验证结构、H200 单卡):UMA-S 1.2.1 力误差约 7.6 meV/Å、1296 原子 NVE 单步约 50 ms;UMA-M 1.1 约 5.8 meV/Å 但单步约 400 ms。

图5:OMol25 验证集上力 MAE(meV/Å)对 1296 原子 ASE NVE 单步耗时(ms)的对比;UMA-S 1.2.1(蓝方块)处于误差低、速度快的区域

性能与扩展:workers=8 把 8000 原子推上 ns/day

多 GPU 并行不需要手动配 MPI/Kokkos:一个workers参数,底层用 Ray 拉起每张 GPU 一个 Actor,并把图切分到各 rank。官方验证过数百张 GPU、最大 1M 原子体系,官方示例是 8000 原子 fcc 晶体上的 Langevin MD(8 卡下约 10 次查询/秒量级):

predictor = pretrained_mlip.get_predict_unit( "uma-s-1p2p1", inference_settings="turbo", device="cuda", workers=8 ) calc = FAIRChemCalculator(predictor, task_name="omat") atoms = get_fcc_crystal_by_num_atoms(8000) # 来自 fairchem.core.datasets

单卡侧的小旋钮也值得知道(详见 docs/core/common_tasks/ase_calculator.md):execution_mode="umas-fast-gpu"对 uma-s 系列约 30–40% 加速;activation_checkpointing在 >1000 原子时建议保持开启;小系统批量弛豫/MD 在单张 H100 上最高可达 4 倍于串行的批处理加速(UMA 1.2 changelog)。

常见问题 FAQ

• 📌模型下载失败?UMA 在 Hugging Face 上是 gated 仓库:先申请 UMA 访问权限,再用huggingface-cli login写入 token;另外 V2 代码无法加载 UMA 1.0/1.1 老 checkpoint(会抛RuntimeError),需要老权重请降级fairchem-core<=2.21.0。
• 推理慢,怎么调?按顺序试:inference_settings="turbo"(MD 场景官方推荐)→execution_mode="umas-fast-gpu"→workers=N多卡。若系统在几百原子内,单卡 turbo 通常已够。
• 结果不收敛或偏差大?先查 task 是否选对:oc20不含色散,大分子在表面上的吸附能会偏;omol必须给 charge/spin;跨 task 的能量差(如用omat能量减oc20能量)没有物理意义。
• 显存不足?确认activation_checkpointing开着(>1000 原子默认建议开启),缩小体系尺寸,或从uma-m-1p1退回uma-s-1p2p1。
• 和 Materials Project 数据混用?OMat24/UMA 的 DFT 设置与 MP 不兼容(赝势、磁态、参考能不同),README 明确要求不要用 MP2020 校正或 MP 参考化合物。

资源索引与下一步

• 入门教程:docs/core/quickstart.md(从 ASE 结构到 UMA 预测)
• 任务与限制详解:docs/core/uma.md、发布记录 docs/core/uma_changelog.md
• 催化完整工作流(表面能、Wulff 形貌、NEB):docs/uma_tutorials/uma_catalysis_tutorial.md
• 核心源码:src/fairchem/core/,UMA 模型定义在 src/fairchem/core/models/uma/
• 基准复现脚本:benchmarks/omol/(speed.py/force_mae.py)
• 训练与评估配置:configs/(Hydra 风格,如configs/uma/benchmark/)

三条进阶建议:你可以先用uma-s-1p2p1把你熟悉的 2–3 个体系跑一遍,与已有 DFT 参考值核对误差量级,再决定是否放大筛选规模;随后把 MD 场景切到 turbo 模式;体系超过千原子时,再引入workers多卡与 LAMMPS 接口(src/fairchem/lammps/)。

写在最后

FAIRChem 把"用 DFT 训练、用 ML 推理"这条链路做成了标准库:一个模型覆盖 7 个化学领域,加速数字(2200 倍/70% 成功率、约 1000 倍单流程提速)都有对应的验证集与复现脚本可对账。你可以从上面 CO/Cu 的 10 行示例开始,先在自己熟悉的体系上建立对误差的信任,再逐步把 DFT 预算集中到真正需要高精度的候选上。

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