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

infini_train/include/generator.h

@@ -0,0 +1,46 @@
+#pragma once
+
+#include <cstdint>
+#include <memory>
+#include <vector>
+
+#include "infini_train/include/device.h"
+
+namespace infini_train {
+
+class GeneratorImpl;
+namespace detail {
+class GeneratorAccessor;
+}
+
+class Generator {
+public:
+    void ManualSeed(uint64_t seed);
+    uint64_t Seed();
+    uint64_t InitialSeed() const;
+    std::vector<uint8_t> GetState() const;
+    void SetState(const std::vector<uint8_t> &state);
+    // The CUDA index is the generator's home index. An explicit CUDA generator
+    // may be consumed by tensors on another CUDA index.
+    Device GetDevice() const;
+
+private:
+    explicit Generator(std::shared_ptr<GeneratorImpl> impl);
+
+    std::shared_ptr<GeneratorImpl> impl_;
+
+    friend class detail::GeneratorAccessor;
+    friend std::shared_ptr<Generator> MakeCPUGenerator(uint64_t seed);
+    friend std::shared_ptr<Generator> MakeCUDAGenerator(int8_t device_index, uint64_t seed);
+};
+
+std::shared_ptr<Generator> MakeCPUGenerator(uint64_t seed = 42);
+std::shared_ptr<Generator> MakeCUDAGenerator(int8_t device_index, uint64_t seed = 42);
+std::shared_ptr<Generator> GetDefaultCPUGenerator();
+std::shared_ptr<Generator> GetDefaultCUDAGenerator(int8_t device_index);
+std::shared_ptr<Generator> GetDefaultGenerator(const Device &device);
+// ManualSeed is an alias for ManualSeedAll.
+void ManualSeed(uint64_t seed);
+void ManualSeedAll(uint64_t seed);
+
+} // namespace infini_train

infini_train/include/nn/functional.h

@@ -4,9 +4,12 @@
 #include <memory>
 #include <vector>
 
+#include "infini_train/include/device.h"
+
 namespace infini_train {
+class Generator;
 class Tensor;
-}
+} // namespace infini_train
 
 namespace infini_train::nn::function {
 
@@ -47,6 +50,12 @@ 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);
 
+std::shared_ptr<Tensor> Rand(const std::vector<int64_t> &size, Device device = Device(),
+                             std::shared_ptr<Generator> generator = nullptr);
+
+std::shared_ptr<Tensor> Randn(const std::vector<int64_t> &size, Device device = Device(),
+                              std::shared_ptr<Generator> generator = nullptr);
+
 // Returns a new tensor with the reciprocal of the elements of input.
 //
 // Args:

infini_train/include/nn/init.h

@@ -1,21 +1,20 @@
 #pragma once
 
+#include <cstdint>
 #include <memory>
-#include <optional>
-#include <random>
 #include <utility>
 
 #include "infini_train/include/datatype.h"
 #include "infini_train/include/device.h"
 
 namespace infini_train {
+class Generator;
 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::shared_ptr<Tensor> Normal(const std::shared_ptr<Tensor> &tensor, float mean = 0.0f, float stddev = 1.0f,
+                               std::shared_ptr<Generator> generator = nullptr);
 
 std::pair<int64_t, int64_t> CalculateFanInAndFanOut(const std::shared_ptr<Tensor> &tensor);
 
@@ -42,10 +41,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::shared_ptr<Generator> generator = nullptr);
 
 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::shared_ptr<Generator> generator = nullptr);
 
 std::shared_ptr<Tensor> Ones(const std::shared_ptr<Tensor> &tensor);
 

infini_train/include/nn/parallel/ddp/distributed_data_parallel.h

@@ -38,6 +38,7 @@ class DistributedDataParallel : public nn::Module {
     const std::vector<std::shared_ptr<ParamAndGradBucketGroup>> &bucket_groups() const { return bucket_groups_; }
 
 private:
+    void SynchronizeModuleState();
     void BuildParamAndGradBuffers();
     void RegisterBackwardHooks();
     void OnGradReady(const std::shared_ptr<Tensor> &param);

infini_train/include/nn/parallel/process_group.h

@@ -134,25 +134,26 @@ class ProcessGroupFactory {
     template <typename Creator, typename = std::enable_if_t<std::is_invocable_v<Creator>>>
     const ProcessGroup *GetOrCreate(const std::string &name, Creator &&creator) {
         std::unique_lock<std::mutex> lock(mutex_);
-        auto [it, inserted] = name_to_group_.emplace(name, nullptr);
+        const bool inserted = name_to_group_.emplace(name, nullptr).second;
         if (!inserted) {
-            while (it->second == nullptr) { cond_.wait(lock); }
-            return it->second.get();
+            cond_.wait(lock, [this, &name]() { return name_to_group_.at(name) != nullptr; });
+            return name_to_group_.at(name).get();
         }
 
         lock.unlock();
         auto new_group = creator();
         lock.lock();
 
-        it->second = std::move(new_group);
+        auto &group = name_to_group_.at(name);
+        group = std::move(new_group);
         cond_.notify_all();
-        return it->second.get();
+        return group.get();
     }
 
 private:
     // TODO(dcj): maybe RWLock later?
     mutable std::mutex mutex_;
-    std::condition_variable cond_;
+    mutable std::condition_variable cond_;
     std::unordered_map<std::string, std::unique_ptr<ProcessGroup>> name_to_group_;
     Device::DeviceType backend_ = Device::DeviceType::kInvalid;
 };

infini_train/include/tensor.h

@@ -4,7 +4,6 @@
 #include <iostream>
 #include <memory>
 #include <optional>
-#include <random>
 #include <vector>
 
 #include "Eigen/Dense"
@@ -15,6 +14,7 @@
 #include "infini_train/include/scalar.h"
 
 namespace infini_train {
+class Generator;
 namespace autograd {
 class Function;
 class AccumulateGrad;
@@ -150,8 +150,7 @@ class Tensor : public std::enable_shared_from_this<Tensor> {
     std::shared_ptr<Tensor> Unsqueeze(int64_t dim);
 
     // distribution
-    std::shared_ptr<Tensor> Uniform(float from = 0.0f, float to = 1.0f,
-                                    std::optional<std::mt19937> generator = std::nullopt);
+    std::shared_ptr<Tensor> Uniform(float from = 0.0f, float to = 1.0f, std::shared_ptr<Generator> generator = nullptr);
 
     std::shared_ptr<Tensor> Matmul(const std::shared_ptr<Tensor> &other);
     std::shared_ptr<Tensor> Outer(const std::shared_ptr<Tensor> &other);