GPTFF 通用 AI 力场在昇腾 NPU 上的部署与推理实践

qq_45368619

85人浏览 · 2026-08-24 17:16:23

qq_45368619 · 2026-08-24 17:16:23 发布

本文记录无机材料通用力场 GPTFF 在华为昇腾（Ascend）NPU 平台上的部署与推理验证过程，涵盖模型原理、整体架构、环境搭建、推理代码解析、运行结果与常见问题，力求给出一份可复现的实操指南。

一、GPTFF 介绍

GPTFF（Graph-based Pre-trained Transformer Force Field，基于图结构的预训练 Transformer 力场）是中国科学院物理研究所孟胜、刘淼团队联合松山湖材料实验室研发的新一代高精度无机材料通用力场，官方代码仓：https://github.com/atomly-materials-research-lab/GPTFF 。

在材料模拟领域，DFT 计算精度高但只能处理数十至数百原子体系；经典分子动力学快但力场精度低、适用体系少。GPTFF 作为原子尺度的通用力场大模型，在两者之间取得平衡：可适用于几乎任意无机化合物的近平衡态，支持大体系及复杂体系的分子动力学模拟，计算速度比 DFT 快 3~4 个数量级。

核心特点：

混合架构：将图神经网络（GNN）与 Transformer 注意力机制融合。GNN 部分自然满足晶体结构的旋转、平移和置换对称性，Transformer 部分增强对原子间长程相互作用的建模能力；
数据基础：基于自主研发的 Atomly 材料数据库（https://atomly.net/）训练，训练集包含约 35 万无机材料的 223 万条 DFT 优化轨迹、3780 万个单点能量、117 亿个原子力向量和 3.4 亿个应力值，数据规模约为 MPtrj 数据集的 27.6 倍；
多任务输出：一次前向同时预测体系总能量、原子力（F = -∂E/∂x）和应力张量（σ = (1/V)∂E/∂ε），力与应力通过对能量的自动微分得到，保证物理自洽；
精度水平：能量 MAE 32 meV/atom，力 MAE 71 meV/Å，应力 MAE 0.365 GPa，在多项指标上优于 M3GNet、CHGNet 等国际主流模型；在 16,653 个全新结构（不在训练集与 Materials Project 中）上力预测 MAE 仍保持 66 meV/Å，泛化能力突出；
ASE 生态接入：以 ASECalculator 形式封装，可直接驱动 ASE 的结构优化（BFGS/FIRE 等）、分子动力学、声子计算等上层工作流；
模型参数量约 50 万（gptff_v1.pth），单卡推理开销极小。

Atomly 数据库中的元素与训练数据分布

图 1：Atomly 数据库中的元素与训练数据分布。训练数据集包含 2,234,767 个化合物，在基于密度泛函理论（DFT）进行结构优化的过程中，共生成约 3780 万个单点计算样本。

典型应用场景：晶体结构弛豫、固态电解质离子输运（如 Li₃YCl₆ 电导率预测与实验值高度吻合）、金属应力相变（如 Ti 的 HCP→FCC 相变模拟）、新材料大规模筛选等。

二、整体架构

GPTFF 将元素、化学键、键角等信息在图块（graph block）中完成表征与更新，并可选地融入 Transformer 模块捕捉长程相互作用：
Atomly 数据库中的元素与训练数据分布
图 2：GPTFF 模型架构示意图。元素、化学键、键角以及 Transformer 模块等信息在图块（graph block）中完成表征与更新。

GPTFF 推理链路由四部分组成：图构建（Cython 加速的近邻搜索与三体对枚举）→ 数据整理（collate）→ 模型前向（GNN 主干）→ 能量/力/应力输出（自动微分）。整体数据流如下：

cif / ASE Atoms（晶体结构）
    │  custom_graph.transform()          ← Cython: find_neighbors + compute_tp_cc
    ▼
原子特征(核电荷数) + 键(r=5.0Å) + 三体角(a=3.5Å) 图数据
    │  collate_fn()                      ← 组装为 torch 张量
    ▼
