彻底解决 GROMACS 模拟中的 CUDA Out of Memory：从域分解与显存分配机制谈起

在进行大体系分子动力学（MD）模拟或使用多卡/多路 CPU 强卡并行的生产环境中，GROMACS 报错 "Out of memory" 导致 CUDA 驱动崩溃是一个非常经典且让人头疼的问题。

这类显存溢出（OOM）通常不是因为 GPU 的物理显存真的不够装下体系本身，而是由于域分解（Domain Decomposition, 简称 DD）参数配置不合理，导致局部子域（Domain）内的粒子数突增、邻居表（Neighbor List）构建过大，或者多 MPI 进程在单张显卡上创建了过多的 CUDA Context。

本文将从域分解的底层机制出发，详细拆解如何通过调整 GROMACS 的 DD 参数及相关运行参数，彻底解决 CUDA OOM 报错。

一、 为什么域分解（DD）会导致 CUDA 显存溢出？

GROMACS 在并行计算时，会将三维模拟盒子切分成多个子域（Cells），每个子域由一个 MPI 进程（Rank）负责。

当引入 GPU 加速时（如 -nb gpu -pme gpu），每个负责特定子域的 MPI 进程都需要在 GPU 上申请显存来存储：

1. 子域内部的粒子坐标与速度。
2. 通信缓冲区（Halo/Ghost Regions）：用于存放相邻子域传过来的粒子数据。
3. 邻居列表（Neighbor List）：这是显存消耗的“大户”。如果某个子域划分得过大，或者粒子密度在局部发生剧烈波动，该子域对应的 GPU 显存占用就会瞬间暴涨，触发 CUDA OOM。

此外，如果你的并行程数（-ntmpi）设置得过高，多个 MPI 进程会同时挤在同一张 GPU 上。每个 MPI 进程初始化 CUDA 时，都会在 GPU 上创建独立的 CUDA Context（通常每个 Context 默认静态占用 300MB~800MB 显存）。这就导致还没开始计算，显存就已经被 Context 挤爆了。

二、 调整域分解参数的黄金法则

解决此问题的核心思想是：优化子域的几何划分，控制单个 GPU 承载的计算边界，并减少不必要的显存开销。

1. 手动指定域分解网格：-dd

默认情况下，GROMACS 会根据总 MPI 进程数自动寻找最优的 3D 划分（例如将 8 个进程划分为 $2 \times 2 \times 2$）。但在某些拉伸体系（如膜蛋白、碳纳米管、一维流体）中，自动划分可能极其不合理，导致某些方向的子域尺寸过小或过大。

• 优化策略：使用 -dd x y z 显式指定三个方向的切分份数。
• 规则：
  • 切分份数的乘积必须等于你的总 MPI 任务数（即 -ntmpi 的值，若有独立 PME 进程则为 总MPI - PME进程）。
  • 避免在有大范围非平衡拉伸（AAMD）的方向切分过细。
  • 原则：尽量让切分后的子域呈正立方体，避免出现扁平或细长状的子域。 细长的子域会引入极大的 Ghost Region，导致显存中缓存的相邻粒子数量暴增。

# 假设有 8 个 PP 进程，体系在 Z 轴方向很长，避免在 Z 轴切分过细
gmx mdrun -deffnm topol -ntmpi 8 -dd 4 2 1

2. 限制最小子域尺寸：-dds 与 -rdd

GROMACS 运行过程中会动态调整子域边界（Dynamic Load Balancing, DLB）。如果体系局部密度波动剧烈（如气液界面、高分子坍缩），DLB 可能会把某个子域压缩得非常小。

当子域尺寸小于非键截断半径（Cut-off）时，GROMACS 为了保证计算正确，会强制扩大通信缓存，这在 GPU 上会导致显存分配异常。

• -dds (Domain Decomposition Screw/Constraint)：默认值通常是 0.8。它表示子域的最小尺寸不能小于 rcoulomb / dds 或 rvdw / dds。
  • 如果频繁因为“cell size too small”或者 OOM 报错，可以尝试将其调低（例如 -dds 0.7 或 -dds 0.6），这会强制约束子域不能收缩得太厉害，平摊 GPU 显存压力。
• -rdd (Real-space limit for DD)：显式设置域分解的几何截断值。
  • 手动将其设置为略大于你的 rvdw 和 rcoulomb 的值（例如，若 Cut-off 为 1.2 nm，可设置 -rdd 1.3），防止 GROMACS 动态估算出一个过大的通信距离导致显存溢出。

