GPU Parallelism issue when using Pytorch as backend in DeePMD-kit 3.0.0b4

[chenggoj]

Hi,

Recently, I compiled and installed the lammps=2024.08.29=cuda118_py310_h572165c_mpi_openmpi_0 and deepmd-kit=3.0.0b4=cuda118py310hd03092b_mpi_openmpi_0. I am testing the GPU Parallelism on one GPU node with 4 NVIDIA Tesla V100 cards and 40 AMD CPU cores.

mpirun -h

mpirun (Open MPI) 5.0.5

When I run DP-GEN using tensorflow as backend using "CUDA_VISIBLE_DEVICES=0,1,2,3 mpirun -np 4 dp --tf", it is working very well and I can make full use of all GPU resources.

However, when I run the dpa2 descriptor using Pytorch as backend using "CUDA_VISIBLE_DEVICES=0,1,2,3 mpirun -np 4 dp --pt" I found I can only use one GPU and the other three GPUs are idle. I

The results are below when I run nvidia-smi
+-----------------------------------------------------------------------------+
| NVIDIA-SMI 525.85.12 Driver Version: 525.85.12 CUDA Version: 12.0 |
|-------------------------------+----------------------+----------------------+
| GPU Name Persistence-M| Bus-Id Disp.A | Volatile Uncorr. ECC |
| Fan Temp Perf Pwr:Usage/Cap| Memory-Usage | GPU-Util Compute M. |
| | | MIG M. |
|===============================+======================+======================|
| 0 Tesla V100-SXM2... On | 00000000:18:00.0 Off | Off |
| N/A 47C P0 154W / 300W | 31983MiB / 32768MiB | 100% Default |
| | | N/A |
+-------------------------------+----------------------+----------------------+
| 1 Tesla V100-SXM2... On | 00000000:3B:00.0 Off | 0 |
| N/A 31C P0 42W / 300W | 3MiB / 32768MiB | 0% Default |
| | | N/A |
+-------------------------------+----------------------+----------------------+
| 2 Tesla V100-SXM2... On | 00000000:86:00.0 Off | Off |
| N/A 35C P0 41W / 300W | 3MiB / 32768MiB | 0% Default |
| | | N/A |
+-------------------------------+----------------------+----------------------+
| 3 Tesla V100-SXM2... On | 00000000:AF:00.0 Off | 0 |
| N/A 35C P0 41W / 300W | 3MiB / 32768MiB | 0% Default |
| | | N/A |
+-------------------------------+----------------------+----------------------+

+-----------------------------------------------------------------------------+
| Processes: |
| GPU GI CI PID Type Process name GPU Memory |
| ID ID Usage |
|=============================================================================|
| 0 N/A N/A 1932458 C ...0b4-GPU-py3.11/bin/python 6732MiB |
| 0 N/A N/A 1932459 C ...0b4-GPU-py3.11/bin/python 2530MiB |
| 0 N/A N/A 1932460 C ...0b4-GPU-py3.11/bin/python 3842MiB |
| 0 N/A N/A 1932461 C ...0b4-GPU-py3.11/bin/python 18874MiB |
+-----------------------------------------------------------------------------+
Here is some information in the train.log. All 4 GPU are detected and visible. But only one GPU is used. I have no idea why. Can you give me any guidance to make full use of all GPUs when using Pytorch as backend? Thanks a lot.

[2024-09-30 22:07:32,823] DEEPMD INFO CUDA_VISIBLE_DEVICES: 0,1,2,3
[2024-09-30 22:07:32,823] DEEPMD INFO computing device: cuda:0
[2024-09-30 22:07:32,823] DEEPMD INFO world size: 4
[2024-09-30 22:07:32,823] DEEPMD INFO Count of visible GPUs: 4
[2024-09-30 22:07:32,823] DEEPMD INFO CUDA_VISIBLE_DEVICES: 0,1,2,3
[2024-09-30 22:07:32,823] DEEPMD INFO node list: exp-4-58
[2024-09-30 22:07:32,823] DEEPMD INFO Count of visible GPUs: 4
[2024-09-30 22:07:32,823] DEEPMD INFO num_intra_threads: 4
[2024-09-30 22:07:32,823] DEEPMD INFO running on: exp-4-58
[2024-09-30 22:07:32,823] DEEPMD INFO num_inter_threads: 4
[2024-09-30 22:07:32,823] DEEPMD INFO computing device: cuda:0
[2024-09-30 22:07:32,823] DEEPMD INFO num_intra_threads: 4
[2024-09-30 22:07:32,823] DEEPMD INFO ------------------------------------------------------------------------------------------------------------------------------------------------
[2024-09-30 22:07:32,823] DEEPMD INFO CUDA_VISIBLE_DEVICES: 0,1,2,3
[2024-09-30 22:07:32,823] DEEPMD INFO num_inter_threads: 4
[2024-09-30 22:07:32,823] DEEPMD INFO Count of visible GPUs: 4
[2024-09-30 22:07:32,823] DEEPMD INFO ------------------------------------------------------------------------------------------------------------------------------------------------
[2024-09-30 22:07:32,823] DEEPMD INFO num_intra_threads: 4
[2024-09-30 22:07:32,823] DEEPMD INFO num_inter_threads: 4

[njzjz]

However, when I run the dpa2 descriptor using Pytorch as backend using "CUDA_VISIBLE_DEVICES=0,1,2,3 mpirun -np 4 dp --pt" I found I can only use one GPU and the other three GPUs are idle. I

MPI training is not supported in the PyTorch backend. xref: #3951 (comment)

[chenggoj]

Thanks a lot. I found this command torchrun --nproc_per_node=4 --no-python dp --pt train input.json should work well on one single GPU node with 4 GPU cards. However, when I tried to use 2 GPU nodes (with 8 GPU cards), it did not work again. So could you please give me any guidance how to Parallel Pytorch training tasks on 8 GPUs across 2 different GPU nodes, if possible?

[njzjz]

Did you set --nnodes=2?

[chenggoj]

Yes, I added. But I deployed it in DP-GEN. It would not work at all when I tried to run the DP-GEN across two GPU nodes.