[2026春季][T2-2-1]ChaoticLuna

infini_train/include/autograd/dropout.h

@@ -0,0 +1,34 @@
+#pragma once
+
+#include <memory>
+#include <optional>
+#include <string>
+#include <utility>
+#include <vector>
+
+#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> generator) : Function(kType), p_(p), generator_(std::move(generator)) {}
+
+    std::vector<std::shared_ptr<Tensor>> Forward(const std::vector<std::shared_ptr<Tensor>> &input_tensors) override;
+    void SetupContext(const std::vector<std::shared_ptr<Tensor>> &input_tensors,
+                      const std::vector<std::shared_ptr<Tensor>> &output_tensors) override;
+    std::vector<std::shared_ptr<Tensor>> Backward(const std::vector<std::shared_ptr<Tensor>> &grad_outputs) override;
+
+private:
+    double p_ = 0.0;
+    std::optional<Generator> generator_;
+};
+
+} // namespace infini_train::autograd

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 <cmath>
+#include <cstdint>
+#include <limits>
+#include <numbers>
+#include <optional>
+#include <type_traits>
+
+#include "glog/logging.h"
+
+namespace infini_train::common::cpu {
+
+template <typename T> 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<T>::max());
+    }
+
+    uniform_real_distribution(const uniform_real_distribution &) = default;
+    uniform_real_distribution &operator=(const uniform_real_distribution &) = delete;
+
+    template <typename RNG> T operator()(RNG *generator) const {
+        if constexpr (std::is_same_v<T, double>) {
+            return transform(generator->random64());
+        } else {
+            return transform(generator->random());
+        }
+    }
+
+private:
+    T from_;
+    T to_;
+
+    template <typename V> T transform(V val) const {
+        constexpr auto MASK = static_cast<V>((static_cast<uint64_t>(1) << std::numeric_limits<T>::digits) - 1);
+        constexpr auto DIVISOR = static_cast<T>(1) / (static_cast<uint64_t>(1) << std::numeric_limits<T>::digits);
+        T x = (val & MASK) * DIVISOR;
+        return x * (to_ - from_) + from_;
+    }
+};
+
+template <typename RNG, typename = decltype(&RNG::next_double_normal_sample),
+          typename = decltype(&RNG::set_next_double_normal_sample)>
+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 <typename RNG, typename = decltype(&RNG::next_float_normal_sample),
+          typename = decltype(&RNG::set_next_float_normal_sample)>
+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 <typename RNG> bool maybe_get_next_normal_sample(RNG * /*generator*/, void * /*ret*/) { return false; }
+
+template <typename RNG, typename = decltype(&RNG::set_next_double_normal_sample)>
+void maybe_set_next_normal_sample(RNG *generator, const double *cache) {
+    generator->set_next_double_normal_sample(*cache);
+}
+
+template <typename RNG, typename = decltype(&RNG::set_next_float_normal_sample)>
+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 <typename RNG> void maybe_set_next_normal_sample(RNG * /*generator*/, const void * /*cache*/) {}
+
+template <typename T> struct normal_distribution {
+    normal_distribution(T mean, T stdv) : mean_(mean), stdv_(stdv) { CHECK_GE(stdv, static_cast<T>(0)); }
+
+    normal_distribution(const normal_distribution &) = default;
+    normal_distribution &operator=(const normal_distribution &) = delete;
+
+    template <typename RNG> T operator()(RNG *generator) const {
+        T ret;
+        if (maybe_get_next_normal_sample(generator, &ret)) {
+            return ret * stdv_ + mean_;
+        }
+
+        uniform_real_distribution<T> uniform(static_cast<T>(0), static_cast<T>(1));
+        const T u1 = uniform(generator);
+        const T u2 = uniform(generator);
+
+        const T r = std::sqrt(static_cast<T>(-2.0) * std::log1p(-u2));
+        constexpr T kTwoPi = static_cast<T>(2.0 * std::numbers::pi_v<double>);
+        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

infini_train/include/generator.h

@@ -0,0 +1,130 @@
+#pragma once
+
+#include <cstdint>
+#include <memory>
+#include <mutex>
+#include <optional>
+#include <stdexcept>
+#include <utility>
+
+#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<Tensor> get_state() const = 0;
+
+    std::shared_ptr<GeneratorImpl> clone() const { return std::shared_ptr<GeneratorImpl>(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<GeneratorImpl> 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<Tensor> 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<T>(); this unchecked accessor assumes a matching backend.
+    template <typename T> T *get() const { return static_cast<T *>(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<GeneratorImpl> impl_;
+};
+
+// Internal factory for backend implementations.
+template <class Impl, class... Args> Generator make_generator(Args &&...args) {
+    return Generator(std::make_shared<Impl>(std::forward<Args>(args)...));
+}
+
+template <typename T> 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<T *>(generator.unsafeGetGeneratorImpl());
+    if (impl == nullptr) {
+        throw std::invalid_argument("Generator implementation does not match the requested backend");
+    }
+    return impl;
+}
+
+template <typename T>
+T *get_generator_or_default(const std::optional<Generator> &generator, const Generator &default_generator) {
+    return generator.has_value() && generator->defined() ? check_generator<T>(*generator)
+                                                         : check_generator<T>(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

infini_train/include/nn/functional.h

@@ -2,8 +2,13 @@
 
 #include <cstdint>
 #include <memory>
+#include <optional>
 #include <vector>
 
+#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<Tensor> Triu(const std::shared_ptr<Tensor> &input, int64_t diago
 //   A tensor of the given shape filled with the scalar value 1.
 std::shared_ptr<Tensor> Ones(const std::vector<int64_t> size);
 
+// Returns a tensor with uniformly distributed random values in [0, 1).
+std::shared_ptr<Tensor> Rand(const std::vector<int64_t> &size, DataType dtype = DataType::kFLOAT32,
+                             Device device = Device(), std::optional<Generator> 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<Tensor> Randn(const std::vector<int64_t> &size, DataType dtype = DataType::kFLOAT32,
+                              Device device = Device(), std::optional<Generator> generator = std::nullopt,
+                              bool requires_grad = false);
+
+std::shared_ptr<Tensor> Dropout(const std::shared_ptr<Tensor> &input, double p = 0.5, bool training = true,
+                                std::optional<Generator> generator = std::nullopt);
+
 // Returns a new tensor with the reciprocal of the elements of input.
 //
 // Args:

infini_train/include/nn/init.h

@@ -1,21 +1,21 @@
 #pragma once
 
+#include <cstdint>
 #include <memory>
 #include <optional>
-#include <random>
 #include <utility>
 
 #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<Tensor> Normal(const std::shared_ptr<Tensor> &tensor, float mean = 0.0, float std = 1.0,
-                               std::optional<std::mt19937> generator = std::nullopt);
+                               std::optional<Generator> generator = std::nullopt);
 
 std::pair<int64_t, int64_t> CalculateFanInAndFanOut(const std::shared_ptr<Tensor> &tensor);
 
@@ -42,10 +42,10 @@ enum class NonLinearityType : int8_t {
 std::shared_ptr<Tensor> KaimingUniform(const std::shared_ptr<Tensor> &tensor, float a = 0.0f,
                                        KaimingMode mode = KaimingMode::kFanIn,
                                        NonLinearityType non_linearity = NonLinearityType::kLeakyReLU,
-                                       std::optional<std::mt19937> generator = std::nullopt);
+                                       std::optional<Generator> generator = std::nullopt);
 
 std::shared_ptr<Tensor> Uniform(const std::shared_ptr<Tensor> &tensor, float a = 0.0f, float b = 1.0f,
-                                std::optional<std::mt19937> generator = std::nullopt);
+                                std::optional<Generator> generator = std::nullopt);
 
 std::shared_ptr<Tensor> Ones(const std::shared_ptr<Tensor> &tensor);
 

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<Tensor> &param);
 
 private:
+    void SynchronizeModuleState();
+
     std::shared_ptr<Reducer> reducer_ = nullptr;
 
     DistributedDataParallelConfig ddp_config_;

infini_train/include/tensor.h

@@ -5,13 +5,15 @@
 #include <memory>
 #include <optional>
 #include <random>
+
 #include <vector>
 
 #include "Eigen/Dense"
 #include "glog/logging.h"
 
 #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<Tensor> {
     void *DataPtr();
     const void *DataPtr() const;
 
+    bool defined() const { return buffer_ != nullptr; }
+
     size_t SizeInBytes() const;
 
     const std::vector<int64_t> &Dims() const;
@@ -151,7 +155,7 @@ class Tensor : public std::enable_shared_from_this<Tensor> {
 
     // distribution
     std::shared_ptr<Tensor> Uniform(float from = 0.0f, float to = 1.0f,
-                                    std::optional<std::mt19937> generator = std::nullopt);
+                                    std::optional<Generator> generator = std::nullopt);
 
     std::shared_ptr<Tensor> Matmul(const std::shared_ptr<Tensor> &other);
     std::shared_ptr<Tensor> Outer(const std::shared_ptr<Tensor> &other);