[WIP] feat(dcu): add BW1000 DCU backend support

CMakeLists.txt

@@ -1,11 +1,17 @@
-cmake_minimum_required(VERSION 3.28)
-
+# Platforms
 option(USE_CUDA "Support NVIDIA CUDA" OFF)
+option(USE_MACA "Support MetaX MACA" OFF)
+option(USE_DCU "Support Hygon DCU through DTK/HIP" OFF)
+
 option(PROFILE_MODE "ENABLE PROFILE MODE" OFF)
 option(USE_OMP "Use OpenMP as backend for Eigen" ON)
-option(USE_NCCL "Build project for distributed running" ON)
+option(USE_NCCL "Build project for distributed running on CUDA using NCCL" ON)
+option(USE_MCCL "Build project for distributed running on MACA using MCCL" ON)
+option(USE_RCCL "Build project for distributed running on DCU using RCCL" OFF)
+option(USE_MPI "Enable MPI for inter-node CPU communication" ON)
+cmake_minimum_required(VERSION 3.28)
 
-project(infini_train VERSION 0.5.0 LANGUAGES CXX)
+project(infini_train VERSION 0.3.0 LANGUAGES CXX)
 
 set(CMAKE_CXX_STANDARD 20)
 set(CMAKE_CXX_STANDARD_REQUIRED ON)
@@ -14,63 +20,104 @@ set(CMAKE_CXX_EXTENSIONS OFF)
 # Generate compile_commands.json
 set(CMAKE_EXPORT_COMPILE_COMMANDS ON)
 
-# ------------------------------------------------------------------------------
-# Third-party deps
-# ------------------------------------------------------------------------------
-
-# gflags
+# Add gflags
 add_subdirectory(third_party/gflags)
 include_directories(${gflags_SOURCE_DIR}/include)
 
-# glog
 set(WITH_GFLAGS OFF CACHE BOOL "Disable glog finding system gflags" FORCE)
 set(WITH_GTEST OFF CACHE BOOL "Disable glog finding system gtest" FORCE)
+
+# Add glog
 add_subdirectory(third_party/glog)
 include_directories(${glog_SOURCE_DIR}/src)
 
-# eigen
+# Add eigen
 if(USE_OMP)
-  find_package(OpenMP REQUIRED)
+    find_package(OpenMP REQUIRED)
+
+    set(INFINI_OMP_LIBS OpenMP::OpenMP_CXX)
+
+    # Under MACA/mxcc, use mxomp instead of original libgomp
+    if(USE_MACA)
+        set(MACA_PATH $ENV{MACA_PATH})
+        find_library(OMP_RUNTIME_LIB
+            NAMES omp iomp5
+            HINTS
+                "${MACA_PATH}/lib"
+                "${MACA_PATH}/mxgpu_llvm/lib"
+                "${MACA_PATH}/mxgpu_llvm/lib64"
+            REQUIRED
+        )
+
+        set(INFINI_OMP_LIBS OpenMP::OpenMP_CXX ${OMP_RUNTIME_LIB})
+    endif()
 endif()
+
+# find_package(OpenBLAS REQUIRED)
+# include_directories(${OpenBLAS_INCLUDE_DIR})
+
 add_subdirectory(third_party/eigen)
 include_directories(${PROJECT_SOURCE_DIR}/third_party/eigen)
+# add_definitions(-DEIGEN_USE_BLAS)
 
 include_directories(${PROJECT_SOURCE_DIR})
