isayevlab/aimnet2-2025：Ascend NPU 推理复现

本仓库提供 isayevlab/aimnet2-2025 的单一 Ascend NPU 推理入口。复现不依赖迁移时的容器目录：代码、模型和模型层 Python 包分别放在同级的 repo/、model/、pydeps/，所有命令均使用相对路径。

固定模型快照

• 模型源：https://gitcode.com/hf_mirrors/isayevlab/aimnet2-2025.git
• 固定 revision：a2728c57db16064effa03e27006953013974db4f
• 许可：MIT
• 运行成员：ensemble_0.safetensors

快照校验值仅用于下载后离线核对，不在推理入口中绑定运行时身份：

git clone https://gitcode.com/2301_78565107/aimnet2-2025.git repo
git clone https://gitcode.com/hf_mirrors/isayevlab/aimnet2-2025.git model
git -C model checkout a2728c57db16064effa03e27006953013974db4f
git -C model lfs pull

环境

已验证平台栈：Ascend 910B3、CANN 8.5.1、Python 3.11、PyTorch 2.9.0、torch_npu 2.9.0.post1。CANN、PyTorch 和 torch_npu 属于平台栈，不写入模型层 requirements，也不得被 pip 覆盖。

cd repo
python3 -m pip install --no-deps --target ../pydeps -r requirements.txt
source /usr/local/Ascend/cann-8.5.1/set_env.sh
export PYTHONNOUSERSITE=1
export PYTHONPATH=../pydeps:$PYTHONPATH

tad-dftd3/tad-mctc 是本模型 D3(BJ) 项的必要辅助包，已固定在唯一的 requirements.txt 中。模型没有额外辅助权重。

NPU 推理

cd repo
source /usr/local/Ascend/cann-8.5.1/set_env.sh
export PYTHONNOUSERSITE=1
export PYTHONPATH=../pydeps:$PYTHONPATH
EVIDENCE_HOLD_SECONDS=0 python3 inference.py \
  --model_path ../model \
  --device npu:0 \
  --seed 20260816 \
  --output_json result.json

截图取证时可将 EVIDENCE_HOLD_SECONDS 设为 1..600。同一进程会在模型、输入和输出仍驻留 NPU 时打印 EVIDENCE_HOLD_START，等待后打印 EVIDENCE_HOLD_END。

入口只接受 npu:N，缺少 torch_npu、NPU 不可用或设备放置不一致都会直接失败，不存在 CPU/CUDA fallback。Warp 没有 Ascend 后端，因此本仓使用纯 torch NPU 张量实现邻居表，并用 tad-dftd3 的纯 torch autograd 路径计算 D3(BJ)；两部分均留在目标 NPU，不经过 CPU 数值桥。

固定用例为 8 个 H/C/N/O 分子，输出每个分子的能量和原子力，共 110 个数值。程序在 NPU 张量上检查模型/输出设备、shape、样本数、总数值数和 finite，成功结果写出 exit_status=0。已验证结果为 8 samples、110 values，requested_device、actual_model_device、actual_tensor_device、actual_output_device 均为 npu:0。

证据文件

assets/ 中包含两套不同时间点的设备调用、结果和实施方案截图：

• agent_workflow.png、agent_workflow_2.png
• npu_device_call.png、npu_device_call_2.png
• model_result.png、model_result_2.png

两张 workflow 均为实施方案，但提示词、章节和叙述顺序不同。

本仓为 2301_76761127 队伍交付件（与备份账号合并提交）。