┌──────────────────────────────────────────────┐
│  tModLodaer (GNN 主干, n_layers=3)           │
│   每层:                                      │
│   ├─ ThreeBody    三体交互(键角余弦嵌入)      │
│   ├─ EdgeUpdate   边特征更新(门控)           │
│   └─ ConvLayer    原子消息传递(门控+径向基)   │
│   (可选 TransformerBlock, 见 2.3)            │
└──────────────────────────────────────────────┘
    │  readout: index_add 汇聚 + MLP
    ▼
体系总能量 E (加原子参考能修正)
    │  torch.autograd.grad
    ├──► 原子力 F = -∂E/∂x        (eV/Å)
    └──► 应力 σ = (1/V)·∂E/∂ε     (×160.218 转为 GPa)

对应源码见 gptff/model/model.py 与 gptff/model/mpredict.py，下面逐层拆解。

2.1 图构建与径向基

custom_graph.transform()（mpredict.py）负责把晶体结构转成图：find_neighbors 以 r_cut=5.0 Å 搜索周期性近邻键，compute_tp_cc 以 a_cut=3.5 Å 枚举三体（键角）对，二者均为 Cython 实现（gptff/utils_/compute_tp.pyx、compute_nb.pyx），这也是部署时需要 python setup.py build_ext --inplace 编译的原因。

原子特征直接使用元素核电荷数，经 nn.Embedding(95, 64) 嵌入（覆盖 95 种元素）。键长与键角统一用 16 维正弦径向基编码（类似 Bessel 基）：

# model.py
def ebf(d_ij, rcut, device):
    filters = torch.arange(0, 16, 1).to(device)
    return torch.sqrt(torch.tensor(2.) / radius) * torch.sin(filters * torch.pi / radius * d_ij) / d_ij

键角信息作为三体相互作用的关键编码，把键角余弦值从一维映射到高维空间，与两侧原子嵌入、边特征拼接后学习三体关联。

2.2 GNN 主干（tModLodaer）

默认权重走纯 GNN 路径 tModLodaer，每层依次执行：

ThreeBody：拼接三元组 (i, j, k) 的原子嵌入与两条边的特征，经 W_fea 投影后与键角嵌入、两条键长基做逐元素门控融合，index_add 聚合回边特征；
EdgeUpdate：以两端原子特征与当前边特征更新边，门控形式 SiLU(W_1·x) * Sigmoid(W_2·x)，再与径向基门控相乘；
ConvLayer：消息传递核心，邻居信息 SiLU(W_1) * Sigmoid(W_2) 门控聚合到中心原子，W_r(径向基) 提供距离加权。

三层堆叠后，index_add 将原子特征按结构汇聚为晶体特征，经两层 SiLU MLP 输出标量能量（扣除线性原子参考能 atom_refs 修正）。

2.3 可选 Transformer 分支（tModLodaer_t）

config 中 transformer_activate=true 时启用 tModLodaer_t：在每层 GNN 卷积后插入 TransformerBlock（nn.TransformerEncoder，d_model=64、2 头、FFN 256），对按结构 padding 组批的原子序列做自注意力，再与 GNN 输出残差融合。该分支利用注意力捕捉长程相互作用，同时保持 GNN 的对称性优势——这也是 GPTFF 名称中 “Transformer” 的来源。随仓发布的 gptff_v1.pth / gptff_v2.pth 为纯 GNN 路径权重。

2.4 力与应力：自动微分

get_efs() 中坐标与应变被设为可导，前向得到能量后一次反向同时求出力与应力：

# mpredict.py
force_pred, stress_pred = torch.autograd.grad(
    ener_pred, [coords, strain], torch.ones_like(ener_pred))
force_pred = -1.0 * force_pred
stress_pred = 1. / volumes[:, None, None] * stress_pred * 160.21766208  # eV/Å³ → GPa

应变通过 coords·(I+ε) 与 lattice·(I+ε) 隐式注入，这一设计同时覆盖了 torch.linalg.det（体积）等算子，是 NPU 上算子覆盖情况的关键路径（见第七节）。

