Repository: cupy/cupy

Issue:
CI: Need to install MPI for distributed tests



Discussion:

Resolving patch:
diff --git a/.pfnci/generate.py b/.pfnci/generate.py
--- a/.pfnci/generate.py
+++ b/.pfnci/generate.py
@@ -102,7 +102,9 @@ def generate_dockerfile(self) -> str:
                 '       libbz2-dev libreadline-dev libsqlite3-dev wget \\',
                 '       curl llvm libncursesw5-dev xz-utils tk-dev \\',
                 '       libxml2-dev libxmlsec1-dev libffi-dev \\',
-                '       liblzma-dev && \\',
+                '       liblzma-dev \\',
+                '       libopenmpi-dev \\' if matrix.mpi4py else '\\',
+                '       && \\',
                 '    apt-get -qqy install ccache git curl && \\',
                 '    apt-get -qqy --allow-change-held-packages \\',
                 '            --allow-downgrades install {}'.format(
@@ -138,6 +140,7 @@ def generate_dockerfile(self) -> str:
                 'ENV PATH "/usr/lib64/ccache:${PATH}"',
                 '',
             ]
+            assert matrix.mpi4py is None, 'mpi4py test unsupported on CentOS'
         else:
             raise AssertionError
 
@@ -169,7 +172,8 @@ def generate_dockerfile(self) -> str:
         # Setup Python libraries.
         pip_args = []
         pip_uninstall_args = []
-        for pylib in ('numpy', 'scipy', 'optuna', 'cython', 'cuda-python'):
+        for pylib in ('numpy', 'scipy', 'optuna', 'mpi4py',
+                      'cython', 'cuda-python'):
             pylib_ver = getattr(matrix, pylib)
             if pylib_ver is None:
                 pip_uninstall_args.append(pylib)
@@ -286,6 +290,10 @@ def generate_script(self) -> str:
             if matrix.test == 'unit':
                 spec = 'not slow and not multi_gpu'
             elif matrix.test == 'unit-multi':
+                lines += [
+                    'export OMPI_ALLOW_RUN_AS_ROOT=1',
+                    'export OMPI_ALLOW_RUN_AS_ROOT_CONFIRM=1'
+                ]
                 spec = 'not slow and multi_gpu'
             elif matrix.test == 'unit-slow':
                 # Slow tests may use multiple GPUs.
