diff --git a/.gitignore b/.gitignore index 50b9fa06..24a3e0aa 100644 --- a/.gitignore +++ b/.gitignore @@ -1,7 +1,34 @@ build/ +build-cpu/ +build-full/ +/build-*/ .cache/ .vscode/ +# Local CMake output +/cmake-build-*/ +/CMakeCache.txt +/CMakeFiles/ +/cmake_install.cmake +/CTestTestfile.cmake +/CPackConfig.cmake +/CPackSourceConfig.cmake +/compile_commands.json +/CMakeUserPresets.json +/Makefile +/Testing/ + +# Python cache +__pycache__/ +*.py[cod] +.pytest_cache/ +.mypy_cache/ + *.log *.report.rank* *.records.log.rank* + +server_code/ +.claude/ +pytorch_ref/ +CLAUDE.md diff --git a/CMakeLists.txt b/CMakeLists.txt index f92d4905..c3920acc 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -219,6 +219,12 @@ if(BUILD_TEST) add_subdirectory(tests) endif() +# Benchmarks +option(BUILD_BENCHMARK "Build InfiniTrain benchmarks" OFF) +if(BUILD_BENCHMARK) + add_subdirectory(benchmarks) +endif() + # Negative compile test: missing dtype registration must fail at compile time. set(DTYPE_DISPATCH_COMPILE_FAIL_SOURCE ${PROJECT_SOURCE_DIR}/tests/dtype/test_dtype_dispatch_compile_fail.cc) diff --git a/benchmarks/CMakeLists.txt b/benchmarks/CMakeLists.txt new file mode 100644 index 00000000..16df2ff4 --- /dev/null +++ b/benchmarks/CMakeLists.txt @@ -0,0 +1,6 @@ +add_executable(generator_benchmark generator/generator_benchmark.cc) +link_infini_train_exe(generator_benchmark) + +set_target_properties(generator_benchmark PROPERTIES + RUNTIME_OUTPUT_DIRECTORY ${CMAKE_BINARY_DIR}/benchmarks +) diff --git a/benchmarks/generator/README.md b/benchmarks/generator/README.md new file mode 100644 index 00000000..306d6e6f --- /dev/null +++ b/benchmarks/generator/README.md @@ -0,0 +1,69 @@ +# Generator benchmark + +This benchmark measures Generator state-management latency and end-to-end +`Uniform` / `Normal` tensor fill throughput. It is intentionally separate from +the GoogleTest correctness suite: benchmark results are informative and do not +make tests flaky. + +## Build + +From the repository root: + +```bash +cmake -S . -B build \ + -DCMAKE_BUILD_TYPE=Release \ + -DUSE_CUDA=ON \ + -DUSE_NCCL=OFF \ + -DBUILD_TEST=ON \ + -DBUILD_BENCHMARK=ON +cmake --build build -j +``` + +For a CPU-only build, use `-DUSE_CUDA=OFF`. + +## Correctness tests + +```bash +ctest --test-dir build -L cpu --output-on-failure +ctest --test-dir build -L cuda --output-on-failure +``` + +To run only this feature's tests: + +```bash +ctest --test-dir build -R Generator --output-on-failure +``` + +## Single benchmark + +```bash +./build/benchmarks/generator_benchmark \ + --device cuda \ + --device-index 0 \ + --op all \ + --generator explicit \ + --elements 1048576 \ + --warmup 10 \ + --iterations 100 \ + --seed 42 +``` + +The output is CSV with average latency, generated samples per second, and +effective output bandwidth. CUDA timing synchronizes the device before and +after the measured loop, so it includes kernel execution rather than only +asynchronous launch overhead. + +## Standard matrix + +```bash +bash benchmarks/generator/run_generator_benchmark.sh \ + ./build/benchmarks/generator_benchmark cpu > generator_cpu.csv + +bash benchmarks/generator/run_generator_benchmark.sh \ + ./build/benchmarks/generator_benchmark cuda > generator_cuda.csv +``` + +For comparable reports, keep the compiler, build type, CPU/GPU model, CUDA +version, warmup count, iteration count, and thread-related environment +variables fixed. Compare before/after InfiniTrain commits rather than requiring +element-wise or performance parity with PyTorch. diff --git a/benchmarks/generator/generator_benchmark.cc b/benchmarks/generator/generator_benchmark.cc new file mode 100644 index 00000000..e9553872 --- /dev/null +++ b/benchmarks/generator/generator_benchmark.cc @@ -0,0 +1,228 @@ +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#include "infini_train/include/core/runtime/device_guard.h" +#include "infini_train/include/generator.h" +#include "infini_train/include/nn/init.h" +#include "infini_train/include/tensor.h" +#include "infini_train/src/core/runtime/cpu/cpu_generator_impl.h" +#if defined(USE_CUDA) +#include "infini_train/src/core/runtime/cuda/cuda_generator_impl.h" +#endif + +namespace { + +using Clock = std::chrono::steady_clock; +using infini_train::DataType; +using infini_train::Device; +using infini_train::Generator; +using infini_train::Tensor; + +struct Options { + std::string device = "cpu"; + std::string operation = "all"; + std::string generator = "explicit"; + int device_index = 0; + int64_t elements = 1 << 20; + int warmup = 10; + int iterations = 100; + uint64_t seed = 42; +}; + +void PrintUsage(const char *program) { + std::cout + << "Usage: " << program << " [options]\n" + << " --device cpu|cuda\n" + << " --device-index N\n" + << " --op uniform|normal|state|all\n" + << " --generator explicit|default\n" + << " --elements N\n" + << " --warmup N\n" + << " --iterations N\n" + << " --seed N\n"; +} + +std::string RequireValue(int argc, char **argv, int &index) { + if (++index >= argc) { + throw std::invalid_argument(std::string("missing value for ") + argv[index - 1]); + } + return argv[index]; +} + +Options ParseOptions(int argc, char **argv) { + Options options; + for (int i = 1; i < argc; ++i) { + const std::string_view argument(argv[i]); + if (argument == "--help" || argument == "-h") { + PrintUsage(argv[0]); + std::exit(0); + } else if (argument == "--device") { + options.device = RequireValue(argc, argv, i); + } else if (argument == "--device-index") { + options.device_index = std::stoi(RequireValue(argc, argv, i)); + } else if (argument == "--op") { + options.operation = RequireValue(argc, argv, i); + } else if (argument == "--generator") { + options.generator = RequireValue(argc, argv, i); + } else if (argument == "--elements") { + options.elements = std::stoll(RequireValue(argc, argv, i)); + } else if (argument == "--warmup") { + options.warmup = std::stoi(RequireValue(argc, argv, i)); + } else if (argument == "--iterations") { + options.iterations = std::stoi(RequireValue(argc, argv, i)); + } else if (argument == "--seed") { + options.seed = std::stoull(RequireValue(argc, argv, i)); + } else { + throw std::invalid_argument(std::string("unknown option: ") + argv[i]); + } + } + + if (options.device != "cpu" && options.device != "cuda") { + throw std::invalid_argument("--device must be cpu or cuda"); + } + if (options.operation != "uniform" && options.operation != "normal" + && options.operation != "state" && options.operation != "all") { + throw std::invalid_argument("--op must be uniform, normal, state, or all"); + } + if (options.generator != "explicit" && options.generator != "default") { + throw std::invalid_argument("--generator must be explicit or default"); + } + if (options.elements <= 0 || options.warmup < 0 || options.iterations <= 0) { + throw std::invalid_argument("elements and iterations must be positive; warmup must be non-negative"); + } + return options; +} + +Device MakeDevice(const Options &options) { + if (options.device == "cpu") { + return Device(Device::DeviceType::kCPU, 0); + } +#if defined(USE_CUDA) + return Device(Device::DeviceType::kCUDA, options.device_index); +#else + throw std::invalid_argument("CUDA benchmark requested, but InfiniTrain was built without USE_CUDA"); +#endif +} + +Generator MakeGenerator(Device device, uint64_t seed) { + if (device.IsCPU()) { + return infini_train::core::cpu::createCPUGenerator(seed); + } +#if defined(USE_CUDA) + return infini_train::core::cuda::createCUDAGenerator(device.index(), seed); +#else + (void)seed; + throw std::invalid_argument("CUDA support is disabled"); +#endif +} + +void Synchronize(Device device) { + infini_train::core::GetDeviceGuardImpl(device.type())->SynchronizeDevice(device); +} + +template +double MeasureMicroseconds(Device device, int warmup, int iterations, Function &&function) { + for (int i = 0; i < warmup; ++i) { + function(); + } + Synchronize(device); + const auto start = Clock::now(); + for (int i = 0; i < iterations; ++i) { + function(); + } + Synchronize(device); + const auto end = Clock::now(); + return std::chrono::duration(end - start).count() / iterations; +} + +void PrintHeader() { + std::cout << "device,device_index,operation,generator,elements,iterations,latency_us,gsamples_s,bandwidth_gbps\n"; +} + +void PrintResult(const Options &options, std::string_view operation, double latency_us) { + const double seconds = latency_us * 1e-6; + const double samples_per_second = static_cast(options.elements) / seconds; + const double gsamples_per_second = samples_per_second / 1e9; + const double bandwidth_gbps = samples_per_second * sizeof(float) / 1e9; + std::cout << options.device << ',' << options.device_index << ',' << operation << ',' + << options.generator << ',' << options.elements << ',' << options.iterations << ',' + << std::fixed << std::setprecision(3) << latency_us << ',' << gsamples_per_second << ',' + << bandwidth_gbps << '\n'; +} + +void RunDistribution(const Options &options, Device device, std::string_view operation) { + Generator explicit_generator = MakeGenerator(device, options.seed); + std::optional generator = options.generator == "explicit" + ? std::optional(explicit_generator) + : std::nullopt; + if (!generator) { + infini_train::manual_seed(options.seed); + } + auto tensor = std::make_shared( + std::vector{options.elements}, DataType::kFLOAT32, device); + + const double latency_us = MeasureMicroseconds(device, options.warmup, options.iterations, [&] { + if (operation == "uniform") { + infini_train::nn::init::Uniform(tensor, 0.0f, 1.0f, generator); + } else { + infini_train::nn::init::Normal(tensor, 0.0f, 1.0f, generator); + } + }); + PrintResult(options, operation, latency_us); +} + +void RunStateBenchmark(const Options &options, Device device) { + Generator generator = MakeGenerator(device, options.seed); + auto state = generator.get_state(); + const double get_state_us = MeasureMicroseconds(device, options.warmup, options.iterations, [&] { + state = generator.get_state(); + }); + const double set_state_us = MeasureMicroseconds(device, options.warmup, options.iterations, [&] { + generator.set_state(*state); + }); + uint64_t seed = options.seed; + const double manual_seed_us = MeasureMicroseconds(device, options.warmup, options.iterations, [&] { + generator.set_current_seed(seed++); + }); + + Options state_options = options; + // State-management operations do not generate tensor elements. Reporting + // zero avoids presenting meaningless throughput numbers for these rows. + state_options.elements = 0; + PrintResult(state_options, "get_state", get_state_us); + PrintResult(state_options, "set_state", set_state_us); + PrintResult(state_options, "manual_seed", manual_seed_us); +} + +} // namespace + +int main(int argc, char **argv) { + try { + const Options options = ParseOptions(argc, argv); + const Device device = MakeDevice(options); + PrintHeader(); + if (options.operation == "uniform" || options.operation == "all") { + RunDistribution(options, device, "uniform"); + } + if (options.operation == "normal" || options.operation == "all") { + RunDistribution(options, device, "normal"); + } + if (options.operation == "state" || options.operation == "all") { + RunStateBenchmark(options, device); + } + return 0; + } catch (const std::exception &error) { + std::cerr << "generator_benchmark: " << error.what() << '\n'; + PrintUsage(argv[0]); + return 2; + } +} diff --git a/benchmarks/generator/run_generator_benchmark.sh b/benchmarks/generator/run_generator_benchmark.sh new file mode 100644 index 00000000..b408a9ed --- /dev/null +++ b/benchmarks/generator/run_generator_benchmark.sh @@ -0,0 +1,46 @@ +#!/usr/bin/env bash +set -euo pipefail + +binary="${1:-./build/benchmarks/generator_benchmark}" +device="${2:-cpu}" +device_index="${DEVICE_INDEX:-0}" +warmup="${WARMUP:-10}" +iterations="${ITERATIONS:-100}" +seed="${SEED:-42}" + +if [[ ! -x "${binary}" ]]; then + echo "benchmark executable not found or not executable: ${binary}" >&2 + exit 2 +fi + +echo "# timestamp=$(date --iso-8601=seconds)" >&2 +echo "# host=$(hostname)" >&2 +echo "# kernel=$(uname -srmo)" >&2 +echo "# binary=${binary}" >&2 +echo "device,device_index,operation,generator,elements,iterations,latency_us,gsamples_s,bandwidth_gbps" + +for generator in explicit default; do + for operation in uniform normal; do + for elements in 1024 1048576 16777216; do + "${binary}" \ + --device "${device}" \ + --device-index "${device_index}" \ + --op "${operation}" \ + --generator "${generator}" \ + --elements "${elements}" \ + --warmup "${warmup}" \ + --iterations "${iterations}" \ + --seed "${seed}" | awk 'NR > 1' + done + done +done + +"${binary}" \ + --device "${device}" \ + --device-index "${device_index}" \ + --op state \ + --generator explicit \ + --elements 1 \ + --warmup "${warmup}" \ + --iterations "${iterations}" \ + --seed "${seed}" | awk 'NR > 1' diff --git a/example/common/tokenizer.cc b/example/common/tokenizer.cc index 9541454a..1f6ca0f4 100644 --- a/example/common/tokenizer.cc +++ b/example/common/tokenizer.cc @@ -10,6 +10,7 @@ #include "glog/logging.h" #include "example/common/utils.h" +#include "infini_train/include/autograd/grad_mode.h" #include "infini_train/include/nn/functional.h" #include "infini_train/include/nn/modules/module.h" #include "infini_train/include/tensor.h" @@ -107,6 +108,7 @@ std::string Tokenizer::Decode(uint32_t token_id) const { void Tokenizer::GenerateText(infini_train::nn::Module &model, uint32_t batch_size, uint32_t sequence_length, uint32_t text_length, Device device) const { + CHECK_LE(text_length, sequence_length) << "text_length must be <= sequence_length"; std::vector dims; dims.assign({batch_size, sequence_length}); // x_tensor (FLAGS_batch_size, FLAGS_sequence_length) eq:(4, 64) @@ -121,20 +123,23 @@ void Tokenizer::GenerateText(infini_train::nn::Module &model, uint32_t batch_siz std::cout << "The meaning of life is"; auto x = std::make_shared(x_tensor.To(device)); - uint64_t kRngState = kRngState; + uint64_t rng_state = kRngState; LOG(INFO) << "start generate text:"; auto cpu_device = Device(); for (int t = prompt_len; t < text_length; ++t) { x = std::make_shared(x->To(device)); // CPU->calc device - // TODO(jym): use no_grad forward later - auto logits = model.Forward({x})[0]; - auto logits_orignal = nn::function::Softmax(logits, -1); + std::shared_ptr logits_orignal; + { + infini_train::autograd::NoGradGuard no_grad; + auto logits = model.Forward({x})[0]; + logits_orignal = nn::function::Softmax(logits, -1); + } auto logits_cpu = logits_orignal->To(cpu_device); auto data = logits_cpu.DataPtr(); - auto vocab_size = logits->Dims()[2]; + auto vocab_size = logits_orignal->Dims()[2]; float *probs = static_cast(data) + (t - 1) * vocab_size; - float coin = RandomF32(kRngState); + float coin = RandomF32(rng_state); int next_token = SampleMult(probs, vocab_size, coin); x = std::make_shared(x->To(cpu_device)); // calc device->CPU diff --git a/example/mnist/dataset.cc b/example/mnist/dataset.cc index ee683f6d..6b8ec08c 100644 --- a/example/mnist/dataset.cc +++ b/example/mnist/dataset.cc @@ -89,7 +89,7 @@ MNISTDataset::MNISTDataset(const std::string &dataset, bool train) std::format("{}/{}-labels-idx1-ubyte", dataset, train ? kTrainPrefix : kTestPrefix))), image_dims_(image_file_.dims.begin() + 1, image_file_.dims.end()), label_dims_(label_file_.dims.begin() + 1, label_file_.dims.end()), - image_size_in_bytes_(kSN3TypeToSize.at(image_file_.type) + image_size_in_bytes_(sizeof(float) * std::accumulate(image_dims_.begin(), image_dims_.end(), 1, std::multiplies())), label_size_in_bytes_(kSN3TypeToSize.at(label_file_.type) * std::accumulate(label_dims_.begin(), label_dims_.end(), 1, std::multiplies())) { diff --git a/infini_train/include/autograd/dropout.h b/infini_train/include/autograd/dropout.h new file mode 100644 index 00000000..9b8451d2 --- /dev/null +++ b/infini_train/include/autograd/dropout.h @@ -0,0 +1,34 @@ +#pragma once + +#include +#include +#include + +#include "infini_train/include/autograd/function.h" +#include "infini_train/include/generator.h" + +namespace infini_train { +class Tensor; +} + +namespace infini_train::autograd { + +class Dropout final : public Function { +public: + static constexpr char kType[] = "DropoutFunction"; + + Dropout(double p, std::optional generator) + : Function(kType), p_(p), generator_(std::move(generator)) {} + + std::vector> Forward(const std::vector> &input_tensors) override; + void SetupContext(const std::vector> &input_tensors, + const std::vector> &output_tensors) override; + std::vector> Backward(const std::vector> &grad_outputs) override; + +private: + double p_ = 0.0; + std::optional generator_; + std::shared_ptr mask_; +}; + +} // namespace infini_train::autograd diff --git a/infini_train/include/autograd/elementwise.h b/infini_train/include/autograd/elementwise.h index c4333b64..16c33c4f 100644 --- a/infini_train/include/autograd/elementwise.h +++ b/infini_train/include/autograd/elementwise.h @@ -104,6 +104,8 @@ class Exp : public Function { explicit Exp() : Function(kType) {} std::vector> Forward(const std::vector> &input_tensors) override; + void SetupContext(const std::vector> &input_tensors, + const std::vector> &output_tensors) override; std::vector> Backward(const std::vector> &grad_outputs) override; }; diff --git a/infini_train/include/core/runtime/dispatch_stub.h b/infini_train/include/core/runtime/dispatch_stub.h new file mode 100644 index 00000000..13e7c4bb --- /dev/null +++ b/infini_train/include/core/runtime/dispatch_stub.h @@ -0,0 +1,68 @@ +#pragma once + +#include +#include + +#include "glog/logging.h" + +#include "infini_train/include/device.h" + +namespace infini_train { + +// ============================================================ +// DispatchStub — 设备无关的函数指针分发（仿 PyTorch DispatchStub） +// ============================================================ +// 每个 stub 按 DeviceType 存一组函数指针，调用时按设备查表分发。 +// 新增后端只需在 DeviceType 枚举加一项，kCount 自动适配。 +// +// 使用： +// 1. DECLARE_DISPATCH(fn_type, name) — 声明 extern 全局 stub +// 2. DEFINE_DISPATCH(name) — 定义全局 stub 实例 +// 3. REGISTER_DISPATCH(name, dt, fn) — 后端注册函数指针 +// 4. name(device_type, args...) — 调用，自动分发 +// +template +class DispatchStub { +public: + using FnType = FnPtr; + + static constexpr size_t kNumDevices = static_cast(Device::DeviceType::kCount); + + void register_kernel(Device::DeviceType dt, FnPtr fn) { + table_[static_cast(dt)] = fn; + } + + template + auto operator()(Device::DeviceType dt, Args&&... args) const { + auto idx = static_cast(dt); + CHECK(idx < kNumDevices) << "Invalid device type " << static_cast(dt); + CHECK(table_[idx] != nullptr) << "Dispatch kernel not registered for device type " + << static_cast(dt); + return (*table_[idx])(std::forward(args)...); + } + +private: + FnPtr table_[kNumDevices] = {}; +}; + +// ---- 宏 ---- + +// 声明：放在头文件里，告诉其他编译单元"这个 stub 存在" +#define DECLARE_DISPATCH(fn_type, name) \ + extern DispatchStub name + +// 定义：放在一个 .cpp 里，分配实际存储空间 +#define DEFINE_DISPATCH(name) \ + DispatchStub name + +// 注册：放在后端 .cpp/.cu 里，静态初始化时把函数填进表 +#define INFINI_TRAIN_CONCAT_IMPL(x, y) x##y +#define INFINI_TRAIN_CONCAT(x, y) INFINI_TRAIN_CONCAT_IMPL(x, y) + +#define REGISTER_DISPATCH(name, device_type, fn) \ + static const bool INFINI_TRAIN_CONCAT(name##_registered_, __COUNTER__) = []() { \ + (name).register_kernel((device_type), (fn)); \ + return true; \ + }(); + +} // namespace infini_train diff --git a/infini_train/include/core/runtime/distribution_kernels.h b/infini_train/include/core/runtime/distribution_kernels.h new file mode 100644 index 00000000..48a6c268 --- /dev/null +++ b/infini_train/include/core/runtime/distribution_kernels.h @@ -0,0 +1,28 @@ +#pragma once + +/// distribution_kernels.h +/// +/// 公共 API：uniform_kernel / normal_kernel +/// +/// init.cc 等上层代码调用这些包装函数，内部通过 dispatch_stub 自动分发到 +/// 正确的设备后端（CPU / 未来 CUDA）。 +/// +/// 仿 PyTorch aten/src/ATen/native/Distributions.cpp 中的 struct UniformStub 等包装层。 + +#include +#include + +#include "infini_train/include/device.h" +#include "infini_train/include/generator.h" + +namespace infini_train { + +class Tensor; + +void uniform_kernel(Tensor &tensor, double from, double to, + const std::optional &gen); + +void normal_kernel(Tensor &tensor, double mean, double std, + const std::optional &gen); + +} // namespace infini_train diff --git a/infini_train/include/core/runtime/distribution_stubs.h b/infini_train/include/core/runtime/distribution_stubs.h new file mode 100644 index 00000000..c6d0d682 --- /dev/null +++ b/infini_train/include/core/runtime/distribution_stubs.h @@ -0,0 +1,29 @@ +#pragma once + +/// distribution_stubs.h +/// +/// 内部细节：声明 uniform 和 normal 两个 DispatchStub。 +/// 上层代码不应直接调用这些 stub，而应使用 distribution_kernels.h +/// 中的 uniform_kernel() / normal_kernel() 包装函数。 +/// +/// 仿 PyTorch aten/src/ATen/native/UnaryOps.h 中的 DECLARE_DISPATCH 声明。 + +#include +#include + +#include "infini_train/include/core/runtime/dispatch_stub.h" +#include "infini_train/include/generator.h" + +namespace infini_train { + +class Tensor; + +DECLARE_DISPATCH(void (*)(Tensor &tensor, double a, double b, + const std::optional &gen), + uniform_stub); + +DECLARE_DISPATCH(void (*)(Tensor &tensor, double a, double b, + const std::optional &gen), + normal_stub); + +} // namespace infini_train diff --git a/infini_train/include/core/runtime/distributions_helper.h b/infini_train/include/core/runtime/distributions_helper.h new file mode 100644 index 00000000..f82f9065 --- /dev/null +++ b/infini_train/include/core/runtime/distributions_helper.h @@ -0,0 +1,172 @@ +#pragma once + +#include +#include +#include +#include +#include +#include + +namespace infini_train { +namespace { + +// ============================================================ +// uniform_real_distribution — 均匀分布函子 +// ============================================================ +// 仿 PyTorch at::uniform_real_distribution +// （aten/src/ATen/core/DistributionsHelper.h:99） +// +// 调用 generator->random()（float）或 generator->random64()（double） +// 获取原始随机比特，变换到 [from, to) 区间。 +// +// 模板参数 T = float | double +// 模板参数 RNG = CPUGeneratorImpl* | CUDAGeneratorImpl* | ... +// +template +struct uniform_real_distribution { + uniform_real_distribution(T from, T to) : from_(from), to_(to) { + assert(from <= to); + assert(to - from <= std::numeric_limits::max()); + } + // + // 不允许重新绑定 from/to，因为分布函子是无状态的 + uniform_real_distribution(const uniform_real_distribution &) = default; + uniform_real_distribution &operator=(const uniform_real_distribution &) = delete; + + template + T operator()(RNG *generator) const { + if constexpr (std::is_same_v) { + return transform(generator->random64()); + } else { + return transform(generator->random()); + } + } + +private: + T from_; + T to_; + + // 变换：raw bits → [0, 1) → [from_, to_) + // 仿 PyTorch at::transformation::uniform_real + // （aten/src/ATen/core/TransformationHelper.h:84） + template + T transform(V val) const { + constexpr auto MASK + = static_cast((static_cast(1) << std::numeric_limits::digits) - 1); + constexpr auto DIVISOR + = static_cast(1) / (static_cast(1) << std::numeric_limits::digits); + T x = (val & MASK) * DIVISOR; + return x * (to_ - from_) + from_; + } +}; + +// ============================================================ +// Box-Muller 正态分布缓存辅助（SFINAE） +// ============================================================ +// 仿 PyTorch at::maybe_get_next_normal_sample / maybe_set_next_normal_sample +// （aten/src/ATen/core/DistributionsHelper.h:120-163） +// +// 如果 RNG 有 next_float_normal_sample() 系列方法（如 CPUGeneratorImpl）， +// 则 Box-Muller 第二个样本被缓存到 Generator 里。 +// 如果 RNG 没有这些方法（如 CUDA curand state），则 SFINAE 回退到 no-op， +// 每次生成两个样本但只返回一个（丢弃另一个）。 +// + +// ---- get: 尝试读取缓存 ---- + +template +bool maybe_get_next_normal_sample(RNG *generator, double *ret) { + const auto sample = generator->next_double_normal_sample(); + if (!sample.has_value()) + return false; + *ret = sample.value(); + generator->set_next_double_normal_sample(std::nullopt); + return true; +} + +template +bool maybe_get_next_normal_sample(RNG *generator, float *ret) { + const auto sample = generator->next_float_normal_sample(); + if (!sample.has_value()) + return false; + *ret = sample.value(); + generator->set_next_float_normal_sample(std::nullopt); + return true; +} + +// 兜底：不支持缓存时总是返回 false +template +bool maybe_get_next_normal_sample(RNG * /*generator*/, void * /*ret*/) { + return false; +} + +// ---- set: 写入缓存 ---- + +template +void maybe_set_next_normal_sample(RNG *generator, const double *cache) { + generator->set_next_double_normal_sample(*cache); +} + +template +void maybe_set_next_normal_sample(RNG *generator, const float *cache) { + generator->set_next_float_normal_sample(*cache); +} + +// 兜底：不支持缓存时 no-op +template +void maybe_set_next_normal_sample(RNG * /*generator*/, const void * /*cache*/) {} + +// ============================================================ +// normal_distribution — 正态分布函子（Box-Muller） +// ============================================================ +// 仿 PyTorch at::normal_distribution +// （aten/src/ATen/core/DistributionsHelper.h:172） +// +// Box-Muller 每次产生两个正态样本，第二个缓存到 Generator。 +// +template +struct normal_distribution {//用struct能记录 + normal_distribution(T mean, T stdv) : mean_(mean), stdv_(stdv) { + assert(stdv >= 0); + } + // + normal_distribution(const normal_distribution &) = default; + normal_distribution &operator=(const normal_distribution &) = delete; + + template + T operator()(RNG *generator) const { + T ret; + // 先检查缓存 + if (maybe_get_next_normal_sample(generator, &ret)) { + return ret * stdv_ + mean_; + } + + // 生成两个 [0, 1) 均匀样本 + uniform_real_distribution uniform(static_cast(0), static_cast(1)); + const T u1 = uniform(generator); + const T u2 = uniform(generator); + + // Box-Muller 变换 + const T r = std::sqrt(static_cast(-2.0) * std::log1p(-u2)); + constexpr T kTwoPi = static_cast(2.0 * M_PI); + const T theta = kTwoPi * u1; + const T sample = r * std::sin(theta); + + // 缓存第二个样本 + maybe_set_next_normal_sample(generator, &sample); + + ret = r * std::cos(theta); + return ret * stdv_ + mean_; + } + +private: + T mean_; + T stdv_; +}; + +} // namespace +} // namespace infini_train diff --git a/infini_train/include/core/runtime/dropout_kernels.h b/infini_train/include/core/runtime/dropout_kernels.h new file mode 100644 index 00000000..f02a3b8b --- /dev/null +++ b/infini_train/include/core/runtime/dropout_kernels.h @@ -0,0 +1,16 @@ +#pragma once + +#include + +#include "infini_train/include/generator.h" + +namespace infini_train { + +class Tensor; + +void dropout_forward_kernel(Tensor &output, Tensor &mask, const Tensor &input, double p, + const std::optional &generator); + +void dropout_backward_kernel(Tensor &grad_input, const Tensor &grad_output, const Tensor &mask, double p); + +} // namespace infini_train diff --git a/infini_train/include/core/runtime/dropout_stubs.h b/infini_train/include/core/runtime/dropout_stubs.h new file mode 100644 index 00000000..5faa91a0 --- /dev/null +++ b/infini_train/include/core/runtime/dropout_stubs.h @@ -0,0 +1,19 @@ +#pragma once + +#include + +#include "infini_train/include/core/runtime/dispatch_stub.h" +#include "infini_train/include/generator.h" + +namespace infini_train { + +class Tensor; + +DECLARE_DISPATCH(void (*)(Tensor &output, Tensor &mask, const Tensor &input, double p, + const std::optional &generator), + dropout_forward_stub); + +DECLARE_DISPATCH(void (*)(Tensor &grad_input, const Tensor &grad_output, const Tensor &mask, double p), + dropout_backward_stub); + +} // namespace infini_train diff --git a/infini_train/include/generator.h b/infini_train/include/generator.h new file mode 100644 index 00000000..c59ecdae --- /dev/null +++ b/infini_train/include/generator.h @@ -0,0 +1,181 @@ +#pragma once + +#include +#include +#include +#include +#include + +#include "infini_train/include/device.h" + +namespace infini_train { + +// 前向声明，避免循环依赖（仿 PyTorch at::Tensor 前向声明） +class Tensor; + +namespace detail { + +// Validates the common Tensor contract for serialized RNG states. +void check_rng_state(const Tensor &state); + +} // namespace detail + +// ============================================================ +// GeneratorImpl — 抽象基类（仿 c10::GeneratorImpl） +// ============================================================ +// 定义所有 RNG 后端必须实现的纯虚接口。 +// +// 拷贝/移动已删除，只能通过 clone() 显式深拷贝。 +// clone 采用 NVI（Non-Virtual Interface）模式： +// - clone() 公有非虚，内部调用 protected clone_impl() +// - 子类只需覆写 clone_impl() 返回堆上分配的同类型拷贝 +// +// 线程安全：mutex_ 是 public 的，调用方对多步原子操作自行加锁。 +// +class GeneratorImpl { +public: + explicit GeneratorImpl(Device device) : device_(device) {} + virtual ~GeneratorImpl() = default; + + // 禁止拷贝 / 移动，避免意外覆盖 RNG 状态 + GeneratorImpl(const GeneratorImpl &other) = delete; + GeneratorImpl(GeneratorImpl &&other) = delete; + GeneratorImpl &operator=(const GeneratorImpl &other) = delete; + GeneratorImpl &operator=(GeneratorImpl &&other) = delete; + + // ---- 纯虚接口（子类必须实现）---- + virtual void set_current_seed(uint64_t seed) = 0; + virtual uint64_t current_seed() const = 0; + virtual uint64_t seed() = 0; + virtual void set_state(const Tensor &state) = 0; + virtual std::shared_ptr get_state() const = 0; + + // ---- NVI clone ---- + std::shared_ptr clone() const { + return std::shared_ptr(clone_impl()); + } + + // ---- 设备 ---- + Device device() const { return device_; } + + // 线程安全（public，调用方自行加锁） + std::mutex mutex_; + +protected: + Device device_; + + // 子类覆写点：返回堆上分配的同类型拷贝 + virtual GeneratorImpl *clone_impl() const = 0; +}; + +// ============================================================ +// Generator — 值类型壳（仿 at::Generator） +// ============================================================ +// 轻量级、可拷贝的值类型。拷贝是浅拷贝——两个 Generator +// 共享同一个 GeneratorImpl，推进一个对另一个可见。 +// +// 默认构造的 Generator 处于 "undefined" 状态（impl_ == nullptr）。 +// 使用 make_generator(seed) 或 Generator(impl) +// 来创建可用的实例。 +// +class Generator { +public: + // 默认种子：一个大数，bit 分布均匀（仿 PyTorch default_rng_seed_val） + static constexpr uint64_t kDefaultSeed = 67280421310721; + + // 默认构造：undefined 状态 + Generator() = default; + + // 从已有 Impl 构造（impl 不能为 nullptr，实现在 .cc 做检查） + explicit Generator(std::shared_ptr impl); + + // 拷贝 / 移动（浅拷贝，共享 impl_） + Generator(const Generator &) = default; + Generator &operator=(const Generator &) = default; + Generator(Generator &&) = default; + Generator &operator=(Generator &&) = default; + + ~Generator() = default; + + // ---- 种子 ---- + void set_current_seed(uint64_t seed) const { impl_->set_current_seed(seed); } + uint64_t current_seed() const { return impl_->current_seed(); } + uint64_t seed() { return impl_->seed(); } + + // ---- 状态序列化（实现在 .cc，需要 Tensor 完整定义）---- + void set_state(const Tensor &state); + std::shared_ptr get_state() const; + + // ---- 设备 ---- + Device device() const { return impl_->device(); } + + // ---- 克隆（深拷贝 Impl）---- + Generator clone() const { return Generator(impl_->clone()); } + + // ---- 线程安全 ---- + std::mutex &mutex() const { return impl_->mutex_; } + + // ---- 非检查 downcast ---- + // 调用方必须先确认 Impl 类型与 Generator 的设备类型匹配。 + // 算子代码应使用 check_generator()。 + template + T *get() const { + return static_cast(impl_.get()); + } + + // ---- 底层访问 ---- + GeneratorImpl *unsafeGetGeneratorImpl() const { return impl_.get(); } + bool defined() const { return impl_ != nullptr; } + + // ---- 比较 ---- + friend bool operator==(const Generator &a, const Generator &b) { + return a.impl_ == b.impl_; + } + friend bool operator!=(const Generator &a, const Generator &b) { + return !(a == b); + } + +private: + std::shared_ptr impl_; +}; + +// ============================================================ +// 工具函数 +// ============================================================ + +// 工厂函数：make_generator(seed) +template +Generator make_generator(Args &&...args) { + return Generator(std::make_shared(std::forward(args)...)); +} + +// 检查 Generator 已定义且属于 T 对应的设备后端，再进行 downcast。 +template +T *check_generator(const Generator &generator) { + if (!generator.defined()) { + throw std::invalid_argument("Generator with undefined implementation is not allowed"); + } + if (T::device_type() != generator.device().type()) { + throw std::invalid_argument("Generator device type does not match the requested backend"); + } + + auto *impl = dynamic_cast(generator.unsafeGetGeneratorImpl()); + if (impl == nullptr) { + throw std::invalid_argument("Generator implementation does not match the requested backend"); + } + return impl; +} + +// 显式 Generator 优先；未提供或未定义时使用该设备的默认 Generator。 +template +T *get_generator_or_default(const std::optional &generator, + const Generator &default_generator) { + return generator.has_value() && generator->defined() + ? check_generator(*generator) + : check_generator(default_generator); +} + +// Reset the default generators for all enabled devices. +void manual_seed(uint64_t seed); + +} // namespace infini_train diff --git a/infini_train/include/nn/functional.h b/infini_train/include/nn/functional.h index e4354fd1..7ed6d047 100644 --- a/infini_train/include/nn/functional.h +++ b/infini_train/include/nn/functional.h @@ -2,8 +2,13 @@ #include #include +#include #include +#include "infini_train/include/datatype.h" +#include "infini_train/include/device.h" +#include "infini_train/include/generator.h" + namespace infini_train { class Tensor; } @@ -47,6 +52,19 @@ std::shared_ptr Triu(const std::shared_ptr &input, int64_t diago // A tensor of the given shape filled with the scalar value 1. std::shared_ptr Ones(const std::vector size); +// Returns a tensor with uniformly distributed random values in [0, 1). +std::shared_ptr Rand(const std::vector &size, DataType dtype = DataType::kFLOAT32, + Device device = Device(), std::optional generator = std::nullopt, + bool requires_grad = false); + +// Returns a tensor with normally distributed random values with mean 0 and standard deviation 1. +std::shared_ptr Randn(const std::vector &size, DataType dtype = DataType::kFLOAT32, + Device device = Device(), std::optional generator = std::nullopt, + bool requires_grad = false); + +std::shared_ptr Dropout(const std::shared_ptr &input, double p = 0.5, bool training = true, + std::optional generator = std::nullopt); + // Returns a new tensor with the reciprocal of the elements of input. // // Args: diff --git a/infini_train/include/nn/init.h b/infini_train/include/nn/init.h index fc6effec..05b95b2c 100644 --- a/infini_train/include/nn/init.h +++ b/infini_train/include/nn/init.h @@ -2,11 +2,11 @@ #include #include -#include #include #include "infini_train/include/datatype.h" #include "infini_train/include/device.h" +#include "infini_train/include/generator.h" namespace infini_train { class Tensor; @@ -15,7 +15,7 @@ class Device; namespace infini_train::nn::init { std::shared_ptr Normal(const std::shared_ptr &tensor, float mean = 0.0, float std = 1.0, - std::optional generator = std::nullopt); + std::optional generator = std::nullopt); std::pair CalculateFanInAndFanOut(const std::shared_ptr &tensor); @@ -42,10 +42,10 @@ enum class NonLinearityType : int8_t { std::shared_ptr KaimingUniform(const std::shared_ptr &tensor, float a = 0.0f, KaimingMode mode = KaimingMode::kFanIn, NonLinearityType non_linearity = NonLinearityType::kLeakyReLU, - std::optional generator = std::nullopt); + std::optional generator = std::nullopt); std::shared_ptr Uniform(const std::shared_ptr &tensor, float a = 0.0f, float b = 1.0f, - std::optional generator = std::nullopt); + std::optional generator = std::nullopt); std::shared_ptr Ones(const std::shared_ptr &tensor); diff --git a/infini_train/include/tensor.h b/infini_train/include/tensor.h index dcfd8927..59679716 100644 --- a/infini_train/include/tensor.h +++ b/infini_train/include/tensor.h @@ -4,7 +4,6 @@ #include #include #include -#include #include #include "Eigen/Dense" @@ -12,6 +11,7 @@ #include "infini_train/include/datatype.h" #include "infini_train/include/device.h" +#include "infini_train/include/generator.h" #include "infini_train/include/scalar.h" namespace infini_train { @@ -79,6 +79,7 @@ class Tensor : public std::enable_shared_from_this { const std::vector &Dims() const; size_t NumElements() const; DataType Dtype() const; + bool defined() const { return buffer_ != nullptr; } std::shared_ptr Detach() const; @@ -151,7 +152,7 @@ class Tensor : public std::enable_shared_from_this { // distribution std::shared_ptr Uniform(float from = 0.0f, float to = 1.0f, - std::optional generator = std::nullopt); + std::optional generator = std::nullopt); std::shared_ptr Matmul(const std::shared_ptr &other); std::shared_ptr Outer(const std::shared_ptr &other); diff --git a/infini_train/src/autograd/dropout.cc b/infini_train/src/autograd/dropout.cc new file mode 100644 index 00000000..4ac818b3 --- /dev/null +++ b/infini_train/src/autograd/dropout.cc @@ -0,0 +1,38 @@ +#include "infini_train/include/autograd/dropout.h" + +#include "infini_train/include/core/runtime/dropout_kernels.h" +#include "infini_train/include/tensor.h" + +namespace infini_train::autograd { + +std::vector> Dropout::Forward(const std::vector> &input_tensors) { + CHECK_EQ(input_tensors.size(), 1); + const auto &input = input_tensors[0]; + + auto output = std::make_shared(input->Dims(), input->Dtype(), input->GetDevice()); + mask_ = std::make_shared(input->Dims(), DataType::kUINT8, input->GetDevice()); + dropout_forward_kernel(*output, *mask_, *input, p_, generator_); + return {output}; +} + +void Dropout::SetupContext(const std::vector> &, + const std::vector> &) { + if (!ctx_.needs_input_grad().empty() && ctx_.needs_input_grad()[0]) { + ctx_.SaveForBackward({mask_}); + } + mask_.reset(); +} + +std::vector> Dropout::Backward(const std::vector> &grad_outputs) { + CHECK_EQ(grad_outputs.size(), 1); + auto saved_tensors = ctx_.GetSavedTensors(); + CHECK_EQ(saved_tensors.size(), 1); + const auto &grad_output = grad_outputs[0]; + const auto &mask = saved_tensors[0]; + + auto grad_input = std::make_shared(grad_output->Dims(), grad_output->Dtype(), grad_output->GetDevice()); + dropout_backward_kernel(*grad_input, *grad_output, *mask, p_); + return {grad_input}; +} + +} // namespace infini_train::autograd diff --git a/infini_train/src/autograd/elementwise.cc b/infini_train/src/autograd/elementwise.cc index 36a2cad7..e54bd232 100644 --- a/infini_train/src/autograd/elementwise.cc +++ b/infini_train/src/autograd/elementwise.cc @@ -182,12 +182,21 @@ std::vector> Exp::Forward(const std::vector>({device, "ExpForward"}, input)}; } +void Exp::SetupContext(const std::vector> &, + const std::vector> &output_tensors) { + const auto &output = output_tensors[0]; + ctx_.SaveForBackward({output}); +} + std::vector> Exp::Backward(const std::vector> &grad_outputs) { + auto saved_tensors = ctx_.GetSavedTensors(); + CHECK_EQ(saved_tensors.size(), 1); + const auto &output = saved_tensors[0]; CHECK_EQ(grad_outputs.size(), 1); const auto &grad_output = grad_outputs[0]; - auto device = grad_output->GetDevice().type(); - return {Dispatcher::Instance().Call>({device, "ExpBackward"}, grad_output)}; + auto device = output->GetDevice().type(); + return {Dispatcher::Instance().Call>({device, "ExpBackward"}, grad_output, output)}; } std::vector> Log::Forward(const std::vector> &input_tensors) { diff --git a/infini_train/src/core/distribution_kernels.cc b/infini_train/src/core/distribution_kernels.cc new file mode 100644 index 00000000..953da8f3 --- /dev/null +++ b/infini_train/src/core/distribution_kernels.cc @@ -0,0 +1,160 @@ +/// distribution_kernels.cpp +/// +/// 桥接层：Generator Handle → CPUGeneratorImpl* → 分布函子 +/// +/// - get_generator_or_default: 解析 Generator → 具体 Impl 指针 +/// - *_cpu_kernel: CPU 后端实现（加锁 + 遍历 + 分布函子） +/// - REGISTER_DISPATCH: 将内核注册到 dispatch 表 +/// +/// 仿 PyTorch aten/src/ATen/native/cpu/DistributionKernels.cpp（253 行） + +#include +#include + +#include "infini_train/include/core/runtime/distribution_kernels.h" +#include "infini_train/include/core/runtime/distribution_stubs.h" +#include "infini_train/include/core/runtime/distributions_helper.h" +#include "infini_train/include/tensor.h" +#include "infini_train/src/core/runtime/cpu/cpu_dispatch.h" +#include "infini_train/src/core/runtime/cpu/cpu_generator_impl.h" + +namespace infini_train { + +DEFINE_DISPATCH(uniform_stub); +DEFINE_DISPATCH(normal_stub); + +namespace { + +void check_distribution_tensor(const Tensor &tensor) { + CHECK(IsFloatingPointDType(tensor.Dtype())) + << "Uniform and Normal initialization support floating-point tensors only"; +} + +struct DistributionBounds { + double lowest; + double max; +}; + +DistributionBounds distribution_bounds(DataType dtype) { + switch (dtype) { + case DataType::kFLOAT16: { + const double max = static_cast(FP16(static_cast(0x7bff), FP16::from_bits())); + return {-max, max}; + } + case DataType::kBFLOAT16: { + const double max = static_cast(BF16(static_cast(0x7f7f), BF16::from_bits())); + return {-max, max}; + } + case DataType::kFLOAT32: + return {-std::numeric_limits::max(), std::numeric_limits::max()}; + case DataType::kFLOAT64: + return {-std::numeric_limits::max(), std::numeric_limits::max()}; + default: + LOG(FATAL) << "Unsupported distribution dtype: " << kDataTypeToDesc.at(dtype); + return {}; + } +} + +void check_uniform_parameters(const Tensor &tensor, double from, double to) { + const auto bounds = distribution_bounds(tensor.Dtype()); + CHECK_GE(from, bounds.lowest) << "uniform expects from to be within the range of " + << kDataTypeToDesc.at(tensor.Dtype()); + CHECK_LE(from, bounds.max) << "uniform expects from to be within the range of " + << kDataTypeToDesc.at(tensor.Dtype()); + CHECK_GE(to, bounds.lowest) << "uniform expects to to be within the range of " + << kDataTypeToDesc.at(tensor.Dtype()); + CHECK_LE(to, bounds.max) << "uniform expects to to be within the range of " + << kDataTypeToDesc.at(tensor.Dtype()); + CHECK_LE(from, to) << "uniform expects a [from, to) range, but found from=" << from << " > to=" << to; + CHECK_LE(to - from, bounds.max) << "uniform expects to - from to fit in " + << kDataTypeToDesc.at(tensor.Dtype()); +} + +void check_normal_parameters(double std) { + CHECK_GE(std, 0.0) << "normal expects std >= 0.0, but found std=" << std; +} + +// ---- CPU 内核实现 ---- + +template +void uniform_cpu_kernel_impl(Tensor &tensor, double from, double to, + core::cpu::CPUGeneratorImpl *generator) { + auto *buf = static_cast(tensor.DataPtr()); + uniform_real_distribution dist(static_cast(from), static_cast(to)); + const storage_t from_value = static_cast(from); + const random_t to_value = static_cast(static_cast(to)); + for (int64_t i = 0; i < tensor.NumElements(); ++i) { + const storage_t value = static_cast(dist(generator)); + buf[i] = static_cast(value) == to_value ? from_value : value; + } +} + +template +void normal_cpu_kernel_impl(Tensor &tensor, double mean, double std, + core::cpu::CPUGeneratorImpl *generator) { + auto *buf = static_cast(tensor.DataPtr()); + normal_distribution dist(static_cast(mean), static_cast(std)); + for (int64_t i = 0; i < tensor.NumElements(); ++i) { + buf[i] = static_cast(dist(generator)); + } +} + +void uniform_cpu_kernel(Tensor &tensor, double from, double to, + const std::optional &gen) { + CHECK(tensor.GetDevice().IsCPU()); + auto *cpu_gen = get_generator_or_default( + gen, core::cpu::getDefaultCPUGenerator()); + + std::lock_guard lock(cpu_gen->mutex_); + core::cpu::DispatchCpuFunc( + tensor.Dtype(), + [&]() { + using random_t = std::conditional_t, double, float>; + uniform_cpu_kernel_impl(tensor, from, to, cpu_gen); + }, + "CPU uniform"); +} + +void normal_cpu_kernel(Tensor &tensor, double mean, double std, + const std::optional &gen) { + CHECK(tensor.GetDevice().IsCPU()); + auto *cpu_gen = get_generator_or_default( + gen, core::cpu::getDefaultCPUGenerator()); + + std::lock_guard lock(cpu_gen->mutex_); + core::cpu::DispatchCpuFunc( + tensor.Dtype(), + [&]() { + using random_t = std::conditional_t, double, float>; + normal_cpu_kernel_impl(tensor, mean, std, cpu_gen); + }, + "CPU normal"); +} + +} // namespace + +// ---- REGISTER_DISPATCH：启动时自动把内核填进 dispatch 表 ---- + +REGISTER_DISPATCH(uniform_stub, Device::DeviceType::kCPU, &uniform_cpu_kernel); +REGISTER_DISPATCH(normal_stub, Device::DeviceType::kCPU, &normal_cpu_kernel); + +// ============================================================ +// 公共 API：包装层（仿 PyTorch Distributions.cpp 中的 uniform_/normal_ 入口函数） +// ============================================================ +// init.cc 等上层代码调用这些函数，不直接碰 uniform_stub()。 + +void uniform_kernel(Tensor &tensor, double from, double to, + const std::optional &gen) { + check_distribution_tensor(tensor); + check_uniform_parameters(tensor, from, to); + uniform_stub(tensor.GetDevice().type(), tensor, from, to, gen); +} + +void normal_kernel(Tensor &tensor, double mean, double std, + const std::optional &gen) { + check_distribution_tensor(tensor); + check_normal_parameters(std); + normal_stub(tensor.GetDevice().type(), tensor, mean, std, gen); +} + +} // namespace infini_train diff --git a/infini_train/src/core/dropout_kernels.cc b/infini_train/src/core/dropout_kernels.cc new file mode 100644 index 00000000..dc3251e1 --- /dev/null +++ b/infini_train/src/core/dropout_kernels.cc @@ -0,0 +1,43 @@ +#include "infini_train/include/core/runtime/dropout_kernels.h" + +#include "infini_train/include/core/runtime/dropout_stubs.h" +#include "infini_train/include/tensor.h" + +namespace infini_train { + +DEFINE_DISPATCH(dropout_forward_stub); +DEFINE_DISPATCH(dropout_backward_stub); + +namespace { + +void check_dropout_probability(double p) { + CHECK_GE(p, 0.0) << "dropout probability has to be between 0 and 1, but got " << p; + CHECK_LE(p, 1.0) << "dropout probability has to be between 0 and 1, but got " << p; +} + +void check_dropout_tensors(const Tensor &output, const Tensor &mask, const Tensor &input) { + CHECK(IsFloatingPointDType(input.Dtype())) << "Dropout supports floating-point tensors only"; + CHECK_EQ(static_cast(output.Dtype()), static_cast(input.Dtype())); + CHECK(output.GetDevice() == input.GetDevice()); + CHECK(output.Dims() == input.Dims()); + CHECK_EQ(static_cast(mask.Dtype()), static_cast(DataType::kUINT8)); + CHECK(mask.GetDevice() == input.GetDevice()); + CHECK(mask.Dims() == input.Dims()); +} + +} // namespace + +void dropout_forward_kernel(Tensor &output, Tensor &mask, const Tensor &input, double p, + const std::optional &generator) { + check_dropout_probability(p); + check_dropout_tensors(output, mask, input); + dropout_forward_stub(input.GetDevice().type(), output, mask, input, p, generator); +} + +void dropout_backward_kernel(Tensor &grad_input, const Tensor &grad_output, const Tensor &mask, double p) { + check_dropout_probability(p); + check_dropout_tensors(grad_input, mask, grad_output); + dropout_backward_stub(grad_output.GetDevice().type(), grad_input, grad_output, mask, p); +} + +} // namespace infini_train diff --git a/infini_train/src/core/runtime/cpu/cpu_generator_impl.cc b/infini_train/src/core/runtime/cpu/cpu_generator_impl.cc new file mode 100644 index 00000000..9a01aa6f --- /dev/null +++ b/infini_train/src/core/runtime/cpu/cpu_generator_impl.cc @@ -0,0 +1,240 @@ +#include "infini_train/src/core/runtime/cpu/cpu_generator_impl.h" + +#include +#include +#include +#include + +#include "glog/logging.h" + +#include "infini_train/include/tensor.h" + +namespace infini_train::core::cpu { +namespace { + +constexpr size_t kStateFooterSize = sizeof(uint64_t) + sizeof(uint8_t) + sizeof(float) + sizeof(uint8_t) + + sizeof(double); + +} // namespace + +// ============================================================ +// 非确定性随机数（仿 c10::detail::getNonDeterministicRandom） +// ============================================================ +// Linux 下读取 /dev/urandom；其他平台 fallback 到 std::random_device +// +static uint64_t getNonDeterministicRandom() { + std::random_device rd; + uint64_t val = (static_cast(rd()) << 32) | rd(); + return val; +} + +// ============================================================ +// CPUGeneratorImpl +// ============================================================ + +CPUGeneratorImpl::CPUGeneratorImpl(uint64_t seed) + : GeneratorImpl(Device(Device::DeviceType::kCPU, 0)) + , engine_(seed) + , seed_(seed) {} + +void CPUGeneratorImpl::set_current_seed(uint64_t seed) { + seed_ = seed; + next_float_normal_sample_.reset(); + next_double_normal_sample_.reset(); + engine_ = std::mt19937(seed); +} + +uint64_t CPUGeneratorImpl::current_seed() const { + return seed_; +} + +uint64_t CPUGeneratorImpl::seed() { + uint64_t random_seed = getNonDeterministicRandom(); + set_current_seed(random_seed); + return random_seed; +} + +// ============================================================ +// 状态序列化 +// ============================================================ +// 二进制格式（仿 PyTorch CPUGeneratorImplState）： +// [engine_stream: N bytes] — mt19937 的 operator<< 输出 +// [seed_: 8 bytes] — 当前种子 +// [has_float: 1 byte] — 是否有缓存的 float 正态样本 +// [float_val: 4 bytes] — 缓存的 float 正态样本 +// [has_double: 1 byte] — 是否有缓存的 double 正态样本 +// [double_val: 8 bytes] — 缓存的 double 正态样本 +// +// 注意：engine 部分使用 operator<< / operator>>，格式是实现定义的。 +// Known limitation: PyTorch 通过自己实现 mt19937_engine 暴露 data()/set_data() +// 来实现固定格式序列化。我们使用 std::mt19937，标准库不提供内部状态访问接口， +// 因此无法做到跨编译器/跨版本的固定格式。此实现在同一构建下是稳定的。 + +void CPUGeneratorImpl::set_state(const Tensor &state) { + ::infini_train::detail::check_rng_state(state); + + const size_t data_size = state.SizeInBytes(); + CHECK_GT(data_size, kStateFooterSize) << "CPU generator state is too small"; + + const uint8_t *data = static_cast(state.DataPtr()); + const size_t engine_size = data_size - kStateFooterSize; + std::string engine_str(reinterpret_cast(data), engine_size); + + std::istringstream iss(engine_str); + std::mt19937 restored_engine; + iss >> restored_engine; + CHECK(!iss.fail()) << "Invalid CPU generator engine state"; + iss >> std::ws; + CHECK(iss.eof()) << "Invalid trailing bytes in CPU generator engine state"; + + size_t offset = engine_size; + uint64_t restored_seed = 0; + std::memcpy(&restored_seed, data + offset, sizeof(restored_seed)); + offset += sizeof(restored_seed); + + const uint8_t has_float = data[offset++]; + CHECK_LE(has_float, 1) << "Invalid CPU generator float normal cache flag"; + float restored_float = 0.0f; + std::memcpy(&restored_float, data + offset, sizeof(restored_float)); + offset += sizeof(restored_float); + + const uint8_t has_double = data[offset++]; + CHECK_LE(has_double, 1) << "Invalid CPU generator double normal cache flag"; + double restored_double = 0.0; + std::memcpy(&restored_double, data + offset, sizeof(restored_double)); + + // Do not change the generator until the complete state has been validated. + engine_ = restored_engine; + seed_ = restored_seed; + next_float_normal_sample_ = has_float ? std::optional(restored_float) : std::nullopt; + next_double_normal_sample_ = has_double ? std::optional(restored_double) : std::nullopt; +} + +std::shared_ptr CPUGeneratorImpl::get_state() const { + // 1. 序列化引擎 + std::ostringstream oss; + oss << engine_; + std::string engine_str = oss.str(); + + // 2. 计算总大小 + const size_t engine_size = engine_str.size(); + const size_t total_size = engine_size + kStateFooterSize; + + auto state_tensor = std::make_shared( + std::vector{static_cast(total_size)}, + DataType::kUINT8, Device(Device::DeviceType::kCPU, 0)); + + uint8_t *data = static_cast(state_tensor->DataPtr()); + size_t offset = 0; + + // 写入引擎 + std::memcpy(data + offset, engine_str.data(), engine_size); + offset += engine_size; + + // 写入种子 + std::memcpy(data + offset, &seed_, sizeof(seed_)); + offset += sizeof(seed_); + + // 写入 float 正态缓存 + bool has_float = next_float_normal_sample_.has_value(); + data[offset++] = has_float ? 1 : 0; + float float_val = has_float ? *next_float_normal_sample_ : 0.0f; + std::memcpy(data + offset, &float_val, sizeof(float_val)); + offset += sizeof(float_val); + + // 写入 double 正态缓存 + bool has_double = next_double_normal_sample_.has_value(); + data[offset++] = has_double ? 1 : 0; + double double_val = has_double ? *next_double_normal_sample_ : 0.0; + std::memcpy(data + offset, &double_val, sizeof(double_val)); + + return state_tensor; +} + +// ============================================================ +// 随机数生成 +// ============================================================ + +uint32_t CPUGeneratorImpl::random() { + return engine_(); +} + +uint64_t CPUGeneratorImpl::random64() { + uint32_t hi = engine_(); + uint32_t lo = engine_(); + return (static_cast(hi) << 32) | lo; +} + +// ============================================================ +// Box-Muller 正态缓存 +// ============================================================ + +std::optional CPUGeneratorImpl::next_float_normal_sample() const { + return next_float_normal_sample_; +} + +std::optional CPUGeneratorImpl::next_double_normal_sample() const { + return next_double_normal_sample_; +} + +void CPUGeneratorImpl::set_next_float_normal_sample(std::optional randn) { + next_float_normal_sample_ = randn; +} + +void CPUGeneratorImpl::set_next_double_normal_sample(std::optional randn) { + next_double_normal_sample_ = randn; +} + +// ============================================================ +// clone +// ============================================================ + +std::shared_ptr CPUGeneratorImpl::clone() const { + return std::shared_ptr(clone_impl()); +} + +CPUGeneratorImpl *CPUGeneratorImpl::clone_impl() const { + auto gen = new CPUGeneratorImpl(seed_); + gen->set_engine(engine_); + gen->set_next_float_normal_sample(next_float_normal_sample_); + gen->set_next_double_normal_sample(next_double_normal_sample_); + return gen; +} + +void CPUGeneratorImpl::set_engine(std::mt19937 engine) { + engine_ = std::move(engine); +} + +// ============================================================ +// 类型标识 +// ============================================================ + +Device::DeviceType CPUGeneratorImpl::device_type() { + return Device::DeviceType::kCPU; +} + +} // namespace infini_train::core::cpu + +// ============================================================ +// 默认 Generator 管理（文件作用域，仿 PyTorch detail 命名空间） +// ============================================================ + +namespace infini_train::core::cpu { + +const Generator &getDefaultCPUGenerator() { + // 使用真随机种子初始化默认 generator（仿 PyTorch） + static auto default_gen = createCPUGenerator(getNonDeterministicRandom()); + return default_gen; +} + +Generator createCPUGenerator(uint64_t seed) { + return make_generator(seed); +} + +void manual_seed(uint64_t seed) { + const auto &default_gen = getDefaultCPUGenerator(); + std::lock_guard lock(default_gen.mutex()); + default_gen.set_current_seed(seed); +} + +} // namespace infini_train::core::cpu diff --git a/infini_train/src/core/runtime/cpu/cpu_generator_impl.h b/infini_train/src/core/runtime/cpu/cpu_generator_impl.h new file mode 100644 index 00000000..0fb2506c --- /dev/null +++ b/infini_train/src/core/runtime/cpu/cpu_generator_impl.h @@ -0,0 +1,64 @@ +#pragma once + +#include +#include +#include +#include + +#include "infini_train/include/generator.h" + +namespace infini_train::core::cpu { + +// ============================================================ +// CPUGeneratorImpl — CPU RNG 后端（仿 at::CPUGeneratorImpl） +// ============================================================ +// 包装 std::mt19937 Mersenne Twister 引擎。 +// 缓存 Box-Muller 正态分布样本以优化性能。 +// +class CPUGeneratorImpl final : public GeneratorImpl { +public: + explicit CPUGeneratorImpl(uint64_t seed = Generator::kDefaultSeed); + ~CPUGeneratorImpl() override = default; + + // ---- GeneratorImpl 接口 ---- + void set_current_seed(uint64_t seed) override; + uint64_t current_seed() const override; + uint64_t seed() override; + void set_state(const Tensor &state) override; + std::shared_ptr get_state() const override; + + // ---- clone（类型安全版本）---- + std::shared_ptr clone() const; + + // ---- 类型标识（用于 check_generator 模板）---- + static Device::DeviceType device_type(); + + // ---- 随机数生成 ---- + uint32_t random(); + uint64_t random64(); + + // ---- Box-Muller 正态缓存 ---- + std::optional next_float_normal_sample() const; + std::optional next_double_normal_sample() const; + void set_next_float_normal_sample(std::optional randn); + void set_next_double_normal_sample(std::optional randn); + +private: + CPUGeneratorImpl *clone_impl() const override; + + // ---- 引擎（private：比赛要求公共接口不暴露 std::mt19937）---- + std::mt19937 engine() const { return engine_; } + void set_engine(std::mt19937 engine); + + std::mt19937 engine_; + uint64_t seed_ = Generator::kDefaultSeed; + std::optional next_float_normal_sample_; + std::optional next_double_normal_sample_; +}; + +// ---- 默认 Generator 管理（仿 PyTorch detail 命名空间）---- +const Generator &getDefaultCPUGenerator(); +Generator createCPUGenerator(uint64_t seed); +void manual_seed(uint64_t seed); + +} // namespace infini_train::core::cpu diff --git a/infini_train/src/core/runtime/cuda/cuda_generator_impl.cc b/infini_train/src/core/runtime/cuda/cuda_generator_impl.cc new file mode 100644 index 00000000..fac84278 --- /dev/null +++ b/infini_train/src/core/runtime/cuda/cuda_generator_impl.cc @@ -0,0 +1,136 @@ +#include "infini_train/src/core/runtime/cuda/cuda_generator_impl.h" + +#include +#include +#include +#include +#include + +#include + +#include "glog/logging.h" + +#include "infini_train/include/common/cuda/common_cuda.h" +#include "infini_train/include/tensor.h" + +namespace infini_train::core::cuda { +namespace { + +constexpr size_t kStateSize = sizeof(uint64_t) * 2; + +std::once_flag default_generators_init_flag; +std::vector default_generators; +std::deque default_generator_init_flags; + +uint64_t get_non_deterministic_random() { + std::random_device random_device; + return (static_cast(random_device()) << 32) | random_device(); +} + +void init_default_generators() { + std::call_once(default_generators_init_flag, [] { + int device_count = 0; + const cudaError_t status = cudaGetDeviceCount(&device_count); + if (status == cudaErrorNoDevice) { + cudaGetLastError(); + return; + } + CHECK_EQ(status, cudaSuccess) << "cudaGetDeviceCount failed: " << cudaGetErrorString(status); + default_generators.resize(device_count); + default_generator_init_flags.resize(device_count); + }); +} + +int resolve_device_index(int8_t device_index) { + init_default_generators(); + int index = device_index; + if (index == -1) { + CUDA_CHECK(cudaGetDevice(&index)); + } + int device_count = 0; + device_count = static_cast(default_generators.size()); + CHECK(index >= 0 && index < device_count) << "Invalid CUDA device index " << index; + return index; +} + +} // namespace + +CUDAGeneratorImpl::CUDAGeneratorImpl(int8_t device_index, uint64_t seed) + : GeneratorImpl(Device(Device::DeviceType::kCUDA, device_index)) + , seed_(seed) {} + +void CUDAGeneratorImpl::set_current_seed(uint64_t seed) { + seed_ = seed; + next_philox_subsequence_ = 0; +} + +uint64_t CUDAGeneratorImpl::current_seed() const { + return seed_; +} + +uint64_t CUDAGeneratorImpl::seed() { + const uint64_t random_seed = get_non_deterministic_random(); + set_current_seed(random_seed); + return random_seed; +} + +void CUDAGeneratorImpl::set_state(const Tensor &state) { + ::infini_train::detail::check_rng_state(state); + CHECK_EQ(state.SizeInBytes(), kStateSize); + + const auto *data = static_cast(state.DataPtr()); + std::memcpy(&seed_, data, sizeof(seed_)); + std::memcpy(&next_philox_subsequence_, data + sizeof(seed_), sizeof(next_philox_subsequence_)); +} + +std::shared_ptr CUDAGeneratorImpl::get_state() const { + auto state = std::make_shared( + std::vector{static_cast(kStateSize)}, + DataType::kUINT8, + Device(Device::DeviceType::kCPU, 0)); + + auto *data = static_cast(state->DataPtr()); + std::memcpy(data, &seed_, sizeof(seed_)); + std::memcpy(data + sizeof(seed_), &next_philox_subsequence_, sizeof(next_philox_subsequence_)); + return state; +} + +Device::DeviceType CUDAGeneratorImpl::device_type() { + return Device::DeviceType::kCUDA; +} + +uint64_t CUDAGeneratorImpl::philox_subsequence(uint64_t increment) { + const uint64_t subsequence = next_philox_subsequence_; + next_philox_subsequence_ += increment; + return subsequence; +} + +CUDAGeneratorImpl *CUDAGeneratorImpl::clone_impl() const { + auto *generator = new CUDAGeneratorImpl(device().index(), seed_); + generator->next_philox_subsequence_ = next_philox_subsequence_; + return generator; +} + +const Generator &getDefaultCUDAGenerator(int8_t device_index) { + const int index = resolve_device_index(device_index); + std::call_once(default_generator_init_flags[index], [index] { + default_generators[index] = createCUDAGenerator(static_cast(index), get_non_deterministic_random()); + }); + return default_generators[index]; +} + +Generator createCUDAGenerator(int8_t device_index, uint64_t seed) { + const int index = resolve_device_index(device_index); + return make_generator(static_cast(index), seed); +} + +void manual_seed_all(uint64_t seed) { + init_default_generators(); + for (size_t index = 0; index < default_generators.size(); ++index) { + const auto &generator = getDefaultCUDAGenerator(static_cast(index)); + std::lock_guard lock(generator.mutex()); + generator.set_current_seed(seed); + } +} + +} // namespace infini_train::core::cuda diff --git a/infini_train/src/core/runtime/cuda/cuda_generator_impl.h b/infini_train/src/core/runtime/cuda/cuda_generator_impl.h new file mode 100644 index 00000000..5ffdbdc8 --- /dev/null +++ b/infini_train/src/core/runtime/cuda/cuda_generator_impl.h @@ -0,0 +1,37 @@ +#pragma once + +#include +#include + +#include "infini_train/include/generator.h" + +namespace infini_train::core::cuda { + +class CUDAGeneratorImpl final : public GeneratorImpl { +public: + explicit CUDAGeneratorImpl(int8_t device_index, uint64_t seed = Generator::kDefaultSeed); + ~CUDAGeneratorImpl() override = default; + + void set_current_seed(uint64_t seed) override; + uint64_t current_seed() const override; + uint64_t seed() override; + void set_state(const Tensor &state) override; + std::shared_ptr get_state() const override; + + static Device::DeviceType device_type(); + + // The caller must hold mutex_ while reserving Philox subsequences. + uint64_t philox_subsequence(uint64_t increment); + +private: + CUDAGeneratorImpl *clone_impl() const override; + + uint64_t seed_ = Generator::kDefaultSeed; + uint64_t next_philox_subsequence_ = 0; +}; + +const Generator &getDefaultCUDAGenerator(int8_t device_index = -1); +Generator createCUDAGenerator(int8_t device_index, uint64_t seed = Generator::kDefaultSeed); +void manual_seed_all(uint64_t seed); + +} // namespace infini_train::core::cuda diff --git a/infini_train/src/generator.cc b/infini_train/src/generator.cc new file mode 100644 index 00000000..a3e2eb2b --- /dev/null +++ b/infini_train/src/generator.cc @@ -0,0 +1,54 @@ +#include "infini_train/include/generator.h" + +#include "glog/logging.h" + +#include "infini_train/include/tensor.h" +#include "infini_train/src/core/runtime/cpu/cpu_generator_impl.h" + +#ifdef USE_CUDA +#include "infini_train/src/core/runtime/cuda/cuda_generator_impl.h" +#endif + +namespace infini_train { + +// ============================================================ +// Generator 构造函数（null 检查需要 glog，放 .cc） +// ============================================================ + +Generator::Generator(std::shared_ptr impl) + : impl_(std::move(impl)) { + CHECK(impl_) << "GeneratorImpl with nullptr is not supported"; +} + +// ============================================================ +// Generator — 状态序列化（需要 Tensor 完整定义，放 .cc） +// ============================================================ + +void Generator::set_state(const Tensor &state) { + CHECK(state.defined()) << "Undefined tensor is not allowed"; + impl_->set_state(state); +} + +std::shared_ptr Generator::get_state() const { + return impl_->get_state(); +} + +namespace detail { + +void check_rng_state(const Tensor &state) { + CHECK(state.GetDevice().IsCPU()) << "RNG state must be a CPU tensor"; + CHECK_EQ(static_cast(state.Dtype()), static_cast(DataType::kUINT8)) + << "RNG state must be a UINT8 tensor"; +} + +} // namespace detail + +void manual_seed(uint64_t seed) { + core::cpu::manual_seed(seed); + +#ifdef USE_CUDA + core::cuda::manual_seed_all(seed); +#endif +} + +} // namespace infini_train diff --git a/infini_train/src/kernels/cpu/cross_entropy.cc b/infini_train/src/kernels/cpu/cross_entropy.cc index f520b2e9..37c7e9cd 100644 --- a/infini_train/src/kernels/cpu/cross_entropy.cc +++ b/infini_train/src/kernels/cpu/cross_entropy.cc @@ -1,6 +1,7 @@ #include #include #include +#include #include #include diff --git a/infini_train/src/kernels/cpu/dropout.cc b/infini_train/src/kernels/cpu/dropout.cc new file mode 100644 index 00000000..ac909ead --- /dev/null +++ b/infini_train/src/kernels/cpu/dropout.cc @@ -0,0 +1,100 @@ +#include +#include + +#include "infini_train/include/core/runtime/dropout_stubs.h" +#include "infini_train/include/core/runtime/distributions_helper.h" +#include "infini_train/include/tensor.h" +#include "infini_train/src/core/runtime/cpu/cpu_dispatch.h" +#include "infini_train/src/core/runtime/cpu/cpu_generator_impl.h" + +namespace infini_train::kernels::cpu { +namespace { + +template +void dropout_forward_cpu_impl(Tensor &output, Tensor &mask, const Tensor &input, double p, + core::cpu::CPUGeneratorImpl *generator) { + auto *output_data = static_cast(output.DataPtr()); + auto *mask_data = static_cast(mask.DataPtr()); + const auto *input_data = static_cast(input.DataPtr()); + const int64_t n = input.NumElements(); + + if (p == 0.0) { + for (int64_t index = 0; index < n; ++index) { + mask_data[index] = 1; + output_data[index] = input_data[index]; + } + return; + } + if (p == 1.0) { + for (int64_t index = 0; index < n; ++index) { + mask_data[index] = 0; + output_data[index] = static_cast(0.0); + } + return; + } + + const random_t scale = static_cast(1.0 / (1.0 - p)); + uniform_real_distribution distribution(static_cast(0), static_cast(1)); + for (int64_t index = 0; index < n; ++index) { + const bool keep = distribution(generator) >= static_cast(p); + mask_data[index] = keep ? 1 : 0; + output_data[index] = keep + ? static_cast(static_cast(input_data[index]) * scale) + : static_cast(0.0); + } +} + +template +void dropout_backward_cpu_impl(Tensor &grad_input, const Tensor &grad_output, const Tensor &mask, double p) { + auto *grad_input_data = static_cast(grad_input.DataPtr()); + const auto *grad_output_data = static_cast(grad_output.DataPtr()); + const auto *mask_data = static_cast(mask.DataPtr()); + const random_t scale = p == 1.0 ? static_cast(0) : static_cast(1.0 / (1.0 - p)); + + for (int64_t index = 0; index < grad_output.NumElements(); ++index) { + grad_input_data[index] = mask_data[index] + ? static_cast(static_cast(grad_output_data[index]) * scale) + : static_cast(0.0); + } +} + +void dropout_forward_cpu(Tensor &output, Tensor &mask, const Tensor &input, double p, + const std::optional &generator) { + CHECK(input.GetDevice().IsCPU()); + + core::cpu::DispatchCpuFunc( + input.Dtype(), + [&]() { + using random_t = std::conditional_t, double, float>; + if (p == 0.0 || p == 1.0) { + dropout_forward_cpu_impl(output, mask, input, p, nullptr); + return; + } + auto *cpu_generator = get_generator_or_default( + generator, core::cpu::getDefaultCPUGenerator()); + std::lock_guard lock(cpu_generator->mutex_); + dropout_forward_cpu_impl(output, mask, input, p, cpu_generator); + }, + "CPU dropout forward"); +} + +void dropout_backward_cpu(Tensor &grad_input, const Tensor &grad_output, const Tensor &mask, double p) { + CHECK(grad_output.GetDevice().IsCPU()); + core::cpu::DispatchCpuFunc( + grad_output.Dtype(), + [&]() { + using random_t = std::conditional_t, double, float>; + dropout_backward_cpu_impl(grad_input, grad_output, mask, p); + }, + "CPU dropout backward"); +} + +} // namespace + +using infini_train::dropout_backward_stub; +using infini_train::dropout_forward_stub; + +REGISTER_DISPATCH(dropout_forward_stub, Device::DeviceType::kCPU, &dropout_forward_cpu); +REGISTER_DISPATCH(dropout_backward_stub, Device::DeviceType::kCPU, &dropout_backward_cpu); + +} // namespace infini_train::kernels::cpu diff --git a/infini_train/src/kernels/cpu/transform.cc b/infini_train/src/kernels/cpu/transform.cc index 1a810b44..04031631 100644 --- a/infini_train/src/kernels/cpu/transform.cc +++ b/infini_train/src/kernels/cpu/transform.cc @@ -1,5 +1,6 @@ #include #include +#include #include "glog/logging.h" diff --git a/infini_train/src/kernels/cuda/distribution.cu b/infini_train/src/kernels/cuda/distribution.cu new file mode 100644 index 00000000..4b181fdc --- /dev/null +++ b/infini_train/src/kernels/cuda/distribution.cu @@ -0,0 +1,148 @@ +#include +#include +#include + +#include + +#include "infini_train/include/common/cuda/common_cuda.h" +#include "infini_train/include/common/cuda/kernel_helper.cuh" +#include "infini_train/include/core/runtime/device_guard.h" +#include "infini_train/include/core/runtime/distribution_stubs.h" +#include "infini_train/include/generator.h" +#include "infini_train/include/tensor.h" +#include "infini_train/src/core/runtime/cuda/cuda_dispatch.h" +#include "infini_train/src/core/runtime/cuda/cuda_generator_impl.h" +#include "infini_train/src/core/runtime/cuda/cuda_runtime_common.h" + +namespace infini_train::kernels::cuda { +namespace { + +constexpr int kThreadsPerBlock = 256; + +template __device__ random_t uniform_sample(curandStatePhilox4_32_10_t *state) { + if constexpr (std::is_same_v) { + return 1.0 - curand_uniform_double(state); + } else { + return static_cast(curand(state)) * 0x1p-32f; + } +} + +template __device__ random_t normal_sample(curandStatePhilox4_32_10_t *state) { + if constexpr (std::is_same_v) { + return curand_normal_double(state); + } else { + return curand_normal(state); + } +} + +template +__global__ void UniformKernel(storage_t *data, int64_t n, random_t from, random_t to, + uint64_t seed, uint64_t subsequence) { + const int64_t index = static_cast(blockIdx.x) * blockDim.x + threadIdx.x; + if (index >= n) { + return; + } + + curandStatePhilox4_32_10_t state; + curand_init(seed, subsequence + static_cast(index), 0, &state); + const storage_t from_value = common::cuda::Cast(from); + const storage_t to_value = common::cuda::Cast(to); + const storage_t value = common::cuda::Cast(from + uniform_sample(&state) * (to - from)); + data[index] = common::cuda::Cast(value) == common::cuda::Cast(to_value) ? from_value : value; +} + +template +__global__ void NormalKernel(storage_t *data, int64_t n, random_t mean, random_t std, + uint64_t seed, uint64_t subsequence) { + const int64_t index = static_cast(blockIdx.x) * blockDim.x + threadIdx.x; + if (index >= n) { + return; + } + + curandStatePhilox4_32_10_t state; + curand_init(seed, subsequence + static_cast(index), 0, &state); + data[index] = common::cuda::Cast(mean + normal_sample(&state) * std); +} + +const core::cuda::CudaStream *get_cuda_stream(const Device &device) { + return dynamic_cast( + core::GetDeviceGuardImpl(device.type())->GetStream(device)); +} + +void uniform_cuda_kernel(Tensor &tensor, double from, double to, + const std::optional &generator) { + const Device device = tensor.GetDevice(); + CHECK(device.IsCUDA()); + const int64_t n = tensor.NumElements(); + if (n == 0) { + return; + } + core::DeviceGuard guard(device); + auto *cuda_generator = get_generator_or_default( + generator, core::cuda::getDefaultCUDAGenerator(device.index())); + + uint64_t seed = 0; + uint64_t subsequence = 0; + { + std::lock_guard lock(cuda_generator->mutex_); + seed = cuda_generator->current_seed(); + subsequence = cuda_generator->philox_subsequence(static_cast(n)); + } + + const int blocks = static_cast((n + kThreadsPerBlock - 1) / kThreadsPerBlock); + const auto *stream = get_cuda_stream(device); + core::cuda::DispatchCudaFunc( + tensor.Dtype(), + [&]() { + using random_t = std::conditional_t, double, float>; + UniformKernel<<cuda_stream()>>>( + static_cast(tensor.DataPtr()), n, static_cast(from), static_cast(to), + seed, subsequence); + }, + "CUDA uniform"); + CUDA_CHECK(cudaGetLastError()); +} + +void normal_cuda_kernel(Tensor &tensor, double mean, double std, + const std::optional &generator) { + const Device device = tensor.GetDevice(); + CHECK(device.IsCUDA()); + const int64_t n = tensor.NumElements(); + if (n == 0) { + return; + } + core::DeviceGuard guard(device); + auto *cuda_generator = get_generator_or_default( + generator, core::cuda::getDefaultCUDAGenerator(device.index())); + + uint64_t seed = 0; + uint64_t subsequence = 0; + { + std::lock_guard lock(cuda_generator->mutex_); + seed = cuda_generator->current_seed(); + subsequence = cuda_generator->philox_subsequence(static_cast(n)); + } + + const int blocks = static_cast((n + kThreadsPerBlock - 1) / kThreadsPerBlock); + const auto *stream = get_cuda_stream(device); + core::cuda::DispatchCudaFunc( + tensor.Dtype(), + [&]() { + using random_t = std::conditional_t, double, float>; + NormalKernel<<cuda_stream()>>>( + static_cast(tensor.DataPtr()), n, static_cast(mean), static_cast(std), + seed, subsequence); + }, + "CUDA normal"); + CUDA_CHECK(cudaGetLastError()); +} + +} // namespace + +using infini_train::normal_stub; +using infini_train::uniform_stub; + +REGISTER_DISPATCH(uniform_stub, Device::DeviceType::kCUDA, &uniform_cuda_kernel); +REGISTER_DISPATCH(normal_stub, Device::DeviceType::kCUDA, &normal_cuda_kernel); + +} // namespace infini_train::kernels::cuda diff --git a/infini_train/src/kernels/cuda/dropout.cu b/infini_train/src/kernels/cuda/dropout.cu new file mode 100644 index 00000000..f1315da3 --- /dev/null +++ b/infini_train/src/kernels/cuda/dropout.cu @@ -0,0 +1,144 @@ +#include +#include +#include +#include + +#include + +#include "infini_train/include/common/cuda/common_cuda.h" +#include "infini_train/include/common/cuda/kernel_helper.cuh" +#include "infini_train/include/core/runtime/device_guard.h" +#include "infini_train/include/core/runtime/dropout_stubs.h" +#include "infini_train/include/generator.h" +#include "infini_train/include/tensor.h" +#include "infini_train/src/core/runtime/cuda/cuda_dispatch.h" +#include "infini_train/src/core/runtime/cuda/cuda_generator_impl.h" +#include "infini_train/src/core/runtime/cuda/cuda_runtime_common.h" + +namespace infini_train::kernels::cuda { +namespace { + +constexpr int kThreadsPerBlock = 256; + +template __device__ random_t uniform_sample(curandStatePhilox4_32_10_t *state) { + if constexpr (std::is_same_v) { + return 1.0 - curand_uniform_double(state); + } else { + return static_cast(curand(state)) * 0x1p-32f; + } +} + +template +__global__ void DropoutForwardKernel(storage_t *output, uint8_t *mask, const storage_t *input, int64_t n, + random_t p, uint64_t seed, uint64_t subsequence) { + const int64_t index = static_cast(blockIdx.x) * blockDim.x + threadIdx.x; + if (index >= n) { + return; + } + if (p == static_cast(0)) { + mask[index] = 1; + output[index] = input[index]; + return; + } + if (p == static_cast(1)) { + mask[index] = 0; + output[index] = common::cuda::Cast(0.0f); + return; + } + + curandStatePhilox4_32_10_t state; + curand_init(seed, subsequence + static_cast(index), 0, &state); + const bool keep = uniform_sample(&state) >= p; + const random_t scale = static_cast(1) / (static_cast(1) - p); + mask[index] = keep ? 1 : 0; + output[index] = keep + ? common::cuda::Cast(common::cuda::Cast(input[index]) * scale) + : common::cuda::Cast(0.0f); +} + +template +__global__ void DropoutBackwardKernel(storage_t *grad_input, const storage_t *grad_output, const uint8_t *mask, + int64_t n, random_t p) { + const int64_t index = static_cast(blockIdx.x) * blockDim.x + threadIdx.x; + if (index >= n) { + return; + } + const random_t scale = p == static_cast(1) + ? static_cast(0) + : static_cast(1) / (static_cast(1) - p); + grad_input[index] = mask[index] + ? common::cuda::Cast(common::cuda::Cast(grad_output[index]) * scale) + : common::cuda::Cast(0.0f); +} + +const core::cuda::CudaStream *get_cuda_stream(const Device &device) { + return dynamic_cast( + core::GetDeviceGuardImpl(device.type())->GetStream(device)); +} + +void dropout_forward_cuda(Tensor &output, Tensor &mask, const Tensor &input, double p, + const std::optional &generator) { + const Device device = input.GetDevice(); + CHECK(device.IsCUDA()); + const int64_t n = input.NumElements(); + if (n == 0) { + return; + } + core::DeviceGuard guard(device); + + uint64_t seed = 0; + uint64_t subsequence = 0; + if (p > 0.0 && p < 1.0) { + auto *cuda_generator = get_generator_or_default( + generator, core::cuda::getDefaultCUDAGenerator(device.index())); + std::lock_guard lock(cuda_generator->mutex_); + seed = cuda_generator->current_seed(); + subsequence = cuda_generator->philox_subsequence(static_cast(n)); + } + + const int blocks = static_cast((n + kThreadsPerBlock - 1) / kThreadsPerBlock); + const auto *stream = get_cuda_stream(device); + core::cuda::DispatchCudaFunc( + input.Dtype(), + [&]() { + using random_t = std::conditional_t, double, float>; + DropoutForwardKernel<<cuda_stream()>>>( + static_cast(output.DataPtr()), static_cast(mask.DataPtr()), + static_cast(input.DataPtr()), n, static_cast(p), seed, subsequence); + }, + "CUDA dropout forward"); + CUDA_CHECK(cudaGetLastError()); +} + +void dropout_backward_cuda(Tensor &grad_input, const Tensor &grad_output, const Tensor &mask, double p) { + const Device device = grad_output.GetDevice(); + CHECK(device.IsCUDA()); + const int64_t n = grad_output.NumElements(); + if (n == 0) { + return; + } + core::DeviceGuard guard(device); + + const int blocks = static_cast((n + kThreadsPerBlock - 1) / kThreadsPerBlock); + const auto *stream = get_cuda_stream(device); + core::cuda::DispatchCudaFunc( + grad_output.Dtype(), + [&]() { + using random_t = std::conditional_t, double, float>; + DropoutBackwardKernel<<cuda_stream()>>>( + static_cast(grad_input.DataPtr()), static_cast(grad_output.DataPtr()), + static_cast(mask.DataPtr()), n, static_cast(p)); + }, + "CUDA dropout backward"); + CUDA_CHECK(cudaGetLastError()); +} + +} // namespace + +using infini_train::dropout_backward_stub; +using infini_train::dropout_forward_stub; + +REGISTER_DISPATCH(dropout_forward_stub, Device::DeviceType::kCUDA, &dropout_forward_cuda); +REGISTER_DISPATCH(dropout_backward_stub, Device::DeviceType::kCUDA, &dropout_backward_cuda); + +} // namespace infini_train::kernels::cuda diff --git a/infini_train/src/kernels/cuda/elementwise.cu b/infini_train/src/kernels/cuda/elementwise.cu index fc423b35..a4b6dae8 100644 --- a/infini_train/src/kernels/cuda/elementwise.cu +++ b/infini_train/src/kernels/cuda/elementwise.cu @@ -1,4 +1,5 @@ #include +#include #include diff --git a/infini_train/src/kernels/cuda/gather.cu b/infini_train/src/kernels/cuda/gather.cu index d8b0cffa..7b3862ac 100644 --- a/infini_train/src/kernels/cuda/gather.cu +++ b/infini_train/src/kernels/cuda/gather.cu @@ -1,3 +1,5 @@ +#include + #include "glog/logging.h" #include "infini_train/include/common/common.h" diff --git a/infini_train/src/kernels/cuda/no_op.cu b/infini_train/src/kernels/cuda/no_op.cu index ef2c9566..d26024b0 100644 --- a/infini_train/src/kernels/cuda/no_op.cu +++ b/infini_train/src/kernels/cuda/no_op.cu @@ -1,3 +1,5 @@ +#include + #include "glog/logging.h" #include "infini_train/include/dispatcher.h" diff --git a/infini_train/src/kernels/cuda/reduction.cu b/infini_train/src/kernels/cuda/reduction.cu index c56470e3..2c080edd 100644 --- a/infini_train/src/kernels/cuda/reduction.cu +++ b/infini_train/src/kernels/cuda/reduction.cu @@ -1,3 +1,5 @@ +#include + #include #include "infini_train/include/common/cuda/common_cuda.h" diff --git a/infini_train/src/nn/functional.cc b/infini_train/src/nn/functional.cc index c33e2368..4af04a2d 100644 --- a/infini_train/src/nn/functional.cc +++ b/infini_train/src/nn/functional.cc @@ -5,6 +5,7 @@ #include #include "infini_train/include/autograd/activations.h" +#include "infini_train/include/autograd/dropout.h" #include "infini_train/include/autograd/elementwise.h" #include "infini_train/include/autograd/reduction.h" #include "infini_train/include/autograd/softmax.h" @@ -26,6 +27,29 @@ std::shared_ptr Ones(const std::vector size) { return init::Ones(ones); } +std::shared_ptr Rand(const std::vector &size, DataType dtype, Device device, + std::optional generator, bool requires_grad) { + auto result = std::make_shared(size, dtype, device, requires_grad); + return init::Uniform(result, 0.0f, 1.0f, generator); +} + +std::shared_ptr Randn(const std::vector &size, DataType dtype, Device device, + std::optional generator, bool requires_grad) { + auto result = std::make_shared(size, dtype, device, requires_grad); + return init::Normal(result, 0.0f, 1.0f, generator); +} + +std::shared_ptr Dropout(const std::shared_ptr &input, double p, bool training, + std::optional generator) { + CHECK(input != nullptr); + CHECK_GE(p, 0.0) << "dropout probability has to be between 0 and 1, but got " << p; + CHECK_LE(p, 1.0) << "dropout probability has to be between 0 and 1, but got " << p; + if (!training || p == 0.0 || input->NumElements() == 0) { + return input; + } + return std::make_shared(p, std::move(generator))->Apply({input})[0]; +} + std::shared_ptr Reciprocal(const std::shared_ptr &input) { return input->Reciprocal(); } std::shared_ptr Sin(const std::shared_ptr &input) { return input->Sin(); } diff --git a/infini_train/src/nn/init.cc b/infini_train/src/nn/init.cc index 79b4b48b..8034a6b6 100644 --- a/infini_train/src/nn/init.cc +++ b/infini_train/src/nn/init.cc @@ -1,67 +1,24 @@ #include "infini_train/include/nn/init.h" -#include -#include #include -#include -#include +#include #include - -#ifdef USE_OMP -#include -#endif +#include #include "glog/logging.h" +#include "infini_train/include/core/runtime/distribution_kernels.h" #include "infini_train/include/core/runtime/device_guard.h" #include "infini_train/include/device.h" #include "infini_train/include/tensor.h" namespace infini_train::nn::init { -namespace { -constexpr int kRandomSeed = 42; - -// FIXME: RNG design is incomplete. -// -// Current implementation lacks: -// - unified Generator abstraction -// - global default generator and seed control -// - reproducible / clonable RNG state -// -// TODO: -// - introduce Generator interface and backend impl -// - add default generator management (per device) -// - refactor random ops to consume Generator -static std::mt19937 gen(kRandomSeed); -} // namespace std::shared_ptr Normal(const std::shared_ptr &tensor, float mean, float std, - std::optional generator) { - const int64_t num_elements = tensor->NumElements(); - std::vector buffer(num_elements); - -#ifdef USE_OMP -#pragma omp parallel - { - std::mt19937 local_gen(kRandomSeed + omp_get_thread_num()); - std::normal_distribution local_dis(mean, std); -#pragma omp for - for (int i = 0; i < buffer.size(); ++i) { - buffer[i] = generator ? local_dis(generator.value()) : local_dis(local_gen); - } - } -#else - std::normal_distribution dis(mean, std); - std::generate(buffer.begin(), buffer.end(), [&]() { return generator ? dis(generator.value()) : dis(gen); }); -#endif - + std::optional generator) { auto device = tensor->GetDevice(); core::DeviceGuard guard(device); - auto impl = core::GetDeviceGuardImpl(device.type()); - - impl->MemcpyAsync(tensor->DataPtr(), buffer.data(), num_elements * sizeof(float), - device.type() == Device::DeviceType::kCPU ? core::MemcpyKind::kD2D : core::MemcpyKind::kH2D, - impl->GetStream(device)); + normal_kernel(*tensor, mean, std, generator); return tensor; } @@ -113,7 +70,7 @@ float CalculateGain(NonLinearityType nonlinearity, std::optional param = } // namespace std::shared_ptr KaimingUniform(const std::shared_ptr &tensor, float a, KaimingMode mode, - NonLinearityType nonlinearity, std::optional generator) { + NonLinearityType nonlinearity, std::optional generator) { for (const auto dim : tensor->Dims()) { if (dim == 0) { LOG(WARNING) << "Initializing zero-element tensors is a no-op"; @@ -128,33 +85,10 @@ std::shared_ptr KaimingUniform(const std::shared_ptr &tensor, fl } std::shared_ptr Uniform(const std::shared_ptr &tensor, float a, float b, - std::optional generator) { - const int64_t num_elements = tensor->NumElements(); - std::vector buffer(num_elements); - -#ifdef USE_OMP -#pragma omp parallel - { - std::mt19937 local_gen(kRandomSeed + omp_get_thread_num()); - std::uniform_real_distribution local_dis(a, b); -#pragma omp for - for (int i = 0; i < buffer.size(); ++i) { - buffer[i] = generator ? local_dis(generator.value()) : local_dis(local_gen); - } - } -#else - std::uniform_real_distribution dis(a, b); - std::generate(buffer.begin(), buffer.end(), [&]() { return generator ? dis(generator.value()) : dis(gen); }); -#endif - + std::optional generator) { auto device = tensor->GetDevice(); - core::DeviceGuard guard(device); - auto impl = core::GetDeviceGuardImpl(device.type()); - - impl->MemcpyAsync(tensor->DataPtr(), buffer.data(), num_elements * sizeof(float), - device.type() == Device::DeviceType::kCPU ? core::MemcpyKind::kD2D : core::MemcpyKind::kH2D, - impl->GetStream(device)); + uniform_kernel(*tensor, a, b, generator); return tensor; } diff --git a/infini_train/src/nn/parallel/ddp/param_and_grad_buffer.cc b/infini_train/src/nn/parallel/ddp/param_and_grad_buffer.cc index ab3a8002..b5afd48b 100644 --- a/infini_train/src/nn/parallel/ddp/param_and_grad_buffer.cc +++ b/infini_train/src/nn/parallel/ddp/param_and_grad_buffer.cc @@ -3,6 +3,7 @@ #include #include #include +#include #include "glog/logging.h" diff --git a/infini_train/src/tensor.cc b/infini_train/src/tensor.cc index 18ca3d22..9c6a4d66 100644 --- a/infini_train/src/tensor.cc +++ b/infini_train/src/tensor.cc @@ -3,6 +3,7 @@ #include #include #include +#include #include #include @@ -128,10 +129,11 @@ Eigen::Map> Tensor::Eig Tensor Tensor::To(Device device) { const auto buffer_device = buffer_->GetDevice(); if (device == buffer_device) { - auto new_tensor = Tensor(*this, offset_, dims_); + auto new_tensor = Tensor(*this, 0, dims_); if (grad_) { - new_tensor.grad_ = std::make_unique(*grad_.get(), grad_->offset_, grad_->dims_); + new_tensor.grad_ = std::make_unique(*grad_.get(), 0, grad_->dims_); } + new_tensor.requires_grad_ = requires_grad_; return new_tensor; } @@ -174,10 +176,11 @@ Tensor Tensor::To(Device device) { Tensor Tensor::To(DataType dtype) { if (dtype == dtype_) { - auto new_tensor = Tensor(*this, offset_, dims_); + auto new_tensor = Tensor(*this, 0, dims_); if (grad_) { - new_tensor.grad_ = std::make_unique(*grad_.get(), grad_->offset_, grad_->dims_); + new_tensor.grad_ = std::make_unique(*grad_.get(), 0, grad_->dims_); } + new_tensor.requires_grad_ = requires_grad_; return new_tensor; } @@ -471,7 +474,7 @@ std::shared_ptr Tensor::Outer(const std::shared_ptr &other) { } // distribution -std::shared_ptr Tensor::Uniform(float from, float to, std::optional generator) { +std::shared_ptr Tensor::Uniform(float from, float to, std::optional generator) { return nn::init::Uniform(shared_from_this(), from, to, generator); } diff --git a/tests/CMakeLists.txt b/tests/CMakeLists.txt index 96776585..9e5eb322 100644 --- a/tests/CMakeLists.txt +++ b/tests/CMakeLists.txt @@ -28,5 +28,8 @@ add_subdirectory(dtype) # Transformer architecture tests add_subdirectory(transformer) +# Generator tests +add_subdirectory(generator) + # Checkpoint tests add_subdirectory(checkpoint) diff --git a/tests/autograd/test_autograd_elementwise_backward.cc b/tests/autograd/test_autograd_elementwise_backward.cc index f7eb0d5f..2b72ec2f 100644 --- a/tests/autograd/test_autograd_elementwise_backward.cc +++ b/tests/autograd/test_autograd_elementwise_backward.cc @@ -4,6 +4,7 @@ #include "gtest/gtest.h" #include "infini_train/include/autograd/elementwise.h" +#include "infini_train/include/core/runtime/device_guard.h" #include "infini_train/include/nn/parallel/global.h" #include "infini_train/include/tensor.h" @@ -133,7 +134,16 @@ TEST_P(AutogradElementwiseBackwardTest, ExpBackward) { auto grad = std::make_shared(std::vector{2, 3}, DataType::kFLOAT32, GetDevice(), true); grad->Fill(1.0f); auto grad_inputs = exp_fn->Backward({grad}); - EXPECT_EQ(grad_inputs.size(), 1); + ASSERT_EQ(grad_inputs.size(), 1); + + auto device = grad_inputs[0]->GetDevice(); + core::GetDeviceGuardImpl(device.type())->SynchronizeDevice(device); + Tensor grad_input_cpu = grad_inputs[0]->To(Device(Device::DeviceType::kCPU, 0)); + core::GetDeviceGuardImpl(device.type())->SynchronizeDevice(device); + const auto *grad_input = static_cast(grad_input_cpu.DataPtr()); + for (size_t i = 0; i < grad_input_cpu.NumElements(); ++i) { + EXPECT_FLOAT_EQ(grad_input[i], std::exp(1.0f)); + } } TEST_P(AutogradElementwiseBackwardTest, LogBackward) { diff --git a/tests/generator/CMakeLists.txt b/tests/generator/CMakeLists.txt new file mode 100644 index 00000000..05063fa6 --- /dev/null +++ b/tests/generator/CMakeLists.txt @@ -0,0 +1,19 @@ +# Shared Generator behavior tests. The existing suite helper builds separate +# CPU and CUDA executables and filters the parameterized instances by label. +infini_train_add_test_suite(test_generator + SOURCES test_generator.cc +) + +# Device-independent handle and validation tests are not parameterized, so +# keep them in a separate executable that is not subject to CPU/* filters. +infini_train_add_test(test_generator_interface + SOURCES test_generator_interface.cc + LABELS cpu +) + +if(USE_CUDA) + infini_train_add_test(test_generator_cuda_only + SOURCES cuda_only/test_cuda_generator.cc + LABELS cuda + ) +endif() diff --git a/tests/generator/cuda_only/test_cuda_generator.cc b/tests/generator/cuda_only/test_cuda_generator.cc new file mode 100644 index 00000000..aea83ff7 --- /dev/null +++ b/tests/generator/cuda_only/test_cuda_generator.cc @@ -0,0 +1,56 @@ +#include +#include +#include + +#include "gtest/gtest.h" + +#include "infini_train/include/generator.h" +#include "infini_train/include/nn/init.h" +#include "infini_train/include/tensor.h" +#include "infini_train/src/core/runtime/cuda/cuda_generator_impl.h" +#include "tests/common/test_utils.h" + +namespace infini_train::test { +namespace { + +constexpr uint64_t kSeed = 0x12345678ULL; + +std::vector StateBytes(const Generator &generator) { + auto state = generator.get_state(); + const auto *data = static_cast(state->DataPtr()); + return std::vector(data, data + state->SizeInBytes()); +} + +TEST(CudaGeneratorTest, StateIsAnOpaqueUint8CpuTensor) { + Generator generator = core::cuda::createCUDAGenerator(0, kSeed); + auto state = generator.get_state(); + + ASSERT_TRUE(state); + EXPECT_TRUE(state->GetDevice().IsCPU()); + EXPECT_EQ(state->Dtype(), DataType::kUINT8); + EXPECT_EQ(state->SizeInBytes(), sizeof(uint64_t) * 2); +} + +TEST(CudaGeneratorTest, MissingGeneratorUsesMatchingDeviceDefaultGenerator) { + REQUIRE_MIN_DEVICES(2); + manual_seed(kSeed); + const Generator &device_zero = core::cuda::getDefaultCUDAGenerator(0); + const Generator &device_one = core::cuda::getDefaultCUDAGenerator(1); + + EXPECT_NE(device_zero, device_one); + EXPECT_EQ(device_zero.device().index(), 0); + EXPECT_EQ(device_one.device().index(), 1); + + const std::vector device_zero_before = StateBytes(device_zero); + const std::vector device_one_before = StateBytes(device_one); + auto tensor = std::make_shared( + std::vector{1024}, DataType::kFLOAT32, + Device(Device::DeviceType::kCUDA, 1)); + nn::init::Uniform(tensor); + + EXPECT_EQ(device_zero_before, StateBytes(device_zero)); + EXPECT_NE(device_one_before, StateBytes(device_one)); +} + +} // namespace +} // namespace infini_train::test diff --git a/tests/generator/test_generator.cc b/tests/generator/test_generator.cc new file mode 100644 index 00000000..1c642a94 --- /dev/null +++ b/tests/generator/test_generator.cc @@ -0,0 +1,282 @@ +#include +#include +#include +#include +#include + +#include "gtest/gtest.h" + +#include "infini_train/include/core/runtime/device_guard.h" +#include "infini_train/include/generator.h" +#include "infini_train/include/nn/functional.h" +#include "infini_train/include/nn/init.h" +#include "infini_train/include/tensor.h" +#include "infini_train/src/core/runtime/cpu/cpu_generator_impl.h" +#if defined(USE_CUDA) +#include "infini_train/src/core/runtime/cuda/cuda_generator_impl.h" +#endif +#include "tests/common/test_utils.h" + +namespace infini_train::test { +namespace { + +constexpr uint64_t kSeed = 0x12345678ULL; +constexpr int64_t kElements = 4096; + +Generator CreateGenerator(Device device, uint64_t seed) { + if (device.IsCPU()) { + return core::cpu::createCPUGenerator(seed); + } +#if defined(USE_CUDA) + return core::cuda::createCUDAGenerator(device.index(), seed); +#else + (void)seed; + throw std::runtime_error("CUDA support is disabled"); +#endif +} + +const Generator &GetDefaultGenerator(Device device) { + if (device.IsCPU()) { + return core::cpu::getDefaultCPUGenerator(); + } +#if defined(USE_CUDA) + return core::cuda::getDefaultCUDAGenerator(device.index()); +#else + throw std::runtime_error("CUDA support is disabled"); +#endif +} + +void Synchronize(Device device) { + core::GetDeviceGuardImpl(device.type())->SynchronizeDevice(device); +} + +std::vector ToHostVector(const std::shared_ptr &tensor) { + const Device device = tensor->GetDevice(); + Synchronize(device); + Tensor cpu = tensor->To(Device(Device::DeviceType::kCPU, 0)); + Synchronize(device); + const auto *data = static_cast(cpu.DataPtr()); + return std::vector(data, data + cpu.NumElements()); +} + +std::vector StateBytes(const Generator &generator) { + auto state = generator.get_state(); + const auto *data = static_cast(state->DataPtr()); + return std::vector(data, data + state->SizeInBytes()); +} + +std::shared_ptr MakeTensor(Device device) { + return std::make_shared(std::vector{kElements}, DataType::kFLOAT32, device); +} + +class GeneratorTest : public InfiniTrainTest {}; + +TEST_P(GeneratorTest, SupportsDeviceAndSeedInterface) { + const Device device = GetDevice(); + Generator generator = CreateGenerator(device, kSeed); + + EXPECT_TRUE(generator.defined()); + EXPECT_EQ(generator.device(), device); + EXPECT_EQ(generator.current_seed(), kSeed); + generator.set_current_seed(kSeed + 1); + EXPECT_EQ(generator.current_seed(), kSeed + 1); +} + +TEST_P(GeneratorTest, UniformIsReproducibleForSameSeed) { + Generator first_generator = CreateGenerator(GetDevice(), kSeed); + Generator second_generator = CreateGenerator(GetDevice(), kSeed); + auto first = MakeTensor(GetDevice()); + auto second = MakeTensor(GetDevice()); + + nn::init::Uniform(first, -3.0f, 7.0f, first_generator); + nn::init::Uniform(second, -3.0f, 7.0f, second_generator); + + EXPECT_EQ(ToHostVector(first), ToHostVector(second)); +} + +TEST_P(GeneratorTest, NormalIsReproducibleForSameSeed) { + Generator first_generator = CreateGenerator(GetDevice(), kSeed); + Generator second_generator = CreateGenerator(GetDevice(), kSeed); + auto first = MakeTensor(GetDevice()); + auto second = MakeTensor(GetDevice()); + + nn::init::Normal(first, 2.0f, 0.5f, first_generator); + nn::init::Normal(second, 2.0f, 0.5f, second_generator); + + EXPECT_EQ(ToHostVector(first), ToHostVector(second)); +} + +TEST_P(GeneratorTest, RandAndRandnSupportExplicitAndDefaultGenerators) { + const Device device = GetDevice(); + Generator first_generator = CreateGenerator(device, kSeed); + Generator second_generator = CreateGenerator(device, kSeed); + + auto first = nn::function::Rand({17, 19}, DataType::kFLOAT32, device, first_generator); + auto second = nn::function::Rand({17, 19}, DataType::kFLOAT32, device, second_generator); + EXPECT_EQ(ToHostVector(first), ToHostVector(second)); + + manual_seed(kSeed); + auto first_default = nn::function::Randn({17, 19}, DataType::kFLOAT32, device); + manual_seed(kSeed); + auto second_default = nn::function::Randn({17, 19}, DataType::kFLOAT32, device); + EXPECT_EQ(ToHostVector(first_default), ToHostVector(second_default)); +} + +TEST_P(GeneratorTest, DropoutMaskIsReproducibleForSameSeed) { + const Device device = GetDevice(); + auto input = MakeTensor(device); + input->Fill(1.0f); + + manual_seed(kSeed); + auto first = nn::function::Dropout(input, 0.25, true); + manual_seed(kSeed); + auto second = nn::function::Dropout(input, 0.25, true); + manual_seed(kSeed + 1); + auto different_seed = nn::function::Dropout(input, 0.25, true); + + EXPECT_EQ(ToHostVector(first), ToHostVector(second)); + EXPECT_NE(ToHostVector(first), ToHostVector(different_seed)); +} + +TEST_P(GeneratorTest, ConsecutiveCallsAdvanceSequence) { + Generator generator = CreateGenerator(GetDevice(), kSeed); + auto first = MakeTensor(GetDevice()); + auto second = MakeTensor(GetDevice()); + + nn::init::Uniform(first, 0.0f, 1.0f, generator); + nn::init::Uniform(second, 0.0f, 1.0f, generator); + + EXPECT_NE(ToHostVector(first), ToHostVector(second)); +} + +TEST_P(GeneratorTest, DifferentSeedsProduceDifferentSequences) { + Generator first_generator = CreateGenerator(GetDevice(), kSeed); + Generator second_generator = CreateGenerator(GetDevice(), kSeed + 1); + auto first = MakeTensor(GetDevice()); + auto second = MakeTensor(GetDevice()); + + nn::init::Uniform(first, 0.0f, 1.0f, first_generator); + nn::init::Uniform(second, 0.0f, 1.0f, second_generator); + + EXPECT_NE(ToHostVector(first), ToHostVector(second)); +} + +TEST_P(GeneratorTest, SameSeedAndCallOrderReproduceResults) { + const Device device = GetDevice(); + auto first_uniform = MakeTensor(device); + auto second_uniform = MakeTensor(device); + auto first_normal = MakeTensor(device); + auto second_normal = MakeTensor(device); + auto input = MakeTensor(device); + input->Fill(1.0f); + + manual_seed(kSeed); + nn::init::Uniform(first_uniform); + auto first_dropout = nn::function::Dropout(input, 0.5, true); + nn::init::Normal(first_normal); + + manual_seed(kSeed); + nn::init::Uniform(second_uniform); + auto second_dropout = nn::function::Dropout(input, 0.5, true); + nn::init::Normal(second_normal); + + EXPECT_EQ(ToHostVector(first_uniform), ToHostVector(second_uniform)); + EXPECT_EQ(ToHostVector(first_dropout), ToHostVector(second_dropout)); + EXPECT_EQ(ToHostVector(first_normal), ToHostVector(second_normal)); +} + +TEST_P(GeneratorTest, StateRestoreReplaysUniformSequence) { + Generator generator = CreateGenerator(GetDevice(), kSeed); + auto prefix = MakeTensor(GetDevice()); + auto expected = MakeTensor(GetDevice()); + auto actual = MakeTensor(GetDevice()); + + nn::init::Uniform(prefix, 0.0f, 1.0f, generator); + auto state = generator.get_state(); + nn::init::Uniform(expected, 0.0f, 1.0f, generator); + generator.set_state(*state); + nn::init::Uniform(actual, 0.0f, 1.0f, generator); + + EXPECT_EQ(ToHostVector(expected), ToHostVector(actual)); +} + +TEST_P(GeneratorTest, StateRestoreReplaysNormalSequence) { + Generator generator = CreateGenerator(GetDevice(), kSeed); + auto prefix = MakeTensor(GetDevice()); + auto expected = MakeTensor(GetDevice()); + auto actual = MakeTensor(GetDevice()); + + nn::init::Normal(prefix, 0.0f, 1.0f, generator); + auto state = generator.get_state(); + nn::init::Normal(expected, 0.0f, 1.0f, generator); + generator.set_state(*state); + nn::init::Normal(actual, 0.0f, 1.0f, generator); + + EXPECT_EQ(ToHostVector(expected), ToHostVector(actual)); +} + +TEST_P(GeneratorTest, ExplicitGeneratorDoesNotAdvanceDefaultGenerator) { + const Device device = GetDevice(); + manual_seed(kSeed); + const Generator &default_generator = GetDefaultGenerator(device); + const std::vector before = StateBytes(default_generator); + Generator explicit_generator = CreateGenerator(device, kSeed + 1); + + auto tensor = MakeTensor(device); + nn::init::Uniform(tensor, 0.0f, 1.0f, explicit_generator); + + EXPECT_EQ(before, StateBytes(default_generator)); +} + +TEST_P(GeneratorTest, MissingGeneratorUsesAndAdvancesDefaultGenerator) { + const Device device = GetDevice(); + manual_seed(kSeed); + const Generator &default_generator = GetDefaultGenerator(device); + const std::vector before = StateBytes(default_generator); + + auto tensor = MakeTensor(device); + nn::init::Uniform(tensor); + + EXPECT_NE(before, StateBytes(default_generator)); +} + +TEST_P(GeneratorTest, RepeatedDefaultGeneratorLookupSharesState) { + const Device device = GetDevice(); + const Generator &first = GetDefaultGenerator(device); + const Generator &second = GetDefaultGenerator(device); + + EXPECT_EQ(first, second); +} + +TEST_P(GeneratorTest, GlobalManualSeedReproducesDefaultGeneratorResults) { + const Device device = GetDevice(); + auto first = MakeTensor(device); + auto second = MakeTensor(device); + + manual_seed(kSeed); + EXPECT_EQ(GetDefaultGenerator(device).current_seed(), kSeed); + nn::init::Uniform(first); + manual_seed(kSeed); + nn::init::Uniform(second); + + EXPECT_EQ(ToHostVector(first), ToHostVector(second)); +} + +TEST_P(GeneratorTest, KaimingUniformUsesExplicitGenerator) { + Generator first_generator = CreateGenerator(GetDevice(), kSeed); + Generator second_generator = CreateGenerator(GetDevice(), kSeed); + auto first = std::make_shared(std::vector{64, 32}, DataType::kFLOAT32, GetDevice()); + auto second = std::make_shared(std::vector{64, 32}, DataType::kFLOAT32, GetDevice()); + + nn::init::KaimingUniform(first, 0.0f, nn::init::KaimingMode::kFanIn, + nn::init::NonLinearityType::kReLU, first_generator); + nn::init::KaimingUniform(second, 0.0f, nn::init::KaimingMode::kFanIn, + nn::init::NonLinearityType::kReLU, second_generator); + + EXPECT_EQ(ToHostVector(first), ToHostVector(second)); +} + +INFINI_TRAIN_REGISTER_TEST(GeneratorTest); + +} // namespace +} // namespace infini_train::test diff --git a/tests/generator/test_generator_interface.cc b/tests/generator/test_generator_interface.cc new file mode 100644 index 00000000..eb8a93a8 --- /dev/null +++ b/tests/generator/test_generator_interface.cc @@ -0,0 +1,37 @@ +#include +#include +#include + +#include "gtest/gtest.h" + +#include "infini_train/include/generator.h" +#include "infini_train/include/tensor.h" +#include "infini_train/src/core/runtime/cpu/cpu_generator_impl.h" + +namespace infini_train::test { +namespace { + +constexpr uint64_t kSeed = 0x12345678ULL; + +TEST(GeneratorInterfaceTest, CpuStateIsAnOpaqueUint8CpuTensor) { + Generator generator = core::cpu::createCPUGenerator(kSeed); + auto state = generator.get_state(); + + ASSERT_TRUE(state); + EXPECT_TRUE(state->GetDevice().IsCPU()); + EXPECT_EQ(state->Dtype(), DataType::kUINT8); + EXPECT_GT(state->SizeInBytes(), 0u); +} + +TEST(GeneratorInterfaceTest, CpuStateRoundTripRestoresSeed) { + Generator generator = core::cpu::createCPUGenerator(kSeed); + auto state = generator.get_state(); + generator.set_current_seed(kSeed + 1); + + generator.set_state(*state); + + EXPECT_EQ(generator.current_seed(), kSeed); +} + +} // namespace +} // namespace infini_train::test