DeePMD-kit Performance Optimization#

Overview#

DeePMD-kit exposes three independent levels of parallelism, each controlled by separate environment variables. Correct tuning requires the product of parallel-process counts not to exceed available CPU cores 1. Performance work splits into two concerns: (1) CPU/GPU thread tuning for single-node inference and training, and (2) multi-GPU / multi-node distributed training.

Thread-Level Parallelism (CPU)#

The full reference for all runtime environment variables is at env.html 2.

1 2

OMP_NUM_THREADS is the broadest lever: it controls DeePMD-kit's custom CPU kernels, the OneDNN path in TensorFlow, and the OpenMP path in PyTorch. GPU users should still set it because some DeePMD-kit OPs lack GPU implementations 1.

DP_INTER_OP_PARALLELISM_THREADS (inter-op) manages multiple CPU streams running independent operators concurrently. For GPU devices, TensorFlow and PyTorch use only one compute stream, so inter-op threads matter mainly for CPU-only workloads 1.

Additional OpenMP knobs such as KMP_BLOCKTIME can help; see Intel TF performance guide (TF) or PyTorch tuning recipes (PyTorch) 1.

Empirical Starting Points#

For 3 cores across 2 CPUs on one node, set OMP_NUM_THREADS=3 and DP_INTRA_OP_PARALLELISM_THREADS=3 with DP_INTER_OP_PARALLELISM_THREADS=1 1.

For a 128-core node, a good starting point is:

OMP_NUM_THREADS=16
DP_INTRA_OP_PARALLELISM_THREADS=16
DP_INTER_OP_PARALLELISM_THREADS=8 # 16 × 8 = 128

1

No single configuration is universally optimal — empirical testing is expected 1.

Other Performance-Relevant Variables#

2

Multi-GPU / Multi-Node Parallel Training#

Parallel training uses data parallelism: each worker sees a different mini-batch; per-batch training time does not decrease, but throughput scales. See Parallel training docs 3.

TensorFlow — Horovod#

Install Horovod and MPI4py, then launch with horovodrun 3:

horovodrun -np 4 dp train input.json

Each process binds to one GPU via CUDA_VISIBLE_DEVICES. DeePMD-kit auto-detects the MPI context and switches to distributed mode (visible in the log as world size > 1, distributed) 3.

Learning rate scaling: With N workers the effective batch size is N × batch_size. The learning rate is automatically scaled by N (linear mode); alternatives sqrt and none are available via scale_by_worker 3.

Observed speedups on 8 GPUs (se_e2_a, water example): 6.72× throughput with 8 GPUs vs. single GPU 3.

PyTorch — torchrun / DDP#

Uses PyTorch DistributedDataParallel (DDP) 3:

# Single node, 4 GPUs
torchrun --nproc_per_node=4 dp train input.json

For multi-node, set --nnodes, --node_rank, and --rdzv_endpoint 3.

ZeRO memory optimization is supported in the PyTorch backend via the zero_stage input config key 3:

Stage 3 incurs 50% more communication (3Ψ vs. 2Ψ); use it only when lower stages still exhaust GPU memory. Stage 2 can slow down models with many small layers — mitigate with torch.compile 3.

NUM_WORKERS (default: min(8, ncpus)) controls the DataLoader subprocess count in the PyTorch backend 3.

Paddle — paddle.distributed.fleet#

Use CUDA_VISIBLE_DEVICES and paddle.distributed.launch or mpirun. Multi-node setup follows the same nnodes/node_rank pattern as torchrun 3.

Inference Parallelism#

MPI is not directly supported for inference in DeePMD-kit; instead, third-party tools like LAMMPS handle distributed MD via their own MPI layer 1. Set CUDA_VISIBLE_DEVICES and thread variables as described above for single-process GPU inference.

Key References#

• Runtime environment variables 2
• How to control the parallelism of a job 1
• Parallel training 3