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

CMakeLists.txt

@@ -107,7 +107,7 @@ if(USE_CUDA)
   file(GLOB_RECURSE CUDA_KERNELS ${PROJECT_SOURCE_DIR}/infini_train/src/*.cu)
 
   add_library(infini_train_cuda_kernels STATIC ${CUDA_KERNELS})
-  set_target_properties(infini_train_cuda_kernels PROPERTIES CUDA_ARCHITECTURES "75;80;90")
+  set_target_properties(infini_train_cuda_kernels PROPERTIES CUDA_ARCHITECTURES "75;80;90;120")
 
   target_link_libraries(infini_train_cuda_kernels
     PUBLIC

example/gpt2/main.cc

@@ -16,6 +16,7 @@
 #include "infini_train/include/core/runtime/device_guard.h"
 #include "infini_train/include/dataloader.h"
 #include "infini_train/include/device.h"
+#include "infini_train/include/generator.h"
 #include "infini_train/include/lr_scheduler.h"
 #include "infini_train/include/nn/lora/lora_utils.h"
 #include "infini_train/include/nn/modules/loss.h"
@@ -216,9 +217,6 @@ void Train(const nn::parallel::Rank &rank) {
     LOG(INFO) << "total desired batch size: " << FLAGS_total_batch_size
               << " => calculated gradient accumulation steps: " << grad_accum_steps;
 
-    // rng / reproducibility
-    // ManualSeed(42);
-
     // init the model, either from scratch or from OpenAI pretrained checkpoint
     nn::TransformerConfig model_config = gpt2::GPT2Config();
     std::shared_ptr<nn::Module> model = nullptr;
@@ -581,6 +579,7 @@ int main(int argc, char *argv[]) {
     nn::parallel::global::InitAllEnv(FLAGS_nthread_per_process, FLAGS_tensor_parallel, FLAGS_sequence_parallel,
                                      FLAGS_pipeline_parallel, FLAGS_virtual_pipeline_parallel);
     utils::PrecisionCheckEnv::Instance().Init(precision_config);
+    infini_train::ManualSeed(42);
 
     LOG(INFO) << nn::parallel::global::ProcessGroupOverview();
 

example/llama3/main.cc

@@ -15,6 +15,7 @@
 #include "infini_train/include/core/runtime/device_guard.h"
 #include "infini_train/include/dataloader.h"
 #include "infini_train/include/device.h"
+#include "infini_train/include/generator.h"
 #include "infini_train/include/lr_scheduler.h"
 #include "infini_train/include/nn/lora/lora_utils.h"
 #include "infini_train/include/nn/modules/loss.h"
@@ -206,9 +207,6 @@ void Train(const nn::parallel::Rank &rank) {
                   << " => calculated gradient accumulation steps: " << grad_accum_steps;
     }
 
-    // rng / reproducibility
-    // ManualSeed(42);
-
     nn::TransformerConfig model_config = llama3::LLaMA3Config();
     std::shared_ptr<nn::Module> model = nullptr;
     if (!FLAGS_llmc_filepath.empty()) {
@@ -560,6 +558,7 @@ int main(int argc, char *argv[]) {
     nn::parallel::global::InitAllEnv(FLAGS_nthread_per_process, FLAGS_tensor_parallel, FLAGS_sequence_parallel,
                                      FLAGS_pipeline_parallel, FLAGS_virtual_pipeline_parallel);
     utils::PrecisionCheckEnv::Instance().Init(precision_config);
+    infini_train::ManualSeed(42);
 
     LOG(INFO) << nn::parallel::global::ProcessGroupOverview();
 

example/mixtral/main.cc

@@ -13,6 +13,7 @@
 #include "infini_train/include/core/runtime/device_guard.h"
 #include "infini_train/include/dataloader.h"
 #include "infini_train/include/device.h"
+#include "infini_train/include/generator.h"
 #include "infini_train/include/nn/modules/loss.h"
 #include "infini_train/include/nn/modules/transformer/transformer.h"
 #include "infini_train/include/nn/parallel/global.h"
@@ -69,6 +70,7 @@ int main(int argc, char *argv[]) {
         /*sequence_parallel_enabled=*/false,
         /*pipeline_parallel_size=*/1,
         /*virtual_pipeline_parallel_size=*/1);
