Add collective communication kernels (#3163)

* comm abstraction

* add custom

* fused rms norm

* refactor

* push-based kernel

* optimize for small hidden dims

* integration

* clean up

* export options & fix things

* allgather2d & VMM allocation

* optimize allgather2d

* remove obsolete comm utils

* handle non-multi-gpu build

* fix lint

* fix lint

* avoid using mscclpp repo (some deps are not needed)

* fix lint

* fix nccl version & clean up deps

* fix lint

* custom -> native

* rename

* fix p-lora

* fix lm head

* log fatal exception explicitly

* fix memory location

* fix lora buff size

* pad max_fwd_token_num

* fix allocation for `context_logits_buf_`

* simplify

Author: lzhangzz

CMakeLists.txt

@@ -17,28 +17,16 @@ project(TurboMind LANGUAGES CXX CUDA)
 
 find_package(CUDA 10.2 REQUIRED)
 
+find_package(CUDAToolkit REQUIRED)
+
 if(${CUDA_VERSION_MAJOR} VERSION_GREATER_EQUAL "11")
   add_definitions("-DENABLE_BF16")
   message("CUDA_VERSION ${CUDA_VERSION_MAJOR}.${CUDA_VERSION_MINOR} is greater or equal than 11.0, enable -DENABLE_BF16 flag")
 endif()
 
-# if((${CUDA_VERSION_MAJOR} VERSION_GREATER_EQUAL "11" AND ${CUDA_VERSION_MINOR} VERSION_GREATER_EQUAL "8") OR (${CUDA_VERSION_MAJOR} VERSION_GREATER_EQUAL "12"))
-#   add_definitions("-DENABLE_FP8")
-#   option(ENABLE_FP8 "ENABLE_FP8" OFF)
-#   if(ENABLE_FP8)
-#     message("CUDA_VERSION ${CUDA_VERSION_MAJOR}.${CUDA_VERSION_MINOR} is greater or equal than 11.8, enable -DENABLE_FP8 flag")
-#   endif()
-# endif()
-
 set(CMAKE_MODULE_PATH ${PROJECT_SOURCE_DIR}/cmake/Modules)
 
-option(BUILD_PYT "Build in PyTorch TorchScript class mode" OFF)
-if(NOT BUILD_MULTI_GPU)
-  option(BUILD_MULTI_GPU "Build project about multi-GPU" OFF)
-endif()
-if(NOT USE_TRITONSERVER_DATATYPE)
-  option(USE_TRITONSERVER_DATATYPE "Build triton backend for triton server" OFF)
-endif()
+option(BUILD_MULTI_GPU "Build multi-gpu support" ON)
 option(BUILD_PY_FFI "Build python ffi" ON)
 option(BUILD_TEST "Build tests" OFF)
 
@@ -89,43 +77,24 @@ if (LMDEPLOY_UBSAN_ENABLE)
 endif ()
 
 if(BUILD_MULTI_GPU)
-  message(STATUS "Add DBUILD_MULTI_GPU, requires MPI and NCCL")
-  add_definitions("-DBUILD_MULTI_GPU")
-  set(CMAKE_MODULE_PATH ${PROJECT_SOURCE_DIR}/cmake/Modules)
-  find_package(MPI REQUIRED)
-  find_package(NCCL REQUIRED)
-  set(CMAKE_MODULE_PATH "") # prevent the bugs for pytorch building
+    add_definitions("-DBUILD_MULTI_GPU=1")
+    set(CMAKE_MODULE_PATH ${PROJECT_SOURCE_DIR}/cmake/Modules)
+    find_package(NCCL)
+    if (NCCL_FOUND)
+        set(USE_NCCL ON)
+        add_definitions("-DUSE_NCCL=1")
+    endif ()
 endif()
 
-if(BUILD_PYT)
-  if(DEFINED ENV{NVIDIA_PYTORCH_VERSION})
-    if($ENV{NVIDIA_PYTORCH_VERSION} VERSION_LESS "20.03")
-      message(FATAL_ERROR "NVIDIA PyTorch image is too old for TorchScript mode.")
-    endif()
-    if($ENV{NVIDIA_PYTORCH_VERSION} VERSION_EQUAL "20.03")
-      add_definitions(-DLEGACY_THS=1)
-    endif()
-  endif()
-endif()
-
-if(USE_TRITONSERVER_DATATYPE)
-  message("-- USE_TRITONSERVER_DATATYPE")
-  add_definitions("-DUSE_TRITONSERVER_DATATYPE")
-endif()
 
 set(CXX_STD "17" CACHE STRING "C++ standard")
 # enable gold linker for binary and .so
 set(CMAKE_EXE_LINKER_FLAGS "${CMAKE_EXE_LINKER_FLAGS} -fuse-ld=gold")
 set(CMAKE_SHARED_LINKER_FLAGS "${CMAKE_SHARED_LINKER_FLAGS} -fuse-ld=gold")
 set(CUDA_PATH ${CUDA_TOOLKIT_ROOT_DIR})
 
-set(TF_PATH "" CACHE STRING "TensorFlow path")
 set(CUSPARSELT_PATH "" CACHE STRING "cuSPARSELt path")
 
-if((BUILD_TF OR BUILD_TF2) AND NOT TF_PATH)
-  message(FATAL_ERROR "TF_PATH must be set if BUILD_TF or BUILD_TF2 (=TensorFlow mode) is on.")
-endif()
-
 list(APPEND CMAKE_MODULE_PATH ${CUDA_PATH}/lib64)
 
 # profiling
@@ -204,64 +173,8 @@ if (SPARSITY_SUPPORT)
   add_definitions(-DSPARSITY_ENABLED=1)
 endif()
 
-if(BUILD_TF)
-  list(APPEND COMMON_HEADER_DIRS ${TF_PATH}/include)
-  list(APPEND COMMON_LIB_DIRS ${TF_PATH})
-  add_definitions(-D_GLIBCXX_USE_CXX11_ABI=0)
-endif()
-
-if(BUILD_TF2)
-  list(APPEND COMMON_HEADER_DIRS ${TF_PATH}/include)
-  list(APPEND COMMON_LIB_DIRS ${TF_PATH})
-  add_definitions(-D_GLIBCXX_USE_CXX11_ABI=1)
-endif()
 
 set(PYTHON_PATH "python" CACHE STRING "Python path")
-if(BUILD_PYT)
-  execute_process(COMMAND ${PYTHON_PATH} "-c" "from __future__ import print_function; import torch; print(torch.__version__,end='');"
-                  RESULT_VARIABLE _PYTHON_SUCCESS
-                  OUTPUT_VARIABLE TORCH_VERSION)
-  if (TORCH_VERSION VERSION_LESS "1.5.0")
-      message(FATAL_ERROR "PyTorch >= 1.5.0 is needed for TorchScript mode.")
-  endif()
-  execute_process(COMMAND ${PYTHON_PATH} "-c" "from __future__ import print_function; import os; import torch;
-print(os.path.dirname(torch.__file__),end='');"
-                  RESULT_VARIABLE _PYTHON_SUCCESS
-                  OUTPUT_VARIABLE TORCH_DIR)
-  if (NOT _PYTHON_SUCCESS MATCHES 0)
-      message(FATAL_ERROR "Torch config Error.")
-  endif()
-  list(APPEND CMAKE_PREFIX_PATH ${TORCH_DIR})
-  find_package(Torch REQUIRED)
-  execute_process(COMMAND ${PYTHON_PATH} "-c" "from __future__ import print_function; from distutils import sysconfig;
-print(sysconfig.get_python_inc());"
-                  RESULT_VARIABLE _PYTHON_SUCCESS
-                  OUTPUT_VARIABLE PY_INCLUDE_DIR)
-  if (NOT _PYTHON_SUCCESS MATCHES 0)
-      message(FATAL_ERROR "Python config Error.")
-  endif()
-  list(APPEND COMMON_HEADER_DIRS ${PY_INCLUDE_DIR})
-  execute_process(COMMAND ${PYTHON_PATH} "-c" "from __future__ import print_function; import torch;
-print(torch._C._GLIBCXX_USE_CXX11_ABI,end='');"
-                  RESULT_VARIABLE _PYTHON_SUCCESS
-                  OUTPUT_VARIABLE USE_CXX11_ABI)
-  message("-- USE_CXX11_ABI=${USE_CXX11_ABI}")
-  if (USE_CXX11_ABI)
-    set(CMAKE_CUDA_FLAGS_RELWITHDEBINFO "${CMAKE_CUDA_FLAGS_RELWITHDEBINFO} -D_GLIBCXX_USE_CXX11_ABI=1")
-    set(CMAKE_CXX_FLAGS_RELWITHDEBINFO "${CMAKE_CXX_FLAGS_RELWITHDEBINFO} -D_GLIBCXX_USE_CXX11_ABI=1")
-    set(CMAKE_CUDA_FLAGS_RELEASE "${CMAKE_CUDA_FLAGS_RELEASE} -D_GLIBCXX_USE_CXX11_ABI=1")
-    set(CMAKE_CXX_FLAGS_RELEASE "${CMAKE_CXX_FLAGS_RELEASE} -D_GLIBCXX_USE_CXX11_ABI=1")
-    set(CMAKE_CUDA_FLAGS_DEBUG "${CMAKE_CUDA_FLAGS_DEBUG} -D_GLIBCXX_USE_CXX11_ABI=1")
-    set(CMAKE_CXX_FLAGS_DEBUG "${CMAKE_CXX_FLAGS_DEBUG} -D_GLIBCXX_USE_CXX11_ABI=1")
-  else()
-    set(CMAKE_CUDA_FLAGS_RELWITHDEBINFO "${CMAKE_CUDA_FLAGS_RELWITHDEBINFO} -D_GLIBCXX_USE_CXX11_ABI=0")
-    set(CMAKE_CXX_FLAGS_RELWITHDEBINFO "${CMAKE_CXX_FLAGS_RELWITHDEBINFO} -D_GLIBCXX_USE_CXX11_ABI=0")
-    set(CMAKE_CUDA_FLAGS_RELEASE "${CMAKE_CUDA_FLAGS_RELEASE} -D_GLIBCXX_USE_CXX11_ABI=0")
-    set(CMAKE_CXX_FLAGS_RELEASE "${CMAKE_CXX_FLAGS_RELEASE} -D_GLIBCXX_USE_CXX11_ABI=0")
-    set(CMAKE_CUDA_FLAGS_DEBUG "${CMAKE_CUDA_FLAGS_DEBUG} -D_GLIBCXX_USE_CXX11_ABI=0")
-    set(CMAKE_CXX_FLAGS_DEBUG "${CMAKE_CXX_FLAGS_DEBUG} -D_GLIBCXX_USE_CXX11_ABI=0")
-  endif()
-endif()
 
 # turn off warnings on windows
 if (MSVC)
@@ -286,14 +199,6 @@ if (MSVC)
   endforeach()
 endif()
 
-if (BUILD_MULTI_GPU)
-  list(APPEND COMMON_HEADER_DIRS ${MPI_INCLUDE_PATH})
-endif()
-
-if(USE_TRITONSERVER_DATATYPE)
-  list(APPEND COMMON_HEADER_DIRS ${PROJECT_SOURCE_DIR}/../repo-core-src/include)
-endif()
-
 include_directories(
   ${COMMON_HEADER_DIRS}
 )
@@ -314,111 +219,3 @@ endif()
 if (BUILD_PY_FFI)
   install(TARGETS _turbomind DESTINATION ${CMAKE_SOURCE_DIR}/lmdeploy/lib)
 endif ()
-
-if (MSVC)
-  return()
-endif ()
-
-# # Mesaure the compile time
-option(MEASURE_BUILD_TIME "Measure the build time of each module" OFF)
-if (MEASURE_BUILD_TIME)
-  set_property(GLOBAL PROPERTY RULE_LAUNCH_COMPILE "${CMAKE_COMMAND} -E time")
-  set_property(GLOBAL PROPERTY RULE_LAUNCH_CUSTOM "${CMAKE_COMMAND} -E time")
-  set_property(GLOBAL PROPERTY RULE_LAUNCH_LINK "${CMAKE_COMMAND} -E time")
-endif()
-
-########################################
-
-add_library(transformer-shared SHARED
-  $<TARGET_OBJECTS:DynamicDecodeLayer>
-  $<TARGET_OBJECTS:Llama>
-  $<TARGET_OBJECTS:LlamaTritonBackend>
-  $<TARGET_OBJECTS:TransformerTritonBackend>
-  $<TARGET_OBJECTS:activation_kernels>
-  $<TARGET_OBJECTS:ban_bad_words>
-  $<TARGET_OBJECTS:cublasAlgoMap>
-  $<TARGET_OBJECTS:cublasMMWrapper>
-  $<TARGET_OBJECTS:cuda_utils>
-  $<TARGET_OBJECTS:custom_ar_comm>
-  $<TARGET_OBJECTS:custom_ar_kernels>
-  $<TARGET_OBJECTS:attention>
-  $<TARGET_OBJECTS:decoding_kernels>
-  $<TARGET_OBJECTS:gpt_kernels>
-  $<TARGET_OBJECTS:logprob_kernels>
-  $<TARGET_OBJECTS:logger>
-  $<TARGET_OBJECTS:memory_utils>
-  $<TARGET_OBJECTS:mpi_utils>
-  $<TARGET_OBJECTS:nccl_utils>
-  $<TARGET_OBJECTS:nvtx_utils>
-  $<TARGET_OBJECTS:anomaly_handler>
-  $<TARGET_OBJECTS:sampling_penalty_kernels>
-  $<TARGET_OBJECTS:sampling_topk_kernels>
-  $<TARGET_OBJECTS:sampling_topp_kernels>
-  $<TARGET_OBJECTS:stop_criteria>
-  $<TARGET_OBJECTS:tensor>
-  $<TARGET_OBJECTS:unfused_attention_kernels>
-)
-
-if (BUILD_MULTI_GPU)
-target_link_libraries(transformer-shared PUBLIC
-  ${MPI_CXX_LIBRARIES}
-  ${NCCL_LIBRARIES}
-)
-endif()
-
-if(USE_NVTX)
-target_link_libraries(transformer-shared PUBLIC
-  -lnvToolsExt
-)
-endif()
-
-set_target_properties(transformer-shared PROPERTIES POSITION_INDEPENDENT_CODE ON)
-set_target_properties(transformer-shared PROPERTIES CUDA_RESOLVE_DEVICE_SYMBOLS ON)
-set_target_properties(transformer-shared PROPERTIES LINKER_LANGUAGE CXX)
-target_link_libraries(transformer-shared PUBLIC -lcudart -lcublas -lcublasLt -lcurand)
-
-include(GNUInstallDirs)
-set(INSTALL_CONFIGDIR ${CMAKE_INSTALL_LIBDIR}/cmake/TurboMind)
-
-include(CMakePackageConfigHelpers)
-configure_package_config_file(
-  ${CMAKE_CURRENT_LIST_DIR}/cmake/TurboMindConfig.cmake.in
-  ${CMAKE_CURRENT_BINARY_DIR}/TurboMindConfig.cmake
-  INSTALL_DESTINATION ${INSTALL_CONFIGDIR}
-)
-
-install(
-  FILES
-  ${CMAKE_CURRENT_BINARY_DIR}/TurboMindConfig.cmake
-  DESTINATION ${INSTALL_CONFIGDIR}
-)
-
-install(
-  TARGETS
-    transformer-shared
-  EXPORT
-    transformer-shared-targets
-  LIBRARY DESTINATION ${CMAKE_INSTALL_PREFIX}/backends/turbomind
-  ARCHIVE DESTINATION ${CMAKE_INSTALL_PREFIX}/backends/turbomind
-  RUNTIME DESTINATION ${CMAKE_INSTALL_PREFIX}/bin
-)
-
-install(
-  EXPORT
-    transformer-shared-targets
-  FILE
-    TurboMindTargets.cmake
-  DESTINATION
-    ${INSTALL_CONFIGDIR}
-)
-
-export(
-  EXPORT
-    transformer-shared-targets
-  FILE
-    ${CMAKE_CURRENT_BINARY_DIR}/TurboMindTargets.cmake
-  NAMESPACE
-    TritonCore::
-)
-
-export(PACKAGE TurboMind)

benchmark/profile_pipeline_api.py

@@ -167,6 +167,7 @@ def parse_args():
     ArgumentHelper.quant_policy(tb_group, default=0)
     ArgumentHelper.num_tokens_per_iter(tb_group)
     ArgumentHelper.max_prefill_iters(tb_group)
+    ArgumentHelper.communicator(tb_group)
 
     args = parser.parse_args()
     return args
@@ -188,6 +189,7 @@ def main():
             num_tokens_per_iter=args.num_tokens_per_iter,
             max_prefill_iters=args.max_prefill_iters,
             enable_prefix_caching=args.enable_prefix_caching,
+            communicator=args.communicator,
         )
     elif args.backend == 'pytorch':
         engine_config = PytorchEngineConfig(

benchmark/profile_throughput.py

@@ -221,6 +221,7 @@ def parse_args():
     ArgumentHelper.model_format(tb_group, default='hf')
     ArgumentHelper.num_tokens_per_iter(tb_group)
     ArgumentHelper.max_prefill_iters(tb_group)
+    ArgumentHelper.communicator(tb_group)
 
     args = parser.parse_args()
     return args
@@ -242,6 +243,7 @@ def main():
             max_prefill_iters=args.max_prefill_iters,
             enable_prefix_caching=args.enable_prefix_caching,
             dtype=args.dtype,
+            communicator=args.communicator,
         )
     elif args.backend == 'pytorch':
         engine_config = PytorchEngineConfig(

lmdeploy/cli/serve.py

@@ -76,6 +76,7 @@ def add_parser_gradio():
         ArgumentHelper.model_format(tb_group)
         ArgumentHelper.quant_policy(tb_group)
         ArgumentHelper.rope_scaling_factor(tb_group)
+        ArgumentHelper.communicator(tb_group)
 
     @staticmethod
     def add_parser_api_server():
@@ -173,6 +174,7 @@ def add_parser_api_server():
         ArgumentHelper.rope_scaling_factor(tb_group)
         ArgumentHelper.num_tokens_per_iter(tb_group)
         ArgumentHelper.max_prefill_iters(tb_group)
+        ArgumentHelper.communicator(tb_group)
 
         # vlm args
         vision_group = parser.add_argument_group('Vision model arguments')
@@ -255,7 +257,8 @@ def gradio(args):
                                                    cache_max_entry_count=args.cache_max_entry_count,
                                                    cache_block_seq_len=args.cache_block_seq_len,
                                                    enable_prefix_caching=args.enable_prefix_caching,
-                                                   max_prefill_token_num=args.max_prefill_token_num)
+                                                   max_prefill_token_num=args.max_prefill_token_num,
+                                                   communicator=args.communicator)
         chat_template_config = get_chat_template(args.chat_template)
         run(args.model_path_or_server,
             server_name=args.server_name,
@@ -305,7 +308,8 @@ def api_server(args):
                                                    cache_max_entry_count=args.cache_max_entry_count,
                                                    cache_block_seq_len=args.cache_block_seq_len,
                                                    enable_prefix_caching=args.enable_prefix_caching,
-                                                   max_prefill_token_num=args.max_prefill_token_num)
+                                                   max_prefill_token_num=args.max_prefill_token_num,
+                                                   communicator=args.communicator)
         chat_template_config = get_chat_template(args.chat_template)
 
         from lmdeploy.messages import VisionConfig

