mpi4py和cupy的联合应用(anaconda环境):GPU-aware MPI + Python GPU arrays
Demo代码:
from mpi4py import MPI import cupy as cp comm = MPI.COMM_WORLD size = comm.Get_size() rank = comm.Get_rank() sendbuf = cp.arange(10, dtype='i') recvbuf = cp.empty_like(sendbuf) assert hasattr(sendbuf, '__cuda_array_interface__') assert hasattr(recvbuf, '__cuda_array_interface__') cp.cuda.get_current_stream().synchronize() comm.Allreduce(sendbuf, recvbuf) assert cp.allclose(recvbuf, sendbuf*size)
使用anaconda环境:
安装mpi4py:
conda install -c conda-forge mpi4py openmpi
安装后的提示:
For Linux 64, Open MPI is built with CUDA awareness but this support is disabled by default.
To enable it, please set the environmental variable OMPI_MCA_opal_cuda_support=true before
launching your MPI processes. Equivalently, you can set the MCA parameter in the command line:
mpiexec --mca opal_cuda_support 1 ...
这个提示是说虽然anaconda安装的mpi4py是支持cuda-aware的,但是默认并没有开启,所以在运行时要设置环境变量 OMPI_MCA_opal_cuda_support=true ,并且运行时使用参数 mpiexec --mca opal_cuda_support 1
安装cupy:
conda install -c conda-forge cupy cudnn cutensor nccl
运行demo代码:
mpiexec --mca opal_cuda_support 1 -N 4 python x.py
==================================================
posted on 2023-06-14 20:22 Angry_Panda 阅读(344) 评论(0) 编辑 收藏 举报