+    infini_train::ManualSeed(42);
 
     infini_train::nn::TransformerConfig model_config = mixtral::TinyMixtralConfig();
     mixtral::SanitizeTinyMixtralConfig(model_config);

example/mnist/main.cc

@@ -11,6 +11,7 @@
 
 #include "infini_train/include/dataloader.h"
 #include "infini_train/include/device.h"
+#include "infini_train/include/generator.h"
 #include "infini_train/include/nn/modules/loss.h"
 #include "infini_train/include/optimizer.h"
 
@@ -39,6 +40,7 @@ DEFINE_validator(device,
 int main(int argc, char *argv[]) {
     gflags::ParseCommandLineFlags(&argc, &argv, true);
     google::InitGoogleLogging(argv[0]);
+    infini_train::ManualSeed(42);
 
     auto train_dataset = std::make_shared<MNISTDataset>(FLAGS_dataset, true);
     DataLoader train_dataloader(train_dataset, FLAGS_bs);

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/core/runtime/generator_backend.h

@@ -0,0 +1,50 @@
+#pragma once
+
+// Internal Generator backend registration interface. Application code should include generator.h.
+
+#include <cstdint>
+#include <memory>
+#include <unordered_map>
+
+#include "infini_train/include/common/common.h"
+#include "infini_train/include/device.h"
+#include "infini_train/include/generator.h"
+
+namespace infini_train {
+
+class GeneratorBackend {
+public:
+    virtual ~GeneratorBackend() = default;
+
+    virtual Device::DeviceType Type() const = 0;
+    virtual Generator Create(const Device &device, uint64_t seed) = 0;
+    virtual const Generator &GetDefault(const Device &device) = 0;
+    virtual void ManualSeedAll(uint64_t seed) = 0;
+};
+
+// Registers one backend per device type during static initialization.
+class GeneratorBackendRegistry {
+public:
+    static GeneratorBackendRegistry &Instance();
+
+    void Register(Device::DeviceType type, std::unique_ptr<GeneratorBackend> backend);
+    // Throws if no backend is registered for the requested device type.
+    GeneratorBackend &Get(Device::DeviceType type) const;
+    void ManualSeedAll(uint64_t seed) const;
+
+private:
+    GeneratorBackendRegistry() = default;
+    GeneratorBackendRegistry(const GeneratorBackendRegistry &) = delete;
+    GeneratorBackendRegistry &operator=(const GeneratorBackendRegistry &) = delete;
+
+    std::unordered_map<Device::DeviceType, std::unique_ptr<GeneratorBackend>> backends_;
+};
+
+} // namespace infini_train
+
+// CAT expands __COUNTER__ before token pasting so each registration has a unique name.
+#define INFINI_TRAIN_REGISTER_GENERATOR_BACKEND(device_type, class_impl)                                               \
+    [[maybe_unused]] static const bool CAT(infini_train_generator_backend_registered_, __COUNTER__) = []() {           \
+        infini_train::GeneratorBackendRegistry::Instance().Register(device_type, std::make_unique<class_impl>());      \
+        return true;                                                                                                   \
+    }();

infini_train/include/generator.h

@@ -0,0 +1,53 @@
+#pragma once
+
+#include <cstdint>
+#include <memory>
+
+#include "infini_train/include/device.h"
+
+namespace infini_train {
+
+class Tensor;
+class GeneratorImpl;
+
+// A lightweight handle with shared-copy semantics. Use Clone() for an independent state.
+class Generator {
+public:
+    static constexpr uint64_t kDefaultSeed = 67280421310721;
+
+    Generator() = delete;
+
+    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;
+    uint64_t current_seed() const;
+    uint64_t Seed();
+    void set_state(const Tensor &state);
+    std::shared_ptr<Tensor> get_state() const;
+    Device device() const;
+    Generator Clone() const;
+
+    friend bool operator==(const Generator &a, const Generator &b) { return a.impl_ == b.impl_; }
+    friend bool operator!=(const Generator &a, const Generator &b) { return !(a == b); }
+
+private:
+    friend class GeneratorAccessor;
+
+    explicit Generator(std::shared_ptr<GeneratorImpl> impl);
+    std::shared_ptr<GeneratorImpl> impl_;
+};
+
+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);
+
+// Initializes or resets the default generators for all enabled devices with the given seed.
+void ManualSeed(uint64_t seed);
+
+} // namespace infini_train

infini_train/include/generator_impl.h

@@ -0,0 +1,90 @@
+#pragma once
+
+// Internal Generator implementation interface. Application code should include generator.h.
+
+#include <cstdint>
+#include <memory>
+#include <mutex>
+#include <optional>
+#include <stdexcept>
+#include <type_traits>
+#include <utility>
+
+#include "infini_train/include/device.h"
+#include "infini_train/include/generator.h"
+#include "infini_train/include/tensor.h"
+
+namespace infini_train {
+
+class Tensor;
+
+class GeneratorImpl {
+public:
+    explicit GeneratorImpl(Device device) : device_(device) {}
+    virtual ~GeneratorImpl() = default;
+
+    GeneratorImpl(const GeneratorImpl &other) = delete;
+    GeneratorImpl(GeneratorImpl &&other) = delete;
+    GeneratorImpl &operator=(const GeneratorImpl &other) = delete;
+    GeneratorImpl &operator=(GeneratorImpl &&other) = delete;
+
+    virtual void set_current_seed(uint64_t seed) = 0;
+    virtual uint64_t current_seed() const = 0;
+    virtual uint64_t Seed() = 0;
+    virtual void set_state(const Tensor &state) = 0;
+    virtual std::shared_ptr<Tensor> get_state() const = 0;
+
+    std::shared_ptr<GeneratorImpl> Clone() const { return std::shared_ptr<GeneratorImpl>(CloneImpl()); }
+    Device device() const { return device_; }
+
+    // Callers must hold this mutex when accessing shared RNG state concurrently.
+    std::mutex mutex_;
+
+protected:
+    Device device_;
+    virtual GeneratorImpl *CloneImpl() const = 0;
+};
+
+class GeneratorAccessor {
+public:
+    static GeneratorImpl &Get(const Generator &generator) {
+        if (!generator.impl_) {
+            throw std::invalid_argument("Undefined Generator");
+        }
+        return *generator.impl_;
+    }
+
+    static std::mutex &Mutex(const Generator &generator) { return Get(generator).mutex_; }
+
+    static Generator FromImpl(std::shared_ptr<GeneratorImpl> impl) { return Generator(std::move(impl)); }
+};
+
+template <typename T> T &CheckedGeneratorImpl(const Generator &generator, const Device &expected_device) {
+    static_assert(std::is_base_of_v<GeneratorImpl, T>);
+    auto &base = GeneratorAccessor::Get(generator);
+    if (base.device() != expected_device) {
+        throw std::invalid_argument("Generator device mismatch");
+    }
+    auto *typed = dynamic_cast<T *>(&base);
+    if (typed == nullptr) {
+        throw std::invalid_argument("Generator backend mismatch");
+    }
+    return *typed;
+}
+
+template <class Impl, class... Args> Generator MakeGenerator(Args &&...args) {
+    return GeneratorAccessor::FromImpl(std::make_shared<Impl>(std::forward<Args>(args)...));
+}
+
+template <typename T>
+T &GetGeneratorOrDefault(const std::optional<Generator> &generator, const Generator &default_generator,
+                         const Device &expected_device) {
+    const Generator &chosen = generator.has_value() ? *generator : default_generator;
+    return CheckedGeneratorImpl<T>(chosen, expected_device);
+}
+
+namespace detail {
+void CheckRngState(const Tensor &state);
+} // namespace detail
+
+} // 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);