-
-if(PROFILE_MODE)
-  add_compile_definitions(PROFILE_MODE=1)
-endif()
-
-# ------------------------------------------------------------------------------
-# Sources
-# ------------------------------------------------------------------------------
-
-# Framework core sources (*.cc), excluding cpu kernels (they are built separately)
 file(GLOB_RECURSE SRC ${PROJECT_SOURCE_DIR}/infini_train/src/*.cc)
 list(FILTER SRC EXCLUDE REGEX ".*kernels/cpu/.*")
-if(NOT USE_NCCL)
-  list(FILTER SRC EXCLUDE REGEX ".*infini_train/src/core/ccl/cuda/.*")
+if(NOT USE_CUDA)
+    list(FILTER SRC EXCLUDE REGEX ".*/(ccl|runtime)/cuda/.*")
+endif()
+if(NOT USE_MACA)
+    list(FILTER SRC EXCLUDE REGEX ".*/(ccl|runtime)/maca/.*")
+endif()
+if(NOT USE_DCU)
+    list(FILTER SRC EXCLUDE REGEX ".*/(ccl|runtime)/dcu/.*")
+elseif(NOT USE_RCCL)
+    list(FILTER SRC EXCLUDE REGEX ".*/ccl/dcu/.*")
 endif()
 
-# CPU kernels (*.cc)
-file(GLOB_RECURSE CPU_KERNELS ${PROJECT_SOURCE_DIR}/infini_train/src/kernels/cpu/*.cc)
-
-# ------------------------------------------------------------------------------
-# CPU kernels library
-# ------------------------------------------------------------------------------
+if(PROFILE_MODE)
+    add_compile_definitions(PROFILE_MODE=1)
+endif()
 
+file (GLOB_RECURSE CPU_KERNELS ${PROJECT_SOURCE_DIR}/infini_train/src/kernels/cpu/*.cc)
 add_library(infini_train_cpu_kernels STATIC ${CPU_KERNELS})
-target_link_libraries(infini_train_cpu_kernels PUBLIC glog Eigen3::Eigen)
-
+target_link_libraries(infini_train_cpu_kernels glog Eigen3::Eigen)
 if(USE_OMP)
-  add_compile_definitions(USE_OMP=1)
-  target_link_libraries(infini_train_cpu_kernels PUBLIC OpenMP::OpenMP_CXX)
+    add_compile_definitions(USE_OMP=1)
+    target_link_libraries(infini_train_cpu_kernels ${INFINI_OMP_LIBS})
 endif()
 
-# ------------------------------------------------------------------------------
-# CUDA kernels library (optional)
-# ------------------------------------------------------------------------------
+# =========================
+# MPI (optional)
+# =========================
+if (USE_MPI)
+    add_compile_definitions(USE_MPI=1)
+    if(USE_MACA AND DEFINED ENV{MACA_PATH} AND EXISTS "$ENV{MACA_PATH}/ompi")
+        set(OPENMPI_ROOT $ENV{MACA_PATH}/ompi CACHE PATH "OpenMPI root directory")
+    else()
+        set(OPENMPI_ROOT /opt/openmpi-4.1.6 CACHE PATH "OpenMPI root directory")
+    endif()
+
+    # ---- MPI include & lib (explicit OpenMPI path) ----
+    set(MPI_INCLUDE_DIR ${OPENMPI_ROOT}/include)
+    set(MPI_LIB_DIR     ${OPENMPI_ROOT}/lib)
+
+    include_directories(${MPI_INCLUDE_DIR})
+    link_directories(${MPI_LIB_DIR})
+
+    # OpenMPI core libs (C++ bindings are deprecated; MPI is C ABI)
+    set(MPI_LIBS mpi)
+
+    # mxcc 不支持 -pthread，用 Threads::Threads（-lpthread）
+    if (USE_MACA)
+        set(THREADS_PREFER_PTHREAD_FLAG OFF)
+        find_package(Threads REQUIRED)
+    endif()
+endif()
 
+# =========================
+# CUDA backend
+# =========================
 if(USE_CUDA)
   add_compile_definitions(USE_CUDA=1)
   enable_language(CUDA)
@@ -94,43 +141,141 @@ if(USE_CUDA)
       CUDA::cuda_driver
   )
 
+  add_library(infini_train STATIC ${SRC})
+  target_link_libraries(infini_train glog gflags infini_train_cpu_kernels infini_train_cuda_kernels)
+
   if(USE_NCCL)
     message(STATUS "Add USE_NCCL, use NCCL with CUDA")
     list(APPEND CMAKE_MODULE_PATH ${PROJECT_SOURCE_DIR}/cmake)
     find_package(NCCL REQUIRED)
     add_compile_definitions(USE_NCCL=1)
     target_link_libraries(infini_train_cuda_kernels PUBLIC nccl)
   endif()
-endif()
 
-# ------------------------------------------------------------------------------
-# Main framework library
-# ------------------------------------------------------------------------------
-
-add_library(infini_train STATIC ${SRC})
-target_link_libraries(infini_train
-  PUBLIC
-    glog
-    gflags
-    infini_train_cpu_kernels
-)
-
-if(USE_CUDA)
-  # infini_train contains cuda runtime wrappers (*.cc) like cuda_blas_handle.cc/cuda_guard.cc
-  # Those may need CUDA runtime/driver/cublas symbols at final link, so attach them here too.
-  target_link_libraries(infini_train
-    PUBLIC
-      infini_train_cuda_kernels
-      CUDA::cudart
-      CUDA::cublas
-      CUDA::cuda_driver
-  )
-
-  if(USE_NCCL)
-    # If your core library code also directly references NCCL symbols (not only kernels),
-    # keep this. Otherwise it's harmless.
-    target_link_libraries(infini_train PUBLIC nccl)
+  if (USE_MPI)
+    target_link_libraries(infini_train ${MPI_LIBS})
   endif()
+
+# =========================
+# DCU backend (Hygon DTK/HIP)
+# =========================
+elseif(USE_DCU)
+    add_compile_definitions(USE_DCU=1)
+
+    set(DCU_PATH "$ENV{DTK_PATH}" CACHE PATH "Hygon DTK installation root")
+    set(DCU_ARCH "" CACHE STRING "Optional HIP offload architecture reported by rocminfo")
+    if(NOT DCU_PATH)
+        set(DCU_PATH /opt/dtk)
+    endif()
+
+    find_program(HIPCC_EXECUTABLE hipcc
+        HINTS "${DCU_PATH}/bin" "${DCU_PATH}/llvm/bin" /opt/rocm/bin
+        REQUIRED)
+    if(NOT CMAKE_CXX_COMPILER MATCHES "hipcc")
+        message(WARNING
+            "DCU kernels must be compiled by hipcc. Reconfigure with "
+            "-DCMAKE_CXX_COMPILER=${HIPCC_EXECUTABLE}")
+    endif()
+
+    include_directories("${DCU_PATH}/include")
+    link_directories("${DCU_PATH}/lib" "${DCU_PATH}/lib64")
+
+    find_library(DCU_RUNTIME_LIB NAMES amdhip64 hip_hcc
+        HINTS "${DCU_PATH}/lib" "${DCU_PATH}/lib64" /opt/rocm/lib /opt/rocm/lib64
+        REQUIRED)
+    find_library(DCU_BLAS_LIB NAMES hipblas
+        HINTS "${DCU_PATH}/lib" "${DCU_PATH}/lib64" /opt/rocm/lib /opt/rocm/lib64
+        REQUIRED)
+
+    file(GLOB_RECURSE DCU_KERNELS ${PROJECT_SOURCE_DIR}/infini_train/src/kernels/dcu/*.hip)
+    set_source_files_properties(${DCU_KERNELS} PROPERTIES
+        LANGUAGE CXX
+        COMPILE_OPTIONS "-x;hip"
+    )
+    add_library(infini_train_dcu_kernels STATIC ${DCU_KERNELS})
+    if(DCU_ARCH)
+        target_compile_options(infini_train_dcu_kernels PRIVATE "--offload-arch=${DCU_ARCH}")
+    endif()
+    target_link_libraries(infini_train_dcu_kernels glog ${DCU_RUNTIME_LIB} ${DCU_BLAS_LIB})
+
+    add_library(infini_train STATIC ${SRC})
+    target_link_libraries(infini_train glog gflags infini_train_cpu_kernels infini_train_dcu_kernels)
+
+    if(USE_RCCL)
+        message(STATUS "Add USE_RCCL under DCU backend")
+        find_library(DCU_COMM_LIB NAMES rccl nccl
+            HINTS
+                "${DCU_PATH}/lib"
+                "${DCU_PATH}/lib64"
+                "${DCU_PATH}/rccl/lib"
+                "${DCU_PATH}/cuda/cuda-12/targets/x86_64-linux/lib"
+                "${DCU_PATH}/cuda/targets/x86_64-linux/lib"
+                /opt/rocm/lib
+                /opt/rocm/lib64
+            REQUIRED)
+        target_compile_definitions(infini_train PRIVATE USE_RCCL=1)
+        target_link_libraries(infini_train ${DCU_COMM_LIB})
+    endif()
+
+    if(USE_MPI)
+        target_link_libraries(infini_train ${MPI_LIBS})
+    endif()
+
+# =========================
+# MACA backend (MetaX)
+# =========================
+elseif(USE_MACA)
+    add_compile_definitions(USE_MACA=1)
+
+    # ---- configure MACA SDK paths ----
+    # Typical: /opt/maca (can be overridden by -DMACA_PATH=...)
+    set(MACA_PATH $ENV{MACA_PATH})
+    set(CMAKE_C_COMPILER ${MACA_PATH}/mxgpu_llvm/bin/mxcc)
+    set(CMAKE_CXX_COMPILER ${MACA_PATH}/mxgpu_llvm/bin/mxcc)
+
+    include_directories("${MACA_PATH}/include")
+    link_directories("${MACA_PATH}/lib")
+
+    # Libraries: mcruntime / mcdnn / mcblas
+    find_library(MACA_RUNTIME_LIB NAMES mcruntime HINTS "${MACA_PATH}/lib" REQUIRED)
+    find_library(MACA_DNN_LIB     NAMES mcdnn     HINTS "${MACA_PATH}/lib" REQUIRED)
+    find_library(MACA_BLAS_LIB    NAMES mcblas    HINTS "${MACA_PATH}/lib" REQUIRED)
+
+    file(GLOB_RECURSE MACA_KERNELS ${PROJECT_SOURCE_DIR}/infini_train/src/kernels/maca/*.maca)
+    set_source_files_properties(${MACA_KERNELS} PROPERTIES
+        LANGUAGE CXX
+        COMPILE_OPTIONS "-x;maca"
+    )
+    add_library(infini_train_maca_kernels STATIC ${MACA_KERNELS})
+    target_link_libraries(infini_train_maca_kernels glog ${MACA_RUNTIME_LIB} ${MACA_DNN_LIB} ${MACA_BLAS_LIB})
+
+    add_library(infini_train STATIC ${SRC})
+    target_link_libraries(infini_train glog gflags infini_train_cpu_kernels infini_train_maca_kernels)
+
+    if (USE_MCCL)
+        message(STATUS "Add USE_MCCL under MACA backend, use MCCL (mccl)")
+        find_library(MACA_COMM_LIB NAMES mccl HINTS "${MACA_PATH}/lib" REQUIRED)
+        add_compile_definitions(USE_MCCL=1)
+        target_link_libraries(infini_train ${MACA_COMM_LIB})
+    endif()
+
+    if (USE_MPI)
+        target_link_libraries(infini_train ${MPI_LIBS} Threads::Threads)
+
+        # 有些 MPI 还需要额外 link flags（比如 -Wl,...），也一并带上
+        if (MPI_CXX_LINK_FLAGS)
+            set_target_properties(infini_train PROPERTIES
+                LINK_FLAGS "${MPI_CXX_LINK_FLAGS}"
+            )
+        endif()
+    endif()
+
+# =========================
+# CPU-only backend
+# =========================
+else()
+    add_library(infini_train STATIC ${SRC})
+    target_link_libraries(infini_train glog gflags infini_train_cpu_kernels)
 endif()
 
 # ------------------------------------------------------------------------------
@@ -148,6 +293,26 @@ function(link_infini_train_exe target_name)
       "-Wl,--no-whole-archive"
       "-Wl,--end-group"
     )
+  elseif(USE_MACA)
+    target_link_libraries(${target_name} PRIVATE
+      "-Wl,--start-group"
+      "-Wl,--whole-archive"
+        infini_train
+        infini_train_cpu_kernels
+        infini_train_maca_kernels
+      "-Wl,--no-whole-archive"
+      "-Wl,--end-group"
+    )
+  elseif(USE_DCU)
+    target_link_libraries(${target_name} PRIVATE
+      "-Wl,--start-group"
+      "-Wl,--whole-archive"
+        infini_train
+        infini_train_cpu_kernels
+        infini_train_dcu_kernels
+      "-Wl,--no-whole-archive"
+      "-Wl,--end-group"
+    )
   else()
     target_link_libraries(${target_name} PRIVATE
       "-Wl,--start-group"
@@ -160,7 +325,6 @@ function(link_infini_train_exe target_name)
   endif()
 endfunction()
 
-
 # ------------------------------------------------------------------------------
 # Examples
 # ------------------------------------------------------------------------------

README.md

@@ -35,8 +35,11 @@ A from-scratch C++ training framework for large-scale models with multi-dimensio
 mkdir build
 cd build
 cmake .. -DUSE_CUDA=ON -DUSE_NCCL=ON
-make -j
-```
+make -j
+```
+
+For Hygon BW1000 / DCU builds and validation, see
+[`docs/dcu_bw1000.md`](docs/dcu_bw1000.md).
 
 Build Options:
 
@@ -171,4 +174,4 @@ Multiple parallelism strategies (DDP, TP, SP, PP) can be freely combined to scal
   Added Autocast, multi-dimensional distributed parallelism
    (DDP, TP, SP, PP with GPipe / 1F1B / vPP),
    multi-node training, `no_grad` mode,
-   and communication–computation overlap with bucketed gradient synchronization.
+   and communication–computation overlap with bucketed gradient synchronization.

example/gpt2/checkpoint_loader.cc

@@ -3,6 +3,7 @@
 #include <cmath>
 #include <cstdlib>
 #include <filesystem>
+#include <format>
 #include <fstream>
 #include <memory>
 #include <random>