Fix MPI related issue in GPT model parallel (#1894)

* Add mpi include dir into cmake.
Support more MPI environment vars.
Fix HOST confict of MPI and Paddle.

* Fix pointer convert in gpt_op.

Author: guoshengCS

paddlenlp/ops/CMakeLists.txt

@@ -38,6 +38,7 @@ option(WITH_PARALLEL    "Compile with model parallel for GPT"
 set(CMAKE_MODULE_PATH ${CMAKE_MODULE_PATH} "${CMAKE_CURRENT_SOURCE_DIR}/cmake")
 if(WITH_PARALLEL)
   # https://cmake.org/cmake/help/latest/module/FindMPI.html#variables-for-locating-mpi
+  # https://github.com/Kitware/CMake/blob/master/Modules/FindMPI.cmake
   find_package(MPI REQUIRED)
   find_package(NCCL REQUIRED)
   add_definitions(-DBUILD_GPT)
@@ -241,6 +242,12 @@ set(COMMON_LIB_DIRS
   ${CUDA_PATH}/lib64
 )
 
+if(WITH_PARALLEL)
+  list(APPEND COMMON_HEADER_DIRS
+    ${NCCL_INCLUDE_PATH}
+    ${MPI_INCLUDE_PATH})
+endif()
+
 set(THIRD_PATH "third-party")
 set(THIRD_PARTY_NAME "fastertransformer")
 

paddlenlp/ops/faster_transformer/src/fusion_gpt_op.cu

@@ -9,12 +9,14 @@
 #include <vector>
 
 // TODO(guosheng): `HOST` conflict exists in float.h of paddle and mpi.h of mpi
+#include "fusion_gpt_op.h"
+#include "pd_traits.h"
+#ifdef HOST
+#undef HOST
+#endif
 #include "fastertransformer/cuda/cub/cub.cuh"
 #include "fastertransformer/gpt.h"
-#include "fastertransformer/open_decoder.h"
 #include "fastertransformer/utils/common.h"
-#include "fusion_gpt_op.h"
-#include "pd_traits.h"
 
 #ifdef BUILD_GPT  // consistent with FasterTransformer
 #include <map>
@@ -274,14 +276,14 @@ std::vector<paddle::Tensor> gpt2_kernel(
   const int hidden_unit = size_per_head * n_head;
 
 #ifdef BUILD_GPT
-  auto* model_para_desc =
-      ModelParaDescFactory::CreateModelParaDesc(n_head,
-                                                size_per_head,
-                                                num_layer,
-                                                tensor_para_size,
-                                                layer_para_size,
-                                                layer_para_batch_size,
-                                                word_emb.data<data_t_>());
+  auto* model_para_desc = ModelParaDescFactory::CreateModelParaDesc(
+      n_head,
+      size_per_head,
+      num_layer,
+      tensor_para_size,
+      layer_para_size,
+      layer_para_batch_size,
+      const_cast<data_t_*>(word_emb.data<data_t_>()));
   auto& tensor_parallel_param = model_para_desc->tensor_parallel_param;
   auto& layer_parallel_param = model_para_desc->layer_parallel_param;
   auto seed = model_para_desc->dist(model_para_desc->gen);

paddlenlp/ops/faster_transformer/transformer/decoding.py

@@ -1175,29 +1175,40 @@ def __init__(self,
                  tensor_para_size=None,
                  layer_para_size=None,
                  layer_para_batch_size=1):
-        # Maybe we should import mpi4py later.
-        self.word_size = int(
-            os.environ.get(
-                "MPI_LOCALNRANKS",  # MPICH
-                os.environ.get("OMPI_COMM_WORLD_SIZE", 1)))  # OpenMPI
-        self.rank = int(
-            os.environ.get(
-                "MPI_LOCALRANKID",  # MPICH
-                os.environ.get("OMPI_COMM_WORLD_RANK", 0)))  # OpenMPI
-        if tensor_para_size is None: tensor_para_size = 1
-        if layer_para_size is None:
-            layer_para_size = self.word_size // tensor_para_size
+        self.world_size = self._env2int(
+            [  # MPICH, OpenMPI, IMPI
+                "MPI_LOCALNRANKS", "OMPI_COMM_WORLD_SIZE", "PMI_SIZE",
+                "MV2_COMM_WORLD_SIZE", "WORLD_SIZE"
+            ],
+            1)
+        self.rank = self._env2int(
+            [  # MPICH, OpenMPI, IMPI
+                "MPI_LOCALRANKID", "OMPI_COMM_WORLD_RANK", "PMI_RANK",
+                "MV2_COMM_WORLD_RANK", "RANK"
+            ],
+            0)
+        if layer_para_size is None: layer_para_size = 1
+        if tensor_para_size is None:
+            tensor_para_size = self.world_size // layer_para_size
         self.no_para = tensor_para_size == 1 and layer_para_size == 1
         self.tensor_para_size = tensor_para_size
         self.layer_para_size = layer_para_size
         self.layer_para_batch_size = layer_para_batch_size
 
-        assert self.word_size == tensor_para_size * layer_para_size, (
+        assert self.world_size == tensor_para_size * layer_para_size, (
             "tensor_para_size * layer_para_size must be equal to world_size.")
         self.tensor_para_rank = self.rank % self.tensor_para_size
         self.layer_para_rank = self.rank // self.tensor_para_size
         self.is_partial_model = False
 
+    @staticmethod
+    def _env2int(env_list, default=-1):
+        for e in env_list:
+            val = int(os.environ.get(e, -1))
+            if val >= 0:
+                return val
+        return default
+
     def is_last_group(self):
         r"""
         For layer parallel, only the process corresponding to the last layer