[BUG] VRAM is wasted when running Lammps with multiple GPUs

### Bug summary

I have been using DP for a long time, and in every version I have used, I have encountered this issue: when running a Lammps MD simulation using multiple GPUs via `mpirun`, each MPI Rank consumes VRAM on all GPUs, even though the computation of each MPI Rank is actually running on only one GPU.

For example, in the picture below, I requested 4 V100-SXM2-16GB GPUs for a single MD job and started 4 MPI Ranks. In reality, each GPU has (4-1)*0.3=0.9GiB of VRAM "wasted". For an 8-GPU job, this would "waste" (8-1)*0.3=2.1GiB of VRAM. If MPS is used, the "wasted" VRAM would be doubled.


On the surface, it seems that this issue arises because the TensorFlow `gpu_device` runtime executes a "create device" operation for each GPU in every MPI Rank (as can be seen in the logs), but I don't know how to avoid this problem. It is noteworthy that TensorFlow "can't see" the GPUs on different nodes, so when running Lammps MD across multiple nodes and each node uses only one GPU, there is no such issue.

### DeePMD-kit Version

3.0.0b4

### Backend and its version

TensorFlow v2.15.2, Lammps 29Aug2024

### How did you download the software?

Offline packages

### Input Files, Running Commands, Error Log, etc.

Running Commands:
`mpirun -np 4 lmp_mpi -in input.lammps`

Part of Log:

```
...
2024-10-01 03:13:12.619343: I tensorflow/core/common_runtime/gpu/gpu_device.cc:1886] Created device /job:localhost/replica:0/task:0/device:GPU:0 with 14529 MB memory: -> device: 0, name: Tesla V100-SXM2-16GB, pci bus id: 0000:84:00.0, compute capability: 7.0
2024-10-01 03:13:12.620016: I tensorflow/core/common_runtime/gpu/gpu_device.cc:1886] Created device /job:localhost/replica:0/task:0/device:GPU:1 with 14529 MB memory: -> device: 1, name: Tesla V100-SXM2-16GB, pci bus id: 0000:85:00.0, compute capability: 7.0
2024-10-01 03:13:12.620570: I tensorflow/core/common_runtime/gpu/gpu_device.cc:1886] Created device /job:localhost/replica:0/task:0/device:GPU:2 with 14529 MB memory: -> device: 2, name: Tesla V100-SXM2-16GB, pci bus id: 0000:c4:00.0, compute capability: 7.0
2024-10-01 03:13:12.621108: I tensorflow/core/common_runtime/gpu/gpu_device.cc:1886] Created device /job:localhost/replica:0/task:0/device:GPU:3 with 14529 MB memory: -> device: 3, name: Tesla V100-SXM2-16GB, pci bus id: 0000:c5:00.0, compute capability: 7.0
2024-10-01 03:13:12.640945: I tensorflow/core/common_runtime/gpu/gpu_device.cc:1886] Created device /job:localhost/replica:0/task:0/device:GPU:0 with 14529 MB memory: -> device: 0, name: Tesla V100-SXM2-16GB, pci bus id: 0000:84:00.0, compute capability: 7.0
2024-10-01 03:13:12.641605: I tensorflow/core/common_runtime/gpu/gpu_device.cc:1886] Created device /job:localhost/replica:0/task:0/device:GPU:1 with 14529 MB memory: -> device: 1, name: Tesla V100-SXM2-16GB, pci bus id: 0000:85:00.0, compute capability: 7.0
2024-10-01 03:13:12.642124: I tensorflow/core/common_runtime/gpu/gpu_device.cc:1886] Created device /job:localhost/replica:0/task:0/device:GPU:2 with 14529 MB memory: -> device: 2, name: Tesla V100-SXM2-16GB, pci bus id: 0000:c4:00.0, compute capability: 7.0
2024-10-01 03:13:12.642635: I tensorflow/core/common_runtime/gpu/gpu_device.cc:1886] Created device /job:localhost/replica:0/task:0/device:GPU:3 with 14529 MB memory: -> device: 3, name: Tesla V100-SXM2-16GB, pci bus id: 0000:c5:00.0, compute capability: 7.0
2024-10-01 03:13:12.659556: I tensorflow/core/common_runtime/gpu/gpu_device.cc:1886] Created device /job:localhost/replica:0/task:0/device:GPU:0 with 14529 MB memory: -> device: 0, name: Tesla V100-SXM2-16GB, pci bus id: 0000:84:00.0, compute capability: 7.0
2024-10-01 03:13:12.660457: I tensorflow/core/common_runtime/gpu/gpu_device.cc:1886] Created device /job:localhost/replica:0/task:0/device:GPU:1 with 14529 MB memory: -> device: 1, name: Tesla V100-SXM2-16GB, pci bus id: 0000:85:00.0, compute capability: 7.0
2024-10-01 03:13:12.661253: I tensorflow/core/common_runtime/gpu/gpu_device.cc:1886] Created device /job:localhost/replica:0/task:0/device:GPU:0 with 14529 MB memory: -> device: 0, name: Tesla V100-SXM2-16GB, pci bus id: 0000:84:00.0, compute capability: 7.0
2024-10-01 03:13:12.661270: I tensorflow/core/common_runtime/gpu/gpu_device.cc:1886] Created device /job:localhost/replica:0/task:0/device:GPU:2 with 14529 MB memory: -> device: 2, name: Tesla V100-SXM2-16GB, pci bus id: 0000:c4:00.0, compute capability: 7.0
2024-10-01 03:13:12.662060: I tensorflow/core/common_runtime/gpu/gpu_device.cc:1886] Created device /job:localhost/replica:0/task:0/device:GPU:1 with 14529 MB memory: -> device: 1, name: Tesla V100-SXM2-16GB, pci bus id: 0000:85:00.0, compute capability: 7.0
2024-10-01 03:13:12.662095: I tensorflow/core/common_runtime/gpu/gpu_device.cc:1886] Created device /job:localhost/replica:0/task:0/device:GPU:3 with 14529 MB memory: -> device: 3, name: Tesla V100-SXM2-16GB, pci bus id: 0000:c5:00.0, compute capability: 7.0
2024-10-01 03:13:12.662639: I tensorflow/core/common_runtime/gpu/gpu_device.cc:1886] Created device /job:localhost/replica:0/task:0/device:GPU:2 with 14529 MB memory: -> device: 2, name: Tesla V100-SXM2-16GB, pci bus id: 0000:c4:00.0, compute capability: 7.0
2024-10-01 03:13:12.663289: I tensorflow/core/common_runtime/gpu/gpu_device.cc:1886] Created device /job:localhost/replica:0/task:0/device:GPU:3 with 14529 MB memory: -> device: 3, name: Tesla V100-SXM2-16GB, pci bus id: 0000:c5:00.0, compute capability: 7.0
...
```

### Steps to Reproduce

N/A

### Further Information, Files, and Links

_No response_