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

CMakeLists.txt

@@ -1,9 +1,11 @@
 cmake_minimum_required(VERSION 3.28)
 
 option(USE_CUDA "Support NVIDIA CUDA" 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_RCCL "Build project for distributed running on DCU using RCCL" OFF)
 option(BUILD_TEST "Build InfiniTrain tests" OFF)
 
 project(infini_train VERSION 0.6.0 LANGUAGES CXX)
@@ -71,6 +73,11 @@ endif()
 if(NOT USE_NCCL)
   list(FILTER SRC EXCLUDE REGEX ".*infini_train/src/core/ccl/cuda/.*")
 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)
@@ -123,6 +130,50 @@ if(USE_CUDA)
   endif()
 endif()
 
+# ------------------------------------------------------------------------------
+# DCU kernels library (optional)
+# ------------------------------------------------------------------------------
+
+if(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 PUBLIC glog ${DCU_RUNTIME_LIB} ${DCU_BLAS_LIB})
+endif()
+
 # ------------------------------------------------------------------------------
 # Main framework library
 # ------------------------------------------------------------------------------
@@ -153,6 +204,24 @@ if(USE_CUDA)
   endif()
 endif()
 
+if(USE_DCU)
+  target_link_libraries(infini_train PUBLIC infini_train_dcu_kernels ${DCU_RUNTIME_LIB} ${DCU_BLAS_LIB})
+
+  if(USE_RCCL)
+    message(STATUS "Add USE_RCCL, use RCCL with DCU")
+    find_library(DCU_COMM_LIB NAMES rccl nccl
+      HINTS
+        "${DCU_PATH}/lib"
+        "${DCU_PATH}/lib64"
+        "${DCU_PATH}/rccl/lib"
+        /opt/rocm/lib
+        /opt/rocm/lib64
+      REQUIRED)
+    add_compile_definitions(USE_RCCL=1)
+    target_link_libraries(infini_train PUBLIC ${DCU_COMM_LIB})
+  endif()
+endif()
+
 # ------------------------------------------------------------------------------
 # Helper: link libraries in a group to fix static lib one-pass resolution
 # (THIS is what fixes "undefined reference" from cuda_kernels -> core symbols)
@@ -168,6 +237,16 @@ function(link_infini_train_exe target_name)
       "-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"

example/gpt2/checkpoint_loader.cc

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

example/gpt2/main.cc

@@ -75,7 +75,7 @@ DEFINE_uint32(sample_every, 0, "how often to sample from the model?");
 // debugging
 DEFINE_bool(overfit_single_batch, true, "overfit just one batch of data");
 // memory management
