DeepMD-kit Environment (RTX 4090 / CUDA 12.8)

1. 简介 (Introduction)

DeepMD-kit 是一个基于深度学习的分子动力学模拟软件包，旨在生成具有量子力学精度的势能面（PES）。 本项目环境配置使用 PyTorch 作为计算后端，针对 NVIDIA RTX 4090 (Ada Lovelace 架构) 和 CUDA 12.8 进行了适配。

2. 环境配置 (Configuration)

当前服务器环境详细参数如下：

组件 (Component)	版本/型号 (Details)	备注
OS	Ubuntu Linux
GPU	NVIDIA GeForce RTX 4090
CUDA Driver	12.8	支持最新 Blackwell/Ada 架构
Python	3.12	Conda 环境
DeepMD-kit	v3.1.2 (Source Build)
Backend	PyTorch (cu128)	关键配置：不使用 TensorFlow

关键环境变量

为了避免 DeepMD-kit 错误地调用 TensorFlow，必须设置以下环境变量：

export DP_BACKEND=pytorch

建议将上述命令加入 ~/.bashrc 或 Conda 的激活脚本中。

3. 快速运行 GPU 测试 (Running GPU Test)

使用官方提供的水分子（Water）案例来验证 GPU 是否正常介入训练。

步骤 1：进入测试目录假设您处于 DeepMD-kit 源码根目录：

cd examples/water/se_e2_a

步骤 2：运行训练执行以下命令开始训练（显式指定 PyTorch 后端以防万一）：

dp --backend pytorch train input.json

或者，如果您已经设置了环境变量 DP_BACKEND=pytorch：

dp train input.json

步骤 3：验证运行状态

1. 检查日志输出 终端应显示类似以下的训练进度，包含 batch、l2_err 等信息：

DEEPMD INFO    start training...
DEEPMD INFO    batch    100   life  ...
DEEPMD INFO    batch    200   life  ...

2. 检查 GPU 占用 在训练过程中，打开一个新的终端窗口，输入：

nvidia-smi

请观察 python 进程是否出现在列表底部，且 GPU-Util（利用率）有明显浮动（非 0%），这证明 RTX 4090 正在工作。

---

4. 常见问题 (Troubleshooting)

Q: 出现 ImportError: cannot import name 'main' from 'deepmd.tf...' 错误？

• 原因：DeepMD-kit 默认尝试加载 TensorFlow，但当前环境未安装 TF。
• 解决：请务必加上 --backend pytorch 参数，或检查 export DP_BACKEND=pytorch 是否已执行。

Q: 出现 torch not compiled with CUDA enabled？

• 原因：安装了 CPU 版的 PyTorch。
• 解决：重新安装支持 CUDA 的 PyTorch 版本（推荐 index-url 指向 cu124 或 cu126）。