DeePMD on AMD GPUs

[Dankomaister]

Hi,
At the beginning of next year our group will get access to a new supercomputer with plenty of AMD Instinct MI250X GPUs. These are monsters with 128GB of GPU memory and I wanted to ask what is needed in order to run DeePMD on AMD GPUs. As far as I understand TensorFlow already runs on the AMD Instinct GPUs, but that doesn’t necessarily mean that DeePMD will work out of the box right?

[wanghan-iapcm]

Yes, our offline package only support cuda. You may have to install from source. please refer to https://github.com/deepmodeling/deepmd-kit/blob/master/doc/install/install-from-source.md.

Basically one needs to set the environment variable DP_VARIANT to rocm to switch on the compiling of deepmd-kit operators with rocm supported. If not set, deepmd-kit operators will be compiled for cpu, but all tensorflow operators are still on GPU.

Please contact us if you have any questions regarding the installation of deepmd-kit for AMD GPU.

[sigbjobo]

Hi, I am attempting to install deepmd on the LUMI supercomputer, which has MI250X GPUs. I have gotten quite far with the python interface but failed at the finishing line when training.

# Make a virtual environment
python3 -m venv deepmd-env
module reset
module load LUMI/22.08 partition/G
module swap PrgEnv-cray PrgEnv-gnu
module load cray-python
module load rocm
module load rccl
source deepmd-env/bin/activate
LDFLAGS='-L/opt/rocm/lib -lamdhip64' pip install tensorflow-rocm


# Test that it works on a GPU node
salloc --nodes=1 --account=XXXX  --partition=dev-g  --time=00:30:00  --gpus-per-node=1
srun python -c "import tensorflow as tf;print(tf.reduce_sum(tf.random.normal([1000, 1000])))" 
# Displays, so we know that rocm is using gpu
# Created device /job:localhost/replica:0/task:0/device:GPU:0 with 63996 MB memory:  -> device: 0, name: , pci bus id: 0000:c1:00.0
# 2023-01-17 12:58:35.139741: I tensorflow/core/common_runtime/gpu_fusion_pass.cc:507] ROCm Fusion is enabled.
# 2023-01-17 12:58:35.820510: I tensorflow/core/common_runtime/gpu_fusion_pass.cc:507] ROCm Fusion is enabled.
# 2023-01-17 12:58:35.912639: I tensorflow/core/common_runtime/gpu_fusion_pass.cc:507] ROCm Fusion is enabled.
# 2023-01-17 12:58:36.001386: I tensorflow/core/common_runtime/gpu_fusion_pass.cc:507] ROCm Fusion is enabled.
# 2023-01-17 12:58:36.222856: I tensorflow/core/common_runtime/gpu_fusion_pass.cc:507] ROCm Fusion is enabled.
# 2023-01-17 12:58:36.280467: I tensorflow/core/common_runtime/gpu_fusion_pass.cc:507] ROCm Fusion is enabled.

#Moving onto the deepmd-kit
git clone --recursive https://github.com/deepmodeling/deepmd-kit.git deepmd-kit 
cd deepmd-kit
export DP_VARIANT=rocm
export HIP_HIPCC_FLAGS="--amdgpu-target=gfx90a" # This is something we have tried
pip install . 

# Test on a gpu node
salloc --nodes=1 --account=XXXX  --partition=dev-g  --time=00:30:00  --gpus-per-node=1
module reset
module load LUMI/22.08 partition/G
module swap PrgEnv-cray PrgEnv-gnu
module load cray-python
module load rocm
module load rccl
source ../deepmd-env/bin/activate
cd examples/water/se_e2_a
srun dp train input.json

# And we get this error:
"hipErrorNoBinaryForGpu: Unable to find code object for all current devices!"
srun: error: nid007540: task 0: Aborted (core dumped)
srun: launch/slurm: _step_signal: Terminating StepId=2583879.1

Seem like the deepmd community also had the same problem: ROCm/ROCm#1623 (comment)

Do you have any advice? Also, the guys at LUMI, said it would be great if you had a rocm docker container for their newer version of ROCM: rocm/5.0.2

The following modules are loaded:

[njzjz]

export HIP_HIPCC_FLAGS="--amdgpu-target=gfx90a" # This is something we have tried

HIP_HIPCC_FLAGS should be a CMake variable. If you want it to be an environmental variable, you can contribute to setup.py.

[sigbjobo]

Fantastic, added:

hipcc_flags=os.environ.get("HIP_HIPCC_FLAGS")
if hipcc_flags:
    cmake_args.append(f"-DHIP_HIPCC_FLAGS:STRING={hipcc_flags}")

And runs without error now.

[njzjz]

It's welcome to submit a pull request (PR) to the devel branch for this feature.