feat: add PrivateUse1 backend extension support

CMakeLists.txt

@@ -3,7 +3,7 @@ cmake_minimum_required(VERSION 3.28)
 option(USE_CUDA "Support NVIDIA CUDA" 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(BUILD_TEST "Build InfiniTrain tests" OFF)
 
 project(infini_train VERSION 0.6.0 LANGUAGES CXX)
@@ -41,7 +41,15 @@ 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_subdirectory(third_party/glog)
+get_property(_infinitrain_build_testing_was_cached CACHE BUILD_TESTING PROPERTY TYPE SET)
+block()
+  set(BUILD_TESTING OFF)
+  add_subdirectory(third_party/glog)
+endblock()
+if(NOT _infinitrain_build_testing_was_cached)
+  unset(BUILD_TESTING CACHE)
+endif()
+unset(_infinitrain_build_testing_was_cached)
 include_directories(${glog_SOURCE_DIR}/src)
 
 # eigen
@@ -64,12 +72,15 @@ endif()
 # 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/.*")
+
+# Exclude backend-specific runtime/ccl translation units when the corresponding
+# backend is disabled. This keeps each build self-contained and avoids pulling
+# in headers (e.g. <cuda_runtime.h> / <mcr/mc_runtime.h>) that are not on the
+# include path.
 if(NOT USE_CUDA)
-  list(FILTER SRC EXCLUDE REGEX ".*runtime/cuda/.*")
-  list(FILTER SRC EXCLUDE REGEX ".*ccl/cuda/.*")
-endif()
-if(NOT USE_NCCL)
-  list(FILTER SRC EXCLUDE REGEX ".*infini_train/src/core/ccl/cuda/.*")
+  list(FILTER SRC EXCLUDE REGEX ".*/(ccl|runtime)/cuda/.*")
+elseif(NOT USE_NCCL)
+  list(FILTER SRC EXCLUDE REGEX ".*/ccl/cuda/.*")
 endif()
 
 # CPU kernels (*.cc)
@@ -128,6 +139,11 @@ endif()
 # ------------------------------------------------------------------------------
 
 add_library(infini_train STATIC ${SRC})
+add_library(InfiniTrain::infini_train ALIAS infini_train)
+add_library(InfiniTrain::cpu_kernels ALIAS infini_train_cpu_kernels)
+target_include_directories(infini_train PUBLIC
+  "$<BUILD_INTERFACE:${CMAKE_CURRENT_SOURCE_DIR}>"
+)
 target_link_libraries(infini_train
   PUBLIC
     glog
@@ -185,6 +201,7 @@ endfunction()
 # Examples
 # ------------------------------------------------------------------------------
 
+if(PROJECT_IS_TOP_LEVEL)
 add_executable(mnist
   example/mnist/main.cc
   example/mnist/dataset.cc
@@ -217,10 +234,13 @@ add_executable(llama3
   example/llama3/checkpoint_loader.cc
 )
 link_infini_train_exe(llama3)
+endif()
 
 # Tools
+if(PROJECT_IS_TOP_LEVEL)
 add_subdirectory(tools/infini_run)
 set_target_properties(infini_run PROPERTIES RUNTIME_OUTPUT_DIRECTORY ${CMAKE_BINARY_DIR})
+endif()
 
 # Tests
 if(BUILD_TEST)

docs/test_infrastructure_design.md

@@ -183,79 +183,20 @@ add_subdirectory(foo)
 | `ONLY_CUDA()` | 只在 CUDA 实例运行 |
 | `REQUIRE_MIN_DEVICES(n)` | 加速器设备不足时 skip |
 
-## 4. 扩展新设备平台（以沐曦 MACA 为例）
+## 4. 扩展新设备平台
 
-当前测试体系围绕 CPU / CUDA 两种设备参数化。如果需要支持新平台（以沐曦 MACA 为例），需要改动以下几处：
+第三方设备统一占用 `DeviceType::kPrivateUse1`，不再向框架枚举和根 CMake
+增加厂商类型或 SDK 选项。厂商仓库负责注册 runtime、kernel 和可选 CCL，
+并将 InfiniTrain 固定为同一链接图中的 submodule。
 
-### 4.1 框架层：注册新设备类型
+框架中的 `tests/backend/test_privateuse1_backend.cc` 使用不依赖硬件的 fake
+backend 验证扩展契约。厂商硬件测试应放在厂商目录中，单独链接并显式初始化
+provider；fake backend 与真实 backend 不能放进同一测试进程，因为一个进程
+只能注册一个 `kPrivateUse1` provider。
 
-在 `infini_train/include/device.h` 的 `DeviceType` 枚举中新增：
+厂商测试建议至少覆盖：
 
-```cpp
-enum class DeviceType : int8_t {
-    kCPU = 0,
-    kCUDA = 1,
-    kMACA = 2,  // 新增
-};
-```
-
-### 4.2 测试工具层：`test_utils.h`
-
-1. 新增 MACA 头文件的编译期引入（和 CUDA 对称）：
-
-```cpp
-#if defined(USE_MACA)
-#include <maca_runtime_api.h>
-#endif
-```
-
-2. 新增 `ONLY_MACA()` 宏：
-
-```cpp
-#define ONLY_MACA() \
-    do { if (GetParam() != infini_train::Device::DeviceType::kMACA) { GTEST_SKIP() << "MACA-only test"; } } while (0)
-```
-
-如果希望有类似 `REQUIRE_MIN_DEVICES(n)` 但针对 MACA 的语义，可以按 `USE_CUDA` 分支的写法增加一个新的宏；同理 `USE_MACA` 不开时该宏直接 skip 即可。
-
-### 4.3 注册宏：新增 MACA 实例
-
-沿用 `USE_CUDA` 的做法，未开启编译开关时不注册对应实例：
-
-```cpp
-#if defined(USE_CUDA) && defined(USE_MACA)
-#define INFINI_TRAIN_REGISTER_TEST(TestName)                                    \
-    INSTANTIATE_TEST_SUITE_P(CPU, TestName,                                     \
-        ::testing::Values(infini_train::Device::DeviceType::kCPU));             \
-    INSTANTIATE_TEST_SUITE_P(CUDA, TestName,                                    \
-        ::testing::Values(infini_train::Device::DeviceType::kCUDA));            \
-    INSTANTIATE_TEST_SUITE_P(MACA, TestName,                                    \
-        ::testing::Values(infini_train::Device::DeviceType::kMACA))
-#elif defined(USE_CUDA)
-#define INFINI_TRAIN_REGISTER_TEST(TestName) /* CPU + CUDA, 同现状 */
-#elif defined(USE_MACA)
-#define INFINI_TRAIN_REGISTER_TEST(TestName)                                    \
-    INSTANTIATE_TEST_SUITE_P(CPU, TestName,                                     \
-        ::testing::Values(infini_train::Device::DeviceType::kCPU));             \
-    INSTANTIATE_TEST_SUITE_P(MACA, TestName,                                    \
-        ::testing::Values(infini_train::Device::DeviceType::kMACA))
-#else
-#define INFINI_TRAIN_REGISTER_TEST(TestName) /* 仅 CPU */
-#endif
-```
-
-运行时如果机器上没有对应设备（例如 `USE_MACA` 编译但无 MACA 硬件），让测试直接报错而不是静默跳过。
-
-### 4.4 CMake 层：`test_macros.cmake`
-
-将默认 label 列表从 `cpu cuda` 扩展为 `cpu cuda maca`
-
-### 4.5 检查清单
-
-| 步骤 | 文件 | 改动 |
-|------|------|------|
-| 1 | `device.h` | `DeviceType` 枚举新增 `kMACA` |
-| 2 | `test_utils.h` | 新增 `USE_MACA` 下的 `<maca_runtime_api.h>` 引入、`ONLY_MACA()` 宏 |
-| 3 | `test_utils.h` | `INFINI_TRAIN_REGISTER_TEST` 按 `USE_MACA` 条件新增 MACA 实例 |
-| 4 | `test_macros.cmake` | 将默认 label 列表扩展为 `cpu cuda maca` |
-| 5 | `CMakeLists.txt`（根） | 新增 `USE_MACA` option + MACA SDK 查找 + kernel 编译 |
+1. `DeviceGuardImpl` 的设备、stream、event、allocator 和 copy 行为；
+2. `Cast`、`Fill`、`NoOpForward`、`NoOpBackward` 等基础 kernel；
+3. CCL 初始化和 collective 行为；
+4. provider 名称解析和重复初始化。

example/gpt2/main.cc

@@ -4,12 +4,19 @@
 #include <format>
 #include <memory>
 #include <optional>
+#include <thread>
 #include <unordered_map>
 #include <unordered_set>
 
 #include "gflags/gflags.h"
 #include "glog/logging.h"
 
+// Out-of-tree builds inject the selected provider's declaration without adding
+// a vendor dependency to the upstream example.
+#ifdef INFINITRAIN_EXAMPLE_EXTERNAL_BACKEND_HEADER
+#include INFINITRAIN_EXAMPLE_EXTERNAL_BACKEND_HEADER
+#endif
+
 #include "infini_train/include/autocast.h"
 #include "infini_train/include/checkpoint/checkpoint.h"
 #include "infini_train/include/core/runtime/device_guard.h"
@@ -75,12 +82,11 @@ 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, useless if using parallel training mode");
 // parallel
-DEFINE_int32(
-    nthread_per_process, 1,
-    "Number of threads to use for each process. "
-    "When set > 1, enables data parallelism with device=cuda on the specified number of visible CUDA devices.");
+DEFINE_int32(nthread_per_process, 1,
+             "Number of threads to use for each process. "
+             "When set > 1, enables data parallelism on the specified accelerator devices.");
 DEFINE_uint32(tensor_parallel, 1, "Tensor Parallel world size");
 DEFINE_bool(sequence_parallel, false, "Whether to enable Sequence Parallel");
 DEFINE_uint32(pipeline_parallel, 1, "Pipeline Parallel world size, specified the number of PP stages.");
@@ -112,8 +118,6 @@ namespace {
 // validation
 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 kDtypeFP32[] = "float32";
 constexpr char kDtypeBF16[] = "bfloat16";
 const std::unordered_set<std::string> kSupportedLRDecayStyles
@@ -130,8 +134,7 @@ 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 Device::ParseType(value).has_value(); });
 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); });
