From 2d8e754a822cc531af84144b49e31efdf1c3440b Mon Sep 17 00:00:00 2001 From: cx Date: Wed, 12 Aug 2026 07:49:16 +0000 Subject: [PATCH] feat: add PrivateUse1 backend extension support - add a provider-neutral PrivateUse1 device type and registration API - validate runtime, kernel, and optional CCL backend registrations - initialize external device runtimes lazily on first use - support provider names in device parsing and display - require explicit autocast dtype for PrivateUse1 devices - allow examples to register an external backend before flag parsing - honor average_in_collective consistently across DDP paths - expose embeddable CMake targets and add fake backend tests --- CMakeLists.txt | 34 +++- docs/test_infrastructure_design.md | 85 ++------ example/gpt2/main.cc | 35 ++-- example/llama3/main.cc | 35 ++-- example/mixtral/main.cc | 25 ++- infini_train/include/autocast.h | 14 +- infini_train/include/core/ccl/ccl.h | 2 + .../include/core/privateuse1_backend.h | 22 +++ .../include/core/runtime/device_guard.h | 6 + infini_train/include/device.h | 8 +- infini_train/include/dispatcher.h | 3 + infini_train/src/core/ccl/ccl.cc | 2 + infini_train/src/core/privateuse1_backend.cc | 97 ++++++++++ infini_train/src/core/runtime/device_guard.cc | 19 +- infini_train/src/device.cc | 32 ++- .../parallel/ddp/distributed_data_parallel.cc | 5 +- infini_train/src/nn/parallel/ddp/reducer.cc | 4 +- scripts/run_models_and_profile.bash | 7 +- tests/CMakeLists.txt | 5 +- tests/backend/CMakeLists.txt | 9 + tests/backend/test_privateuse1_backend.cc | 182 ++++++++++++++++++ .../test_privateuse1_backend_validation.cc | 31 +++ tests/common/CMakeLists.txt | 2 +- 23 files changed, 530 insertions(+), 134 deletions(-) create mode 100644 infini_train/include/core/privateuse1_backend.h create mode 100644 infini_train/src/core/privateuse1_backend.cc create mode 100644 tests/backend/CMakeLists.txt create mode 100644 tests/backend/test_privateuse1_backend.cc create mode 100644 tests/backend/test_privateuse1_backend_validation.cc diff --git a/CMakeLists.txt b/CMakeLists.txt index 6bd8069d..f0eac734 100644 --- a/CMakeLists.txt +++ b/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. / ) 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 + "$" +) 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) diff --git a/docs/test_infrastructure_design.md b/docs/test_infrastructure_design.md index 8aa210ce..a3ee1830 100644 --- a/docs/test_infrastructure_design.md +++ b/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 -#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` 下的 `` 引入、`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 名称解析和重复初始化。 diff --git a/example/gpt2/main.cc b/example/gpt2/main.cc index 60c0c908..36481a38 100644 --- a/example/gpt2/main.cc +++ b/example/gpt2/main.cc @@ -4,12 +4,19 @@ #include #include #include +#include #include #include #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 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 kSupportedLRDecayStyles @@ -130,8 +134,7 @@ const std::unordered_map 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, diff --git a/example/llama3/main.cc b/example/llama3/main.cc index 302e0808..d8aec557 100644 --- a/example/llama3/main.cc +++ b/example/llama3/main.cc @@ -3,11 +3,18 @@ #include #include #include +#include #include #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,8 +114,6 @@ using namespace infini_train; namespace { // validation const std::unordered_set kSupportedModels = {"llama3"}; -constexpr char kDeviceCPU[] = "cpu"; -constexpr char kDeviceCUDA[] = "cuda"; constexpr char kDtypeFP32[] = "float32"; constexpr char kDtypeBF16[] = "bfloat16"; const std::unordered_set kSupportedLRDecayStyles @@ -117,8 +121,7 @@ const std::unordered_set kSupportedLRDecayStyles } // 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, diff --git a/example/mixtral/main.cc b/example/mixtral/main.cc index 84853865..18f4d60e 100644 --- a/example/mixtral/main.cc +++ b/example/mixtral/main.cc @@ -9,6 +9,12 @@ #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/core/runtime/device_guard.h" #include "infini_train/include/dataloader.h" @@ -34,7 +40,7 @@ DEFINE_uint32(num_iteration, 10, "number of training iterations"); DEFINE_double(learning_rate, 1e-4, "Adam learning rate"); DEFINE_string(llmc_filepath, "", "optional PyTorch-generated tiny Mixtral LLMC model file path to load before training"); -DEFINE_string(device, "cpu", "Training device: cpu or cuda."); +DEFINE_string(device, "cpu", "Training device."); DEFINE_string(dtype, "float32", "precision used in training (float32/bfloat16)"); DEFINE_uint32(log_interval, 1, "Print train loss every N steps. 0 disables step loss logging."); DEFINE_bool(print_timing, false, "Print training-loop elapsed time and token throughput."); @@ -60,8 +66,13 @@ void ValidateRuntimeFlags(const infini_train::nn::TransformerConfig &config) { } // namespace 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); infini_train::nn::parallel::global::InitAllEnv( /*nthread_per_process=*/1, @@ -80,13 +91,11 @@ int main(int argc, char *argv[]) { } ValidateRuntimeFlags(model_config); - Device train_device; - if (FLAGS_device == "cuda") { - train_device = Device(Device::DeviceType::kCUDA, 0); + const auto device_type = Device::ParseType(FLAGS_device); + CHECK(device_type.has_value()) << "Unsupported training device: " << FLAGS_device; + Device train_device(*device_type, 0); + if (!train_device.IsCPU()) { model->To(train_device); - } else { - CHECK_EQ(FLAGS_device, "cpu") << "Unsupported training device: " << FLAGS_device; - train_device = Device(); } infini_train::DistributedDataLoader train_loader( diff --git a/infini_train/include/autocast.h b/infini_train/include/autocast.h index 0ce83e98..f6191800 100644 --- a/infini_train/include/autocast.h +++ b/infini_train/include/autocast.h @@ -70,8 +70,8 @@ inline const std::unordered_map kOpCastPolicyMap = {"Layernorm", CastPolicy::kFP32}, }; -// Default autocast data types for each device type -inline constexpr std::array(Device::DeviceType::kCount)> kDeviceDefaultDtype = { +// Default autocast data types for built-in device types +inline constexpr std::array kDeviceDefaultDtype = { DataType::kBFLOAT16, // CPU DataType::kFLOAT16, // CUDA. }; @@ -160,8 +160,16 @@ class AutocastGuard { tls_autocast_context.autocast_dtype = autocast_dtype; } + // PrivateUse1 is provider-defined, so the framework cannot choose its default autocast dtype. + // Callers must use the overload above and pass the dtype explicitly. AutocastGuard(Device::DeviceType device_type) - : AutocastGuard(device_type, kDeviceDefaultDtype[static_cast(device_type)]) {} + : AutocastGuard(device_type, [device_type] { + CHECK(device_type != Device::DeviceType::kPrivateUse1) + << "PrivateUse1 autocast requires an explicit dtype"; + const auto index = static_cast(device_type); + CHECK_LT(index, kDeviceDefaultDtype.size()) << "Invalid device type for autocast"; + return kDeviceDefaultDtype[index]; + }()) {} // Disable autocast (restore previous state) ~AutocastGuard() { tls_autocast_context = saved_context_; } diff --git a/infini_train/include/core/ccl/ccl.h b/infini_train/include/core/ccl/ccl.h index 626cb078..3fb471be 100644 --- a/infini_train/include/core/ccl/ccl.h +++ b/infini_train/include/core/ccl/ccl.h @@ -79,6 +79,8 @@ class CclImplRegistry { void Register(Device::DeviceType type, std::unique_ptr impl); + bool Has(Device::DeviceType type) const; + CclImpl *Get(Device::DeviceType type) const; private: diff --git a/infini_train/include/core/privateuse1_backend.h b/infini_train/include/core/privateuse1_backend.h new file mode 100644 index 00000000..21b15d73 --- /dev/null +++ b/infini_train/include/core/privateuse1_backend.h @@ -0,0 +1,22 @@ +#pragma once + +#include + +namespace infini_train::core { + +using PrivateUse1RegistrationCallback = void (*)(); + +struct PrivateUse1BackendRegistration { + std::string name; + PrivateUse1RegistrationCallback register_runtime = nullptr; + PrivateUse1RegistrationCallback register_kernels = nullptr; + PrivateUse1RegistrationCallback register_ccl = nullptr; +}; + +void RegisterPrivateUse1Backend(const PrivateUse1BackendRegistration ®istration); + +bool HasPrivateUse1Backend(); + +std::string GetPrivateUse1BackendName(); + +} // namespace infini_train::core diff --git a/infini_train/include/core/runtime/device_guard.h b/infini_train/include/core/runtime/device_guard.h index dc56fc6f..e25ba8c1 100644 --- a/infini_train/include/core/runtime/device_guard.h +++ b/infini_train/include/core/runtime/device_guard.h @@ -62,6 +62,10 @@ class DeviceGuardImpl { virtual ~DeviceGuardImpl() = default; + // Performs runtime initialization on first backend use. Implementations + // that need initialization should make this method idempotent. + virtual void Initialize(); + // ---------------------------------------------------------------------- // Device management // ---------------------------------------------------------------------- @@ -202,6 +206,8 @@ class DeviceGuardImplRegistry { void Register(Device::DeviceType type, std::unique_ptr impl); + bool Has(Device::DeviceType type) const; + DeviceGuardImpl *Get(Device::DeviceType type) const; private: diff --git a/infini_train/include/device.h b/infini_train/include/device.h index 28db395f..001f41ee 100644 --- a/infini_train/include/device.h +++ b/infini_train/include/device.h @@ -1,8 +1,10 @@ #pragma once #include +#include #include #include +#include #include "infini_train/include/nn/parallel/rank.h" @@ -13,7 +15,8 @@ class Device { enum class DeviceType : int8_t { kCPU = 0, kCUDA = 1, - kCount = 2, + kPrivateUse1 = 2, + kCount = 3, kInvalid = -1, }; @@ -30,6 +33,9 @@ class Device { bool IsCPU() const; bool IsCUDA() const; + bool IsPrivateUse1() const; + + static std::optional ParseType(std::string_view name); std::string ToString() const; diff --git a/infini_train/include/dispatcher.h b/infini_train/include/dispatcher.h index 2c390ec8..51ccb662 100644 --- a/infini_train/include/dispatcher.h +++ b/infini_train/include/dispatcher.h @@ -1,6 +1,7 @@ #pragma once #include +#include #include #include @@ -55,6 +56,8 @@ class Dispatcher { return instance; } + bool HasKernel(const KeyT &key) const { return key_to_kernel_map_.contains(key); } + const KernelFunction &GetKernel(KeyT key) const { CHECK(key_to_kernel_map_.contains(key)) << "Kernel not found: " << key.second << " on device: " << static_cast(key.first); diff --git a/infini_train/src/core/ccl/ccl.cc b/infini_train/src/core/ccl/ccl.cc index 92c14cc6..1bddee0e 100644 --- a/infini_train/src/core/ccl/ccl.cc +++ b/infini_train/src/core/ccl/ccl.cc @@ -85,6 +85,8 @@ void CclImplRegistry::Register(Device::DeviceType type, std::unique_ptr impls_[type] = std::move(impl); } +bool CclImplRegistry::Has(Device::DeviceType type) const { return impls_.contains(type); } + CclImpl *CclImplRegistry::Get(Device::DeviceType type) const { auto it = impls_.find(type); if (it == impls_.end()) { diff --git a/infini_train/src/core/privateuse1_backend.cc b/infini_train/src/core/privateuse1_backend.cc new file mode 100644 index 00000000..42abac8e --- /dev/null +++ b/infini_train/src/core/privateuse1_backend.cc @@ -0,0 +1,97 @@ +#include "infini_train/include/core/privateuse1_backend.h" + +#include +#include +#include +#include + +#include "glog/logging.h" + +#include "infini_train/include/core/ccl/ccl.h" +#include "infini_train/include/core/runtime/device_guard.h" +#include "infini_train/include/dispatcher.h" + +namespace infini_train::core { +namespace { + +struct PrivateUse1BackendState { + enum class Status { + kUnregistered, + kRegistering, + kRegistered, + }; + + Status status = Status::kUnregistered; + std::string name; +}; + +PrivateUse1BackendState g_backend_state; +std::mutex g_backend_mutex; + +constexpr std::array kRequiredKernels = { + "Cast", + "Fill", + "NoOpForward", + "NoOpBackward", +}; + +bool IsValidBackendName(std::string_view name) { + return !name.empty() && name != "cpu" && name != "cuda" + && std::all_of(name.begin(), name.end(), + [](unsigned char c) { return (c >= 'a' && c <= 'z') || (c >= '0' && c <= '9') || c == '_'; }); +} + +} // namespace + +void RegisterPrivateUse1Backend(const PrivateUse1BackendRegistration ®istration) { + CHECK(IsValidBackendName(registration.name)) + << "PrivateUse1 backend name must be non-reserved and contain only lowercase ASCII letters, digits, and " + "underscores"; + CHECK(registration.register_runtime != nullptr) << "PrivateUse1 backend must register a runtime"; + CHECK(registration.register_kernels != nullptr) << "PrivateUse1 backend must register kernels"; + + { + std::lock_guard lock(g_backend_mutex); + CHECK(g_backend_state.status == PrivateUse1BackendState::Status::kUnregistered) + << "PrivateUse1 backend is already registered or registration is in progress as " << g_backend_state.name; + g_backend_state.status = PrivateUse1BackendState::Status::kRegistering; + g_backend_state.name = registration.name; + } + + // Backend callbacks may query the provider metadata, so they must run + // outside g_backend_mutex. + registration.register_runtime(); + CHECK(DeviceGuardImplRegistry::Instance().Has(Device::DeviceType::kPrivateUse1)) + << "PrivateUse1 runtime callback did not register DeviceGuardImpl"; + + if (registration.register_ccl != nullptr) { + registration.register_ccl(); + CHECK(CclImplRegistry::Instance().Has(Device::DeviceType::kPrivateUse1)) + << "PrivateUse1 CCL callback did not register CclImpl"; + } + + registration.register_kernels(); + for (const auto kernel : kRequiredKernels) { + CHECK(Dispatcher::Instance().HasKernel({Device::DeviceType::kPrivateUse1, std::string(kernel)})) + << "PrivateUse1 backend is missing required kernel " << kernel; + } + + { + std::lock_guard lock(g_backend_mutex); + CHECK(g_backend_state.status == PrivateUse1BackendState::Status::kRegistering); + g_backend_state.status = PrivateUse1BackendState::Status::kRegistered; + } +} + +bool HasPrivateUse1Backend() { + std::lock_guard lock(g_backend_mutex); + return g_backend_state.status == PrivateUse1BackendState::Status::kRegistered; +} + +std::string GetPrivateUse1BackendName() { + std::lock_guard lock(g_backend_mutex); + return g_backend_state.status == PrivateUse1BackendState::Status::kUnregistered ? "privateuse1" + : g_backend_state.name; +} + +} // namespace infini_train::core diff --git a/infini_train/src/core/runtime/device_guard.cc b/infini_train/src/core/runtime/device_guard.cc index fbcb316f..dd2ffe1e 100644 --- a/infini_train/src/core/runtime/device_guard.cc +++ b/infini_train/src/core/runtime/device_guard.cc @@ -11,6 +11,8 @@ namespace infini_train::core { // DeviceGuardImpl (base fallback: FATAL only) +void DeviceGuardImpl::Initialize() {} + void DeviceGuardImpl::SetDevice(Device) const { LOG(FATAL) << "DeviceGuardImpl::SetDevice is not implemented."; } int DeviceGuardImpl::DeviceCount() const { @@ -143,18 +145,11 @@ void DeviceGuardImplRegistry::Register(Device::DeviceType type, std::unique_ptr< LOG(FATAL) << std::format("DeviceGuardImpl for type {} already registrered", static_cast(type)); } - if (!impls_.empty()) { - for (auto &kv : impls_) { - if (kv.first != Device::DeviceType::kCPU) { - LOG(FATAL) << std::format("Only CPU and one GPU backend allowed. Already have GPU={}, new={} rejected.", - static_cast(kv.first), static_cast(type)); - } - } - } - impls_[type] = std::move(impl); } +bool DeviceGuardImplRegistry::Has(Device::DeviceType type) const { return impls_.contains(type); } + DeviceGuardImpl *DeviceGuardImplRegistry::Get(Device::DeviceType type) const { auto it = impls_.find(type); if (it == impls_.end()) { @@ -163,6 +158,10 @@ DeviceGuardImpl *DeviceGuardImplRegistry::Get(Device::DeviceType type) const { return it->second.get(); } -DeviceGuardImpl *GetDeviceGuardImpl(Device::DeviceType type) { return DeviceGuardImplRegistry::Instance().Get(type); } +DeviceGuardImpl *GetDeviceGuardImpl(Device::DeviceType type) { + auto *impl = DeviceGuardImplRegistry::Instance().Get(type); + impl->Initialize(); + return impl; +} } // namespace infini_train::core diff --git a/infini_train/src/device.cc b/infini_train/src/device.cc index 1bb3aaad..455c882f 100644 --- a/infini_train/src/device.cc +++ b/infini_train/src/device.cc @@ -7,6 +7,7 @@ #include "glog/logging.h" +#include "infini_train/include/core/privateuse1_backend.h" #include "infini_train/include/nn/parallel/global.h" namespace infini_train { @@ -26,9 +27,38 @@ bool Device::IsCPU() const { return type_ == DeviceType::kCPU; } bool Device::IsCUDA() const { return type_ == DeviceType::kCUDA; } +bool Device::IsPrivateUse1() const { return type_ == DeviceType::kPrivateUse1; } + +std::optional Device::ParseType(std::string_view name) { + if (name == "cpu") { + return DeviceType::kCPU; + } + if (name == "cuda") { + return DeviceType::kCUDA; + } + if (name == "privateuse1" || (core::HasPrivateUse1Backend() && name == core::GetPrivateUse1BackendName())) { + return DeviceType::kPrivateUse1; + } + return std::nullopt; +} + std::string Device::ToString() const { + std::string type_str = "Unknown"; + switch (type_) { + case DeviceType::kCPU: + type_str = "CPU"; + break; + case DeviceType::kCUDA: + type_str = "CUDA"; + break; + case DeviceType::kPrivateUse1: + type_str = core::GetPrivateUse1BackendName(); + break; + default: + break; + } std::ostringstream oss; - oss << std::format("Device({}, {})", type_ == DeviceType::kCPU ? "CPU" : "CUDA", index_); + oss << std::format("Device({}, {})", type_str, index_); return oss.str(); } diff --git a/infini_train/src/nn/parallel/ddp/distributed_data_parallel.cc b/infini_train/src/nn/parallel/ddp/distributed_data_parallel.cc index b08b64fb..c53e7d91 100644 --- a/infini_train/src/nn/parallel/ddp/distributed_data_parallel.cc +++ b/infini_train/src/nn/parallel/ddp/distributed_data_parallel.cc @@ -34,8 +34,9 @@ DistributedDataParallel::DistributedDataParallel(std::shared_ptr mod auto device = param->GetDevice(); CHECK_EQ(device.index(), rank.thread_rank()) << "All parameters must be on the same device as the module"; if (!ddp_config.gradient_bucketing_enabled && ddp_config.zero_stage < 1) { - auto hook = std::make_unique( - function::ReduceOpType::kAvg, ddp_pg_); + const auto reduce_op + = ddp_config.average_in_collective ? function::ReduceOpType::kAvg : function::ReduceOpType::kSum; + auto hook = std::make_unique(reduce_op, ddp_pg_); param->RegisterPostAccumulateGradHook(std::move(hook)); } } diff --git a/infini_train/src/nn/parallel/ddp/reducer.cc b/infini_train/src/nn/parallel/ddp/reducer.cc index a13a3937..80115b8e 100644 --- a/infini_train/src/nn/parallel/ddp/reducer.cc +++ b/infini_train/src/nn/parallel/ddp/reducer.cc @@ -401,7 +401,9 @@ void Reducer::FinalizeBucketDense(size_t bucket_index) { // FIXME(zbl): support custom hook later LOG(FATAL) << "Custom hook is not supported now"; } else { - bucket.work = ddp_pg->AllReduce(bucket.contents, function::ReduceOpType::kAvg, true); + const auto reduce_op + = ddp_config_.average_in_collective ? function::ReduceOpType::kAvg : function::ReduceOpType::kSum; + bucket.work = ddp_pg->AllReduce(bucket.contents, reduce_op, true); } } diff --git a/scripts/run_models_and_profile.bash b/scripts/run_models_and_profile.bash index b7a7b5cc..1c01f9d5 100755 --- a/scripts/run_models_and_profile.bash +++ b/scripts/run_models_and_profile.bash @@ -78,6 +78,7 @@ MIXTRAL_LLMC_FILEPATH="$(read_var MIXTRAL_LLMC_FILEPATH)"; : "${MIXTRAL_LLMC_FIL GPT2_TEST_GROUPS="$(read_var GPT2_TEST_GROUPS)"; : "${GPT2_TEST_GROUPS:=basic,zero,lora,checkpoint}" LLAMA3_TEST_GROUPS="$(read_var LLAMA3_TEST_GROUPS)"; : "${LLAMA3_TEST_GROUPS:=basic,zero,lora,checkpoint}" MIXTRAL_TEST_GROUPS="$(read_var MIXTRAL_TEST_GROUPS)"; : "${MIXTRAL_TEST_GROUPS:=moe}" +DEVICE_BACKEND="$(read_var DEVICE_BACKEND)"; : "${DEVICE_BACKEND:=cuda}" # export custom variables from config first. LOG_DIR/PROFILE_LOG_DIR are normalized below. while IFS="=" read -r k v; do @@ -513,20 +514,20 @@ for ((id=0; id +#include +#include +#include + +#include "gtest/gtest.h" + +#include "infini_train/include/autocast.h" +#include "infini_train/include/core/privateuse1_backend.h" +#include "infini_train/include/core/runtime/device_guard.h" +#include "infini_train/include/core/runtime/runtime_common.h" +#include "infini_train/include/dispatcher.h" +#include "infini_train/include/tensor.h" + +namespace infini_train::test { +namespace { + +int g_fake_runtime_initialize_count = 0; + +class FakeStream final : public core::Stream {}; +class FakeEvent final : public core::Event {}; + +class FakePrivateUse1GuardImpl final : public core::DeviceGuardImpl { +public: + void Initialize() override { + std::call_once(initialize_once_, [] { ++g_fake_runtime_initialize_count; }); + } + + Device GetDevice() const override { return Device(Type(), current_device_); } + + void SetDevice(Device device) const override { + CHECK(device.type() == Type()); + CHECK_EQ(device.index(), 0); + current_device_ = device.index(); + } + + int DeviceCount() const override { return 1; } + + Device::DeviceType Type() const override { return Device::DeviceType::kPrivateUse1; } + + core::Stream *GetStream(Device device) const override { + CHECK(device.type() == Type()); + return const_cast(&default_stream_); + } + + core::Stream *CreateStream(Device device) const override { + CHECK(device.type() == Type()); + return new FakeStream(); + } + + core::Stream *CreateStreamWithPriority(Device device, int) const override { return CreateStream(device); } + + void DestroyStream(core::Stream *stream) const override { delete stream; } + + void GetStreamPriorityRange(int *low, int *high) const override { + *low = 0; + *high = 0; + } + + void EventCreate(core::Event **event) const override { *event = new FakeEvent(); } + + void EventCreateWithFlags(core::Event **event, core::EventFlag) const override { EventCreate(event); } + + void EventDestroy(core::Event *event) const override { delete event; } + + void EventRecord(core::Event *, core::Stream *) const override {} + + void StreamWaitEvent(core::Stream *, core::Event *, uint32_t) const override {} + + core::RuntimeStatus EventSynchronize(core::Event *) const override { return core::RuntimeStatus::kSuccess; } + + core::RuntimeStatus EventQuery(core::Event *) const override { return core::RuntimeStatus::kSuccess; } + + float EventElapsedTime(core::Event *, core::Event *) const override { return 0.0F; } + + void SynchronizeDevice(Device) const override {} + + void SynchronizeStream(core::Stream *) const override {} + + void Malloc(void **dev_ptr, size_t size) override { *dev_ptr = ::operator new(size == 0 ? 1 : size); } + + void MallocAsync(void **dev_ptr, size_t size, core::Stream *) override { Malloc(dev_ptr, size); } + + void Free(void *dev_ptr) override { ::operator delete(dev_ptr); } + + void FreeAsync(void *dev_ptr, core::Stream *) override { Free(dev_ptr); } + + void Memcpy(void *dst, const void *src, size_t count, core::MemcpyKind) override { std::memcpy(dst, src, count); } + + void MemcpyAsync(void *dst, const void *src, size_t count, core::MemcpyKind kind, core::Stream *) override { + Memcpy(dst, src, count, kind); + } + + void ResetMemPoolHighWatermarks(Device) const override {} + + std::pair GetMemPoolPeakMB(Device) const override { return {0, 0}; } + +private: + inline static thread_local int8_t current_device_ = 0; + std::once_flag initialize_once_; + FakeStream default_stream_; +}; + +void FakeFill(std::shared_ptr tensor, Scalar value) { + CHECK(tensor->Dtype() == DataType::kFLOAT32); + auto *data = static_cast(tensor->DataPtr()); + for (size_t i = 0; i < tensor->NumElements(); ++i) { data[i] = value.to(); } +} + +std::shared_ptr FakeCast(std::shared_ptr input, DataType dtype) { + CHECK(input->Dtype() == dtype); + auto output = std::make_shared(input->Dims(), dtype, input->GetDevice()); + std::memcpy(output->DataPtr(), input->DataPtr(), input->SizeInBytes()); + return output; +} + +std::shared_ptr FakeNoOpForward(const std::shared_ptr &input, const std::vector &dims) { + return std::make_shared(*input, 0, dims); +} + +std::shared_ptr FakeNoOpBackward(const std::vector &dims, const std::shared_ptr &grad_output) { + return std::make_shared(*grad_output, 0, dims); +} + +void RegisterFakeRuntime() { + CHECK_EQ(core::GetPrivateUse1BackendName(), "fake"); + CHECK_EQ(Device(Device::DeviceType::kPrivateUse1, 0).ToString(), "Device(fake, 0)"); + INFINI_TRAIN_REGISTER_DEVICE_GUARD_IMPL(Device::DeviceType::kPrivateUse1, FakePrivateUse1GuardImpl) +} + +void RegisterFakeKernels() { + REGISTER_KERNEL(Device::DeviceType::kPrivateUse1, Cast, FakeCast) + REGISTER_KERNEL(Device::DeviceType::kPrivateUse1, Fill, FakeFill) + REGISTER_KERNEL(Device::DeviceType::kPrivateUse1, NoOpForward, FakeNoOpForward) + REGISTER_KERNEL(Device::DeviceType::kPrivateUse1, NoOpBackward, FakeNoOpBackward) +} + +void InitializeFakeBackend() { + static std::once_flag once; + std::call_once(once, [] { + core::PrivateUse1BackendRegistration registration; + registration.name = "fake"; + registration.register_runtime = &RegisterFakeRuntime; + registration.register_kernels = &RegisterFakeKernels; + core::RegisterPrivateUse1Backend(registration); + }); +} + +TEST(PrivateUse1BackendTest, RegistersRuntimeMetadataAndBasicKernels) { + InitializeFakeBackend(); + + EXPECT_EQ(g_fake_runtime_initialize_count, 0); + EXPECT_TRUE(core::HasPrivateUse1Backend()); + EXPECT_EQ(core::GetPrivateUse1BackendName(), "fake"); + EXPECT_EQ(Device::ParseType("fake"), Device::DeviceType::kPrivateUse1); + EXPECT_EQ(Device::ParseType("privateuse1"), Device::DeviceType::kPrivateUse1); + EXPECT_FALSE(Device::ParseType("missing").has_value()); + + Device device(Device::DeviceType::kPrivateUse1, 0); + EXPECT_EQ(device.ToString(), "Device(fake, 0)"); + EXPECT_TRUE(core::DeviceGuardImplRegistry::Instance().Has(device.type())); + EXPECT_EQ(core::GetDeviceGuardImpl(device.type())->Type(), device.type()); + EXPECT_EQ(g_fake_runtime_initialize_count, 1); + + { + AutocastGuard guard(device.type(), DataType::kBFLOAT16); + EXPECT_EQ(GetCurrentAutocastContext().autocast_dtype, DataType::kBFLOAT16); + } + EXPECT_DEATH({ AutocastGuard guard(device.type()); }, "requires an explicit dtype"); + + auto tensor = std::make_shared(std::vector{2, 3}, DataType::kFLOAT32, device); + tensor->Fill(Scalar(3.0F)); + const auto *data = static_cast(tensor->DataPtr()); + for (size_t i = 0; i < tensor->NumElements(); ++i) { EXPECT_FLOAT_EQ(data[i], 3.0F); } + + auto view = tensor->View({3, 2}); + EXPECT_EQ(view->Dims(), (std::vector{3, 2})); + EXPECT_EQ(view->DataPtr(), tensor->DataPtr()); +} + +} // namespace +} // namespace infini_train::test diff --git a/tests/backend/test_privateuse1_backend_validation.cc b/tests/backend/test_privateuse1_backend_validation.cc new file mode 100644 index 00000000..717cf953 --- /dev/null +++ b/tests/backend/test_privateuse1_backend_validation.cc @@ -0,0 +1,31 @@ +#include "gtest/gtest.h" + +#include "infini_train/include/core/privateuse1_backend.h" + +namespace infini_train::test { +namespace { + +void UnusedRegistrationCallback() {} + +core::PrivateUse1BackendRegistration MinimalRegistration() { + core::PrivateUse1BackendRegistration registration; + registration.name = "valid_name"; + registration.register_runtime = &UnusedRegistrationCallback; + registration.register_kernels = &UnusedRegistrationCallback; + return registration; +} + +TEST(PrivateUse1BackendValidationTest, RejectsReservedBackendNames) { + auto registration = MinimalRegistration(); + registration.name = "cuda"; + EXPECT_DEATH(core::RegisterPrivateUse1Backend(registration), "non-reserved"); +} + +TEST(PrivateUse1BackendValidationTest, RejectsNonAsciiBackendNames) { + auto registration = MinimalRegistration(); + registration.name = "invalid-name"; + EXPECT_DEATH(core::RegisterPrivateUse1Backend(registration), "lowercase ASCII"); +} + +} // namespace +} // namespace infini_train::test diff --git a/tests/common/CMakeLists.txt b/tests/common/CMakeLists.txt index 84c92e51..6f91e891 100644 --- a/tests/common/CMakeLists.txt +++ b/tests/common/CMakeLists.txt @@ -1,7 +1,7 @@ # Common test utilities + custom test main (initializes GlobalEnv before RUN_ALL_TESTS). add_library(test_main OBJECT test_main.cc) target_include_directories(test_main PUBLIC - ${CMAKE_SOURCE_DIR} + ${PROJECT_SOURCE_DIR} ${glog_SOURCE_DIR}/src ) target_link_libraries(test_main PUBLIC GTest::gtest)