gmx mdrun -deffnm topol -dds 0.75 -rdd 1.3

3. 分离 PP（粒子-粒子）与 PME 进程：-npme

如果你开启了 GPU 上的 PME 加速（-pme gpu），默认情况下，所有的 GPU 既要算 PP，又要算 PME。PP 需要大量的邻居表显存，PME 需要大量的 FFT 网格显存，两者叠加极易 OOM。

• 优化策略：通过 -npme 参数将专门负责 PME 计算的进程剥离出来。
  • 比如你用 2 张 GPU（MPI 进程数设为 4）。你可以指定 3 个进程做 PP，1 个进程专门做 PME。
  • 更彻底的方法：将 PME 扔回 CPU 计算，而 GPU 只负责 PP。 这样可以释放 GPU 极大的显存空间。

# 将 PME 限制在 CPU 上运行，GPU 仅负责非键相互作用（PP），极大缓解显存压力
gmx mdrun -deffnm topol -nb gpu -pme cpu

三、 配合 DD 调整的其他关键“续命”参数

仅仅调整 -dd 有时还不够，必须配合以下进程与显存管理参数：

1. 减少 MPI 进程数，增加 OpenMP 线程数（最有效的降显存手段）

如前文所述，MPI 进程越多，单卡上的 CUDA Context 越多，显存浪费越严重。

• 错误示范（在一张卡上跑 8 个 MPI 进程）：
  -ntmpi 8 -ntomp 1 -> 产生 8 个 CUDA Context，显存直接白白浪费 4GB+。
• 正确示范（在一张卡上跑 1 或 2 个 MPI 进程，用多线程填满 CPU）：
  -ntmpi 1 -ntomp 8 或 -ntmpi 2 -ntomp 4。

此时，域分解只会在 1 或 2 个进程间进行，子域数量减少，每个子域的边界通信极大简化，显存占用呈断崖式下跌。

# 推荐的高效低显存配置（单卡 8 核 CPU 示例）
gmx mdrun -deffnm topol -ntmpi 1 -ntomp 8 -nb gpu -pme gpu

2. 限制邻居列表的更新频率：-nstlist

GROMACS 的 GPU 算法会动态调节 -nstlist（邻居表更新步数）。如果体系压力或温度剧烈波动，-nstlist 自动调整得太小，会导致频繁重新分配显存。

• 可以在 mdp 文件中，显式锁死 nstlist = 20 或 40。
• 配合 mdrun -nstlist 40，防止因为动态搜索机制在运行初期分配出过大的 GPU 缓冲区。

四、 黄金排查与解决工作流

当你遇到 CUDA Out of memory 报错时，请按以下步骤依次调整并排查：

[ 发生 CUDA OOM 报错 ]
       │
       ▼
1. 检查物理显存与进程数 ───> 是否在一张卡上跑了太多的 MPI 进程 (-ntmpi)？
       │                    如果是：减少 -ntmpi，增加 -ntomp (推荐 1卡对应 1~2 个 MPI)
       ▼
2. 检查 PME 绑定 ────────> 是否开启了 -pme gpu？
       │                    如果是：尝试改用 -pme cpu，把显存留给 PP 的邻居表
       ▼
3. 调整域分解几何约束 ───> 体系是否在某一方向极长，或者局部密度不均匀？
       │                    如果是：
       │                    - 使用 -dd x y z 显式指定均衡的 3D 网格
       │                    - 添加 -dds 0.7 限制子域过度收缩
       │                    - 添加 -rdd [cut-off + 0.1] 锁定通信边界
       ▼
[ 成功稳定运行 ]

最佳实战命令模版

对于一个 8 核 CPU、1 张 GPU 的普通工作站，运行一个 20 万原子左右的蛋白质-水盒子体系，若遇到 OOM，最稳健的启动命令为：

gmx mdrun -v -deffnm prod \
          -ntmpi 1 \
          -ntomp 8 \
          -nb gpu \
          -pme cpu \
          -dds 0.75 \
          -dlb yes

注：这里将 PME 放回 CPU（-pme cpu），并将 -ntmpi 降为 1，彻底消除了多 Context 竞争和域分解带来的通信显存开销，通常可释放 60% 以上的 GPU 显存，是解决 OOM 的终极底牌。