三、实验环境

组件	版本
硬件	Ascend 910（64G HBM）x16，x86_64 Docker 容器
操作系统	openEuler 24.03 LTS-SP2
HDK	25.5.1
CANN	8.5.1
Python	3.11（conda）
torch / torch_npu	2.12.0 / 2.12.0
ase	3.21.1（必须，见 7.1）
numpy / scipy / pymatgen	2.4.6 / 1.17.1 / 2026.5.4
Cython	3.3.0

注：torch_npu 版本需与 torch 主版本严格一致，且与 CANN 版本配套；GPTFF 官方仓标注的验证组合为 CANN 8.2.RC1 + torch 2.6.0，本次实践在 CANN 8.5.1 + torch 2.12.0 上验证通过，说明其 GNN 算子面在昇腾后向兼容性良好。

四、环境搭建

4.1 拉取模型代码

git clone https://atomgit.com/AI4Science/gptff.git
cd gptff

预训练权重 pretrained/gptff_v1.pth、gptff_v2.pth 随仓携带，无需另行下载。

4.2 创建 conda 环境

conda create -n gptff python=3.11 -y
conda activate gptff

注：若后续 import torch_npu 报 libstdc++.so.6: version 'CXXABI_1.3.15' not found，是 conda 自带 libstdc++ 与系统版本冲突，需将 conda lib 目录前插：

export LD_LIBRARY_PATH=$CONDA_PREFIX/lib:$LD_LIBRARY_PATH

4.3 安装依赖

pip install -e .
pip install torch==2.12.0 torch-npu==2.12.0
pip install pyyaml decorator attrs psutil absl-py cloudpickle ml-dtypes scipy tornado ase==3.21.1
pip install cython
python setup.py build_ext --inplace

要点：

python setup.py build_ext --inplace 编译 Cython 近邻搜索/三体对模块，必须在本机执行（仓内自带的 .so 为他人环境产物，x86_64 与 aarch64 不通用）；
ase 必须为 3.21.1：pip install -e . 会按 pyproject 声明拉取新版 ase（≥3.22.1），装完后务必再执行一次 pip install ase==3.21.1 覆盖（见 7.1）。

4.4 验证 PyTorch 与 torch_npu

source /usr/local/Ascend/ascend-toolkit/set_env.sh
python3 -c "import torch;import torch_npu; a = torch.randn(3, 4).npu(); print(a + a);"

输出 device='npu:0' 张量即成功。报错时排查顺序：set_env.sh 是否已 source → decorator/pyyaml 等运行时依赖是否安装 → CANN 与 torch_npu 版本是否匹配。

4.5 流水优化（可选）

export CPU_AFFINITY_CONF=1
export TASK_QUEUE_ENABLE=2

五、推理代码解析（model_inference.py）

入口脚本 model_inference.py 演示"加载力场 → 读结构 → 算能量/力/应力"的完整流程：

from gptff.model.mpredict import ASECalculator
from pymatgen.core import Structure
from pymatgen.io.ase import AseAtomsAdaptor
import torch_npu
from torch_npu.contrib import transfer_to_npu

model_weight = "pretrained/gptff_v1.pth"
device = 'npu' # or cpu
p = ASECalculator(model_weight, device)   # 初始化模型并加载权重

adp = AseAtomsAdaptor()
struc = Structure.from_file('notebooks/data/NaCl.cif')  # 读入 NaCl 晶胞
atoms = adp.get_atoms(struc)
atoms.set_calculator(p)                   # 挂载力场，_atoms 被 GPTFF 接管

energy = atoms.get_potential_energy()     # 势能 (eV)
forces = atoms.get_forces()               # 原子受力 (eV/Å)
stress = atoms.get_stress()               # 应力张量 (GPa)

关键点：