@@ -157,6 +160,7 @@ void Train(const nn::parallel::Rank &rank) {
 
     // select the device
     Device device;
+    const auto device_type = Device::ParseType(FLAGS_device).value();
 
     int ddp_world_size = global::GetDataParallelSize();
     int tp_world_size = global::GetTensorParallelSize();
@@ -181,7 +185,9 @@ void Train(const nn::parallel::Rank &rank) {
     const ProcessGroup *pp_pg = nullptr;
 
     if (rank.IsParallel()) {
-        device = Device(Device::DeviceType::kCUDA, rank.thread_rank());
+        CHECK(device_type != Device::DeviceType::kCPU) << "Parallel training requires an accelerator backend";
+        device = Device(device_type, rank.thread_rank());
+
         auto *pg_factory = ProcessGroupFactory::Instance(device.type());
 
         if (ddp_world_size > 1) {
@@ -206,7 +212,7 @@ void Train(const nn::parallel::Rank &rank) {
             nn::parallel::pp_rank = pp_rank;
         }
     } else {
-        device = FLAGS_device == kDeviceCPU ? Device() : Device(Device::DeviceType::kCUDA, 0);
+        device = Device(device_type, 0);
     }
 
     // calculate gradient accumulation from the desired total batch size and the current run configuration
@@ -558,8 +564,13 @@ void Train(const nn::parallel::Rank &rank) {
 }
 
 int main(int argc, char *argv[]) {
-    gflags::ParseCommandLineFlags(&argc, &argv, true);
     google::InitGoogleLogging(argv[0]);
+    // Register provider metadata and implementations before gflags validates
+    // --device. The device runtime initializes lazily on first DeviceGuard use.
+#ifdef INFINITRAIN_EXAMPLE_EXTERNAL_BACKEND_REGISTRAR
+    INFINITRAIN_EXAMPLE_EXTERNAL_BACKEND_REGISTRAR();
+#endif
+    gflags::ParseCommandLineFlags(&argc, &argv, true);
 
     auto precision_config = utils::PrecisionCheckConfig::Parse(FLAGS_precision_check);
     nn::parallel::global::InitAllEnv(FLAGS_nthread_per_process, FLAGS_tensor_parallel, FLAGS_sequence_parallel,

example/llama3/main.cc

@@ -3,11 +3,18 @@
 #include <format>
 #include <memory>
 #include <optional>
+#include <thread>
 #include <unordered_set>
 
 #include "gflags/gflags.h"
 #include "glog/logging.h"
 
+// Out-of-tree builds inject the selected provider's declaration without adding
+// a vendor dependency to the upstream example.
+#ifdef INFINITRAIN_EXAMPLE_EXTERNAL_BACKEND_HEADER
+#include INFINITRAIN_EXAMPLE_EXTERNAL_BACKEND_HEADER
+#endif
+
 #include "infini_train/include/autocast.h"
 #include "infini_train/include/checkpoint/checkpoint.h"
 #include "infini_train/include/checkpoint/checkpoint_manager.h"
@@ -74,12 +81,11 @@ 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, useless if using parallel training mode");
 // parallel
-DEFINE_int32(
-    nthread_per_process, 1,
-    "Number of threads to use for each process. "
-    "When set > 1, enables data parallelism with device=cuda on the specified number of visible CUDA devices.");
+DEFINE_int32(nthread_per_process, 1,
+             "Number of threads to use for each process. "
+             "When set > 1, enables data parallelism on the specified accelerator devices.");
 DEFINE_uint32(tensor_parallel, 1, "Tensor Parallel world size");
 DEFINE_bool(sequence_parallel, false, "Whether to enable Sequence Parallel");
 DEFINE_uint32(pipeline_parallel, 1, "Pipeline Parallel world size, specified the number of PP stages.");
@@ -108,17 +114,14 @@ using namespace infini_train;
 namespace {
 // validation
 const std::unordered_set<std::string> kSupportedModels = {"llama3"};
-constexpr char kDeviceCPU[] = "cpu";
-constexpr char kDeviceCUDA[] = "cuda";
 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 Device::ParseType(value).has_value(); });
 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); });
@@ -144,6 +147,7 @@ void Train(const nn::parallel::Rank &rank) {
 
     // select the device
     Device device;
+    const auto device_type = Device::ParseType(FLAGS_device).value();
 
     int ddp_world_size = global::GetDataParallelSize();
     int tp_world_size = global::GetTensorParallelSize();
@@ -167,7 +171,9 @@ void Train(const nn::parallel::Rank &rank) {
     const ProcessGroup *pp_pg = nullptr;
 
     if (rank.IsParallel()) {
-        device = Device(Device::DeviceType::kCUDA, rank.thread_rank());
+        CHECK(device_type != Device::DeviceType::kCPU) << "Parallel training requires an accelerator backend";
+        device = Device(device_type, rank.thread_rank());
+
         auto *pg_factory = ProcessGroupFactory::Instance(device.type());
 
         if (ddp_world_size > 1) {
@@ -192,7 +198,7 @@ void Train(const nn::parallel::Rank &rank) {
             nn::parallel::pp_rank = pp_rank;
         }
     } else {
-        device = FLAGS_device == kDeviceCPU ? Device() : Device(Device::DeviceType::kCUDA, 0);
+        device = Device(device_type, 0);
     }
 
     // calculate gradient accumulation from the desired total batch size and the current run configuration
@@ -535,8 +541,13 @@ void Train(const nn::parallel::Rank &rank) {
 }
 
 int main(int argc, char *argv[]) {
-    gflags::ParseCommandLineFlags(&argc, &argv, true);
     google::InitGoogleLogging(argv[0]);
+    // Register provider metadata and implementations before gflags validates
+    // --device. The device runtime initializes lazily on first DeviceGuard use.
+#ifdef INFINITRAIN_EXAMPLE_EXTERNAL_BACKEND_REGISTRAR
+    INFINITRAIN_EXAMPLE_EXTERNAL_BACKEND_REGISTRAR();
+#endif
+    gflags::ParseCommandLineFlags(&argc, &argv, true);
 
     auto precision_config = utils::PrecisionCheckConfig::Parse(FLAGS_precision_check);
     nn::parallel::global::InitAllEnv(FLAGS_nthread_per_process, FLAGS_tensor_parallel, FLAGS_sequence_parallel,