【Comm】Add FlagCX Comm Context (#71924)

Author: mikethegoblin

.gitmodules

@@ -129,3 +129,7 @@
 	path = third_party/openvino
 	url = https://github.com/openvinotoolkit/openvino.git
 	ignore = dirty
+[submodule "third_party/flagcx"]
+	path = third_party/flagcx
+	url = https://github.com/FlagOpen/FlagCX.git
+	ignore = dirty

CMakeLists.txt

@@ -313,6 +313,7 @@ option(
   OFF)
 option(WITH_CINN "Compile PaddlePaddle with CINN" OFF)
 option(WITH_NCCL "Compile PaddlePaddle with NCCL support" ON)
+option(WITH_FLAGCX "Compile PaddlePaddle with FLAGCX support" OFF)
 option(WITH_RCCL "Compile PaddlePaddle with RCCL support" ON)
 option(WITH_XPU_BKCL "Compile PaddlePaddle with BAIDU KUNLUN XPU BKCL" OFF)
 option(WITH_CRYPTO "Compile PaddlePaddle with crypto support" ON)
@@ -538,6 +539,10 @@ else()
   endif()
 endif()
 
+if(WITH_FLAGCX)
+  add_definitions("-DPADDLE_WITH_FLAGCX")
+endif()
+
 if(WITH_HETERPS AND WITH_PSLIB)
   set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} -D_GLIBCXX_USE_CXX11_ABI=0")
 endif()

cmake/external/flagcx.cmake

@@ -0,0 +1,47 @@
+set(CMAKE_FIND_DEBUG_MODE ON)
+# flagcx.cmake
+if(NOT WITH_FLAGCX)
+  return()
+endif()
+
+set(FLAGCX_SOURCE_DIR "${PADDLE_SOURCE_DIR}/third_party/flagcx")
+set(FLAGCX_BINARY_DIR "${PADDLE_SOURCE_DIR}/build/third_party/flagcx")
+set(THIRD_PARTY_DIR "${PADDLE_SOURCE_DIR}/build/third_party")
+set(FLAGCX_ROOT "/usr/local/flagcx")
+set(FLAGCX_LIB_DIR "${FLAGCX_BINARY_DIR}/build/lib")
+set(USR_LOCAL_DIR "/usr/local")
+
+file(REMOVE_RECURSE ${FLAGCX_BINARY_DIR})
+message(STATUS "removed old flagcx dir")
+message(STATUS "Copying third-party source to build directory")
+execute_process(COMMAND cp -r ${FLAGCX_SOURCE_DIR} ${THIRD_PARTY_DIR}
+                RESULT_VARIABLE COPY_RESULT)
+
+if(NOT COPY_RESULT EQUAL 0)
+  message(FATAL_ERROR "Failed to copy third-party source to build directory")
+endif()
+
+# Create a custom target to build the third-party library
+message(STATUS "Building third-party library with its Makefile")
+execute_process(
+  COMMAND make
+  WORKING_DIRECTORY ${FLAGCX_BINARY_DIR}
+  RESULT_VARIABLE BUILD_RESULT)
+
+find_path(
+  FLAGCX_INCLUDE_DIR flagcx.h
+  PATHS ${FLAGCX_SOURCE_DIR}/flagcx/include
+  NO_DEFAULT_PATH)
+
+message(STATUS "FLAGCX_INCLUDE_DIR is ${FLAGCX_INCLUDE_DIR}")
+include_directories(SYSTEM ${FLAGCX_INCLUDE_DIR})
+
+add_library(flagcx INTERFACE)
+find_library(
+  FLAGCX_LIB
+  NAMES flagcx libflagcx
+  PATHS ${FLAGCX_LIB_DIR}
+  DOC "My custom library")
+
+add_dependencies(flagcx FLAGCX_LIB)
+message(STATUS "FLAGCX_LIB is ${FLAGCX_LIB}")

cmake/third_party.cmake

@@ -470,6 +470,11 @@ if(WITH_TESTING OR WITH_DISTRIBUTE)
   list(APPEND third_party_deps extern_gtest)
 endif()
 
+if(WITH_FLAGCX)
+  include(external/flagcx)
+  list(APPEND third_party_deps flagcx)
+endif()
+
 if(WITH_ONNXRUNTIME)
   include(external/onnxruntime
   )# download, build, install onnxruntime、paddle2onnx

paddle/phi/CMakeLists.txt

@@ -87,6 +87,10 @@ if(WITH_GLOO)
   list(APPEND PHI_DEPS gloo)
 endif()
 
+if(WITH_FLAGCX)
+  list(APPEND PHI_DEPS flagcx)
+endif()
+
 if(WITH_CUDNN_FRONTEND)
   list(APPEND PHI_DEPS cudnn-frontend)
 endif()

paddle/phi/backends/dynload/CMakeLists.txt

@@ -80,6 +80,10 @@ if(WITH_XPU)
   collect_srcs(backends_srcs SRCS xpti.cc)
 endif()
 