lmdeploy/cli/utils.py

@@ -463,3 +463,11 @@ def eager_mode(parser):
                                    default=False,
                                    help='Whether to enable eager mode. '
                                    'If True, cuda graph would be disabled')
+
+    @staticmethod
+    def communicator(parser):
+        return parser.add_argument('--communicator',
+                                   type=str,
+                                   default='nccl',
+                                   choices=['nccl', 'native'],
+                                   help='Communication backend for multi-GPU inference')

lmdeploy/messages.py

@@ -223,6 +223,7 @@ class TurbomindEngineConfig:
     max_prefill_token_num: int = 8192
     num_tokens_per_iter: int = 0
     max_prefill_iters: int = 1
+    communicator: str = 'nccl'
 
     def __post_init__(self):
         """Check input validation."""

lmdeploy/turbomind/chat.py

@@ -69,6 +69,7 @@ def main(model_path: str,
          stream_output: bool = True,
          request_output_len: int = 1024,
          chat_template_config: ChatTemplateConfig = None,
+         communicator: str = 'nccl',
          **kwargs):
     """An example to perform model inference through the command line
     interface.
@@ -130,7 +131,8 @@ def main(model_path: str,
                                        quant_policy=quant_policy,
                                        rope_scaling_factor=rope_scaling_factor,
                                        dtype=dtype,
-                                       tp=tp)
+                                       tp=tp,
+                                       communicator=communicator)
     print('engine_cfg:\n', engine_cfg, sep='', flush=True)
     tokenizer = Tokenizer(model_path)
     from lmdeploy import turbomind as tm