HPC 集群上 OpenMM GPU 多版本安装实战指南

HPC 集群上 OpenMM GPU 多版本安装实战指南

本文记录了在 HPC 集群共享存储（NFS）上，为不同 GPU 架构（RTX 4090 / RTX 5090）安装 OpenMM GPU 版本的完整流程，涵盖踩坑经验和修复方案。

---

目录

• [1. 环境背景](#1. 环境背景)
• [2. 最终成果概览](#2. 最终成果概览)
• [3. 安装前的准备工作](#3. 安装前的准备工作)
• [4. 正确的安装流程（完整代码）](#4. 正确的安装流程（完整代码）)
  • [4.1 CUDA 12.1 环境（RTX 4090）](#4.1 CUDA 12.1 环境（RTX 4090）)
  • [4.2 CUDA 12.8 环境（RTX 5090）](#4.2 CUDA 12.8 环境（RTX 5090）)
• [5. Module 文件配置](#5. Module 文件配置)
• [6. 验证测试](#6. 验证测试)
• [7. 踩坑记录与经验总结](#7. 踩坑记录与经验总结)
  • [坑1：conda 安装 CUDA toolkit 在 NFS 上卡死](#坑1：conda 安装 CUDA toolkit 在 NFS 上卡死)
  • [坑2：ambertools 拉高 cuda-version 导致 PTX 版本不兼容](#坑2：ambertools 拉高 cuda-version 导致 PTX 版本不兼容)
  • [坑3：pip PyTorch 与 conda OpenMM 的 C++ ABI 不兼容](#坑3：pip PyTorch 与 conda OpenMM 的 C++ ABI 不兼容)
  • [坑4：openmmforcefields 缺少 lxml 依赖](#坑4：openmmforcefields 缺少 lxml 依赖)
  • [坑5：conda 与 pip 混装时 conda 强制降级 PyTorch](#坑5：conda 与 pip 混装时 conda 强制降级 PyTorch)
  • [坑6：openmm-ml 安装超时或失败的多种解决方案](#坑6：openmm-ml 安装超时或失败的多种解决方案)
• [8. 使用方法](#8. 使用方法)

---

1. 环境背景

项目	说明
集群类型	HPC 共享存储（NFS），登录节点无 GPU
GPU 节点	cn01-cn09: RTX 4090（驱动 555.42.02, CUDA ≤12.5）；cn10-cn19: RTX 5090（驱动 590.48.01, CUDA 13.1）
系统 CUDA	/public/software/cuda-12.1 和 /public/software/cuda-12.8
包管理器	mamba（位于 /public/software/mamba）
Module 系统	Environment Modules（TCL），modulefiles 在 /public/software/modules/modulefiles/
安装目标	/public/software/openmm/（NFS 共享，所有节点可见）

约束条件：

• 不允许 apt-get install
• 不允许修改已有软件
• 不允许重复安装系统已有的 CUDA toolkit

---

2. 最终成果概览

环境	CUDA 12.1（RTX 4090）	CUDA 12.8（RTX 5090）
路径	/public/software/openmm/openmm-cu121	/public/software/openmm/openmm-cu128
OpenMM	8.2.0 (conda)	8.4.0 (conda)
PyTorch	2.5.1+cu121 (pip)	2.5.1+cu124 (pip)
AmberTools	24.8 (conda)	24.8 (conda)
OpenFF-toolkit	0.18.0 (conda)	0.18.0 (conda)
OpenMM-PLUMED	2.1 (conda)	2.1 (conda)
OpenMM-ML	1.5 (conda --no-deps)	1.5 (conda --no-deps)
MACE-torch	0.3.16 (pip)	0.3.16 (pip)
openmmforcefields	0.15.1 (pip)	0.15.1 (pip)
mdtraj	1.11.1 (conda/pip)	1.11.1 (conda/pip)
pdbfixer	1.12.0 (pip)	1.12.0 (pip)
lxml	6.1.1 (pip)	6.1.1 (pip)
CUDA Platform	✅ 已验证	✅ 已验证

---

---

3. 安装前的准备工作

3.1 确认 GPU 驱动版本

在 GPU 计算节点上执行：

nvidia-smi | head -4

输出示例（4090 节点）：

| NVIDIA-SMI 555.42.02    Driver Version: 555.42.02    CUDA Version: 12.5 |

关键规则 ：nvidia-smi 显示的 CUDA Version 是该驱动支持的最高 CUDA 版本。PyTorch 编译时使用的 CUDA 版本必须 ≤ 这个值。例如 4090 驱动 555.42.02 支持 CUDA ≤12.5，所以不能用 CUDA 12.6 编译的 PyTorch。

3.2 确认系统 CUDA 路径

ls /public/software/cuda-12.1/bin/nvcc   # 存在
ls /public/software/cuda-12.8/bin/nvcc   # 存在

3.3 确认 mamba 可用

/public/software/mamba/bin/mamba --version

---

4. 正确的安装流程（完整代码）

核心思路

conda 装 CPU/科学计算包 → pip 装 CUDA 相关包 → 降级 cuda-version 修复 PTX 兼容性

为什么不让 conda 装 CUDA？ 因为 conda 安装 CUDA toolkit 需要解析数百个依赖，在 NFS 共享存储上极慢（可能卡死数小时）。而 pip 的 torch+cu121 wheel 自带 CUDA runtime，体积小、安装快。

4.1 CUDA 12.1 环境（RTX 4090）

Step 1：创建基础环境 + ambertools + openff-toolkit

# 先装约束最多的包（ambertools 和 openff-toolkit），让 mamba 自动锁定兼容版本
/public/software/mamba/bin/mamba create -y \
    -p /public/software/openmm/openmm-cu121 \
    -c conda-forge \
    python=3.11 ambertools openff-toolkit

为什么要先装 ambertools？ ambertools 对 numpy、libboost、icu 等有严格的版本上限要求。如果先装了 openmm（这些依赖版本较新），再装 ambertools 就会冲突。先装 ambertools 可以让它锁定兼容的底层版本。

Step 2：安装 OpenMM + OpenMM-PLUMED

source /public/software/mamba/bin/activate /public/software/openmm/openmm-cu121

/public/software/mamba/bin/mamba install -y \
    -c conda-forge \
    openmm openmm-plumed

Step 3：pip 安装 PyTorch CUDA 版

# pip 的 torch+cu121 自带 CUDA 12.1 runtime，不需要系统安装 CUDA toolkit
# pip 不检测 __cuda 虚拟包，所以在无 GPU 的登录节点也能安装
pip install torch==2.5.1 \
    --index-url https://download.pytorch.org/whl/cu121

为什么用 pip 而不是 conda？

1. conda 安装 pytorch=*cuda* 需要 GPU 节点上有 __cuda 虚拟包（登录节点没有）
2. conda 安装 CUDA 相关包在 NFS 上极慢，可能卡死
3. pip wheel 小巧快速，自带完整 CUDA runtime

Step 4：pip 安装其他科学计算包

pip install mace-torch openmmforcefields mdtraj pdbfixer lxml

lxml 是 openmmforcefields 的硬依赖 ，调用 GAFFTemplateGenerator 时必须用到，但 openmmforcefields 的 setup.py 没有声明这个依赖，需要手动安装。

Step 5：conda --no-deps 安装 openmm-ml

# openmm-ml 是 MACE 力场集成到 OpenMM 的桥梁（提供 MLPotential 接口）
# 必须用 --no-deps，否则 conda 会强制降级 pip 安装的 PyTorch
/public/software/mamba/bin/mamba install -y \
    -c conda-forge \
    --no-deps openmm-ml

为什么必须用 --no-deps？ openmm-ml 依赖 PyTorch，如果让 conda 自动解析依赖，它会把 pip 安装的 torch 2.5.1+cu121 替换为 conda 的 CPU 版本。

Step 6：降级 cuda-version（关键步骤！）

# ambertools 安装时会把 cuda-version 拉到 13.x
# 如果不降级，NVRTC 生成的 PTX 版本会超出 GPU 驱动支持范围
# 导致运行时报 CUDA_ERROR_UNSUPPORTED_PTX_VERSION
/public/software/mamba/bin/mamba install -y \
    -c conda-forge \
    cuda-version=12.1 cuda-nvrtc=12.1

这是最容易遗漏的步骤！ 详见坑2。

Step 7：验证

# 登录节点验证导入（不需要 GPU）
python -c "
import torch;              print('torch:', torch.__version__)
import openmm;              print('openmm:', openmm.__version__)
from openmmplumed import PlumedForce;       print('plumed: OK')
from openff.toolkit.topology import Molecule; print('openff: OK')
from openmmforcefields.generators import GAFFTemplateGenerator; print('GAFF: OK')
import mdtraj;              print('mdtraj:', mdtraj.__version__)
import pdbfixer;            print('pdbfixer: OK')
import mace;                print('mace: OK')
print('All imports OK')
"

---

4.2 CUDA 12.8 环境（RTX 5090）

流程完全相同，只需替换 CUDA 版本号和 torch 的 index：

# Step 1: 创建环境
/public/software/mamba/bin/mamba create -y \
    -p /public/software/openmm/openmm-cu128 \
    -c conda-forge \
    python=3.11 ambertools openff-toolkit

# Step 2: 安装 OpenMM
source /public/software/mamba/bin/activate /public/software/openmm/openmm-cu128
/public/software/mamba/bin/mamba install -y -c conda-forge openmm openmm-plumed

# Step 3: pip 安装 PyTorch CUDA 12.4 版本
# 注意：5090 驱动 590 支持 CUDA 13.1，所以 cu124 完全兼容
pip install torch==2.5.1 \
    --index-url https://download.pytorch.org/whl/cu124

# Step 4: 其他包
pip install mace-torch openmmforcefields mdtraj pdbfixer lxml

# Step 5: openmm-ml（必须 --no-deps）
$MAMBA install -y -c conda-forge --no-deps openmm-ml

# Step 6: 降级 cuda-version（关键！）
/public/software/mamba/bin/mamba install -y \
    -c conda-forge \
    cuda-version=12.8 cuda-nvrtc=12.8

# Step 7: 验证（同上）

---

5. Module 文件配置

cu121 的 module 文件

文件路径：/public/software/modules/modulefiles/openmm/8.4-cu121

#%Module1.0

proc ModulesHelp {} {
        global version modroot
        puts stdout "\t loads OpenMM with CUDA 12.1 support\n"
        puts stdout "\t Includes: OpenMM-PLUMED, MACE-torch, AmberTools, OpenFF-toolkit\n"
        puts stdout "\t Use: source /public/software/mamba/bin/activate \$OPENMM_DIR\n"
}

module-whatis   "OpenMM with CUDA 12.1 (PLUMED, MACE, AmberTools)"

conflict openmm

set VERSION 8.4-cu121
set OPENMM_DIR /public/software/openmm/openmm-cu121
set CUDA_DIR /public/software/cuda-12.1

prepend-path    PATH               ${OPENMM_DIR}/bin
prepend-path    LD_LIBRARY_PATH    ${OPENMM_DIR}/lib
prepend-path    LD_LIBRARY_PATH    ${CUDA_DIR}/lib64
prepend-path    PYTHONPATH         ${OPENMM_DIR}/lib/python3.11/site-packages
prepend-path    OPENMM_DEFAULT_PLUGIN_DIR ${OPENMM_DIR}/lib/plugins
setenv          OPENMM_DIR         ${OPENMM_DIR}

if [ module-info mode load ] {
    system echo "OpenMM CUDA 12.1 loaded"
    system echo "  Activate: source /public/software/mamba/bin/activate ${OPENMM_DIR}"
}
if [ module-info mode remove ] {
    system echo "OpenMM CUDA 12.1 unloaded"
}

cu128 的 module 文件

文件路径：/public/software/modules/modulefiles/openmm/8.4-cu128

#%Module1.0

proc ModulesHelp {} {
        global version modroot
        puts stdout "\t loads OpenMM with CUDA 12.8 support\n"
        puts stdout "\t Includes: OpenMM-PLUMED, MACE-torch, AmberTools, OpenFF-toolkit\n"
        puts stdout "\t Use: source /public/software/mamba/bin/activate \$OPENMM_DIR\n"
}

module-whatis   "OpenMM with CUDA 12.8 (PLUMED, MACE, AmberTools)"

conflict openmm

set VERSION 8.4-cu128
set OPENMM_DIR /public/software/openmm/openmm-cu128
set CUDA_DIR /public/software/cuda-12.8

prepend-path    PATH               ${OPENMM_DIR}/bin
prepend-path    PATH               ${CUDA_DIR}/bin
prepend-path    LD_LIBRARY_PATH    ${OPENMM_DIR}/lib
prepend-path    LD_LIBRARY_PATH    ${CUDA_DIR}/lib64
prepend-path    PYTHONPATH         ${OPENMM_DIR}/lib/python3.11/site-packages
prepend-path    OPENMM_DEFAULT_PLUGIN_DIR ${OPENMM_DIR}/lib/plugins
setenv          OPENMM_DIR         ${OPENMM_DIR}

if [ module-info mode load ] {
    system echo "OpenMM CUDA 12.8 loaded"
    system echo "  Activate: source /public/software/mamba/bin/activate ${OPENMM_DIR}"
}
if [ module-info mode remove ] {
    system echo "OpenMM CUDA 12.8 unloaded"
}

注意 ：cu128 的 module 文件中额外加了 PATH ${CUDA_DIR}/bin，因为 5090 节点可能需要 nvcc 做 runtime 编译。

---

6. 验证测试

登录节点（无 GPU）------ 测试包导入

source /public/software/mamba/bin/activate /public/software/openmm/openmm-cu121

python -c "
import torch; print('torch:', torch.__version__)
import openmm; print('openmm:', openmm.__version__)
from openmmplumed import PlumedForce; print('plumed: OK')
from openff.toolkit.topology import Molecule; print('openff: OK')
from openmmforcefields.generators import GAFFTemplateGenerator; print('GAFF: OK')
import mdtraj, pdbfixer, mace, lxml
print('All imports OK')
"

GPU 计算节点 ------ 测试 CUDA 功能

# SSH 到 GPU 节点（4090 示例）
ssh cn01

source /public/software/mamba/bin/activate /public/software/openmm/openmm-cu121

python -c "
from openmm import *
from openmm.unit import *
import torch

# 测试 PyTorch CUDA
print('torch:', torch.__version__, 'CUDA:', torch.cuda.is_available())
print('GPU:', torch.cuda.get_device_name(0))

# 测试 OpenMM CUDA Platform
system = System()
system.addParticle(1.0)
force = NonbondedForce()
force.addParticle(0.0, 1.0, 0.0)
system.addForce(force)
integrator = LangevinMiddleIntegrator(300*kelvin, 1/picosecond, 0.002*picoseconds)
platform = Platform.getPlatformByName('CUDA')
context = Context(system, integrator, platform)
context.setPositions([Vec3(0, 0, 0)])
print('OpenMM CUDA Platform: SUCCESS!')

# 查看所有可用平台
print('Platforms:', [Platform.getPlatform(i).getName() for i in range(Platform.getNumPlatforms())])
"

预期输出：

torch: 2.5.1+cu121 CUDA: True
GPU: NVIDIA GeForce RTX 4090
OpenMM CUDA Platform: SUCCESS!
Platforms: ['Reference', 'CPU', 'CUDA', 'OpenCL']

---

7. 踩坑记录与经验总结

坑1：conda 安装 CUDA toolkit 在 NFS 上卡死

现象 ：执行 mamba install pytorch=2.5.1=*cuda* 时，mamba 进程长时间处于 D 状态（disk I/O sleep），持续数小时无进展。

原因：conda 安装 CUDA 相关包（pytorch CUDA、cuda-toolkit、libcufft 等）时需要解压数百 MB 的大文件（MKL 119MB、libcufft 184MB 等），NFS 共享存储的随机 IO 性能极差。

解决方案：用 pip 安装 PyTorch CUDA 版。pip 的 wheel 是预编译好的单文件，解压极快。

# ❌ 错误做法：conda 装 CUDA 包
mamba install -c conda-forge pytorch=2.5.1=*cuda* cuda-version=12.6

# ✅ 正确做法：pip 装 torch+cuXX
pip install torch==2.5.1 --index-url https://download.pytorch.org/whl/cu121

---

坑2：ambertools 拉高 cuda-version 导致 PTX 版本不兼容

现象：在 GPU 节点上运行 OpenMM 时报错：

openmm.OpenMMException: Error loading CUDA module: CUDA_ERROR_UNSUPPORTED_PTX_VERSION (222)

原因 ：这是最隐蔽的一个坑。安装 ambertools 时，conda 会自动安装 cuda-version=13.2 和 cuda-nvrtc=13.2.78 作为依赖。虽然环境名叫 cu121，但实际的 NVRTC（CUDA runtime compiler）是 13.2 版本，它生成的 PTX 代码版本超出了 GPU 驱动支持范围。

验证方法：

conda list | grep cuda-version
# 输出: cuda-version  13.2  ← 这是问题所在！

conda list | grep cuda-nvrtc
# 输出: cuda-nvrtc  13.2.78  ← NVRTC 版本也是 13.2

解决方案：安装完所有包后，强制降级 cuda-version 和 cuda-nvrtc：

# cu121 环境
mamba install -y -c conda-forge cuda-version=12.1 cuda-nvrtc=12.1

# cu128 环境
mamba install -y -c conda-forge cuda-version=12.8 cuda-nvrtc=12.8

经验法则 ：只要用了 ambertools，就必须在最后一步检查并降级 cuda-version 和 cuda-nvrtc。

---

坑3：pip PyTorch 与 conda OpenMM 的 C++ ABI 不兼容

现象：编译或加载 openmm-torch 时报错：

undefined symbol: _ZN6OpenMM18SerializationProxyC2ERKNSt7__cxx1112basic_stringIcSt11char_traitsIcESaIcEEE

原因：

• conda-forge 编译的 OpenMM 使用 新 ABI （_GLIBCXX_USE_CXX11_ABI=1）
• pip 安装的 PyTorch 2.5.1+cu121 使用 旧 ABI （_GLIBCXX_USE_CXX11_ABI=0）
• openmm-torch 是 C++ 扩展，链接时要求两端 ABI 一致

影响范围：openmm-torch 无法使用。但 openmm-ml、MACE-torch、openmmforcefields 等纯 Python 包不受影响。

解决方案：不装 openmm-torch。如果确实需要 openmm-torch，则必须用 conda 安装 PyTorch CUDA 版（但会遇到坑1的 NFS 问题）。这是目前 conda + pip 混合安装方案的已知限制。

---

坑4：openmmforcefields 缺少 lxml 依赖

现象 ：调用 GAFFTemplateGenerator 时报错：

ModuleNotFoundError: No module named 'lxml'

原因 ：openmmforcefields 在 generate_residue_template 函数中 from lxml import etree，但 setup.py 没有声明 lxml 为依赖。

解决方案：

pip install lxml

---

坑5：conda 与 pip 混装时 conda 强制降级 PyTorch

现象 ：用 pip 装好 torch==2.5.1+cu121 后，再执行 mamba install xxx，PyTorch 被降级为 CPU 版本。

原因：conda 的依赖解析器不知道 pip 安装的包。当安装新包时，conda 会根据自己的规则重新解析环境，发现 pytorch 没有满足约束就装一个 CPU 版本。

解决方案：

# 方法1：所有 conda 操作在 pip install torch 之前完成（推荐）
# 先 conda 装完所有需要的包，最后再 pip install torch

# 方法2：使用 --no-deps 避免依赖解析
mamba install --no-deps -c conda-forge openmm-ml

# 方法3：用 conda-meta/pinned 锁定版本（不完全可靠）
echo "pytorch ==2.5.1" >> /path/to/env/conda-meta/pinned

---

坑6：openmm-ml 安装超时或失败的多种解决方案

现象 ：执行 mamba install openmm-ml 时超时，或安装后 conda 降级了 PyTorch。

原因：

1. 多个 mamba 进程同时运行竞争锁文件（NFS 上常见）
2. 不加 --no-deps 时 conda 会降级 pip 安装的 PyTorch

解决方案（按优先级排序）：

# 方案 A：conda --no-deps（首选，环境干净时几乎必成功）
mamba install -y -c conda-forge --no-deps openmm-ml

# 方案 B：pip 从 GitHub 安装（需网络）
pip install git+https://github.com/openmm/openmm-ml.git

# 方案 C：手动下载源码 + pip 本地安装
# 在有网络的机器上下载 https://github.com/openmm/openmm-ml/archive/refs/heads/main.zip
# 传到服务器后：
unzip openmm-ml-main.zip
cd openmm-ml-main
pip install .

# 方案 D：从 conda 缓存直接拷贝（openmm-ml 是纯 Python 包）
# 适用于之前下载过但安装失败的情况
cp -r /public/software/mamba/pkgs/openmm-ml-*/site-packages/openmmml \
    /public/software/openmm/openmm-cu121/lib/python3.11/site-packages/

经验 ：之前安装失败是因为环境状态混乱 + 多个 mamba 进程竞争锁。环境重建干净后，方案 A 秒成功。所以遇到超时时，先确认没有其他 mamba 进程在跑（ps aux | grep mamba）。

---

8. 使用方法

加载环境

# 在 4090 节点上
module load openmm/8.4-cu121
source /public/software/mamba/bin/activate /public/software/openmm/openmm-cu121

# 在 5090 节点上
module load openmm/8.4-cu128
source /public/software/mamba/bin/activate /public/software/openmm/openmm-cu128

MACE 力场模型路径

预训练模型位于共享目录：

/public/software/openmm/models/mace-off-main/
├── mace_off23/
│   ├── MACE-OFF23_small.model      (7.1MB)
│   ├── MACE-OFF23_medium.model     (18MB)
│   ├── MACE-OFF23_large.model      (53MB)
│   └── MACE-OFF23b_medium.model    (18MB)
└── mace_off24/
    └── MACE-OFF24_medium.model     (18MB)

常用 import 头文件

from openmm import *
from openmm.app import *
from openmm.unit import *
from openmmplumed import PlumedForce
from openmmforcefields.generators import GAFFTemplateGenerator
from openff.toolkit.topology import Molecule
import mdtraj as md
import pdbfixer
import numpy as np

---

附录：安装命令速查表

cu121（RTX 4090，CUDA ≤12.5）

MAMBA=/public/software/mamba/bin/mamba
ENV=/public/software/openmm/openmm-cu121

$MAMBA create -y -p $ENV -c conda-forge python=3.11 ambertools openff-toolkit
source $MAMBA activate $ENV
$MAMBA install -y -c conda-forge openmm openmm-plumed
pip install torch==2.5.1 --index-url https://download.pytorch.org/whl/cu121
pip install mace-torch openmmforcefields mdtraj pdbfixer lxml
$MAMBA install -y -c conda-forge --no-deps openmm-ml
$MAMBA install -y -c conda-forge cuda-version=12.1 cuda-nvrtc=12.1

cu128（RTX 5090，CUDA ≤13.1）

MAMBA=/public/software/mamba/bin/mamba
ENV=/public/software/openmm/openmm-cu128

$MAMBA create -y -p $ENV -c conda-forge python=3.11 ambertools openff-toolkit
source $MAMBA activate $ENV
$MAMBA install -y -c conda-forge openmm openmm-plumed
pip install torch==2.5.1 --index-url https://download.pytorch.org/whl/cu124
pip install mace-torch openmmforcefields mdtraj pdbfixer lxml
$MAMBA install -y -c conda-forge --no-deps openmm-ml
$MAMBA install -y -c conda-forge cuda-version=12.8 cuda-nvrtc=12.8