diff --git a/cupyx/distributed/_nccl_comm.py b/cupyx/distributed/_nccl_comm.py
--- a/cupyx/distributed/_nccl_comm.py
+++ b/cupyx/distributed/_nccl_comm.py
@@ -523,7 +523,10 @@ def _exchange_shape_and_sizes(
                 comm._mpi_comm.Recv(sizes_shape, source=peer, tag=1)
                 return sizes_shape
             elif method == 'bcast':
-                sizes_shape = numpy.array(sizes_shape, dtype='q')
+                if comm.rank == peer:
+                    sizes_shape = numpy.array(sizes_shape, dtype='q')
+                else:
+                    sizes_shape = numpy.empty(5, dtype='q')
                 comm._mpi_comm.Bcast(sizes_shape, root=peer)
                 return sizes_shape
             elif method == 'gather':
@@ -531,6 +534,11 @@ def _exchange_shape_and_sizes(
                 recv_buf = numpy.empty([comm._n_devices, 5], dtype='q')
                 comm._mpi_comm.Gather(sizes_shape, recv_buf, peer)
                 return recv_buf
+            elif method == 'alltoall':
+                sizes_shape = numpy.array(sizes_shape, dtype='q')
+                recv_buf = numpy.empty([comm._n_devices, 5], dtype='q')
+                comm._mpi_comm.Alltoall(sizes_shape, recv_buf)
+                return recv_buf
             else:
                 raise RuntimeError('Unsupported method')
         else:
@@ -566,6 +574,12 @@ def _exchange_shape_and_sizes(
                 _DenseNCCLCommunicator.gather(
                     comm, sizes_shape, recv_buf, root=peer, stream=stream)
                 return cupy.asnumpy(recv_buf)
+            elif method == 'alltoall':
+                sizes_shape = cupy.array(sizes_shape, dtype='q')
+                recv_buf = cupy.empty((comm._n_devices, 5), dtype='q')
+                _DenseNCCLCommunicator.all_to_all(
+                    comm, sizes_shape, recv_buf, stream=stream)
+                return cupy.asnumpy(recv_buf)
             else:
                 raise RuntimeError('Unsupported method')
 
@@ -798,18 +812,32 @@ def all_to_all(cls, comm, in_array, out_array, stream=None):
             raise RuntimeError(
                 f'all_to_all requires in_array to have {comm._n_devices}'
                 f'elements, found {len(in_array)}')
-        # TODO(ecastill) assert the types
-        for _ in range(comm._n_devices):
-            out_array.append(_make_sparse_empty(
-                in_array[comm.rank].dtype,
-                _get_sparse_type(in_array[comm.rank])))
-        # TODO check out dtypes are the same as in dtypes
+
+        # Exchange metadata
+        shape_and_sizes = []
+        recv_shape_and_sizes = []
+        for i, a in enumerate(in_array):
+            arrays = cls._get_internal_arrays(a)
+            shape_and_sizes.append(cls._get_shape_and_sizes(arrays, a.shape))
+
+        recv_shape_and_sizes = cls._exchange_shape_and_sizes(
+            comm, i, shape_and_sizes, 'alltoall', stream)
+
+        # prepare the arrays to recv the data
         for i in range(comm._n_devices):
-            if i != comm.rank:
-                cls.send(comm, in_array[i], i, stream)
-                cls.recv(comm, out_array[i], i, stream)
-            else:
-                cls._assign_arrays(
-                    out_array[i],
-                    cls._get_internal_arrays(in_array[i]),
-                    in_array[i].shape)
+            shape = tuple(recv_shape_and_sizes[i][0:2])
+            sizes = recv_shape_and_sizes[i][2:]
+            s_arrays = cls._get_internal_arrays(in_array[i])
+            # TODO(use the out_array datatypes)
+            r_arrays = [
+                cupy.empty(s, dtype=a.dtype) for s, a in zip(sizes, s_arrays)]
+            nccl.groupStart()
+            for a in s_arrays:
+                cls._send(comm, a, i, a.dtype, a.size, stream)
+            for a in r_arrays:
+                cls._recv(comm, a, i, a.dtype, a.size, stream)
+            nccl.groupEnd()
+            out_array.append(_make_sparse_empty(
+                in_array[i].dtype,
+                _get_sparse_type(in_array[i])))
+            cls._assign_arrays(out_array[i], r_arrays, shape)


Test patch:
diff --git a/.pfnci/linux/tests/benchmark.Dockerfile b/.pfnci/linux/tests/benchmark.Dockerfile
--- a/.pfnci/linux/tests/benchmark.Dockerfile
+++ b/.pfnci/linux/tests/benchmark.Dockerfile
@@ -9,7 +9,9 @@ RUN export DEBIAN_FRONTEND=noninteractive && \
        libbz2-dev libreadline-dev libsqlite3-dev wget \
        curl llvm libncursesw5-dev xz-utils tk-dev \
        libxml2-dev libxmlsec1-dev libffi-dev \
-       liblzma-dev && \
+       liblzma-dev \
+\
+       && \
     apt-get -qqy install ccache git curl && \
     apt-get -qqy --allow-change-held-packages \
             --allow-downgrades install 'libnccl2=2.10.*+cuda11.4' 'libnccl-dev=2.10.*+cuda11.4' 'libcutensor1=1.4.*' 'libcutensor-dev=1.4.*' 'libcusparselt0=0.2.0.*' 'libcusparselt-dev=0.2.0.*' 'libcudnn8=8.3.*+cuda11.5' 'libcudnn8-dev=8.3.*+cuda11.5'
@@ -24,5 +26,5 @@ RUN pyenv install 3.10.0 && \
     pip install -U setuptools pip
 
 RUN pip install -U 'numpy==1.22.*' 'scipy==1.7.*' 'optuna==2.*' 'cython==0.29.*'
-RUN pip uninstall -y cuda-python && \
+RUN pip uninstall -y mpi4py cuda-python && \
     pip check
diff --git a/.pfnci/linux/tests/benchmark.head.Dockerfile b/.pfnci/linux/tests/benchmark.head.Dockerfile
--- a/.pfnci/linux/tests/benchmark.head.Dockerfile
+++ b/.pfnci/linux/tests/benchmark.head.Dockerfile
@@ -9,7 +9,9 @@ RUN export DEBIAN_FRONTEND=noninteractive && \
        libbz2-dev libreadline-dev libsqlite3-dev wget \
        curl llvm libncursesw5-dev xz-utils tk-dev \
        libxml2-dev libxmlsec1-dev libffi-dev \
-       liblzma-dev && \
+       liblzma-dev \
+\
+       && \
     apt-get -qqy install ccache git curl && \
     apt-get -qqy --allow-change-held-packages \
             --allow-downgrades install 'libnccl2=2.10.*+cuda11.4' 'libnccl-dev=2.10.*+cuda11.4' 'libcutensor1=1.4.*' 'libcutensor-dev=1.4.*' 'libcusparselt0=0.2.0.*' 'libcusparselt-dev=0.2.0.*' 'libcudnn8=8.3.*+cuda11.5' 'libcudnn8-dev=8.3.*+cuda11.5'
@@ -24,5 +26,5 @@ RUN pyenv install 3.10.0 && \
     pip install -U setuptools pip
 
 RUN pip install -U 'numpy==1.22.*' 'scipy==1.7.*' 'optuna==2.*' 'cython==0.29.*'
-RUN pip uninstall -y cuda-python && \
+RUN pip uninstall -y mpi4py cuda-python && \
     pip check
diff --git a/.pfnci/linux/tests/cuda-example.Dockerfile b/.pfnci/linux/tests/cuda-example.Dockerfile
--- a/.pfnci/linux/tests/cuda-example.Dockerfile
+++ b/.pfnci/linux/tests/cuda-example.Dockerfile
@@ -9,7 +9,9 @@ RUN export DEBIAN_FRONTEND=noninteractive && \
        libbz2-dev libreadline-dev libsqlite3-dev wget \
        curl llvm libncursesw5-dev xz-utils tk-dev \
        libxml2-dev libxmlsec1-dev libffi-dev \
-       liblzma-dev && \
+       liblzma-dev \
+\
+       && \
     apt-get -qqy install ccache git curl && \
     apt-get -qqy --allow-change-held-packages \
             --allow-downgrades install 'libnccl2=2.10.*+cuda11.4' 'libnccl-dev=2.10.*+cuda11.4' 'libcutensor1=1.5.*' 'libcutensor-dev=1.5.*' 'libcusparselt0=0.2.0.*' 'libcusparselt-dev=0.2.0.*' 'libcudnn8=8.4.*+cuda11.6' 'libcudnn8-dev=8.4.*+cuda11.6'
@@ -24,5 +26,5 @@ RUN pyenv install 3.8.11 && \
     pip install -U setuptools pip
 
 RUN pip install -U 'numpy==1.20.*' 'scipy==1.7.*' 'optuna==2.*' 'cython==0.29.*'
-RUN pip uninstall -y cuda-python && \
+RUN pip uninstall -y mpi4py cuda-python && \
     pip check
diff --git a/.pfnci/linux/tests/cuda-head.Dockerfile b/.pfnci/linux/tests/cuda-head.Dockerfile
--- a/.pfnci/linux/tests/cuda-head.Dockerfile
+++ b/.pfnci/linux/tests/cuda-head.Dockerfile
@@ -9,7 +9,9 @@ RUN export DEBIAN_FRONTEND=noninteractive && \
        libbz2-dev libreadline-dev libsqlite3-dev wget \
        curl llvm libncursesw5-dev xz-utils tk-dev \
        libxml2-dev libxmlsec1-dev libffi-dev \
-       liblzma-dev && \
+       liblzma-dev \
+\
+       && \
     apt-get -qqy install ccache git curl && \
     apt-get -qqy --allow-change-held-packages \
             --allow-downgrades install 'libcutensor1=1.5.*' 'libcutensor-dev=1.5.*' 'libcusparselt0=0.2.0.*' 'libcusparselt-dev=0.2.0.*' 'libcudnn8=8.4.*+cuda11.6' 'libcudnn8-dev=8.4.*+cuda11.6'
@@ -24,5 +26,5 @@ RUN pyenv install 3.10.0 && \
     pip install -U setuptools pip
 
 RUN pip install -U 'numpy>=0a0' 'scipy>=0a0' 'optuna>=0a0' 'cython==0.29.*'
-RUN pip uninstall -y cuda-python && \
+RUN pip uninstall -y mpi4py cuda-python && \
     pip check
diff --git a/.pfnci/linux/tests/cuda-slow.Dockerfile b/.pfnci/linux/tests/cuda-slow.Dockerfile
--- a/.pfnci/linux/tests/cuda-slow.Dockerfile
+++ b/.pfnci/linux/tests/cuda-slow.Dockerfile
@@ -9,7 +9,9 @@ RUN export DEBIAN_FRONTEND=noninteractive && \
        libbz2-dev libreadline-dev libsqlite3-dev wget \
        curl llvm libncursesw5-dev xz-utils tk-dev \
        libxml2-dev libxmlsec1-dev libffi-dev \
-       liblzma-dev && \
+       liblzma-dev \
+\
+       && \
     apt-get -qqy install ccache git curl && \
     apt-get -qqy --allow-change-held-packages \
             --allow-downgrades install 'libnccl2=2.10.*+cuda11.4' 'libnccl-dev=2.10.*+cuda11.4' 'libcutensor1=1.5.*' 'libcutensor-dev=1.5.*' 'libcusparselt0=0.2.0.*' 'libcusparselt-dev=0.2.0.*' 'libcudnn8=8.4.*+cuda11.6' 'libcudnn8-dev=8.4.*+cuda11.6'
@@ -24,5 +26,5 @@ RUN pyenv install 3.9.6 && \
     pip install -U setuptools pip
 
 RUN pip install -U 'numpy==1.21.*' 'scipy==1.7.*' 'optuna==2.*' 'cython==0.29.*'
-RUN pip uninstall -y cuda-python && \
+RUN pip uninstall -y mpi4py cuda-python && \
     pip check
diff --git a/.pfnci/linux/tests/cuda102.Dockerfile b/.pfnci/linux/tests/cuda102.Dockerfile
--- a/.pfnci/linux/tests/cuda102.Dockerfile
+++ b/.pfnci/linux/tests/cuda102.Dockerfile
@@ -15,7 +15,9 @@ RUN export DEBIAN_FRONTEND=noninteractive && \
        libbz2-dev libreadline-dev libsqlite3-dev wget \
        curl llvm libncursesw5-dev xz-utils tk-dev \
        libxml2-dev libxmlsec1-dev libffi-dev \
-       liblzma-dev && \
+       liblzma-dev \
+\
+       && \
     apt-get -qqy install ccache git curl && \
     apt-get -qqy --allow-change-held-packages \
             --allow-downgrades install 'libnccl2=2.8.*+cuda10.2' 'libnccl-dev=2.8.*+cuda10.2' 'libcudnn7=7.6.*+cuda10.2' 'libcudnn7-dev=7.6.*+cuda10.2'
@@ -30,5 +32,5 @@ RUN pyenv install 3.7.11 && \
     pip install -U setuptools pip
 
 RUN pip install -U 'numpy==1.20.*' 'scipy==1.6.*' 'optuna==2.*' 'cython==0.29.*'
-RUN pip uninstall -y cuda-python && \
+RUN pip uninstall -y mpi4py cuda-python && \
     pip check
diff --git a/.pfnci/linux/tests/cuda102.multi.Dockerfile b/.pfnci/linux/tests/cuda102.multi.Dockerfile
--- a/.pfnci/linux/tests/cuda102.multi.Dockerfile
+++ b/.pfnci/linux/tests/cuda102.multi.Dockerfile
@@ -15,7 +15,9 @@ RUN export DEBIAN_FRONTEND=noninteractive && \
        libbz2-dev libreadline-dev libsqlite3-dev wget \
        curl llvm libncursesw5-dev xz-utils tk-dev \
        libxml2-dev libxmlsec1-dev libffi-dev \
-       liblzma-dev && \
+       liblzma-dev \
+       libopenmpi-dev \
+       && \
     apt-get -qqy install ccache git curl && \
     apt-get -qqy --allow-change-held-packages \
             --allow-downgrades install 'libnccl2=2.8.*+cuda10.2' 'libnccl-dev=2.8.*+cuda10.2' 'libcudnn7=7.6.*+cuda10.2' 'libcudnn7-dev=7.6.*+cuda10.2'
@@ -29,6 +31,6 @@ RUN pyenv install 3.7.11 && \
     pyenv global 3.7.11 && \
     pip install -U setuptools pip
 
-RUN pip install -U 'numpy==1.20.*' 'scipy==1.6.*' 'optuna==2.*' 'cython==0.29.*'
+RUN pip install -U 'numpy==1.20.*' 'scipy==1.6.*' 'optuna==2.*' 'mpi4py==3.*' 'cython==0.29.*'
 RUN pip uninstall -y cuda-python && \
     pip check
diff --git a/.pfnci/linux/tests/cuda102.multi.sh b/.pfnci/linux/tests/cuda102.multi.sh
--- a/.pfnci/linux/tests/cuda102.multi.sh
+++ b/.pfnci/linux/tests/cuda102.multi.sh
@@ -10,5 +10,7 @@ ACTIONS="$(dirname $0)/actions"
 export NVCC="ccache nvcc"
 
 "$ACTIONS/build.sh"
+export OMPI_ALLOW_RUN_AS_ROOT=1
+export OMPI_ALLOW_RUN_AS_ROOT_CONFIRM=1
 "$ACTIONS/unittest.sh" "not slow and multi_gpu"
 "$ACTIONS/cleanup.sh"
diff --git a/.pfnci/linux/tests/cuda110.Dockerfile b/.pfnci/linux/tests/cuda110.Dockerfile
--- a/.pfnci/linux/tests/cuda110.Dockerfile
+++ b/.pfnci/linux/tests/cuda110.Dockerfile
@@ -9,7 +9,9 @@ RUN export DEBIAN_FRONTEND=noninteractive && \
        libbz2-dev libreadline-dev libsqlite3-dev wget \
        curl llvm libncursesw5-dev xz-utils tk-dev \
        libxml2-dev libxmlsec1-dev libffi-dev \
-       liblzma-dev && \
+       liblzma-dev \
+\
+       && \
     apt-get -qqy install ccache git curl && \
     apt-get -qqy --allow-change-held-packages \
             --allow-downgrades install 'libnccl2=2.9.*+cuda11.0' 'libnccl-dev=2.9.*+cuda11.0' 'libcutensor1=1.4.*' 'libcutensor-dev=1.4.*' 'libcudnn8=8.2.*+cuda11.4' 'libcudnn8-dev=8.2.*+cuda11.4'
@@ -24,5 +26,5 @@ RUN pyenv install 3.9.6 && \
     pip install -U setuptools pip
 
 RUN pip install -U 'numpy==1.20.*' 'scipy==1.7.*' 'optuna==2.*' 'cython==0.29.*'
-RUN pip uninstall -y cuda-python && \
+RUN pip uninstall -y mpi4py cuda-python && \
     pip check
diff --git a/.pfnci/linux/tests/cuda110.multi.Dockerfile b/.pfnci/linux/tests/cuda110.multi.Dockerfile
--- a/.pfnci/linux/tests/cuda110.multi.Dockerfile
+++ b/.pfnci/linux/tests/cuda110.multi.Dockerfile
@@ -9,7 +9,9 @@ RUN export DEBIAN_FRONTEND=noninteractive && \
        libbz2-dev libreadline-dev libsqlite3-dev wget \
        curl llvm libncursesw5-dev xz-utils tk-dev \
        libxml2-dev libxmlsec1-dev libffi-dev \
-       liblzma-dev && \
+       liblzma-dev \
+       libopenmpi-dev \
+       && \
     apt-get -qqy install ccache git curl && \
     apt-get -qqy --allow-change-held-packages \
             --allow-downgrades install 'libnccl2=2.9.*+cuda11.0' 'libnccl-dev=2.9.*+cuda11.0' 'libcutensor1=1.4.*' 'libcutensor-dev=1.4.*' 'libcudnn8=8.2.*+cuda11.4' 'libcudnn8-dev=8.2.*+cuda11.4'
@@ -23,6 +25,6 @@ RUN pyenv install 3.9.6 && \
     pyenv global 3.9.6 && \
     pip install -U setuptools pip
 
-RUN pip install -U 'numpy==1.20.*' 'scipy==1.7.*' 'optuna==2.*' 'cython==0.29.*'
+RUN pip install -U 'numpy==1.20.*' 'scipy==1.7.*' 'optuna==2.*' 'mpi4py==3.*' 'cython==0.29.*'
 RUN pip uninstall -y cuda-python && \
     pip check
diff --git a/.pfnci/linux/tests/cuda110.multi.sh b/.pfnci/linux/tests/cuda110.multi.sh
--- a/.pfnci/linux/tests/cuda110.multi.sh
+++ b/.pfnci/linux/tests/cuda110.multi.sh
@@ -12,5 +12,7 @@ export NVCC="ccache nvcc"
 export CUPY_ACCELERATORS="cub,cutensor"
 
 "$ACTIONS/build.sh"
+export OMPI_ALLOW_RUN_AS_ROOT=1
+export OMPI_ALLOW_RUN_AS_ROOT_CONFIRM=1
 "$ACTIONS/unittest.sh" "not slow and multi_gpu"
 "$ACTIONS/cleanup.sh"
diff --git a/.pfnci/linux/tests/cuda111.Dockerfile b/.pfnci/linux/tests/cuda111.Dockerfile
--- a/.pfnci/linux/tests/cuda111.Dockerfile
+++ b/.pfnci/linux/tests/cuda111.Dockerfile
@@ -25,5 +25,5 @@ RUN pyenv install 3.7.11 && \
     pip install -U setuptools pip
 
 RUN pip install -U 'numpy==1.20.*' 'cython==0.29.*'
-RUN pip uninstall -y scipy optuna cuda-python && \
+RUN pip uninstall -y scipy optuna mpi4py cuda-python && \
     pip check
diff --git a/.pfnci/linux/tests/cuda111.multi.Dockerfile b/.pfnci/linux/tests/cuda111.multi.Dockerfile
--- a/.pfnci/linux/tests/cuda111.multi.Dockerfile
+++ b/.pfnci/linux/tests/cuda111.multi.Dockerfile
@@ -25,5 +25,5 @@ RUN pyenv install 3.7.11 && \
     pip install -U setuptools pip
 
 RUN pip install -U 'numpy==1.20.*' 'cython==0.29.*'
-RUN pip uninstall -y scipy optuna cuda-python && \
+RUN pip uninstall -y scipy optuna mpi4py cuda-python && \
     pip check
diff --git a/.pfnci/linux/tests/cuda111.multi.sh b/.pfnci/linux/tests/cuda111.multi.sh
--- a/.pfnci/linux/tests/cuda111.multi.sh
+++ b/.pfnci/linux/tests/cuda111.multi.sh
@@ -12,5 +12,7 @@ export NVCC="ccache nvcc"
 export CUPY_ACCELERATORS="cutensor,cub"
 
 "$ACTIONS/build.sh"
+export OMPI_ALLOW_RUN_AS_ROOT=1
+export OMPI_ALLOW_RUN_AS_ROOT_CONFIRM=1
 "$ACTIONS/unittest.sh" "not slow and multi_gpu"
 "$ACTIONS/cleanup.sh"
diff --git a/.pfnci/linux/tests/cuda112.Dockerfile b/.pfnci/linux/tests/cuda112.Dockerfile
--- a/.pfnci/linux/tests/cuda112.Dockerfile
+++ b/.pfnci/linux/tests/cuda112.Dockerfile
@@ -25,5 +25,5 @@ RUN pyenv install 3.7.11 && \
     pip install -U setuptools pip
 
 RUN pip install -U 'numpy==1.21.*' 'scipy==1.7.*' 'optuna==2.*' 'cython==0.29.*'
-RUN pip uninstall -y cuda-python && \
+RUN pip uninstall -y mpi4py cuda-python && \
     pip check
diff --git a/.pfnci/linux/tests/cuda112.multi.Dockerfile b/.pfnci/linux/tests/cuda112.multi.Dockerfile
--- a/.pfnci/linux/tests/cuda112.multi.Dockerfile
+++ b/.pfnci/linux/tests/cuda112.multi.Dockerfile
@@ -25,5 +25,5 @@ RUN pyenv install 3.7.11 && \
     pip install -U setuptools pip
 
 RUN pip install -U 'numpy==1.21.*' 'scipy==1.7.*' 'optuna==2.*' 'cython==0.29.*'
-RUN pip uninstall -y cuda-python && \
+RUN pip uninstall -y mpi4py cuda-python && \
     pip check
diff --git a/.pfnci/linux/tests/cuda112.multi.sh b/.pfnci/linux/tests/cuda112.multi.sh
--- a/.pfnci/linux/tests/cuda112.multi.sh
+++ b/.pfnci/linux/tests/cuda112.multi.sh
@@ -12,5 +12,7 @@ export NVCC="ccache nvcc"
 export CUPY_ACCELERATORS=""
 
 "$ACTIONS/build.sh"
+export OMPI_ALLOW_RUN_AS_ROOT=1
+export OMPI_ALLOW_RUN_AS_ROOT_CONFIRM=1
 "$ACTIONS/unittest.sh" "not slow and multi_gpu"
 "$ACTIONS/cleanup.sh"
diff --git a/.pfnci/linux/tests/cuda113.Dockerfile b/.pfnci/linux/tests/cuda113.Dockerfile
--- a/.pfnci/linux/tests/cuda113.Dockerfile
+++ b/.pfnci/linux/tests/cuda113.Dockerfile
@@ -9,7 +9,9 @@ RUN export DEBIAN_FRONTEND=noninteractive && \
        libbz2-dev libreadline-dev libsqlite3-dev wget \
        curl llvm libncursesw5-dev xz-utils tk-dev \
        libxml2-dev libxmlsec1-dev libffi-dev \
-       liblzma-dev && \
+       liblzma-dev \
+\
+       && \
     apt-get -qqy install ccache git curl && \
     apt-get -qqy --allow-change-held-packages \
             --allow-downgrades install 'libnccl2=2.9.*+cuda11.3' 'libnccl-dev=2.9.*+cuda11.3' 'libcutensor1=1.5.*' 'libcutensor-dev=1.5.*' 'libcudnn8=8.2.*+cuda11.4' 'libcudnn8-dev=8.2.*+cuda11.4'
@@ -24,5 +26,5 @@ RUN pyenv install 3.8.11 && \
     pip install -U setuptools pip
 
 RUN pip install -U 'numpy==1.20.*' 'scipy==1.6.*' 'optuna==2.*' 'cython==0.29.*'
-RUN pip uninstall -y cuda-python && \
+RUN pip uninstall -y mpi4py cuda-python && \
     pip check
diff --git a/.pfnci/linux/tests/cuda113.multi.Dockerfile b/.pfnci/linux/tests/cuda113.multi.Dockerfile
--- a/.pfnci/linux/tests/cuda113.multi.Dockerfile
+++ b/.pfnci/linux/tests/cuda113.multi.Dockerfile
@@ -9,7 +9,9 @@ RUN export DEBIAN_FRONTEND=noninteractive && \
        libbz2-dev libreadline-dev libsqlite3-dev wget \
        curl llvm libncursesw5-dev xz-utils tk-dev \
        libxml2-dev libxmlsec1-dev libffi-dev \
-       liblzma-dev && \
+       liblzma-dev \
+       libopenmpi-dev \
+       && \
     apt-get -qqy install ccache git curl && \
     apt-get -qqy --allow-change-held-packages \
             --allow-downgrades install 'libnccl2=2.9.*+cuda11.3' 'libnccl-dev=2.9.*+cuda11.3' 'libcutensor1=1.5.*' 'libcutensor-dev=1.5.*' 'libcudnn8=8.2.*+cuda11.4' 'libcudnn8-dev=8.2.*+cuda11.4'
@@ -23,6 +25,6 @@ RUN pyenv install 3.8.11 && \
     pyenv global 3.8.11 && \
     pip install -U setuptools pip
 
-RUN pip install -U 'numpy==1.20.*' 'scipy==1.6.*' 'optuna==2.*' 'cython==0.29.*'
+RUN pip install -U 'numpy==1.20.*' 'scipy==1.6.*' 'optuna==2.*' 'mpi4py==3.*' 'cython==0.29.*'
 RUN pip uninstall -y cuda-python && \
     pip check
diff --git a/.pfnci/linux/tests/cuda113.multi.sh b/.pfnci/linux/tests/cuda113.multi.sh
--- a/.pfnci/linux/tests/cuda113.multi.sh
+++ b/.pfnci/linux/tests/cuda113.multi.sh
@@ -12,5 +12,7 @@ export NVCC="ccache nvcc"
 export CUPY_ACCELERATORS="cub,cutensor"
 
 "$ACTIONS/build.sh"
+export OMPI_ALLOW_RUN_AS_ROOT=1
+export OMPI_ALLOW_RUN_AS_ROOT_CONFIRM=1
 "$ACTIONS/unittest.sh" "not slow and multi_gpu"
 "$ACTIONS/cleanup.sh"
diff --git a/.pfnci/linux/tests/cuda114.Dockerfile b/.pfnci/linux/tests/cuda114.Dockerfile
--- a/.pfnci/linux/tests/cuda114.Dockerfile
+++ b/.pfnci/linux/tests/cuda114.Dockerfile
@@ -9,7 +9,9 @@ RUN export DEBIAN_FRONTEND=noninteractive && \
        libbz2-dev libreadline-dev libsqlite3-dev wget \
        curl llvm libncursesw5-dev xz-utils tk-dev \
        libxml2-dev libxmlsec1-dev libffi-dev \
-       liblzma-dev && \
+       liblzma-dev \
+\
+       && \
     apt-get -qqy install ccache git curl && \
     apt-get -qqy --allow-change-held-packages \
             --allow-downgrades install 'libnccl2=2.10.*+cuda11.4' 'libnccl-dev=2.10.*+cuda11.4' 'libcutensor1=1.4.*' 'libcutensor-dev=1.4.*' 'libcusparselt0=0.2.0.*' 'libcusparselt-dev=0.2.0.*' 'libcudnn8=8.3.*+cuda11.5' 'libcudnn8-dev=8.3.*+cuda11.5'
@@ -24,5 +26,5 @@ RUN pyenv install 3.10.0 && \
     pip install -U setuptools pip
 
 RUN pip install -U 'numpy==1.22.*' 'scipy==1.7.*' 'optuna==2.*' 'cython==0.29.*'
-RUN pip uninstall -y cuda-python && \
+RUN pip uninstall -y mpi4py cuda-python && \
     pip check
diff --git a/.pfnci/linux/tests/cuda114.multi.Dockerfile b/.pfnci/linux/tests/cuda114.multi.Dockerfile
--- a/.pfnci/linux/tests/cuda114.multi.Dockerfile
+++ b/.pfnci/linux/tests/cuda114.multi.Dockerfile
@@ -9,7 +9,9 @@ RUN export DEBIAN_FRONTEND=noninteractive && \
        libbz2-dev libreadline-dev libsqlite3-dev wget \
        curl llvm libncursesw5-dev xz-utils tk-dev \
        libxml2-dev libxmlsec1-dev libffi-dev \
-       liblzma-dev && \
+       liblzma-dev \
+       libopenmpi-dev \
+       && \
     apt-get -qqy install ccache git curl && \
     apt-get -qqy --allow-change-held-packages \
             --allow-downgrades install 'libnccl2=2.10.*+cuda11.4' 'libnccl-dev=2.10.*+cuda11.4' 'libcutensor1=1.4.*' 'libcutensor-dev=1.4.*' 'libcusparselt0=0.2.0.*' 'libcusparselt-dev=0.2.0.*' 'libcudnn8=8.3.*+cuda11.5' 'libcudnn8-dev=8.3.*+cuda11.5'
@@ -23,6 +25,6 @@ RUN pyenv install 3.10.0 && \
     pyenv global 3.10.0 && \
     pip install -U setuptools pip
 
-RUN pip install -U 'numpy==1.22.*' 'scipy==1.7.*' 'optuna==2.*' 'cython==0.29.*'
+RUN pip install -U 'numpy==1.22.*' 'scipy==1.7.*' 'optuna==2.*' 'mpi4py==3.*' 'cython==0.29.*'
 RUN pip uninstall -y cuda-python && \
     pip check
diff --git a/.pfnci/linux/tests/cuda114.multi.sh b/.pfnci/linux/tests/cuda114.multi.sh
--- a/.pfnci/linux/tests/cuda114.multi.sh
+++ b/.pfnci/linux/tests/cuda114.multi.sh
@@ -12,5 +12,7 @@ export NVCC="ccache nvcc"
 export CUPY_ACCELERATORS="cutensor,cub"
 
 "$ACTIONS/build.sh"
+export OMPI_ALLOW_RUN_AS_ROOT=1
+export OMPI_ALLOW_RUN_AS_ROOT_CONFIRM=1
 "$ACTIONS/unittest.sh" "not slow and multi_gpu"
 "$ACTIONS/cleanup.sh"
diff --git a/.pfnci/linux/tests/cuda115.Dockerfile b/.pfnci/linux/tests/cuda115.Dockerfile
--- a/.pfnci/linux/tests/cuda115.Dockerfile
+++ b/.pfnci/linux/tests/cuda115.Dockerfile
@@ -9,7 +9,9 @@ RUN export DEBIAN_FRONTEND=noninteractive && \
        libbz2-dev libreadline-dev libsqlite3-dev wget \
        curl llvm libncursesw5-dev xz-utils tk-dev \
        libxml2-dev libxmlsec1-dev libffi-dev \
-       liblzma-dev && \
+       liblzma-dev \
+\
+       && \
     apt-get -qqy install ccache git curl && \
     apt-get -qqy --allow-change-held-packages \
             --allow-downgrades install 'libnccl2=2.11.*+cuda11.5' 'libnccl-dev=2.11.*+cuda11.5' 'libcutensor1=1.5.*' 'libcutensor-dev=1.5.*' 'libcusparselt0=0.2.0.*' 'libcusparselt-dev=0.2.0.*' 'libcudnn8=8.4.*+cuda11.6' 'libcudnn8-dev=8.4.*+cuda11.6'
@@ -24,5 +26,5 @@ RUN pyenv install 3.10.0 && \
     pip install -U setuptools pip
 
 RUN pip install -U 'numpy==1.23.*' 'scipy==1.8.*' 'optuna==2.*' 'cython==0.29.*'
-RUN pip uninstall -y cuda-python && \
+RUN pip uninstall -y mpi4py cuda-python && \
     pip check
diff --git a/.pfnci/linux/tests/cuda115.multi.Dockerfile b/.pfnci/linux/tests/cuda115.multi.Dockerfile
--- a/.pfnci/linux/tests/cuda115.multi.Dockerfile
+++ b/.pfnci/linux/tests/cuda115.multi.Dockerfile
@@ -9,7 +9,9 @@ RUN export DEBIAN_FRONTEND=noninteractive && \
        libbz2-dev libreadline-dev libsqlite3-dev wget \
        curl llvm libncursesw5-dev xz-utils tk-dev \
        libxml2-dev libxmlsec1-dev libffi-dev \
-       liblzma-dev && \
+       liblzma-dev \
+       libopenmpi-dev \
+       && \
     apt-get -qqy install ccache git curl && \
     apt-get -qqy --allow-change-held-packages \
             --allow-downgrades install 'libnccl2=2.11.*+cuda11.5' 'libnccl-dev=2.11.*+cuda11.5' 'libcutensor1=1.5.*' 'libcutensor-dev=1.5.*' 'libcusparselt0=0.2.0.*' 'libcusparselt-dev=0.2.0.*' 'libcudnn8=8.4.*+cuda11.6' 'libcudnn8-dev=8.4.*+cuda11.6'
@@ -23,6 +25,6 @@ RUN pyenv install 3.10.0 && \
     pyenv global 3.10.0 && \
     pip install -U setuptools pip
 
-RUN pip install -U 'numpy==1.23.*' 'scipy==1.8.*' 'optuna==2.*' 'cython==0.29.*'
+RUN pip install -U 'numpy==1.23.*' 'scipy==1.8.*' 'optuna==2.*' 'mpi4py==3.*' 'cython==0.29.*'
 RUN pip uninstall -y cuda-python && \
     pip check
diff --git a/.pfnci/linux/tests/cuda115.multi.sh b/.pfnci/linux/tests/cuda115.multi.sh
--- a/.pfnci/linux/tests/cuda115.multi.sh
+++ b/.pfnci/linux/tests/cuda115.multi.sh
@@ -12,5 +12,7 @@ export NVCC="ccache nvcc"
 export CUPY_ACCELERATORS="cutensor,cub"
 
 "$ACTIONS/build.sh"
+export OMPI_ALLOW_RUN_AS_ROOT=1
+export OMPI_ALLOW_RUN_AS_ROOT_CONFIRM=1
 "$ACTIONS/unittest.sh" "not slow and multi_gpu"
 "$ACTIONS/cleanup.sh"
diff --git a/.pfnci/linux/tests/cuda116.Dockerfile b/.pfnci/linux/tests/cuda116.Dockerfile
--- a/.pfnci/linux/tests/cuda116.Dockerfile
+++ b/.pfnci/linux/tests/cuda116.Dockerfile
@@ -9,7 +9,9 @@ RUN export DEBIAN_FRONTEND=noninteractive && \
        libbz2-dev libreadline-dev libsqlite3-dev wget \
        curl llvm libncursesw5-dev xz-utils tk-dev \
        libxml2-dev libxmlsec1-dev libffi-dev \
-       liblzma-dev && \
+       liblzma-dev \
+\
+       && \
     apt-get -qqy install ccache git curl && \
     apt-get -qqy --allow-change-held-packages \
             --allow-downgrades install 'libnccl2=2.11.*+cuda11.6' 'libnccl-dev=2.11.*+cuda11.6' 'libcutensor1=1.5.*' 'libcutensor-dev=1.5.*' 'libcusparselt0=0.2.0.*' 'libcusparselt-dev=0.2.0.*' 'libcudnn8=8.4.*+cuda11.6' 'libcudnn8-dev=8.4.*+cuda11.6'
@@ -24,5 +26,5 @@ RUN pyenv install 3.10.0 && \
     pip install -U setuptools pip
 
 RUN pip install -U 'numpy==1.22.*' 'scipy==1.7.*' 'optuna==2.*' 'cython==0.29.*'
-RUN pip uninstall -y cuda-python && \
+RUN pip uninstall -y mpi4py cuda-python && \
     pip check
diff --git a/.pfnci/linux/tests/cuda116.multi.Dockerfile b/.pfnci/linux/tests/cuda116.multi.Dockerfile
--- a/.pfnci/linux/tests/cuda116.multi.Dockerfile
+++ b/.pfnci/linux/tests/cuda116.multi.Dockerfile
@@ -9,7 +9,9 @@ RUN export DEBIAN_FRONTEND=noninteractive && \
        libbz2-dev libreadline-dev libsqlite3-dev wget \
        curl llvm libncursesw5-dev xz-utils tk-dev \
        libxml2-dev libxmlsec1-dev libffi-dev \
-       liblzma-dev && \
+       liblzma-dev \
+       libopenmpi-dev \
+       && \
     apt-get -qqy install ccache git curl && \
     apt-get -qqy --allow-change-held-packages \
             --allow-downgrades install 'libnccl2=2.11.*+cuda11.6' 'libnccl-dev=2.11.*+cuda11.6' 'libcutensor1=1.5.*' 'libcutensor-dev=1.5.*' 'libcusparselt0=0.2.0.*' 'libcusparselt-dev=0.2.0.*' 'libcudnn8=8.4.*+cuda11.6' 'libcudnn8-dev=8.4.*+cuda11.6'
@@ -23,6 +25,6 @@ RUN pyenv install 3.10.0 && \
     pyenv global 3.10.0 && \
     pip install -U setuptools pip
 
-RUN pip install -U 'numpy==1.22.*' 'scipy==1.7.*' 'optuna==2.*' 'cython==0.29.*'
+RUN pip install -U 'numpy==1.22.*' 'scipy==1.7.*' 'optuna==2.*' 'mpi4py==3.*' 'cython==0.29.*'
 RUN pip uninstall -y cuda-python && \
     pip check
diff --git a/.pfnci/linux/tests/cuda116.multi.sh b/.pfnci/linux/tests/cuda116.multi.sh
--- a/.pfnci/linux/tests/cuda116.multi.sh
+++ b/.pfnci/linux/tests/cuda116.multi.sh
@@ -12,5 +12,7 @@ export NVCC="ccache nvcc"
 export CUPY_ACCELERATORS="cutensor,cub"
 
 "$ACTIONS/build.sh"
+export OMPI_ALLOW_RUN_AS_ROOT=1
+export OMPI_ALLOW_RUN_AS_ROOT_CONFIRM=1
 "$ACTIONS/unittest.sh" "not slow and multi_gpu"
 "$ACTIONS/cleanup.sh"
diff --git a/.pfnci/linux/tests/cuda117.Dockerfile b/.pfnci/linux/tests/cuda117.Dockerfile
--- a/.pfnci/linux/tests/cuda117.Dockerfile
+++ b/.pfnci/linux/tests/cuda117.Dockerfile
@@ -9,7 +9,9 @@ RUN export DEBIAN_FRONTEND=noninteractive && \
        libbz2-dev libreadline-dev libsqlite3-dev wget \
        curl llvm libncursesw5-dev xz-utils tk-dev \
        libxml2-dev libxmlsec1-dev libffi-dev \
-       liblzma-dev && \
+       liblzma-dev \
+\
+       && \
     apt-get -qqy install ccache git curl && \
     apt-get -qqy --allow-change-held-packages \
             --allow-downgrades install 'libcutensor1=1.5.*' 'libcutensor-dev=1.5.*' 'libcusparselt0=0.2.0.*' 'libcusparselt-dev=0.2.0.*' 'libcudnn8=8.4.*+cuda11.6' 'libcudnn8-dev=8.4.*+cuda11.6'
@@ -24,5 +26,5 @@ RUN pyenv install 3.10.0 && \
     pip install -U setuptools pip
 
 RUN pip install -U 'numpy==1.23.*' 'scipy==1.9.*' 'optuna==2.*' 'cython==0.29.*'
-RUN pip uninstall -y cuda-python && \
+RUN pip uninstall -y mpi4py cuda-python && \
     pip check
diff --git a/.pfnci/linux/tests/cuda117.multi.Dockerfile b/.pfnci/linux/tests/cuda117.multi.Dockerfile
--- a/.pfnci/linux/tests/cuda117.multi.Dockerfile
+++ b/.pfnci/linux/tests/cuda117.multi.Dockerfile
@@ -9,7 +9,9 @@ RUN export DEBIAN_FRONTEND=noninteractive && \
        libbz2-dev libreadline-dev libsqlite3-dev wget \
        curl llvm libncursesw5-dev xz-utils tk-dev \
        libxml2-dev libxmlsec1-dev libffi-dev \
-       liblzma-dev && \
+       liblzma-dev \
+       libopenmpi-dev \
+       && \
     apt-get -qqy install ccache git curl && \
     apt-get -qqy --allow-change-held-packages \
             --allow-downgrades install 'libcutensor1=1.5.*' 'libcutensor-dev=1.5.*' 'libcusparselt0=0.2.0.*' 'libcusparselt-dev=0.2.0.*' 'libcudnn8=8.4.*+cuda11.6' 'libcudnn8-dev=8.4.*+cuda11.6'
@@ -23,6 +25,6 @@ RUN pyenv install 3.10.0 && \
     pyenv global 3.10.0 && \
     pip install -U setuptools pip
 
-RUN pip install -U 'numpy==1.23.*' 'scipy==1.9.*' 'optuna==2.*' 'cython==0.29.*'
+RUN pip install -U 'numpy==1.23.*' 'scipy==1.9.*' 'optuna==2.*' 'mpi4py==3.*' 'cython==0.29.*'
 RUN pip uninstall -y cuda-python && \
     pip check
diff --git a/.pfnci/linux/tests/cuda117.multi.sh b/.pfnci/linux/tests/cuda117.multi.sh
--- a/.pfnci/linux/tests/cuda117.multi.sh
+++ b/.pfnci/linux/tests/cuda117.multi.sh
@@ -12,5 +12,7 @@ export NVCC="ccache nvcc"
 export CUPY_ACCELERATORS="cutensor,cub"
 
 "$ACTIONS/build.sh"
+export OMPI_ALLOW_RUN_AS_ROOT=1
+export OMPI_ALLOW_RUN_AS_ROOT_CONFIRM=1
 "$ACTIONS/unittest.sh" "not slow and multi_gpu"
 "$ACTIONS/cleanup.sh"
diff --git a/.pfnci/linux/tests/cuda11x-cuda-python.Dockerfile b/.pfnci/linux/tests/cuda11x-cuda-python.Dockerfile
--- a/.pfnci/linux/tests/cuda11x-cuda-python.Dockerfile
+++ b/.pfnci/linux/tests/cuda11x-cuda-python.Dockerfile
@@ -9,7 +9,9 @@ RUN export DEBIAN_FRONTEND=noninteractive && \
        libbz2-dev libreadline-dev libsqlite3-dev wget \
        curl llvm libncursesw5-dev xz-utils tk-dev \
        libxml2-dev libxmlsec1-dev libffi-dev \
-       liblzma-dev && \
+       liblzma-dev \
+\
+       && \
     apt-get -qqy install ccache git curl && \
     apt-get -qqy --allow-change-held-packages \
             --allow-downgrades install 'libnccl2=2.11.*+cuda11.5' 'libnccl-dev=2.11.*+cuda11.5' 'libcutensor1=1.5.*' 'libcutensor-dev=1.5.*' 'libcusparselt0=0.2.0.*' 'libcusparselt-dev=0.2.0.*' 'libcudnn8=8.4.*+cuda11.6' 'libcudnn8-dev=8.4.*+cuda11.6'
@@ -24,3 +26,5 @@ RUN pyenv install 3.10.0 && \
     pip install -U setuptools pip
 
 RUN pip install -U 'numpy==1.21.*' 'scipy==1.7.*' 'optuna==2.*' 'cython==0.29.*' 'cuda-python==11.*'
+RUN pip uninstall -y mpi4py && \
+    pip check
diff --git a/.pfnci/linux/tests/rocm-4-3.Dockerfile b/.pfnci/linux/tests/rocm-4-3.Dockerfile
--- a/.pfnci/linux/tests/rocm-4-3.Dockerfile
+++ b/.pfnci/linux/tests/rocm-4-3.Dockerfile
@@ -13,7 +13,9 @@ RUN export DEBIAN_FRONTEND=noninteractive && \
        libbz2-dev libreadline-dev libsqlite3-dev wget \
        curl llvm libncursesw5-dev xz-utils tk-dev \
        libxml2-dev libxmlsec1-dev libffi-dev \
-       liblzma-dev && \
+       liblzma-dev \
+\
+       && \
     apt-get -qqy install ccache git curl && \
     apt-get -qqy --allow-change-held-packages \
             --allow-downgrades install rocm-dev hipblas hipfft hipsparse hipcub rocsparse rocrand rocthrust rocsolver rocfft rocprim rccl
@@ -33,5 +35,5 @@ RUN pyenv install 3.7.11 && \
     pip install -U setuptools pip
 
 RUN pip install -U 'numpy==1.20.*' 'scipy==1.6.*' 'optuna==2.*' 'cython==0.29.*'
-RUN pip uninstall -y cuda-python && \
+RUN pip uninstall -y mpi4py cuda-python && \
     pip check
diff --git a/.pfnci/linux/tests/rocm-5-0.Dockerfile b/.pfnci/linux/tests/rocm-5-0.Dockerfile
--- a/.pfnci/linux/tests/rocm-5-0.Dockerfile
+++ b/.pfnci/linux/tests/rocm-5-0.Dockerfile
@@ -13,7 +13,9 @@ RUN export DEBIAN_FRONTEND=noninteractive && \
        libbz2-dev libreadline-dev libsqlite3-dev wget \
        curl llvm libncursesw5-dev xz-utils tk-dev \
        libxml2-dev libxmlsec1-dev libffi-dev \
-       liblzma-dev && \
+       liblzma-dev \
+\
+       && \
     apt-get -qqy install ccache git curl && \
     apt-get -qqy --allow-change-held-packages \
             --allow-downgrades install rocm-dev hipblas hipfft hipsparse hipcub rocsparse rocrand rocthrust rocsolver rocfft rocprim rccl
@@ -33,5 +35,5 @@ RUN pyenv install 3.9.6 && \
     pip install -U setuptools pip
 
 RUN pip install -U 'numpy==1.23.*' 'scipy==1.9.*' 'optuna==2.*' 'cython==0.29.*'
-RUN pip uninstall -y cuda-python && \
+RUN pip uninstall -y mpi4py cuda-python && \
     pip check
diff --git a/.pfnci/linux/tests/rocm-5-1.Dockerfile b/.pfnci/linux/tests/rocm-5-1.Dockerfile
--- a/.pfnci/linux/tests/rocm-5-1.Dockerfile
+++ b/.pfnci/linux/tests/rocm-5-1.Dockerfile
@@ -13,7 +13,9 @@ RUN export DEBIAN_FRONTEND=noninteractive && \
        libbz2-dev libreadline-dev libsqlite3-dev wget \
        curl llvm libncursesw5-dev xz-utils tk-dev \
        libxml2-dev libxmlsec1-dev libffi-dev \
-       liblzma-dev && \
+       liblzma-dev \
+\
+       && \
     apt-get -qqy install ccache git curl && \
     apt-get -qqy --allow-change-held-packages \
             --allow-downgrades install rocm-dev hipblas hipfft hipsparse hipcub rocsparse rocrand rocthrust rocsolver rocfft rocprim rccl
@@ -33,5 +35,5 @@ RUN pyenv install 3.9.6 && \
     pip install -U setuptools pip
 
 RUN pip install -U 'numpy==1.23.*' 'scipy==1.8.*' 'optuna==2.*' 'cython==0.29.*'
-RUN pip uninstall -y cuda-python && \
+RUN pip uninstall -y mpi4py cuda-python && \
     pip check
diff --git a/.pfnci/linux/tests/rocm-5-2.Dockerfile b/.pfnci/linux/tests/rocm-5-2.Dockerfile
--- a/.pfnci/linux/tests/rocm-5-2.Dockerfile
+++ b/.pfnci/linux/tests/rocm-5-2.Dockerfile
@@ -13,7 +13,9 @@ RUN export DEBIAN_FRONTEND=noninteractive && \
        libbz2-dev libreadline-dev libsqlite3-dev wget \
        curl llvm libncursesw5-dev xz-utils tk-dev \
        libxml2-dev libxmlsec1-dev libffi-dev \
-       liblzma-dev && \
+       liblzma-dev \
+\
+       && \
     apt-get -qqy install ccache git curl && \
     apt-get -qqy --allow-change-held-packages \
             --allow-downgrades install rocm-dev hipblas hipfft hipsparse hipcub rocsparse rocrand rocthrust rocsolver rocfft rocprim rccl
@@ -33,5 +35,5 @@ RUN pyenv install 3.9.6 && \
     pip install -U setuptools pip
 
 RUN pip install -U 'numpy==1.23.*' 'scipy==1.8.*' 'optuna==2.*' 'cython==0.29.*'
-RUN pip uninstall -y cuda-python && \
+RUN pip uninstall -y mpi4py cuda-python && \
     pip check
diff --git a/tests/cupyx_tests/distributed_tests/comm_runner.py b/tests/cupyx_tests/distributed_tests/comm_runner.py
--- a/tests/cupyx_tests/distributed_tests/comm_runner.py
+++ b/tests/cupyx_tests/distributed_tests/comm_runner.py
@@ -76,8 +76,8 @@ def run_reduce(rank, root, dtype, use_mpi=False):
     if use_mpi:
         from mpi4py import MPI
         # This process was run with mpiexec
-        run_reduce(MPI.COMM_WORLD.Get_rank(), 0, dtype, False)
-        run_reduce(MPI.COMM_WORLD.Get_rank(), 1, dtype, False)
+        run_reduce(MPI.COMM_WORLD.Get_rank(), 0, dtype, True)
+        run_reduce(MPI.COMM_WORLD.Get_rank(), 1, dtype, True)
     else:
         _launch_workers(run_reduce, (0, dtype))
         _launch_workers(run_reduce, (1, dtype))
@@ -100,7 +100,7 @@ def run_all_reduce(rank, dtype, use_mpi=False):
     if use_mpi:
         from mpi4py import MPI
         # This process was run with mpiexec
-        run_all_reduce(MPI.COMM_WORLD.Get_rank(), dtype, False)
+        run_all_reduce(MPI.COMM_WORLD.Get_rank(), dtype, True)
     else:
         _launch_workers(run_all_reduce, (dtype,))
 
@@ -123,7 +123,7 @@ def run_reduce_scatter(rank, dtype, use_mpi=False):
     if use_mpi:
         from mpi4py import MPI
         # This process was run with mpiexec
-        run_reduce_scatter(MPI.COMM_WORLD.Get_rank(), dtype, False)
+        run_reduce_scatter(MPI.COMM_WORLD.Get_rank(), dtype, True)
     else:
         _launch_workers(run_reduce_scatter, (dtype,))
 
@@ -148,7 +148,7 @@ def run_all_gather(rank, dtype, use_mpi=False):
     if use_mpi:
         from mpi4py import MPI
         # This process was run with mpiexec
-        run_all_gather(MPI.COMM_WORLD.Get_rank(), dtype, False)
+        run_all_gather(MPI.COMM_WORLD.Get_rank(), dtype, True)
     else:
         _launch_workers(run_all_gather, (dtype,))
 
@@ -172,7 +172,7 @@ def run_send_and_recv(rank, dtype, use_mpi=False):
     if use_mpi:
         from mpi4py import MPI
         # This process was run with mpiexec
-        run_send_and_recv(MPI.COMM_WORLD.Get_rank(), dtype, False)
+        run_send_and_recv(MPI.COMM_WORLD.Get_rank(), dtype, True)
     else:
         _launch_workers(run_send_and_recv, (dtype,))
 
@@ -194,7 +194,7 @@ def run_send_recv(rank, dtype, use_mpi=False):
     if use_mpi:
         from mpi4py import MPI
         # This process was run with mpiexec
-        run_send_recv(MPI.COMM_WORLD.Get_rank(), dtype, False)
+        run_send_recv(MPI.COMM_WORLD.Get_rank(), dtype, True)
     else:
         _launch_workers(run_send_recv, (dtype,))
 
@@ -218,8 +218,8 @@ def run_scatter(rank, root, dtype, use_mpi=False):
     if use_mpi:
         from mpi4py import MPI
         # This process was run with mpiexec
-        run_scatter(MPI.COMM_WORLD.Get_rank(), 0, dtype, False)
-        run_scatter(MPI.COMM_WORLD.Get_rank(), 1, dtype, False)
+        run_scatter(MPI.COMM_WORLD.Get_rank(), 0, dtype, True)
+        run_scatter(MPI.COMM_WORLD.Get_rank(), 1, dtype, True)
     else:
         _launch_workers(run_scatter, (0, dtype))
         _launch_workers(run_scatter, (1, dtype))
@@ -245,8 +245,8 @@ def run_gather(rank, root, dtype, use_mpi=False):
     if use_mpi:
         from mpi4py import MPI
         # This process was run with mpiexec
-        run_gather(MPI.COMM_WORLD.Get_rank(), 0, dtype, False)
-        run_gather(MPI.COMM_WORLD.Get_rank(), 1, dtype, False)
+        run_gather(MPI.COMM_WORLD.Get_rank(), 0, dtype, True)
+        run_gather(MPI.COMM_WORLD.Get_rank(), 1, dtype, True)
     else:
         _launch_workers(run_gather, (0, dtype))
         _launch_workers(run_gather, (1, dtype))
@@ -271,7 +271,7 @@ def run_all_to_all(rank, dtype, use_mpi=False):
     if use_mpi:
         from mpi4py import MPI
         # This process was run with mpiexec
-        run_all_to_all(MPI.COMM_WORLD.Get_rank(), dtype, False)
+        run_all_to_all(MPI.COMM_WORLD.Get_rank(), dtype, True)
     else:
         _launch_workers(run_all_to_all, (dtype,))
 
@@ -292,7 +292,7 @@ def run_barrier(rank, use_mpi=False):
     if use_mpi:
         from mpi4py import MPI
         # This process was run with mpiexec
-        run_barrier(MPI.COMM_WORLD.Get_rank(), False)
+        run_barrier(MPI.COMM_WORLD.Get_rank(), True)
     else:
         _launch_workers(run_barrier)
 
@@ -312,8 +312,8 @@ def run_init(rank, use_mpi=False):
     if use_mpi:
         from mpi4py import MPI
         # This process was run with mpiexec
-        run_init(MPI.COMM_WORLD.Get_rank(), dtype, False)
-        run_init(MPI.COMM_WORLD.Get_rank(), dtype, False)
+        run_init(MPI.COMM_WORLD.Get_rank(), dtype, True)
+        run_init(MPI.COMM_WORLD.Get_rank(), dtype, True)
     else:
         _launch_workers(run_init)
 
@@ -350,7 +350,7 @@ def run_send_and_recv(rank, dtype, use_mpi=False):
     if use_mpi:
         from mpi4py import MPI
         # This process was run with mpiexec
-        run_send_and_recv(MPI.COMM_WORLD.Get_rank(), dtype, False)
+        run_send_and_recv(MPI.COMM_WORLD.Get_rank(), dtype, True)
     else:
         _launch_workers(run_send_and_recv, (dtype,))
 
@@ -373,7 +373,7 @@ def run_send_recv(rank, dtype, use_mpi=False):
     if use_mpi:
         from mpi4py import MPI
         # This process was run with mpiexec
-        run_send_recv(MPI.COMM_WORLD.Get_rank(), dtype, False)
+        run_send_recv(MPI.COMM_WORLD.Get_rank(), dtype, True)
     else:
         _launch_workers(run_send_recv, (dtype,))
 
diff --git a/tests/cupyx_tests/distributed_tests/test_comm.py b/tests/cupyx_tests/distributed_tests/test_comm.py
--- a/tests/cupyx_tests/distributed_tests/test_comm.py
+++ b/tests/cupyx_tests/distributed_tests/test_comm.py
@@ -2,6 +2,7 @@
 import subprocess
 import sys
 import unittest
+import os
 
 import numpy
 import pytest
@@ -36,14 +37,16 @@ def _run_test_with_mpi(test_name, dtype=None):
     # subprocess is required not to interfere with cupy module imported in top
     # of this file
     runner_path = pathlib.Path(__file__).parent / 'comm_runner.py'
-    args = ['mpiexec', '-n', '2',
+    args = ['mpiexec', '-n', '2', '--allow-run-as-root',
             sys.executable, runner_path, test_name, 'mpi']
     if dtype is not None:
         args.append(numpy.dtype(dtype).char)
     proc = subprocess.Popen(
         args,
         stdout=subprocess.PIPE,
-        stderr=subprocess.PIPE)
+        stderr=subprocess.PIPE,
+        env=os.environ
+    )
     stdoutdata, stderrdata = proc.communicate()
     assert stderrdata.decode() == ''
     assert proc.returncode == 0