+if(WITH_FLAGCX)
+  collect_srcs(backends_srcs SRCS flagcx.cc)
+endif()
+
 if(WITH_FLASHATTN)
   list(APPEND DYNLOAD_COMMON_SRCS flashattn.cc)
 endif()

paddle/phi/backends/dynload/dynamic_loader.cc

@@ -68,6 +68,17 @@ PHI_DEFINE_string(rccl_dir,
                   "dlopen will search rccl from LD_LIBRARY_PATH");
 #endif
 
+#ifdef PADDLE_WITH_FLAGCX
+COMMON_DECLARE_string(flagcx_dir);
+#endif
+
+PHI_DEFINE_EXPORTED_string(
+    flagcx_dir,  // NOLINT
+    "",
+    "Specify path for loading libflagcx.so. For instance, "
+    "For instance, /usr/local/flagcx/lib. If default, "
+    "dlopen will search flagcx from LD_LIBRARY_PATH");
+
 #ifdef PADDLE_WITH_XPU
 PD_DEFINE_string(xpti_dir, "", "Specify path for loading libxpti.so.");
 #endif
@@ -777,6 +788,14 @@ void* GetNCCLDsoHandle() {
 #endif
 }
 
+void* GetFLAGCXDsoHandle() {
+#ifdef PADDLE_WITH_FLAGCX
+  return GetDsoHandleFromSearchPath(FLAGS_flagcx_dir, "libflagcx.so");
+#else
+  return nullptr;
+#endif
+}
+
 void* GetTensorRtDsoHandle() {
 #if defined(__APPLE__) || defined(__OSX__)
   return GetDsoHandleFromSearchPath(FLAGS_tensorrt_dir, "libnvinfer.dylib");

paddle/phi/backends/dynload/dynamic_loader.h

@@ -39,6 +39,7 @@ void* GetWarpRNNTDsoHandle();
 void* GetFlashAttnDsoHandle();
 void* GetFlashAttnV3DsoHandle();
 void* GetNCCLDsoHandle();
+void* GetFLAGCXDsoHandle();
 void* GetTensorRtDsoHandle();
 void* GetMKLMLDsoHandle();
 void* GetLAPACKDsoHandle();

paddle/phi/backends/dynload/flagcx.cc

@@ -0,0 +1,25 @@
+/* Copyright (c) 2016 PaddlePaddle Authors. All Rights Reserved.
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+    http://www.apache.org/licenses/LICENSE-2.0
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License. */
+
+#include "paddle/phi/backends/dynload/flagcx.h"
+
+namespace phi {
+namespace dynload {
+
+std::once_flag flagcx_dso_flag;
+void* flagcx_dso_handle = nullptr;
+
+#define DEFINE_WRAP(__name) DynLoad__##__name __name
+
+FLAGCX_RAND_ROUTINE_EACH(DEFINE_WRAP);
+
+}  // namespace dynload
+}  // namespace phi

paddle/phi/backends/dynload/flagcx.h

@@ -0,0 +1,67 @@
+/* Copyright (c) 2016 PaddlePaddle Authors. All Rights Reserved.
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+    http://www.apache.org/licenses/LICENSE-2.0
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License. */
+#pragma once
+
+#include <flagcx.h>
+
+#include <mutex>  // NOLINT
+
+#include "paddle/phi/backends/dynload/dynamic_loader.h"
+#include "paddle/phi/common/port.h"
+
+namespace phi {
+namespace dynload {
+
+extern std::once_flag flagcx_dso_flag;
+extern void* flagcx_dso_handle;
+
+#define DECLARE_DYNAMIC_LOAD_FLAGCX_WRAP(__name)                     \
+  struct DynLoad__##__name {                                         \
+    template <typename... Args>                                      \
+    auto operator()(Args... args) -> DECLARE_TYPE(__name, args...) { \
+      using flagcx_func = decltype(&::__name);                       \
+      std::call_once(flagcx_dso_flag, []() {                         \
+        flagcx_dso_handle = phi::dynload::GetFLAGCXDsoHandle();      \
+      });                                                            \
+      static void* p_##__name = dlsym(flagcx_dso_handle, #__name);   \
+      return reinterpret_cast<flagcx_func>(p_##__name)(args...);     \
+    }                                                                \
+  };                                                                 \
+  extern struct DynLoad__##__name __name
+
+#define FLAGCX_RAND_ROUTINE_EACH(__macro) \
+  __macro(flagcxGetUniqueId);             \
+  __macro(flagcxCommInitRank);            \
+  __macro(flagcxGetVersion);              \
+  __macro(flagcxCommAbort);               \
+  __macro(flagcxCommDestroy);             \
+  __macro(flagcxCommCount);               \
+  __macro(flagcxCommUserRank);            \
+  __macro(flagcxAllReduce);               \
+  __macro(flagcxBroadcast);               \
+  __macro(flagcxAllGather);               \
+  __macro(flagcxGroupStart);              \
+  __macro(flagcxGroupEnd);                \
+  __macro(flagcxReduce);                  \
+  __macro(flagcxReduceScatter);           \
+  __macro(flagcxCommGetAsyncError);       \
+  __macro(flagcxSend);                    \
+  __macro(flagcxRecv);                    \
+  __macro(flagcxHandleInit);              \
+  __macro(flagcxHandleFree);              \
+  __macro(flagcxGetErrorString);
+
+FLAGCX_RAND_ROUTINE_EACH(DECLARE_DYNAMIC_LOAD_FLAGCX_WRAP)
+
+#undef DECLARE_DYNAMIC_LOAD_FLAGCX_WRAP
+
+}  // namespace dynload
+}  // namespace phi

paddle/phi/core/distributed/CMakeLists.txt

@@ -23,4 +23,8 @@ if(WITH_XPU_BKCL)
   list(APPEND DISTRIBUTED_COMMON_SRCS bkcl_comm_context.cc)
 endif()
 
+if(WITH_FLAGCX)
+  list(APPEND DISTRIBUTED_COMMON_SRCS flagcx_comm_context.cc flagcx_tools.cc)
+endif()
+
 collect_srcs(core_srcs SRCS ${DISTRIBUTED_COMMON_SRCS})

paddle/phi/core/distributed/comm_context_manager.cc

@@ -44,6 +44,11 @@
 #include "paddle/phi/core/distributed/xccl_comm_context.h"
 #endif
 
+#if defined(PADDLE_WITH_FLAGCX)
+#include "paddle/phi/core/distributed/flagcx_comm_context.h"
+#include "paddle/phi/core/distributed/flagcx_tools.h"
+#endif
+
 namespace phi::distributed {
 
 int CommContextManager::device_id = -1;
@@ -261,6 +266,51 @@ void CommContextManager::CreateBKCLCommContext(
   comm_context_manager.Emplace(unique_comm_key, std::move(bkcl_comm_context));
 }
 #endif
+
+#if defined(PADDLE_WITH_FLAGCX)
+void CommContextManager::CreateFlagcxCommContext(
+    const std::shared_ptr<Store>& store,
+    const std::string& unique_comm_key,
+    int rank,
+    int size,
+    const std::string& hash_key) {
+  auto& comm_context_manager = CommContextManager::GetInstance();
+  if (comm_context_manager.Has(unique_comm_key)) {
+    return;
+  }
+  flagcxHandlerGroup_t flagcx_handler;
+  phi::dynload::flagcxHandleInit(&flagcx_handler);
+  if (rank == 0) {
+    phi::dynload::flagcxGetUniqueId(&flagcx_handler->uniqueId);
+  }
+
+  std::string unique_key = "FlagcxCommContext/" + unique_comm_key + hash_key;
+  if (rank == 0) {
+    std::vector<uint8_t> flagcx_id_wrapper(
+        reinterpret_cast<uint8_t*>(flagcx_handler->uniqueId),
+        reinterpret_cast<uint8_t*>(flagcx_handler->uniqueId) +
+            sizeof(flagcxUniqueId));
+    store->set(unique_key, flagcx_id_wrapper);
+  } else {
+    const auto& flagcx_id_wrapper = store->get(unique_key);
+    std::memcpy(reinterpret_cast<uint8_t*>(flagcx_handler->uniqueId),
+                flagcx_id_wrapper.data(),
+                flagcx_id_wrapper.size());
+  }
+
+  VLOG(3) << "init FlagcxCommContext rank: " << rank << ", size: " << size
+          << ", unique_comm_key: " << unique_comm_key
+          << ", unique_key: " << unique_key << ", flagcx_id: "
+          << SerializeFlagcxUniqueId(*flagcx_handler->uniqueId);
+  auto flagcx_comm_context =
+      std::make_unique<FlagcxCommContext>(rank, size, flagcx_handler);
+  // TODO(changtao): find a way to manage different device context,
+  //  now we use cuda device context as default
+  comm_context_manager.SetStore(store);
+  comm_context_manager.Emplace(unique_comm_key, std::move(flagcx_comm_context));
+}
+#endif
+
 CommContext* CommContextManager::Emplace(
     const std::string& unique_comm_key,
     std::unique_ptr<CommContext> comm_context) {

paddle/phi/core/distributed/comm_context_manager.h

@@ -28,6 +28,10 @@
 #include "paddle/phi/backends/gpu/forwards.h"
 #endif
 
+#if defined(PADDLE_WITH_FLAGCX)
+#include <flagcx.h>
+#endif
+
 namespace phi {
 namespace distributed {
 
@@ -105,6 +109,14 @@ class CommContextManager {
                                     const std::string& hash_key = "");
 #endif
 
+#if defined(PADDLE_WITH_FLAGCX)
+  static void CreateFlagcxCommContext(const std::shared_ptr<Store>& store,
+                                      const std::string& unique_comm_key,
+                                      int rank,
+                                      int size,
+                                      const std::string& hash_key = "");
+#endif
+
  private:
   DISABLE_COPY_AND_ASSIGN(CommContextManager);