GPU Memory usage in inference

[dingye18]

I’m currently profiling the simulation performance of OpenMM+DP and comparing it with LAMMPS+DP. Interestingly, I found that the simulation speed of OpenMM is similar to LAMMPS. However, there is a significant difference in GPU memory usage between OpenMM+DP and LAMMPS+DP. While LAMMPS+DP fetches the neighbor list information from LAMMPS, OpenMM+DP relies on the neighbor list construction algorithm in DeePMD-kit.
To illustrate this point, I’ve created a table below showing the GPU memory usage for simulations with 1024 water molecules:

Hardware (Single GPU+ 4 CPU cores)	1060Ti	A40
OpenMM+DP	2.6G	2.98G
LAMMPS+DP	1.5G	1.68G

Is it reasonable to expect these results for simulations with DP models? If so, are there any plans to reduce the GPU memory usage during inference in the future? Currently, the high memory usage of DP limits the maximum number of atoms that can be simulated, particularly when compared with the neighbor list construction implementation in OpenMM which can typically support simulations with millions of atoms.

Are there any ongoing efforts to address this issue and improve the scalability of DP simulations in terms of the number of atoms that can be simulated?

[njzjz]

The DP native neighbor list OP calculates distances in $O(N^2)$ ($N$ is the number of atoms), which should be a bottleneck in performance and memory when $N$ is quite large.

deepmd-kit/source/lib/src/cuda/neighbor_list.cu

Lines 67 to 94 in 92ca097

	template <typename FPTYPE>
	__global__ void build_nlist(int *ilist,
	int *temp_nlist,
	const FPTYPE *c_cpy,
	const FPTYPE rcut2,
	const int nloc,
	const int nall,
	const int mem_size) {
	const unsigned int atom_idx = blockIdx.x;
	const unsigned int neighbor_idx = blockIdx.y * blockDim.y + threadIdx.y;
	if (neighbor_idx < nall) {
	int *neighbor_row = temp_nlist + atom_idx * mem_size;
	if (neighbor_idx == atom_idx) {
	ilist[atom_idx] = atom_idx;
	} else {
	const FPTYPE *ccoord = c_cpy + atom_idx * 3;
	const FPTYPE *ncoord = c_cpy + neighbor_idx * 3;
	FPTYPE diff[3];
	for (int kk = 0; kk < 3; kk++) {
	diff[kk] = ccoord[kk] - ncoord[kk];
	}
	FPTYPE r2 = dev_dot(diff, diff);
	if (r2 < rcut2) {
	neighbor_row[neighbor_idx] = neighbor_idx;
	}
	}
	}
	}

So I suggest passing the neighbor list from OpenMM to DP, or do you want to implement a more effect way to calculate distances when $N$ is large?

P.S. I am surprised that you said the speed is similar.

[dingye18]

The above profiling results were running on A40 + 4CPU cores. Time step is 0.2 fs.

It’s worth noting that the neighbor list implementation in OpenMM’s CUDA platform is different from that in LAMMPS, and currently there is no available approach to access the neighbor list variable in CUDA context.
As a result, I plan to call the CPU version of the neighbor list in OpenMM, which utilizes an O(N) algorithm.

BTW, I believe that DeePMD-kit utilizes an O(N) version of the neighbor list algorithm when running on GPU devices, as the functions are parallelized across multiple GPU cores.

[njzjz]

as the functions are parallelized across multiple GPU cores.

This only happens when $N$ is smaller than the number of GPU cores.