diff --git a/infini_train/include/autograd/dropout.h b/infini_train/include/autograd/dropout.h new file mode 100644 index 00000000..a6573e1d --- /dev/null +++ b/infini_train/include/autograd/dropout.h @@ -0,0 +1,34 @@ +#pragma once + +#include +#include +#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/common/cpu/distributions_helper.h b/infini_train/include/common/cpu/distributions_helper.h new file mode 100644 index 00000000..c8604892 --- /dev/null +++ b/infini_train/include/common/cpu/distributions_helper.h @@ -0,0 +1,118 @@ +#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 + +#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 * std::numbers::pi_v); + 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/generator.h b/infini_train/include/generator.h new file mode 100644 index 00000000..ee107c71 --- /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/include/nn/functional.h b/infini_train/include/nn/functional.h index e4354fd1..7ed6d047 100644 --- a/infini_train/include/nn/functional.h +++ b/infini_train/include/nn/functional.h @@ -2,8 +2,13 @@ #include #include +#include #include +#include "infini_train/include/datatype.h" +#include "infini_train/include/device.h" +#include "infini_train/include/generator.h" + namespace infini_train { class Tensor; } @@ -47,6 +52,19 @@ std::shared_ptr Triu(const std::shared_ptr &input, int64_t diago // A tensor of the given shape filled with the scalar value 1. std::shared_ptr Ones(const std::vector size); +// Returns a tensor with uniformly distributed random values in [0, 1). +std::shared_ptr Rand(const std::vector &size, DataType dtype = DataType::kFLOAT32, + Device device = Device(), std::optional generator = std::nullopt, + bool requires_grad = false); + +// Returns a tensor with normally distributed random values with mean 0 and standard deviation 1. +std::shared_ptr Randn(const std::vector &size, DataType dtype = DataType::kFLOAT32, + Device device = Device(), std::optional generator = std::nullopt, + bool requires_grad = false); + +std::shared_ptr Dropout(const std::shared_ptr &input, double p = 0.5, bool training = true, + std::optional generator = std::nullopt); + // Returns a new tensor with the reciprocal of the elements of input. // // Args: diff --git a/infini_train/include/nn/init.h b/infini_train/include/nn/init.h index fc6effec..bbc76c8a 100644 --- a/infini_train/include/nn/init.h +++ b/infini_train/include/nn/init.h @@ -1,21 +1,21 @@ #pragma once +#include #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; -class Device; } // namespace infini_train 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/nn/parallel/ddp/distributed_data_parallel.h b/infini_train/include/nn/parallel/ddp/distributed_data_parallel.h index 823ae82b..1b2d3ffa 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/include/tensor.h b/infini_train/include/tensor.h index dcfd8927..05b29263 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/autograd/dropout.cc b/infini_train/src/autograd/dropout.cc new file mode 100644 index 00000000..5b40c7f6 --- /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/core/runtime/cpu/cpu_generator_impl.cc b/infini_train/src/core/runtime/cpu/cpu_generator_impl.cc new file mode 100644 index 00000000..cb49fb3b --- /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 00000000..22a766ec --- /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 00000000..0b44e301 --- /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 00000000..5ffdbdc8 --- /dev/null +++ b/infini_train/src/core/runtime/cuda/cuda_generator_impl.h @@ -0,0 +1,37 @@ +#pragma once + +#include +#include + +#include "infini_train/include/generator.h" + +namespace infini_train::core::cuda { + +class CUDAGeneratorImpl final : public GeneratorImpl { +public: + explicit CUDAGeneratorImpl(int8_t device_index, uint64_t seed = Generator::kDefaultSeed); + ~CUDAGeneratorImpl() override = default; + + void set_current_seed(uint64_t seed) override; + uint64_t current_seed() const override; + uint64_t seed() override; + void set_state(const Tensor &state) override; + std::shared_ptr get_state() const override; + + static Device::DeviceType device_type(); + + // The caller must hold mutex_ while reserving Philox subsequences. + uint64_t philox_subsequence(uint64_t increment); + +private: + CUDAGeneratorImpl *clone_impl() const override; + + uint64_t seed_ = Generator::kDefaultSeed; + uint64_t next_philox_subsequence_ = 0; +}; + +const Generator &getDefaultCUDAGenerator(int8_t device_index = -1); +Generator createCUDAGenerator(int8_t device_index, uint64_t seed = Generator::kDefaultSeed); +void manual_seed_all(uint64_t seed); + +} // namespace infini_train::core::cuda diff --git a/infini_train/src/generator.cc b/infini_train/src/generator.cc new file mode 100644 index 00000000..e4a928c0 --- /dev/null +++ b/infini_train/src/generator.cc @@ -0,0 +1,81 @@ +#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" + +#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 diff --git a/infini_train/src/kernels/cpu/distribution.cc b/infini_train/src/kernels/cpu/distribution.cc new file mode 100644 index 00000000..94761400 --- /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/cpu/dropout.cc b/infini_train/src/kernels/cpu/dropout.cc new file mode 100644 index 00000000..d4e62eb7 --- /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/distribution.cu b/infini_train/src/kernels/cuda/distribution.cu new file mode 100644 index 00000000..6469c344 --- /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/kernels/cuda/dropout.cu b/infini_train/src/kernels/cuda/dropout.cu new file mode 100644 index 00000000..da106585 --- /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 c33e2368..4af04a2d 100644 --- a/infini_train/src/nn/functional.cc +++ b/infini_train/src/nn/functional.cc @@ -5,6 +5,7 @@ #include #include "infini_train/include/autograd/activations.h" +#include "infini_train/include/autograd/dropout.h" #include "infini_train/include/autograd/elementwise.h" #include "infini_train/include/autograd/reduction.h" #include "infini_train/include/autograd/softmax.h" @@ -26,6 +27,29 @@ std::shared_ptr Ones(const std::vector size) { return init::Ones(ones); } +std::shared_ptr Rand(const std::vector &size, DataType dtype, Device device, + std::optional generator, bool requires_grad) { + auto result = std::make_shared(size, dtype, device, requires_grad); + return init::Uniform(result, 0.0f, 1.0f, generator); +} + +std::shared_ptr Randn(const std::vector &size, DataType dtype, Device device, + std::optional generator, bool requires_grad) { + auto result = std::make_shared(size, dtype, device, requires_grad); + return init::Normal(result, 0.0f, 1.0f, generator); +} + +std::shared_ptr Dropout(const std::shared_ptr &input, double p, bool training, + std::optional generator) { + CHECK(input != nullptr); + CHECK_GE(p, 0.0) << "dropout probability has to be between 0 and 1, but got " << p; + CHECK_LE(p, 1.0) << "dropout probability has to be between 0 and 1, but got " << p; + if (!training || p == 0.0 || input->NumElements() == 0) { + return input; + } + return std::make_shared(p, std::move(generator))->Apply({input})[0]; +} + std::shared_ptr Reciprocal(const std::shared_ptr &input) { return input->Reciprocal(); } std::shared_ptr Sin(const std::shared_ptr &input) { return input->Sin(); } diff --git a/infini_train/src/nn/init.cc b/infini_train/src/nn/init.cc index 79b4b48b..b93e5980 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/nn/parallel/ddp/distributed_data_parallel.cc b/infini_train/src/nn/parallel/ddp/distributed_data_parallel.cc index b08b64fb..91a87596 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; diff --git a/infini_train/src/tensor.cc b/infini_train/src/tensor.cc index 18ca3d22..0ccf96fb 100644 --- a/infini_train/src/tensor.cc +++ b/infini_train/src/tensor.cc @@ -3,6 +3,7 @@ #include #include #include +#include #include #include @@ -471,7 +472,7 @@ std::shared_ptr Tensor::Outer(const std::shared_ptr &other) { } // distribution -std::shared_ptr Tensor::Uniform(float from, float to, std::optional generator) { +std::shared_ptr Tensor::Uniform(float from, float to, std::optional generator) { return nn::init::Uniform(shared_from_this(), from, to, generator); } diff --git a/tests/CMakeLists.txt b/tests/CMakeLists.txt index 96776585..9e5eb322 100644 --- a/tests/CMakeLists.txt +++ b/tests/CMakeLists.txt @@ -28,5 +28,8 @@ add_subdirectory(dtype) # Transformer architecture tests add_subdirectory(transformer) +# Generator tests +add_subdirectory(generator) + # Checkpoint tests add_subdirectory(checkpoint) diff --git a/tests/generator/CMakeLists.txt b/tests/generator/CMakeLists.txt new file mode 100644 index 00000000..c0107ce2 --- /dev/null +++ b/tests/generator/CMakeLists.txt @@ -0,0 +1,23 @@ +# ========================================================================== +# Generator tests +# ========================================================================== + +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 new file mode 100644 index 00000000..07e36ffc --- /dev/null +++ b/tests/generator/test_generator_core.cc @@ -0,0 +1,256 @@ +#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 "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, 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 = [&] { + 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, 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_ddp.cc b/tests/generator/test_generator_ddp.cc new file mode 100644 index 00000000..51ee2a4b --- /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 new file mode 100644 index 00000000..a807e8d6 --- /dev/null +++ b/tests/generator/test_generator_random_ops.cc @@ -0,0 +1,349 @@ +#include +#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> RunRandomScript(const Device &device, std::optional 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(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)); + const auto second = TensorBytes(nn::function::Rand({256}, DataType::kFLOAT32, device, generator)); + 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 + = {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);