From 22506dc2025965ac40c49951fc4da89ad1c55e99 Mon Sep 17 00:00:00 2001 From: chaos Date: Thu, 11 Jun 2026 22:00:19 +0800 Subject: [PATCH 01/11] Fix example training sanity checks --- .gitignore | 5 ++++ example/common/tokenizer.cc | 17 +++++++++----- example/gpt2/main.cc | 47 +++++++++++++++++++++++-------------- example/mnist/dataset.cc | 2 +- infini_train/src/tensor.cc | 10 ++++---- 5 files changed, 53 insertions(+), 28 deletions(-) diff --git a/.gitignore b/.gitignore index 50b9fa06..ac84e5d4 100644 --- a/.gitignore +++ b/.gitignore @@ -1,7 +1,12 @@ build/ +build-cpu/ +build-full/ .cache/ .vscode/ *.log *.report.rank* *.records.log.rank* + +server_code/ +.claude/ diff --git a/example/common/tokenizer.cc b/example/common/tokenizer.cc index 9541454a..1f6ca0f4 100644 --- a/example/common/tokenizer.cc +++ b/example/common/tokenizer.cc @@ -10,6 +10,7 @@ #include "glog/logging.h" #include "example/common/utils.h" +#include "infini_train/include/autograd/grad_mode.h" #include "infini_train/include/nn/functional.h" #include "infini_train/include/nn/modules/module.h" #include "infini_train/include/tensor.h" @@ -107,6 +108,7 @@ std::string Tokenizer::Decode(uint32_t token_id) const { void Tokenizer::GenerateText(infini_train::nn::Module &model, uint32_t batch_size, uint32_t sequence_length, uint32_t text_length, Device device) const { + CHECK_LE(text_length, sequence_length) << "text_length must be <= sequence_length"; std::vector dims; dims.assign({batch_size, sequence_length}); // x_tensor (FLAGS_batch_size, FLAGS_sequence_length) eq:(4, 64) @@ -121,20 +123,23 @@ void Tokenizer::GenerateText(infini_train::nn::Module &model, uint32_t batch_siz std::cout << "The meaning of life is"; auto x = std::make_shared(x_tensor.To(device)); - uint64_t kRngState = kRngState; + uint64_t rng_state = kRngState; LOG(INFO) << "start generate text:"; auto cpu_device = Device(); for (int t = prompt_len; t < text_length; ++t) { x = std::make_shared(x->To(device)); // CPU->calc device - // TODO(jym): use no_grad forward later - auto logits = model.Forward({x})[0]; - auto logits_orignal = nn::function::Softmax(logits, -1); + std::shared_ptr logits_orignal; + { + infini_train::autograd::NoGradGuard no_grad; + auto logits = model.Forward({x})[0]; + logits_orignal = nn::function::Softmax(logits, -1); + } auto logits_cpu = logits_orignal->To(cpu_device); auto data = logits_cpu.DataPtr(); - auto vocab_size = logits->Dims()[2]; + auto vocab_size = logits_orignal->Dims()[2]; float *probs = static_cast(data) + (t - 1) * vocab_size; - float coin = RandomF32(kRngState); + float coin = RandomF32(rng_state); int next_token = SampleMult(probs, vocab_size, coin); x = std::make_shared(x->To(cpu_device)); // calc device->CPU diff --git a/example/gpt2/main.cc b/example/gpt2/main.cc index c12b5a28..739446fd 100644 --- a/example/gpt2/main.cc +++ b/example/gpt2/main.cc @@ -3,8 +3,10 @@ #include #include #include +#include #include #include +#include #include "gflags/gflags.h" #include "glog/logging.h" @@ -57,12 +59,13 @@ DEFINE_uint32(freq_generate_txt, 10, "frequency of text generation"); DEFINE_uint32(text_length, 64, "the length of the generated text"); // optimization DEFINE_double(learning_rate, 1e-4, "learning rate warmup iterations"); +DEFINE_string(optimizer, "sgd", "optimizer type (sgd/adam)"); DEFINE_bool(use_distributed_optimizer, false, "Whether to enable DistributedOptimizer(only take effects when DP>1)"); // evaluation DEFINE_uint32(val_loss_every, 0, "every how many steps to evaluate val loss?"); DEFINE_uint32(sample_every, 0, "how often to sample from the model?"); // debugging -DEFINE_bool(overfit_single_batch, true, "overfit just one batch of data"); +DEFINE_bool(overfit_single_batch, false, "overfit just one batch of data"); // memory management DEFINE_string(device, "cuda", "device type (cpu/cuda), useless if using parallel training mode"); // parallel @@ -100,6 +103,8 @@ constexpr char kDeviceCPU[] = "cpu"; constexpr char kDeviceCUDA[] = "cuda"; constexpr char kDtypeFP32[] = "float32"; constexpr char kDtypeBF16[] = "bfloat16"; +constexpr char kOptimizerSGD[] = "sgd"; +constexpr char kOptimizerAdam[] = "adam"; // const std::unordered_map kModelToConfigs = { @@ -114,6 +119,9 @@ const std::unordered_map kModelToConfigs = { DEFINE_validator(model, [](const char *, const std::string &value) { return kSupportedModels.contains(value); }); DEFINE_validator(device, [](const char *, const std::string &value) { return value == kDeviceCPU || value == kDeviceCUDA; }); +DEFINE_validator(optimizer, [](const char *, const std::string &value) { + return value == kOptimizerSGD || value == kOptimizerAdam; +}); void Train(const nn::parallel::Rank &rank) { using namespace nn::parallel; @@ -290,9 +298,8 @@ void Train(const nn::parallel::Rank &rank) { tokenizer = std::make_unique(FLAGS_tokenizer_bin); } - // TODO(dcj): support more complex optimizer later - // auto optimizer = optimizers::SGD(model->Parameters(), FLAGS_learning_rate); - auto optimizer_creator = optimizers::SGD::Create(FLAGS_learning_rate); + auto optimizer_creator = FLAGS_optimizer == kOptimizerAdam ? optimizers::Adam::Create(FLAGS_learning_rate) + : optimizers::SGD::Create(FLAGS_learning_rate); std::shared_ptr optimizer = nullptr; if (FLAGS_use_distributed_optimizer) { @@ -306,6 +313,7 @@ void Train(const nn::parallel::Rank &rank) { } auto train_iter = train_loader.begin(); + std::vector, std::shared_ptr>> overfit_batches; std::shared_ptr loss_fn = (tp_world_size > 1) ? std::static_pointer_cast( std::make_shared(model_config.original_vocab_size)) @@ -353,20 +361,20 @@ void Train(const nn::parallel::Rank &rank) { if (pp_world_size == 1) { optimizer->ZeroGrad(); - // if we are trying to overfit a single batch, we reset the loader here - if (FLAGS_overfit_single_batch) { - // train_loader.Reset(); - } - for (int micro_step = 0; micro_step < grad_accum_steps; ++micro_step) { // enable autocast for the current step infini_train::AutocastGuard autocast_guard(device.type(), dtype); // (bs, seq_len), (bs, seq_len) - auto [x, y] = *train_iter; - // if we are trying to overfit a single batch, we reset the loader here by commenting out the line below - // TODO(dcj): support dataloader.reset() later - ++train_iter; + std::shared_ptr x; + std::shared_ptr y; + if (FLAGS_overfit_single_batch && overfit_batches.size() == grad_accum_steps) { + std::tie(x, y) = overfit_batches[micro_step]; + } else { + std::tie(x, y) = *train_iter; + ++train_iter; + if (FLAGS_overfit_single_batch) { overfit_batches.emplace_back(x, y); } + } x = std::make_shared(x->To(device)); y = std::make_shared(y->To(device)); @@ -393,10 +401,15 @@ void Train(const nn::parallel::Rank &rank) { optimizer->Step(); } else { - auto [x, y] = *train_iter; - // if we are trying to overfit a single batch, we reset the loader here by commenting out the line below - // TODO(dcj): support dataloader.reset() later - ++train_iter; + std::shared_ptr x; + std::shared_ptr y; + if (FLAGS_overfit_single_batch && !overfit_batches.empty()) { + std::tie(x, y) = overfit_batches[0]; + } else { + std::tie(x, y) = *train_iter; + ++train_iter; + if (FLAGS_overfit_single_batch) { overfit_batches.emplace_back(x, y); } + } x = std::make_shared(x->To(device)); y = std::make_shared(y->To(device)); diff --git a/example/mnist/dataset.cc b/example/mnist/dataset.cc index ee683f6d..6b8ec08c 100644 --- a/example/mnist/dataset.cc +++ b/example/mnist/dataset.cc @@ -89,7 +89,7 @@ MNISTDataset::MNISTDataset(const std::string &dataset, bool train) std::format("{}/{}-labels-idx1-ubyte", dataset, train ? kTrainPrefix : kTestPrefix))), image_dims_(image_file_.dims.begin() + 1, image_file_.dims.end()), label_dims_(label_file_.dims.begin() + 1, label_file_.dims.end()), - image_size_in_bytes_(kSN3TypeToSize.at(image_file_.type) + image_size_in_bytes_(sizeof(float) * std::accumulate(image_dims_.begin(), image_dims_.end(), 1, std::multiplies())), label_size_in_bytes_(kSN3TypeToSize.at(label_file_.type) * std::accumulate(label_dims_.begin(), label_dims_.end(), 1, std::multiplies())) { diff --git a/infini_train/src/tensor.cc b/infini_train/src/tensor.cc index f7947030..bf7210ad 100644 --- a/infini_train/src/tensor.cc +++ b/infini_train/src/tensor.cc @@ -125,10 +125,11 @@ Eigen::Map> Tensor::Eig Tensor Tensor::To(Device device) { const auto buffer_device = buffer_->GetDevice(); if (device == buffer_device) { - auto new_tensor = Tensor(*this, offset_, dims_); + auto new_tensor = Tensor(*this, 0, dims_); if (grad_) { - new_tensor.grad_ = std::make_unique(*grad_.get(), grad_->offset_, grad_->dims_); + new_tensor.grad_ = std::make_unique(*grad_.get(), 0, grad_->dims_); } + new_tensor.requires_grad_ = requires_grad_; return new_tensor; } @@ -171,10 +172,11 @@ Tensor Tensor::To(Device device) { Tensor Tensor::To(DataType dtype) { if (dtype == dtype_) { - auto new_tensor = Tensor(*this, offset_, dims_); + auto new_tensor = Tensor(*this, 0, dims_); if (grad_) { - new_tensor.grad_ = std::make_unique(*grad_.get(), grad_->offset_, grad_->dims_); + new_tensor.grad_ = std::make_unique(*grad_.get(), 0, grad_->dims_); } + new_tensor.requires_grad_ = requires_grad_; return new_tensor; } From f20308945115a8269561f9e80e10252fa47ffd81 Mon Sep 17 00:00:00 2001 From: chaos Date: Wed, 17 Jun 2026 00:23:48 +0800 Subject: [PATCH 02/11] Add ref agent definition for PyTorch reference research --- .claude/agents/ref.md | 23 +++++++++++++++++++++++ 1 file changed, 23 insertions(+) create mode 100644 .claude/agents/ref.md diff --git a/.claude/agents/ref.md b/.claude/agents/ref.md new file mode 100644 index 00000000..decae7e7 --- /dev/null +++ b/.claude/agents/ref.md @@ -0,0 +1,23 @@ +--- +name: ref +description: 查找参考实现和官方文档，用于对比学习。搜索 PyTorch 源码、API 文档等。 +tools: Read, Grep, Glob, WebSearch, WebFetch +model: sonnet +--- + +你是参考研究员。你的职责是： +1. 在 PyTorch 源码中找到对应的实现 +2. 查阅官方文档，找到 API 说明和设计动机 +3. 对比不同框架的实现差异 +4. 提供参考链接和源码位置 + +输出格式： +- 给出 PyTorch 对应代码的文件路径和关键片段 +- 解释 PyTorch 为什么这样设计 +- 如果有多种实现方式，列出对比 +- 不要修改任何文件 + +使用场景： + +▎ "让 ref 查一下 PyTorch 的 Optimizer 是怎么实现 step() 的" +▎ "让 ref 找找 DistributedDataParallel 的梯度同步逻辑" From 117dcb61796452d152973922d1a03f6eeb32fa08 Mon Sep 17 00:00:00 2001 From: chaos Date: Sat, 11 Jul 2026 17:24:27 +0800 Subject: [PATCH 03/11] feat: add CPU generator infrastructure --- .gitignore | 2 + ...75\350\261\241\351\200\211\351\242\230.md" | 243 ++++++++++++++++++ .../include/core/runtime/dispatch_stub.h | 68 +++++ .../core/runtime/distribution_kernels.h | 27 ++ .../include/core/runtime/distribution_stubs.h | 28 ++ .../core/runtime/distributions_helper.h | 172 +++++++++++++ infini_train/include/generator.h | 141 ++++++++++ infini_train/include/nn/init.h | 8 +- infini_train/include/tensor.h | 4 +- infini_train/src/core/distribution_kernels.cc | 99 +++++++ .../core/runtime/cpu/cpu_generator_impl.cc | 229 +++++++++++++++++ .../src/core/runtime/cpu/cpu_generator_impl.h | 63 +++++ infini_train/src/generator.cc | 30 +++ infini_train/src/kernels/cpu/cross_entropy.cc | 1 + infini_train/src/kernels/cpu/transform.cc | 1 + infini_train/src/nn/init.cc | 65 +---- .../nn/parallel/ddp/param_and_grad_buffer.cc | 1 + infini_train/src/tensor.cc | 3 +- 18 files changed, 1123 insertions(+), 62 deletions(-) create mode 100644 "docs/\343\200\2202026 \346\230\245\345\255\243\344\272\272\345\267\245\346\231\272\350\203\275\345\244\247\350\265\233\343\200\221Generator \346\212\275\350\261\241\351\200\211\351\242\230.md" create mode 100644 infini_train/include/core/runtime/dispatch_stub.h create mode 100644 infini_train/include/core/runtime/distribution_kernels.h create mode 100644 infini_train/include/core/runtime/distribution_stubs.h create mode 100644 infini_train/include/core/runtime/distributions_helper.h create mode 100644 infini_train/include/generator.h create mode 100644 infini_train/src/core/distribution_kernels.cc create mode 100644 infini_train/src/core/runtime/cpu/cpu_generator_impl.cc create mode 100644 infini_train/src/core/runtime/cpu/cpu_generator_impl.h create mode 100644 infini_train/src/generator.cc diff --git a/.gitignore b/.gitignore index ac84e5d4..6ec79bff 100644 --- a/.gitignore +++ b/.gitignore @@ -10,3 +10,5 @@ build-full/ server_code/ .claude/ +pytorch_ref/ +CLAUDE.md diff --git "a/docs/\343\200\2202026 \346\230\245\345\255\243\344\272\272\345\267\245\346\231\272\350\203\275\345\244\247\350\265\233\343\200\221Generator \346\212\275\350\261\241\351\200\211\351\242\230.md" "b/docs/\343\200\2202026 \346\230\245\345\255\243\344\272\272\345\267\245\346\231\272\350\203\275\345\244\247\350\265\233\343\200\221Generator \346\212\275\350\261\241\351\200\211\351\242\230.md" new file mode 100644 index 00000000..6e170aca --- /dev/null +++ "b/docs/\343\200\2202026 \346\230\245\345\255\243\344\272\272\345\267\245\346\231\272\350\203\275\345\244\247\350\265\233\343\200\221Generator \346\212\275\350\261\241\351\200\211\351\242\230.md" @@ -0,0 +1,243 @@ +# 中期报告要求 + +- 比赛期：2026/05/18 - 2026/07/12 （56 天） + - 比赛期第一天，即 2026/05/18 的 12:00 公布赛题 + - 每个参赛小组需在特定的截止日期前提交中期报告。中期报告需包含： + - 小组名称与所有成员 + - 所选的赛题 + - 截至撰写时各个赛题的进度与完成情况 + - 中期报告没有字数要求，简明扼要、表述清晰即可 + - 提交要求与方式： + - 报告文件命名要求：<小组名称>_<赛道名称>_中期报告.pdf + - 提交方式：InfiniTensor 网站上提交，提交流程后续补充 + +# 【2026 春季人工智能大赛】Generator 抽象选题 + +## 一、题目背景 + +在深度学习训练与推理框架中，随机数生成器（Generator）是支撑参数初始化、随机采样、Dropout、噪声注入等功能的重要基础设施。当前框架中相关能力尚不完善，缺少统一的 Generator 抽象、后端实现以及全局随机种子管理机制，导致随机相关算子的行为与主流框架存在差距，也难以满足后续功能扩展需求。 + +当前框架在随机数相关能力上仍不完善，主要表现为： + +- [x] 缺少统一的 Generator 抽象，随机数状态、种子设置、状态获取与恢复等能力未形成标准接口； +- [x] CPU 与 CUDA 后端缺少统一的 Generator 实现，不同设备上的随机行为缺乏一致的调用方式；（CPU ✅，CUDA ❌） +- [x] 缺少全局默认 Generator 机制，随机算子在未显式传入 generator 时无法方便地使用当前设备的默认随机源；（CPU 默认 Generator 已实现，算子层未接入） +- [ ] 缺少统一的全局随机种子控制入口，导致参数初始化、Dropout 等随机算子在多次运行间难以稳定复现； +- [ ] 随机数相关能力与主流框架（如 PyTorch）存在差距，不利于后续算子扩展、模型对齐与测试验证。 + +为补齐这一基础能力，本题目要求参赛者参考主流深度学习框架的设计思路，完成一套可扩展的随机数生成器基础设施，实现 CPU 与 CUDA 后端的 Generator，并支持基于 Generator 的随机数生成与种子控制能力。 + +## 二、题目目标 + +参赛者需要围绕框架内的随机数生成体系，完成以下几个方面的工作： + +- [x] 1. 设计并实现统一的 Generator 抽象接口，支持随机数状态管理、种子设置、状态获取与恢复等基础能力。 +- [x] 2. 分别实现 CPU 与 CUDA 后端对应的 GeneratorImpl，使不同设备上的随机数生成具备统一的调用方式。（CPU ✅，CUDA ❌） +- [ ] 3. 建立全局随机种子控制机制，支持通过统一入口固定随机种子，并使随机相关算子的行为可复现。 +- [ ] 4. 改造随机数相关算子，使其在未显式传入 generator 参数时，能够自动使用当前设备上的默认全局 Generator。 +- [ ] 5. 验证实现结果与主流框架在基本行为上的一致性，包括随机性、可复现性以及跨设备下的接口统一性。 + +## 三、任务拆解 + +### 1. Generator 抽象与状态管理设计 + +在框架中设计统一的 Generator 抽象层，用于屏蔽不同设备后端的随机数实现差异。该抽象应支持： + +- [x] 设置随机种子，如 `ManualSeed()` / `SetCurrentSeed()`； +- [x] 获取当前种子或初始种子，如 `Seed()` / `InitialSeed()`； +- [x] 获取当前随机数状态，如 `GetState()`； +- [x] 恢复随机数状态，如 `SetState(state)`； +- [x] 查询所属设备类型； +- [x] 提供必要的线程安全或状态访问保护机制。 + +实现时需注意区分 seed 与 state：seed 只用于初始化随机序列，state 表示随机序列当前推进到的位置。随机算子每次使用 Generator 后，应推进其内部状态，避免重复生成相同随机序列。 + +推荐设计思路： + +- [x] 提供用户侧可持有的 `Generator` 句柄类； +- [x] 底层以 `GeneratorImpl` 作为多态实现基类； +- [x] CPU / CUDA 分别派生对应实现；（CPU 已实现，CUDA 未实现） +- [x] 公共接口不暴露 `std::mt19937`、curand、Philox 等后端细节； +- [x] 为后续其他平台 Generator 接入预留扩展空间。 + +### 2. CPU / CUDA 后端 Generator 实现 + +基于现有设备后端，分别实现： + +- [x] `CPUGeneratorImpl` +- [ ] `CUDAGeneratorImpl` + +两者应保持统一接口语义，但内部状态组织方式可以不同，且**不要求 CPU 与 CUDA 在相同 seed 下生成逐元素一致的随机结果**。 + +实现时需重点关注： + +- [x] CPU 与 CUDA 随机数状态如何组织与保存；（CPU 已实现二进制序列化格式） +- [ ] CUDA 侧应按设备维度维护独立 Generator； +- [x] `GetState()` / `SetState()` 的状态序列化、反序列化与格式校验；（CPU 已实现） +- [x] 多次随机调用后，随机序列是否连续推进；（CPU 已实现，`random()`/`random64()` 持续推进 `engine_`） +- [x] 同一 seed、同一设备、同一调用顺序下结果是否可复现；（CPU 已实现，mt19937 确定性） +- [ ] 显式传入 Generator 与使用默认 Generator 时，行为是否符合预期。（算子层尚未接入） + +CUDA 后端建议在状态中考虑 seed、offset、counter 等信息，避免不同 kernel 调用之间随机序列重叠。此部分需要保证功能正确、语义清晰、状态可管理、接口一致，为后续随机算子和分布式扩展提供基础。 + +### 3. 默认 Generator 与统一随机种子入口 + +建立框架级默认 Generator 管理机制，支持两种使用方式： + +- 用户显式传入某个 Generator； +- 用户不传入 Generator，系统自动使用当前设备上的默认 Generator。 + +该机制应包括： + +- [x] CPU 默认全局 Generator； +- [ ] 各 CUDA 设备对应的默认 Generator； +- [x] 获取默认 Generator 的统一入口；（`getDefaultCPUGenerator()` 已实现） +- [ ] 统一的全局随机种子设置入口； +- [ ] 设置全局随机种子时，同步更新默认 Generator 的初始状态。 + +要求： + +- [x] 多次获取同一设备默认 Generator 时，应返回同一随机状态来源；（static 局部变量，单例模式） +- [ ] 不同 CUDA 设备的默认 Generator 应相互独立； +- [x] 用户显式传入 Generator 时，必须使用用户指定 Generator；（API 支持 `std::optional`） +- [ ] 未传入 Generator 时，才使用当前设备默认 Generator；（算子层尚未接入） +- [ ] 设置统一 seed 后，参数初始化、Dropout、随机采样等行为应可稳定复现。（缺少全局 seed 入口） + +这一部分是本题的核心验收点。随机数系统最终服务于训练可复现性，因此必须通过统一入口稳定控制随机行为。 + +### 4. 随机相关算子接入改造 + +改造框架中已有的随机相关使用，使其接入 Generator 机制。建议至少覆盖两类场景： + +- [ ] 初始化类随机算子，如 uniform、normal、kaiming 等；（init.cc 仍使用 raw std::mt19937） +- [ ] 训练过程中的随机算子，如 dropout、rand、randn 等； + +改造后的算子应满足： + +- [ ] 支持显式传入 Generator； +- [ ] 未传入 Generator 时，自动使用当前设备默认 Generator； +- [ ] 随机结果受统一 seed 入口控制； +- [x] 多次调用会推进 Generator 状态；（CPUGeneratorImpl 引擎层已支持，算子层未接入） +- [x] 同一 seed、同一调用顺序下结果可复现。（引擎层已支持，算子层未接入） + +不要求一次性改造所有随机算子，但应至少完成一个初始化类算子和一个训练期随机算子，以证明 Generator 机制能够贯通框架层与算子层。 + +### 5. 测试与对齐验证 + +为验证实现结果，需要补充系统化测试。测试重点应放在"语义是否正确"和"是否可复现"，而不是与 PyTorch 逐元素数值一致。 + +建议至少覆盖以下内容： + +#### （1）接口一致性测试 + +验证 CPU / CUDA Generator 是否支持统一接口，包括： + +- [ ] seed 设置与获取； +- [ ] state 获取与恢复； +- [ ] 设备类型查询； +- [ ] 默认 Generator 获取； +- [ ] 显式 Generator 与默认 Generator 两条调用路径。 + +#### （2）种子可复现测试 + +验证统一随机种子入口是否生效，包括： + +- [ ] 同一 seed 下，多次运行参数初始化结果一致； +- [ ] 同一 seed 下，多次运行 Dropout mask 一致； +- [ ] 不同 seed 下结果应发生变化； +- [ ] 同一 seed、同一调用顺序下结果应一致。 + +#### （3）状态恢复测试 + +验证 `GetState()` / `SetState()` 是否真正恢复随机序列，包括： + +- [ ] 保存 state； +- [ ] 继续生成一段随机结果； +- [ ] 恢复 state； +- [ ] 再次生成随机结果； +- [ ] 验证恢复后的结果与原序列对齐。 + +CPU 后端必须覆盖该测试；CUDA 后端若已实现，也应覆盖。 + +#### （4）默认 Generator 行为测试 + +验证随机算子是否正确使用默认 Generator，包括： + +- [ ] 不传 Generator 时，是否使用当前设备默认 Generator； +- [ ] CPU tensor 是否使用 CPU 默认 Generator； +- [ ] CUDA tensor 是否使用对应 CUDA 设备默认 Generator； +- [ ] 显式传入 Generator 时，是否不会误用默认 Generator。 + +#### （5）主流框架语义对齐验证 + +可选取典型场景与 PyTorch 进行语义对比，包括： + +- [ ] 手动设置 seed 后结果可复现； +- [ ] Dropout 等随机运算在同 seed 下可复现； +- [ ] 随机张量生成接口支持显式 Generator 与默认 Generator； +- [ ] state 保存与恢复语义一致。 + +需要明确：本项目不要求底层随机算法与 PyTorch 逐 bit 一致，也不要求 CPU 与 CUDA 随机结果彼此一致。验收重点是接口语义一致、可复现逻辑一致、默认 Generator 行为一致。 + +## 四、提交报告要求 + +除代码外，参赛者应提交以下内容作为技术报告： + +### 1. 功能正确性验证 + +需要提供上述 Generator 基础设施功能验证结果 + +### 2. 对齐性与行为分析报告 + +需要说明本项目实现与主流框架的对齐情况。建议从以下角度展开： + +- [ ] 参考了哪些主流框架设计； +- [ ] 当前实现与 PyTorch 在接口语义上的对应关系； +- [ ] 哪些行为做到基本一致； +- [ ] 哪些行为暂未完全对齐，以及原因分析； +- [ ] 当前实现范围与后续可扩展方向。 + +此部分重点不是要求一切细节完全一致，而是要求留档明确设计思路，能够清晰说明： + +- [ ] 自己实现了什么； +- [ ] 为什么这样设计； +- [ ] 与主流框架相比的设计异同及原因； +- [ ] 后续演进路径是什么。 + +### 3. 测试与可复现性说明 + +需要提供完整测试脚本与说明，保证 reviewer 可在相同环境下复现主要结果。 + +## 五、验收要求 + +### 验收要求 + +项目需满足以下基本要求： + +- [ ] 代码以 PR 形式提交，结构清晰，具备基本可 review 性； +- [x] 完成统一 Generator 抽象设计； +- [x] 实现 CPU Generator，CUDA Generator 建议实现；（CPU ✅，CUDA ❌） +- [x] 建立默认 Generator 管理机制；（CPU 侧已完成） +- [ ] 提供统一的全局随机种子设置入口； +- [ ] 支持默认 Generator 与显式 Generator 两种使用方式；（API 层已支持，算子层未接入） +- [x] 支持 Generator 状态的获取与恢复（get_state / set_state）； +- [ ] 至少改造一类初始化算子和一类框架内随机数生成调用，使其接入 Generator； +- [ ] 提供测试或运行日志，验证随机行为具备基本可复现性（同 seed 一致、不同 seed 不同、状态恢复有效）。 + +### 加分项 + +在满足验收要求的基础上，具备以下内容可作为加分项： + +- [ ] 设计清晰，接口与实现分层合理，便于后续扩展 +- [ ] 测试覆盖充分，包含 seed、state、默认/显式 generator、跨设备等关键场景； +- [ ] 调用处改造完全，接口风格统一； +- [ ] 与 PyTorch 的接口语义和行为分析完整，报告质量高 +- [ ] PR 经过完整 review 流程，达到可合入标准。 + +## 参考链接 + +- https://docs.pytorch.org/docs/stable/generated/torch.Generator.html#generator +- https://github.com/pytorch/pytorch/blob/main/aten/src/ATen/core/Generator.h +- https://github.com/pytorch/pytorch/blob/main/aten/src/ATen/CPUGeneratorImpl.cpp +- https://github.com/pytorch/pytorch/blob/main/aten/src/ATen/cuda/CUDAGeneratorImpl.cpp +- https://docs.pytorch.org/docs/stable/notes/randomness.html#pytorch-random-number-generator diff --git a/infini_train/include/core/runtime/dispatch_stub.h b/infini_train/include/core/runtime/dispatch_stub.h new file mode 100644 index 00000000..13e7c4bb --- /dev/null +++ b/infini_train/include/core/runtime/dispatch_stub.h @@ -0,0 +1,68 @@ +#pragma once + +#include +#include + +#include "glog/logging.h" + +#include "infini_train/include/device.h" + +namespace infini_train { + +// ============================================================ +// DispatchStub — 设备无关的函数指针分发（仿 PyTorch DispatchStub） +// ============================================================ +// 每个 stub 按 DeviceType 存一组函数指针，调用时按设备查表分发。 +// 新增后端只需在 DeviceType 枚举加一项，kCount 自动适配。 +// +// 使用： +// 1. DECLARE_DISPATCH(fn_type, name) — 声明 extern 全局 stub +// 2. DEFINE_DISPATCH(name) — 定义全局 stub 实例 +// 3. REGISTER_DISPATCH(name, dt, fn) — 后端注册函数指针 +// 4. name(device_type, args...) — 调用，自动分发 +// +template +class DispatchStub { +public: + using FnType = FnPtr; + + static constexpr size_t kNumDevices = static_cast(Device::DeviceType::kCount); + + void register_kernel(Device::DeviceType dt, FnPtr fn) { + table_[static_cast(dt)] = fn; + } + + template + auto operator()(Device::DeviceType dt, Args&&... args) const { + auto idx = static_cast(dt); + CHECK(idx < kNumDevices) << "Invalid device type " << static_cast(dt); + CHECK(table_[idx] != nullptr) << "Dispatch kernel not registered for device type " + << static_cast(dt); + return (*table_[idx])(std::forward(args)...); + } + +private: + FnPtr table_[kNumDevices] = {}; +}; + +// ---- 宏 ---- + +// 声明：放在头文件里，告诉其他编译单元"这个 stub 存在" +#define DECLARE_DISPATCH(fn_type, name) \ + extern DispatchStub name + +// 定义：放在一个 .cpp 里，分配实际存储空间 +#define DEFINE_DISPATCH(name) \ + DispatchStub name + +// 注册：放在后端 .cpp/.cu 里，静态初始化时把函数填进表 +#define INFINI_TRAIN_CONCAT_IMPL(x, y) x##y +#define INFINI_TRAIN_CONCAT(x, y) INFINI_TRAIN_CONCAT_IMPL(x, y) + +#define REGISTER_DISPATCH(name, device_type, fn) \ + static const bool INFINI_TRAIN_CONCAT(name##_registered_, __COUNTER__) = []() { \ + (name).register_kernel((device_type), (fn)); \ + return true; \ + }(); + +} // namespace infini_train diff --git a/infini_train/include/core/runtime/distribution_kernels.h b/infini_train/include/core/runtime/distribution_kernels.h new file mode 100644 index 00000000..e779f0fe --- /dev/null +++ b/infini_train/include/core/runtime/distribution_kernels.h @@ -0,0 +1,27 @@ +#pragma once + +/// distribution_kernels.h +/// +/// 公共 API：uniform_kernel / normal_kernel +/// +/// init.cc 等上层代码调用这些包装函数，内部通过 dispatch_stub 自动分发到 +/// 正确的设备后端（CPU / 未来 CUDA）。 +/// +/// 仿 PyTorch aten/src/ATen/native/Distributions.cpp 中的 struct UniformStub 等包装层。 + +#include +#include + +#include "infini_train/include/device.h" +#include "infini_train/include/generator.h" + +namespace infini_train { + +// 填 n 个 float 到 data +void uniform_kernel(void *data, int64_t n, double from, double to, + Device::DeviceType device_type, const std::optional &gen); + +void normal_kernel(void *data, int64_t n, double mean, double std, + Device::DeviceType device_type, const std::optional &gen); + +} // namespace infini_train diff --git a/infini_train/include/core/runtime/distribution_stubs.h b/infini_train/include/core/runtime/distribution_stubs.h new file mode 100644 index 00000000..94d49339 --- /dev/null +++ b/infini_train/include/core/runtime/distribution_stubs.h @@ -0,0 +1,28 @@ +#pragma once + +/// distribution_stubs.h +/// +/// 内部细节：声明 uniform 和 normal 两个 DispatchStub。 +/// 上层代码不应直接调用这些 stub，而应使用 distribution_kernels.h +/// 中的 uniform_kernel() / normal_kernel() 包装函数。 +/// +/// 仿 PyTorch aten/src/ATen/native/UnaryOps.h 中的 DECLARE_DISPATCH 声明。 + +#include +#include + +#include "infini_train/include/core/runtime/dispatch_stub.h" +#include "infini_train/include/generator.h" + +namespace infini_train { + +// 填 n 个 float 到 data 指向的缓冲区 +DECLARE_DISPATCH(void (*)(void *data, int64_t n, double a, double b, + const std::optional &gen), + uniform_stub); + +DECLARE_DISPATCH(void (*)(void *data, int64_t n, double a, double b, + const std::optional &gen), + normal_stub); + +} // namespace infini_train diff --git a/infini_train/include/core/runtime/distributions_helper.h b/infini_train/include/core/runtime/distributions_helper.h new file mode 100644 index 00000000..f82f9065 --- /dev/null +++ b/infini_train/include/core/runtime/distributions_helper.h @@ -0,0 +1,172 @@ +#pragma once + +#include +#include +#include +#include +#include +#include + +namespace infini_train { +namespace { + +// ============================================================ +// uniform_real_distribution — 均匀分布函子 +// ============================================================ +// 仿 PyTorch at::uniform_real_distribution +// （aten/src/ATen/core/DistributionsHelper.h:99） +// +// 调用 generator->random()（float）或 generator->random64()（double） +// 获取原始随机比特，变换到 [from, to) 区间。 +// +// 模板参数 T = float | double +// 模板参数 RNG = CPUGeneratorImpl* | CUDAGeneratorImpl* | ... +// +template +struct uniform_real_distribution { + uniform_real_distribution(T from, T to) : from_(from), to_(to) { + assert(from <= to); + assert(to - from <= std::numeric_limits::max()); + } + // + // 不允许重新绑定 from/to，因为分布函子是无状态的 + 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_; + + // 变换：raw bits → [0, 1) → [from_, to_) + // 仿 PyTorch at::transformation::uniform_real + // （aten/src/ATen/core/TransformationHelper.h:84） + 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_; + } +}; + +// ============================================================ +// Box-Muller 正态分布缓存辅助（SFINAE） +// ============================================================ +// 仿 PyTorch at::maybe_get_next_normal_sample / maybe_set_next_normal_sample +// （aten/src/ATen/core/DistributionsHelper.h:120-163） +// +// 如果 RNG 有 next_float_normal_sample() 系列方法（如 CPUGeneratorImpl）， +// 则 Box-Muller 第二个样本被缓存到 Generator 里。 +// 如果 RNG 没有这些方法（如 CUDA curand state），则 SFINAE 回退到 no-op， +// 每次生成两个样本但只返回一个（丢弃另一个）。 +// + +// ---- get: 尝试读取缓存 ---- + +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; +} + +// 兜底：不支持缓存时总是返回 false +template +bool maybe_get_next_normal_sample(RNG * /*generator*/, void * /*ret*/) { + return false; +} + +// ---- set: 写入缓存 ---- + +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); +} + +// 兜底：不支持缓存时 no-op +template +void maybe_set_next_normal_sample(RNG * /*generator*/, const void * /*cache*/) {} + +// ============================================================ +// normal_distribution — 正态分布函子（Box-Muller） +// ============================================================ +// 仿 PyTorch at::normal_distribution +// （aten/src/ATen/core/DistributionsHelper.h:172） +// +// Box-Muller 每次产生两个正态样本，第二个缓存到 Generator。 +// +template +struct normal_distribution {//用struct能记录 + normal_distribution(T mean, T stdv) : mean_(mean), stdv_(stdv) { + assert(stdv >= 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_; + } + + // 生成两个 [0, 1) 均匀样本 + uniform_real_distribution uniform(static_cast(0), static_cast(1)); + const T u1 = uniform(generator); + const T u2 = uniform(generator); + + // Box-Muller 变换 + const T r = std::sqrt(static_cast(-2.0) * std::log1p(-u2)); + constexpr T kTwoPi = static_cast(2.0 * M_PI); + const T theta = kTwoPi * u1; + const T sample = r * std::sin(theta); + + // 缓存第二个样本 + maybe_set_next_normal_sample(generator, &sample); + + ret = r * std::cos(theta); + return ret * stdv_ + mean_; + } + +private: + T mean_; + T stdv_; +}; + +} // namespace +} // namespace infini_train diff --git a/infini_train/include/generator.h b/infini_train/include/generator.h new file mode 100644 index 00000000..be3e13ca --- /dev/null +++ b/infini_train/include/generator.h @@ -0,0 +1,141 @@ +#pragma once + +#include +#include +#include + +#include "infini_train/include/device.h" + +namespace infini_train { + +// 前向声明，避免循环依赖（仿 PyTorch at::Tensor 前向声明） +class Tensor; + +// ============================================================ +// GeneratorImpl — 抽象基类（仿 c10::GeneratorImpl） +// ============================================================ +// 定义所有 RNG 后端必须实现的纯虚接口。 +// +// 拷贝/移动已删除，只能通过 clone() 显式深拷贝。 +// clone 采用 NVI（Non-Virtual Interface）模式： +// - clone() 公有非虚，内部调用 protected clone_impl() +// - 子类只需覆写 clone_impl() 返回堆上分配的同类型拷贝 +// +// 线程安全：mutex_ 是 public 的，调用方对多步原子操作自行加锁。 +// +class GeneratorImpl { +public: + explicit GeneratorImpl(Device device) : device_(device) {} + virtual ~GeneratorImpl() = default; + + // 禁止拷贝 / 移动，避免意外覆盖 RNG 状态 + 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; + + // ---- NVI clone ---- + std::shared_ptr clone() const { + return std::shared_ptr(clone_impl()); + } + + // ---- 设备 ---- + Device device() const { return device_; } + + // 线程安全（public，调用方自行加锁） + std::mutex mutex_; + +protected: + Device device_; + + // 子类覆写点：返回堆上分配的同类型拷贝 + virtual GeneratorImpl *clone_impl() const = 0; +}; + +// ============================================================ +// Generator — 值类型壳（仿 at::Generator） +// ============================================================ +// 轻量级、可拷贝的值类型。拷贝是浅拷贝——两个 Generator +// 共享同一个 GeneratorImpl，推进一个对另一个可见。 +// +// 默认构造的 Generator 处于 "undefined" 状态（impl_ == nullptr）。 +// 使用 make_generator(seed) 或 Generator(impl) +// 来创建可用的实例。 +// +class Generator { +public: + // 默认种子：一个大数，bit 分布均匀（仿 PyTorch default_rng_seed_val） + static constexpr uint64_t kDefaultSeed = 67280421310721; + + // 默认构造：undefined 状态 + Generator() = default; + + // 从已有 Impl 构造（impl 不能为 nullptr，实现在 .cc 做检查） + explicit Generator(std::shared_ptr impl); + + // 拷贝 / 移动（浅拷贝，共享 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) { impl_->set_current_seed(seed); } + uint64_t current_seed() const { return impl_->current_seed(); } + uint64_t seed() { return impl_->seed(); } + + // ---- 状态序列化（实现在 .cc，需要 Tensor 完整定义）---- + void set_state(const Tensor &state); + std::shared_ptr get_state() const; + + // ---- 设备 ---- + Device device() const { return impl_->device(); } + + // ---- 克隆（深拷贝 Impl）---- + Generator clone() const { return Generator(impl_->clone()); } + + // ---- 线程安全 ---- + std::mutex &mutex() { return impl_->mutex_; } + + // ---- 类型安全 downcast ---- + 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_; +}; + +// ============================================================ +// 工具函数 +// ============================================================ + +// 工厂函数：make_generator(seed) +template +Generator make_generator(Args &&...args) { + return Generator(std::make_shared(std::forward(args)...)); +} + +} // namespace infini_train diff --git a/infini_train/include/nn/init.h b/infini_train/include/nn/init.h index fc6effec..05b95b2c 100644 --- a/infini_train/include/nn/init.h +++ b/infini_train/include/nn/init.h @@ -2,11 +2,11 @@ #include #include -#include #include #include "infini_train/include/datatype.h" #include "infini_train/include/device.h" +#include "infini_train/include/generator.h" namespace infini_train { class Tensor; @@ -15,7 +15,7 @@ class Device; namespace infini_train::nn::init { std::shared_ptr Normal(const std::shared_ptr &tensor, float mean = 0.0, float std = 1.0, - std::optional generator = std::nullopt); + std::optional generator = std::nullopt); std::pair CalculateFanInAndFanOut(const std::shared_ptr &tensor); @@ -42,10 +42,10 @@ enum class NonLinearityType : int8_t { std::shared_ptr KaimingUniform(const std::shared_ptr &tensor, float a = 0.0f, KaimingMode mode = KaimingMode::kFanIn, NonLinearityType non_linearity = NonLinearityType::kLeakyReLU, - std::optional generator = std::nullopt); + std::optional generator = std::nullopt); std::shared_ptr Uniform(const std::shared_ptr &tensor, float a = 0.0f, float b = 1.0f, - std::optional generator = std::nullopt); + std::optional generator = std::nullopt); std::shared_ptr Ones(const std::shared_ptr &tensor); diff --git a/infini_train/include/tensor.h b/infini_train/include/tensor.h index 12f45f57..78f0af58 100644 --- a/infini_train/include/tensor.h +++ b/infini_train/include/tensor.h @@ -4,7 +4,6 @@ #include #include #include -#include #include #include "Eigen/Dense" @@ -12,6 +11,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 { @@ -148,7 +148,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/core/distribution_kernels.cc b/infini_train/src/core/distribution_kernels.cc new file mode 100644 index 00000000..4ecbae5c --- /dev/null +++ b/infini_train/src/core/distribution_kernels.cc @@ -0,0 +1,99 @@ +/// distribution_kernels.cpp +/// +/// 桥接层：Generator Handle → CPUGeneratorImpl* → 分布函子 +/// +/// - get_generator_or_default: 解析 Generator → 具体 Impl 指针 +/// - *_cpu_kernel: CPU 后端实现（加锁 + 遍历 + 分布函子） +/// - REGISTER_DISPATCH: 将内核注册到 dispatch 表 +/// +/// 仿 PyTorch aten/src/ATen/native/cpu/DistributionKernels.cpp（253 行） + +#include +#include + +#include "infini_train/include/core/runtime/distribution_kernels.h" +#include "infini_train/include/core/runtime/distribution_stubs.h" +#include "infini_train/include/core/runtime/distributions_helper.h" +#include "infini_train/src/core/runtime/cpu/cpu_generator_impl.h" + +namespace infini_train { +namespace { + +// ============================================================ +// get_generator_or_default — Generator Handle → 具体 Impl 指针 +// ============================================================ +// 仿 PyTorch at::native::get_generator_or_default +// +// 如果调用方显式传了 Generator，就用那个； +// 否则 fallback 到当前设备的默认 Generator。 +// +template +T *get_generator_or_default(const std::optional &gen, const Generator &default_gen) { + if (gen.has_value() && gen->defined()) { + return gen->get(); + } + return default_gen.get(); +} + +// ---- DEFINE_DISPATCH：为 stub 分配全局存储 ---- +// 注意：必须在 namespace infini_train 层级，与 DECLARE_DISPATCH 声明的 extern 匹配。 +// 不能放在匿名 namespace 里，否则会创建两个不同的变量。 + +} // namespace + +DEFINE_DISPATCH(uniform_stub); +DEFINE_DISPATCH(normal_stub); + +namespace { + +// ---- CPU 内核实现 ---- + +void uniform_cpu_kernel(void *data, int64_t n, double from, double to, + const std::optional &gen) { + auto *cpu_gen = get_generator_or_default( + gen, core::cpu::getDefaultCPUGenerator()); + + std::lock_guard lock(cpu_gen->mutex_); + auto *buf = static_cast(data); + uniform_real_distribution dist(static_cast(from), static_cast(to)); + for (int64_t i = 0; i < n; ++i) { + buf[i] = dist(cpu_gen); + } +} + +void normal_cpu_kernel(void *data, int64_t n, double mean, double std, + const std::optional &gen) { + auto *cpu_gen = get_generator_or_default( + gen, core::cpu::getDefaultCPUGenerator()); + + std::lock_guard lock(cpu_gen->mutex_); + auto *buf = static_cast(data); + normal_distribution dist(static_cast(mean), static_cast(std)); + for (int64_t i = 0; i < n; ++i) { + buf[i] = dist(cpu_gen); + } +} + +} // namespace + +// ---- REGISTER_DISPATCH：启动时自动把内核填进 dispatch 表 ---- + +REGISTER_DISPATCH(uniform_stub, Device::DeviceType::kCPU, &uniform_cpu_kernel); +REGISTER_DISPATCH(normal_stub, Device::DeviceType::kCPU, &normal_cpu_kernel); + +// ============================================================ +// 公共 API：包装层（仿 PyTorch Distributions.cpp 中的 uniform_/normal_ 入口函数） +// ============================================================ +// init.cc 等上层代码调用这些函数，不直接碰 uniform_stub()。 + +void uniform_kernel(void *data, int64_t n, double from, double to, + Device::DeviceType device_type, const std::optional &gen) { + uniform_stub(device_type, data, n, from, to, gen); +} + +void normal_kernel(void *data, int64_t n, double mean, double std, + Device::DeviceType device_type, const std::optional &gen) { + normal_stub(device_type, data, n, mean, std, gen); +} + +} // namespace infini_train diff --git a/infini_train/src/core/runtime/cpu/cpu_generator_impl.cc b/infini_train/src/core/runtime/cpu/cpu_generator_impl.cc new file mode 100644 index 00000000..0292a4e8 --- /dev/null +++ b/infini_train/src/core/runtime/cpu/cpu_generator_impl.cc @@ -0,0 +1,229 @@ +#include "infini_train/src/core/runtime/cpu/cpu_generator_impl.h" + +#include +#include +#include + +#include "glog/logging.h" + +#include "infini_train/include/tensor.h" + +namespace infini_train::core::cpu { + +// ============================================================ +// 非确定性随机数（仿 c10::detail::getNonDeterministicRandom） +// ============================================================ +// Linux 下读取 /dev/urandom；其他平台 fallback 到 std::random_device +// +static uint64_t getNonDeterministicRandom() { + std::random_device rd; + uint64_t val = (static_cast(rd()) << 32) | rd(); + return val; +} + +// ============================================================ +// CPUGeneratorImpl +// ============================================================ + +CPUGeneratorImpl::CPUGeneratorImpl(uint64_t seed) + : GeneratorImpl(Device(Device::DeviceType::kCPU, 0)) + , 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_ = std::mt19937(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; +} + +// ============================================================ +// 状态序列化 +// ============================================================ +// 二进制格式（仿 PyTorch CPUGeneratorImplState）： +// [engine_stream: N bytes] — mt19937 的 operator<< 输出 +// [seed_: 8 bytes] — 当前种子 +// [has_float: 1 byte] — 是否有缓存的 float 正态样本 +// [float_val: 4 bytes] — 缓存的 float 正态样本 +// [has_double: 1 byte] — 是否有缓存的 double 正态样本 +// [double_val: 8 bytes] — 缓存的 double 正态样本 +// +// 注意：engine 部分使用 operator<< / operator>>，格式是实现定义的。 +// Known limitation: PyTorch 通过自己实现 mt19937_engine 暴露 data()/set_data() +// 来实现固定格式序列化。我们使用 std::mt19937，标准库不提供内部状态访问接口， +// 因此无法做到跨编译器/跨版本的固定格式。此实现在同一构建下是稳定的。 + +void CPUGeneratorImpl::set_state(const Tensor &state) { + const uint8_t *data = static_cast(state.DataPtr()); + const size_t data_size = state.SizeInBytes(); + size_t offset = 0; + + // 1. 恢复引擎（变长部分直到 seed_ 字段前 ~ 22 字节的固定尾部） + // 引擎的 operator<< 输出是变长的，我们把剩下的 data 一起给 stream + // 但流可能会多读。改用精确长度：data_size 减去尾部固定字段长度。 + constexpr size_t kFooterSize = 8 + 1 + 4 + 1 + 8; // seed + has_float + float + has_double + double + + std::string engine_str; + if (data_size >= kFooterSize) { + engine_str.assign(reinterpret_cast(data), data_size - kFooterSize); + offset = data_size - kFooterSize; + } else { + // 旧格式：没有 footer（向后兼容），整个 data 就是 engine 状态 + engine_str.assign(reinterpret_cast(data), data_size); + offset = data_size; + } + + std::istringstream iss(engine_str); + iss >> engine_; + + // 2. 恢复种子和正态缓存 + if (offset + kFooterSize <= data_size) { + std::memcpy(&seed_, data + offset, sizeof(seed_)); + offset += sizeof(seed_); + + bool has_float = (data[offset++] != 0); + float float_val; + std::memcpy(&float_val, data + offset, sizeof(float_val)); + offset += sizeof(float_val); + next_float_normal_sample_ = has_float ? std::optional(float_val) : std::nullopt; + + bool has_double = (data[offset++] != 0); + double double_val; + std::memcpy(&double_val, data + offset, sizeof(double_val)); + next_double_normal_sample_ = has_double ? std::optional(double_val) : std::nullopt; + } +} + +std::shared_ptr CPUGeneratorImpl::get_state() const { + // 1. 序列化引擎 + std::ostringstream oss; + oss << engine_; + std::string engine_str = oss.str(); + + // 2. 计算总大小 + const size_t engine_size = engine_str.size(); + constexpr size_t kFooterSize = 8 + 1 + 4 + 1 + 8; + const size_t total_size = engine_size + kFooterSize; + + 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, engine_str.data(), engine_size); + offset += engine_size; + + // 写入种子 + std::memcpy(data + offset, &seed_, sizeof(seed_)); + offset += sizeof(seed_); + + // 写入 float 正态缓存 + 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); + + // 写入 double 正态缓存 + 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; +} + +// ============================================================ +// Box-Muller 正态缓存 +// ============================================================ + +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; +} + +// ============================================================ +// clone +// ============================================================ + +std::shared_ptr CPUGeneratorImpl::clone() const { + return std::shared_ptr(clone_impl()); +} + +CPUGeneratorImpl *CPUGeneratorImpl::clone_impl() const { + auto gen = new CPUGeneratorImpl(seed_); + gen->set_engine(engine_); + gen->set_next_float_normal_sample(next_float_normal_sample_); + gen->set_next_double_normal_sample(next_double_normal_sample_); + return gen; +} + +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 + +// ============================================================ +// 默认 Generator 管理（文件作用域，仿 PyTorch detail 命名空间） +// ============================================================ + +namespace infini_train::core::cpu { + +const Generator &getDefaultCPUGenerator() { + // 使用真随机种子初始化默认 generator（仿 PyTorch） + static auto default_gen = createCPUGenerator(getNonDeterministicRandom()); + return default_gen; +} + +Generator createCPUGenerator(uint64_t seed) { + return make_generator(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..05c65e0a --- /dev/null +++ b/infini_train/src/core/runtime/cpu/cpu_generator_impl.h @@ -0,0 +1,63 @@ +#pragma once + +#include +#include +#include +#include + +#include "infini_train/include/generator.h" + +namespace infini_train::core::cpu { + +// ============================================================ +// CPUGeneratorImpl — CPU RNG 后端（仿 at::CPUGeneratorImpl） +// ============================================================ +// 包装 std::mt19937 Mersenne Twister 引擎。 +// 缓存 Box-Muller 正态分布样本以优化性能。 +// +class CPUGeneratorImpl final : public GeneratorImpl { +public: + explicit CPUGeneratorImpl(uint64_t seed = Generator::kDefaultSeed); + ~CPUGeneratorImpl() override = default; + + // ---- GeneratorImpl 接口 ---- + 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; + + // ---- clone（类型安全版本）---- + std::shared_ptr clone() const; + + // ---- 类型标识（用于 check_generator 模板）---- + static Device::DeviceType device_type(); + + // ---- 随机数生成 ---- + uint32_t random(); + uint64_t random64(); + + // ---- Box-Muller 正态缓存 ---- + 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; + + // ---- 引擎（private：比赛要求公共接口不暴露 std::mt19937）---- + 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_; +}; + +// ---- 默认 Generator 管理（仿 PyTorch detail 命名空间）---- +const Generator &getDefaultCPUGenerator(); +Generator createCPUGenerator(uint64_t seed); + +} // namespace infini_train::core::cpu diff --git a/infini_train/src/generator.cc b/infini_train/src/generator.cc new file mode 100644 index 00000000..d066d2a1 --- /dev/null +++ b/infini_train/src/generator.cc @@ -0,0 +1,30 @@ +#include "infini_train/include/generator.h" + +#include "glog/logging.h" + +#include "infini_train/include/tensor.h" + +namespace infini_train { + +// ============================================================ +// Generator 构造函数（null 检查需要 glog，放 .cc） +// ============================================================ + +Generator::Generator(std::shared_ptr impl) + : impl_(std::move(impl)) { + CHECK(impl_) << "GeneratorImpl with nullptr is not supported"; +} + +// ============================================================ +// Generator — 状态序列化（需要 Tensor 完整定义，放 .cc） +// ============================================================ + +void Generator::set_state(const Tensor &state) { + impl_->set_state(state); +} + +std::shared_ptr Generator::get_state() const { + return impl_->get_state(); +} + +} // namespace infini_train diff --git a/infini_train/src/kernels/cpu/cross_entropy.cc b/infini_train/src/kernels/cpu/cross_entropy.cc index f520b2e9..37c7e9cd 100644 --- a/infini_train/src/kernels/cpu/cross_entropy.cc +++ b/infini_train/src/kernels/cpu/cross_entropy.cc @@ -1,6 +1,7 @@ #include #include #include +#include #include #include diff --git a/infini_train/src/kernels/cpu/transform.cc b/infini_train/src/kernels/cpu/transform.cc index 1a810b44..04031631 100644 --- a/infini_train/src/kernels/cpu/transform.cc +++ b/infini_train/src/kernels/cpu/transform.cc @@ -1,5 +1,6 @@ #include #include +#include #include "glog/logging.h" diff --git a/infini_train/src/nn/init.cc b/infini_train/src/nn/init.cc index 79b4b48b..7299598d 100644 --- a/infini_train/src/nn/init.cc +++ b/infini_train/src/nn/init.cc @@ -1,59 +1,26 @@ #include "infini_train/include/nn/init.h" -#include -#include #include -#include -#include +#include #include - -#ifdef USE_OMP -#include -#endif +#include #include "glog/logging.h" +#include "infini_train/include/core/runtime/distribution_kernels.h" #include "infini_train/include/core/runtime/device_guard.h" #include "infini_train/include/device.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) { + 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); - } - } -#else - std::normal_distribution dis(mean, std); - std::generate(buffer.begin(), buffer.end(), [&]() { return generator ? dis(generator.value()) : dis(gen); }); -#endif + // 始终在 CPU 上生成随机数，后续 MemcpyAsync 负责搬运到目标设备 + normal_kernel(buffer.data(), num_elements, mean, std, Device::DeviceType::kCPU, generator); auto device = tensor->GetDevice(); core::DeviceGuard guard(device); @@ -113,7 +80,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,24 +95,12 @@ 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) { + 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 + // 始终在 CPU 上生成随机数，后续 MemcpyAsync 负责搬运到目标设备 + uniform_kernel(buffer.data(), num_elements, a, b, Device::DeviceType::kCPU, generator); auto device = tensor->GetDevice(); diff --git a/infini_train/src/nn/parallel/ddp/param_and_grad_buffer.cc b/infini_train/src/nn/parallel/ddp/param_and_grad_buffer.cc index 75a21f63..9ffd2ad0 100644 --- a/infini_train/src/nn/parallel/ddp/param_and_grad_buffer.cc +++ b/infini_train/src/nn/parallel/ddp/param_and_grad_buffer.cc @@ -3,6 +3,7 @@ #include #include #include +#include #include "glog/logging.h" diff --git a/infini_train/src/tensor.cc b/infini_train/src/tensor.cc index bf7210ad..1bcfbf3c 100644 --- a/infini_train/src/tensor.cc +++ b/infini_train/src/tensor.cc @@ -3,6 +3,7 @@ #include #include #include +#include #include #include @@ -470,7 +471,7 @@ std::shared_ptr Tensor::Outer(const std::shared_ptr &other) { } // distribution -std::shared_ptr Tensor::Uniform(float from, float to, std::optional generator) { +std::shared_ptr Tensor::Uniform(float from, float to, std::optional generator) { return nn::init::Uniform(shared_from_this(), from, to, generator); } From c0b7c839f96c4c22e86da46dc9d5ef7b9437ed6a Mon Sep 17 00:00:00 2001 From: chaos Date: Sat, 11 Jul 2026 18:35:04 +0800 Subject: [PATCH 04/11] feat: add global manual seed control --- infini_train/include/generator.h | 8 ++++++-- infini_train/src/core/runtime/cpu/cpu_generator_impl.cc | 9 ++++++++- infini_train/src/core/runtime/cpu/cpu_generator_impl.h | 1 + infini_train/src/generator.cc | 5 +++++ 4 files changed, 20 insertions(+), 3 deletions(-) diff --git a/infini_train/include/generator.h b/infini_train/include/generator.h index be3e13ca..dbc4d6f5 100644 --- a/infini_train/include/generator.h +++ b/infini_train/include/generator.h @@ -89,7 +89,7 @@ class Generator { ~Generator() = default; // ---- 种子 ---- - void set_current_seed(uint64_t seed) { impl_->set_current_seed(seed); } + 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(); } @@ -104,7 +104,7 @@ class Generator { Generator clone() const { return Generator(impl_->clone()); } // ---- 线程安全 ---- - std::mutex &mutex() { return impl_->mutex_; } + std::mutex &mutex() const { return impl_->mutex_; } // ---- 类型安全 downcast ---- template @@ -138,4 +138,8 @@ Generator make_generator(Args &&...args) { return Generator(std::make_shared(std::forward(args)...)); } +// Reset the default generators for all supported devices. CPU is currently +// supported; CUDA default generators will be added to this entry point later. +void manual_seed(uint64_t seed); + } // namespace infini_train diff --git a/infini_train/src/core/runtime/cpu/cpu_generator_impl.cc b/infini_train/src/core/runtime/cpu/cpu_generator_impl.cc index 0292a4e8..a63f1c4a 100644 --- a/infini_train/src/core/runtime/cpu/cpu_generator_impl.cc +++ b/infini_train/src/core/runtime/cpu/cpu_generator_impl.cc @@ -1,8 +1,9 @@ #include "infini_train/src/core/runtime/cpu/cpu_generator_impl.h" #include +#include #include -#include +#include #include "glog/logging.h" @@ -226,4 +227,10 @@ 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 index 05c65e0a..0fb2506c 100644 --- a/infini_train/src/core/runtime/cpu/cpu_generator_impl.h +++ b/infini_train/src/core/runtime/cpu/cpu_generator_impl.h @@ -59,5 +59,6 @@ class CPUGeneratorImpl final : public GeneratorImpl { // ---- 默认 Generator 管理（仿 PyTorch detail 命名空间）---- 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/generator.cc b/infini_train/src/generator.cc index d066d2a1..3bc5906e 100644 --- a/infini_train/src/generator.cc +++ b/infini_train/src/generator.cc @@ -3,6 +3,7 @@ #include "glog/logging.h" #include "infini_train/include/tensor.h" +#include "infini_train/src/core/runtime/cpu/cpu_generator_impl.h" namespace infini_train { @@ -27,4 +28,8 @@ std::shared_ptr Generator::get_state() const { return impl_->get_state(); } +void manual_seed(uint64_t seed) { + core::cpu::manual_seed(seed); +} + } // namespace infini_train From 32cb797956ee709f7ffec6afe177b2cf80969d40 Mon Sep 17 00:00:00 2001 From: chaos Date: Sat, 11 Jul 2026 18:57:54 +0800 Subject: [PATCH 05/11] fix: validate generator backend before dispatch --- infini_train/include/generator.h | 32 ++++++++++++++++++- infini_train/src/core/distribution_kernels.cc | 24 -------------- 2 files changed, 31 insertions(+), 25 deletions(-) diff --git a/infini_train/include/generator.h b/infini_train/include/generator.h index dbc4d6f5..fc69fcab 100644 --- a/infini_train/include/generator.h +++ b/infini_train/include/generator.h @@ -3,6 +3,8 @@ #include #include #include +#include +#include #include "infini_train/include/device.h" @@ -106,7 +108,9 @@ class Generator { // ---- 线程安全 ---- std::mutex &mutex() const { return impl_->mutex_; } - // ---- 类型安全 downcast ---- + // ---- 非检查 downcast ---- + // 调用方必须先确认 Impl 类型与 Generator 的设备类型匹配。 + // 算子代码应使用 check_generator()。 template T *get() const { return static_cast(impl_.get()); @@ -138,6 +142,32 @@ Generator make_generator(Args &&...args) { return Generator(std::make_shared(std::forward(args)...)); } +// 检查 Generator 已定义且属于 T 对应的设备后端，再进行 downcast。 +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; +} + +// 显式 Generator 优先；未提供或未定义时使用该设备的默认 Generator。 +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); +} + // Reset the default generators for all supported devices. CPU is currently // supported; CUDA default generators will be added to this entry point later. void manual_seed(uint64_t seed); diff --git a/infini_train/src/core/distribution_kernels.cc b/infini_train/src/core/distribution_kernels.cc index 4ecbae5c..eec6e55c 100644 --- a/infini_train/src/core/distribution_kernels.cc +++ b/infini_train/src/core/distribution_kernels.cc @@ -9,7 +9,6 @@ /// 仿 PyTorch aten/src/ATen/native/cpu/DistributionKernels.cpp（253 行） #include -#include #include "infini_train/include/core/runtime/distribution_kernels.h" #include "infini_train/include/core/runtime/distribution_stubs.h" @@ -17,29 +16,6 @@ #include "infini_train/src/core/runtime/cpu/cpu_generator_impl.h" namespace infini_train { -namespace { - -// ============================================================ -// get_generator_or_default — Generator Handle → 具体 Impl 指针 -// ============================================================ -// 仿 PyTorch at::native::get_generator_or_default -// -// 如果调用方显式传了 Generator，就用那个； -// 否则 fallback 到当前设备的默认 Generator。 -// -template -T *get_generator_or_default(const std::optional &gen, const Generator &default_gen) { - if (gen.has_value() && gen->defined()) { - return gen->get(); - } - return default_gen.get(); -} - -// ---- DEFINE_DISPATCH：为 stub 分配全局存储 ---- -// 注意：必须在 namespace infini_train 层级，与 DECLARE_DISPATCH 声明的 extern 匹配。 -// 不能放在匿名 namespace 里，否则会创建两个不同的变量。 - -} // namespace DEFINE_DISPATCH(uniform_stub); DEFINE_DISPATCH(normal_stub); From 86bc2eaf7e01520b0f2ee9a8e433b763017982ce Mon Sep 17 00:00:00 2001 From: chaos Date: Sat, 11 Jul 2026 19:54:23 +0800 Subject: [PATCH 06/11] feat: add CUDA generator distribution support --- .../core/runtime/distribution_kernels.h | 11 +- .../include/core/runtime/distribution_stubs.h | 7 +- infini_train/src/core/distribution_kernels.cc | 34 +++-- .../core/runtime/cuda/cuda_generator_impl.cc | 137 ++++++++++++++++++ .../core/runtime/cuda/cuda_generator_impl.h | 37 +++++ infini_train/src/generator.cc | 8 + infini_train/src/kernels/cuda/distribution.cu | 112 ++++++++++++++ infini_train/src/kernels/cuda/elementwise.cu | 1 + infini_train/src/kernels/cuda/gather.cu | 2 + infini_train/src/kernels/cuda/no_op.cu | 2 + infini_train/src/kernels/cuda/reduction.cu | 2 + infini_train/src/nn/init.cc | 25 +--- 12 files changed, 335 insertions(+), 43 deletions(-) create mode 100644 infini_train/src/core/runtime/cuda/cuda_generator_impl.cc create mode 100644 infini_train/src/core/runtime/cuda/cuda_generator_impl.h create mode 100644 infini_train/src/kernels/cuda/distribution.cu diff --git a/infini_train/include/core/runtime/distribution_kernels.h b/infini_train/include/core/runtime/distribution_kernels.h index e779f0fe..48a6c268 100644 --- a/infini_train/include/core/runtime/distribution_kernels.h +++ b/infini_train/include/core/runtime/distribution_kernels.h @@ -17,11 +17,12 @@ namespace infini_train { -// 填 n 个 float 到 data -void uniform_kernel(void *data, int64_t n, double from, double to, - Device::DeviceType device_type, const std::optional &gen); +class Tensor; -void normal_kernel(void *data, int64_t n, double mean, double std, - Device::DeviceType device_type, const std::optional &gen); +void uniform_kernel(Tensor &tensor, double from, double to, + const std::optional &gen); + +void normal_kernel(Tensor &tensor, double mean, double std, + const std::optional &gen); } // namespace infini_train diff --git a/infini_train/include/core/runtime/distribution_stubs.h b/infini_train/include/core/runtime/distribution_stubs.h index 94d49339..c6d0d682 100644 --- a/infini_train/include/core/runtime/distribution_stubs.h +++ b/infini_train/include/core/runtime/distribution_stubs.h @@ -16,12 +16,13 @@ namespace infini_train { -// 填 n 个 float 到 data 指向的缓冲区 -DECLARE_DISPATCH(void (*)(void *data, int64_t n, double a, double b, +class Tensor; + +DECLARE_DISPATCH(void (*)(Tensor &tensor, double a, double b, const std::optional &gen), uniform_stub); -DECLARE_DISPATCH(void (*)(void *data, int64_t n, double a, double b, +DECLARE_DISPATCH(void (*)(Tensor &tensor, double a, double b, const std::optional &gen), normal_stub); diff --git a/infini_train/src/core/distribution_kernels.cc b/infini_train/src/core/distribution_kernels.cc index eec6e55c..f5a438f2 100644 --- a/infini_train/src/core/distribution_kernels.cc +++ b/infini_train/src/core/distribution_kernels.cc @@ -13,6 +13,7 @@ #include "infini_train/include/core/runtime/distribution_kernels.h" #include "infini_train/include/core/runtime/distribution_stubs.h" #include "infini_train/include/core/runtime/distributions_helper.h" +#include "infini_train/include/tensor.h" #include "infini_train/src/core/runtime/cpu/cpu_generator_impl.h" namespace infini_train { @@ -22,30 +23,37 @@ DEFINE_DISPATCH(normal_stub); namespace { +void check_distribution_tensor(const Tensor &tensor) { + CHECK_EQ(static_cast(tensor.Dtype()), static_cast(DataType::kFLOAT32)) + << "Uniform and Normal initialization currently support FLOAT32 tensors only"; +} + // ---- CPU 内核实现 ---- -void uniform_cpu_kernel(void *data, int64_t n, double from, double to, +void uniform_cpu_kernel(Tensor &tensor, double from, double to, const std::optional &gen) { + CHECK(tensor.GetDevice().IsCPU()); auto *cpu_gen = get_generator_or_default( gen, core::cpu::getDefaultCPUGenerator()); std::lock_guard lock(cpu_gen->mutex_); - auto *buf = static_cast(data); + auto *buf = static_cast(tensor.DataPtr()); uniform_real_distribution dist(static_cast(from), static_cast(to)); - for (int64_t i = 0; i < n; ++i) { + for (int64_t i = 0; i < tensor.NumElements(); ++i) { buf[i] = dist(cpu_gen); } } -void normal_cpu_kernel(void *data, int64_t n, double mean, double std, +void normal_cpu_kernel(Tensor &tensor, double mean, double std, const std::optional &gen) { + CHECK(tensor.GetDevice().IsCPU()); auto *cpu_gen = get_generator_or_default( gen, core::cpu::getDefaultCPUGenerator()); std::lock_guard lock(cpu_gen->mutex_); - auto *buf = static_cast(data); + auto *buf = static_cast(tensor.DataPtr()); normal_distribution dist(static_cast(mean), static_cast(std)); - for (int64_t i = 0; i < n; ++i) { + for (int64_t i = 0; i < tensor.NumElements(); ++i) { buf[i] = dist(cpu_gen); } } @@ -62,14 +70,16 @@ REGISTER_DISPATCH(normal_stub, Device::DeviceType::kCPU, &normal_cpu_kernel); // ============================================================ // init.cc 等上层代码调用这些函数，不直接碰 uniform_stub()。 -void uniform_kernel(void *data, int64_t n, double from, double to, - Device::DeviceType device_type, const std::optional &gen) { - uniform_stub(device_type, data, n, from, to, gen); +void uniform_kernel(Tensor &tensor, double from, double to, + const std::optional &gen) { + check_distribution_tensor(tensor); + uniform_stub(tensor.GetDevice().type(), tensor, from, to, gen); } -void normal_kernel(void *data, int64_t n, double mean, double std, - Device::DeviceType device_type, const std::optional &gen) { - normal_stub(device_type, data, n, mean, std, gen); +void normal_kernel(Tensor &tensor, double mean, double std, + const std::optional &gen) { + check_distribution_tensor(tensor); + normal_stub(tensor.GetDevice().type(), tensor, mean, std, gen); } } // namespace infini_train 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..a4ca0e9b --- /dev/null +++ b/infini_train/src/core/runtime/cuda/cuda_generator_impl.cc @@ -0,0 +1,137 @@ +#include "infini_train/src/core/runtime/cuda/cuda_generator_impl.h" + +#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 { + +constexpr size_t kStateSize = 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) { + CHECK(state.GetDevice().IsCPU()); + CHECK_EQ(static_cast(state.Dtype()), static_cast(DataType::kUINT8)); + CHECK_EQ(state.SizeInBytes(), kStateSize); + + const auto *data = static_cast(state.DataPtr()); + std::memcpy(&seed_, data, sizeof(seed_)); + std::memcpy(&next_philox_subsequence_, data + 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, &seed_, sizeof(seed_)); + std::memcpy(data + 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) { + const uint64_t subsequence = next_philox_subsequence_; + next_philox_subsequence_ += increment; + return subsequence; +} + +CUDAGeneratorImpl *CUDAGeneratorImpl::clone_impl() const { + auto *generator = new CUDAGeneratorImpl(device().index(), seed_); + generator->next_philox_subsequence_ = next_philox_subsequence_; + return generator; +} + +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 index 3bc5906e..e1b5b746 100644 --- a/infini_train/src/generator.cc +++ b/infini_train/src/generator.cc @@ -5,6 +5,10 @@ #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 { // ============================================================ @@ -30,6 +34,10 @@ std::shared_ptr Generator::get_state() const { 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/cuda/distribution.cu b/infini_train/src/kernels/cuda/distribution.cu new file mode 100644 index 00000000..b5ab9903 --- /dev/null +++ b/infini_train/src/kernels/cuda/distribution.cu @@ -0,0 +1,112 @@ +#include +#include +#include + +#include + +#include "infini_train/include/common/cuda/common_cuda.h" +#include "infini_train/include/core/runtime/device_guard.h" +#include "infini_train/include/core/runtime/distribution_stubs.h" +#include "infini_train/include/generator.h" +#include "infini_train/include/tensor.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; + +__global__ void UniformKernel(float *data, int64_t n, float from, float 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 float unit = static_cast(curand(&state)) * 0x1p-32f; + data[index] = from + unit * (to - from); +} + +__global__ void NormalKernel(float *data, int64_t n, float mean, float 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] = mean + curand_normal(&state) * std; +} + +const core::cuda::CudaStream *get_cuda_stream(const Device &device) { + return dynamic_cast( + core::GetDeviceGuardImpl(device.type())->GetStream(device)); +} + +void uniform_cuda_kernel(Tensor &tensor, double from, double to, + const std::optional &generator) { + 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( + generator, 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); + UniformKernel<<cuda_stream()>>>( + static_cast(tensor.DataPtr()), n, static_cast(from), static_cast(to), seed, subsequence); + CUDA_CHECK(cudaGetLastError()); +} + +void normal_cuda_kernel(Tensor &tensor, double mean, double std, + const std::optional &generator) { + 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( + generator, 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); + NormalKernel<<cuda_stream()>>>( + static_cast(tensor.DataPtr()), n, static_cast(mean), static_cast(std), seed, subsequence); + CUDA_CHECK(cudaGetLastError()); +} + +} // namespace + +using infini_train::normal_stub; +using infini_train::uniform_stub; + +REGISTER_DISPATCH(uniform_stub, Device::DeviceType::kCUDA, &uniform_cuda_kernel); +REGISTER_DISPATCH(normal_stub, Device::DeviceType::kCUDA, &normal_cuda_kernel); + +} // namespace infini_train::kernels::cuda diff --git a/infini_train/src/kernels/cuda/elementwise.cu b/infini_train/src/kernels/cuda/elementwise.cu index fe63e0b2..f0fe3298 100644 --- a/infini_train/src/kernels/cuda/elementwise.cu +++ b/infini_train/src/kernels/cuda/elementwise.cu @@ -1,4 +1,5 @@ #include +#include #include diff --git a/infini_train/src/kernels/cuda/gather.cu b/infini_train/src/kernels/cuda/gather.cu index 12d0567d..dc791b9f 100644 --- a/infini_train/src/kernels/cuda/gather.cu +++ b/infini_train/src/kernels/cuda/gather.cu @@ -1,3 +1,5 @@ +#include + #include "glog/logging.h" #include "infini_train/include/common/cuda/common_cuda.h" diff --git a/infini_train/src/kernels/cuda/no_op.cu b/infini_train/src/kernels/cuda/no_op.cu index ef2c9566..d26024b0 100644 --- a/infini_train/src/kernels/cuda/no_op.cu +++ b/infini_train/src/kernels/cuda/no_op.cu @@ -1,3 +1,5 @@ +#include + #include "glog/logging.h" #include "infini_train/include/dispatcher.h" diff --git a/infini_train/src/kernels/cuda/reduction.cu b/infini_train/src/kernels/cuda/reduction.cu index c56470e3..2c080edd 100644 --- a/infini_train/src/kernels/cuda/reduction.cu +++ b/infini_train/src/kernels/cuda/reduction.cu @@ -1,3 +1,5 @@ +#include + #include #include "infini_train/include/common/cuda/common_cuda.h" diff --git a/infini_train/src/nn/init.cc b/infini_train/src/nn/init.cc index 7299598d..8034a6b6 100644 --- a/infini_train/src/nn/init.cc +++ b/infini_train/src/nn/init.cc @@ -16,19 +16,9 @@ namespace infini_train::nn::init { 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); - - // 始终在 CPU 上生成随机数，后续 MemcpyAsync 负责搬运到目标设备 - normal_kernel(buffer.data(), num_elements, mean, std, Device::DeviceType::kCPU, generator); - 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)); + normal_kernel(*tensor, mean, std, generator); return tensor; } @@ -96,20 +86,9 @@ 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); - - // 始终在 CPU 上生成随机数，后续 MemcpyAsync 负责搬运到目标设备 - uniform_kernel(buffer.data(), num_elements, a, b, Device::DeviceType::kCPU, generator); - 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)); + uniform_kernel(*tensor, a, b, generator); return tensor; } From 31bf2802eac790edd3fc3d445a1b736c025e3e6d Mon Sep 17 00:00:00 2001 From: chaos Date: Sun, 12 Jul 2026 11:34:59 +0800 Subject: [PATCH 07/11] fix: validate generator state inputs --- .claude/agents/ref.md | 23 ------ .gitignore | 19 +++++ ...75\350\261\241\351\200\211\351\242\230.md" | 77 ++++++++++--------- infini_train/include/generator.h | 10 ++- infini_train/include/tensor.h | 1 + .../core/runtime/cpu/cpu_generator_impl.cc | 76 +++++++++--------- .../core/runtime/cuda/cuda_generator_impl.cc | 3 +- infini_train/src/generator.cc | 11 +++ 8 files changed, 121 insertions(+), 99 deletions(-) delete mode 100644 .claude/agents/ref.md diff --git a/.claude/agents/ref.md b/.claude/agents/ref.md deleted file mode 100644 index decae7e7..00000000 --- a/.claude/agents/ref.md +++ /dev/null @@ -1,23 +0,0 @@ ---- -name: ref -description: 查找参考实现和官方文档，用于对比学习。搜索 PyTorch 源码、API 文档等。 -tools: Read, Grep, Glob, WebSearch, WebFetch -model: sonnet ---- - -你是参考研究员。你的职责是： -1. 在 PyTorch 源码中找到对应的实现 -2. 查阅官方文档，找到 API 说明和设计动机 -3. 对比不同框架的实现差异 -4. 提供参考链接和源码位置 - -输出格式： -- 给出 PyTorch 对应代码的文件路径和关键片段 -- 解释 PyTorch 为什么这样设计 -- 如果有多种实现方式，列出对比 -- 不要修改任何文件 - -使用场景： - -▎ "让 ref 查一下 PyTorch 的 Optimizer 是怎么实现 step() 的" -▎ "让 ref 找找 DistributedDataParallel 的梯度同步逻辑" diff --git a/.gitignore b/.gitignore index 6ec79bff..b37295bb 100644 --- a/.gitignore +++ b/.gitignore @@ -4,6 +4,25 @@ build-full/ .cache/ .vscode/ +# Local CMake output +/cmake-build-*/ +/CMakeCache.txt +/CMakeFiles/ +/cmake_install.cmake +/CTestTestfile.cmake +/CPackConfig.cmake +/CPackSourceConfig.cmake +/compile_commands.json +/CMakeUserPresets.json +/Makefile +/Testing/ + +# Python cache +__pycache__/ +*.py[cod] +.pytest_cache/ +.mypy_cache/ + *.log *.report.rank* *.records.log.rank* diff --git "a/docs/\343\200\2202026 \346\230\245\345\255\243\344\272\272\345\267\245\346\231\272\350\203\275\345\244\247\350\265\233\343\200\221Generator \346\212\275\350\261\241\351\200\211\351\242\230.md" "b/docs/\343\200\2202026 \346\230\245\345\255\243\344\272\272\345\267\245\346\231\272\350\203\275\345\244\247\350\265\233\343\200\221Generator \346\212\275\350\261\241\351\200\211\351\242\230.md" index 6e170aca..f550bee3 100644 --- "a/docs/\343\200\2202026 \346\230\245\345\255\243\344\272\272\345\267\245\346\231\272\350\203\275\345\244\247\350\265\233\343\200\221Generator \346\212\275\350\261\241\351\200\211\351\242\230.md" +++ "b/docs/\343\200\2202026 \346\230\245\345\255\243\344\272\272\345\267\245\346\231\272\350\203\275\345\244\247\350\265\233\343\200\221Generator \346\212\275\350\261\241\351\200\211\351\242\230.md" @@ -17,13 +17,13 @@ 在深度学习训练与推理框架中，随机数生成器（Generator）是支撑参数初始化、随机采样、Dropout、噪声注入等功能的重要基础设施。当前框架中相关能力尚不完善，缺少统一的 Generator 抽象、后端实现以及全局随机种子管理机制，导致随机相关算子的行为与主流框架存在差距，也难以满足后续功能扩展需求。 -当前框架在随机数相关能力上仍不完善，主要表现为： +当前已完成 Generator 基础设施与初始化算子的首轮接入，进度如下： -- [x] 缺少统一的 Generator 抽象，随机数状态、种子设置、状态获取与恢复等能力未形成标准接口； -- [x] CPU 与 CUDA 后端缺少统一的 Generator 实现，不同设备上的随机行为缺乏一致的调用方式；（CPU ✅，CUDA ❌） -- [x] 缺少全局默认 Generator 机制，随机算子在未显式传入 generator 时无法方便地使用当前设备的默认随机源；（CPU 默认 Generator 已实现，算子层未接入） -- [ ] 缺少统一的全局随机种子控制入口，导致参数初始化、Dropout 等随机算子在多次运行间难以稳定复现； -- [ ] 随机数相关能力与主流框架（如 PyTorch）存在差距，不利于后续算子扩展、模型对齐与测试验证。 +- [x] 已建立统一的 Generator 抽象，支持随机数状态、种子设置、状态获取与恢复； +- [x] 已实现 CPU 与 CUDA 后端 Generator，并保持统一的调用接口； +- [x] 已建立按设备维护的默认 Generator；未显式传入 generator 时，初始化算子使用目标 Tensor 所在设备的默认随机源； +- [x] 已提供全局 `manual_seed(uint64_t)` 入口，重置 CPU 及所有 CUDA 默认 Generator； +- [ ] Dropout、`rand`、`randn` 等训练期或通用随机算子尚未接入；distribution 当前仅支持 `FLOAT32`，也未实现 CUDA Graph 语义。 为补齐这一基础能力，本题目要求参赛者参考主流深度学习框架的设计思路，完成一套可扩展的随机数生成器基础设施，实现 CPU 与 CUDA 后端的 Generator，并支持基于 Generator 的随机数生成与种子控制能力。 @@ -32,10 +32,10 @@ 参赛者需要围绕框架内的随机数生成体系，完成以下几个方面的工作： - [x] 1. 设计并实现统一的 Generator 抽象接口，支持随机数状态管理、种子设置、状态获取与恢复等基础能力。 -- [x] 2. 分别实现 CPU 与 CUDA 后端对应的 GeneratorImpl，使不同设备上的随机数生成具备统一的调用方式。（CPU ✅，CUDA ❌） -- [ ] 3. 建立全局随机种子控制机制，支持通过统一入口固定随机种子，并使随机相关算子的行为可复现。 -- [ ] 4. 改造随机数相关算子，使其在未显式传入 generator 参数时，能够自动使用当前设备上的默认全局 Generator。 -- [ ] 5. 验证实现结果与主流框架在基本行为上的一致性，包括随机性、可复现性以及跨设备下的接口统一性。 +- [x] 2. 分别实现 CPU 与 CUDA 后端对应的 GeneratorImpl，使不同设备上的随机数生成具备统一的调用方式。 +- [x] 3. 建立全局随机种子控制机制，支持通过统一入口固定随机种子；已接入的初始化算子可稳定复现。 +- [x] 4. 改造初始化类随机算子，使其在未显式传入 generator 参数时，自动使用目标设备的默认全局 Generator。 +- [x] 5. 已完成 CPU/CUDA 初始化路径的基本行为、可复现性和接口语义冒烟验证；仓库内正式测试仍待补充。 ## 三、任务拆解 @@ -56,7 +56,7 @@ - [x] 提供用户侧可持有的 `Generator` 句柄类； - [x] 底层以 `GeneratorImpl` 作为多态实现基类； -- [x] CPU / CUDA 分别派生对应实现；（CPU 已实现，CUDA 未实现） +- [x] CPU / CUDA 分别派生对应实现； - [x] 公共接口不暴露 `std::mt19937`、curand、Philox 等后端细节； - [x] 为后续其他平台 Generator 接入预留扩展空间。 @@ -65,18 +65,18 @@ 基于现有设备后端，分别实现： - [x] `CPUGeneratorImpl` -- [ ] `CUDAGeneratorImpl` +- [x] `CUDAGeneratorImpl` 两者应保持统一接口语义，但内部状态组织方式可以不同，且**不要求 CPU 与 CUDA 在相同 seed 下生成逐元素一致的随机结果**。 实现时需重点关注： -- [x] CPU 与 CUDA 随机数状态如何组织与保存；（CPU 已实现二进制序列化格式） -- [ ] CUDA 侧应按设备维度维护独立 Generator； -- [x] `GetState()` / `SetState()` 的状态序列化、反序列化与格式校验；（CPU 已实现） -- [x] 多次随机调用后，随机序列是否连续推进；（CPU 已实现，`random()`/`random64()` 持续推进 `engine_`） -- [x] 同一 seed、同一设备、同一调用顺序下结果是否可复现；（CPU 已实现，mt19937 确定性） -- [ ] 显式传入 Generator 与使用默认 Generator 时，行为是否符合预期。（算子层尚未接入） +- [x] CPU 与 CUDA 随机数状态如何组织与保存；CPU 保存 mt19937 状态与正态缓存，CUDA 保存 Philox seed 与 subsequence； +- [x] CUDA 侧按设备维度维护独立默认 Generator； +- [x] `GetState()` / `SetState()` 的状态序列化、反序列化与格式校验；CPU/CUDA 均校验 CPU `UINT8` state，CPU 额外校验 footer 与 mt19937 反序列化； +- [x] 多次随机调用后，随机序列连续推进；CPU 推进 mt19937，CUDA 为每次 kernel 保留 Philox subsequence； +- [x] 同一 seed、同一设备、同一调用顺序下结果可复现； +- [x] 初始化算子已支持显式 Generator 和默认 Generator 两条调用路径。 CUDA 后端建议在状态中考虑 seed、offset、counter 等信息，避免不同 kernel 调用之间随机序列重叠。此部分需要保证功能正确、语义清晰、状态可管理、接口一致，为后续随机算子和分布式扩展提供基础。 @@ -90,18 +90,19 @@ CUDA 后端建议在状态中考虑 seed、offset、counter 等信息，避免 该机制应包括： - [x] CPU 默认全局 Generator； -- [ ] 各 CUDA 设备对应的默认 Generator； -- [x] 获取默认 Generator 的统一入口；（`getDefaultCPUGenerator()` 已实现） -- [ ] 统一的全局随机种子设置入口； -- [ ] 设置全局随机种子时，同步更新默认 Generator 的初始状态。 +- [x] 各 CUDA 设备对应的默认 Generator； +- [x] CPU/CUDA 均提供获取默认 Generator 的后端入口； +- [x] 统一的全局随机种子设置入口：`manual_seed(uint64_t)`； +- [x] 设置全局 seed 时同步重置 CPU 与所有 CUDA 默认 Generator 的状态。 要求： - [x] 多次获取同一设备默认 Generator 时，应返回同一随机状态来源；（static 局部变量，单例模式） -- [ ] 不同 CUDA 设备的默认 Generator 应相互独立； -- [x] 用户显式传入 Generator 时，必须使用用户指定 Generator；（API 支持 `std::optional`） -- [ ] 未传入 Generator 时，才使用当前设备默认 Generator；（算子层尚未接入） -- [ ] 设置统一 seed 后，参数初始化、Dropout、随机采样等行为应可稳定复现。（缺少全局 seed 入口） +- [x] 不同 CUDA 设备的默认 Generator 相互独立； +- [x] 用户显式传入 Generator 时，初始化算子使用用户指定 Generator； +- [x] 未传入 Generator 时，初始化算子使用目标设备默认 Generator； +- [x] 设置统一 seed 后，已接入的参数初始化和随机采样可稳定复现； +- [ ] Dropout 等尚未接入的随机算子仍待完成。 这一部分是本题的核心验收点。随机数系统最终服务于训练可复现性，因此必须通过统一入口稳定控制随机行为。 @@ -109,23 +110,27 @@ CUDA 后端建议在状态中考虑 seed、offset、counter 等信息，避免 改造框架中已有的随机相关使用，使其接入 Generator 机制。建议至少覆盖两类场景： -- [ ] 初始化类随机算子，如 uniform、normal、kaiming 等；（init.cc 仍使用 raw std::mt19937） +- [x] 初始化类随机算子：`Uniform`、`Normal`、`KaimingUniform` 已接入 Generator，并按 Tensor device 分发 CPU/CUDA kernel； - [ ] 训练过程中的随机算子，如 dropout、rand、randn 等； 改造后的算子应满足： -- [ ] 支持显式传入 Generator； -- [ ] 未传入 Generator 时，自动使用当前设备默认 Generator； -- [ ] 随机结果受统一 seed 入口控制； -- [x] 多次调用会推进 Generator 状态；（CPUGeneratorImpl 引擎层已支持，算子层未接入） -- [x] 同一 seed、同一调用顺序下结果可复现。（引擎层已支持，算子层未接入） +- [x] 已接入初始化算子支持显式传入 Generator； +- [x] 已接入初始化算子未传入 Generator 时，自动使用目标设备默认 Generator； +- [x] 已接入初始化算子的随机结果受统一 `manual_seed` 入口控制； +- [x] 已接入初始化算子多次调用会推进 Generator 状态； +- [x] 已接入初始化算子在同一 seed、同一调用顺序下可复现。 不要求一次性改造所有随机算子，但应至少完成一个初始化类算子和一个训练期随机算子，以证明 Generator 机制能够贯通框架层与算子层。 +当前仅完成初始化类算子接入；训练期随机算子仍是后续工作。 + ### 5. 测试与对齐验证 为验证实现结果，需要补充系统化测试。测试重点应放在"语义是否正确"和"是否可复现"，而不是与 PyTorch 逐元素数值一致。 +当前已在 `/tmp` 完成 CPU/CUDA 构建和 Generator、初始化可复现性的冒烟验证；按项目约定尚未新增仓库内测试，因此以下正式测试清单仍未完成。 + 建议至少覆盖以下内容： #### （1）接口一致性测试 @@ -216,10 +221,10 @@ CPU 后端必须覆盖该测试；CUDA 后端若已实现，也应覆盖。 - [ ] 代码以 PR 形式提交，结构清晰，具备基本可 review 性； - [x] 完成统一 Generator 抽象设计； -- [x] 实现 CPU Generator，CUDA Generator 建议实现；（CPU ✅，CUDA ❌） -- [x] 建立默认 Generator 管理机制；（CPU 侧已完成） -- [ ] 提供统一的全局随机种子设置入口； -- [ ] 支持默认 Generator 与显式 Generator 两种使用方式；（API 层已支持，算子层未接入） +- [x] 实现 CPU 与 CUDA Generator； +- [x] 建立 CPU 与各 CUDA device 的默认 Generator 管理机制； +- [x] 提供统一的全局随机种子设置入口； +- [x] 初始化算子支持默认 Generator 与显式 Generator 两种使用方式； - [x] 支持 Generator 状态的获取与恢复（get_state / set_state）； - [ ] 至少改造一类初始化算子和一类框架内随机数生成调用，使其接入 Generator； - [ ] 提供测试或运行日志，验证随机行为具备基本可复现性（同 seed 一致、不同 seed 不同、状态恢复有效）。 diff --git a/infini_train/include/generator.h b/infini_train/include/generator.h index fc69fcab..c59ecdae 100644 --- a/infini_train/include/generator.h +++ b/infini_train/include/generator.h @@ -13,6 +13,13 @@ namespace infini_train { // 前向声明，避免循环依赖（仿 PyTorch at::Tensor 前向声明） class Tensor; +namespace detail { + +// Validates the common Tensor contract for serialized RNG states. +void check_rng_state(const Tensor &state); + +} // namespace detail + // ============================================================ // GeneratorImpl — 抽象基类（仿 c10::GeneratorImpl） // ============================================================ @@ -168,8 +175,7 @@ T *get_generator_or_default(const std::optional &generator, : check_generator(default_generator); } -// Reset the default generators for all supported devices. CPU is currently -// supported; CUDA default generators will be added to this entry point later. +// Reset the default generators for all enabled devices. void manual_seed(uint64_t seed); } // namespace infini_train diff --git a/infini_train/include/tensor.h b/infini_train/include/tensor.h index 78f0af58..beabe12e 100644 --- a/infini_train/include/tensor.h +++ b/infini_train/include/tensor.h @@ -78,6 +78,7 @@ class Tensor : public std::enable_shared_from_this { const std::vector &Dims() const; size_t NumElements() const; DataType Dtype() const; + bool defined() const { return buffer_ != nullptr; } void Fill(Scalar value); diff --git a/infini_train/src/core/runtime/cpu/cpu_generator_impl.cc b/infini_train/src/core/runtime/cpu/cpu_generator_impl.cc index a63f1c4a..9a01aa6f 100644 --- a/infini_train/src/core/runtime/cpu/cpu_generator_impl.cc +++ b/infini_train/src/core/runtime/cpu/cpu_generator_impl.cc @@ -10,6 +10,12 @@ #include "infini_train/include/tensor.h" namespace infini_train::core::cpu { +namespace { + +constexpr size_t kStateFooterSize = sizeof(uint64_t) + sizeof(uint8_t) + sizeof(float) + sizeof(uint8_t) + + sizeof(double); + +} // namespace // ============================================================ // 非确定性随机数（仿 c10::detail::getNonDeterministicRandom） @@ -65,44 +71,43 @@ uint64_t CPUGeneratorImpl::seed() { // 因此无法做到跨编译器/跨版本的固定格式。此实现在同一构建下是稳定的。 void CPUGeneratorImpl::set_state(const Tensor &state) { - const uint8_t *data = static_cast(state.DataPtr()); + ::infini_train::detail::check_rng_state(state); + const size_t data_size = state.SizeInBytes(); - size_t offset = 0; + CHECK_GT(data_size, kStateFooterSize) << "CPU generator state is too small"; - // 1. 恢复引擎（变长部分直到 seed_ 字段前 ~ 22 字节的固定尾部） - // 引擎的 operator<< 输出是变长的，我们把剩下的 data 一起给 stream - // 但流可能会多读。改用精确长度：data_size 减去尾部固定字段长度。 - constexpr size_t kFooterSize = 8 + 1 + 4 + 1 + 8; // seed + has_float + float + has_double + double - - std::string engine_str; - if (data_size >= kFooterSize) { - engine_str.assign(reinterpret_cast(data), data_size - kFooterSize); - offset = data_size - kFooterSize; - } else { - // 旧格式：没有 footer（向后兼容），整个 data 就是 engine 状态 - engine_str.assign(reinterpret_cast(data), data_size); - offset = data_size; - } + const uint8_t *data = static_cast(state.DataPtr()); + const size_t engine_size = data_size - kStateFooterSize; + std::string engine_str(reinterpret_cast(data), engine_size); std::istringstream iss(engine_str); - iss >> engine_; - - // 2. 恢复种子和正态缓存 - if (offset + kFooterSize <= data_size) { - std::memcpy(&seed_, data + offset, sizeof(seed_)); - offset += sizeof(seed_); - - bool has_float = (data[offset++] != 0); - float float_val; - std::memcpy(&float_val, data + offset, sizeof(float_val)); - offset += sizeof(float_val); - next_float_normal_sample_ = has_float ? std::optional(float_val) : std::nullopt; - - bool has_double = (data[offset++] != 0); - double double_val; - std::memcpy(&double_val, data + offset, sizeof(double_val)); - next_double_normal_sample_ = has_double ? std::optional(double_val) : std::nullopt; - } + 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 = 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 { @@ -113,8 +118,7 @@ std::shared_ptr CPUGeneratorImpl::get_state() const { // 2. 计算总大小 const size_t engine_size = engine_str.size(); - constexpr size_t kFooterSize = 8 + 1 + 4 + 1 + 8; - const size_t total_size = engine_size + kFooterSize; + const size_t total_size = engine_size + kStateFooterSize; auto state_tensor = std::make_shared( std::vector{static_cast(total_size)}, diff --git a/infini_train/src/core/runtime/cuda/cuda_generator_impl.cc b/infini_train/src/core/runtime/cuda/cuda_generator_impl.cc index a4ca0e9b..fac84278 100644 --- a/infini_train/src/core/runtime/cuda/cuda_generator_impl.cc +++ b/infini_train/src/core/runtime/cuda/cuda_generator_impl.cc @@ -75,8 +75,7 @@ uint64_t CUDAGeneratorImpl::seed() { } void CUDAGeneratorImpl::set_state(const Tensor &state) { - CHECK(state.GetDevice().IsCPU()); - CHECK_EQ(static_cast(state.Dtype()), static_cast(DataType::kUINT8)); + ::infini_train::detail::check_rng_state(state); CHECK_EQ(state.SizeInBytes(), kStateSize); const auto *data = static_cast(state.DataPtr()); diff --git a/infini_train/src/generator.cc b/infini_train/src/generator.cc index e1b5b746..a3e2eb2b 100644 --- a/infini_train/src/generator.cc +++ b/infini_train/src/generator.cc @@ -25,6 +25,7 @@ Generator::Generator(std::shared_ptr impl) // ============================================================ void Generator::set_state(const Tensor &state) { + CHECK(state.defined()) << "Undefined tensor is not allowed"; impl_->set_state(state); } @@ -32,6 +33,16 @@ 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 + void manual_seed(uint64_t seed) { core::cpu::manual_seed(seed); From 35020b483beadb7b6aaa5bdf4226e7d9188d3944 Mon Sep 17 00:00:00 2001 From: chaos Date: Sun, 12 Jul 2026 18:11:27 +0800 Subject: [PATCH 08/11] feat: support floating dtypes in random distributions --- infini_train/include/nn/functional.h | 15 +++ infini_train/src/core/distribution_kernels.cc | 99 ++++++++++++++++--- infini_train/src/kernels/cuda/distribution.cu | 54 ++++++++-- infini_train/src/nn/functional.cc | 12 +++ 4 files changed, 159 insertions(+), 21 deletions(-) diff --git a/infini_train/include/nn/functional.h b/infini_train/include/nn/functional.h index e4354fd1..124e914e 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,16 @@ 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); + // Returns a new tensor with the reciprocal of the elements of input. // // Args: diff --git a/infini_train/src/core/distribution_kernels.cc b/infini_train/src/core/distribution_kernels.cc index f5a438f2..953da8f3 100644 --- a/infini_train/src/core/distribution_kernels.cc +++ b/infini_train/src/core/distribution_kernels.cc @@ -8,12 +8,14 @@ /// /// 仿 PyTorch aten/src/ATen/native/cpu/DistributionKernels.cpp（253 行） +#include #include #include "infini_train/include/core/runtime/distribution_kernels.h" #include "infini_train/include/core/runtime/distribution_stubs.h" #include "infini_train/include/core/runtime/distributions_helper.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 { @@ -24,12 +26,79 @@ DEFINE_DISPATCH(normal_stub); namespace { void check_distribution_tensor(const Tensor &tensor) { - CHECK_EQ(static_cast(tensor.Dtype()), static_cast(DataType::kFLOAT32)) - << "Uniform and Normal initialization currently support FLOAT32 tensors only"; + CHECK(IsFloatingPointDType(tensor.Dtype())) + << "Uniform and Normal initialization support floating-point tensors only"; +} + +struct DistributionBounds { + double lowest; + double max; +}; + +DistributionBounds distribution_bounds(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 {}; + } +} + +void check_uniform_parameters(const Tensor &tensor, double from, double to) { + const auto bounds = distribution_bounds(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 check_normal_parameters(double std) { + CHECK_GE(std, 0.0) << "normal expects std >= 0.0, but found std=" << std; } // ---- CPU 内核实现 ---- +template +void uniform_cpu_kernel_impl(Tensor &tensor, double from, double to, + core::cpu::CPUGeneratorImpl *generator) { + auto *buf = static_cast(tensor.DataPtr()); + 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)); + buf[i] = static_cast(value) == to_value ? from_value : value; + } +} + +template +void normal_cpu_kernel_impl(Tensor &tensor, double mean, double std, + core::cpu::CPUGeneratorImpl *generator) { + auto *buf = static_cast(tensor.DataPtr()); + normal_distribution dist(static_cast(mean), static_cast(std)); + for (int64_t i = 0; i < tensor.NumElements(); ++i) { + buf[i] = static_cast(dist(generator)); + } +} + void uniform_cpu_kernel(Tensor &tensor, double from, double to, const std::optional &gen) { CHECK(tensor.GetDevice().IsCPU()); @@ -37,11 +106,13 @@ void uniform_cpu_kernel(Tensor &tensor, double from, double to, gen, core::cpu::getDefaultCPUGenerator()); std::lock_guard lock(cpu_gen->mutex_); - auto *buf = static_cast(tensor.DataPtr()); - uniform_real_distribution dist(static_cast(from), static_cast(to)); - for (int64_t i = 0; i < tensor.NumElements(); ++i) { - buf[i] = dist(cpu_gen); - } + core::cpu::DispatchCpuFunc( + tensor.Dtype(), + [&]() { + using random_t = std::conditional_t, double, float>; + uniform_cpu_kernel_impl(tensor, from, to, cpu_gen); + }, + "CPU uniform"); } void normal_cpu_kernel(Tensor &tensor, double mean, double std, @@ -51,11 +122,13 @@ void normal_cpu_kernel(Tensor &tensor, double mean, double std, gen, core::cpu::getDefaultCPUGenerator()); std::lock_guard lock(cpu_gen->mutex_); - auto *buf = static_cast(tensor.DataPtr()); - normal_distribution dist(static_cast(mean), static_cast(std)); - for (int64_t i = 0; i < tensor.NumElements(); ++i) { - buf[i] = dist(cpu_gen); - } + core::cpu::DispatchCpuFunc( + tensor.Dtype(), + [&]() { + using random_t = std::conditional_t, double, float>; + normal_cpu_kernel_impl(tensor, mean, std, cpu_gen); + }, + "CPU normal"); } } // namespace @@ -73,12 +146,14 @@ REGISTER_DISPATCH(normal_stub, Device::DeviceType::kCPU, &normal_cpu_kernel); void uniform_kernel(Tensor &tensor, double from, double to, const std::optional &gen) { check_distribution_tensor(tensor); + check_uniform_parameters(tensor, from, to); uniform_stub(tensor.GetDevice().type(), tensor, from, to, gen); } void normal_kernel(Tensor &tensor, double mean, double std, const std::optional &gen) { check_distribution_tensor(tensor); + check_normal_parameters(std); normal_stub(tensor.GetDevice().type(), tensor, mean, std, gen); } diff --git a/infini_train/src/kernels/cuda/distribution.cu b/infini_train/src/kernels/cuda/distribution.cu index b5ab9903..4b181fdc 100644 --- a/infini_train/src/kernels/cuda/distribution.cu +++ b/infini_train/src/kernels/cuda/distribution.cu @@ -5,10 +5,12 @@ #include #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/core/runtime/distribution_stubs.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" @@ -17,7 +19,24 @@ namespace { constexpr int kThreadsPerBlock = 256; -__global__ void UniformKernel(float *data, int64_t n, float from, float to, +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) { @@ -26,11 +45,14 @@ __global__ void UniformKernel(float *data, int64_t n, float from, float to, curandStatePhilox4_32_10_t state; curand_init(seed, subsequence + static_cast(index), 0, &state); - const float unit = static_cast(curand(&state)) * 0x1p-32f; - data[index] = from + unit * (to - from); + 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)); + data[index] = common::cuda::Cast(value) == common::cuda::Cast(to_value) ? from_value : value; } -__global__ void NormalKernel(float *data, int64_t n, float mean, float std, +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) { @@ -39,7 +61,7 @@ __global__ void NormalKernel(float *data, int64_t n, float mean, float std, curandStatePhilox4_32_10_t state; curand_init(seed, subsequence + static_cast(index), 0, &state); - data[index] = mean + curand_normal(&state) * std; + data[index] = common::cuda::Cast(mean + normal_sample(&state) * std); } const core::cuda::CudaStream *get_cuda_stream(const Device &device) { @@ -69,8 +91,15 @@ void uniform_cuda_kernel(Tensor &tensor, double from, double to, const int blocks = static_cast((n + kThreadsPerBlock - 1) / kThreadsPerBlock); const auto *stream = get_cuda_stream(device); - UniformKernel<<cuda_stream()>>>( - static_cast(tensor.DataPtr()), n, static_cast(from), static_cast(to), seed, subsequence); + 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()); } @@ -96,8 +125,15 @@ void normal_cuda_kernel(Tensor &tensor, double mean, double std, const int blocks = static_cast((n + kThreadsPerBlock - 1) / kThreadsPerBlock); const auto *stream = get_cuda_stream(device); - NormalKernel<<cuda_stream()>>>( - static_cast(tensor.DataPtr()), n, static_cast(mean), static_cast(std), seed, subsequence); + 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()); } diff --git a/infini_train/src/nn/functional.cc b/infini_train/src/nn/functional.cc index b02f185a..a80c995c 100644 --- a/infini_train/src/nn/functional.cc +++ b/infini_train/src/nn/functional.cc @@ -27,6 +27,18 @@ 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 Reciprocal(const std::shared_ptr &input) { return input->Reciprocal(); } std::shared_ptr Sin(const std::shared_ptr &input) { return input->Sin(); } From 10834ae579de148afd758a3d656ff9dba7e73955 Mon Sep 17 00:00:00 2001 From: chaos Date: Sun, 12 Jul 2026 18:39:55 +0800 Subject: [PATCH 09/11] docs: update generator task progress --- ...75\350\261\241\351\200\211\351\242\230.md" | 31 ++++++++++--------- 1 file changed, 16 insertions(+), 15 deletions(-) diff --git "a/docs/\343\200\2202026 \346\230\245\345\255\243\344\272\272\345\267\245\346\231\272\350\203\275\345\244\247\350\265\233\343\200\221Generator \346\212\275\350\261\241\351\200\211\351\242\230.md" "b/docs/\343\200\2202026 \346\230\245\345\255\243\344\272\272\345\267\245\346\231\272\350\203\275\345\244\247\350\265\233\343\200\221Generator \346\212\275\350\261\241\351\200\211\351\242\230.md" index f550bee3..2528d3f3 100644 --- "a/docs/\343\200\2202026 \346\230\245\345\255\243\344\272\272\345\267\245\346\231\272\350\203\275\345\244\247\350\265\233\343\200\221Generator \346\212\275\350\261\241\351\200\211\351\242\230.md" +++ "b/docs/\343\200\2202026 \346\230\245\345\255\243\344\272\272\345\267\245\346\231\272\350\203\275\345\244\247\350\265\233\343\200\221Generator \346\212\275\350\261\241\351\200\211\351\242\230.md" @@ -17,13 +17,14 @@ 在深度学习训练与推理框架中，随机数生成器（Generator）是支撑参数初始化、随机采样、Dropout、噪声注入等功能的重要基础设施。当前框架中相关能力尚不完善，缺少统一的 Generator 抽象、后端实现以及全局随机种子管理机制，导致随机相关算子的行为与主流框架存在差距，也难以满足后续功能扩展需求。 -当前已完成 Generator 基础设施与初始化算子的首轮接入，进度如下： +当前已完成 Generator 基础设施、初始化算子与随机张量创建算子的首轮接入，进度如下： - [x] 已建立统一的 Generator 抽象，支持随机数状态、种子设置、状态获取与恢复； - [x] 已实现 CPU 与 CUDA 后端 Generator，并保持统一的调用接口； - [x] 已建立按设备维护的默认 Generator；未显式传入 generator 时，初始化算子使用目标 Tensor 所在设备的默认随机源； - [x] 已提供全局 `manual_seed(uint64_t)` 入口，重置 CPU 及所有 CUDA 默认 Generator； -- [ ] Dropout、`rand`、`randn` 等训练期或通用随机算子尚未接入；distribution 当前仅支持 `FLOAT32`，也未实现 CUDA Graph 语义。 +- [x] 已接入 `Rand`、`Randn`，支持显式/默认 Generator、CPU/CUDA 及 `FLOAT16`、`BFLOAT16`、`FLOAT32`、`FLOAT64`； +- [ ] Dropout 与 CUDA Graph RNG 语义尚未实现。 为补齐这一基础能力，本题目要求参赛者参考主流深度学习框架的设计思路，完成一套可扩展的随机数生成器基础设施，实现 CPU 与 CUDA 后端的 Generator，并支持基于 Generator 的随机数生成与种子控制能力。 @@ -33,9 +34,9 @@ - [x] 1. 设计并实现统一的 Generator 抽象接口，支持随机数状态管理、种子设置、状态获取与恢复等基础能力。 - [x] 2. 分别实现 CPU 与 CUDA 后端对应的 GeneratorImpl，使不同设备上的随机数生成具备统一的调用方式。 -- [x] 3. 建立全局随机种子控制机制，支持通过统一入口固定随机种子；已接入的初始化算子可稳定复现。 -- [x] 4. 改造初始化类随机算子，使其在未显式传入 generator 参数时，自动使用目标设备的默认全局 Generator。 -- [x] 5. 已完成 CPU/CUDA 初始化路径的基本行为、可复现性和接口语义冒烟验证；仓库内正式测试仍待补充。 +- [x] 3. 建立全局随机种子控制机制，支持通过统一入口固定随机种子；已接入的初始化算子和随机张量创建算子可稳定复现。 +- [x] 4. 改造初始化类随机算子和随机张量创建算子，使其在未显式传入 generator 参数时，自动使用目标设备的默认全局 Generator。 +- [x] 5. 已完成 CPU/CUDA 初始化、`Rand/Randn`、多 dtype 分布和参数校验的基本行为、可复现性和接口语义冒烟验证；仓库内正式测试仍待补充。 ## 三、任务拆解 @@ -76,7 +77,7 @@ - [x] `GetState()` / `SetState()` 的状态序列化、反序列化与格式校验；CPU/CUDA 均校验 CPU `UINT8` state，CPU 额外校验 footer 与 mt19937 反序列化； - [x] 多次随机调用后，随机序列连续推进；CPU 推进 mt19937，CUDA 为每次 kernel 保留 Philox subsequence； - [x] 同一 seed、同一设备、同一调用顺序下结果可复现； -- [x] 初始化算子已支持显式 Generator 和默认 Generator 两条调用路径。 +- [x] 初始化算子和 `Rand/Randn` 已支持显式 Generator 和默认 Generator 两条调用路径。 CUDA 后端建议在状态中考虑 seed、offset、counter 等信息，避免不同 kernel 调用之间随机序列重叠。此部分需要保证功能正确、语义清晰、状态可管理、接口一致，为后续随机算子和分布式扩展提供基础。 @@ -101,7 +102,7 @@ CUDA 后端建议在状态中考虑 seed、offset、counter 等信息，避免 - [x] 不同 CUDA 设备的默认 Generator 相互独立； - [x] 用户显式传入 Generator 时，初始化算子使用用户指定 Generator； - [x] 未传入 Generator 时，初始化算子使用目标设备默认 Generator； -- [x] 设置统一 seed 后，已接入的参数初始化和随机采样可稳定复现； +- [x] 设置统一 seed 后，已接入的参数初始化和 `Rand/Randn` 随机采样可稳定复现； - [ ] Dropout 等尚未接入的随机算子仍待完成。 这一部分是本题的核心验收点。随机数系统最终服务于训练可复现性，因此必须通过统一入口稳定控制随机行为。 @@ -111,25 +112,25 @@ CUDA 后端建议在状态中考虑 seed、offset、counter 等信息，避免 改造框架中已有的随机相关使用，使其接入 Generator 机制。建议至少覆盖两类场景： - [x] 初始化类随机算子：`Uniform`、`Normal`、`KaimingUniform` 已接入 Generator，并按 Tensor device 分发 CPU/CUDA kernel； -- [ ] 训练过程中的随机算子，如 dropout、rand、randn 等； +- [x] 通用随机张量创建算子：`Rand`、`Randn` 已接入 Generator，支持 CPU/CUDA 与四种浮点 dtype；Dropout 仍待接入； 改造后的算子应满足： -- [x] 已接入初始化算子支持显式传入 Generator； -- [x] 已接入初始化算子未传入 Generator 时，自动使用目标设备默认 Generator； -- [x] 已接入初始化算子的随机结果受统一 `manual_seed` 入口控制； -- [x] 已接入初始化算子多次调用会推进 Generator 状态； -- [x] 已接入初始化算子在同一 seed、同一调用顺序下可复现。 +- [x] 已接入初始化算子和 `Rand/Randn` 支持显式传入 Generator； +- [x] 已接入初始化算子和 `Rand/Randn` 未传入 Generator 时，自动使用目标设备默认 Generator； +- [x] 已接入初始化算子和 `Rand/Randn` 的随机结果受统一 `manual_seed` 入口控制； +- [x] 已接入初始化算子和 `Rand/Randn` 多次调用会推进 Generator 状态； +- [x] 已接入初始化算子和 `Rand/Randn` 在同一 seed、同一调用顺序下可复现。 不要求一次性改造所有随机算子，但应至少完成一个初始化类算子和一个训练期随机算子，以证明 Generator 机制能够贯通框架层与算子层。 -当前仅完成初始化类算子接入；训练期随机算子仍是后续工作。 +当前已完成初始化类算子与通用随机张量创建算子接入；Dropout 等训练期随机算子仍是后续工作。 ### 5. 测试与对齐验证 为验证实现结果，需要补充系统化测试。测试重点应放在"语义是否正确"和"是否可复现"，而不是与 PyTorch 逐元素数值一致。 -当前已在 `/tmp` 完成 CPU/CUDA 构建和 Generator、初始化可复现性的冒烟验证；按项目约定尚未新增仓库内测试，因此以下正式测试清单仍未完成。 +当前已在 `/tmp` 完成 CPU/CUDA 构建，以及 Generator 状态恢复、初始化、`Rand/Randn`、四种浮点 dtype 分布、参数校验和显式/默认 Generator 路径的冒烟验证；按项目约定尚未新增仓库内测试，因此以下正式测试清单仍未完成。 建议至少覆盖以下内容： From 1e2fe858eda570cdc8e18540ff6de3778f2528d1 Mon Sep 17 00:00:00 2001 From: chaos Date: Sun, 12 Jul 2026 21:26:23 +0800 Subject: [PATCH 10/11] feat: add generator dropout support --- .gitignore | 1 + CMakeLists.txt | 6 + benchmarks/CMakeLists.txt | 6 + benchmarks/generator/README.md | 69 ++++ benchmarks/generator/generator_benchmark.cc | 228 +++++++++++++ .../generator/run_generator_benchmark.sh | 46 +++ ...75\350\261\241\351\200\211\351\242\230.md" | 33 +- infini_train/include/autograd/dropout.h | 34 ++ infini_train/include/autograd/elementwise.h | 2 + .../include/core/runtime/dropout_kernels.h | 16 + .../include/core/runtime/dropout_stubs.h | 19 ++ infini_train/include/nn/functional.h | 3 + infini_train/src/autograd/dropout.cc | 37 +++ infini_train/src/autograd/elementwise.cc | 12 +- infini_train/src/core/dropout_kernels.cc | 43 +++ infini_train/src/kernels/cpu/dropout.cc | 100 ++++++ infini_train/src/kernels/cuda/dropout.cu | 144 +++++++++ infini_train/src/nn/functional.cc | 12 + tests/CMakeLists.txt | 3 + .../test_autograd_elementwise_backward.cc | 12 +- tests/generator/CMakeLists.txt | 19 ++ .../cuda_only/test_cuda_generator.cc | 72 +++++ tests/generator/test_generator.cc | 305 ++++++++++++++++++ tests/generator/test_generator_interface.cc | 64 ++++ 24 files changed, 1267 insertions(+), 19 deletions(-) create mode 100644 benchmarks/CMakeLists.txt create mode 100644 benchmarks/generator/README.md create mode 100644 benchmarks/generator/generator_benchmark.cc create mode 100644 benchmarks/generator/run_generator_benchmark.sh create mode 100644 infini_train/include/autograd/dropout.h create mode 100644 infini_train/include/core/runtime/dropout_kernels.h create mode 100644 infini_train/include/core/runtime/dropout_stubs.h create mode 100644 infini_train/src/autograd/dropout.cc create mode 100644 infini_train/src/core/dropout_kernels.cc create mode 100644 infini_train/src/kernels/cpu/dropout.cc create mode 100644 infini_train/src/kernels/cuda/dropout.cu create mode 100644 tests/generator/CMakeLists.txt create mode 100644 tests/generator/cuda_only/test_cuda_generator.cc create mode 100644 tests/generator/test_generator.cc create mode 100644 tests/generator/test_generator_interface.cc diff --git a/.gitignore b/.gitignore index b37295bb..24a3e0aa 100644 --- a/.gitignore +++ b/.gitignore @@ -1,6 +1,7 @@ build/ build-cpu/ build-full/ +/build-*/ .cache/ .vscode/ diff --git a/CMakeLists.txt b/CMakeLists.txt index 4c6da822..3125a75a 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -217,6 +217,12 @@ if(BUILD_TEST) add_subdirectory(tests) endif() +# Benchmarks +option(BUILD_BENCHMARK "Build InfiniTrain benchmarks" OFF) +if(BUILD_BENCHMARK) + add_subdirectory(benchmarks) +endif() + # Negative compile test: missing dtype registration must fail at compile time. set(DTYPE_DISPATCH_COMPILE_FAIL_SOURCE ${PROJECT_SOURCE_DIR}/tests/dtype/test_dtype_dispatch_compile_fail.cc) diff --git a/benchmarks/CMakeLists.txt b/benchmarks/CMakeLists.txt new file mode 100644 index 00000000..16df2ff4 --- /dev/null +++ b/benchmarks/CMakeLists.txt @@ -0,0 +1,6 @@ +add_executable(generator_benchmark generator/generator_benchmark.cc) +link_infini_train_exe(generator_benchmark) + +set_target_properties(generator_benchmark PROPERTIES + RUNTIME_OUTPUT_DIRECTORY ${CMAKE_BINARY_DIR}/benchmarks +) diff --git a/benchmarks/generator/README.md b/benchmarks/generator/README.md new file mode 100644 index 00000000..306d6e6f --- /dev/null +++ b/benchmarks/generator/README.md @@ -0,0 +1,69 @@ +# Generator benchmark + +This benchmark measures Generator state-management latency and end-to-end +`Uniform` / `Normal` tensor fill throughput. It is intentionally separate from +the GoogleTest correctness suite: benchmark results are informative and do not +make tests flaky. + +## Build + +From the repository root: + +```bash +cmake -S . -B build \ + -DCMAKE_BUILD_TYPE=Release \ + -DUSE_CUDA=ON \ + -DUSE_NCCL=OFF \ + -DBUILD_TEST=ON \ + -DBUILD_BENCHMARK=ON +cmake --build build -j +``` + +For a CPU-only build, use `-DUSE_CUDA=OFF`. + +## Correctness tests + +```bash +ctest --test-dir build -L cpu --output-on-failure +ctest --test-dir build -L cuda --output-on-failure +``` + +To run only this feature's tests: + +```bash +ctest --test-dir build -R Generator --output-on-failure +``` + +## Single benchmark + +```bash +./build/benchmarks/generator_benchmark \ + --device cuda \ + --device-index 0 \ + --op all \ + --generator explicit \ + --elements 1048576 \ + --warmup 10 \ + --iterations 100 \ + --seed 42 +``` + +The output is CSV with average latency, generated samples per second, and +effective output bandwidth. CUDA timing synchronizes the device before and +after the measured loop, so it includes kernel execution rather than only +asynchronous launch overhead. + +## Standard matrix + +```bash +bash benchmarks/generator/run_generator_benchmark.sh \ + ./build/benchmarks/generator_benchmark cpu > generator_cpu.csv + +bash benchmarks/generator/run_generator_benchmark.sh \ + ./build/benchmarks/generator_benchmark cuda > generator_cuda.csv +``` + +For comparable reports, keep the compiler, build type, CPU/GPU model, CUDA +version, warmup count, iteration count, and thread-related environment +variables fixed. Compare before/after InfiniTrain commits rather than requiring +element-wise or performance parity with PyTorch. diff --git a/benchmarks/generator/generator_benchmark.cc b/benchmarks/generator/generator_benchmark.cc new file mode 100644 index 00000000..e9553872 --- /dev/null +++ b/benchmarks/generator/generator_benchmark.cc @@ -0,0 +1,228 @@ +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#include "infini_train/include/core/runtime/device_guard.h" +#include "infini_train/include/generator.h" +#include "infini_train/include/nn/init.h" +#include "infini_train/include/tensor.h" +#include "infini_train/src/core/runtime/cpu/cpu_generator_impl.h" +#if defined(USE_CUDA) +#include "infini_train/src/core/runtime/cuda/cuda_generator_impl.h" +#endif + +namespace { + +using Clock = std::chrono::steady_clock; +using infini_train::DataType; +using infini_train::Device; +using infini_train::Generator; +using infini_train::Tensor; + +struct Options { + std::string device = "cpu"; + std::string operation = "all"; + std::string generator = "explicit"; + int device_index = 0; + int64_t elements = 1 << 20; + int warmup = 10; + int iterations = 100; + uint64_t seed = 42; +}; + +void PrintUsage(const char *program) { + std::cout + << "Usage: " << program << " [options]\n" + << " --device cpu|cuda\n" + << " --device-index N\n" + << " --op uniform|normal|state|all\n" + << " --generator explicit|default\n" + << " --elements N\n" + << " --warmup N\n" + << " --iterations N\n" + << " --seed N\n"; +} + +std::string RequireValue(int argc, char **argv, int &index) { + if (++index >= argc) { + throw std::invalid_argument(std::string("missing value for ") + argv[index - 1]); + } + return argv[index]; +} + +Options ParseOptions(int argc, char **argv) { + Options options; + for (int i = 1; i < argc; ++i) { + const std::string_view argument(argv[i]); + if (argument == "--help" || argument == "-h") { + PrintUsage(argv[0]); + std::exit(0); + } else if (argument == "--device") { + options.device = RequireValue(argc, argv, i); + } else if (argument == "--device-index") { + options.device_index = std::stoi(RequireValue(argc, argv, i)); + } else if (argument == "--op") { + options.operation = RequireValue(argc, argv, i); + } else if (argument == "--generator") { + options.generator = RequireValue(argc, argv, i); + } else if (argument == "--elements") { + options.elements = std::stoll(RequireValue(argc, argv, i)); + } else if (argument == "--warmup") { + options.warmup = std::stoi(RequireValue(argc, argv, i)); + } else if (argument == "--iterations") { + options.iterations = std::stoi(RequireValue(argc, argv, i)); + } else if (argument == "--seed") { + options.seed = std::stoull(RequireValue(argc, argv, i)); + } else { + throw std::invalid_argument(std::string("unknown option: ") + argv[i]); + } + } + + if (options.device != "cpu" && options.device != "cuda") { + throw std::invalid_argument("--device must be cpu or cuda"); + } + if (options.operation != "uniform" && options.operation != "normal" + && options.operation != "state" && options.operation != "all") { + throw std::invalid_argument("--op must be uniform, normal, state, or all"); + } + if (options.generator != "explicit" && options.generator != "default") { + throw std::invalid_argument("--generator must be explicit or default"); + } + if (options.elements <= 0 || options.warmup < 0 || options.iterations <= 0) { + throw std::invalid_argument("elements and iterations must be positive; warmup must be non-negative"); + } + return options; +} + +Device MakeDevice(const Options &options) { + if (options.device == "cpu") { + return Device(Device::DeviceType::kCPU, 0); + } +#if defined(USE_CUDA) + return Device(Device::DeviceType::kCUDA, options.device_index); +#else + throw std::invalid_argument("CUDA benchmark requested, but InfiniTrain was built without USE_CUDA"); +#endif +} + +Generator MakeGenerator(Device device, uint64_t seed) { + if (device.IsCPU()) { + return infini_train::core::cpu::createCPUGenerator(seed); + } +#if defined(USE_CUDA) + return infini_train::core::cuda::createCUDAGenerator(device.index(), seed); +#else + (void)seed; + throw std::invalid_argument("CUDA support is disabled"); +#endif +} + +void Synchronize(Device device) { + infini_train::core::GetDeviceGuardImpl(device.type())->SynchronizeDevice(device); +} + +template +double MeasureMicroseconds(Device device, int warmup, int iterations, Function &&function) { + for (int i = 0; i < warmup; ++i) { + function(); + } + Synchronize(device); + const auto start = Clock::now(); + for (int i = 0; i < iterations; ++i) { + function(); + } + Synchronize(device); + const auto end = Clock::now(); + return std::chrono::duration(end - start).count() / iterations; +} + +void PrintHeader() { + std::cout << "device,device_index,operation,generator,elements,iterations,latency_us,gsamples_s,bandwidth_gbps\n"; +} + +void PrintResult(const Options &options, std::string_view operation, double latency_us) { + const double seconds = latency_us * 1e-6; + const double samples_per_second = static_cast(options.elements) / seconds; + const double gsamples_per_second = samples_per_second / 1e9; + const double bandwidth_gbps = samples_per_second * sizeof(float) / 1e9; + std::cout << options.device << ',' << options.device_index << ',' << operation << ',' + << options.generator << ',' << options.elements << ',' << options.iterations << ',' + << std::fixed << std::setprecision(3) << latency_us << ',' << gsamples_per_second << ',' + << bandwidth_gbps << '\n'; +} + +void RunDistribution(const Options &options, Device device, std::string_view operation) { + Generator explicit_generator = MakeGenerator(device, options.seed); + std::optional generator = options.generator == "explicit" + ? std::optional(explicit_generator) + : std::nullopt; + if (!generator) { + infini_train::manual_seed(options.seed); + } + auto tensor = std::make_shared( + std::vector{options.elements}, DataType::kFLOAT32, device); + + const double latency_us = MeasureMicroseconds(device, options.warmup, options.iterations, [&] { + if (operation == "uniform") { + infini_train::nn::init::Uniform(tensor, 0.0f, 1.0f, generator); + } else { + infini_train::nn::init::Normal(tensor, 0.0f, 1.0f, generator); + } + }); + PrintResult(options, operation, latency_us); +} + +void RunStateBenchmark(const Options &options, Device device) { + Generator generator = MakeGenerator(device, options.seed); + auto state = generator.get_state(); + const double get_state_us = MeasureMicroseconds(device, options.warmup, options.iterations, [&] { + state = generator.get_state(); + }); + const double set_state_us = MeasureMicroseconds(device, options.warmup, options.iterations, [&] { + generator.set_state(*state); + }); + uint64_t seed = options.seed; + const double manual_seed_us = MeasureMicroseconds(device, options.warmup, options.iterations, [&] { + generator.set_current_seed(seed++); + }); + + Options state_options = options; + // State-management operations do not generate tensor elements. Reporting + // zero avoids presenting meaningless throughput numbers for these rows. + state_options.elements = 0; + PrintResult(state_options, "get_state", get_state_us); + PrintResult(state_options, "set_state", set_state_us); + PrintResult(state_options, "manual_seed", manual_seed_us); +} + +} // namespace + +int main(int argc, char **argv) { + try { + const Options options = ParseOptions(argc, argv); + const Device device = MakeDevice(options); + PrintHeader(); + if (options.operation == "uniform" || options.operation == "all") { + RunDistribution(options, device, "uniform"); + } + if (options.operation == "normal" || options.operation == "all") { + RunDistribution(options, device, "normal"); + } + if (options.operation == "state" || options.operation == "all") { + RunStateBenchmark(options, device); + } + return 0; + } catch (const std::exception &error) { + std::cerr << "generator_benchmark: " << error.what() << '\n'; + PrintUsage(argv[0]); + return 2; + } +} diff --git a/benchmarks/generator/run_generator_benchmark.sh b/benchmarks/generator/run_generator_benchmark.sh new file mode 100644 index 00000000..b408a9ed --- /dev/null +++ b/benchmarks/generator/run_generator_benchmark.sh @@ -0,0 +1,46 @@ +#!/usr/bin/env bash +set -euo pipefail + +binary="${1:-./build/benchmarks/generator_benchmark}" +device="${2:-cpu}" +device_index="${DEVICE_INDEX:-0}" +warmup="${WARMUP:-10}" +iterations="${ITERATIONS:-100}" +seed="${SEED:-42}" + +if [[ ! -x "${binary}" ]]; then + echo "benchmark executable not found or not executable: ${binary}" >&2 + exit 2 +fi + +echo "# timestamp=$(date --iso-8601=seconds)" >&2 +echo "# host=$(hostname)" >&2 +echo "# kernel=$(uname -srmo)" >&2 +echo "# binary=${binary}" >&2 +echo "device,device_index,operation,generator,elements,iterations,latency_us,gsamples_s,bandwidth_gbps" + +for generator in explicit default; do + for operation in uniform normal; do + for elements in 1024 1048576 16777216; do + "${binary}" \ + --device "${device}" \ + --device-index "${device_index}" \ + --op "${operation}" \ + --generator "${generator}" \ + --elements "${elements}" \ + --warmup "${warmup}" \ + --iterations "${iterations}" \ + --seed "${seed}" | awk 'NR > 1' + done + done +done + +"${binary}" \ + --device "${device}" \ + --device-index "${device_index}" \ + --op state \ + --generator explicit \ + --elements 1 \ + --warmup "${warmup}" \ + --iterations "${iterations}" \ + --seed "${seed}" | awk 'NR > 1' diff --git "a/docs/\343\200\2202026 \346\230\245\345\255\243\344\272\272\345\267\245\346\231\272\350\203\275\345\244\247\350\265\233\343\200\221Generator \346\212\275\350\261\241\351\200\211\351\242\230.md" "b/docs/\343\200\2202026 \346\230\245\345\255\243\344\272\272\345\267\245\346\231\272\350\203\275\345\244\247\350\265\233\343\200\221Generator \346\212\275\350\261\241\351\200\211\351\242\230.md" index 2528d3f3..67b1ddcc 100644 --- "a/docs/\343\200\2202026 \346\230\245\345\255\243\344\272\272\345\267\245\346\231\272\350\203\275\345\244\247\350\265\233\343\200\221Generator \346\212\275\350\261\241\351\200\211\351\242\230.md" +++ "b/docs/\343\200\2202026 \346\230\245\345\255\243\344\272\272\345\267\245\346\231\272\350\203\275\345\244\247\350\265\233\343\200\221Generator \346\212\275\350\261\241\351\200\211\351\242\230.md" @@ -17,14 +17,14 @@ 在深度学习训练与推理框架中，随机数生成器（Generator）是支撑参数初始化、随机采样、Dropout、噪声注入等功能的重要基础设施。当前框架中相关能力尚不完善，缺少统一的 Generator 抽象、后端实现以及全局随机种子管理机制，导致随机相关算子的行为与主流框架存在差距，也难以满足后续功能扩展需求。 -当前已完成 Generator 基础设施、初始化算子与随机张量创建算子的首轮接入，进度如下： +当前已完成 Generator 基础设施、初始化算子、随机张量创建算子与普通 Dropout 的首轮接入，进度如下： - [x] 已建立统一的 Generator 抽象，支持随机数状态、种子设置、状态获取与恢复； - [x] 已实现 CPU 与 CUDA 后端 Generator，并保持统一的调用接口； - [x] 已建立按设备维护的默认 Generator；未显式传入 generator 时，初始化算子使用目标 Tensor 所在设备的默认随机源； - [x] 已提供全局 `manual_seed(uint64_t)` 入口，重置 CPU 及所有 CUDA 默认 Generator； - [x] 已接入 `Rand`、`Randn`，支持显式/默认 Generator、CPU/CUDA 及 `FLOAT16`、`BFLOAT16`、`FLOAT32`、`FLOAT64`； -- [ ] Dropout 与 CUDA Graph RNG 语义尚未实现。 +- [x] 已接入普通 `Dropout`，训练态保存随机 mask 供反向传播复用，支持显式/默认 Generator、CPU/CUDA 及四种浮点 dtype； 为补齐这一基础能力，本题目要求参赛者参考主流深度学习框架的设计思路，完成一套可扩展的随机数生成器基础设施，实现 CPU 与 CUDA 后端的 Generator，并支持基于 Generator 的随机数生成与种子控制能力。 @@ -34,9 +34,9 @@ - [x] 1. 设计并实现统一的 Generator 抽象接口，支持随机数状态管理、种子设置、状态获取与恢复等基础能力。 - [x] 2. 分别实现 CPU 与 CUDA 后端对应的 GeneratorImpl，使不同设备上的随机数生成具备统一的调用方式。 -- [x] 3. 建立全局随机种子控制机制，支持通过统一入口固定随机种子；已接入的初始化算子和随机张量创建算子可稳定复现。 -- [x] 4. 改造初始化类随机算子和随机张量创建算子，使其在未显式传入 generator 参数时，自动使用目标设备的默认全局 Generator。 -- [x] 5. 已完成 CPU/CUDA 初始化、`Rand/Randn`、多 dtype 分布和参数校验的基本行为、可复现性和接口语义冒烟验证；仓库内正式测试仍待补充。 +- [x] 3. 建立全局随机种子控制机制，支持通过统一入口固定随机种子；已接入的初始化算子、随机张量创建算子和 Dropout 可稳定复现。 +- [x] 4. 改造初始化类随机算子、随机张量创建算子和 Dropout，使其在未显式传入 generator 参数时，自动使用目标设备的默认全局 Generator。 +- [x] 5. 已完成 CPU/CUDA 初始化、`Rand/Randn`、Dropout、多 dtype 分布和参数校验的基本行为、可复现性和接口语义冒烟验证；仓库内正式测试仍待补充。 ## 三、任务拆解 @@ -77,7 +77,7 @@ - [x] `GetState()` / `SetState()` 的状态序列化、反序列化与格式校验；CPU/CUDA 均校验 CPU `UINT8` state，CPU 额外校验 footer 与 mt19937 反序列化； - [x] 多次随机调用后，随机序列连续推进；CPU 推进 mt19937，CUDA 为每次 kernel 保留 Philox subsequence； - [x] 同一 seed、同一设备、同一调用顺序下结果可复现； -- [x] 初始化算子和 `Rand/Randn` 已支持显式 Generator 和默认 Generator 两条调用路径。 +- [x] 初始化算子、`Rand/Randn` 和 Dropout 已支持显式 Generator 和默认 Generator 两条调用路径。 CUDA 后端建议在状态中考虑 seed、offset、counter 等信息，避免不同 kernel 调用之间随机序列重叠。此部分需要保证功能正确、语义清晰、状态可管理、接口一致，为后续随机算子和分布式扩展提供基础。 @@ -102,8 +102,8 @@ CUDA 后端建议在状态中考虑 seed、offset、counter 等信息，避免 - [x] 不同 CUDA 设备的默认 Generator 相互独立； - [x] 用户显式传入 Generator 时，初始化算子使用用户指定 Generator； - [x] 未传入 Generator 时，初始化算子使用目标设备默认 Generator； -- [x] 设置统一 seed 后，已接入的参数初始化和 `Rand/Randn` 随机采样可稳定复现； -- [ ] Dropout 等尚未接入的随机算子仍待完成。 +- [x] 设置统一 seed 后，已接入的参数初始化、`Rand/Randn` 随机采样和 Dropout mask 可稳定复现； +- [ ] CUDA Graph RNG 状态管理仍待完成。 这一部分是本题的核心验收点。随机数系统最终服务于训练可复现性，因此必须通过统一入口稳定控制随机行为。 @@ -112,25 +112,26 @@ CUDA 后端建议在状态中考虑 seed、offset、counter 等信息，避免 改造框架中已有的随机相关使用，使其接入 Generator 机制。建议至少覆盖两类场景： - [x] 初始化类随机算子：`Uniform`、`Normal`、`KaimingUniform` 已接入 Generator，并按 Tensor device 分发 CPU/CUDA kernel； -- [x] 通用随机张量创建算子：`Rand`、`Randn` 已接入 Generator，支持 CPU/CUDA 与四种浮点 dtype；Dropout 仍待接入； +- [x] 通用随机张量创建算子：`Rand`、`Randn` 已接入 Generator，支持 CPU/CUDA 与四种浮点 dtype； +- [x] 训练过程随机算子：普通 Dropout 已接入 Generator，前向保存 mask，反向复用同一 mask，不重新消耗随机状态； 改造后的算子应满足： -- [x] 已接入初始化算子和 `Rand/Randn` 支持显式传入 Generator； -- [x] 已接入初始化算子和 `Rand/Randn` 未传入 Generator 时，自动使用目标设备默认 Generator； -- [x] 已接入初始化算子和 `Rand/Randn` 的随机结果受统一 `manual_seed` 入口控制； -- [x] 已接入初始化算子和 `Rand/Randn` 多次调用会推进 Generator 状态； -- [x] 已接入初始化算子和 `Rand/Randn` 在同一 seed、同一调用顺序下可复现。 +- [x] 已接入初始化算子、`Rand/Randn` 和 Dropout 支持显式传入 Generator； +- [x] 已接入初始化算子、`Rand/Randn` 和 Dropout 未传入 Generator 时，自动使用目标设备默认 Generator； +- [x] 已接入初始化算子、`Rand/Randn` 和 Dropout 的随机结果受统一 `manual_seed` 入口控制； +- [x] 已接入初始化算子、`Rand/Randn` 和 Dropout 多次调用会推进 Generator 状态； +- [x] 已接入初始化算子、`Rand/Randn` 和 Dropout 在同一 seed、同一调用顺序下可复现。 不要求一次性改造所有随机算子，但应至少完成一个初始化类算子和一个训练期随机算子，以证明 Generator 机制能够贯通框架层与算子层。 -当前已完成初始化类算子与通用随机张量创建算子接入；Dropout 等训练期随机算子仍是后续工作。 +当前已完成初始化类算子、通用随机张量创建算子和普通 Dropout 接入；后续可扩展 AlphaDropout、FeatureDropout、CUDA Graph RNG 等能力。 ### 5. 测试与对齐验证 为验证实现结果，需要补充系统化测试。测试重点应放在"语义是否正确"和"是否可复现"，而不是与 PyTorch 逐元素数值一致。 -当前已在 `/tmp` 完成 CPU/CUDA 构建，以及 Generator 状态恢复、初始化、`Rand/Randn`、四种浮点 dtype 分布、参数校验和显式/默认 Generator 路径的冒烟验证；按项目约定尚未新增仓库内测试，因此以下正式测试清单仍未完成。 +当前已在 `/tmp` 完成 CPU/CUDA 构建，以及 Generator 状态恢复、初始化、`Rand/Randn`、Dropout 前后向、四种浮点 dtype 分布、参数校验和显式/默认 Generator 路径的冒烟验证；按项目约定尚未新增仓库内测试，因此以下正式测试清单仍未完成。 建议至少覆盖以下内容： diff --git a/infini_train/include/autograd/dropout.h b/infini_train/include/autograd/dropout.h new file mode 100644 index 00000000..9b8451d2 --- /dev/null +++ b/infini_train/include/autograd/dropout.h @@ -0,0 +1,34 @@ +#pragma once + +#include +#include +#include + +#include "infini_train/include/autograd/function.h" +#include "infini_train/include/generator.h" + +namespace infini_train { +class Tensor; +} + +namespace infini_train::autograd { + +class Dropout final : public Function { +public: + static constexpr char kType[] = "DropoutFunction"; + + Dropout(double p, std::optional generator) + : Function(kType), p_(p), generator_(std::move(generator)) {} + + std::vector> Forward(const std::vector> &input_tensors) override; + void SetupContext(const std::vector> &input_tensors, + const std::vector> &output_tensors) override; + std::vector> Backward(const std::vector> &grad_outputs) override; + +private: + double p_ = 0.0; + std::optional generator_; + std::shared_ptr mask_; +}; + +} // namespace infini_train::autograd diff --git a/infini_train/include/autograd/elementwise.h b/infini_train/include/autograd/elementwise.h index c4333b64..16c33c4f 100644 --- a/infini_train/include/autograd/elementwise.h +++ b/infini_train/include/autograd/elementwise.h @@ -104,6 +104,8 @@ class Exp : public Function { explicit Exp() : Function(kType) {} 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; }; diff --git a/infini_train/include/core/runtime/dropout_kernels.h b/infini_train/include/core/runtime/dropout_kernels.h new file mode 100644 index 00000000..f02a3b8b --- /dev/null +++ b/infini_train/include/core/runtime/dropout_kernels.h @@ -0,0 +1,16 @@ +#pragma once + +#include + +#include "infini_train/include/generator.h" + +namespace infini_train { + +class Tensor; + +void dropout_forward_kernel(Tensor &output, Tensor &mask, const Tensor &input, double p, + const std::optional &generator); + +void dropout_backward_kernel(Tensor &grad_input, const Tensor &grad_output, const Tensor &mask, double p); + +} // namespace infini_train diff --git a/infini_train/include/core/runtime/dropout_stubs.h b/infini_train/include/core/runtime/dropout_stubs.h new file mode 100644 index 00000000..5faa91a0 --- /dev/null +++ b/infini_train/include/core/runtime/dropout_stubs.h @@ -0,0 +1,19 @@ +#pragma once + +#include + +#include "infini_train/include/core/runtime/dispatch_stub.h" +#include "infini_train/include/generator.h" + +namespace infini_train { + +class Tensor; + +DECLARE_DISPATCH(void (*)(Tensor &output, Tensor &mask, const Tensor &input, double p, + const std::optional &generator), + dropout_forward_stub); + +DECLARE_DISPATCH(void (*)(Tensor &grad_input, const Tensor &grad_output, const Tensor &mask, double p), + dropout_backward_stub); + +} // namespace infini_train diff --git a/infini_train/include/nn/functional.h b/infini_train/include/nn/functional.h index 124e914e..7ed6d047 100644 --- a/infini_train/include/nn/functional.h +++ b/infini_train/include/nn/functional.h @@ -62,6 +62,9 @@ std::shared_ptr Randn(const std::vector &size, DataType dtype = Device device = Device(), std::optional generator = std::nullopt, bool requires_grad = false); +std::shared_ptr Dropout(const std::shared_ptr &input, double p = 0.5, bool training = true, + std::optional generator = std::nullopt); + // Returns a new tensor with the reciprocal of the elements of input. // // Args: diff --git a/infini_train/src/autograd/dropout.cc b/infini_train/src/autograd/dropout.cc new file mode 100644 index 00000000..55b6a40a --- /dev/null +++ b/infini_train/src/autograd/dropout.cc @@ -0,0 +1,37 @@ +#include "infini_train/include/autograd/dropout.h" + +#include "infini_train/include/core/runtime/dropout_kernels.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 output = std::make_shared(input->Dims(), input->Dtype(), input->GetDevice()); + mask_ = std::make_shared(input->Dims(), DataType::kUINT8, input->GetDevice()); + dropout_forward_kernel(*output, *mask_, *input, p_, generator_); + return {output}; +} + +void Dropout::SetupContext(const std::vector> &, + const std::vector> &) { + if (!needs_input_grad_.empty() && needs_input_grad_[0]) { + saved_tensors_ = {mask_}; + } + mask_.reset(); +} + +std::vector> Dropout::Backward(const std::vector> &grad_outputs) { + CHECK_EQ(grad_outputs.size(), 1); + CHECK_EQ(saved_tensors_.size(), 1); + const auto &grad_output = grad_outputs[0]; + const auto &mask = saved_tensors_[0]; + + auto grad_input = std::make_shared(grad_output->Dims(), grad_output->Dtype(), grad_output->GetDevice()); + dropout_backward_kernel(*grad_input, *grad_output, *mask, p_); + return {grad_input}; +} + +} // namespace infini_train::autograd diff --git a/infini_train/src/autograd/elementwise.cc b/infini_train/src/autograd/elementwise.cc index 655cd309..d37d8ecb 100644 --- a/infini_train/src/autograd/elementwise.cc +++ b/infini_train/src/autograd/elementwise.cc @@ -176,12 +176,20 @@ std::vector> Exp::Forward(const std::vector>({device, "ExpForward"}, input)}; } +void Exp::SetupContext(const std::vector> &, + const std::vector> &output_tensors) { + const auto &output = output_tensors[0]; + saved_tensors_ = {output}; +} + std::vector> Exp::Backward(const std::vector> &grad_outputs) { + CHECK_EQ(saved_tensors_.size(), 1); + const auto &output = saved_tensors_[0]; CHECK_EQ(grad_outputs.size(), 1); const auto &grad_output = grad_outputs[0]; - auto device = grad_output->GetDevice().type(); - return {Dispatcher::Instance().Call>({device, "ExpBackward"}, grad_output)}; + auto device = output->GetDevice().type(); + return {Dispatcher::Instance().Call>({device, "ExpBackward"}, grad_output, output)}; } std::vector> Log::Forward(const std::vector> &input_tensors) { diff --git a/infini_train/src/core/dropout_kernels.cc b/infini_train/src/core/dropout_kernels.cc new file mode 100644 index 00000000..dc3251e1 --- /dev/null +++ b/infini_train/src/core/dropout_kernels.cc @@ -0,0 +1,43 @@ +#include "infini_train/include/core/runtime/dropout_kernels.h" + +#include "infini_train/include/core/runtime/dropout_stubs.h" +#include "infini_train/include/tensor.h" + +namespace infini_train { + +DEFINE_DISPATCH(dropout_forward_stub); +DEFINE_DISPATCH(dropout_backward_stub); + +namespace { + +void check_dropout_probability(double p) { + 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; +} + +void check_dropout_tensors(const Tensor &output, const Tensor &mask, const Tensor &input) { + CHECK(IsFloatingPointDType(input.Dtype())) << "Dropout supports floating-point tensors only"; + CHECK_EQ(static_cast(output.Dtype()), static_cast(input.Dtype())); + CHECK(output.GetDevice() == input.GetDevice()); + CHECK(output.Dims() == input.Dims()); + CHECK_EQ(static_cast(mask.Dtype()), static_cast(DataType::kUINT8)); + CHECK(mask.GetDevice() == input.GetDevice()); + CHECK(mask.Dims() == input.Dims()); +} + +} // namespace + +void dropout_forward_kernel(Tensor &output, Tensor &mask, const Tensor &input, double p, + const std::optional &generator) { + check_dropout_probability(p); + check_dropout_tensors(output, mask, input); + dropout_forward_stub(input.GetDevice().type(), output, mask, input, p, generator); +} + +void dropout_backward_kernel(Tensor &grad_input, const Tensor &grad_output, const Tensor &mask, double p) { + check_dropout_probability(p); + check_dropout_tensors(grad_input, mask, grad_output); + dropout_backward_stub(grad_output.GetDevice().type(), grad_input, grad_output, mask, p); +} + +} // namespace infini_train diff --git a/infini_train/src/kernels/cpu/dropout.cc b/infini_train/src/kernels/cpu/dropout.cc new file mode 100644 index 00000000..ac909ead --- /dev/null +++ b/infini_train/src/kernels/cpu/dropout.cc @@ -0,0 +1,100 @@ +#include +#include + +#include "infini_train/include/core/runtime/dropout_stubs.h" +#include "infini_train/include/core/runtime/distributions_helper.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 dropout_forward_cpu_impl(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)); + 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 dropout_backward_cpu_impl(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); + } +} + +void dropout_forward_cpu(Tensor &output, Tensor &mask, const Tensor &input, double p, + const std::optional &generator) { + CHECK(input.GetDevice().IsCPU()); + + core::cpu::DispatchCpuFunc( + input.Dtype(), + [&]() { + using random_t = std::conditional_t, double, float>; + if (p == 0.0 || p == 1.0) { + dropout_forward_cpu_impl(output, mask, input, p, nullptr); + return; + } + auto *cpu_generator = get_generator_or_default( + generator, core::cpu::getDefaultCPUGenerator()); + std::lock_guard lock(cpu_generator->mutex_); + dropout_forward_cpu_impl(output, mask, input, p, cpu_generator); + }, + "CPU dropout forward"); +} + +void dropout_backward_cpu(Tensor &grad_input, const Tensor &grad_output, const Tensor &mask, double p) { + CHECK(grad_output.GetDevice().IsCPU()); + core::cpu::DispatchCpuFunc( + grad_output.Dtype(), + [&]() { + using random_t = std::conditional_t, double, float>; + dropout_backward_cpu_impl(grad_input, grad_output, mask, p); + }, + "CPU dropout backward"); +} + +} // namespace + +using infini_train::dropout_backward_stub; +using infini_train::dropout_forward_stub; + +REGISTER_DISPATCH(dropout_forward_stub, Device::DeviceType::kCPU, &dropout_forward_cpu); +REGISTER_DISPATCH(dropout_backward_stub, Device::DeviceType::kCPU, &dropout_backward_cpu); + +} // namespace infini_train::kernels::cpu diff --git a/infini_train/src/kernels/cuda/dropout.cu b/infini_train/src/kernels/cuda/dropout.cu new file mode 100644 index 00000000..f1315da3 --- /dev/null +++ b/infini_train/src/kernels/cuda/dropout.cu @@ -0,0 +1,144 @@ +#include +#include +#include +#include + +#include + +#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/core/runtime/dropout_stubs.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)); +} + +void dropout_forward_cuda(Tensor &output, Tensor &mask, const Tensor &input, double p, + const std::optional &generator) { + const Device device = input.GetDevice(); + CHECK(device.IsCUDA()); + const int64_t n = input.NumElements(); + if (n == 0) { + return; + } + core::DeviceGuard guard(device); + + uint64_t seed = 0; + uint64_t subsequence = 0; + if (p > 0.0 && p < 1.0) { + auto *cuda_generator = get_generator_or_default( + generator, 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()); +} + +void dropout_backward_cuda(Tensor &grad_input, const Tensor &grad_output, const Tensor &mask, double p) { + const Device device = grad_output.GetDevice(); + CHECK(device.IsCUDA()); + const int64_t n = grad_output.NumElements(); + if (n == 0) { + return; + } + core::DeviceGuard guard(device); + + 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()); +} + +} // namespace + +using infini_train::dropout_backward_stub; +using infini_train::dropout_forward_stub; + +REGISTER_DISPATCH(dropout_forward_stub, Device::DeviceType::kCUDA, &dropout_forward_cuda); +REGISTER_DISPATCH(dropout_backward_stub, Device::DeviceType::kCUDA, &dropout_backward_cuda); + +} // namespace infini_train::kernels::cuda diff --git a/infini_train/src/nn/functional.cc b/infini_train/src/nn/functional.cc index a80c995c..50622d3f 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/misc.h" #include "infini_train/include/autograd/reduction.h" @@ -39,6 +40,17 @@ std::shared_ptr Randn(const std::vector &size, DataType dtype, 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/tests/CMakeLists.txt b/tests/CMakeLists.txt index b2ed5009..07061a00 100644 --- a/tests/CMakeLists.txt +++ b/tests/CMakeLists.txt @@ -27,3 +27,6 @@ add_subdirectory(dtype) # Transformer architecture tests add_subdirectory(transformer) + +# Generator tests +add_subdirectory(generator) diff --git a/tests/autograd/test_autograd_elementwise_backward.cc b/tests/autograd/test_autograd_elementwise_backward.cc index 296094e9..f79751ba 100644 --- a/tests/autograd/test_autograd_elementwise_backward.cc +++ b/tests/autograd/test_autograd_elementwise_backward.cc @@ -4,6 +4,7 @@ #include "gtest/gtest.h" #include "infini_train/include/autograd/elementwise.h" +#include "infini_train/include/core/runtime/device_guard.h" #include "infini_train/include/nn/parallel/global.h" #include "infini_train/include/tensor.h" @@ -117,7 +118,16 @@ TEST_P(AutogradElementwiseBackwardTest, ExpBackward) { auto grad = std::make_shared(std::vector{2, 3}, DataType::kFLOAT32, GetDevice(), true); grad->Fill(1.0f); auto grad_inputs = exp_fn->Backward({grad}); - EXPECT_EQ(grad_inputs.size(), 1); + ASSERT_EQ(grad_inputs.size(), 1); + + auto device = grad_inputs[0]->GetDevice(); + core::GetDeviceGuardImpl(device.type())->SynchronizeDevice(device); + Tensor grad_input_cpu = grad_inputs[0]->To(Device(Device::DeviceType::kCPU, 0)); + core::GetDeviceGuardImpl(device.type())->SynchronizeDevice(device); + const auto *grad_input = static_cast(grad_input_cpu.DataPtr()); + for (size_t i = 0; i < grad_input_cpu.NumElements(); ++i) { + EXPECT_FLOAT_EQ(grad_input[i], std::exp(1.0f)); + } } TEST_P(AutogradElementwiseBackwardTest, LogBackward) { diff --git a/tests/generator/CMakeLists.txt b/tests/generator/CMakeLists.txt new file mode 100644 index 00000000..05063fa6 --- /dev/null +++ b/tests/generator/CMakeLists.txt @@ -0,0 +1,19 @@ +# Shared Generator behavior tests. The existing suite helper builds separate +# CPU and CUDA executables and filters the parameterized instances by label. +infini_train_add_test_suite(test_generator + SOURCES test_generator.cc +) + +# Device-independent handle and validation tests are not parameterized, so +# keep them in a separate executable that is not subject to CPU/* filters. +infini_train_add_test(test_generator_interface + SOURCES test_generator_interface.cc + LABELS cpu +) + +if(USE_CUDA) + infini_train_add_test(test_generator_cuda_only + SOURCES cuda_only/test_cuda_generator.cc + LABELS cuda + ) +endif() diff --git a/tests/generator/cuda_only/test_cuda_generator.cc b/tests/generator/cuda_only/test_cuda_generator.cc new file mode 100644 index 00000000..0cdd468b --- /dev/null +++ b/tests/generator/cuda_only/test_cuda_generator.cc @@ -0,0 +1,72 @@ +#include +#include +#include + +#include "gtest/gtest.h" + +#include "infini_train/include/generator.h" +#include "infini_train/include/nn/init.h" +#include "infini_train/include/tensor.h" +#include "infini_train/src/core/runtime/cpu/cpu_generator_impl.h" +#include "infini_train/src/core/runtime/cuda/cuda_generator_impl.h" +#include "tests/common/test_utils.h" + +namespace infini_train::test { +namespace { + +constexpr uint64_t kSeed = 0x12345678ULL; + +std::vector StateBytes(const Generator &generator) { + auto state = generator.get_state(); + const auto *data = static_cast(state->DataPtr()); + return std::vector(data, data + state->SizeInBytes()); +} + +TEST(CudaGeneratorTest, StateIsAnOpaqueUint8CpuTensor) { + Generator generator = core::cuda::createCUDAGenerator(0, kSeed); + auto state = generator.get_state(); + + ASSERT_TRUE(state); + EXPECT_TRUE(state->GetDevice().IsCPU()); + EXPECT_EQ(state->Dtype(), DataType::kUINT8); + EXPECT_EQ(state->SizeInBytes(), sizeof(uint64_t) * 2); +} + +TEST(CudaGeneratorTest, RejectsStateWithInvalidSize) { + Generator generator = core::cuda::createCUDAGenerator(0, kSeed); + Tensor invalid_state(std::vector{1}, DataType::kUINT8, + Device(Device::DeviceType::kCPU, 0)); + + EXPECT_DEATH(generator.set_state(invalid_state), ""); +} + +TEST(CudaGeneratorTest, CpuAndCudaStatesAreNotInterchangeable) { + Generator cpu_generator = core::cpu::createCPUGenerator(kSeed); + Generator cuda_generator = core::cuda::createCUDAGenerator(0, kSeed); + auto cpu_state = cpu_generator.get_state(); + auto cuda_state = cuda_generator.get_state(); + + EXPECT_DEATH(cuda_generator.set_state(*cpu_state), ""); + EXPECT_DEATH(cpu_generator.set_state(*cuda_state), ""); +} + +TEST(CudaGeneratorTest, DefaultGeneratorsArePerDevice) { + REQUIRE_MIN_DEVICES(2); + const Generator &device_zero = core::cuda::getDefaultCUDAGenerator(0); + const Generator &device_one = core::cuda::getDefaultCUDAGenerator(1); + + EXPECT_NE(device_zero, device_one); + EXPECT_EQ(device_zero.device().index(), 0); + EXPECT_EQ(device_one.device().index(), 1); + + const std::vector device_one_before = StateBytes(device_one); + auto tensor = std::make_shared( + std::vector{1024}, DataType::kFLOAT32, + Device(Device::DeviceType::kCUDA, 0)); + nn::init::Uniform(tensor, 0.0f, 1.0f, device_zero); + + EXPECT_EQ(device_one_before, StateBytes(device_one)); +} + +} // namespace +} // namespace infini_train::test diff --git a/tests/generator/test_generator.cc b/tests/generator/test_generator.cc new file mode 100644 index 00000000..cf0b2867 --- /dev/null +++ b/tests/generator/test_generator.cc @@ -0,0 +1,305 @@ +#include +#include +#include +#include +#include +#include + +#include "gtest/gtest.h" + +#include "infini_train/include/core/runtime/device_guard.h" +#include "infini_train/include/generator.h" +#include "infini_train/include/nn/functional.h" +#include "infini_train/include/nn/init.h" +#include "infini_train/include/tensor.h" +#include "infini_train/src/core/runtime/cpu/cpu_generator_impl.h" +#if defined(USE_CUDA) +#include "infini_train/src/core/runtime/cuda/cuda_generator_impl.h" +#endif +#include "tests/common/test_utils.h" + +namespace infini_train::test { +namespace { + +constexpr uint64_t kSeed = 0x12345678ULL; +constexpr int64_t kElements = 4096; + +Generator CreateGenerator(Device device, uint64_t seed) { + if (device.IsCPU()) { + return core::cpu::createCPUGenerator(seed); + } +#if defined(USE_CUDA) + return core::cuda::createCUDAGenerator(device.index(), seed); +#else + (void)seed; + throw std::runtime_error("CUDA support is disabled"); +#endif +} + +const Generator &GetDefaultGenerator(Device device) { + if (device.IsCPU()) { + return core::cpu::getDefaultCPUGenerator(); + } +#if defined(USE_CUDA) + return core::cuda::getDefaultCUDAGenerator(device.index()); +#else + throw std::runtime_error("CUDA support is disabled"); +#endif +} + +void Synchronize(Device device) { + core::GetDeviceGuardImpl(device.type())->SynchronizeDevice(device); +} + +std::vector ToHostVector(const std::shared_ptr &tensor) { + const Device device = tensor->GetDevice(); + Synchronize(device); + Tensor cpu = tensor->To(Device(Device::DeviceType::kCPU, 0)); + Synchronize(device); + const auto *data = static_cast(cpu.DataPtr()); + return std::vector(data, data + cpu.NumElements()); +} + +std::vector ToHostBytes(const std::shared_ptr &tensor) { + const Device device = tensor->GetDevice(); + Synchronize(device); + Tensor cpu = tensor->To(Device(Device::DeviceType::kCPU, 0)); + Synchronize(device); + const auto *data = static_cast(cpu.DataPtr()); + return std::vector(data, data + cpu.SizeInBytes()); +} + +std::vector StateBytes(const Generator &generator) { + auto state = generator.get_state(); + const auto *data = static_cast(state->DataPtr()); + return std::vector(data, data + state->SizeInBytes()); +} + +std::shared_ptr MakeTensor(Device device) { + return std::make_shared(std::vector{kElements}, DataType::kFLOAT32, device); +} + +class GeneratorTest : public InfiniTrainTest {}; + +TEST_P(GeneratorTest, ReportsDeviceAndInitialSeed) { + const Device device = GetDevice(); + Generator generator = CreateGenerator(device, kSeed); + + EXPECT_TRUE(generator.defined()); + EXPECT_EQ(generator.device(), device); + EXPECT_EQ(generator.current_seed(), kSeed); +} + +TEST_P(GeneratorTest, ShallowCopySharesStateAndCloneDoesNot) { + Generator generator = CreateGenerator(GetDevice(), kSeed); + Generator alias = generator; + Generator clone = generator.clone(); + + EXPECT_EQ(alias, generator); + EXPECT_NE(clone, generator); + EXPECT_EQ(StateBytes(alias), StateBytes(generator)); + EXPECT_EQ(StateBytes(clone), StateBytes(generator)); + + auto tensor = MakeTensor(GetDevice()); + nn::init::Uniform(tensor, -1.0f, 1.0f, alias); + + EXPECT_EQ(StateBytes(alias), StateBytes(generator)); + EXPECT_NE(StateBytes(clone), StateBytes(generator)); +} + +TEST_P(GeneratorTest, UniformIsReproducibleForSameSeed) { + Generator first_generator = CreateGenerator(GetDevice(), kSeed); + Generator second_generator = CreateGenerator(GetDevice(), kSeed); + auto first = MakeTensor(GetDevice()); + auto second = MakeTensor(GetDevice()); + + nn::init::Uniform(first, -3.0f, 7.0f, first_generator); + nn::init::Uniform(second, -3.0f, 7.0f, second_generator); + + EXPECT_EQ(ToHostVector(first), ToHostVector(second)); +} + +TEST_P(GeneratorTest, NormalIsReproducibleForSameSeed) { + Generator first_generator = CreateGenerator(GetDevice(), kSeed); + Generator second_generator = CreateGenerator(GetDevice(), kSeed); + auto first = MakeTensor(GetDevice()); + auto second = MakeTensor(GetDevice()); + + nn::init::Normal(first, 2.0f, 0.5f, first_generator); + nn::init::Normal(second, 2.0f, 0.5f, second_generator); + + EXPECT_EQ(ToHostVector(first), ToHostVector(second)); +} + +TEST_P(GeneratorTest, UniformSupportsEveryFloatingPointDtype) { + for (const DataType dtype : {DataType::kFLOAT16, DataType::kBFLOAT16, + DataType::kFLOAT32, DataType::kFLOAT64}) { + Generator first_generator = CreateGenerator(GetDevice(), kSeed); + Generator second_generator = CreateGenerator(GetDevice(), kSeed); + auto first = std::make_shared( + std::vector{kElements}, dtype, GetDevice()); + auto second = std::make_shared( + std::vector{kElements}, dtype, GetDevice()); + + nn::init::Uniform(first, -3.0f, 7.0f, first_generator); + nn::init::Uniform(second, -3.0f, 7.0f, second_generator); + + EXPECT_EQ(ToHostBytes(first), ToHostBytes(second)) + << "dtype=" << static_cast(dtype); + } +} + +TEST_P(GeneratorTest, NormalSupportsEveryFloatingPointDtype) { + for (const DataType dtype : {DataType::kFLOAT16, DataType::kBFLOAT16, + DataType::kFLOAT32, DataType::kFLOAT64}) { + Generator first_generator = CreateGenerator(GetDevice(), kSeed); + Generator second_generator = CreateGenerator(GetDevice(), kSeed); + auto first = std::make_shared( + std::vector{kElements}, dtype, GetDevice()); + auto second = std::make_shared( + std::vector{kElements}, dtype, GetDevice()); + + nn::init::Normal(first, 2.0f, 0.5f, first_generator); + nn::init::Normal(second, 2.0f, 0.5f, second_generator); + + EXPECT_EQ(ToHostBytes(first), ToHostBytes(second)) + << "dtype=" << static_cast(dtype); + } +} + +TEST_P(GeneratorTest, RandAndRandnPreserveRequestedTensorProperties) { + Generator rand_generator = CreateGenerator(GetDevice(), kSeed); + Generator randn_generator = CreateGenerator(GetDevice(), kSeed); + + auto uniform = nn::function::Rand({17, 19}, DataType::kFLOAT64, GetDevice(), + rand_generator, true); + auto normal = nn::function::Randn({17, 19}, DataType::kFLOAT16, GetDevice(), + randn_generator, true); + + EXPECT_EQ(uniform->Dims(), (std::vector{17, 19})); + EXPECT_EQ(uniform->Dtype(), DataType::kFLOAT64); + EXPECT_EQ(uniform->GetDevice(), GetDevice()); + EXPECT_TRUE(uniform->requires_grad()); + EXPECT_EQ(normal->Dims(), (std::vector{17, 19})); + EXPECT_EQ(normal->Dtype(), DataType::kFLOAT16); + EXPECT_EQ(normal->GetDevice(), GetDevice()); + EXPECT_TRUE(normal->requires_grad()); +} + +TEST_P(GeneratorTest, ConsecutiveCallsAdvanceSequence) { + Generator generator = CreateGenerator(GetDevice(), kSeed); + auto first = MakeTensor(GetDevice()); + auto second = MakeTensor(GetDevice()); + + nn::init::Uniform(first, 0.0f, 1.0f, generator); + nn::init::Uniform(second, 0.0f, 1.0f, generator); + + EXPECT_NE(ToHostVector(first), ToHostVector(second)); +} + +TEST_P(GeneratorTest, DifferentSeedsProduceDifferentSequences) { + Generator first_generator = CreateGenerator(GetDevice(), kSeed); + Generator second_generator = CreateGenerator(GetDevice(), kSeed + 1); + auto first = MakeTensor(GetDevice()); + auto second = MakeTensor(GetDevice()); + + nn::init::Uniform(first, 0.0f, 1.0f, first_generator); + nn::init::Uniform(second, 0.0f, 1.0f, second_generator); + + EXPECT_NE(ToHostVector(first), ToHostVector(second)); +} + +TEST_P(GeneratorTest, StateRestoreReplaysUniformSequence) { + Generator generator = CreateGenerator(GetDevice(), kSeed); + auto prefix = MakeTensor(GetDevice()); + auto expected = MakeTensor(GetDevice()); + auto actual = MakeTensor(GetDevice()); + + nn::init::Uniform(prefix, 0.0f, 1.0f, generator); + auto state = generator.get_state(); + nn::init::Uniform(expected, 0.0f, 1.0f, generator); + generator.set_state(*state); + nn::init::Uniform(actual, 0.0f, 1.0f, generator); + + EXPECT_EQ(ToHostVector(expected), ToHostVector(actual)); +} + +TEST_P(GeneratorTest, StateRestoreReplaysNormalSequence) { + Generator generator = CreateGenerator(GetDevice(), kSeed); + auto prefix = MakeTensor(GetDevice()); + auto expected = MakeTensor(GetDevice()); + auto actual = MakeTensor(GetDevice()); + + nn::init::Normal(prefix, 0.0f, 1.0f, generator); + auto state = generator.get_state(); + nn::init::Normal(expected, 0.0f, 1.0f, generator); + generator.set_state(*state); + nn::init::Normal(actual, 0.0f, 1.0f, generator); + + EXPECT_EQ(ToHostVector(expected), ToHostVector(actual)); +} + +TEST_P(GeneratorTest, ExplicitGeneratorDoesNotAdvanceDefaultGenerator) { + const Device device = GetDevice(); + manual_seed(kSeed); + const Generator &default_generator = GetDefaultGenerator(device); + const std::vector before = StateBytes(default_generator); + Generator explicit_generator = CreateGenerator(device, kSeed + 1); + + auto tensor = MakeTensor(device); + nn::init::Uniform(tensor, 0.0f, 1.0f, explicit_generator); + + EXPECT_EQ(before, StateBytes(default_generator)); +} + +TEST_P(GeneratorTest, MissingGeneratorUsesAndAdvancesDefaultGenerator) { + const Device device = GetDevice(); + manual_seed(kSeed); + const Generator &default_generator = GetDefaultGenerator(device); + const std::vector before = StateBytes(default_generator); + + auto tensor = MakeTensor(device); + nn::init::Uniform(tensor); + + EXPECT_NE(before, StateBytes(default_generator)); +} + +TEST_P(GeneratorTest, RepeatedDefaultGeneratorLookupSharesState) { + const Device device = GetDevice(); + const Generator &first = GetDefaultGenerator(device); + const Generator &second = GetDefaultGenerator(device); + + EXPECT_EQ(first, second); +} + +TEST_P(GeneratorTest, GlobalManualSeedReproducesDefaultGeneratorResults) { + const Device device = GetDevice(); + auto first = MakeTensor(device); + auto second = MakeTensor(device); + + manual_seed(kSeed); + nn::init::Uniform(first); + manual_seed(kSeed); + nn::init::Uniform(second); + + EXPECT_EQ(ToHostVector(first), ToHostVector(second)); +} + +TEST_P(GeneratorTest, KaimingUniformUsesExplicitGenerator) { + Generator first_generator = CreateGenerator(GetDevice(), kSeed); + Generator second_generator = CreateGenerator(GetDevice(), kSeed); + auto first = std::make_shared(std::vector{64, 32}, DataType::kFLOAT32, GetDevice()); + auto second = std::make_shared(std::vector{64, 32}, DataType::kFLOAT32, GetDevice()); + + nn::init::KaimingUniform(first, 0.0f, nn::init::KaimingMode::kFanIn, + nn::init::NonLinearityType::kReLU, first_generator); + nn::init::KaimingUniform(second, 0.0f, nn::init::KaimingMode::kFanIn, + nn::init::NonLinearityType::kReLU, second_generator); + + EXPECT_EQ(ToHostVector(first), ToHostVector(second)); +} + +INFINI_TRAIN_REGISTER_TEST(GeneratorTest); + +} // namespace +} // namespace infini_train::test diff --git a/tests/generator/test_generator_interface.cc b/tests/generator/test_generator_interface.cc new file mode 100644 index 00000000..7a698f9c --- /dev/null +++ b/tests/generator/test_generator_interface.cc @@ -0,0 +1,64 @@ +#include +#include +#include +#include + +#include "gtest/gtest.h" + +#include "infini_train/include/generator.h" +#include "infini_train/include/tensor.h" +#include "infini_train/src/core/runtime/cpu/cpu_generator_impl.h" + +namespace infini_train::test { +namespace { + +constexpr uint64_t kSeed = 0x12345678ULL; + +TEST(GeneratorInterfaceTest, DefaultConstructedGeneratorIsUndefined) { + Generator generator; + EXPECT_FALSE(generator.defined()); +} + +TEST(GeneratorInterfaceTest, CheckGeneratorRejectsUndefinedGenerator) { + Generator generator; + EXPECT_THROW(check_generator(generator), std::invalid_argument); +} + +TEST(GeneratorInterfaceTest, CpuStateIsAnOpaqueUint8CpuTensor) { + Generator generator = core::cpu::createCPUGenerator(kSeed); + auto state = generator.get_state(); + + ASSERT_TRUE(state); + EXPECT_TRUE(state->GetDevice().IsCPU()); + EXPECT_EQ(state->Dtype(), DataType::kUINT8); + EXPECT_GT(state->SizeInBytes(), 0u); +} + +TEST(GeneratorInterfaceTest, CpuStateRoundTripRestoresSeed) { + Generator generator = core::cpu::createCPUGenerator(kSeed); + auto state = generator.get_state(); + generator.set_current_seed(kSeed + 1); + + generator.set_state(*state); + + EXPECT_EQ(generator.current_seed(), kSeed); +} + +TEST(GeneratorInterfaceTest, RejectsStateWithWrongDtype) { + Generator generator = core::cpu::createCPUGenerator(kSeed); + Tensor invalid_state(std::vector{64}, DataType::kFLOAT32, + Device(Device::DeviceType::kCPU, 0)); + + EXPECT_DEATH(generator.set_state(invalid_state), "RNG state must be a UINT8 tensor"); +} + +TEST(GeneratorInterfaceTest, RejectsTruncatedCpuState) { + Generator generator = core::cpu::createCPUGenerator(kSeed); + Tensor invalid_state(std::vector{1}, DataType::kUINT8, + Device(Device::DeviceType::kCPU, 0)); + + EXPECT_DEATH(generator.set_state(invalid_state), "CPU generator state is too small"); +} + +} // namespace +} // namespace infini_train::test From 405b333f6e1e98b08362247cbd235529dc22e2f1 Mon Sep 17 00:00:00 2001 From: chaos Date: Sun, 12 Jul 2026 22:06:11 +0800 Subject: [PATCH 11/11] test: cover generator reproducibility semantics --- ...75\350\261\241\351\200\211\351\242\230.md" | 250 ------------------ .../cuda_only/test_cuda_generator.cc | 30 +-- tests/generator/test_generator.cc | 133 ++++------ tests/generator/test_generator_interface.cc | 27 -- 4 files changed, 62 insertions(+), 378 deletions(-) delete mode 100644 "docs/\343\200\2202026 \346\230\245\345\255\243\344\272\272\345\267\245\346\231\272\350\203\275\345\244\247\350\265\233\343\200\221Generator \346\212\275\350\261\241\351\200\211\351\242\230.md" diff --git "a/docs/\343\200\2202026 \346\230\245\345\255\243\344\272\272\345\267\245\346\231\272\350\203\275\345\244\247\350\265\233\343\200\221Generator \346\212\275\350\261\241\351\200\211\351\242\230.md" "b/docs/\343\200\2202026 \346\230\245\345\255\243\344\272\272\345\267\245\346\231\272\350\203\275\345\244\247\350\265\233\343\200\221Generator \346\212\275\350\261\241\351\200\211\351\242\230.md" deleted file mode 100644 index 67b1ddcc..00000000 --- "a/docs/\343\200\2202026 \346\230\245\345\255\243\344\272\272\345\267\245\346\231\272\350\203\275\345\244\247\350\265\233\343\200\221Generator \346\212\275\350\261\241\351\200\211\351\242\230.md" +++ /dev/null @@ -1,250 +0,0 @@ -# 中期报告要求 - -- 比赛期：2026/05/18 - 2026/07/12 （56 天） - - 比赛期第一天，即 2026/05/18 的 12:00 公布赛题 - - 每个参赛小组需在特定的截止日期前提交中期报告。中期报告需包含： - - 小组名称与所有成员 - - 所选的赛题 - - 截至撰写时各个赛题的进度与完成情况 - - 中期报告没有字数要求，简明扼要、表述清晰即可 - - 提交要求与方式： - - 报告文件命名要求：<小组名称>_<赛道名称>_中期报告.pdf - - 提交方式：InfiniTensor 网站上提交，提交流程后续补充 - -# 【2026 春季人工智能大赛】Generator 抽象选题 - -## 一、题目背景 - -在深度学习训练与推理框架中，随机数生成器（Generator）是支撑参数初始化、随机采样、Dropout、噪声注入等功能的重要基础设施。当前框架中相关能力尚不完善，缺少统一的 Generator 抽象、后端实现以及全局随机种子管理机制，导致随机相关算子的行为与主流框架存在差距，也难以满足后续功能扩展需求。 - -当前已完成 Generator 基础设施、初始化算子、随机张量创建算子与普通 Dropout 的首轮接入，进度如下： - -- [x] 已建立统一的 Generator 抽象，支持随机数状态、种子设置、状态获取与恢复； -- [x] 已实现 CPU 与 CUDA 后端 Generator，并保持统一的调用接口； -- [x] 已建立按设备维护的默认 Generator；未显式传入 generator 时，初始化算子使用目标 Tensor 所在设备的默认随机源； -- [x] 已提供全局 `manual_seed(uint64_t)` 入口，重置 CPU 及所有 CUDA 默认 Generator； -- [x] 已接入 `Rand`、`Randn`，支持显式/默认 Generator、CPU/CUDA 及 `FLOAT16`、`BFLOAT16`、`FLOAT32`、`FLOAT64`； -- [x] 已接入普通 `Dropout`，训练态保存随机 mask 供反向传播复用，支持显式/默认 Generator、CPU/CUDA 及四种浮点 dtype； - -为补齐这一基础能力，本题目要求参赛者参考主流深度学习框架的设计思路，完成一套可扩展的随机数生成器基础设施，实现 CPU 与 CUDA 后端的 Generator，并支持基于 Generator 的随机数生成与种子控制能力。 - -## 二、题目目标 - -参赛者需要围绕框架内的随机数生成体系，完成以下几个方面的工作： - -- [x] 1. 设计并实现统一的 Generator 抽象接口，支持随机数状态管理、种子设置、状态获取与恢复等基础能力。 -- [x] 2. 分别实现 CPU 与 CUDA 后端对应的 GeneratorImpl，使不同设备上的随机数生成具备统一的调用方式。 -- [x] 3. 建立全局随机种子控制机制，支持通过统一入口固定随机种子；已接入的初始化算子、随机张量创建算子和 Dropout 可稳定复现。 -- [x] 4. 改造初始化类随机算子、随机张量创建算子和 Dropout，使其在未显式传入 generator 参数时，自动使用目标设备的默认全局 Generator。 -- [x] 5. 已完成 CPU/CUDA 初始化、`Rand/Randn`、Dropout、多 dtype 分布和参数校验的基本行为、可复现性和接口语义冒烟验证；仓库内正式测试仍待补充。 - -## 三、任务拆解 - -### 1. Generator 抽象与状态管理设计 - -在框架中设计统一的 Generator 抽象层，用于屏蔽不同设备后端的随机数实现差异。该抽象应支持： - -- [x] 设置随机种子，如 `ManualSeed()` / `SetCurrentSeed()`； -- [x] 获取当前种子或初始种子，如 `Seed()` / `InitialSeed()`； -- [x] 获取当前随机数状态，如 `GetState()`； -- [x] 恢复随机数状态，如 `SetState(state)`； -- [x] 查询所属设备类型； -- [x] 提供必要的线程安全或状态访问保护机制。 - -实现时需注意区分 seed 与 state：seed 只用于初始化随机序列，state 表示随机序列当前推进到的位置。随机算子每次使用 Generator 后，应推进其内部状态，避免重复生成相同随机序列。 - -推荐设计思路： - -- [x] 提供用户侧可持有的 `Generator` 句柄类； -- [x] 底层以 `GeneratorImpl` 作为多态实现基类； -- [x] CPU / CUDA 分别派生对应实现； -- [x] 公共接口不暴露 `std::mt19937`、curand、Philox 等后端细节； -- [x] 为后续其他平台 Generator 接入预留扩展空间。 - -### 2. CPU / CUDA 后端 Generator 实现 - -基于现有设备后端，分别实现： - -- [x] `CPUGeneratorImpl` -- [x] `CUDAGeneratorImpl` - -两者应保持统一接口语义，但内部状态组织方式可以不同，且**不要求 CPU 与 CUDA 在相同 seed 下生成逐元素一致的随机结果**。 - -实现时需重点关注： - -- [x] CPU 与 CUDA 随机数状态如何组织与保存；CPU 保存 mt19937 状态与正态缓存，CUDA 保存 Philox seed 与 subsequence； -- [x] CUDA 侧按设备维度维护独立默认 Generator； -- [x] `GetState()` / `SetState()` 的状态序列化、反序列化与格式校验；CPU/CUDA 均校验 CPU `UINT8` state，CPU 额外校验 footer 与 mt19937 反序列化； -- [x] 多次随机调用后，随机序列连续推进；CPU 推进 mt19937，CUDA 为每次 kernel 保留 Philox subsequence； -- [x] 同一 seed、同一设备、同一调用顺序下结果可复现； -- [x] 初始化算子、`Rand/Randn` 和 Dropout 已支持显式 Generator 和默认 Generator 两条调用路径。 - -CUDA 后端建议在状态中考虑 seed、offset、counter 等信息，避免不同 kernel 调用之间随机序列重叠。此部分需要保证功能正确、语义清晰、状态可管理、接口一致，为后续随机算子和分布式扩展提供基础。 - -### 3. 默认 Generator 与统一随机种子入口 - -建立框架级默认 Generator 管理机制，支持两种使用方式： - -- 用户显式传入某个 Generator； -- 用户不传入 Generator，系统自动使用当前设备上的默认 Generator。 - -该机制应包括： - -- [x] CPU 默认全局 Generator； -- [x] 各 CUDA 设备对应的默认 Generator； -- [x] CPU/CUDA 均提供获取默认 Generator 的后端入口； -- [x] 统一的全局随机种子设置入口：`manual_seed(uint64_t)`； -- [x] 设置全局 seed 时同步重置 CPU 与所有 CUDA 默认 Generator 的状态。 - -要求： - -- [x] 多次获取同一设备默认 Generator 时，应返回同一随机状态来源；（static 局部变量，单例模式） -- [x] 不同 CUDA 设备的默认 Generator 相互独立； -- [x] 用户显式传入 Generator 时，初始化算子使用用户指定 Generator； -- [x] 未传入 Generator 时，初始化算子使用目标设备默认 Generator； -- [x] 设置统一 seed 后，已接入的参数初始化、`Rand/Randn` 随机采样和 Dropout mask 可稳定复现； -- [ ] CUDA Graph RNG 状态管理仍待完成。 - -这一部分是本题的核心验收点。随机数系统最终服务于训练可复现性，因此必须通过统一入口稳定控制随机行为。 - -### 4. 随机相关算子接入改造 - -改造框架中已有的随机相关使用，使其接入 Generator 机制。建议至少覆盖两类场景： - -- [x] 初始化类随机算子：`Uniform`、`Normal`、`KaimingUniform` 已接入 Generator，并按 Tensor device 分发 CPU/CUDA kernel； -- [x] 通用随机张量创建算子：`Rand`、`Randn` 已接入 Generator，支持 CPU/CUDA 与四种浮点 dtype； -- [x] 训练过程随机算子：普通 Dropout 已接入 Generator，前向保存 mask，反向复用同一 mask，不重新消耗随机状态； - -改造后的算子应满足： - -- [x] 已接入初始化算子、`Rand/Randn` 和 Dropout 支持显式传入 Generator； -- [x] 已接入初始化算子、`Rand/Randn` 和 Dropout 未传入 Generator 时，自动使用目标设备默认 Generator； -- [x] 已接入初始化算子、`Rand/Randn` 和 Dropout 的随机结果受统一 `manual_seed` 入口控制； -- [x] 已接入初始化算子、`Rand/Randn` 和 Dropout 多次调用会推进 Generator 状态； -- [x] 已接入初始化算子、`Rand/Randn` 和 Dropout 在同一 seed、同一调用顺序下可复现。 - -不要求一次性改造所有随机算子，但应至少完成一个初始化类算子和一个训练期随机算子，以证明 Generator 机制能够贯通框架层与算子层。 - -当前已完成初始化类算子、通用随机张量创建算子和普通 Dropout 接入；后续可扩展 AlphaDropout、FeatureDropout、CUDA Graph RNG 等能力。 - -### 5. 测试与对齐验证 - -为验证实现结果，需要补充系统化测试。测试重点应放在"语义是否正确"和"是否可复现"，而不是与 PyTorch 逐元素数值一致。 - -当前已在 `/tmp` 完成 CPU/CUDA 构建，以及 Generator 状态恢复、初始化、`Rand/Randn`、Dropout 前后向、四种浮点 dtype 分布、参数校验和显式/默认 Generator 路径的冒烟验证；按项目约定尚未新增仓库内测试，因此以下正式测试清单仍未完成。 - -建议至少覆盖以下内容： - -#### （1）接口一致性测试 - -验证 CPU / CUDA Generator 是否支持统一接口，包括： - -- [ ] seed 设置与获取； -- [ ] state 获取与恢复； -- [ ] 设备类型查询； -- [ ] 默认 Generator 获取； -- [ ] 显式 Generator 与默认 Generator 两条调用路径。 - -#### （2）种子可复现测试 - -验证统一随机种子入口是否生效，包括： - -- [ ] 同一 seed 下，多次运行参数初始化结果一致； -- [ ] 同一 seed 下，多次运行 Dropout mask 一致； -- [ ] 不同 seed 下结果应发生变化； -- [ ] 同一 seed、同一调用顺序下结果应一致。 - -#### （3）状态恢复测试 - -验证 `GetState()` / `SetState()` 是否真正恢复随机序列，包括： - -- [ ] 保存 state； -- [ ] 继续生成一段随机结果； -- [ ] 恢复 state； -- [ ] 再次生成随机结果； -- [ ] 验证恢复后的结果与原序列对齐。 - -CPU 后端必须覆盖该测试；CUDA 后端若已实现，也应覆盖。 - -#### （4）默认 Generator 行为测试 - -验证随机算子是否正确使用默认 Generator，包括： - -- [ ] 不传 Generator 时，是否使用当前设备默认 Generator； -- [ ] CPU tensor 是否使用 CPU 默认 Generator； -- [ ] CUDA tensor 是否使用对应 CUDA 设备默认 Generator； -- [ ] 显式传入 Generator 时，是否不会误用默认 Generator。 - -#### （5）主流框架语义对齐验证 - -可选取典型场景与 PyTorch 进行语义对比，包括： - -- [ ] 手动设置 seed 后结果可复现； -- [ ] Dropout 等随机运算在同 seed 下可复现； -- [ ] 随机张量生成接口支持显式 Generator 与默认 Generator； -- [ ] state 保存与恢复语义一致。 - -需要明确：本项目不要求底层随机算法与 PyTorch 逐 bit 一致，也不要求 CPU 与 CUDA 随机结果彼此一致。验收重点是接口语义一致、可复现逻辑一致、默认 Generator 行为一致。 - -## 四、提交报告要求 - -除代码外，参赛者应提交以下内容作为技术报告： - -### 1. 功能正确性验证 - -需要提供上述 Generator 基础设施功能验证结果 - -### 2. 对齐性与行为分析报告 - -需要说明本项目实现与主流框架的对齐情况。建议从以下角度展开： - -- [ ] 参考了哪些主流框架设计； -- [ ] 当前实现与 PyTorch 在接口语义上的对应关系； -- [ ] 哪些行为做到基本一致； -- [ ] 哪些行为暂未完全对齐，以及原因分析； -- [ ] 当前实现范围与后续可扩展方向。 - -此部分重点不是要求一切细节完全一致，而是要求留档明确设计思路，能够清晰说明： - -- [ ] 自己实现了什么； -- [ ] 为什么这样设计； -- [ ] 与主流框架相比的设计异同及原因； -- [ ] 后续演进路径是什么。 - -### 3. 测试与可复现性说明 - -需要提供完整测试脚本与说明，保证 reviewer 可在相同环境下复现主要结果。 - -## 五、验收要求 - -### 验收要求 - -项目需满足以下基本要求： - -- [ ] 代码以 PR 形式提交，结构清晰，具备基本可 review 性； -- [x] 完成统一 Generator 抽象设计； -- [x] 实现 CPU 与 CUDA Generator； -- [x] 建立 CPU 与各 CUDA device 的默认 Generator 管理机制； -- [x] 提供统一的全局随机种子设置入口； -- [x] 初始化算子支持默认 Generator 与显式 Generator 两种使用方式； -- [x] 支持 Generator 状态的获取与恢复（get_state / set_state）； -- [ ] 至少改造一类初始化算子和一类框架内随机数生成调用，使其接入 Generator； -- [ ] 提供测试或运行日志，验证随机行为具备基本可复现性（同 seed 一致、不同 seed 不同、状态恢复有效）。 - -### 加分项 - -在满足验收要求的基础上，具备以下内容可作为加分项： - -- [ ] 设计清晰，接口与实现分层合理，便于后续扩展 -- [ ] 测试覆盖充分，包含 seed、state、默认/显式 generator、跨设备等关键场景； -- [ ] 调用处改造完全，接口风格统一； -- [ ] 与 PyTorch 的接口语义和行为分析完整，报告质量高 -- [ ] PR 经过完整 review 流程，达到可合入标准。 - -## 参考链接 - -- https://docs.pytorch.org/docs/stable/generated/torch.Generator.html#generator -- https://github.com/pytorch/pytorch/blob/main/aten/src/ATen/core/Generator.h -- https://github.com/pytorch/pytorch/blob/main/aten/src/ATen/CPUGeneratorImpl.cpp -- https://github.com/pytorch/pytorch/blob/main/aten/src/ATen/cuda/CUDAGeneratorImpl.cpp -- https://docs.pytorch.org/docs/stable/notes/randomness.html#pytorch-random-number-generator diff --git a/tests/generator/cuda_only/test_cuda_generator.cc b/tests/generator/cuda_only/test_cuda_generator.cc index 0cdd468b..aea83ff7 100644 --- a/tests/generator/cuda_only/test_cuda_generator.cc +++ b/tests/generator/cuda_only/test_cuda_generator.cc @@ -7,7 +7,6 @@ #include "infini_train/include/generator.h" #include "infini_train/include/nn/init.h" #include "infini_train/include/tensor.h" -#include "infini_train/src/core/runtime/cpu/cpu_generator_impl.h" #include "infini_train/src/core/runtime/cuda/cuda_generator_impl.h" #include "tests/common/test_utils.h" @@ -32,26 +31,9 @@ TEST(CudaGeneratorTest, StateIsAnOpaqueUint8CpuTensor) { EXPECT_EQ(state->SizeInBytes(), sizeof(uint64_t) * 2); } -TEST(CudaGeneratorTest, RejectsStateWithInvalidSize) { - Generator generator = core::cuda::createCUDAGenerator(0, kSeed); - Tensor invalid_state(std::vector{1}, DataType::kUINT8, - Device(Device::DeviceType::kCPU, 0)); - - EXPECT_DEATH(generator.set_state(invalid_state), ""); -} - -TEST(CudaGeneratorTest, CpuAndCudaStatesAreNotInterchangeable) { - Generator cpu_generator = core::cpu::createCPUGenerator(kSeed); - Generator cuda_generator = core::cuda::createCUDAGenerator(0, kSeed); - auto cpu_state = cpu_generator.get_state(); - auto cuda_state = cuda_generator.get_state(); - - EXPECT_DEATH(cuda_generator.set_state(*cpu_state), ""); - EXPECT_DEATH(cpu_generator.set_state(*cuda_state), ""); -} - -TEST(CudaGeneratorTest, DefaultGeneratorsArePerDevice) { +TEST(CudaGeneratorTest, MissingGeneratorUsesMatchingDeviceDefaultGenerator) { REQUIRE_MIN_DEVICES(2); + manual_seed(kSeed); const Generator &device_zero = core::cuda::getDefaultCUDAGenerator(0); const Generator &device_one = core::cuda::getDefaultCUDAGenerator(1); @@ -59,13 +41,15 @@ TEST(CudaGeneratorTest, DefaultGeneratorsArePerDevice) { EXPECT_EQ(device_zero.device().index(), 0); EXPECT_EQ(device_one.device().index(), 1); + const std::vector device_zero_before = StateBytes(device_zero); const std::vector device_one_before = StateBytes(device_one); auto tensor = std::make_shared( std::vector{1024}, DataType::kFLOAT32, - Device(Device::DeviceType::kCUDA, 0)); - nn::init::Uniform(tensor, 0.0f, 1.0f, device_zero); + Device(Device::DeviceType::kCUDA, 1)); + nn::init::Uniform(tensor); - EXPECT_EQ(device_one_before, StateBytes(device_one)); + EXPECT_EQ(device_zero_before, StateBytes(device_zero)); + EXPECT_NE(device_one_before, StateBytes(device_one)); } } // namespace diff --git a/tests/generator/test_generator.cc b/tests/generator/test_generator.cc index cf0b2867..1c642a94 100644 --- a/tests/generator/test_generator.cc +++ b/tests/generator/test_generator.cc @@ -1,4 +1,3 @@ -#include #include #include #include @@ -60,15 +59,6 @@ std::vector ToHostVector(const std::shared_ptr &tensor) { return std::vector(data, data + cpu.NumElements()); } -std::vector ToHostBytes(const std::shared_ptr &tensor) { - const Device device = tensor->GetDevice(); - Synchronize(device); - Tensor cpu = tensor->To(Device(Device::DeviceType::kCPU, 0)); - Synchronize(device); - const auto *data = static_cast(cpu.DataPtr()); - return std::vector(data, data + cpu.SizeInBytes()); -} - std::vector StateBytes(const Generator &generator) { auto state = generator.get_state(); const auto *data = static_cast(state->DataPtr()); @@ -81,30 +71,15 @@ std::shared_ptr MakeTensor(Device device) { class GeneratorTest : public InfiniTrainTest {}; -TEST_P(GeneratorTest, ReportsDeviceAndInitialSeed) { +TEST_P(GeneratorTest, SupportsDeviceAndSeedInterface) { const Device device = GetDevice(); Generator generator = CreateGenerator(device, kSeed); EXPECT_TRUE(generator.defined()); EXPECT_EQ(generator.device(), device); EXPECT_EQ(generator.current_seed(), kSeed); -} - -TEST_P(GeneratorTest, ShallowCopySharesStateAndCloneDoesNot) { - Generator generator = CreateGenerator(GetDevice(), kSeed); - Generator alias = generator; - Generator clone = generator.clone(); - - EXPECT_EQ(alias, generator); - EXPECT_NE(clone, generator); - EXPECT_EQ(StateBytes(alias), StateBytes(generator)); - EXPECT_EQ(StateBytes(clone), StateBytes(generator)); - - auto tensor = MakeTensor(GetDevice()); - nn::init::Uniform(tensor, -1.0f, 1.0f, alias); - - EXPECT_EQ(StateBytes(alias), StateBytes(generator)); - EXPECT_NE(StateBytes(clone), StateBytes(generator)); + generator.set_current_seed(kSeed + 1); + EXPECT_EQ(generator.current_seed(), kSeed + 1); } TEST_P(GeneratorTest, UniformIsReproducibleForSameSeed) { @@ -131,59 +106,36 @@ TEST_P(GeneratorTest, NormalIsReproducibleForSameSeed) { EXPECT_EQ(ToHostVector(first), ToHostVector(second)); } -TEST_P(GeneratorTest, UniformSupportsEveryFloatingPointDtype) { - for (const DataType dtype : {DataType::kFLOAT16, DataType::kBFLOAT16, - DataType::kFLOAT32, DataType::kFLOAT64}) { - Generator first_generator = CreateGenerator(GetDevice(), kSeed); - Generator second_generator = CreateGenerator(GetDevice(), kSeed); - auto first = std::make_shared( - std::vector{kElements}, dtype, GetDevice()); - auto second = std::make_shared( - std::vector{kElements}, dtype, GetDevice()); - - nn::init::Uniform(first, -3.0f, 7.0f, first_generator); - nn::init::Uniform(second, -3.0f, 7.0f, second_generator); - - EXPECT_EQ(ToHostBytes(first), ToHostBytes(second)) - << "dtype=" << static_cast(dtype); - } -} +TEST_P(GeneratorTest, RandAndRandnSupportExplicitAndDefaultGenerators) { + const Device device = GetDevice(); + Generator first_generator = CreateGenerator(device, kSeed); + Generator second_generator = CreateGenerator(device, kSeed); -TEST_P(GeneratorTest, NormalSupportsEveryFloatingPointDtype) { - for (const DataType dtype : {DataType::kFLOAT16, DataType::kBFLOAT16, - DataType::kFLOAT32, DataType::kFLOAT64}) { - Generator first_generator = CreateGenerator(GetDevice(), kSeed); - Generator second_generator = CreateGenerator(GetDevice(), kSeed); - auto first = std::make_shared( - std::vector{kElements}, dtype, GetDevice()); - auto second = std::make_shared( - std::vector{kElements}, dtype, GetDevice()); - - nn::init::Normal(first, 2.0f, 0.5f, first_generator); - nn::init::Normal(second, 2.0f, 0.5f, second_generator); - - EXPECT_EQ(ToHostBytes(first), ToHostBytes(second)) - << "dtype=" << static_cast(dtype); - } + auto first = nn::function::Rand({17, 19}, DataType::kFLOAT32, device, first_generator); + auto second = nn::function::Rand({17, 19}, DataType::kFLOAT32, device, second_generator); + EXPECT_EQ(ToHostVector(first), ToHostVector(second)); + + manual_seed(kSeed); + auto first_default = nn::function::Randn({17, 19}, DataType::kFLOAT32, device); + manual_seed(kSeed); + auto second_default = nn::function::Randn({17, 19}, DataType::kFLOAT32, device); + EXPECT_EQ(ToHostVector(first_default), ToHostVector(second_default)); } -TEST_P(GeneratorTest, RandAndRandnPreserveRequestedTensorProperties) { - Generator rand_generator = CreateGenerator(GetDevice(), kSeed); - Generator randn_generator = CreateGenerator(GetDevice(), kSeed); - - auto uniform = nn::function::Rand({17, 19}, DataType::kFLOAT64, GetDevice(), - rand_generator, true); - auto normal = nn::function::Randn({17, 19}, DataType::kFLOAT16, GetDevice(), - randn_generator, true); - - EXPECT_EQ(uniform->Dims(), (std::vector{17, 19})); - EXPECT_EQ(uniform->Dtype(), DataType::kFLOAT64); - EXPECT_EQ(uniform->GetDevice(), GetDevice()); - EXPECT_TRUE(uniform->requires_grad()); - EXPECT_EQ(normal->Dims(), (std::vector{17, 19})); - EXPECT_EQ(normal->Dtype(), DataType::kFLOAT16); - EXPECT_EQ(normal->GetDevice(), GetDevice()); - EXPECT_TRUE(normal->requires_grad()); +TEST_P(GeneratorTest, DropoutMaskIsReproducibleForSameSeed) { + const Device device = GetDevice(); + auto input = MakeTensor(device); + input->Fill(1.0f); + + manual_seed(kSeed); + auto first = nn::function::Dropout(input, 0.25, true); + manual_seed(kSeed); + auto second = nn::function::Dropout(input, 0.25, true); + manual_seed(kSeed + 1); + auto different_seed = nn::function::Dropout(input, 0.25, true); + + EXPECT_EQ(ToHostVector(first), ToHostVector(second)); + EXPECT_NE(ToHostVector(first), ToHostVector(different_seed)); } TEST_P(GeneratorTest, ConsecutiveCallsAdvanceSequence) { @@ -209,6 +161,30 @@ TEST_P(GeneratorTest, DifferentSeedsProduceDifferentSequences) { EXPECT_NE(ToHostVector(first), ToHostVector(second)); } +TEST_P(GeneratorTest, SameSeedAndCallOrderReproduceResults) { + const Device device = GetDevice(); + auto first_uniform = MakeTensor(device); + auto second_uniform = MakeTensor(device); + auto first_normal = MakeTensor(device); + auto second_normal = MakeTensor(device); + auto input = MakeTensor(device); + input->Fill(1.0f); + + manual_seed(kSeed); + nn::init::Uniform(first_uniform); + auto first_dropout = nn::function::Dropout(input, 0.5, true); + nn::init::Normal(first_normal); + + manual_seed(kSeed); + nn::init::Uniform(second_uniform); + auto second_dropout = nn::function::Dropout(input, 0.5, true); + nn::init::Normal(second_normal); + + EXPECT_EQ(ToHostVector(first_uniform), ToHostVector(second_uniform)); + EXPECT_EQ(ToHostVector(first_dropout), ToHostVector(second_dropout)); + EXPECT_EQ(ToHostVector(first_normal), ToHostVector(second_normal)); +} + TEST_P(GeneratorTest, StateRestoreReplaysUniformSequence) { Generator generator = CreateGenerator(GetDevice(), kSeed); auto prefix = MakeTensor(GetDevice()); @@ -278,6 +254,7 @@ TEST_P(GeneratorTest, GlobalManualSeedReproducesDefaultGeneratorResults) { auto second = MakeTensor(device); manual_seed(kSeed); + EXPECT_EQ(GetDefaultGenerator(device).current_seed(), kSeed); nn::init::Uniform(first); manual_seed(kSeed); nn::init::Uniform(second); diff --git a/tests/generator/test_generator_interface.cc b/tests/generator/test_generator_interface.cc index 7a698f9c..eb8a93a8 100644 --- a/tests/generator/test_generator_interface.cc +++ b/tests/generator/test_generator_interface.cc @@ -1,6 +1,5 @@ #include #include -#include #include #include "gtest/gtest.h" @@ -14,16 +13,6 @@ namespace { constexpr uint64_t kSeed = 0x12345678ULL; -TEST(GeneratorInterfaceTest, DefaultConstructedGeneratorIsUndefined) { - Generator generator; - EXPECT_FALSE(generator.defined()); -} - -TEST(GeneratorInterfaceTest, CheckGeneratorRejectsUndefinedGenerator) { - Generator generator; - EXPECT_THROW(check_generator(generator), std::invalid_argument); -} - TEST(GeneratorInterfaceTest, CpuStateIsAnOpaqueUint8CpuTensor) { Generator generator = core::cpu::createCPUGenerator(kSeed); auto state = generator.get_state(); @@ -44,21 +33,5 @@ TEST(GeneratorInterfaceTest, CpuStateRoundTripRestoresSeed) { EXPECT_EQ(generator.current_seed(), kSeed); } -TEST(GeneratorInterfaceTest, RejectsStateWithWrongDtype) { - Generator generator = core::cpu::createCPUGenerator(kSeed); - Tensor invalid_state(std::vector{64}, DataType::kFLOAT32, - Device(Device::DeviceType::kCPU, 0)); - - EXPECT_DEATH(generator.set_state(invalid_state), "RNG state must be a UINT8 tensor"); -} - -TEST(GeneratorInterfaceTest, RejectsTruncatedCpuState) { - Generator generator = core::cpu::createCPUGenerator(kSeed); - Tensor invalid_state(std::vector{1}, DataType::kUINT8, - Device(Device::DeviceType::kCPU, 0)); - - EXPECT_DEATH(generator.set_state(invalid_state), "CPU generator state is too small"); -} - } // namespace } // namespace infini_train::test