import torch_npu + transfer_to_npu：注册 NPU 后端并将 torch.cuda.* 调用自动重映射为 torch.npu.*，脚本内无需任何 device 相关改动，device='npu' 直接可用；
ASECalculator 封装：GPTFF 实现了 ASE 的 Calculator 接口（implemented_properties: energy / free_energy / forces / stress），挂载后即可用 ASE 全套工作流（优化器、MD、Phonons 等在 mpredict.py 中均已 import，可直接使用）；
单结构推理：collate_fn 按单条结构组装批数据（代码预留了批处理接口），每次 calculate() 触发一次建图 + 前向 + 反向。

六、运行推理与结果展示

6.1 执行脚本

source /usr/local/Ascend/ascend-toolkit/set_env.sh
export ASCEND_RT_VISIBLE_DEVICES=0
export CPU_AFFINITY_CONF=1
export TASK_QUEUE_ENABLE=2
python model_inference.py

6.2 运行结果

==================================================
运行设备：npu
NaCl 势能：-27.4335 eV
原子受力（前2行）：
[[-3.5710400e-08  1.6210834e-08  1.0128133e-08]
 [-1.6705599e-08  1.3578756e-08 -4.7501089e-08]]
应力张量：
[-1.2516550e+00 -1.2516550e+00 -1.2516547e+00  2.3527781e-08
  2.8091482e-09  8.1194633e-08]
==================================================

结果物理合理性自检：NaCl 为立方晶体且输入为平衡晶格，原子受力应在 0 附近（实测 ~1e-8 eV/Å，数值噪声量级，符合预期）；应力三个正分量相等、剪切分量近 0（实测对角 -1.2517 GPa，偏差来自力场对平衡体积的预测与 DFT 优化体积的微小差异，属正常）。

6.3 性能数据（单卡 910，NaCl 8 原子单点）

阶段	耗时
import torch + torch_npu + transfer_to_npu	2.63 s
模型加载（含权重加载与 .to(npu)）	0.93 s
单点推理（建图 + 前向 + 反向）	0.31 s

单点推理中建图（Cython 近邻搜索）与模型前向均在亚秒级完成；小体系下 NPU 与 CPU 差距不大，NPU 优势需在批量结构或 MD 长轨迹场景下体现。

七、常见问题与告警说明

7.1 ImportError: cannot import name ‘ExpCellFilter’ from ‘ase.constraints’

根因：ASE 3.23 起 ExpCellFilter / StrainFilter 从 ase.constraints 移至 ase.filters，而 gptff/model/mpredict.py:19 仍从旧路径导入。pip install -e . 会按 pyproject 声明安装 ase≥3.22.1 的新版本，触发此错。

解决：安装依赖后强制覆盖为旧版：

pip install ase==3.21.1

7.2 libstdc++ CXXABI 版本报错

ImportError: libstdc++.so.6: version 'CXXABI_1.3.15' not found：conda 环境的 libicu 需要新版 libstdc++，而运行时加载到了系统的旧版。解法见 4.2（LD_LIBRARY_PATH 前插 $CONDA_PREFIX/lib）。

7.3 运行期告警（非报错，不影响结果）

aten::_linalg_det.result 回落 CPU：get_efs() 中 torch.linalg.det（晶胞体积）暂不受 NPU 后端支持，自动 host 回落执行，精度无影响；单点推理调用次数少可忽略，但 MD / 结构优化等每步都要算 stress 的长轨迹场景会引入 NPU↔CPU 数据搬运开销，需关注耗时占比；
fp64 → fp32 替换：transfer_to_npu 将 torch.cuda.DoubleTensor 替换为 FloatTensor（NPU 当前不支持 double），全部张量按 float32 计算。GPTFF 侧 collate_fn 本就构造 float32 张量，影响可控；但与 GPU/原始 fp64 实现做逐位精度对比时需注意该差异；
jit script 禁用：transfer_to_npu 不支持 torch.jit.script，GPTFF 未使用 jit，无影响；
ase / pymatgen 在 numpy 2.x 下的 __array__ copy=False DeprecationWarning：旧版库走 numpy 兼容路径，结果正确。