-DEFINE_string(device, "cuda", "device type (cpu/cuda), useless if using parallel training mode");
+DEFINE_string(device, "cuda", "device type (cpu/cuda/dcu), useless if using parallel training mode");
 // parallel
 DEFINE_int32(
     nthread_per_process, 1,
@@ -114,6 +114,7 @@ const std::unordered_set<std::string> kSupportedModels
     = {"gpt2", "gpt2-medium", "gpt2-large", "gpt2-xl", "d12", "d24", "d36", "d48"};
 constexpr char kDeviceCPU[] = "cpu";
 constexpr char kDeviceCUDA[] = "cuda";
+constexpr char kDeviceDCU[] = "dcu";
 constexpr char kDtypeFP32[] = "float32";
 constexpr char kDtypeBF16[] = "bfloat16";
 const std::unordered_set<std::string> kSupportedLRDecayStyles
@@ -130,8 +131,9 @@ const std::unordered_map<std::string, nn::TransformerConfig> kModelToConfigs = {
 } // namespace
 
 DEFINE_validator(model, [](const char *, const std::string &value) { return kSupportedModels.contains(value); });
-DEFINE_validator(device,
-                 [](const char *, const std::string &value) { return value == kDeviceCPU || value == kDeviceCUDA; });
+DEFINE_validator(device, [](const char *, const std::string &value) {
+    return value == kDeviceCPU || value == kDeviceCUDA || value == kDeviceDCU;
+});
 DEFINE_validator(zero_stage, [](const char *, int32_t value) { return value >= 0 && value <= 3; });
 DEFINE_validator(lr_decay_style,
                  [](const char *, const std::string &value) { return kSupportedLRDecayStyles.contains(value); });
@@ -181,7 +183,8 @@ void Train(const nn::parallel::Rank &rank) {
     const ProcessGroup *pp_pg = nullptr;
 
     if (rank.IsParallel()) {
-        device = Device(Device::DeviceType::kCUDA, rank.thread_rank());
+        auto parallel_device_type = FLAGS_device == kDeviceDCU ? Device::DeviceType::kDCU : Device::DeviceType::kCUDA;
+        device = Device(parallel_device_type, rank.thread_rank());
         auto *pg_factory = ProcessGroupFactory::Instance(device.type());
 
         if (ddp_world_size > 1) {
@@ -205,8 +208,12 @@ void Train(const nn::parallel::Rank &rank) {
 
             nn::parallel::pp_rank = pp_rank;
         }
+    } else if (FLAGS_device == kDeviceCPU) {
+        device = Device();
+    } else if (FLAGS_device == kDeviceDCU) {
+        device = Device(Device::DeviceType::kDCU, 0);
     } else {
-        device = FLAGS_device == kDeviceCPU ? Device() : Device(Device::DeviceType::kCUDA, 0);
+        device = Device(Device::DeviceType::kCUDA, 0);
     }
 
     // calculate gradient accumulation from the desired total batch size and the current run configuration

example/llama3/checkpoint_loader.cc

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

example/llama3/main.cc

@@ -74,7 +74,7 @@ DEFINE_uint32(sample_every, 0, "how often to sample from the model?");
 // debugging
 DEFINE_bool(overfit_single_batch, true, "overfit just one batch of data");
 // memory management
-DEFINE_string(device, "cuda", "device type (cpu/cuda), useless if using parallel training mode");
+DEFINE_string(device, "cuda", "device type (cpu/cuda/dcu), useless if using parallel training mode");
 // parallel
 DEFINE_int32(
     nthread_per_process, 1,
@@ -110,15 +110,17 @@ namespace {
 const std::unordered_set<std::string> kSupportedModels = {"llama3"};
 constexpr char kDeviceCPU[] = "cpu";
 constexpr char kDeviceCUDA[] = "cuda";
+constexpr char kDeviceDCU[] = "dcu";
 constexpr char kDtypeFP32[] = "float32";
 constexpr char kDtypeBF16[] = "bfloat16";
 const std::unordered_set<std::string> kSupportedLRDecayStyles
     = {"none", "constant", "linear", "cosine", "inverse-square-root"};
 } // namespace
 
 DEFINE_validator(model, [](const char *, const std::string &value) { return kSupportedModels.contains(value); });
-DEFINE_validator(device,
-                 [](const char *, const std::string &value) { return value == kDeviceCPU || value == kDeviceCUDA; });
+DEFINE_validator(device, [](const char *, const std::string &value) {
+    return value == kDeviceCPU || value == kDeviceCUDA || value == kDeviceDCU;
+});
 DEFINE_validator(zero_stage, [](const char *, int32_t value) { return value >= 0 && value <= 3; });
 DEFINE_validator(lr_decay_style,
                  [](const char *, const std::string &value) { return kSupportedLRDecayStyles.contains(value); });
@@ -167,7 +169,8 @@ void Train(const nn::parallel::Rank &rank) {
     const ProcessGroup *pp_pg = nullptr;
 
     if (rank.IsParallel()) {
-        device = Device(Device::DeviceType::kCUDA, rank.thread_rank());
+        auto parallel_device_type = FLAGS_device == kDeviceDCU ? Device::DeviceType::kDCU : Device::DeviceType::kCUDA;
+        device = Device(parallel_device_type, rank.thread_rank());
         auto *pg_factory = ProcessGroupFactory::Instance(device.type());
 
         if (ddp_world_size > 1) {
@@ -191,8 +194,12 @@ void Train(const nn::parallel::Rank &rank) {
 
             nn::parallel::pp_rank = pp_rank;
         }
+    } else if (FLAGS_device == kDeviceCPU) {
+        device = Device();
+    } else if (FLAGS_device == kDeviceDCU) {
+        device = Device(Device::DeviceType::kDCU, 0);
     } else {
-        device = FLAGS_device == kDeviceCPU ? Device() : Device(Device::DeviceType::kCUDA, 0);
+        device = Device(Device::DeviceType::kCUDA, 0);
     }
 
     // calculate gradient accumulation from the desired total batch size and the current run configuration

format

@@ -0,0 +1,147 @@
+#pragma once
+
+#include <iomanip>
+#include <sstream>
+#include <string>
+#include <string_view>
+#include <type_traits>
+#include <utility>
+#include <vector>
+
+namespace std {
+namespace infini_train_format_compat {
+
+struct FormatSpec {
+    int width = 0;
+    int precision = -1;
+    bool left = false;
+    char type = '\0';
+};
+
+inline FormatSpec ParseSpec(std::string_view spec) {
+    FormatSpec parsed;
+    if (!spec.empty() && spec.front() == ':') {
+        spec.remove_prefix(1);
+    }
+    if (!spec.empty() && spec.front() == '<') {
+        parsed.left = true;
+        spec.remove_prefix(1);
+    } else if (!spec.empty() && spec.front() == '>') {
+        spec.remove_prefix(1);
+    }
+
+    while (!spec.empty() && spec.front() >= '0' && spec.front() <= '9') {
+        parsed.width = parsed.width * 10 + (spec.front() - '0');
+        spec.remove_prefix(1);
+    }
+
+    if (!spec.empty() && spec.front() == '.') {
+        spec.remove_prefix(1);
+        parsed.precision = 0;
+        while (!spec.empty() && spec.front() >= '0' && spec.front() <= '9') {
+            parsed.precision = parsed.precision * 10 + (spec.front() - '0');
+            spec.remove_prefix(1);
+        }
+    }
+
+    if (!spec.empty()) {
+        parsed.type = spec.front();
+    }
+    return parsed;
+}
+
+template <typename T> std::string FormatOne(std::string_view spec_text, T &&value) {
+    const auto spec = ParseSpec(spec_text);
+    std::ostringstream oss;
+    if (spec.left) {
+        oss << std::left;
+    }
+    if (spec.width > 0) {
+        oss << std::setw(spec.width);
+    }
+    if (spec.precision >= 0) {
+        oss << std::setprecision(spec.precision);
+    }
+    if (spec.type == 'f') {
+        oss << std::fixed;
+    } else if (spec.type == 'e') {
+        oss << std::scientific;
+    }
+    oss << std::forward<T>(value);
+    return oss.str();
+}
+
+inline void CollectFormatArgs(std::vector<std::string> &) {}
+
+template <typename T, typename... Rest>
+void CollectFormatArgs(std::vector<std::string> &out, std::string_view spec, T &&value, Rest &&...rest) {
+    out.push_back(FormatOne(spec, std::forward<T>(value)));
+    if constexpr (sizeof...(Rest) > 0) {
+        CollectFormatArgs(out, "", std::forward<Rest>(rest)...);
+    }
+}
+
+template <typename T> void PushFormattedArg(std::vector<std::string> &out, std::string_view spec, T &&value) {
+    out.push_back(FormatOne(spec, std::forward<T>(value)));
+}
+
+template <typename... Args> std::vector<std::string> MakeFormatArgs(std::string_view fmt, Args &&...args) {
+    std::vector<std::string> result;
+    result.reserve(sizeof...(Args));
+
+    size_t arg_index = 0;
+    auto specs = std::vector<std::string_view>{};
+    for (size_t i = 0; i < fmt.size(); ++i) {
+        if (fmt[i] != '{' || (i + 1 < fmt.size() && fmt[i + 1] == '{')) {
+            if (fmt[i] == '{') {
+                ++i;
+            }
+            continue;
+        }
+        const size_t close = fmt.find('}', i + 1);
+        if (close == std::string_view::npos) {
+            break;
+        }
+        specs.push_back(fmt.substr(i + 1, close - i - 1));
+        i = close;
+        ++arg_index;
+    }
+
+    size_t spec_index = 0;
+    (PushFormattedArg(result, spec_index < specs.size() ? specs[spec_index++] : std::string_view{}, std::forward<Args>(args)), ...);
+    return result;
+}
+
+} // namespace infini_train_format_compat
+
+template <typename... Args> std::string format(std::string_view fmt, Args &&...args) {
+    const auto formatted_args = infini_train_format_compat::MakeFormatArgs(fmt, std::forward<Args>(args)...);
+    std::ostringstream out;
+    size_t arg_index = 0;
+
+    for (size_t i = 0; i < fmt.size(); ++i) {
+        if (fmt[i] == '{') {
+            if (i + 1 < fmt.size() && fmt[i + 1] == '{') {
+                out << '{';
+                ++i;
+                continue;
+            }
+            const size_t close = fmt.find('}', i + 1);
+            if (close != std::string_view::npos) {
+                if (arg_index < formatted_args.size()) {
+                    out << formatted_args[arg_index++];
+                }
+                i = close;
+                continue;
+            }
+        } else if (fmt[i] == '}' && i + 1 < fmt.size() && fmt[i + 1] == '}') {
+            out << '}';
+            ++i;
+            continue;
+        }
+        out << fmt[i];
+    }
+    return out.str();
+}
+
+} // namespace std

infini_train/include/autograd/linear.h

@@ -12,6 +12,12 @@ class Tensor;
 
 namespace infini_train::autograd {
 
+struct LinearGradFlags {
+    bool input = false;
+    bool weight = false;
+    bool bias = false;
+};
+
 class Linear : public Function {
 public:
     static constexpr char kType[] = "LinearFunction";