From df7308cdaf90699e81d4c41336e732a64a2d8943 Mon Sep 17 00:00:00 2001 From: chaos Date: Sun, 26 Jul 2026 23:44:46 +0800 Subject: [PATCH 1/8] feat(generator): add CPU and CUDA generator infrastructure --- infini_train/include/generator.h | 130 ++++++++++++ .../core/runtime/cpu/cpu_generator_impl.cc | 190 ++++++++++++++++++ .../src/core/runtime/cpu/cpu_generator_impl.h | 52 +++++ .../core/runtime/cuda/cuda_generator_impl.cc | 140 +++++++++++++ .../core/runtime/cuda/cuda_generator_impl.h | 37 ++++ infini_train/src/generator.cc | 78 +++++++ 6 files changed, 627 insertions(+) create mode 100644 infini_train/include/generator.h create mode 100644 infini_train/src/core/runtime/cpu/cpu_generator_impl.cc create mode 100644 infini_train/src/core/runtime/cpu/cpu_generator_impl.h create mode 100644 infini_train/src/core/runtime/cuda/cuda_generator_impl.cc create mode 100644 infini_train/src/core/runtime/cuda/cuda_generator_impl.h create mode 100644 infini_train/src/generator.cc diff --git a/infini_train/include/generator.h b/infini_train/include/generator.h new file mode 100644 index 000000000..ee107c71e --- /dev/null +++ b/infini_train/include/generator.h @@ -0,0 +1,130 @@ +#pragma once + +#include +#include +#include +#include +#include +#include + +#include "infini_train/include/device.h" + +namespace infini_train { + +class Tensor; + +namespace detail { + +// Validates the common Tensor contract for serialized RNG states. +void check_rng_state(const Tensor &state); + +} // namespace detail + +// Base interface for device-specific random number generators. +class GeneratorImpl { +public: + explicit GeneratorImpl(Device device) : device_(device) {} + virtual ~GeneratorImpl() = default; + + 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; + + std::shared_ptr clone() const { return std::shared_ptr(clone_impl()); } + + Device device() const { return device_; } + + // Callers must lock this mutex when an operation spans multiple generator calls. + std::mutex mutex_; + +protected: + Device device_; + + virtual GeneratorImpl *clone_impl() const = 0; +}; + +// A lightweight handle with shared-copy semantics. Use clone() for an independent state. +class Generator { +public: + static constexpr uint64_t kDefaultSeed = 67280421310721; + + Generator() = default; + + explicit Generator(std::shared_ptr 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(); } + + void set_state(const Tensor &state); + std::shared_ptr get_state() const; + + Device device() const { return impl_->device(); } + + Generator clone() const { return Generator(impl_->clone()); } + + std::mutex &mutex() const { return impl_->mutex_; } + + // Prefer check_generator(); this unchecked accessor assumes a matching backend. + 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_; +}; + +// Internal factory for backend implementations. +template Generator make_generator(Args &&...args) { + return Generator(std::make_shared(std::forward(args)...)); +} + +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; +} + +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); +} + +// Creates a generator for the requested device without exposing its backend implementation. +Generator CreateGenerator(const Device &device, uint64_t seed = Generator::kDefaultSeed); + +// Returns the lazily initialized default generator for the requested device. +const Generator &GetDefaultGenerator(const Device &device); + +// Reset the default generators for all enabled devices. +void manual_seed(uint64_t seed); + +} // 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 000000000..cb49fb3b8 --- /dev/null +++ b/infini_train/src/core/runtime/cpu/cpu_generator_impl.cc @@ -0,0 +1,190 @@ +#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 { + +// Backend tag used to reject states from other generator implementations. +constexpr char kCPUStateMagic[] = "ITRNGCPU"; +constexpr size_t kStateMagicSize = sizeof(kCPUStateMagic) - 1; + +constexpr size_t kStateFooterSize + = sizeof(uint64_t) + sizeof(uint8_t) + sizeof(float) + sizeof(uint8_t) + sizeof(double); + +} // namespace + +// std::mt19937(seed) truncates to 32 bits, so seed both halves explicitly. +static std::mt19937 make_seeded_engine(uint64_t seed) { + std::seed_seq seq{static_cast(seed), static_cast(seed >> 32)}; + std::mt19937 engine; + engine.seed(seq); + return engine; +} + +static uint64_t getNonDeterministicRandom() { + std::random_device rd; + uint64_t val = (static_cast(rd()) << 32) | rd(); + return val; +} + +CPUGeneratorImpl::CPUGeneratorImpl(uint64_t seed) + : GeneratorImpl(Device(Device::DeviceType::kCPU, 0)), engine_(make_seeded_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_ = make_seeded_engine(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; +} + +// State layout: magic, serialized engine, seed, and cached float/double normal samples. +// The std::mt19937 stream format is only portable across compatible standard-library builds. + +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, kStateMagicSize + kStateFooterSize) << "CPU generator state is too small"; + + const uint8_t *data = static_cast(state.DataPtr()); + CHECK_EQ(std::memcmp(data, kCPUStateMagic, kStateMagicSize), 0) + << "Invalid RNG state: not a CPU generator state (backend magic mismatch)"; + + const size_t engine_size = data_size - kStateMagicSize - kStateFooterSize; + std::string engine_str(reinterpret_cast(data + kStateMagicSize), 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 = kStateMagicSize + 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 { + std::ostringstream oss; + oss << engine_; + std::string engine_str = oss.str(); + + const size_t engine_size = engine_str.size(); + const size_t total_size = kStateMagicSize + 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, kCPUStateMagic, kStateMagicSize); + offset += kStateMagicSize; + + std::memcpy(data + offset, engine_str.data(), engine_size); + offset += engine_size; + + std::memcpy(data + offset, &seed_, sizeof(seed_)); + offset += sizeof(seed_); + + 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); + + 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; +} + +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; +} + +std::shared_ptr CPUGeneratorImpl::clone() const { + return std::shared_ptr(clone_impl()); +} + +CPUGeneratorImpl *CPUGeneratorImpl::clone_impl() const { + auto *clone = new CPUGeneratorImpl(seed_); + clone->set_engine(engine_); + clone->set_next_float_normal_sample(next_float_normal_sample_); + clone->set_next_double_normal_sample(next_double_normal_sample_); + return clone; +} + +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 + +namespace infini_train::core::cpu { + +const Generator &getDefaultCPUGenerator() { + 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 000000000..22a766ec7 --- /dev/null +++ b/infini_train/src/core/runtime/cpu/cpu_generator_impl.h @@ -0,0 +1,52 @@ +#pragma once + +#include +#include +#include +#include + +#include "infini_train/include/generator.h" + +namespace infini_train::core::cpu { + +// CPU generator backed by std::mt19937 with cached Box-Muller samples. +class CPUGeneratorImpl final : public GeneratorImpl { +public: + explicit CPUGeneratorImpl(uint64_t seed = Generator::kDefaultSeed); + ~CPUGeneratorImpl() 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; + + std::shared_ptr clone() const; + + static Device::DeviceType device_type(); + + uint32_t random(); + uint64_t random64(); + + 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; + + 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_; +}; + +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 000000000..0b44e3011 --- /dev/null +++ b/infini_train/src/core/runtime/cuda/cuda_generator_impl.cc @@ -0,0 +1,140 @@ +#include "infini_train/src/core/runtime/cuda/cuda_generator_impl.h" + +#include +#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 { + +// Backend tag used to reject states from other generator implementations. +constexpr char kCUDAStateMagic[] = "ITRNGCUD"; +constexpr size_t kStateMagicSize = sizeof(kCUDAStateMagic) - 1; + +// [magic: 8B] + [seed_: 8B] + [next_philox_subsequence_: 8B] +constexpr size_t kStateSize = kStateMagicSize + 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()); + CHECK_EQ(std::memcmp(data, kCUDAStateMagic, kStateMagicSize), 0) + << "Invalid RNG state: not a CUDA generator state (backend magic mismatch)"; + std::memcpy(&seed_, data + kStateMagicSize, sizeof(seed_)); + std::memcpy(&next_philox_subsequence_, data + kStateMagicSize + 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, kCUDAStateMagic, kStateMagicSize); + std::memcpy(data + kStateMagicSize, &seed_, sizeof(seed_)); + std::memcpy(data + kStateMagicSize + 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) { + CHECK_LE(increment, std::numeric_limits::max() - next_philox_subsequence_) + << "Philox subsequence counter overflow"; + const uint64_t subsequence = next_philox_subsequence_; + next_philox_subsequence_ += increment; + return subsequence; +} + +CUDAGeneratorImpl *CUDAGeneratorImpl::clone_impl() const { + auto *clone = new CUDAGeneratorImpl(device().index(), seed_); + clone->next_philox_subsequence_ = next_philox_subsequence_; + return clone; +} + +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 000000000..5ffdbdc82 --- /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 000000000..e29ff29e9 --- /dev/null +++ b/infini_train/src/generator.cc @@ -0,0 +1,78 @@ +#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::Generator(std::shared_ptr impl) : impl_(std::move(impl)) { + CHECK(impl_) << "GeneratorImpl with nullptr is not supported"; +} + +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 + +Generator CreateGenerator(const Device &device, uint64_t seed) { + if (device.IsCPU()) { + return core::cpu::createCPUGenerator(seed); + } + +#ifdef USE_CUDA + if (device.IsCUDA()) { + return core::cuda::createCUDAGenerator(device.index(), seed); + } +#else + if (device.IsCUDA()) { + throw std::invalid_argument("CUDA generator requested but CUDA support is not enabled"); + } +#endif + + throw std::invalid_argument("Generator can only be created for CPU or CUDA devices"); +} + +const Generator &GetDefaultGenerator(const Device &device) { + if (device.IsCPU()) { + return core::cpu::getDefaultCPUGenerator(); + } + +#ifdef USE_CUDA + if (device.IsCUDA()) { + return core::cuda::getDefaultCUDAGenerator(device.index()); + } +#else + if (device.IsCUDA()) { + throw std::invalid_argument("CUDA default generator requested but CUDA support is not enabled"); + } +#endif + + throw std::invalid_argument("Default generator can only be requested for CPU or CUDA devices"); +} + +void manual_seed(uint64_t seed) { + core::cpu::manual_seed(seed); + +#ifdef USE_CUDA + core::cuda::manual_seed_all(seed); +#endif +} + +} // namespace infini_train From 53675989ff58e6937d3a883dad16be881532d7e9 Mon Sep 17 00:00:00 2001 From: chaos Date: Sun, 26 Jul 2026 23:52:13 +0800 Subject: [PATCH 2/8] feat(random): add distributions and generator-aware initialization --- .../include/common/cpu/distributions_helper.h | 117 ++++++++++++++ infini_train/include/nn/init.h | 8 +- infini_train/include/tensor.h | 6 +- infini_train/src/kernels/cpu/distribution.cc | 80 +++++++++ infini_train/src/kernels/cuda/distribution.cu | 152 ++++++++++++++++++ infini_train/src/nn/init.cc | 135 +++++++--------- infini_train/src/tensor.cc | 2 +- 7 files changed, 421 insertions(+), 79 deletions(-) create mode 100644 infini_train/include/common/cpu/distributions_helper.h create mode 100644 infini_train/src/kernels/cpu/distribution.cc create mode 100644 infini_train/src/kernels/cuda/distribution.cu diff --git a/infini_train/include/common/cpu/distributions_helper.h b/infini_train/include/common/cpu/distributions_helper.h new file mode 100644 index 000000000..9cc9f7a62 --- /dev/null +++ b/infini_train/include/common/cpu/distributions_helper.h @@ -0,0 +1,117 @@ +#pragma once + +// Host-side uniform and normal distributions for generators exposing random() +// and random64(). Box-Muller's second sample is cached when supported by the generator. + +#include +#include +#include +#include +#include + +#include "glog/logging.h" + +namespace infini_train::common::cpu { + +template struct uniform_real_distribution { + uniform_real_distribution(T from, T to) : from_(from), to_(to) { + CHECK_LE(from, to); + CHECK_LE(to - from, std::numeric_limits::max()); + } + + 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_; + + 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_; + } +}; + +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; +} + +// Fallback: RNG without cache support never has a cached sample. +template bool maybe_get_next_normal_sample(RNG * /*generator*/, void * /*ret*/) { return false; } + +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); +} + +// Fallback: RNG without cache support discards the second sample. +template void maybe_set_next_normal_sample(RNG * /*generator*/, const void * /*cache*/) {} + +template struct normal_distribution { + normal_distribution(T mean, T stdv) : mean_(mean), stdv_(stdv) { CHECK_GE(stdv, static_cast(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_; + } + + uniform_real_distribution uniform(static_cast(0), static_cast(1)); + const T u1 = uniform(generator); + const T u2 = uniform(generator); + + 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 infini_train::common::cpu diff --git a/infini_train/include/nn/init.h b/infini_train/include/nn/init.h index fc6effec9..05b95b2ca 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 dcfd8927f..05b292630 100644 --- a/infini_train/include/tensor.h +++ b/infini_train/include/tensor.h @@ -5,6 +5,7 @@ #include #include #include + #include #include "Eigen/Dense" @@ -12,6 +13,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 { @@ -74,6 +76,8 @@ class Tensor : public std::enable_shared_from_this { void *DataPtr(); const void *DataPtr() const; + bool defined() const { return buffer_ != nullptr; } + size_t SizeInBytes() const; const std::vector &Dims() const; @@ -151,7 +155,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/kernels/cpu/distribution.cc b/infini_train/src/kernels/cpu/distribution.cc new file mode 100644 index 000000000..947614002 --- /dev/null +++ b/infini_train/src/kernels/cpu/distribution.cc @@ -0,0 +1,80 @@ +#include +#include +#include +#include +#include + +#include "glog/logging.h" + +#include "infini_train/include/common/cpu/distributions_helper.h" +#include "infini_train/include/datatype.h" +#include "infini_train/include/dispatcher.h" +#include "infini_train/include/generator.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 UniformImpl(Tensor &tensor, double from, double to, core::cpu::CPUGeneratorImpl *generator) { + auto *buf = static_cast(tensor.DataPtr()); + common::cpu::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)); + // [from, to) is half-open: a sample landing exactly on `to` is mapped back to `from`. + buf[i] = static_cast(value) == to_value ? from_value : value; + } +} + +template +void NormalImpl(Tensor &tensor, double mean, double std, core::cpu::CPUGeneratorImpl *generator) { + auto *buf = static_cast(tensor.DataPtr()); + common::cpu::normal_distribution dist(static_cast(mean), static_cast(std)); + for (int64_t i = 0; i < tensor.NumElements(); ++i) { buf[i] = static_cast(dist(generator)); } +} + +} // namespace + +void Uniform(const std::shared_ptr &tensor, double from, double to, const std::optional &gen) { + CHECK(tensor->GetDevice().IsCPU()); + auto *cpu_generator + = get_generator_or_default(gen, core::cpu::getDefaultCPUGenerator()); + + std::lock_guard lock(cpu_generator->mutex_); + core::cpu::DispatchCpuFunc( + tensor->Dtype(), + [&]() { + using random_t = std::conditional_t, double, float>; + UniformImpl(*tensor, from, to, cpu_generator); + }, + "CPU uniform"); +} + +void Normal(const std::shared_ptr &tensor, double mean, double std, const std::optional &gen) { + CHECK(tensor->GetDevice().IsCPU()); + auto *cpu_generator + = get_generator_or_default(gen, core::cpu::getDefaultCPUGenerator()); + + std::lock_guard lock(cpu_generator->mutex_); + core::cpu::DispatchCpuFunc( + tensor->Dtype(), + [&]() { + using random_t = std::conditional_t, double, float>; + NormalImpl(*tensor, mean, std, cpu_generator); + }, + "CPU normal"); +} + +} // namespace infini_train::kernels::cpu + +#define REGISTER_CPU_DISTRIBUTION_KERNEL(kernel_name) \ + REGISTER_KERNEL(infini_train::Device::DeviceType::kCPU, kernel_name, infini_train::kernels::cpu::kernel_name) + +REGISTER_CPU_DISTRIBUTION_KERNEL(Uniform) +REGISTER_CPU_DISTRIBUTION_KERNEL(Normal) + +#undef REGISTER_CPU_DISTRIBUTION_KERNEL diff --git a/infini_train/src/kernels/cuda/distribution.cu b/infini_train/src/kernels/cuda/distribution.cu new file mode 100644 index 000000000..6469c3447 --- /dev/null +++ b/infini_train/src/kernels/cuda/distribution.cu @@ -0,0 +1,152 @@ +#include +#include +#include +#include +#include + +#include + +#include "glog/logging.h" + +#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/dispatcher.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)); + // [from, to) is half-open: a sample landing exactly on `to` is mapped back 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)); +} + +} // namespace + +void Uniform(const std::shared_ptr &tensor, double from, double to, const std::optional &gen) { + 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( + gen, 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(const std::shared_ptr &tensor, double mean, double std, const std::optional &gen) { + 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( + gen, 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 infini_train::kernels::cuda + +#define REGISTER_CUDA_DISTRIBUTION_KERNEL(kernel_name) \ + REGISTER_KERNEL(infini_train::Device::DeviceType::kCUDA, kernel_name, infini_train::kernels::cuda::kernel_name) + +REGISTER_CUDA_DISTRIBUTION_KERNEL(Uniform) +REGISTER_CUDA_DISTRIBUTION_KERNEL(Normal) + +#undef REGISTER_CUDA_DISTRIBUTION_KERNEL diff --git a/infini_train/src/nn/init.cc b/infini_train/src/nn/init.cc index 79b4b48b1..b93e59806 100644 --- a/infini_train/src/nn/init.cc +++ b/infini_train/src/nn/init.cc @@ -1,67 +1,77 @@ #include "infini_train/include/nn/init.h" -#include -#include +#include #include -#include -#include +#include #include - -#ifdef USE_OMP -#include -#endif +#include #include "glog/logging.h" #include "infini_train/include/core/runtime/device_guard.h" +#include "infini_train/include/datatype.h" #include "infini_train/include/device.h" +#include "infini_train/include/dispatcher.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); - } +struct DistributionBounds { + double lowest; + double max; +}; + +DistributionBounds GetDistributionBounds(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 {}; } -#else - std::normal_distribution dis(mean, std); - std::generate(buffer.begin(), buffer.end(), [&]() { return generator ? dis(generator.value()) : dis(gen); }); -#endif +} + +void CheckDistributionTensor(const Tensor &tensor) { + CHECK(IsFloatingPointDType(tensor.Dtype())) + << "Uniform and Normal initialization support floating-point tensors only"; +} + +void CheckUniformParameters(const Tensor &tensor, double from, double to) { + const auto bounds = GetDistributionBounds(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 CheckNormalParameters(double std) { CHECK_GE(std, 0.0) << "normal expects std >= 0.0, but found std=" << std; } + +} // namespace +std::shared_ptr Normal(const std::shared_ptr &tensor, float mean, float std, + std::optional generator) { + CheckDistributionTensor(*tensor); + CheckNormalParameters(std); 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)); + Dispatcher::Instance().Call({device.type(), "Normal"}, tensor, static_cast(mean), + static_cast(std), generator); return tensor; } @@ -113,7 +123,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,34 +138,13 @@ 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) { + CheckDistributionTensor(*tensor); + CheckUniformParameters(*tensor, a, b); 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)); - + Dispatcher::Instance().Call({device.type(), "Uniform"}, tensor, static_cast(a), + static_cast(b), generator); return tensor; } diff --git a/infini_train/src/tensor.cc b/infini_train/src/tensor.cc index 18ca3d22b..44282193a 100644 --- a/infini_train/src/tensor.cc +++ b/infini_train/src/tensor.cc @@ -471,7 +471,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); } From fece48d740093bee0f0761e718787ba79e928986 Mon Sep 17 00:00:00 2001 From: chaos Date: Sun, 26 Jul 2026 23:54:04 +0800 Subject: [PATCH 3/8] feat(random): add rand randn and dropout operators --- infini_train/include/autograd/dropout.h | 32 +++++ infini_train/include/nn/functional.h | 18 +++ infini_train/src/autograd/dropout.cc | 44 +++++++ infini_train/src/kernels/cpu/dropout.cc | 127 ++++++++++++++++++ infini_train/src/kernels/cuda/dropout.cu | 161 +++++++++++++++++++++++ infini_train/src/nn/functional.cc | 24 ++++ 6 files changed, 406 insertions(+) create mode 100644 infini_train/include/autograd/dropout.h create mode 100644 infini_train/src/autograd/dropout.cc create mode 100644 infini_train/src/kernels/cpu/dropout.cc create mode 100644 infini_train/src/kernels/cuda/dropout.cu diff --git a/infini_train/include/autograd/dropout.h b/infini_train/include/autograd/dropout.h new file mode 100644 index 000000000..96d1c4991 --- /dev/null +++ b/infini_train/include/autograd/dropout.h @@ -0,0 +1,32 @@ +#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_; +}; + +} // namespace infini_train::autograd diff --git a/infini_train/include/nn/functional.h b/infini_train/include/nn/functional.h index e4354fd10..7ed6d047b 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/src/autograd/dropout.cc b/infini_train/src/autograd/dropout.cc new file mode 100644 index 000000000..5b40c7f6b --- /dev/null +++ b/infini_train/src/autograd/dropout.cc @@ -0,0 +1,44 @@ +#include "infini_train/include/autograd/dropout.h" + +#include + +#include "glog/logging.h" + +#include "infini_train/include/dispatcher.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 device = input->GetDevice().type(); + auto outputs = Dispatcher::Instance().Call, std::shared_ptr>>( + {device, "DropoutForward"}, input, p_, generator_); + // Keep the mask as a non-differentiable output so SetupContext can save it for backward. + return {std::get<0>(outputs), std::get<1>(outputs)}; +} + +void Dropout::SetupContext(const std::vector> &, + const std::vector> &output_tensors) { + CHECK_EQ(output_tensors.size(), 2); + ctx_.MarkNonDifferentiable({output_tensors[1]}); + if (!ctx_.needs_input_grad().empty() && ctx_.needs_input_grad()[0]) { + ctx_.SaveForBackward({output_tensors[1]}); + } +} + +std::vector> Dropout::Backward(const std::vector> &grad_outputs) { + // The mask output is non-differentiable, so its gradient slot is empty. + CHECK_EQ(grad_outputs.size(), 2); + 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 device = grad_output->GetDevice().type(); + return {Dispatcher::Instance().Call>({device, "DropoutBackward"}, grad_output, mask, p_)}; +} + +} // namespace infini_train::autograd diff --git a/infini_train/src/kernels/cpu/dropout.cc b/infini_train/src/kernels/cpu/dropout.cc new file mode 100644 index 000000000..d4e62eb78 --- /dev/null +++ b/infini_train/src/kernels/cpu/dropout.cc @@ -0,0 +1,127 @@ +#include +#include +#include +#include +#include +#include + +#include "glog/logging.h" + +#include "infini_train/include/common/cpu/distributions_helper.h" +#include "infini_train/include/datatype.h" +#include "infini_train/include/dispatcher.h" +#include "infini_train/include/generator.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 DropoutForwardImpl(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)); + common::cpu::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 DropoutBackwardImpl(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); + } +} + +} // namespace + +std::tuple, std::shared_ptr> +DropoutForward(const std::shared_ptr &input, double p, const std::optional &gen) { + CHECK(input->GetDevice().IsCPU()); + CHECK(IsFloatingPointDType(input->Dtype())) << "Dropout supports floating-point tensors only"; + 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; + + auto output = std::make_shared(input->Dims(), input->Dtype(), input->GetDevice()); + auto mask = std::make_shared(input->Dims(), DataType::kUINT8, input->GetDevice()); + + if (p == 0.0 || p == 1.0) { + core::cpu::DispatchCpuFunc( + input->Dtype(), + [&]() { + using random_t = std::conditional_t, double, float>; + DropoutForwardImpl(*output, *mask, *input, p, nullptr); + }, + "CPU dropout forward"); + return {output, mask}; + } + + auto *cpu_generator + = get_generator_or_default(gen, core::cpu::getDefaultCPUGenerator()); + std::lock_guard lock(cpu_generator->mutex_); + core::cpu::DispatchCpuFunc( + input->Dtype(), + [&]() { + using random_t = std::conditional_t, double, float>; + DropoutForwardImpl(*output, *mask, *input, p, cpu_generator); + }, + "CPU dropout forward"); + return {output, mask}; +} + +std::shared_ptr DropoutBackward(const std::shared_ptr &grad_output, const std::shared_ptr &mask, + double p) { + CHECK(grad_output->GetDevice().IsCPU()); + auto grad_input = std::make_shared(grad_output->Dims(), grad_output->Dtype(), grad_output->GetDevice()); + + core::cpu::DispatchCpuFunc( + grad_output->Dtype(), + [&]() { + using random_t = std::conditional_t, double, float>; + DropoutBackwardImpl(*grad_input, *grad_output, *mask, p); + }, + "CPU dropout backward"); + return grad_input; +} + +} // namespace infini_train::kernels::cpu + +#define REGISTER_CPU_DROPOUT_KERNEL(kernel_name) \ + REGISTER_KERNEL(infini_train::Device::DeviceType::kCPU, kernel_name, infini_train::kernels::cpu::kernel_name) + +REGISTER_CPU_DROPOUT_KERNEL(DropoutForward) +REGISTER_CPU_DROPOUT_KERNEL(DropoutBackward) + +#undef REGISTER_CPU_DROPOUT_KERNEL diff --git a/infini_train/src/kernels/cuda/dropout.cu b/infini_train/src/kernels/cuda/dropout.cu new file mode 100644 index 000000000..da106585e --- /dev/null +++ b/infini_train/src/kernels/cuda/dropout.cu @@ -0,0 +1,161 @@ +#include +#include +#include +#include +#include +#include + +#include + +#include "glog/logging.h" + +#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/datatype.h" +#include "infini_train/include/dispatcher.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)); +} + +} // namespace + +std::tuple, std::shared_ptr> +DropoutForward(const std::shared_ptr &input, double p, const std::optional &gen) { + const Device device = input->GetDevice(); + CHECK(device.IsCUDA()); + CHECK(IsFloatingPointDType(input->Dtype())) << "Dropout supports floating-point tensors only"; + 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; + core::DeviceGuard guard(device); + + auto output = std::make_shared(input->Dims(), input->Dtype(), device); + auto mask = std::make_shared(input->Dims(), DataType::kUINT8, device); + + const int64_t n = input->NumElements(); + if (n == 0) { + return {output, mask}; + } + + uint64_t seed = 0; + uint64_t subsequence = 0; + if (p > 0.0 && p < 1.0) { + auto *cuda_generator = get_generator_or_default( + gen, 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()); + return {output, mask}; +} + +std::shared_ptr DropoutBackward(const std::shared_ptr &grad_output, const std::shared_ptr &mask, + double p) { + const Device device = grad_output->GetDevice(); + CHECK(device.IsCUDA()); + core::DeviceGuard guard(device); + + auto grad_input = std::make_shared(grad_output->Dims(), grad_output->Dtype(), device); + + const int64_t n = grad_output->NumElements(); + if (n == 0) { + return grad_input; + } + + 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()); + return grad_input; +} + +} // namespace infini_train::kernels::cuda + +#define REGISTER_CUDA_DROPOUT_KERNEL(kernel_name) \ + REGISTER_KERNEL(infini_train::Device::DeviceType::kCUDA, kernel_name, infini_train::kernels::cuda::kernel_name) + +REGISTER_CUDA_DROPOUT_KERNEL(DropoutForward) +REGISTER_CUDA_DROPOUT_KERNEL(DropoutBackward) + +#undef REGISTER_CUDA_DROPOUT_KERNEL diff --git a/infini_train/src/nn/functional.cc b/infini_train/src/nn/functional.cc index c33e23684..4af04a2d8 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(); } From 5551fe3744185449f27417ca9b9d210a360efb7e Mon Sep 17 00:00:00 2001 From: chaos Date: Sun, 26 Jul 2026 23:55:55 +0800 Subject: [PATCH 4/8] test(random): add generator and random operator coverage --- tests/CMakeLists.txt | 3 + tests/generator/CMakeLists.txt | 10 + tests/generator/test_generator_core.cc | 187 ++++++++++++ tests/generator/test_generator_random_ops.cc | 283 +++++++++++++++++++ 4 files changed, 483 insertions(+) create mode 100644 tests/generator/CMakeLists.txt create mode 100644 tests/generator/test_generator_core.cc create mode 100644 tests/generator/test_generator_random_ops.cc diff --git a/tests/CMakeLists.txt b/tests/CMakeLists.txt index 96776585d..9e5eb3225 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/generator/CMakeLists.txt b/tests/generator/CMakeLists.txt new file mode 100644 index 000000000..7c70ceb0f --- /dev/null +++ b/tests/generator/CMakeLists.txt @@ -0,0 +1,10 @@ +# ========================================================================== +# Generator tests +# ========================================================================== + +file(GLOB GENERATOR_SOURCES ${CMAKE_CURRENT_SOURCE_DIR}/test_*.cc) + +infini_train_add_test_suite(test_generator + SOURCES ${GENERATOR_SOURCES} + TEST_TIMEOUT 40 +) diff --git a/tests/generator/test_generator_core.cc b/tests/generator/test_generator_core.cc new file mode 100644 index 000000000..ac09bdcbb --- /dev/null +++ b/tests/generator/test_generator_core.cc @@ -0,0 +1,187 @@ +#include +#include +#include + +#include "gtest/gtest.h" + +#include "infini_train/include/core/runtime/device_guard.h" +#include "infini_train/include/datatype.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 "tests/common/test_utils.h" + +using namespace infini_train; + +namespace { + +std::shared_ptr CopyToCPU(const std::shared_ptr &tensor) { + auto host = std::make_shared(tensor->Dims(), tensor->Dtype(), Device()); + host->CopyFrom(tensor); + if (tensor->GetDevice().IsCUDA()) { + core::GetDeviceGuardImpl(tensor->GetDevice().type())->SynchronizeDevice(tensor->GetDevice()); + } + return host; +} + +std::vector TensorBytes(const std::shared_ptr &tensor) { + const auto host = CopyToCPU(tensor); + const auto *data = static_cast(host->DataPtr()); + return {data, data + host->SizeInBytes()}; +} + +} // namespace + +class GeneratorCoreTest : public infini_train::test::InfiniTrainTest {}; + +TEST_P(GeneratorCoreTest, PublicInterfaceCopyCloneAndSeed) { + const Device device = GetDevice(); + auto generator = CreateGenerator(device); + EXPECT_TRUE(generator.defined()); + EXPECT_EQ(generator.device(), device); + EXPECT_EQ(generator.current_seed(), Generator::kDefaultSeed); + + generator.set_current_seed(1234); + EXPECT_EQ(generator.current_seed(), 1234U); + const auto first = TensorBytes(nn::function::Rand({128}, DataType::kFLOAT32, device, generator)); + generator.set_current_seed(1234); + EXPECT_EQ(first, TensorBytes(nn::function::Rand({128}, DataType::kFLOAT32, device, generator))); + + generator.set_current_seed(2345); + Generator alias = generator; + EXPECT_EQ(alias, generator); + const auto state_before = TensorBytes(alias.get_state()); + nn::function::Rand({32}, DataType::kFLOAT32, device, generator); + EXPECT_NE(state_before, TensorBytes(alias.get_state())); + + auto clone = generator.clone(); + EXPECT_NE(clone, generator); + EXPECT_EQ(TensorBytes(clone.get_state()), TensorBytes(generator.get_state())); + clone.set_current_seed(3456); + EXPECT_EQ(generator.current_seed(), 2345U); + EXPECT_EQ(clone.current_seed(), 3456U); + + const uint64_t generated_seed = clone.seed(); + EXPECT_EQ(clone.current_seed(), generated_seed); +} + +TEST_P(GeneratorCoreTest, HighBitsOfCPUSeedAffectTheSequence) { + ONLY_CPU(); + const Device device = GetDevice(); + constexpr uint64_t low_seed = 17; + constexpr uint64_t high_seed = low_seed + (uint64_t{1} << 32); + EXPECT_NE(TensorBytes(nn::function::Rand({256}, DataType::kFLOAT32, device, CreateGenerator(device, low_seed))), + TensorBytes(nn::function::Rand({256}, DataType::kFLOAT32, device, CreateGenerator(device, high_seed)))); +} + +TEST_P(GeneratorCoreTest, StateRestoresUniformAndNormalSequence) { + const Device device = GetDevice(); + auto generator = CreateGenerator(device, 4567); + const auto state = generator.get_state(); + EXPECT_EQ(state->Dtype(), DataType::kUINT8); + EXPECT_TRUE(state->GetDevice().IsCPU()); + + const auto uniform = TensorBytes(nn::function::Rand({129}, DataType::kFLOAT32, device, generator)); + generator.set_state(*state); + EXPECT_EQ(uniform, TensorBytes(nn::function::Rand({129}, DataType::kFLOAT32, device, generator))); + + if (device.IsCPU()) { + nn::function::Randn({3}, DataType::kFLOAT32, device, generator); + const auto state_after_odd_normal = generator.get_state(); + const auto normal = TensorBytes(nn::function::Randn({127}, DataType::kFLOAT32, device, generator)); + generator.set_state(*state_after_odd_normal); + EXPECT_EQ(normal, TensorBytes(nn::function::Randn({127}, DataType::kFLOAT32, device, generator))); + } +} + +TEST_P(GeneratorCoreTest, StateRejectsMalformedOrForeignInput) { + ::testing::GTEST_FLAG(death_test_style) = "threadsafe"; + const Device device = GetDevice(); + auto generator = CreateGenerator(device, 4601); + auto wrong_dtype = std::make_shared(std::vector{1}, DataType::kFLOAT32, Device()); + EXPECT_DEATH(generator.set_state(*wrong_dtype), "UINT8"); + + const auto valid_state = generator.get_state(); + auto truncated = std::make_shared( + std::vector{static_cast(valid_state->SizeInBytes() - 1)}, DataType::kUINT8, Device()); + EXPECT_DEATH(generator.set_state(*truncated), "Check failed"); + + if (device.IsCUDA()) { + auto wrong_device = std::make_shared(std::vector{1}, DataType::kUINT8, device); + EXPECT_DEATH(generator.set_state(*wrong_device), "CPU"); + const auto cpu_state = CreateGenerator(Device(), 4601).get_state(); + EXPECT_DEATH(generator.set_state(*cpu_state), "Check failed"); + auto cpu_generator = CreateGenerator(Device(), 4601); + EXPECT_DEATH(cpu_generator.set_state(*valid_state), "Check failed"); + } +} + +TEST_P(GeneratorCoreTest, DefaultExplicitAndUndefinedPathsHaveExpectedState) { + const Device device = GetDevice(); + constexpr uint64_t seed = 7890; + + manual_seed(seed); + const auto *first_default = &GetDefaultGenerator(device); + const auto *second_default = &GetDefaultGenerator(device); + EXPECT_EQ(first_default, second_default); + EXPECT_EQ(first_default->device(), device); + + const auto default_state = TensorBytes(first_default->get_state()); + const auto explicit_result + = TensorBytes(nn::function::Rand({256}, DataType::kFLOAT32, device, CreateGenerator(device, seed))); + EXPECT_EQ(default_state, TensorBytes(first_default->get_state())); + EXPECT_EQ(explicit_result, TensorBytes(nn::function::Rand({256}, DataType::kFLOAT32, device))); + + manual_seed(seed + 1); + const auto default_result = TensorBytes(nn::function::Rand({256}, DataType::kFLOAT32, device)); + manual_seed(seed + 1); + Generator undefined; + EXPECT_EQ(default_result, TensorBytes(nn::function::Rand({256}, DataType::kFLOAT32, device, undefined))); +} + +TEST_P(GeneratorCoreTest, InitializerEntryPointsUseExplicitGenerator) { + const Device device = GetDevice(); + auto make_tensor = [&] { + return std::make_shared(std::vector{16, 16}, DataType::kFLOAT32, device); + }; + + auto normal_a = make_tensor(); + auto normal_b = make_tensor(); + nn::init::Normal(normal_a, 1.5f, 0.5f, CreateGenerator(device, 8101)); + nn::init::Normal(normal_b, 1.5f, 0.5f, CreateGenerator(device, 8101)); + EXPECT_EQ(TensorBytes(normal_a), TensorBytes(normal_b)); + + auto uniform_a = make_tensor(); + auto uniform_b = make_tensor(); + nn::init::Uniform(uniform_a, -2.0f, 3.0f, CreateGenerator(device, 8102)); + uniform_b->Uniform(-2.0f, 3.0f, CreateGenerator(device, 8102)); + EXPECT_EQ(TensorBytes(uniform_a), TensorBytes(uniform_b)); + + auto kaiming_a = make_tensor(); + auto kaiming_b = make_tensor(); + nn::init::KaimingUniform(kaiming_a, 0.1f, nn::init::KaimingMode::kFanIn, nn::init::NonLinearityType::kLeakyReLU, + CreateGenerator(device, 8103)); + nn::init::KaimingUniform(kaiming_b, 0.1f, nn::init::KaimingMode::kFanIn, nn::init::NonLinearityType::kLeakyReLU, + CreateGenerator(device, 8103)); + EXPECT_EQ(TensorBytes(kaiming_a), TensorBytes(kaiming_b)); +} + +TEST_P(GeneratorCoreTest, DefaultCUDAGeneratorsAreIndependentPerDevice) { + ONLY_CUDA(); + REQUIRE_MIN_DEVICES(2); + const Device cuda0(Device::DeviceType::kCUDA, 0); + const Device cuda1(Device::DeviceType::kCUDA, 1); + manual_seed(8501); + + EXPECT_EQ(GetDefaultGenerator(cuda0).device(), cuda0); + EXPECT_EQ(GetDefaultGenerator(cuda1).device(), cuda1); + EXPECT_EQ(GetDefaultGenerator(cuda0).current_seed(), 8501U); + EXPECT_EQ(GetDefaultGenerator(cuda1).current_seed(), 8501U); + const auto cuda1_before = TensorBytes(GetDefaultGenerator(cuda1).get_state()); + nn::function::Rand({256}, DataType::kFLOAT32, cuda0); + EXPECT_EQ(cuda1_before, TensorBytes(GetDefaultGenerator(cuda1).get_state())); +} + +INFINI_TRAIN_REGISTER_TEST(GeneratorCoreTest); diff --git a/tests/generator/test_generator_random_ops.cc b/tests/generator/test_generator_random_ops.cc new file mode 100644 index 000000000..11e2a5d6d --- /dev/null +++ b/tests/generator/test_generator_random_ops.cc @@ -0,0 +1,283 @@ +#include +#include +#include +#include + +#include "gtest/gtest.h" + +#include "infini_train/include/autograd/dropout.h" +#include "infini_train/include/core/runtime/device_guard.h" +#include "infini_train/include/datatype.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 "tests/common/test_utils.h" + +using namespace infini_train; + +namespace { + +std::shared_ptr CopyToCPU(const std::shared_ptr &tensor) { + auto host = std::make_shared(tensor->Dims(), tensor->Dtype(), Device()); + host->CopyFrom(tensor); + if (tensor->GetDevice().IsCUDA()) { + core::GetDeviceGuardImpl(tensor->GetDevice().type())->SynchronizeDevice(tensor->GetDevice()); + } + return host; +} + +std::vector TensorBytes(const std::shared_ptr &tensor) { + const auto host = CopyToCPU(tensor); + const auto *data = static_cast(host->DataPtr()); + return {data, data + host->SizeInBytes()}; +} + +template void ExpectUniformRangeAndFinite(const Tensor &tensor) { + const auto *data = static_cast(tensor.DataPtr()); + for (int64_t index = 0; index < tensor.NumElements(); ++index) { + const double value = static_cast(data[index]); + EXPECT_TRUE(std::isfinite(value)); + EXPECT_GE(value, 0.0); + EXPECT_LT(value, 1.0); + } +} + +template void ExpectFinite(const Tensor &tensor) { + const auto *data = static_cast(tensor.DataPtr()); + for (int64_t index = 0; index < tensor.NumElements(); ++index) { + EXPECT_TRUE(std::isfinite(static_cast(data[index]))); + } +} + +void CheckUniformRangeAndFinite(const std::shared_ptr &tensor) { + const auto host = CopyToCPU(tensor); + switch (host->Dtype()) { + case DataType::kFLOAT16: + ExpectUniformRangeAndFinite(*host); + break; + case DataType::kBFLOAT16: + ExpectUniformRangeAndFinite(*host); + break; + case DataType::kFLOAT32: + ExpectUniformRangeAndFinite(*host); + break; + case DataType::kFLOAT64: + ExpectUniformRangeAndFinite(*host); + break; + default: + FAIL() << "Unexpected dtype"; + } +} + +void CheckFinite(const std::shared_ptr &tensor) { + const auto host = CopyToCPU(tensor); + switch (host->Dtype()) { + case DataType::kFLOAT16: + ExpectFinite(*host); + break; + case DataType::kBFLOAT16: + ExpectFinite(*host); + break; + case DataType::kFLOAT32: + ExpectFinite(*host); + break; + case DataType::kFLOAT64: + ExpectFinite(*host); + break; + default: + FAIL() << "Unexpected dtype"; + } +} + +std::vector> RunFixedScript(const Device &device, Generator generator) { + auto input = std::make_shared(std::vector{257}, DataType::kFLOAT32, device); + input->Fill(1.0f); + return { + TensorBytes(nn::function::Rand({257}, DataType::kFLOAT32, device, generator)), + TensorBytes(nn::function::Randn({257}, DataType::kFLOAT32, device, generator)), + TensorBytes(nn::function::Dropout(input, 0.25, true, generator)), + TensorBytes(nn::function::Rand({257}, DataType::kFLOAT32, device, generator)), + }; +} + +void ExpectMaskValues(const std::shared_ptr &mask) { + EXPECT_EQ(mask->Dtype(), DataType::kUINT8); + const auto host = CopyToCPU(mask); + const auto *data = static_cast(host->DataPtr()); + bool saw_zero = false; + bool saw_one = false; + for (int64_t index = 0; index < host->NumElements(); ++index) { + EXPECT_TRUE(data[index] == 0 || data[index] == 1); + saw_zero |= data[index] == 0; + saw_one |= data[index] == 1; + } + EXPECT_TRUE(saw_zero); + EXPECT_TRUE(saw_one); +} + +template void ExpectUnitInputDropoutValues(const Tensor &output, const Tensor &mask, double p) { + const auto *output_data = static_cast(output.DataPtr()); + const auto *mask_data = static_cast(mask.DataPtr()); + const double kept_value = 1.0 / (1.0 - p); + for (int64_t index = 0; index < output.NumElements(); ++index) { + const double expected = mask_data[index] ? kept_value : 0.0; + EXPECT_NEAR(static_cast(output_data[index]), expected, 0.02); + } +} + +void ExpectUnitInputDropoutValues(const std::shared_ptr &output, const std::shared_ptr &mask, + double p) { + const auto host_output = CopyToCPU(output); + const auto host_mask = CopyToCPU(mask); + switch (host_output->Dtype()) { + case DataType::kFLOAT16: + ExpectUnitInputDropoutValues(*host_output, *host_mask, p); + break; + case DataType::kBFLOAT16: + ExpectUnitInputDropoutValues(*host_output, *host_mask, p); + break; + case DataType::kFLOAT32: + ExpectUnitInputDropoutValues(*host_output, *host_mask, p); + break; + case DataType::kFLOAT64: + ExpectUnitInputDropoutValues(*host_output, *host_mask, p); + break; + default: + FAIL() << "Unexpected dtype"; + } +} + +} // namespace + +class GeneratorRandomOpsTest : public infini_train::test::InfiniTrainTest {}; + +TEST_P(GeneratorRandomOpsTest, RandAndRandnSupportAllFloatingDtypes) { + const std::vector dtypes + = {DataType::kFLOAT16, DataType::kBFLOAT16, DataType::kFLOAT32, DataType::kFLOAT64}; + const Device device = GetDevice(); + + for (const auto dtype : dtypes) { + auto uniform_first = CreateGenerator(device, 1001); + auto uniform_replay = CreateGenerator(device, 1001); + auto uniform = nn::function::Rand({4097}, dtype, device, uniform_first); + EXPECT_EQ(TensorBytes(uniform), TensorBytes(nn::function::Rand({4097}, dtype, device, uniform_replay))); + CheckUniformRangeAndFinite(uniform); + + auto normal_first = CreateGenerator(device, 2002); + auto normal_replay = CreateGenerator(device, 2002); + auto normal = nn::function::Randn({4097}, dtype, device, normal_first); + EXPECT_EQ(TensorBytes(normal), TensorBytes(nn::function::Randn({4097}, dtype, device, normal_replay))); + CheckFinite(normal); + } +} + +TEST_P(GeneratorRandomOpsTest, DropoutSupportsAllFloatingDtypesWithExplicitGenerator) { + const std::vector dtypes + = {DataType::kFLOAT16, DataType::kBFLOAT16, DataType::kFLOAT32, DataType::kFLOAT64}; + const Device device = GetDevice(); + + for (const auto dtype : dtypes) { + auto input = nn::function::Rand({4097}, dtype, device, CreateGenerator(device, 3003)); + manual_seed(4004); + const auto default_state_before = TensorBytes(GetDefaultGenerator(device).get_state()); + + auto dropout_first = CreateGenerator(device, 5005); + auto dropout_replay = CreateGenerator(device, 5005); + auto output = nn::function::Dropout(input, 0.25, true, dropout_first); + EXPECT_EQ(TensorBytes(output), TensorBytes(nn::function::Dropout(input, 0.25, true, dropout_replay))); + CheckFinite(output); + EXPECT_EQ(default_state_before, TensorBytes(GetDefaultGenerator(device).get_state())); + } +} + +TEST_P(GeneratorRandomOpsTest, SameSeedReplaysRandRandnAndDropoutScript) { + const Device device = GetDevice(); + EXPECT_EQ(RunFixedScript(device, CreateGenerator(device, 5678)), + RunFixedScript(device, CreateGenerator(device, 5678))); + EXPECT_NE(RunFixedScript(device, CreateGenerator(device, 5678)), + RunFixedScript(device, CreateGenerator(device, 5679))); + + auto generator = CreateGenerator(device, 6789); + const auto first = TensorBytes(nn::function::Rand({256}, DataType::kFLOAT32, device, generator)); + const auto second = TensorBytes(nn::function::Rand({256}, DataType::kFLOAT32, device, generator)); + EXPECT_NE(first, second); +} + +TEST_P(GeneratorRandomOpsTest, DropoutBackwardSupportsAllFloatingDtypes) { + const Device device = GetDevice(); + const std::vector dtypes + = {DataType::kFLOAT16, DataType::kBFLOAT16, DataType::kFLOAT32, DataType::kFLOAT64}; + + for (const auto dtype : dtypes) { + auto input = std::make_shared(std::vector{4097}, dtype, device, true); + input->Fill(1.0f); + auto dropout = std::make_shared(0.25, CreateGenerator(device, 8201)); + const auto outputs = dropout->Apply({input}); + ASSERT_EQ(outputs.size(), 2U); + ExpectMaskValues(outputs[1]); + ExpectUnitInputDropoutValues(outputs[0], outputs[1], 0.25); + + auto grad_output = std::make_shared(input->Dims(), dtype, device); + grad_output->Fill(1.0f); + const auto grad_inputs = dropout->Backward({grad_output, nullptr}); + ASSERT_EQ(grad_inputs.size(), 1U); + EXPECT_EQ(TensorBytes(outputs[0]), TensorBytes(grad_inputs[0])); + } +} + +TEST_P(GeneratorRandomOpsTest, DropoutIdentityBoundariesDoNotAdvanceGenerator) { + const Device device = GetDevice(); + auto input = std::make_shared(std::vector{64}, DataType::kFLOAT32, device, true); + input->Fill(1.0f); + auto generator = CreateGenerator(device, 8301); + + auto expect_unchanged = [&](const auto &operation) { + const auto before = TensorBytes(generator.get_state()); + operation(); + EXPECT_EQ(before, TensorBytes(generator.get_state())); + }; + + expect_unchanged([&] { EXPECT_EQ(nn::function::Dropout(input, 0.5, false, generator), input); }); + expect_unchanged([&] { EXPECT_EQ(nn::function::Dropout(input, 0.0, true, generator), input); }); +} + +TEST_P(GeneratorRandomOpsTest, DropoutP1ProducesZeroForwardAndBackwardWithoutAdvancing) { + const Device device = GetDevice(); + auto input = std::make_shared(std::vector{64}, DataType::kFLOAT32, device, true); + input->Fill(1.0f); + auto generator = CreateGenerator(device, 8301); + const auto before = TensorBytes(generator.get_state()); + + auto all_dropped = nn::function::Dropout(input, 1.0, true, generator); + EXPECT_EQ(before, TensorBytes(generator.get_state())); + auto zeros = std::make_shared(input->Dims(), input->Dtype(), device); + zeros->Fill(0.0f); + EXPECT_EQ(TensorBytes(all_dropped), TensorBytes(zeros)); + + auto dropout = std::make_shared(1.0, generator); + const auto outputs = dropout->Apply({input}); + auto grad_output = std::make_shared(input->Dims(), input->Dtype(), device); + grad_output->Fill(1.0f); + const auto grad_inputs = dropout->Backward({grad_output, nullptr}); + EXPECT_EQ(TensorBytes(grad_inputs[0]), TensorBytes(zeros)); +} + +TEST_P(GeneratorRandomOpsTest, InvalidDistributionParametersAreRejected) { + ::testing::GTEST_FLAG(death_test_style) = "threadsafe"; + const Device device = GetDevice(); + auto tensor = std::make_shared(std::vector{8}, DataType::kFLOAT32, device); + auto integer_tensor = std::make_shared(std::vector{8}, DataType::kINT32, device); + auto generator = CreateGenerator(device, 9001); + + EXPECT_DEATH(nn::init::Uniform(tensor, 2.0f, 1.0f, generator), "uniform expects a"); + EXPECT_DEATH(nn::init::Normal(tensor, 0.0f, -1.0f, generator), "std >= 0.0"); + EXPECT_DEATH(nn::function::Dropout(tensor, -0.1, true, generator), "between 0 and 1"); + EXPECT_DEATH(nn::function::Dropout(tensor, 1.1, true, generator), "between 0 and 1"); + EXPECT_DEATH(nn::function::Rand({8}, DataType::kINT32, device, generator), "floating-point"); + EXPECT_DEATH(nn::init::Uniform(integer_tensor, 0.0f, 1.0f, generator), "floating-point"); + EXPECT_DEATH(nn::init::Normal(integer_tensor, 0.0f, 1.0f, generator), "floating-point"); +} + +INFINI_TRAIN_REGISTER_TEST(GeneratorRandomOpsTest); From 0c5aa4452029c432711d9388cfa5b64c8ff27308 Mon Sep 17 00:00:00 2001 From: chaos Date: Mon, 27 Jul 2026 05:27:10 +0800 Subject: [PATCH 5/8] feat(ddp): synchronize module state at initialization --- .../parallel/ddp/distributed_data_parallel.h | 2 ++ .../parallel/ddp/distributed_data_parallel.cc | 27 +++++++++++++++++++ 2 files changed, 29 insertions(+) diff --git a/infini_train/include/nn/parallel/ddp/distributed_data_parallel.h b/infini_train/include/nn/parallel/ddp/distributed_data_parallel.h index 823ae82b5..1b2d3ffa3 100644 --- a/infini_train/include/nn/parallel/ddp/distributed_data_parallel.h +++ b/infini_train/include/nn/parallel/ddp/distributed_data_parallel.h @@ -43,6 +43,8 @@ class DistributedDataParallel : public nn::Module { void OnGradReady(const std::shared_ptr ¶m); private: + void SynchronizeModuleState(); + std::shared_ptr reducer_ = nullptr; DistributedDataParallelConfig ddp_config_; diff --git a/infini_train/src/nn/parallel/ddp/distributed_data_parallel.cc b/infini_train/src/nn/parallel/ddp/distributed_data_parallel.cc index b08b64fb0..91a875969 100644 --- a/infini_train/src/nn/parallel/ddp/distributed_data_parallel.cc +++ b/infini_train/src/nn/parallel/ddp/distributed_data_parallel.cc @@ -1,8 +1,10 @@ #include "infini_train/include/nn/parallel/ddp/distributed_data_parallel.h" +#include #include #include #include +#include #include #include @@ -44,6 +46,7 @@ DistributedDataParallel::DistributedDataParallel(std::shared_ptr mod << "All buffers must be on the same device as the module"; } modules_[kModuleName] = std::move(module); + SynchronizeModuleState(); if (ddp_config.zero_stage >= 1) { BuildParamAndGradBuffers(); @@ -61,6 +64,30 @@ DistributedDataParallel::DistributedDataParallel(std::shared_ptr mod } } +void DistributedDataParallel::SynchronizeModuleState() { + auto state = modules_.at(kModuleName)->StateDict(); + std::vector names; + names.reserve(state.size()); + for (const auto &[name, _] : state) { names.push_back(name); } + std::sort(names.begin(), names.end()); + + std::vector> tensors; + tensors.reserve(names.size()); + for (const auto &name : names) { + auto tensor = state.at(name); + CHECK(tensor != nullptr) << "DDP module state tensor '" << name << "' is null"; + if (!tensors.empty()) { + CHECK_EQ(tensor->GetDevice(), tensors.front()->GetDevice()) + << "DDP module state tensors must be on one device"; + } + tensors.push_back(std::move(tensor)); + } + + if (!tensors.empty()) { + ddp_pg_->Broadcast(tensors, 0); + } +} + void DistributedDataParallel::BuildParamAndGradBuffers() { // (param_dtype, grad_dtype) using DTypePair = std::pair; From 41cb40684ff40b8ca38182f9fd39b9188f254640 Mon Sep 17 00:00:00 2001 From: chaos Date: Mon, 27 Jul 2026 00:23:38 +0800 Subject: [PATCH 6/8] chore(generator): add supplementary tests and IWYU fixes --- infini_train/include/autograd/dropout.h | 2 + .../include/common/cpu/distributions_helper.h | 3 +- infini_train/include/nn/init.h | 2 +- infini_train/src/generator.cc | 3 + infini_train/src/tensor.cc | 1 + tests/generator/CMakeLists.txt | 13 +++ tests/generator/test_generator_core.cc | 91 ++++++++++++++++--- tests/generator/test_generator_ddp.cc | 91 +++++++++++++++++++ tests/generator/test_generator_random_ops.cc | 76 +++++++++++++++- 9 files changed, 264 insertions(+), 18 deletions(-) create mode 100644 tests/generator/test_generator_ddp.cc diff --git a/infini_train/include/autograd/dropout.h b/infini_train/include/autograd/dropout.h index 96d1c4991..a6573e1d5 100644 --- a/infini_train/include/autograd/dropout.h +++ b/infini_train/include/autograd/dropout.h @@ -2,6 +2,8 @@ #include #include +#include +#include #include #include "infini_train/include/autograd/function.h" diff --git a/infini_train/include/common/cpu/distributions_helper.h b/infini_train/include/common/cpu/distributions_helper.h index 9cc9f7a62..c8604892b 100644 --- a/infini_train/include/common/cpu/distributions_helper.h +++ b/infini_train/include/common/cpu/distributions_helper.h @@ -6,6 +6,7 @@ #include #include #include +#include #include #include @@ -99,7 +100,7 @@ template struct normal_distribution { const T u2 = uniform(generator); const T r = std::sqrt(static_cast(-2.0) * std::log1p(-u2)); - constexpr T kTwoPi = static_cast(2.0 * M_PI); + constexpr T kTwoPi = static_cast(2.0 * std::numbers::pi_v); const T theta = kTwoPi * u1; const T sample = r * std::sin(theta); diff --git a/infini_train/include/nn/init.h b/infini_train/include/nn/init.h index 05b95b2ca..bbc76c8ae 100644 --- a/infini_train/include/nn/init.h +++ b/infini_train/include/nn/init.h @@ -1,5 +1,6 @@ #pragma once +#include #include #include #include @@ -10,7 +11,6 @@ namespace infini_train { class Tensor; -class Device; } // namespace infini_train namespace infini_train::nn::init { diff --git a/infini_train/src/generator.cc b/infini_train/src/generator.cc index e29ff29e9..e4a928c0f 100644 --- a/infini_train/src/generator.cc +++ b/infini_train/src/generator.cc @@ -1,7 +1,10 @@ #include "infini_train/include/generator.h" +#include + #include "glog/logging.h" +#include "infini_train/include/datatype.h" #include "infini_train/include/tensor.h" #include "infini_train/src/core/runtime/cpu/cpu_generator_impl.h" diff --git a/infini_train/src/tensor.cc b/infini_train/src/tensor.cc index 44282193a..0ccf96fb6 100644 --- a/infini_train/src/tensor.cc +++ b/infini_train/src/tensor.cc @@ -3,6 +3,7 @@ #include #include #include +#include #include #include diff --git a/tests/generator/CMakeLists.txt b/tests/generator/CMakeLists.txt index 7c70ceb0f..c0107ce2d 100644 --- a/tests/generator/CMakeLists.txt +++ b/tests/generator/CMakeLists.txt @@ -3,8 +3,21 @@ # ========================================================================== file(GLOB GENERATOR_SOURCES ${CMAKE_CURRENT_SOURCE_DIR}/test_*.cc) +set(GENERATOR_DDP_SOURCE ${CMAKE_CURRENT_SOURCE_DIR}/test_generator_ddp.cc) +list(REMOVE_ITEM GENERATOR_SOURCES ${GENERATOR_DDP_SOURCE}) infini_train_add_test_suite(test_generator SOURCES ${GENERATOR_SOURCES} TEST_TIMEOUT 40 ) + +if(USE_CUDA) + add_executable(test_generator_ddp ${GENERATOR_DDP_SOURCE}) + target_compile_options(test_generator_ddp PRIVATE -Wno-error) + target_link_libraries(test_generator_ddp PRIVATE GTest::gtest) + target_include_directories(test_generator_ddp PRIVATE ${glog_SOURCE_DIR}/src) + link_infini_train_exe(test_generator_ddp) + + add_test(NAME test_generator_ddp COMMAND $) + set_tests_properties(test_generator_ddp PROPERTIES LABELS "cuda;distributed" TIMEOUT 40) +endif() diff --git a/tests/generator/test_generator_core.cc b/tests/generator/test_generator_core.cc index ac09bdcbb..07e36ffcf 100644 --- a/tests/generator/test_generator_core.cc +++ b/tests/generator/test_generator_core.cc @@ -1,5 +1,6 @@ #include #include +#include #include #include "gtest/gtest.h" @@ -141,6 +142,47 @@ TEST_P(GeneratorCoreTest, DefaultExplicitAndUndefinedPathsHaveExpectedState) { EXPECT_EQ(default_result, TensorBytes(nn::function::Rand({256}, DataType::kFLOAT32, device, undefined))); } +TEST_P(GeneratorCoreTest, CrossBackendGeneratorsAreRejected) { + ONLY_CPU(); +#if defined(USE_CUDA) + const Device cpu; + const Device cuda(Device::DeviceType::kCUDA, 0); + const auto cpu_generator = CreateGenerator(cpu, 8001); + const auto cuda_generator = CreateGenerator(cuda, 8002); + + auto cpu_input = std::make_shared(std::vector{8}, DataType::kFLOAT32, cpu); + auto cuda_input = std::make_shared(std::vector{8}, DataType::kFLOAT32, cuda); + cpu_input->Fill(1.0f); + cuda_input->Fill(1.0f); + + EXPECT_THROW(nn::function::Rand({8}, DataType::kFLOAT32, cpu, cuda_generator), std::invalid_argument); + EXPECT_THROW(nn::function::Rand({8}, DataType::kFLOAT32, cuda, cpu_generator), std::invalid_argument); + EXPECT_THROW(nn::function::Dropout(cpu_input, 0.25, true, cuda_generator), std::invalid_argument); + EXPECT_THROW(nn::function::Dropout(cuda_input, 0.25, true, cpu_generator), std::invalid_argument); +#else + GTEST_SKIP() << "CUDA disabled"; +#endif +} + +TEST_P(GeneratorCoreTest, CUDAGeneratorCanDriveAnotherCUDADevice) { + ONLY_CUDA(); +#if defined(USE_CUDA) + REQUIRE_MIN_DEVICES(2); + const Device cuda0(Device::DeviceType::kCUDA, 0); + const Device cuda1(Device::DeviceType::kCUDA, 1); + auto generator = CreateGenerator(cuda0, 8401); + const auto state_before = TensorBytes(generator.get_state()); + + const auto output = nn::function::Rand({257}, DataType::kFLOAT32, cuda1, generator); + EXPECT_EQ(output->GetDevice(), cuda1); + EXPECT_NE(state_before, TensorBytes(generator.get_state())); + EXPECT_EQ(TensorBytes(output), + TensorBytes(nn::function::Rand({257}, DataType::kFLOAT32, cuda1, CreateGenerator(cuda0, 8401)))); +#else + GTEST_SKIP() << "CUDA disabled"; +#endif +} + TEST_P(GeneratorCoreTest, InitializerEntryPointsUseExplicitGenerator) { const Device device = GetDevice(); auto make_tensor = [&] { @@ -168,20 +210,47 @@ TEST_P(GeneratorCoreTest, InitializerEntryPointsUseExplicitGenerator) { EXPECT_EQ(TensorBytes(kaiming_a), TensorBytes(kaiming_b)); } -TEST_P(GeneratorCoreTest, DefaultCUDAGeneratorsAreIndependentPerDevice) { +TEST_P(GeneratorCoreTest, DefaultCUDAGeneratorsAreIndependentAcrossAllDevices) { ONLY_CUDA(); - REQUIRE_MIN_DEVICES(2); - const Device cuda0(Device::DeviceType::kCUDA, 0); - const Device cuda1(Device::DeviceType::kCUDA, 1); +#if defined(USE_CUDA) + int device_count = 0; + ASSERT_EQ(cudaGetDeviceCount(&device_count), cudaSuccess); + if (device_count < 2) { + GTEST_SKIP() << "requires at least 2 GPUs (found " << device_count << ")"; + } + manual_seed(8501); + std::vector devices; + devices.reserve(device_count); + for (int index = 0; index < device_count; ++index) { + const Device device(Device::DeviceType::kCUDA, index); + devices.push_back(device); + EXPECT_EQ(GetDefaultGenerator(device).device(), device); + EXPECT_EQ(GetDefaultGenerator(device).current_seed(), 8501U); + } - EXPECT_EQ(GetDefaultGenerator(cuda0).device(), cuda0); - EXPECT_EQ(GetDefaultGenerator(cuda1).device(), cuda1); - EXPECT_EQ(GetDefaultGenerator(cuda0).current_seed(), 8501U); - EXPECT_EQ(GetDefaultGenerator(cuda1).current_seed(), 8501U); - const auto cuda1_before = TensorBytes(GetDefaultGenerator(cuda1).get_state()); - nn::function::Rand({256}, DataType::kFLOAT32, cuda0); - EXPECT_EQ(cuda1_before, TensorBytes(GetDefaultGenerator(cuda1).get_state())); + for (int active_index = 0; active_index < device_count; ++active_index) { + std::vector> states_before; + states_before.reserve(device_count); + for (const auto &device : devices) { + states_before.push_back(TensorBytes(GetDefaultGenerator(device).get_state())); + } + + auto tensor = std::make_shared(std::vector{64}, DataType::kFLOAT32, devices[active_index]); + nn::init::Uniform(tensor); + + for (int observed_index = 0; observed_index < device_count; ++observed_index) { + const auto state_after = TensorBytes(GetDefaultGenerator(devices[observed_index]).get_state()); + if (observed_index == active_index) { + EXPECT_NE(states_before[observed_index], state_after); + } else { + EXPECT_EQ(states_before[observed_index], state_after); + } + } + } +#else + GTEST_SKIP() << "CUDA disabled"; +#endif } INFINI_TRAIN_REGISTER_TEST(GeneratorCoreTest); diff --git a/tests/generator/test_generator_ddp.cc b/tests/generator/test_generator_ddp.cc new file mode 100644 index 000000000..51ee2a4b5 --- /dev/null +++ b/tests/generator/test_generator_ddp.cc @@ -0,0 +1,91 @@ +#include +#include +#include +#include + +#include "gtest/gtest.h" +#include + +#include "infini_train/include/core/runtime/device_guard.h" +#include "infini_train/include/device.h" +#include "infini_train/include/nn/modules/linear.h" +#include "infini_train/include/nn/parallel/ddp/distributed_data_parallel.h" +#include "infini_train/include/nn/parallel/ddp/distributed_data_parallel_config.h" +#include "infini_train/include/nn/parallel/global.h" +#include "infini_train/include/nn/parallel/process_group.h" +#include "infini_train/include/nn/parallel/rank.h" +#include "infini_train/include/nn/parallel/utils.h" +#include "infini_train/include/tensor.h" + +using namespace infini_train; + +namespace { + +std::vector CopyValuesToCPU(const std::shared_ptr &tensor) { + Tensor host(tensor->Dims(), tensor->Dtype(), Device()); + host.CopyFrom(*tensor); + core::GetDeviceGuardImpl(tensor->GetDevice().type())->SynchronizeDevice(tensor->GetDevice()); + + const auto *values = static_cast(host.DataPtr()); + return {values, values + host.NumElements()}; +} + +TEST(DistributedDataParallelTest, SynchronizesParametersFromRankZeroAcrossAvailableGPUs) { + int device_count = 0; + ASSERT_EQ(cudaGetDeviceCount(&device_count), cudaSuccess); + if (device_count < 2) { + GTEST_SKIP() << "requires at least 2 GPUs (found " << device_count << ")"; + } + + const int world_size = nn::parallel::global::GetDataParallelSize(); + ASSERT_EQ(world_size, device_count); + + auto *factory = nn::parallel::ProcessGroupFactory::Instance(Device::DeviceType::kCUDA); + factory->GetOrCreate(nn::parallel::GetDataParallelProcessGroupName(0), nn::parallel::GetDataParallelGroupRanks(0)); + + std::vector> synchronized_values(static_cast(world_size)); + std::vector workers; + workers.reserve(static_cast(world_size)); + for (int thread_rank = 0; thread_rank < world_size; ++thread_rank) { + workers.emplace_back([&, thread_rank] { + nn::parallel::global::thread_global_rank = thread_rank; + const Device device(Device::DeviceType::kCUDA, thread_rank); + core::DeviceGuard guard(device); + + auto module = std::make_shared(8, 4); + module->To(device); + for (const auto ¶meter : module->Parameters()) { parameter->Fill(static_cast(thread_rank + 1)); } + + const nn::parallel::Rank rank(0, thread_rank, 1, world_size); + auto ddp = std::make_shared( + std::move(module), rank, nn::parallel::DistributedDataParallelConfig{}); + + for (const auto ¶meter : ddp->module()->Parameters()) { + auto values = CopyValuesToCPU(parameter); + synchronized_values[static_cast(thread_rank)].insert( + synchronized_values[static_cast(thread_rank)].end(), values.begin(), values.end()); + } + }); + } + for (auto &worker : workers) { worker.join(); } + + ASSERT_FALSE(synchronized_values.front().empty()); + for (const auto &rank_values : synchronized_values) { + EXPECT_EQ(rank_values, synchronized_values.front()); + for (float value : rank_values) { EXPECT_FLOAT_EQ(value, 1.0f); } + } +} + +} // namespace + +int main(int argc, char **argv) { + ::testing::InitGoogleTest(&argc, argv); + + int device_count = 0; + if (cudaGetDeviceCount(&device_count) != cudaSuccess || device_count < 1) { + device_count = 1; + } + nn::parallel::global::GlobalEnv::Instance().Init(device_count, 1, false, 1, 1); + + return RUN_ALL_TESTS(); +} diff --git a/tests/generator/test_generator_random_ops.cc b/tests/generator/test_generator_random_ops.cc index 11e2a5d6d..a807e8d69 100644 --- a/tests/generator/test_generator_random_ops.cc +++ b/tests/generator/test_generator_random_ops.cc @@ -1,6 +1,7 @@ #include #include #include +#include #include #include "gtest/gtest.h" @@ -91,7 +92,7 @@ void CheckFinite(const std::shared_ptr &tensor) { } } -std::vector> RunFixedScript(const Device &device, Generator generator) { +std::vector> RunRandomScript(const Device &device, std::optional generator) { auto input = std::make_shared(std::vector{257}, DataType::kFLOAT32, device); input->Fill(1.0f); return { @@ -194,10 +195,10 @@ TEST_P(GeneratorRandomOpsTest, DropoutSupportsAllFloatingDtypesWithExplicitGener TEST_P(GeneratorRandomOpsTest, SameSeedReplaysRandRandnAndDropoutScript) { const Device device = GetDevice(); - EXPECT_EQ(RunFixedScript(device, CreateGenerator(device, 5678)), - RunFixedScript(device, CreateGenerator(device, 5678))); - EXPECT_NE(RunFixedScript(device, CreateGenerator(device, 5678)), - RunFixedScript(device, CreateGenerator(device, 5679))); + EXPECT_EQ(RunRandomScript(device, CreateGenerator(device, 5678)), + RunRandomScript(device, CreateGenerator(device, 5678))); + EXPECT_NE(RunRandomScript(device, CreateGenerator(device, 5678)), + RunRandomScript(device, CreateGenerator(device, 5679))); auto generator = CreateGenerator(device, 6789); const auto first = TensorBytes(nn::function::Rand({256}, DataType::kFLOAT32, device, generator)); @@ -205,6 +206,71 @@ TEST_P(GeneratorRandomOpsTest, SameSeedReplaysRandRandnAndDropoutScript) { EXPECT_NE(first, second); } +TEST_P(GeneratorRandomOpsTest, DefaultGeneratorReplaysRandRandnAndDropoutScript) { + const Device device = GetDevice(); + + manual_seed(7001); + const auto expected = RunRandomScript(device, std::nullopt); + manual_seed(7001); + EXPECT_EQ(expected, RunRandomScript(device, std::nullopt)); + + manual_seed(7001); + Generator undefined; + EXPECT_EQ(expected, RunRandomScript(device, undefined)); +} + +TEST_P(GeneratorRandomOpsTest, ExplicitRandomGeneratorsDoNotAdvanceDefaultGenerator) { + const Device device = GetDevice(); + manual_seed(7101); + const auto default_state_before = TensorBytes(GetDefaultGenerator(device).get_state()); + + RunRandomScript(device, CreateGenerator(device, 7102)); + + EXPECT_EQ(default_state_before, TensorBytes(GetDefaultGenerator(device).get_state())); +} + +TEST_P(GeneratorRandomOpsTest, OffsetViewsUseTheirOwnDataPointers) { + const Device device = GetDevice(); + constexpr int64_t kStorageElements = 12; + constexpr int64_t kOffsetElements = 4; + constexpr int64_t kViewElements = 4; + constexpr size_t kOffsetBytes = kOffsetElements * sizeof(float); + + auto storage = std::make_shared(std::vector{kStorageElements}, DataType::kFLOAT32, device); + storage->Fill(-1.0f); + auto uniform_view = std::make_shared(*storage, kOffsetBytes, std::vector{kViewElements}); + nn::init::Uniform(uniform_view, 0.0f, 1.0f, CreateGenerator(device, 7201)); + + const auto uniform_storage = CopyToCPU(storage); + const auto *uniform_data = static_cast(uniform_storage->DataPtr()); + for (int64_t index = 0; index < kOffsetElements; ++index) { EXPECT_FLOAT_EQ(uniform_data[index], -1.0f); } + for (int64_t index = kOffsetElements; index < kOffsetElements + kViewElements; ++index) { + EXPECT_GE(uniform_data[index], 0.0f); + EXPECT_LT(uniform_data[index], 1.0f); + } + for (int64_t index = kOffsetElements + kViewElements; index < kStorageElements; ++index) { + EXPECT_FLOAT_EQ(uniform_data[index], -1.0f); + } + + auto input_storage = std::make_shared(std::vector{kStorageElements}, DataType::kFLOAT32, device); + input_storage->Fill(-3.0f); + auto input_view = std::make_shared(*input_storage, kOffsetBytes, std::vector{kViewElements}); + input_view->Fill(1.0f); + const auto output = nn::function::Dropout(input_view, 0.5, true, CreateGenerator(device, 7202)); + const auto host_output = CopyToCPU(output); + const auto *output_data = static_cast(host_output->DataPtr()); + for (int64_t index = 0; index < kViewElements; ++index) { + EXPECT_TRUE(output_data[index] == 0.0f || std::abs(output_data[index] - 2.0f) < 1e-6f); + } + + const auto input_storage_host = CopyToCPU(input_storage); + const auto *input_storage_data = static_cast(input_storage_host->DataPtr()); + for (int64_t index = 0; index < kOffsetElements; ++index) { EXPECT_FLOAT_EQ(input_storage_data[index], -3.0f); } + for (int64_t index = kOffsetElements + kViewElements; index < kStorageElements; ++index) { + EXPECT_FLOAT_EQ(input_storage_data[index], -3.0f); + } +} + TEST_P(GeneratorRandomOpsTest, DropoutBackwardSupportsAllFloatingDtypes) { const Device device = GetDevice(); const std::vector dtypes From 383657e69ee217309a54d2f05b860440793050b3 Mon Sep 17 00:00:00 2001 From: chaos Date: Thu, 30 Jul 2026 23:13:54 +0800 Subject: [PATCH 7/8] fix(random): pass registered kernel handles by value --- infini_train/src/kernels/cpu/distribution.cc | 4 ++-- infini_train/src/kernels/cpu/dropout.cc | 4 ++-- infini_train/src/kernels/cuda/distribution.cu | 4 ++-- infini_train/src/kernels/cuda/dropout.cu | 4 ++-- 4 files changed, 8 insertions(+), 8 deletions(-) diff --git a/infini_train/src/kernels/cpu/distribution.cc b/infini_train/src/kernels/cpu/distribution.cc index 947614002..c0c19fbac 100644 --- a/infini_train/src/kernels/cpu/distribution.cc +++ b/infini_train/src/kernels/cpu/distribution.cc @@ -39,7 +39,7 @@ void NormalImpl(Tensor &tensor, double mean, double std, core::cpu::CPUGenerator } // namespace -void Uniform(const std::shared_ptr &tensor, double from, double to, const std::optional &gen) { +void Uniform(const std::shared_ptr tensor, double from, double to, const std::optional gen) { CHECK(tensor->GetDevice().IsCPU()); auto *cpu_generator = get_generator_or_default(gen, core::cpu::getDefaultCPUGenerator()); @@ -54,7 +54,7 @@ void Uniform(const std::shared_ptr &tensor, double from, double to, cons "CPU uniform"); } -void Normal(const std::shared_ptr &tensor, double mean, double std, const std::optional &gen) { +void Normal(const std::shared_ptr tensor, double mean, double std, const std::optional gen) { CHECK(tensor->GetDevice().IsCPU()); auto *cpu_generator = get_generator_or_default(gen, core::cpu::getDefaultCPUGenerator()); diff --git a/infini_train/src/kernels/cpu/dropout.cc b/infini_train/src/kernels/cpu/dropout.cc index d4e62eb78..a23045a8b 100644 --- a/infini_train/src/kernels/cpu/dropout.cc +++ b/infini_train/src/kernels/cpu/dropout.cc @@ -68,7 +68,7 @@ void DropoutBackwardImpl(Tensor &grad_input, const Tensor &grad_output, const Te } // namespace std::tuple, std::shared_ptr> -DropoutForward(const std::shared_ptr &input, double p, const std::optional &gen) { +DropoutForward(const std::shared_ptr input, double p, const std::optional gen) { CHECK(input->GetDevice().IsCPU()); CHECK(IsFloatingPointDType(input->Dtype())) << "Dropout supports floating-point tensors only"; CHECK_GE(p, 0.0) << "dropout probability has to be between 0 and 1, but got " << p; @@ -101,7 +101,7 @@ DropoutForward(const std::shared_ptr &input, double p, const std::option return {output, mask}; } -std::shared_ptr DropoutBackward(const std::shared_ptr &grad_output, const std::shared_ptr &mask, +std::shared_ptr DropoutBackward(const std::shared_ptr grad_output, const std::shared_ptr mask, double p) { CHECK(grad_output->GetDevice().IsCPU()); auto grad_input = std::make_shared(grad_output->Dims(), grad_output->Dtype(), grad_output->GetDevice()); diff --git a/infini_train/src/kernels/cuda/distribution.cu b/infini_train/src/kernels/cuda/distribution.cu index 6469c3447..a8a50514e 100644 --- a/infini_train/src/kernels/cuda/distribution.cu +++ b/infini_train/src/kernels/cuda/distribution.cu @@ -75,7 +75,7 @@ const core::cuda::CudaStream *get_cuda_stream(const Device &device) { } // namespace -void Uniform(const std::shared_ptr &tensor, double from, double to, const std::optional &gen) { +void Uniform(const std::shared_ptr tensor, double from, double to, const std::optional gen) { const Device device = tensor->GetDevice(); CHECK(device.IsCUDA()); const int64_t n = tensor->NumElements(); @@ -108,7 +108,7 @@ void Uniform(const std::shared_ptr &tensor, double from, double to, cons CUDA_CHECK(cudaGetLastError()); } -void Normal(const std::shared_ptr &tensor, double mean, double std, const std::optional &gen) { +void Normal(const std::shared_ptr tensor, double mean, double std, const std::optional gen) { const Device device = tensor->GetDevice(); CHECK(device.IsCUDA()); const int64_t n = tensor->NumElements(); diff --git a/infini_train/src/kernels/cuda/dropout.cu b/infini_train/src/kernels/cuda/dropout.cu index da106585e..fd3b468a9 100644 --- a/infini_train/src/kernels/cuda/dropout.cu +++ b/infini_train/src/kernels/cuda/dropout.cu @@ -81,7 +81,7 @@ const core::cuda::CudaStream *get_cuda_stream(const Device &device) { } // namespace std::tuple, std::shared_ptr> -DropoutForward(const std::shared_ptr &input, double p, const std::optional &gen) { +DropoutForward(const std::shared_ptr input, double p, const std::optional gen) { const Device device = input->GetDevice(); CHECK(device.IsCUDA()); CHECK(IsFloatingPointDType(input->Dtype())) << "Dropout supports floating-point tensors only"; @@ -122,7 +122,7 @@ DropoutForward(const std::shared_ptr &input, double p, const std::option return {output, mask}; } -std::shared_ptr DropoutBackward(const std::shared_ptr &grad_output, const std::shared_ptr &mask, +std::shared_ptr DropoutBackward(const std::shared_ptr grad_output, const std::shared_ptr mask, double p) { const Device device = grad_output->GetDevice(); CHECK(device.IsCUDA()); From 34cbc914947e75fd922c2dc0b77bdbca27fe2264 Mon Sep 17 00:00:00 2001 From: chaos <2242119@zju.edu.cn> Date: Sun, 13 Sep 2026 01:40:27 +0800 Subject: [PATCH 8/8] refactor(generator): address PR review feedback Separate the public Generator handle from implementation access and dispatch through registered CPU/CUDA backends. Initialize default generators once, validate device indices, and align operation and helper naming. Use bounded CPU/CUDA uniform resampling and curand_uniform2_double. Seed training examples, split backend-specific tests, and clean up includes and comments. Add RNG state and constant/empty uniform regression coverage. Validated on CPU-only and CUDA builds, including dual RTX 3090 tests and two-process RNG replay and communication under infini_run and torchrun. Preserve the approved sm_120 build target. --- CMakeLists.txt | 2 +- example/gpt2/main.cc | 5 +- example/llama3/main.cc | 5 +- example/mixtral/main.cc | 2 + example/mnist/main.cc | 2 + .../include/core/runtime/generator_backend.h | 50 +++++ infini_train/include/generator.h | 103 ++-------- infini_train/include/generator_impl.h | 90 +++++++++ infini_train/include/tensor.h | 2 - infini_train/src/autograd/dropout.cc | 1 + .../core/runtime/cpu/cpu_generator_impl.cc | 187 ++++++++++++------ .../src/core/runtime/cpu/cpu_generator_impl.h | 20 +- .../core/runtime/cuda/cuda_generator_impl.cc | 79 +++++--- .../core/runtime/cuda/cuda_generator_impl.h | 15 +- .../src/core/runtime/generator_backend.cc | 42 ++++ infini_train/src/generator.cc | 61 ++---- infini_train/src/kernels/cpu/cross_entropy.cc | 1 + infini_train/src/kernels/cpu/distribution.cc | 63 +++--- .../kernels}/cpu/distributions_helper.h | 50 ++--- infini_train/src/kernels/cpu/dropout.cc | 59 +++--- infini_train/src/kernels/cpu/transform.cc | 1 + infini_train/src/kernels/cuda/distribution.cu | 97 +++++---- infini_train/src/kernels/cuda/dropout.cu | 81 ++++---- infini_train/src/kernels/cuda/elementwise.cu | 1 + infini_train/src/kernels/cuda/gather.cu | 2 + infini_train/src/kernels/cuda/no_op.cu | 2 + infini_train/src/kernels/cuda/reduction.cu | 2 + infini_train/src/nn/functional.cc | 3 + infini_train/src/nn/init.cc | 6 + .../nn/parallel/ddp/param_and_grad_buffer.cc | 1 + tests/common/test_utils.h | 16 ++ tests/generator/CMakeLists.txt | 24 ++- .../generator/cpu_only/test_generator_cpu.cc | 56 ++++++ .../cuda_only/test_generator_cuda.cc | 110 +++++++++++ tests/generator/test_generator_core.cc | 153 ++------------ tests/generator/test_generator_random_ops.cc | 144 +++++++++----- 36 files changed, 949 insertions(+), 589 deletions(-) create mode 100644 infini_train/include/core/runtime/generator_backend.h create mode 100644 infini_train/include/generator_impl.h create mode 100644 infini_train/src/core/runtime/generator_backend.cc rename infini_train/{include/common => src/kernels}/cpu/distributions_helper.h (57%) create mode 100644 tests/generator/cpu_only/test_generator_cpu.cc create mode 100644 tests/generator/cuda_only/test_generator_cuda.cc diff --git a/CMakeLists.txt b/CMakeLists.txt index 6bd8069d4..3f7d72546 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -104,7 +104,7 @@ if(USE_CUDA) file(GLOB_RECURSE CUDA_KERNELS ${PROJECT_SOURCE_DIR}/infini_train/src/*.cu) add_library(infini_train_cuda_kernels STATIC ${CUDA_KERNELS}) - set_target_properties(infini_train_cuda_kernels PROPERTIES CUDA_ARCHITECTURES "75;80;90") + set_target_properties(infini_train_cuda_kernels PROPERTIES CUDA_ARCHITECTURES "75;80;90;120") target_link_libraries(infini_train_cuda_kernels PUBLIC diff --git a/example/gpt2/main.cc b/example/gpt2/main.cc index 60c0c9080..0e243174b 100644 --- a/example/gpt2/main.cc +++ b/example/gpt2/main.cc @@ -15,6 +15,7 @@ #include "infini_train/include/core/runtime/device_guard.h" #include "infini_train/include/dataloader.h" #include "infini_train/include/device.h" +#include "infini_train/include/generator.h" #include "infini_train/include/lr_scheduler.h" #include "infini_train/include/nn/lora/lora_utils.h" #include "infini_train/include/nn/modules/loss.h" @@ -216,9 +217,6 @@ void Train(const nn::parallel::Rank &rank) { LOG(INFO) << "total desired batch size: " << FLAGS_total_batch_size << " => calculated gradient accumulation steps: " << grad_accum_steps; - // rng / reproducibility - // ManualSeed(42); - // init the model, either from scratch or from OpenAI pretrained checkpoint nn::TransformerConfig model_config = gpt2::GPT2Config(); std::shared_ptr model = nullptr; @@ -565,6 +563,7 @@ int main(int argc, char *argv[]) { nn::parallel::global::InitAllEnv(FLAGS_nthread_per_process, FLAGS_tensor_parallel, FLAGS_sequence_parallel, FLAGS_pipeline_parallel, FLAGS_virtual_pipeline_parallel); utils::PrecisionCheckEnv::Instance().Init(precision_config); + infini_train::ManualSeed(42); LOG(INFO) << nn::parallel::global::ProcessGroupOverview(); diff --git a/example/llama3/main.cc b/example/llama3/main.cc index 302e08080..aaf89c6e1 100644 --- a/example/llama3/main.cc +++ b/example/llama3/main.cc @@ -14,6 +14,7 @@ #include "infini_train/include/core/runtime/device_guard.h" #include "infini_train/include/dataloader.h" #include "infini_train/include/device.h" +#include "infini_train/include/generator.h" #include "infini_train/include/lr_scheduler.h" #include "infini_train/include/nn/lora/lora_utils.h" #include "infini_train/include/nn/modules/loss.h" @@ -204,9 +205,6 @@ void Train(const nn::parallel::Rank &rank) { << " => calculated gradient accumulation steps: " << grad_accum_steps; } - // rng / reproducibility - // ManualSeed(42); - nn::TransformerConfig model_config = llama3::LLaMA3Config(); std::shared_ptr model = nullptr; if (!FLAGS_llmc_filepath.empty()) { @@ -542,6 +540,7 @@ int main(int argc, char *argv[]) { nn::parallel::global::InitAllEnv(FLAGS_nthread_per_process, FLAGS_tensor_parallel, FLAGS_sequence_parallel, FLAGS_pipeline_parallel, FLAGS_virtual_pipeline_parallel); utils::PrecisionCheckEnv::Instance().Init(precision_config); + infini_train::ManualSeed(42); LOG(INFO) << nn::parallel::global::ProcessGroupOverview(); diff --git a/example/mixtral/main.cc b/example/mixtral/main.cc index 848538650..367966b91 100644 --- a/example/mixtral/main.cc +++ b/example/mixtral/main.cc @@ -13,6 +13,7 @@ #include "infini_train/include/core/runtime/device_guard.h" #include "infini_train/include/dataloader.h" #include "infini_train/include/device.h" +#include "infini_train/include/generator.h" #include "infini_train/include/nn/modules/loss.h" #include "infini_train/include/nn/modules/transformer/transformer.h" #include "infini_train/include/nn/parallel/global.h" @@ -69,6 +70,7 @@ int main(int argc, char *argv[]) { /*sequence_parallel_enabled=*/false, /*pipeline_parallel_size=*/1, /*virtual_pipeline_parallel_size=*/1); + infini_train::ManualSeed(42); infini_train::nn::TransformerConfig model_config = mixtral::TinyMixtralConfig(); mixtral::SanitizeTinyMixtralConfig(model_config); diff --git a/example/mnist/main.cc b/example/mnist/main.cc index e62257d7f..3d7139ba7 100644 --- a/example/mnist/main.cc +++ b/example/mnist/main.cc @@ -11,6 +11,7 @@ #include "infini_train/include/dataloader.h" #include "infini_train/include/device.h" +#include "infini_train/include/generator.h" #include "infini_train/include/nn/modules/loss.h" #include "infini_train/include/optimizer.h" @@ -39,6 +40,7 @@ DEFINE_validator(device, int main(int argc, char *argv[]) { gflags::ParseCommandLineFlags(&argc, &argv, true); google::InitGoogleLogging(argv[0]); + infini_train::ManualSeed(42); auto train_dataset = std::make_shared(FLAGS_dataset, true); DataLoader train_dataloader(train_dataset, FLAGS_bs); diff --git a/infini_train/include/core/runtime/generator_backend.h b/infini_train/include/core/runtime/generator_backend.h new file mode 100644 index 000000000..71798b1ce --- /dev/null +++ b/infini_train/include/core/runtime/generator_backend.h @@ -0,0 +1,50 @@ +#pragma once + +// Internal Generator backend registration interface. Application code should include generator.h. + +#include +#include +#include + +#include "infini_train/include/common/common.h" +#include "infini_train/include/device.h" +#include "infini_train/include/generator.h" + +namespace infini_train { + +class GeneratorBackend { +public: + virtual ~GeneratorBackend() = default; + + virtual Device::DeviceType Type() const = 0; + virtual Generator Create(const Device &device, uint64_t seed) = 0; + virtual const Generator &GetDefault(const Device &device) = 0; + virtual void ManualSeedAll(uint64_t seed) = 0; +}; + +// Registers one backend per device type during static initialization. +class GeneratorBackendRegistry { +public: + static GeneratorBackendRegistry &Instance(); + + void Register(Device::DeviceType type, std::unique_ptr backend); + // Throws if no backend is registered for the requested device type. + GeneratorBackend &Get(Device::DeviceType type) const; + void ManualSeedAll(uint64_t seed) const; + +private: + GeneratorBackendRegistry() = default; + GeneratorBackendRegistry(const GeneratorBackendRegistry &) = delete; + GeneratorBackendRegistry &operator=(const GeneratorBackendRegistry &) = delete; + + std::unordered_map> backends_; +}; + +} // namespace infini_train + +// CAT expands __COUNTER__ before token pasting so each registration has a unique name. +#define INFINI_TRAIN_REGISTER_GENERATOR_BACKEND(device_type, class_impl) \ + [[maybe_unused]] static const bool CAT(infini_train_generator_backend_registered_, __COUNTER__) = []() { \ + infini_train::GeneratorBackendRegistry::Instance().Register(device_type, std::make_unique()); \ + return true; \ + }(); diff --git a/infini_train/include/generator.h b/infini_train/include/generator.h index ee107c71e..00c508554 100644 --- a/infini_train/include/generator.h +++ b/infini_train/include/generator.h @@ -2,62 +2,20 @@ #include #include -#include -#include -#include -#include #include "infini_train/include/device.h" namespace infini_train { class Tensor; +class GeneratorImpl; -namespace detail { - -// Validates the common Tensor contract for serialized RNG states. -void check_rng_state(const Tensor &state); - -} // namespace detail - -// Base interface for device-specific random number generators. -class GeneratorImpl { -public: - explicit GeneratorImpl(Device device) : device_(device) {} - virtual ~GeneratorImpl() = default; - - 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; - - std::shared_ptr clone() const { return std::shared_ptr(clone_impl()); } - - Device device() const { return device_; } - - // Callers must lock this mutex when an operation spans multiple generator calls. - std::mutex mutex_; - -protected: - Device device_; - - virtual GeneratorImpl *clone_impl() const = 0; -}; - -// A lightweight handle with shared-copy semantics. Use clone() for an independent state. +// A lightweight handle with shared-copy semantics. Use Clone() for an independent state. class Generator { public: static constexpr uint64_t kDefaultSeed = 67280421310721; - Generator() = default; - - explicit Generator(std::shared_ptr impl); + Generator() = delete; Generator(const Generator &) = default; Generator &operator=(const Generator &) = default; @@ -66,65 +24,30 @@ class Generator { ~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(); } - + void set_current_seed(uint64_t seed) const; + uint64_t current_seed() const; + uint64_t Seed(); void set_state(const Tensor &state); std::shared_ptr get_state() const; - - Device device() const { return impl_->device(); } - - Generator clone() const { return Generator(impl_->clone()); } - - std::mutex &mutex() const { return impl_->mutex_; } - - // Prefer check_generator(); this unchecked accessor assumes a matching backend. - template T *get() const { return static_cast(impl_.get()); } - - GeneratorImpl *unsafeGetGeneratorImpl() const { return impl_.get(); } - bool defined() const { return impl_ != nullptr; } + Device device() const; + Generator Clone() const; 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: + friend class GeneratorAccessor; + + explicit Generator(std::shared_ptr impl); std::shared_ptr impl_; }; -// Internal factory for backend implementations. -template Generator make_generator(Args &&...args) { - return Generator(std::make_shared(std::forward(args)...)); -} - -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; -} - -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); -} - -// Creates a generator for the requested device without exposing its backend implementation. Generator CreateGenerator(const Device &device, uint64_t seed = Generator::kDefaultSeed); // Returns the lazily initialized default generator for the requested device. const Generator &GetDefaultGenerator(const Device &device); -// Reset the default generators for all enabled devices. -void manual_seed(uint64_t seed); +// Initializes or resets the default generators for all enabled devices with the given seed. +void ManualSeed(uint64_t seed); } // namespace infini_train diff --git a/infini_train/include/generator_impl.h b/infini_train/include/generator_impl.h new file mode 100644 index 000000000..9a59ee7f6 --- /dev/null +++ b/infini_train/include/generator_impl.h @@ -0,0 +1,90 @@ +#pragma once + +// Internal Generator implementation interface. Application code should include generator.h. + +#include +#include +#include +#include +#include +#include +#include + +#include "infini_train/include/device.h" +#include "infini_train/include/generator.h" +#include "infini_train/include/tensor.h" + +namespace infini_train { + +class Tensor; + +class GeneratorImpl { +public: + explicit GeneratorImpl(Device device) : device_(device) {} + virtual ~GeneratorImpl() = default; + + 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; + + std::shared_ptr Clone() const { return std::shared_ptr(CloneImpl()); } + Device device() const { return device_; } + + // Callers must hold this mutex when accessing shared RNG state concurrently. + std::mutex mutex_; + +protected: + Device device_; + virtual GeneratorImpl *CloneImpl() const = 0; +}; + +class GeneratorAccessor { +public: + static GeneratorImpl &Get(const Generator &generator) { + if (!generator.impl_) { + throw std::invalid_argument("Undefined Generator"); + } + return *generator.impl_; + } + + static std::mutex &Mutex(const Generator &generator) { return Get(generator).mutex_; } + + static Generator FromImpl(std::shared_ptr impl) { return Generator(std::move(impl)); } +}; + +template T &CheckedGeneratorImpl(const Generator &generator, const Device &expected_device) { + static_assert(std::is_base_of_v); + auto &base = GeneratorAccessor::Get(generator); + if (base.device() != expected_device) { + throw std::invalid_argument("Generator device mismatch"); + } + auto *typed = dynamic_cast(&base); + if (typed == nullptr) { + throw std::invalid_argument("Generator backend mismatch"); + } + return *typed; +} + +template Generator MakeGenerator(Args &&...args) { + return GeneratorAccessor::FromImpl(std::make_shared(std::forward(args)...)); +} + +template +T &GetGeneratorOrDefault(const std::optional &generator, const Generator &default_generator, + const Device &expected_device) { + const Generator &chosen = generator.has_value() ? *generator : default_generator; + return CheckedGeneratorImpl(chosen, expected_device); +} + +namespace detail { +void CheckRngState(const Tensor &state); +} // namespace detail + +} // namespace infini_train diff --git a/infini_train/include/tensor.h b/infini_train/include/tensor.h index 05b292630..e5ef60ec3 100644 --- a/infini_train/include/tensor.h +++ b/infini_train/include/tensor.h @@ -4,8 +4,6 @@ #include #include #include -#include - #include #include "Eigen/Dense" diff --git a/infini_train/src/autograd/dropout.cc b/infini_train/src/autograd/dropout.cc index 5b40c7f6b..531710e9f 100644 --- a/infini_train/src/autograd/dropout.cc +++ b/infini_train/src/autograd/dropout.cc @@ -4,6 +4,7 @@ #include "glog/logging.h" +#include "infini_train/include/device.h" #include "infini_train/include/dispatcher.h" #include "infini_train/include/tensor.h" diff --git a/infini_train/src/core/runtime/cpu/cpu_generator_impl.cc b/infini_train/src/core/runtime/cpu/cpu_generator_impl.cc index cb49fb3b8..cb827533a 100644 --- a/infini_train/src/core/runtime/cpu/cpu_generator_impl.cc +++ b/infini_train/src/core/runtime/cpu/cpu_generator_impl.cc @@ -1,12 +1,21 @@ #include "infini_train/src/core/runtime/cpu/cpu_generator_impl.h" +#include #include #include +#include #include #include +#include +#include +#include #include "glog/logging.h" +#include "infini_train/include/core/runtime/generator_backend.h" +#include "infini_train/include/datatype.h" +#include "infini_train/include/device.h" +#include "infini_train/include/generator_impl.h" #include "infini_train/include/tensor.h" namespace infini_train::core::cpu { @@ -14,41 +23,71 @@ namespace { // Backend tag used to reject states from other generator implementations. constexpr char kCPUStateMagic[] = "ITRNGCPU"; -constexpr size_t kStateMagicSize = sizeof(kCPUStateMagic) - 1; - -constexpr size_t kStateFooterSize - = sizeof(uint64_t) + sizeof(uint8_t) + sizeof(float) + sizeof(uint8_t) + sizeof(double); - -} // namespace - -// std::mt19937(seed) truncates to 32 bits, so seed both halves explicitly. -static std::mt19937 make_seeded_engine(uint64_t seed) { +constexpr std::size_t kStateMagicSize = sizeof(kCPUStateMagic) - 1; + +// Fixed footer after the variable-length mt19937 stream. Cache flags are +// serialized as uint8_t; absent cache values are serialized as zero. +constexpr std::size_t kSerializedSeedSize = sizeof(uint64_t); +constexpr std::size_t kSerializedFloatCacheFlagSize = sizeof(uint8_t); +constexpr std::size_t kSerializedFloatCacheValueSize = sizeof(float); +constexpr std::size_t kSerializedDoubleCacheFlagSize = sizeof(uint8_t); +constexpr std::size_t kSerializedDoubleCacheValueSize = sizeof(double); +constexpr std::size_t kStateFooterSize = kSerializedSeedSize + kSerializedFloatCacheFlagSize + + kSerializedFloatCacheValueSize + kSerializedDoubleCacheFlagSize + + kSerializedDoubleCacheValueSize; + +// Seed through seed_seq using both 32-bit halves of the seed. This produces a +// different sequence from direct std::mt19937(seed) initialization, even when +// the seed fits in 32 bits. +std::mt19937 MakeSeededEngine(uint64_t seed) { std::seed_seq seq{static_cast(seed), static_cast(seed >> 32)}; std::mt19937 engine; engine.seed(seq); return engine; } -static uint64_t getNonDeterministicRandom() { +uint64_t GenerateNonDeterministicSeed() { std::random_device rd; uint64_t val = (static_cast(rd()) << 32) | rd(); return val; } +// Copies state bytes and advances offset; data and value must reference valid +// buffers of data_size and value_size bytes, and offset must be a valid pointer. +void WriteBytes(uint8_t *data, std::size_t data_size, std::size_t *offset, const void *value, std::size_t value_size) { + CHECK_LE(*offset, data_size) << "CPU generator state write offset out of bounds"; + CHECK_LE(value_size, data_size - *offset) << "CPU generator state write exceeds buffer"; + if (value_size != 0) { + std::memcpy(data + *offset, value, value_size); + } + *offset += value_size; +} + +void ReadBytes(const uint8_t *data, std::size_t data_size, std::size_t *offset, void *value, std::size_t value_size) { + CHECK_LE(*offset, data_size) << "CPU generator state read offset out of bounds"; + CHECK_LE(value_size, data_size - *offset) << "CPU generator state read exceeds buffer"; + if (value_size != 0) { + std::memcpy(value, data + *offset, value_size); + } + *offset += value_size; +} + +} // namespace + CPUGeneratorImpl::CPUGeneratorImpl(uint64_t seed) - : GeneratorImpl(Device(Device::DeviceType::kCPU, 0)), engine_(make_seeded_engine(seed)), seed_(seed) {} + : GeneratorImpl(Device(Device::DeviceType::kCPU, 0)), engine_(MakeSeededEngine(seed)), seed_(seed) {} void CPUGeneratorImpl::set_current_seed(uint64_t seed) { seed_ = seed; next_float_normal_sample_.reset(); next_double_normal_sample_.reset(); - engine_ = make_seeded_engine(seed); + engine_ = MakeSeededEngine(seed); } uint64_t CPUGeneratorImpl::current_seed() const { return seed_; } -uint64_t CPUGeneratorImpl::seed() { - uint64_t random_seed = getNonDeterministicRandom(); +uint64_t CPUGeneratorImpl::Seed() { + uint64_t random_seed = GenerateNonDeterministicSeed(); set_current_seed(random_seed); return random_seed; } @@ -57,7 +96,7 @@ uint64_t CPUGeneratorImpl::seed() { // The std::mt19937 stream format is only portable across compatible standard-library builds. void CPUGeneratorImpl::set_state(const Tensor &state) { - ::infini_train::detail::check_rng_state(state); + ::infini_train::detail::CheckRngState(state); const size_t data_size = state.SizeInBytes(); CHECK_GT(data_size, kStateMagicSize + kStateFooterSize) << "CPU generator state is too small"; @@ -78,19 +117,18 @@ void CPUGeneratorImpl::set_state(const Tensor &state) { size_t offset = kStateMagicSize + 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++]; + ReadBytes(data, data_size, &offset, &restored_seed, kSerializedSeedSize); + uint8_t has_float = 0; + ReadBytes(data, data_size, &offset, &has_float, kSerializedFloatCacheFlagSize); 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++]; + ReadBytes(data, data_size, &offset, &restored_float, kSerializedFloatCacheValueSize); + uint8_t has_double = 0; + ReadBytes(data, data_size, &offset, &has_double, kSerializedDoubleCacheFlagSize); 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)); + ReadBytes(data, data_size, &offset, &restored_double, kSerializedDoubleCacheValueSize); + CHECK_EQ(offset, data_size) << "CPU generator state size mismatch"; // Do not change the generator until the complete state has been validated. engine_ = restored_engine; @@ -113,32 +151,26 @@ std::shared_ptr CPUGeneratorImpl::get_state() const { uint8_t *data = static_cast(state_tensor->DataPtr()); size_t offset = 0; - std::memcpy(data + offset, kCPUStateMagic, kStateMagicSize); - offset += kStateMagicSize; + const uint8_t has_float = next_float_normal_sample_.has_value() ? 1 : 0; + const float float_val = has_float ? *next_float_normal_sample_ : 0.0f; + const uint8_t has_double = next_double_normal_sample_.has_value() ? 1 : 0; + const double double_val = has_double ? *next_double_normal_sample_ : 0.0; - std::memcpy(data + offset, engine_str.data(), engine_size); - offset += engine_size; - - std::memcpy(data + offset, &seed_, sizeof(seed_)); - offset += sizeof(seed_); - - 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); - - 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)); + WriteBytes(data, total_size, &offset, kCPUStateMagic, kStateMagicSize); + WriteBytes(data, total_size, &offset, engine_str.data(), engine_size); + WriteBytes(data, total_size, &offset, &seed_, kSerializedSeedSize); + WriteBytes(data, total_size, &offset, &has_float, kSerializedFloatCacheFlagSize); + WriteBytes(data, total_size, &offset, &float_val, kSerializedFloatCacheValueSize); + WriteBytes(data, total_size, &offset, &has_double, kSerializedDoubleCacheFlagSize); + WriteBytes(data, total_size, &offset, &double_val, kSerializedDoubleCacheValueSize); + CHECK_EQ(offset, total_size) << "CPU generator state size mismatch"; return state_tensor; } -uint32_t CPUGeneratorImpl::random() { return engine_(); } +uint32_t CPUGeneratorImpl::Random() { return engine_(); } -uint64_t CPUGeneratorImpl::random64() { +uint64_t CPUGeneratorImpl::Random64() { uint32_t hi = engine_(); uint32_t lo = engine_(); return (static_cast(hi) << 32) | lo; @@ -154,11 +186,11 @@ void CPUGeneratorImpl::set_next_double_normal_sample(std::optional randn next_double_normal_sample_ = randn; } -std::shared_ptr CPUGeneratorImpl::clone() const { - return std::shared_ptr(clone_impl()); +std::shared_ptr CPUGeneratorImpl::Clone() const { + return std::shared_ptr(CloneImpl()); } -CPUGeneratorImpl *CPUGeneratorImpl::clone_impl() const { +CPUGeneratorImpl *CPUGeneratorImpl::CloneImpl() const { auto *clone = new CPUGeneratorImpl(seed_); clone->set_engine(engine_); clone->set_next_float_normal_sample(next_float_normal_sample_); @@ -168,23 +200,64 @@ CPUGeneratorImpl *CPUGeneratorImpl::clone_impl() const { 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 namespace infini_train::core::cpu { +namespace { + +struct DefaultCpuGeneratorState { + std::once_flag init_once; + std::optional generator; +}; -const Generator &getDefaultCPUGenerator() { - static auto default_gen = createCPUGenerator(getNonDeterministicRandom()); - return default_gen; +DefaultCpuGeneratorState &GetDefaultCpuGeneratorState() { + static DefaultCpuGeneratorState state; + return state; } -Generator createCPUGenerator(uint64_t seed) { return make_generator(seed); } +} // namespace + +const Generator &GetDefaultCpuGenerator() { + auto &state = GetDefaultCpuGeneratorState(); + std::call_once(state.init_once, + [&state] { state.generator.emplace(CreateCpuGenerator(GenerateNonDeterministicSeed())); }); + return *state.generator; +} -void manual_seed(uint64_t seed) { - const auto &default_gen = getDefaultCPUGenerator(); - std::lock_guard lock(default_gen.mutex()); - default_gen.set_current_seed(seed); +Generator CreateCpuGenerator(uint64_t seed) { return MakeGenerator(seed); } + +void ManualSeed(uint64_t seed) { + auto &state = GetDefaultCpuGeneratorState(); + bool initialized_here = false; + std::call_once(state.init_once, [&state, seed, &initialized_here] { + state.generator.emplace(CreateCpuGenerator(seed)); + initialized_here = true; + }); + if (initialized_here) { + return; + } + const Generator &generator = *state.generator; + std::lock_guard lock(GeneratorAccessor::Mutex(generator)); + generator.set_current_seed(seed); } } // namespace infini_train::core::cpu + +namespace infini_train::core::cpu { +namespace { + +class CPUGeneratorBackend final : public GeneratorBackend { +public: + Device::DeviceType Type() const override { return Device::DeviceType::kCPU; } + + Generator Create(const Device & /*device*/, uint64_t seed) override { return CreateCpuGenerator(seed); } + + const Generator &GetDefault(const Device & /*device*/) override { return GetDefaultCpuGenerator(); } + + void ManualSeedAll(uint64_t seed) override { ::infini_train::core::cpu::ManualSeed(seed); } +}; + +INFINI_TRAIN_REGISTER_GENERATOR_BACKEND(Device::DeviceType::kCPU, CPUGeneratorBackend); + +} // namespace +} // 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 index 22a766ec7..738b20c74 100644 --- a/infini_train/src/core/runtime/cpu/cpu_generator_impl.h +++ b/infini_train/src/core/runtime/cpu/cpu_generator_impl.h @@ -6,6 +6,8 @@ #include #include "infini_train/include/generator.h" +#include "infini_train/include/generator_impl.h" +#include "infini_train/include/tensor.h" namespace infini_train::core::cpu { @@ -17,16 +19,14 @@ class CPUGeneratorImpl final : public GeneratorImpl { void set_current_seed(uint64_t seed) override; uint64_t current_seed() const override; - uint64_t seed() override; + uint64_t Seed() override; void set_state(const Tensor &state) override; std::shared_ptr get_state() const override; - std::shared_ptr clone() const; + std::shared_ptr Clone() const; - static Device::DeviceType device_type(); - - uint32_t random(); - uint64_t random64(); + uint32_t Random(); + uint64_t Random64(); std::optional next_float_normal_sample() const; std::optional next_double_normal_sample() const; @@ -34,7 +34,7 @@ class CPUGeneratorImpl final : public GeneratorImpl { void set_next_double_normal_sample(std::optional randn); private: - CPUGeneratorImpl *clone_impl() const override; + CPUGeneratorImpl *CloneImpl() const override; std::mt19937 engine() const { return engine_; } void set_engine(std::mt19937 engine); @@ -45,8 +45,8 @@ class CPUGeneratorImpl final : public GeneratorImpl { std::optional next_double_normal_sample_; }; -const Generator &getDefaultCPUGenerator(); -Generator createCPUGenerator(uint64_t seed); -void manual_seed(uint64_t seed); +const Generator &GetDefaultCpuGenerator(); +Generator CreateCpuGenerator(uint64_t seed); +void ManualSeed(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 index 0b44e3011..9abd97a1f 100644 --- a/infini_train/src/core/runtime/cuda/cuda_generator_impl.cc +++ b/infini_train/src/core/runtime/cuda/cuda_generator_impl.cc @@ -3,7 +3,10 @@ #include #include #include +#include #include +#include +#include #include #include @@ -12,6 +15,10 @@ #include "glog/logging.h" #include "infini_train/include/common/cuda/common_cuda.h" +#include "infini_train/include/core/runtime/generator_backend.h" +#include "infini_train/include/datatype.h" +#include "infini_train/include/device.h" +#include "infini_train/include/generator_impl.h" #include "infini_train/include/tensor.h" namespace infini_train::core::cuda { @@ -25,15 +32,15 @@ constexpr size_t kStateMagicSize = sizeof(kCUDAStateMagic) - 1; constexpr size_t kStateSize = kStateMagicSize + sizeof(uint64_t) * 2; std::once_flag default_generators_init_flag; -std::vector default_generators; +std::vector> default_generators; std::deque default_generator_init_flags; -uint64_t get_non_deterministic_random() { +uint64_t GenerateNonDeterministicSeed() { std::random_device random_device; return (static_cast(random_device()) << 32) | random_device(); } -void init_default_generators() { +void InitDefaultGenerators() { std::call_once(default_generators_init_flag, [] { int device_count = 0; const cudaError_t status = cudaGetDeviceCount(&device_count); @@ -47,8 +54,8 @@ void init_default_generators() { }); } -int resolve_device_index(int8_t device_index) { - init_default_generators(); +int ResolveDeviceIndex(int8_t device_index) { + InitDefaultGenerators(); int index = device_index; if (index == -1) { CUDA_CHECK(cudaGetDevice(&index)); @@ -71,14 +78,14 @@ void CUDAGeneratorImpl::set_current_seed(uint64_t seed) { uint64_t CUDAGeneratorImpl::current_seed() const { return seed_; } -uint64_t CUDAGeneratorImpl::seed() { - const uint64_t random_seed = get_non_deterministic_random(); +uint64_t CUDAGeneratorImpl::Seed() { + const uint64_t random_seed = GenerateNonDeterministicSeed(); set_current_seed(random_seed); return random_seed; } void CUDAGeneratorImpl::set_state(const Tensor &state) { - ::infini_train::detail::check_rng_state(state); + ::infini_train::detail::CheckRngState(state); CHECK_EQ(state.SizeInBytes(), kStateSize); const auto *data = static_cast(state.DataPtr()); @@ -99,9 +106,7 @@ std::shared_ptr CUDAGeneratorImpl::get_state() const { return state; } -Device::DeviceType CUDAGeneratorImpl::device_type() { return Device::DeviceType::kCUDA; } - -uint64_t CUDAGeneratorImpl::philox_subsequence(uint64_t increment) { +uint64_t CUDAGeneratorImpl::ReservePhiloxSubsequence(uint64_t increment) { CHECK_LE(increment, std::numeric_limits::max() - next_philox_subsequence_) << "Philox subsequence counter overflow"; const uint64_t subsequence = next_philox_subsequence_; @@ -109,32 +114,60 @@ uint64_t CUDAGeneratorImpl::philox_subsequence(uint64_t increment) { return subsequence; } -CUDAGeneratorImpl *CUDAGeneratorImpl::clone_impl() const { +CUDAGeneratorImpl *CUDAGeneratorImpl::CloneImpl() const { auto *clone = new CUDAGeneratorImpl(device().index(), seed_); clone->next_philox_subsequence_ = next_philox_subsequence_; return clone; } -const Generator &getDefaultCUDAGenerator(int8_t device_index) { - const int index = resolve_device_index(device_index); +const Generator &GetDefaultCudaGenerator(int8_t device_index) { + const int index = ResolveDeviceIndex(device_index); std::call_once(default_generator_init_flags[index], [index] { - default_generators[index] = createCUDAGenerator(static_cast(index), get_non_deterministic_random()); + default_generators[index].emplace( + CreateCudaGenerator(static_cast(index), GenerateNonDeterministicSeed())); }); - return default_generators[index]; + 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); +Generator CreateCudaGenerator(int8_t device_index, uint64_t seed) { + const int index = ResolveDeviceIndex(device_index); + return MakeGenerator(static_cast(index), seed); } -void manual_seed_all(uint64_t seed) { - init_default_generators(); +void ManualSeedAll(uint64_t seed) { + InitDefaultGenerators(); for (size_t index = 0; index < default_generators.size(); ++index) { - const auto &generator = getDefaultCUDAGenerator(static_cast(index)); - std::lock_guard lock(generator.mutex()); + bool initialized_here = false; + std::call_once(default_generator_init_flags[index], [index, seed, &initialized_here] { + default_generators[index].emplace(CreateCudaGenerator(static_cast(index), seed)); + initialized_here = true; + }); + if (initialized_here) { + continue; + } + const Generator &generator = *default_generators[index]; + std::lock_guard lock(GeneratorAccessor::Mutex(generator)); generator.set_current_seed(seed); } } } // namespace infini_train::core::cuda + +namespace infini_train::core::cuda { +namespace { + +class CUDAGeneratorBackend final : public GeneratorBackend { +public: + Device::DeviceType Type() const override { return Device::DeviceType::kCUDA; } + + Generator Create(const Device &device, uint64_t seed) override { return CreateCudaGenerator(device.index(), seed); } + + const Generator &GetDefault(const Device &device) override { return GetDefaultCudaGenerator(device.index()); } + + void ManualSeedAll(uint64_t seed) override { ::infini_train::core::cuda::ManualSeedAll(seed); } +}; + +INFINI_TRAIN_REGISTER_GENERATOR_BACKEND(Device::DeviceType::kCUDA, CUDAGeneratorBackend); + +} // namespace +} // 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 index 5ffdbdc82..ad9577ad0 100644 --- a/infini_train/src/core/runtime/cuda/cuda_generator_impl.h +++ b/infini_train/src/core/runtime/cuda/cuda_generator_impl.h @@ -4,6 +4,7 @@ #include #include "infini_train/include/generator.h" +#include "infini_train/include/generator_impl.h" namespace infini_train::core::cuda { @@ -14,24 +15,22 @@ class CUDAGeneratorImpl final : public GeneratorImpl { void set_current_seed(uint64_t seed) override; uint64_t current_seed() const override; - uint64_t seed() 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); + uint64_t ReservePhiloxSubsequence(uint64_t increment); private: - CUDAGeneratorImpl *clone_impl() const override; + CUDAGeneratorImpl *CloneImpl() 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); +const Generator &GetDefaultCudaGenerator(int8_t device_index = -1); +Generator CreateCudaGenerator(int8_t device_index, uint64_t seed = Generator::kDefaultSeed); +void ManualSeedAll(uint64_t seed); } // namespace infini_train::core::cuda diff --git a/infini_train/src/core/runtime/generator_backend.cc b/infini_train/src/core/runtime/generator_backend.cc new file mode 100644 index 000000000..1283f989a --- /dev/null +++ b/infini_train/src/core/runtime/generator_backend.cc @@ -0,0 +1,42 @@ +#include "infini_train/include/core/runtime/generator_backend.h" + +#include +#include +#include +#include + +#include "glog/logging.h" + +namespace infini_train { + +GeneratorBackendRegistry &GeneratorBackendRegistry::Instance() { + static GeneratorBackendRegistry instance; + return instance; +} + +void GeneratorBackendRegistry::Register(Device::DeviceType type, std::unique_ptr backend) { + CHECK(backend != nullptr) << "Registering a null GeneratorBackend"; + // DeviceType has no stream insertion operator. + CHECK_EQ(static_cast(type), static_cast(backend->Type())) + << "GeneratorBackend type mismatch: registered as " << static_cast(type) << " but Type() is " + << static_cast(backend->Type()); + CHECK(!backends_.contains(type)) << "GeneratorBackend for device type " << static_cast(type) + << " is already registered"; + + backends_.emplace(type, std::move(backend)); +} + +GeneratorBackend &GeneratorBackendRegistry::Get(Device::DeviceType type) const { + auto it = backends_.find(type); + if (it == backends_.end()) { + throw std::invalid_argument("No GeneratorBackend registered for device type " + + std::to_string(static_cast(type))); + } + return *it->second; +} + +void GeneratorBackendRegistry::ManualSeedAll(uint64_t seed) const { + for (const auto &entry : backends_) { entry.second->ManualSeedAll(seed); } +} + +} // namespace infini_train diff --git a/infini_train/src/generator.cc b/infini_train/src/generator.cc index e4a928c0f..f0693264a 100644 --- a/infini_train/src/generator.cc +++ b/infini_train/src/generator.cc @@ -1,16 +1,15 @@ #include "infini_train/include/generator.h" +#include "infini_train/include/generator_impl.h" +#include #include +#include #include "glog/logging.h" +#include "infini_train/include/core/runtime/generator_backend.h" #include "infini_train/include/datatype.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 { @@ -18,6 +17,12 @@ Generator::Generator(std::shared_ptr impl) : impl_(std::move(impl CHECK(impl_) << "GeneratorImpl with nullptr is not supported"; } +void Generator::set_current_seed(uint64_t seed) const { impl_->set_current_seed(seed); } + +uint64_t Generator::current_seed() const { return impl_->current_seed(); } + +uint64_t Generator::Seed() { return impl_->Seed(); } + void Generator::set_state(const Tensor &state) { CHECK(state.defined()) << "Undefined tensor is not allowed"; impl_->set_state(state); @@ -25,9 +30,13 @@ void Generator::set_state(const Tensor &state) { std::shared_ptr Generator::get_state() const { return impl_->get_state(); } +Device Generator::device() const { return impl_->device(); } + +Generator Generator::Clone() const { return Generator(impl_->Clone()); } + namespace detail { -void check_rng_state(const Tensor &state) { +void CheckRngState(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"; } @@ -35,47 +44,13 @@ void check_rng_state(const Tensor &state) { } // namespace detail Generator CreateGenerator(const Device &device, uint64_t seed) { - if (device.IsCPU()) { - return core::cpu::createCPUGenerator(seed); - } - -#ifdef USE_CUDA - if (device.IsCUDA()) { - return core::cuda::createCUDAGenerator(device.index(), seed); - } -#else - if (device.IsCUDA()) { - throw std::invalid_argument("CUDA generator requested but CUDA support is not enabled"); - } -#endif - - throw std::invalid_argument("Generator can only be created for CPU or CUDA devices"); + return GeneratorBackendRegistry::Instance().Get(device.type()).Create(device, seed); } const Generator &GetDefaultGenerator(const Device &device) { - if (device.IsCPU()) { - return core::cpu::getDefaultCPUGenerator(); - } - -#ifdef USE_CUDA - if (device.IsCUDA()) { - return core::cuda::getDefaultCUDAGenerator(device.index()); - } -#else - if (device.IsCUDA()) { - throw std::invalid_argument("CUDA default generator requested but CUDA support is not enabled"); - } -#endif - - throw std::invalid_argument("Default generator can only be requested for CPU or CUDA devices"); + return GeneratorBackendRegistry::Instance().Get(device.type()).GetDefault(device); } -void manual_seed(uint64_t seed) { - core::cpu::manual_seed(seed); - -#ifdef USE_CUDA - core::cuda::manual_seed_all(seed); -#endif -} +void ManualSeed(uint64_t seed) { GeneratorBackendRegistry::Instance().ManualSeedAll(seed); } } // namespace infini_train diff --git a/infini_train/src/kernels/cpu/cross_entropy.cc b/infini_train/src/kernels/cpu/cross_entropy.cc index f520b2e96..5cf95cf21 100644 --- a/infini_train/src/kernels/cpu/cross_entropy.cc +++ b/infini_train/src/kernels/cpu/cross_entropy.cc @@ -2,6 +2,7 @@ #include #include #include +#include #include #include "glog/logging.h" diff --git a/infini_train/src/kernels/cpu/distribution.cc b/infini_train/src/kernels/cpu/distribution.cc index c0c19fbac..56e332efb 100644 --- a/infini_train/src/kernels/cpu/distribution.cc +++ b/infini_train/src/kernels/cpu/distribution.cc @@ -6,65 +6,78 @@ #include "glog/logging.h" -#include "infini_train/include/common/cpu/distributions_helper.h" #include "infini_train/include/datatype.h" +#include "infini_train/include/device.h" #include "infini_train/include/dispatcher.h" #include "infini_train/include/generator.h" +#include "infini_train/include/generator_impl.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" +#include "infini_train/src/kernels/cpu/distributions_helper.h" namespace infini_train::kernels::cpu { namespace { -template +constexpr int kMaxUniformAttempts = 3; + +template void UniformImpl(Tensor &tensor, double from, double to, core::cpu::CPUGeneratorImpl *generator) { - auto *buf = static_cast(tensor.DataPtr()); - common::cpu::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)); + auto *buf = static_cast(tensor.DataPtr()); + common::cpu::UniformRealDistribution dist(static_cast(from), static_cast(to)); + const StorageT from_value = static_cast(from); + if (from == to) { + for (int64_t i = 0; i < tensor.NumElements(); ++i) { buf[i] = from_value; } + return; + } for (int64_t i = 0; i < tensor.NumElements(); ++i) { - const storage_t value = static_cast(dist(generator)); - // [from, to) is half-open: a sample landing exactly on `to` is mapped back to `from`. - buf[i] = static_cast(value) == to_value ? from_value : value; + StorageT value; + int attempt = 0; + do { + value = static_cast(dist(generator)); + ++attempt; + } while (attempt < kMaxUniformAttempts + && (static_cast(value) < from || static_cast(value) >= to)); + // Bounded retries may still leave out-of-range values in very narrow intervals. + buf[i] = static_cast(value) == to ? from_value : value; } } -template +template void NormalImpl(Tensor &tensor, double mean, double std, core::cpu::CPUGeneratorImpl *generator) { - auto *buf = static_cast(tensor.DataPtr()); - common::cpu::normal_distribution dist(static_cast(mean), static_cast(std)); - for (int64_t i = 0; i < tensor.NumElements(); ++i) { buf[i] = static_cast(dist(generator)); } + auto *buf = static_cast(tensor.DataPtr()); + common::cpu::NormalDistribution dist(static_cast(mean), static_cast(std)); + for (int64_t i = 0; i < tensor.NumElements(); ++i) { buf[i] = static_cast(dist(generator)); } } } // namespace void Uniform(const std::shared_ptr tensor, double from, double to, const std::optional gen) { CHECK(tensor->GetDevice().IsCPU()); - auto *cpu_generator - = get_generator_or_default(gen, core::cpu::getDefaultCPUGenerator()); + auto &cpu_generator = GetGeneratorOrDefault(gen, core::cpu::GetDefaultCpuGenerator(), + tensor->GetDevice()); - std::lock_guard lock(cpu_generator->mutex_); + std::lock_guard lock(cpu_generator.mutex_); core::cpu::DispatchCpuFunc( tensor->Dtype(), - [&]() { - using random_t = std::conditional_t, double, float>; - UniformImpl(*tensor, from, to, cpu_generator); + [&]() { + using RandomT = std::conditional_t, double, float>; + UniformImpl(*tensor, from, to, &cpu_generator); }, "CPU uniform"); } void Normal(const std::shared_ptr tensor, double mean, double std, const std::optional gen) { CHECK(tensor->GetDevice().IsCPU()); - auto *cpu_generator - = get_generator_or_default(gen, core::cpu::getDefaultCPUGenerator()); + auto &cpu_generator = GetGeneratorOrDefault(gen, core::cpu::GetDefaultCpuGenerator(), + tensor->GetDevice()); - std::lock_guard lock(cpu_generator->mutex_); + std::lock_guard lock(cpu_generator.mutex_); core::cpu::DispatchCpuFunc( tensor->Dtype(), - [&]() { - using random_t = std::conditional_t, double, float>; - NormalImpl(*tensor, mean, std, cpu_generator); + [&]() { + using RandomT = std::conditional_t, double, float>; + NormalImpl(*tensor, mean, std, &cpu_generator); }, "CPU normal"); } diff --git a/infini_train/include/common/cpu/distributions_helper.h b/infini_train/src/kernels/cpu/distributions_helper.h similarity index 57% rename from infini_train/include/common/cpu/distributions_helper.h rename to infini_train/src/kernels/cpu/distributions_helper.h index c8604892b..6e2c97e65 100644 --- a/infini_train/include/common/cpu/distributions_helper.h +++ b/infini_train/src/kernels/cpu/distributions_helper.h @@ -1,7 +1,7 @@ #pragma once -// Host-side uniform and normal distributions for generators exposing random() -// and random64(). Box-Muller's second sample is cached when supported by the generator. +// Host-side uniform and normal distributions for generators exposing Random() +// and Random64(). Box-Muller's second sample is cached when supported by the generator. #include #include @@ -14,20 +14,20 @@ namespace infini_train::common::cpu { -template struct uniform_real_distribution { - uniform_real_distribution(T from, T to) : from_(from), to_(to) { +template struct UniformRealDistribution { + UniformRealDistribution(T from, T to) : from_(from), to_(to) { CHECK_LE(from, to); CHECK_LE(to - from, std::numeric_limits::max()); } - uniform_real_distribution(const uniform_real_distribution &) = default; - uniform_real_distribution &operator=(const uniform_real_distribution &) = delete; + UniformRealDistribution(const UniformRealDistribution &) = default; + UniformRealDistribution &operator=(const UniformRealDistribution &) = delete; template T operator()(RNG *generator) const { if constexpr (std::is_same_v) { - return transform(generator->random64()); + return Transform(generator->Random64()); } else { - return transform(generator->random()); + return Transform(generator->Random()); } } @@ -35,17 +35,17 @@ template struct uniform_real_distribution { T from_; T to_; - 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; + template T Transform(V val) const { + constexpr auto kMask = static_cast((static_cast(1) << std::numeric_limits::digits) - 1); + constexpr auto kDivisor = static_cast(1) / (static_cast(1) << std::numeric_limits::digits); + T x = (val & kMask) * kDivisor; return x * (to_ - from_) + from_; } }; template -bool maybe_get_next_normal_sample(RNG *generator, double *ret) { +bool MaybeGetNextNormalSample(RNG *generator, double *ret) { const auto sample = generator->next_double_normal_sample(); if (!sample.has_value()) { return false; @@ -57,7 +57,7 @@ bool maybe_get_next_normal_sample(RNG *generator, double *ret) { template -bool maybe_get_next_normal_sample(RNG *generator, float *ret) { +bool MaybeGetNextNormalSample(RNG *generator, float *ret) { const auto sample = generator->next_float_normal_sample(); if (!sample.has_value()) { return false; @@ -68,34 +68,34 @@ bool maybe_get_next_normal_sample(RNG *generator, float *ret) { } // Fallback: RNG without cache support never has a cached sample. -template bool maybe_get_next_normal_sample(RNG * /*generator*/, void * /*ret*/) { return false; } +template bool MaybeGetNextNormalSample(RNG * /*generator*/, void * /*ret*/) { return false; } template -void maybe_set_next_normal_sample(RNG *generator, const double *cache) { +void MaybeSetNextNormalSample(RNG *generator, const double *cache) { generator->set_next_double_normal_sample(*cache); } template -void maybe_set_next_normal_sample(RNG *generator, const float *cache) { +void MaybeSetNextNormalSample(RNG *generator, const float *cache) { generator->set_next_float_normal_sample(*cache); } // Fallback: RNG without cache support discards the second sample. -template void maybe_set_next_normal_sample(RNG * /*generator*/, const void * /*cache*/) {} +template void MaybeSetNextNormalSample(RNG * /*generator*/, const void * /*cache*/) {} -template struct normal_distribution { - normal_distribution(T mean, T stdv) : mean_(mean), stdv_(stdv) { CHECK_GE(stdv, static_cast(0)); } +template struct NormalDistribution { + NormalDistribution(T mean, T stdv) : mean_(mean), stdv_(stdv) { CHECK_GE(stdv, static_cast(0)); } - normal_distribution(const normal_distribution &) = default; - normal_distribution &operator=(const normal_distribution &) = delete; + NormalDistribution(const NormalDistribution &) = default; + NormalDistribution &operator=(const NormalDistribution &) = delete; template T operator()(RNG *generator) const { T ret; - if (maybe_get_next_normal_sample(generator, &ret)) { + if (MaybeGetNextNormalSample(generator, &ret)) { return ret * stdv_ + mean_; } - uniform_real_distribution uniform(static_cast(0), static_cast(1)); + UniformRealDistribution uniform(static_cast(0), static_cast(1)); const T u1 = uniform(generator); const T u2 = uniform(generator); @@ -104,7 +104,7 @@ template struct normal_distribution { const T theta = kTwoPi * u1; const T sample = r * std::sin(theta); - maybe_set_next_normal_sample(generator, &sample); + MaybeSetNextNormalSample(generator, &sample); ret = r * std::cos(theta); return ret * stdv_ + mean_; diff --git a/infini_train/src/kernels/cpu/dropout.cc b/infini_train/src/kernels/cpu/dropout.cc index a23045a8b..5ea78574d 100644 --- a/infini_train/src/kernels/cpu/dropout.cc +++ b/infini_train/src/kernels/cpu/dropout.cc @@ -2,28 +2,31 @@ #include #include #include +#include #include #include #include "glog/logging.h" -#include "infini_train/include/common/cpu/distributions_helper.h" #include "infini_train/include/datatype.h" +#include "infini_train/include/device.h" #include "infini_train/include/dispatcher.h" #include "infini_train/include/generator.h" +#include "infini_train/include/generator_impl.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" +#include "infini_train/src/kernels/cpu/distributions_helper.h" namespace infini_train::kernels::cpu { namespace { -template +template void DropoutForwardImpl(Tensor &output, Tensor &mask, const Tensor &input, double p, core::cpu::CPUGeneratorImpl *generator) { - auto *output_data = static_cast(output.DataPtr()); + auto *output_data = static_cast(output.DataPtr()); auto *mask_data = static_cast(mask.DataPtr()); - const auto *input_data = static_cast(input.DataPtr()); + const auto *input_data = static_cast(input.DataPtr()); const int64_t n = input.NumElements(); if (p == 0.0) { @@ -36,32 +39,32 @@ void DropoutForwardImpl(Tensor &output, Tensor &mask, const Tensor &input, doubl if (p == 1.0) { for (int64_t index = 0; index < n; ++index) { mask_data[index] = 0; - output_data[index] = static_cast(0.0); + output_data[index] = static_cast(0.0); } return; } - const random_t scale = static_cast(1.0 / (1.0 - p)); - common::cpu::uniform_real_distribution distribution(static_cast(0), static_cast(1)); + const RandomT scale = static_cast(1.0 / (1.0 - p)); + common::cpu::UniformRealDistribution distribution(static_cast(0), static_cast(1)); for (int64_t index = 0; index < n; ++index) { - const bool keep = distribution(generator) >= static_cast(p); + 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); + output_data[index] = keep ? static_cast(static_cast(input_data[index]) * scale) + : static_cast(0.0); } } -template +template void DropoutBackwardImpl(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()); + 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)); + const RandomT 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); + ? static_cast(static_cast(grad_output_data[index]) * scale) + : static_cast(0.0); } } @@ -80,22 +83,22 @@ DropoutForward(const std::shared_ptr input, double p, const std::optiona if (p == 0.0 || p == 1.0) { core::cpu::DispatchCpuFunc( input->Dtype(), - [&]() { - using random_t = std::conditional_t, double, float>; - DropoutForwardImpl(*output, *mask, *input, p, nullptr); + [&]() { + using RandomT = std::conditional_t, double, float>; + DropoutForwardImpl(*output, *mask, *input, p, nullptr); }, "CPU dropout forward"); return {output, mask}; } - auto *cpu_generator - = get_generator_or_default(gen, core::cpu::getDefaultCPUGenerator()); - std::lock_guard lock(cpu_generator->mutex_); + auto &cpu_generator = GetGeneratorOrDefault(gen, core::cpu::GetDefaultCpuGenerator(), + input->GetDevice()); + std::lock_guard lock(cpu_generator.mutex_); core::cpu::DispatchCpuFunc( input->Dtype(), - [&]() { - using random_t = std::conditional_t, double, float>; - DropoutForwardImpl(*output, *mask, *input, p, cpu_generator); + [&]() { + using RandomT = std::conditional_t, double, float>; + DropoutForwardImpl(*output, *mask, *input, p, &cpu_generator); }, "CPU dropout forward"); return {output, mask}; @@ -108,9 +111,9 @@ std::shared_ptr DropoutBackward(const std::shared_ptr grad_outpu core::cpu::DispatchCpuFunc( grad_output->Dtype(), - [&]() { - using random_t = std::conditional_t, double, float>; - DropoutBackwardImpl(*grad_input, *grad_output, *mask, p); + [&]() { + using RandomT = std::conditional_t, double, float>; + DropoutBackwardImpl(*grad_input, *grad_output, *mask, p); }, "CPU dropout backward"); return grad_input; diff --git a/infini_train/src/kernels/cpu/transform.cc b/infini_train/src/kernels/cpu/transform.cc index 48063c7a0..a76102b19 100644 --- a/infini_train/src/kernels/cpu/transform.cc +++ b/infini_train/src/kernels/cpu/transform.cc @@ -2,6 +2,7 @@ #include #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 index a8a50514e..9d17d7ad4 100644 --- a/infini_train/src/kernels/cuda/distribution.cu +++ b/infini_train/src/kernels/cuda/distribution.cu @@ -13,6 +13,7 @@ #include "infini_train/include/core/runtime/device_guard.h" #include "infini_train/include/dispatcher.h" #include "infini_train/include/generator.h" +#include "infini_train/include/generator_impl.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" @@ -22,42 +23,59 @@ 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); +constexpr int kMaxUniformAttempts = 3; + +template __device__ RandomT UniformSample(curandStatePhilox4_32_10_t *state) { + if constexpr (std::is_same_v) { + double sample; + do { + const double2 values = curand_uniform2_double(state); + sample = 1.0 - values.x; + // Subtraction can round a tiny positive draw to 1.0. + } while (sample >= 1.0); + return sample; } 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) { +template __device__ RandomT NormalSample(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) { +template +__global__ void UniformKernel(StorageT *data, int64_t n, double from, double to, uint64_t seed, uint64_t subsequence) { const int64_t index = static_cast(blockIdx.x) * blockDim.x + threadIdx.x; if (index >= n) { return; } + const StorageT from_value = common::cuda::Cast(from); + if (from == to) { + data[index] = from_value; + 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)); - // [from, to) is half-open: a sample landing exactly on `to` is mapped back to `from`. - data[index] = common::cuda::Cast(value) == common::cuda::Cast(to_value) ? from_value : value; + const RandomT sample_from = static_cast(from); + const RandomT sample_to = static_cast(to); + StorageT value; + int attempt = 0; + do { + value = common::cuda::Cast(sample_from + UniformSample(&state) * (sample_to - sample_from)); + ++attempt; + } while (attempt < kMaxUniformAttempts + && (common::cuda::Cast(value) < from || common::cuda::Cast(value) >= to)); + // Bounded retries may still leave out-of-range values in very narrow intervals. + data[index] = common::cuda::Cast(value) == to ? from_value : value; } -template -__global__ void NormalKernel(storage_t *data, int64_t n, random_t mean, random_t std, uint64_t seed, +template +__global__ void NormalKernel(StorageT *data, int64_t n, RandomT mean, RandomT std, uint64_t seed, uint64_t subsequence) { const int64_t index = static_cast(blockIdx.x) * blockDim.x + threadIdx.x; if (index >= n) { @@ -66,10 +84,10 @@ __global__ void NormalKernel(storage_t *data, int64_t n, random_t mean, random_t curandStatePhilox4_32_10_t state; curand_init(seed, subsequence + static_cast(index), 0, &state); - data[index] = common::cuda::Cast(mean + normal_sample(&state) * std); + data[index] = common::cuda::Cast(mean + NormalSample(&state) * std); } -const core::cuda::CudaStream *get_cuda_stream(const Device &device) { +const core::cuda::CudaStream *GetCudaStream(const Device &device) { return dynamic_cast(core::GetDeviceGuardImpl(device.type())->GetStream(device)); } @@ -83,26 +101,25 @@ void Uniform(const std::shared_ptr tensor, double from, double to, const return; } core::DeviceGuard guard(device); - auto *cuda_generator = get_generator_or_default( - gen, core::cuda::getDefaultCUDAGenerator(device.index())); + auto &cuda_generator = GetGeneratorOrDefault( + gen, core::cuda::GetDefaultCudaGenerator(device.index()), device); 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)); + if (from != to) { + std::lock_guard lock(cuda_generator.mutex_); + seed = cuda_generator.current_seed(); + subsequence = cuda_generator.ReservePhiloxSubsequence(static_cast(n)); } const int blocks = static_cast((n + kThreadsPerBlock - 1) / kThreadsPerBlock); - const auto *stream = get_cuda_stream(device); + const auto *stream = GetCudaStream(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); + [&]() { + using RandomT = std::conditional_t, double, float>; + UniformKernel<<cuda_stream()>>>( + static_cast(tensor->DataPtr()), n, from, to, seed, subsequence); }, "CUDA uniform"); CUDA_CHECK(cudaGetLastError()); @@ -116,25 +133,25 @@ void Normal(const std::shared_ptr tensor, double mean, double std, const return; } core::DeviceGuard guard(device); - auto *cuda_generator = get_generator_or_default( - gen, core::cuda::getDefaultCUDAGenerator(device.index())); + auto &cuda_generator = GetGeneratorOrDefault( + gen, core::cuda::GetDefaultCudaGenerator(device.index()), device); 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)); + std::lock_guard lock(cuda_generator.mutex_); + seed = cuda_generator.current_seed(); + subsequence = cuda_generator.ReservePhiloxSubsequence(static_cast(n)); } const int blocks = static_cast((n + kThreadsPerBlock - 1) / kThreadsPerBlock); - const auto *stream = get_cuda_stream(device); + const auto *stream = GetCudaStream(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), + [&]() { + using RandomT = std::conditional_t, double, float>; + NormalKernel<<cuda_stream()>>>( + static_cast(tensor->DataPtr()), n, static_cast(mean), static_cast(std), seed, subsequence); }, "CUDA normal"); diff --git a/infini_train/src/kernels/cuda/dropout.cu b/infini_train/src/kernels/cuda/dropout.cu index fd3b468a9..ad212267a 100644 --- a/infini_train/src/kernels/cuda/dropout.cu +++ b/infini_train/src/kernels/cuda/dropout.cu @@ -15,6 +15,7 @@ #include "infini_train/include/datatype.h" #include "infini_train/include/dispatcher.h" #include "infini_train/include/generator.h" +#include "infini_train/include/generator_impl.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" @@ -25,56 +26,62 @@ 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); +template __device__ RandomT UniformSample(curandStatePhilox4_32_10_t *state) { + if constexpr (std::is_same_v) { + double sample; + do { + const double2 values = curand_uniform2_double(state); + sample = 1.0 - values.x; + // Subtraction can round a tiny positive draw to 1.0. + } while (sample >= 1.0); + return sample; } 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, +template +__global__ void DropoutForwardKernel(StorageT *output, uint8_t *mask, const StorageT *input, int64_t n, RandomT 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)) { + if (p == static_cast(0)) { mask[index] = 1; output[index] = input[index]; return; } - if (p == static_cast(1)) { + if (p == static_cast(1)) { mask[index] = 0; - output[index] = common::cuda::Cast(0.0f); + 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); + const bool keep = UniformSample(&state) >= p; + const RandomT 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); + 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) { +template +__global__ void DropoutBackwardKernel(StorageT *grad_input, const StorageT *grad_output, const uint8_t *mask, int64_t n, + RandomT 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); + const RandomT 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); + ? 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) { +const core::cuda::CudaStream *GetCudaStream(const Device &device) { return dynamic_cast(core::GetDeviceGuardImpl(device.type())->GetStream(device)); } @@ -100,22 +107,22 @@ DropoutForward(const std::shared_ptr input, double p, const std::optiona uint64_t seed = 0; uint64_t subsequence = 0; if (p > 0.0 && p < 1.0) { - auto *cuda_generator = get_generator_or_default( - gen, 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)); + auto &cuda_generator = GetGeneratorOrDefault( + gen, core::cuda::GetDefaultCudaGenerator(device.index()), device); + std::lock_guard lock(cuda_generator.mutex_); + seed = cuda_generator.current_seed(); + subsequence = cuda_generator.ReservePhiloxSubsequence(static_cast(n)); } const int blocks = static_cast((n + kThreadsPerBlock - 1) / kThreadsPerBlock); - const auto *stream = get_cuda_stream(device); + const auto *stream = GetCudaStream(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); + [&]() { + using RandomT = 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()); @@ -136,14 +143,14 @@ std::shared_ptr DropoutBackward(const std::shared_ptr grad_outpu } const int blocks = static_cast((n + kThreadsPerBlock - 1) / kThreadsPerBlock); - const auto *stream = get_cuda_stream(device); + const auto *stream = GetCudaStream(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)); + [&]() { + using RandomT = 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()); diff --git a/infini_train/src/kernels/cuda/elementwise.cu b/infini_train/src/kernels/cuda/elementwise.cu index fc423b35f..a4b6dae83 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 d8b0cffa5..7b3862ac5 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 ef2c95668..d26024b0d 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 c56470e3b..2c080eddc 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 4af04a2d8..a28c46c1d 100644 --- a/infini_train/src/nn/functional.cc +++ b/infini_train/src/nn/functional.cc @@ -2,8 +2,11 @@ #include #include +#include #include +#include "glog/logging.h" + #include "infini_train/include/autograd/activations.h" #include "infini_train/include/autograd/dropout.h" #include "infini_train/include/autograd/elementwise.h" diff --git a/infini_train/src/nn/init.cc b/infini_train/src/nn/init.cc index b93e59806..e64abfa42 100644 --- a/infini_train/src/nn/init.cc +++ b/infini_train/src/nn/init.cc @@ -1,8 +1,14 @@ #include "infini_train/include/nn/init.h" +#include +#include #include #include #include +#include +#include +#include +#include #include #include 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 ab3a80024..b5afd48ba 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/tests/common/test_utils.h b/tests/common/test_utils.h index 06a8bf048..9d5a38e3c 100644 --- a/tests/common/test_utils.h +++ b/tests/common/test_utils.h @@ -1,10 +1,16 @@ #pragma once +#include +#include +#include + #if defined(USE_CUDA) #include #endif +#include "infini_train/include/core/runtime/device_guard.h" #include "infini_train/include/device.h" +#include "infini_train/include/tensor.h" #include "gtest/gtest.h" namespace infini_train { @@ -45,6 +51,16 @@ namespace test { } \ } while (0) +// Returns tensor bytes on the host for exact comparisons. +// Synchronizes the source device before reading the copied data. +inline std::vector TensorBytes(const std::shared_ptr &tensor) { + const Tensor host = tensor->To(Device()); + const Device source_device = tensor->GetDevice(); + core::GetDeviceGuardImpl(source_device.type())->SynchronizeDevice(source_device); + const auto *data = static_cast(host.DataPtr()); + return {data, data + host.SizeInBytes()}; +} + class InfiniTrainTest : public ::testing::TestWithParam { protected: Device GetDevice() const { return Device(GetParam(), 0); } diff --git a/tests/generator/CMakeLists.txt b/tests/generator/CMakeLists.txt index c0107ce2d..ce48367fc 100644 --- a/tests/generator/CMakeLists.txt +++ b/tests/generator/CMakeLists.txt @@ -1,7 +1,4 @@ -# ========================================================================== -# Generator tests -# ========================================================================== - +# Shared CPU/CUDA tests; DDP is built separately below. file(GLOB GENERATOR_SOURCES ${CMAKE_CURRENT_SOURCE_DIR}/test_*.cc) set(GENERATOR_DDP_SOURCE ${CMAKE_CURRENT_SOURCE_DIR}/test_generator_ddp.cc) list(REMOVE_ITEM GENERATOR_SOURCES ${GENERATOR_DDP_SOURCE}) @@ -11,7 +8,26 @@ infini_train_add_test_suite(test_generator TEST_TIMEOUT 40 ) +file(GLOB GENERATOR_CPU_ONLY_SOURCES ${CMAKE_CURRENT_SOURCE_DIR}/cpu_only/test_*.cc) +if(GENERATOR_CPU_ONLY_SOURCES) + infini_train_add_test(test_generator_cpu_only + SOURCES ${GENERATOR_CPU_ONLY_SOURCES} + LABELS cpu + TEST_TIMEOUT 40 + ) +endif() + if(USE_CUDA) + file(GLOB GENERATOR_CUDA_ONLY_SOURCES ${CMAKE_CURRENT_SOURCE_DIR}/cuda_only/test_*.cc) + if(GENERATOR_CUDA_ONLY_SOURCES) + infini_train_add_test(test_generator_cuda_only + SOURCES ${GENERATOR_CUDA_ONLY_SOURCES} + LABELS cuda + TEST_TIMEOUT 40 + ) + endif() + + # DDP uses a separate executable with its own main. add_executable(test_generator_ddp ${GENERATOR_DDP_SOURCE}) target_compile_options(test_generator_ddp PRIVATE -Wno-error) target_link_libraries(test_generator_ddp PRIVATE GTest::gtest) diff --git a/tests/generator/cpu_only/test_generator_cpu.cc b/tests/generator/cpu_only/test_generator_cpu.cc new file mode 100644 index 000000000..5a080ee41 --- /dev/null +++ b/tests/generator/cpu_only/test_generator_cpu.cc @@ -0,0 +1,56 @@ +#include +#include + +#include "gtest/gtest.h" + +#include "infini_train/include/datatype.h" +#include "infini_train/include/device.h" +#include "infini_train/include/generator.h" +#include "infini_train/include/nn/functional.h" + +#include "tests/common/test_utils.h" + +using namespace infini_train; +using infini_train::test::TensorBytes; + +class GeneratorCpuTest : public ::testing::Test {}; + +TEST_F(GeneratorCpuTest, HighBitsOfCPUSeedAffectTheSequence) { + const Device cpu; + constexpr uint64_t low_seed = 17; + constexpr uint64_t high_seed = low_seed + (uint64_t{1} << 32); + EXPECT_NE(TensorBytes(nn::function::Rand({256}, DataType::kFLOAT32, cpu, CreateGenerator(cpu, low_seed))), + TensorBytes(nn::function::Rand({256}, DataType::kFLOAT32, cpu, CreateGenerator(cpu, high_seed)))); +} + +TEST_F(GeneratorCpuTest, StateRestoresCachedNormalSequence) { + const Device cpu; + auto generator = CreateGenerator(cpu, 4567); + nn::function::Randn({3}, DataType::kFLOAT32, cpu, generator); + nn::function::Randn({3}, DataType::kFLOAT64, cpu, generator); + const auto state_after_odd_normal = generator.get_state(); + const auto normal_float = TensorBytes(nn::function::Randn({127}, DataType::kFLOAT32, cpu, generator)); + const auto normal_double = TensorBytes(nn::function::Randn({127}, DataType::kFLOAT64, cpu, generator)); + const auto state_after_replay = TensorBytes(generator.get_state()); + + generator.set_state(*state_after_odd_normal); + EXPECT_EQ(normal_float, TensorBytes(nn::function::Randn({127}, DataType::kFLOAT32, cpu, generator))); + EXPECT_EQ(normal_double, TensorBytes(nn::function::Randn({127}, DataType::kFLOAT64, cpu, generator))); + EXPECT_EQ(state_after_replay, TensorBytes(generator.get_state())); +} + +TEST_F(GeneratorCpuTest, ReseedingClearsBothNormalCaches) { + const Device cpu; + auto generator = CreateGenerator(cpu, 4567); + nn::function::Randn({3}, DataType::kFLOAT32, cpu, generator); + nn::function::Randn({3}, DataType::kFLOAT64, cpu, generator); + + generator.set_current_seed(5678); + auto fresh = CreateGenerator(cpu, 5678); + EXPECT_EQ(TensorBytes(generator.get_state()), TensorBytes(fresh.get_state())); + EXPECT_EQ(TensorBytes(nn::function::Randn({127}, DataType::kFLOAT32, cpu, generator)), + TensorBytes(nn::function::Randn({127}, DataType::kFLOAT32, cpu, fresh))); + EXPECT_EQ(TensorBytes(nn::function::Randn({127}, DataType::kFLOAT64, cpu, generator)), + TensorBytes(nn::function::Randn({127}, DataType::kFLOAT64, cpu, fresh))); + EXPECT_EQ(TensorBytes(generator.get_state()), TensorBytes(fresh.get_state())); +} diff --git a/tests/generator/cuda_only/test_generator_cuda.cc b/tests/generator/cuda_only/test_generator_cuda.cc new file mode 100644 index 000000000..19672ef4a --- /dev/null +++ b/tests/generator/cuda_only/test_generator_cuda.cc @@ -0,0 +1,110 @@ +#include +#include +#include +#include + +#include + +#include "gtest/gtest.h" + +#include "infini_train/include/datatype.h" +#include "infini_train/include/device.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 "tests/common/test_utils.h" + +using namespace infini_train; +using infini_train::test::TensorBytes; + +class GeneratorCudaTest : public ::testing::Test { +protected: + const Device cuda0_{Device::DeviceType::kCUDA, 0}; +}; + +TEST_F(GeneratorCudaTest, StateRejectsCudaAndForeignInput) { + ::testing::GTEST_FLAG(death_test_style) = "threadsafe"; + auto generator = CreateGenerator(cuda0_, 4601); + const auto valid_state = generator.get_state(); + + auto wrong_device = std::make_shared(std::vector{1}, DataType::kUINT8, cuda0_); + EXPECT_DEATH(generator.set_state(*wrong_device), "CPU"); + + const auto cpu_state = CreateGenerator(Device(), 4601).get_state(); + EXPECT_DEATH(generator.set_state(*cpu_state), "Check failed"); + + auto cpu_generator = CreateGenerator(Device(), 4601); + EXPECT_DEATH(cpu_generator.set_state(*valid_state), "Check failed"); +} + +TEST_F(GeneratorCudaTest, CrossBackendGeneratorsAreRejected) { + const Device cpu; + const auto cpu_generator = CreateGenerator(cpu, 8001); + const auto cuda_generator = CreateGenerator(cuda0_, 8002); + + auto cpu_input = std::make_shared(std::vector{8}, DataType::kFLOAT32, cpu); + auto cuda_input = std::make_shared(std::vector{8}, DataType::kFLOAT32, cuda0_); + cpu_input->Fill(1.0f); + cuda_input->Fill(1.0f); + + EXPECT_THROW(nn::function::Rand({8}, DataType::kFLOAT32, cpu, cuda_generator), std::invalid_argument); + EXPECT_THROW(nn::function::Rand({8}, DataType::kFLOAT32, cuda0_, cpu_generator), std::invalid_argument); + EXPECT_THROW(nn::function::Dropout(cpu_input, 0.25, true, cuda_generator), std::invalid_argument); + EXPECT_THROW(nn::function::Dropout(cuda_input, 0.25, true, cpu_generator), std::invalid_argument); +} + +// Regression for the device-index check: a generator bound to cuda:0 must not +// drive a cuda:1 tensor, and the rejected calls must not advance its state. +TEST_F(GeneratorCudaTest, CUDAGeneratorRejectsAnotherCUDADevice) { + REQUIRE_MIN_DEVICES(2); + const Device cuda1(Device::DeviceType::kCUDA, 1); + auto generator = CreateGenerator(cuda0_, 8401); + const auto state_before = TensorBytes(generator.get_state()); + + auto input = std::make_shared(std::vector{257}, DataType::kFLOAT32, cuda1); + input->Fill(1.0f); + + EXPECT_THROW(nn::function::Rand({257}, DataType::kFLOAT32, cuda1, generator), std::invalid_argument); + EXPECT_THROW(nn::function::Randn({257}, DataType::kFLOAT32, cuda1, generator), std::invalid_argument); + EXPECT_THROW(nn::function::Dropout(input, 0.25, true, generator), std::invalid_argument); + + EXPECT_EQ(state_before, TensorBytes(generator.get_state())); +} + +TEST_F(GeneratorCudaTest, DefaultCUDAGeneratorsAreIndependentAcrossAllDevices) { + REQUIRE_MIN_DEVICES(2); + int device_count = 0; + ASSERT_EQ(cudaGetDeviceCount(&device_count), cudaSuccess); + + ManualSeed(8501); + std::vector devices; + devices.reserve(device_count); + for (int index = 0; index < device_count; ++index) { + const Device device(Device::DeviceType::kCUDA, index); + devices.push_back(device); + EXPECT_EQ(GetDefaultGenerator(device).device(), device); + EXPECT_EQ(GetDefaultGenerator(device).current_seed(), 8501U); + } + + for (int active_index = 0; active_index < device_count; ++active_index) { + std::vector> states_before; + states_before.reserve(device_count); + for (const auto &device : devices) { + states_before.push_back(TensorBytes(GetDefaultGenerator(device).get_state())); + } + + auto tensor = std::make_shared(std::vector{64}, DataType::kFLOAT32, devices[active_index]); + nn::init::Uniform(tensor); + + for (int observed_index = 0; observed_index < device_count; ++observed_index) { + const auto state_after = TensorBytes(GetDefaultGenerator(devices[observed_index]).get_state()); + if (observed_index == active_index) { + EXPECT_NE(states_before[observed_index], state_after); + } else { + EXPECT_EQ(states_before[observed_index], state_after); + } + } + } +} diff --git a/tests/generator/test_generator_core.cc b/tests/generator/test_generator_core.cc index 07e36ffcf..87b0a4862 100644 --- a/tests/generator/test_generator_core.cc +++ b/tests/generator/test_generator_core.cc @@ -1,46 +1,27 @@ #include #include -#include +#include #include #include "gtest/gtest.h" -#include "infini_train/include/core/runtime/device_guard.h" #include "infini_train/include/datatype.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/include/device.h" #include "tests/common/test_utils.h" using namespace infini_train; - -namespace { - -std::shared_ptr CopyToCPU(const std::shared_ptr &tensor) { - auto host = std::make_shared(tensor->Dims(), tensor->Dtype(), Device()); - host->CopyFrom(tensor); - if (tensor->GetDevice().IsCUDA()) { - core::GetDeviceGuardImpl(tensor->GetDevice().type())->SynchronizeDevice(tensor->GetDevice()); - } - return host; -} - -std::vector TensorBytes(const std::shared_ptr &tensor) { - const auto host = CopyToCPU(tensor); - const auto *data = static_cast(host->DataPtr()); - return {data, data + host->SizeInBytes()}; -} - -} // namespace +using infini_train::test::TensorBytes; class GeneratorCoreTest : public infini_train::test::InfiniTrainTest {}; TEST_P(GeneratorCoreTest, PublicInterfaceCopyCloneAndSeed) { const Device device = GetDevice(); auto generator = CreateGenerator(device); - EXPECT_TRUE(generator.defined()); EXPECT_EQ(generator.device(), device); EXPECT_EQ(generator.current_seed(), Generator::kDefaultSeed); @@ -57,47 +38,30 @@ TEST_P(GeneratorCoreTest, PublicInterfaceCopyCloneAndSeed) { nn::function::Rand({32}, DataType::kFLOAT32, device, generator); EXPECT_NE(state_before, TensorBytes(alias.get_state())); - auto clone = generator.clone(); + auto clone = generator.Clone(); EXPECT_NE(clone, generator); EXPECT_EQ(TensorBytes(clone.get_state()), TensorBytes(generator.get_state())); clone.set_current_seed(3456); EXPECT_EQ(generator.current_seed(), 2345U); EXPECT_EQ(clone.current_seed(), 3456U); - const uint64_t generated_seed = clone.seed(); + const uint64_t generated_seed = clone.Seed(); EXPECT_EQ(clone.current_seed(), generated_seed); } -TEST_P(GeneratorCoreTest, HighBitsOfCPUSeedAffectTheSequence) { - ONLY_CPU(); - const Device device = GetDevice(); - constexpr uint64_t low_seed = 17; - constexpr uint64_t high_seed = low_seed + (uint64_t{1} << 32); - EXPECT_NE(TensorBytes(nn::function::Rand({256}, DataType::kFLOAT32, device, CreateGenerator(device, low_seed))), - TensorBytes(nn::function::Rand({256}, DataType::kFLOAT32, device, CreateGenerator(device, high_seed)))); -} - -TEST_P(GeneratorCoreTest, StateRestoresUniformAndNormalSequence) { +TEST_P(GeneratorCoreTest, StateRestoresUniformSequence) { const Device device = GetDevice(); auto generator = CreateGenerator(device, 4567); const auto state = generator.get_state(); EXPECT_EQ(state->Dtype(), DataType::kUINT8); - EXPECT_TRUE(state->GetDevice().IsCPU()); + EXPECT_EQ(state->GetDevice(), Device()); const auto uniform = TensorBytes(nn::function::Rand({129}, DataType::kFLOAT32, device, generator)); generator.set_state(*state); EXPECT_EQ(uniform, TensorBytes(nn::function::Rand({129}, DataType::kFLOAT32, device, generator))); - - if (device.IsCPU()) { - nn::function::Randn({3}, DataType::kFLOAT32, device, generator); - const auto state_after_odd_normal = generator.get_state(); - const auto normal = TensorBytes(nn::function::Randn({127}, DataType::kFLOAT32, device, generator)); - generator.set_state(*state_after_odd_normal); - EXPECT_EQ(normal, TensorBytes(nn::function::Randn({127}, DataType::kFLOAT32, device, generator))); - } } -TEST_P(GeneratorCoreTest, StateRejectsMalformedOrForeignInput) { +TEST_P(GeneratorCoreTest, StateRejectsMalformedInput) { ::testing::GTEST_FLAG(death_test_style) = "threadsafe"; const Device device = GetDevice(); auto generator = CreateGenerator(device, 4601); @@ -108,22 +72,13 @@ TEST_P(GeneratorCoreTest, StateRejectsMalformedOrForeignInput) { auto truncated = std::make_shared( std::vector{static_cast(valid_state->SizeInBytes() - 1)}, DataType::kUINT8, Device()); EXPECT_DEATH(generator.set_state(*truncated), "Check failed"); - - if (device.IsCUDA()) { - auto wrong_device = std::make_shared(std::vector{1}, DataType::kUINT8, device); - EXPECT_DEATH(generator.set_state(*wrong_device), "CPU"); - const auto cpu_state = CreateGenerator(Device(), 4601).get_state(); - EXPECT_DEATH(generator.set_state(*cpu_state), "Check failed"); - auto cpu_generator = CreateGenerator(Device(), 4601); - EXPECT_DEATH(cpu_generator.set_state(*valid_state), "Check failed"); - } } TEST_P(GeneratorCoreTest, DefaultExplicitAndUndefinedPathsHaveExpectedState) { const Device device = GetDevice(); constexpr uint64_t seed = 7890; - manual_seed(seed); + ManualSeed(seed); const auto *first_default = &GetDefaultGenerator(device); const auto *second_default = &GetDefaultGenerator(device); EXPECT_EQ(first_default, second_default); @@ -135,54 +90,13 @@ TEST_P(GeneratorCoreTest, DefaultExplicitAndUndefinedPathsHaveExpectedState) { EXPECT_EQ(default_state, TensorBytes(first_default->get_state())); EXPECT_EQ(explicit_result, TensorBytes(nn::function::Rand({256}, DataType::kFLOAT32, device))); - manual_seed(seed + 1); + ManualSeed(seed + 1); const auto default_result = TensorBytes(nn::function::Rand({256}, DataType::kFLOAT32, device)); - manual_seed(seed + 1); - Generator undefined; + ManualSeed(seed + 1); + std::optional undefined = std::nullopt; EXPECT_EQ(default_result, TensorBytes(nn::function::Rand({256}, DataType::kFLOAT32, device, undefined))); } -TEST_P(GeneratorCoreTest, CrossBackendGeneratorsAreRejected) { - ONLY_CPU(); -#if defined(USE_CUDA) - const Device cpu; - const Device cuda(Device::DeviceType::kCUDA, 0); - const auto cpu_generator = CreateGenerator(cpu, 8001); - const auto cuda_generator = CreateGenerator(cuda, 8002); - - auto cpu_input = std::make_shared(std::vector{8}, DataType::kFLOAT32, cpu); - auto cuda_input = std::make_shared(std::vector{8}, DataType::kFLOAT32, cuda); - cpu_input->Fill(1.0f); - cuda_input->Fill(1.0f); - - EXPECT_THROW(nn::function::Rand({8}, DataType::kFLOAT32, cpu, cuda_generator), std::invalid_argument); - EXPECT_THROW(nn::function::Rand({8}, DataType::kFLOAT32, cuda, cpu_generator), std::invalid_argument); - EXPECT_THROW(nn::function::Dropout(cpu_input, 0.25, true, cuda_generator), std::invalid_argument); - EXPECT_THROW(nn::function::Dropout(cuda_input, 0.25, true, cpu_generator), std::invalid_argument); -#else - GTEST_SKIP() << "CUDA disabled"; -#endif -} - -TEST_P(GeneratorCoreTest, CUDAGeneratorCanDriveAnotherCUDADevice) { - ONLY_CUDA(); -#if defined(USE_CUDA) - REQUIRE_MIN_DEVICES(2); - const Device cuda0(Device::DeviceType::kCUDA, 0); - const Device cuda1(Device::DeviceType::kCUDA, 1); - auto generator = CreateGenerator(cuda0, 8401); - const auto state_before = TensorBytes(generator.get_state()); - - const auto output = nn::function::Rand({257}, DataType::kFLOAT32, cuda1, generator); - EXPECT_EQ(output->GetDevice(), cuda1); - EXPECT_NE(state_before, TensorBytes(generator.get_state())); - EXPECT_EQ(TensorBytes(output), - TensorBytes(nn::function::Rand({257}, DataType::kFLOAT32, cuda1, CreateGenerator(cuda0, 8401)))); -#else - GTEST_SKIP() << "CUDA disabled"; -#endif -} - TEST_P(GeneratorCoreTest, InitializerEntryPointsUseExplicitGenerator) { const Device device = GetDevice(); auto make_tensor = [&] { @@ -210,47 +124,4 @@ TEST_P(GeneratorCoreTest, InitializerEntryPointsUseExplicitGenerator) { EXPECT_EQ(TensorBytes(kaiming_a), TensorBytes(kaiming_b)); } -TEST_P(GeneratorCoreTest, DefaultCUDAGeneratorsAreIndependentAcrossAllDevices) { - ONLY_CUDA(); -#if defined(USE_CUDA) - int device_count = 0; - ASSERT_EQ(cudaGetDeviceCount(&device_count), cudaSuccess); - if (device_count < 2) { - GTEST_SKIP() << "requires at least 2 GPUs (found " << device_count << ")"; - } - - manual_seed(8501); - std::vector devices; - devices.reserve(device_count); - for (int index = 0; index < device_count; ++index) { - const Device device(Device::DeviceType::kCUDA, index); - devices.push_back(device); - EXPECT_EQ(GetDefaultGenerator(device).device(), device); - EXPECT_EQ(GetDefaultGenerator(device).current_seed(), 8501U); - } - - for (int active_index = 0; active_index < device_count; ++active_index) { - std::vector> states_before; - states_before.reserve(device_count); - for (const auto &device : devices) { - states_before.push_back(TensorBytes(GetDefaultGenerator(device).get_state())); - } - - auto tensor = std::make_shared(std::vector{64}, DataType::kFLOAT32, devices[active_index]); - nn::init::Uniform(tensor); - - for (int observed_index = 0; observed_index < device_count; ++observed_index) { - const auto state_after = TensorBytes(GetDefaultGenerator(devices[observed_index]).get_state()); - if (observed_index == active_index) { - EXPECT_NE(states_before[observed_index], state_after); - } else { - EXPECT_EQ(states_before[observed_index], state_after); - } - } - } -#else - GTEST_SKIP() << "CUDA disabled"; -#endif -} - INFINI_TRAIN_REGISTER_TEST(GeneratorCoreTest); diff --git a/tests/generator/test_generator_random_ops.cc b/tests/generator/test_generator_random_ops.cc index a807e8d69..9cff02523 100644 --- a/tests/generator/test_generator_random_ops.cc +++ b/tests/generator/test_generator_random_ops.cc @@ -1,5 +1,7 @@ #include +#include #include +#include #include #include #include @@ -14,27 +16,14 @@ #include "infini_train/include/nn/init.h" #include "infini_train/include/tensor.h" +#include "infini_train/include/device.h" #include "tests/common/test_utils.h" using namespace infini_train; +using infini_train::test::TensorBytes; namespace { -std::shared_ptr CopyToCPU(const std::shared_ptr &tensor) { - auto host = std::make_shared(tensor->Dims(), tensor->Dtype(), Device()); - host->CopyFrom(tensor); - if (tensor->GetDevice().IsCUDA()) { - core::GetDeviceGuardImpl(tensor->GetDevice().type())->SynchronizeDevice(tensor->GetDevice()); - } - return host; -} - -std::vector TensorBytes(const std::shared_ptr &tensor) { - const auto host = CopyToCPU(tensor); - const auto *data = static_cast(host->DataPtr()); - return {data, data + host->SizeInBytes()}; -} - template void ExpectUniformRangeAndFinite(const Tensor &tensor) { const auto *data = static_cast(tensor.DataPtr()); for (int64_t index = 0; index < tensor.NumElements(); ++index) { @@ -53,19 +42,21 @@ template void ExpectFinite(const Tensor &tensor) { } void CheckUniformRangeAndFinite(const std::shared_ptr &tensor) { - const auto host = CopyToCPU(tensor); - switch (host->Dtype()) { + const Tensor host = tensor->To(Device()); + const Device source_device = tensor->GetDevice(); + core::GetDeviceGuardImpl(source_device.type())->SynchronizeDevice(source_device); + switch (host.Dtype()) { case DataType::kFLOAT16: - ExpectUniformRangeAndFinite(*host); + ExpectUniformRangeAndFinite(host); break; case DataType::kBFLOAT16: - ExpectUniformRangeAndFinite(*host); + ExpectUniformRangeAndFinite(host); break; case DataType::kFLOAT32: - ExpectUniformRangeAndFinite(*host); + ExpectUniformRangeAndFinite(host); break; case DataType::kFLOAT64: - ExpectUniformRangeAndFinite(*host); + ExpectUniformRangeAndFinite(host); break; default: FAIL() << "Unexpected dtype"; @@ -73,19 +64,21 @@ void CheckUniformRangeAndFinite(const std::shared_ptr &tensor) { } void CheckFinite(const std::shared_ptr &tensor) { - const auto host = CopyToCPU(tensor); - switch (host->Dtype()) { + const Tensor host = tensor->To(Device()); + const Device source_device = tensor->GetDevice(); + core::GetDeviceGuardImpl(source_device.type())->SynchronizeDevice(source_device); + switch (host.Dtype()) { case DataType::kFLOAT16: - ExpectFinite(*host); + ExpectFinite(host); break; case DataType::kBFLOAT16: - ExpectFinite(*host); + ExpectFinite(host); break; case DataType::kFLOAT32: - ExpectFinite(*host); + ExpectFinite(host); break; case DataType::kFLOAT64: - ExpectFinite(*host); + ExpectFinite(host); break; default: FAIL() << "Unexpected dtype"; @@ -105,11 +98,13 @@ std::vector> RunRandomScript(const Device &device, std::opt void ExpectMaskValues(const std::shared_ptr &mask) { EXPECT_EQ(mask->Dtype(), DataType::kUINT8); - const auto host = CopyToCPU(mask); - const auto *data = static_cast(host->DataPtr()); + const Tensor host = mask->To(Device()); + const Device source_device = mask->GetDevice(); + core::GetDeviceGuardImpl(source_device.type())->SynchronizeDevice(source_device); + const auto *data = static_cast(host.DataPtr()); bool saw_zero = false; bool saw_one = false; - for (int64_t index = 0; index < host->NumElements(); ++index) { + for (int64_t index = 0; index < host.NumElements(); ++index) { EXPECT_TRUE(data[index] == 0 || data[index] == 1); saw_zero |= data[index] == 0; saw_one |= data[index] == 1; @@ -130,20 +125,25 @@ template void ExpectUnitInputDropoutValues(const Tensor &output, co void ExpectUnitInputDropoutValues(const std::shared_ptr &output, const std::shared_ptr &mask, double p) { - const auto host_output = CopyToCPU(output); - const auto host_mask = CopyToCPU(mask); - switch (host_output->Dtype()) { + const Tensor host_output = output->To(Device()); + const Tensor host_mask = mask->To(Device()); + // Each async D2H copy must complete before its host buffer is read. + const Device output_device = output->GetDevice(); + core::GetDeviceGuardImpl(output_device.type())->SynchronizeDevice(output_device); + const Device mask_source_device = mask->GetDevice(); + core::GetDeviceGuardImpl(mask_source_device.type())->SynchronizeDevice(mask_source_device); + switch (host_output.Dtype()) { case DataType::kFLOAT16: - ExpectUnitInputDropoutValues(*host_output, *host_mask, p); + ExpectUnitInputDropoutValues(host_output, host_mask, p); break; case DataType::kBFLOAT16: - ExpectUnitInputDropoutValues(*host_output, *host_mask, p); + ExpectUnitInputDropoutValues(host_output, host_mask, p); break; case DataType::kFLOAT32: - ExpectUnitInputDropoutValues(*host_output, *host_mask, p); + ExpectUnitInputDropoutValues(host_output, host_mask, p); break; case DataType::kFLOAT64: - ExpectUnitInputDropoutValues(*host_output, *host_mask, p); + ExpectUnitInputDropoutValues(host_output, host_mask, p); break; default: FAIL() << "Unexpected dtype"; @@ -174,6 +174,46 @@ TEST_P(GeneratorRandomOpsTest, RandAndRandnSupportAllFloatingDtypes) { } } +TEST_P(GeneratorRandomOpsTest, UniformConstantDoesNotAdvanceGenerator) { + const Device device = GetDevice(); + for (const auto dtype : {DataType::kFLOAT16, DataType::kBFLOAT16, DataType::kFLOAT32, DataType::kFLOAT64}) { + SCOPED_TRACE(static_cast(dtype)); + auto generator = CreateGenerator(device, 2101); + nn::function::Rand({11}, dtype, device, generator); + const auto state_before = TensorBytes(generator.get_state()); + auto replay = generator.Clone(); + auto tensor = std::make_shared(std::vector{257}, dtype, device); + auto expected = std::make_shared(tensor->Dims(), dtype, Device()); + expected->Fill(0.1f); + + nn::init::Uniform(tensor, 0.1f, 0.1f, generator); + + EXPECT_EQ(TensorBytes(tensor), TensorBytes(expected)); + EXPECT_EQ(TensorBytes(generator.get_state()), state_before); + EXPECT_EQ(TensorBytes(nn::function::Rand({129}, dtype, device, generator)), + TensorBytes(nn::function::Rand({129}, dtype, device, replay))); + } +} + +TEST_P(GeneratorRandomOpsTest, UniformEmptyTensorDoesNotAdvanceGenerator) { + const Device device = GetDevice(); + for (const auto dtype : {DataType::kFLOAT16, DataType::kBFLOAT16, DataType::kFLOAT32, DataType::kFLOAT64}) { + SCOPED_TRACE(static_cast(dtype)); + auto generator = CreateGenerator(device, 2102); + nn::function::Rand({11}, dtype, device, generator); + const auto state_before = TensorBytes(generator.get_state()); + auto replay = generator.Clone(); + auto tensor = std::make_shared(std::vector{2, 0, 3}, dtype, device); + + nn::init::Uniform(tensor, -2.0f, 3.0f, generator); + + EXPECT_EQ(tensor->NumElements(), 0); + EXPECT_EQ(TensorBytes(generator.get_state()), state_before); + EXPECT_EQ(TensorBytes(nn::function::Rand({129}, dtype, device, generator)), + TensorBytes(nn::function::Rand({129}, dtype, device, replay))); + } +} + TEST_P(GeneratorRandomOpsTest, DropoutSupportsAllFloatingDtypesWithExplicitGenerator) { const std::vector dtypes = {DataType::kFLOAT16, DataType::kBFLOAT16, DataType::kFLOAT32, DataType::kFLOAT64}; @@ -181,7 +221,7 @@ TEST_P(GeneratorRandomOpsTest, DropoutSupportsAllFloatingDtypesWithExplicitGener for (const auto dtype : dtypes) { auto input = nn::function::Rand({4097}, dtype, device, CreateGenerator(device, 3003)); - manual_seed(4004); + ManualSeed(4004); const auto default_state_before = TensorBytes(GetDefaultGenerator(device).get_state()); auto dropout_first = CreateGenerator(device, 5005); @@ -209,19 +249,19 @@ TEST_P(GeneratorRandomOpsTest, SameSeedReplaysRandRandnAndDropoutScript) { TEST_P(GeneratorRandomOpsTest, DefaultGeneratorReplaysRandRandnAndDropoutScript) { const Device device = GetDevice(); - manual_seed(7001); + ManualSeed(7001); const auto expected = RunRandomScript(device, std::nullopt); - manual_seed(7001); + ManualSeed(7001); EXPECT_EQ(expected, RunRandomScript(device, std::nullopt)); - manual_seed(7001); - Generator undefined; + ManualSeed(7001); + std::optional undefined = std::nullopt; EXPECT_EQ(expected, RunRandomScript(device, undefined)); } TEST_P(GeneratorRandomOpsTest, ExplicitRandomGeneratorsDoNotAdvanceDefaultGenerator) { const Device device = GetDevice(); - manual_seed(7101); + ManualSeed(7101); const auto default_state_before = TensorBytes(GetDefaultGenerator(device).get_state()); RunRandomScript(device, CreateGenerator(device, 7102)); @@ -241,8 +281,10 @@ TEST_P(GeneratorRandomOpsTest, OffsetViewsUseTheirOwnDataPointers) { auto uniform_view = std::make_shared(*storage, kOffsetBytes, std::vector{kViewElements}); nn::init::Uniform(uniform_view, 0.0f, 1.0f, CreateGenerator(device, 7201)); - const auto uniform_storage = CopyToCPU(storage); - const auto *uniform_data = static_cast(uniform_storage->DataPtr()); + const Tensor uniform_storage = storage->To(Device()); + const Device storage_device = storage->GetDevice(); + core::GetDeviceGuardImpl(storage_device.type())->SynchronizeDevice(storage_device); + const auto *uniform_data = static_cast(uniform_storage.DataPtr()); for (int64_t index = 0; index < kOffsetElements; ++index) { EXPECT_FLOAT_EQ(uniform_data[index], -1.0f); } for (int64_t index = kOffsetElements; index < kOffsetElements + kViewElements; ++index) { EXPECT_GE(uniform_data[index], 0.0f); @@ -257,14 +299,18 @@ TEST_P(GeneratorRandomOpsTest, OffsetViewsUseTheirOwnDataPointers) { auto input_view = std::make_shared(*input_storage, kOffsetBytes, std::vector{kViewElements}); input_view->Fill(1.0f); const auto output = nn::function::Dropout(input_view, 0.5, true, CreateGenerator(device, 7202)); - const auto host_output = CopyToCPU(output); - const auto *output_data = static_cast(host_output->DataPtr()); + const Tensor host_output = output->To(Device()); + const Device output_device = output->GetDevice(); + core::GetDeviceGuardImpl(output_device.type())->SynchronizeDevice(output_device); + const auto *output_data = static_cast(host_output.DataPtr()); for (int64_t index = 0; index < kViewElements; ++index) { EXPECT_TRUE(output_data[index] == 0.0f || std::abs(output_data[index] - 2.0f) < 1e-6f); } - const auto input_storage_host = CopyToCPU(input_storage); - const auto *input_storage_data = static_cast(input_storage_host->DataPtr()); + const Tensor input_storage_host = input_storage->To(Device()); + const Device input_storage_device = input_storage->GetDevice(); + core::GetDeviceGuardImpl(input_storage_device.type())->SynchronizeDevice(input_storage_device); + const auto *input_storage_data = static_cast(input_storage_host.DataPtr()); for (int64_t index = 0; index < kOffsetElements; ++index) { EXPECT_FLOAT_EQ(input_storage_data[index], -3.0f); } for (int64_t index = kOffsetElements + kViewElements; index < kStorageElements; ++index) { EXPECT_FLOAT_EQ(input_storage_data[index], -3.0f);