From a4fe4703efc481ff347b521acce382c5515ffe9d Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E7=8B=AC=E8=A7=92=E4=BB=99=E4=B8=8E=E6=98=9F?= Date: Fri, 18 Sep 2026 15:39:43 +0800 Subject: [PATCH 1/8] feat: small-model CNN training support (Conv2d v1, ReLU, DataLoader, MNIST demo, DDP) - Conv2d v1: FP32/NCHW-OIHW, scalar kernel/stride/padding API, Kh/Kw-independent internals; Conv2dMeta centralizes shape arithmetic; selective-save autograd - Backends: CPU direct reference + CUDA im2col + framework GEMM (gather col2im, cub bias reduction) - DataLoader shape-preserving Stack ([B,*sample_dims]); MNIST sample [1,28,28]; stride bug fix; MLP Flatten(1) adaptation - MNIST: MnistCnn + --model (default mlp unchanged); shared_ptr migration for enable_shared_from_this Modules - DDP: infini_run multiprocess, Broadcast-before-wrap, sharded eval + epoch-level AllReduce(SUM), no per-step logging collectives - Tests: kernel exact-value, gradcheck, torch golden parity, im2col/CUDA grids, dataset/dataloader, network parity dumper --- docs/small_model_patch.md | 151 ++++++ example/mnist/dataset.cc | 13 +- example/mnist/main.cc | 122 ++++- example/mnist/net.cc | 28 +- example/mnist/net.h | 10 + infini_train/include/autograd/activations.h | 12 + infini_train/include/autograd/conv.h | 34 ++ infini_train/include/nn/modules/activations.h | 7 + infini_train/include/nn/modules/conv.h | 35 ++ infini_train/src/autograd/activations.cc | 26 + infini_train/src/autograd/conv.cc | 73 +++ infini_train/src/dataloader.cc | 15 +- infini_train/src/kernels/common/conv.h | 69 +++ infini_train/src/kernels/cpu/conv.cc | 240 +++++++++ infini_train/src/kernels/cpu/relu.cc | 43 ++ infini_train/src/kernels/cuda/conv.cu | 495 ++++++++++++++++++ infini_train/src/kernels/cuda/relu.cu | 86 +++ infini_train/src/nn/modules/activations.cc | 4 + infini_train/src/nn/modules/conv.cc | 44 ++ tests/CMakeLists.txt | 3 + tests/autograd/test_autograd_conv_cpu.cc | 150 ++++++ tests/autograd/test_autograd_conv_cuda.cc | 283 ++++++++++ .../autograd/test_autograd_conv_gradcheck.cc | 128 +++++ tests/autograd/test_autograd_conv_im2col.cc | 62 +++ tests/autograd/test_autograd_conv_parity.cc | 105 ++++ tests/autograd/test_autograd_relu.cc | 49 ++ tests/dataloader/CMakeLists.txt | 5 + tests/dataloader/test_dataloader_stack.cc | 64 +++ tests/example/CMakeLists.txt | 17 + tests/example/mnist_parity.cc | 262 +++++++++ tests/example/test_mnist_dataset.cc | 113 ++++ 31 files changed, 2716 insertions(+), 32 deletions(-) create mode 100644 docs/small_model_patch.md create mode 100644 infini_train/include/autograd/conv.h create mode 100644 infini_train/include/nn/modules/conv.h create mode 100644 infini_train/src/autograd/conv.cc create mode 100644 infini_train/src/kernels/common/conv.h create mode 100644 infini_train/src/kernels/cpu/conv.cc create mode 100644 infini_train/src/kernels/cpu/relu.cc create mode 100644 infini_train/src/kernels/cuda/conv.cu create mode 100644 infini_train/src/kernels/cuda/relu.cu create mode 100644 infini_train/src/nn/modules/conv.cc create mode 100644 tests/autograd/test_autograd_conv_cpu.cc create mode 100644 tests/autograd/test_autograd_conv_cuda.cc create mode 100644 tests/autograd/test_autograd_conv_gradcheck.cc create mode 100644 tests/autograd/test_autograd_conv_im2col.cc create mode 100644 tests/autograd/test_autograd_conv_parity.cc create mode 100644 tests/autograd/test_autograd_relu.cc create mode 100644 tests/dataloader/test_dataloader_stack.cc create mode 100644 tests/example/CMakeLists.txt create mode 100644 tests/example/mnist_parity.cc create mode 100644 tests/example/test_mnist_dataset.cc diff --git a/docs/small_model_patch.md b/docs/small_model_patch.md new file mode 100644 index 000000000..b52f9ebd1 --- /dev/null +++ b/docs/small_model_patch.md @@ -0,0 +1,151 @@ +# 小模型训练支持 — 技术方案与结果摘要（分支 `work/small-model-cnn-impl`） + +> 来源：本文件 §5–§6 逐字移植自上游提案文档 `/docs/small_model_patch.md` +> （只读来源，未修改）；§7 为本分支实测结果摘要。 + +--- + +## 五、技术方案（v1 锁定） + +结合上游 PR #220（kernel / 验证 / PR discipline）与 #219（数据形状语义 / DDP metric 设计）的对比结论，本项目做第三版：**用 #219 的语义设计 + #220 的 kernel 与验证方法**，并去掉两边共有的 demo-level special case。 + +### 5.1 Conv2d v1 scope + +```text +dtype: FP32 only（入口显式 CHECK，不静默 fallback） +layout: input NCHW，weight OIHW +API: scalar kernel_size / stride / padding + optional bias +内部语义: Kh/Kw 独立（读 weight.Dims()[2/3]），不加 Kh == Kw 限制 +不支持: tuple kernel/stride/padding、dilation、groups、FP16/BF16 +``` + +即：**public API 收敛到任务 scope，internal semantics 不做无意义限制**。非 square 权重只在 kernel / autograd 层单测覆盖，Module 层不暴露 tuple API。 + +### 5.2 双后端职责划分 + +```text +CPU: direct reference implementation（forward + dX + dW + db），暂不用 Eigen / im2col +CUDA: im2col + framework GEMM（Dispatcher "Gemm"），gather 风格 col2im，bias reduction +``` + +CPU 与现有 `cpu/matmul.cc` 朴素循环风格一致，职责是可审计的数学基准；CUDA 复用 `cuda/common/gemm.cu` 的 strided-batched GEMM，与 `cuda/matmul.cu`、`cuda/linear.cu` 的调用方式一致。两条独立实现路径互为 cross-check，避免 im2col layout 错误在两边同时成立。 + +`Conv2dMeta`（common 层）集中保存 stride / padding / batch / 通道 / `kernel_h,w` / `output_h,w` / `patches` / `kernel_elems`，shape arithmetic 全局只有一份，CPU / CUDA 不再重复定义。 + +Autograd 沿用两边已验证的 selective-save（需 dW 才存 input，需 dX 才存 weight），`Conv2dBackwardInput / Weight / Bias` 保持三个独立 backend op，不合一。 + +### 5.3 实现次序 + +```text +latest upstream + → MNIST stride bug 修复（image_size_in_bytes_ 按 FLOAT32 计） + → shape-preserving Stack（[B,*sample_dims]）+ 回归测试，不碰 DDP sharding + → ReLU + → Conv2dMeta + CPU direct（forward/dW/dX/db，stride/padding） + → CPU finite-difference gradcheck + PyTorch parity（数学 correctness 结束） + → im2col 单测（tiny 人工可算例子，隔离 layout/transpose bug） + → CUDA forward（先 per-image GEMM，再 strided-batched 优化，语义不变） + → CUDA dW → dX（gather col2im）→ db + → CNN demo（--model 默认 mlp 不变，cnn 显式 opt-in） + → PyTorch 网络级 parity + → infini_run 多进程 DDP + sharded eval + epoch 级 metric reduce + → 1GPU vs 2GPU 数值等价 + 性能清理 +``` + +### 5.4 DataLoader 与 Demo + +Dataset 单样本为 `[1,28,28]`，`Stack` 保持形状得 `[B,1,28,28]`，模型内不再 `View` 恢复形状；MLP 自行 `Flatten(1)`。DDP 采用 `To(device) → Broadcast → DDP wrap` 顺序，训练 step 只同步梯度，logging 不引入额外 collective。 + +### 5.5 Oracle 策略 + +```text + PyTorch + │ + ▼ +CPU direct ─────── CUDA GEMM + │ │ + └── finite diff ────┘ +``` + +CPU numerical gradcheck（tiny tensor，dX/dW/db）+ PyTorch 对标（logits / loss / grad / 单步更新 / 多步趋势）两套独立 oracle，不依赖 CUDA direct 中间实现。 + +--- + +## 六、环境与算力需求 + +### 6.1 开销估算 + +参考网络 `Conv2d(1,16,3) → ReLU → Conv2d(16,32,3) → ReLU → Flatten → Linear(18432,10)`（s=1/p=0）： + +- 参数约 189K（conv 4.8K + fc 184K），权重 < 1MB； +- 约 2.9M MAC/sample，单 epoch（60K）forward 约 0.35 TFLOP，fwd+bwd 约 1 TFLOP； +- batch=64 时激活约 7.7MB，整网 VRAM < 100MB。 + +结论：单卡开发全够，瓶颈在 kernel 启动与 DataLoader，不在 FLOPs；CPU direct 跑 60K 为分钟级，可接受。 + +### 6.2 本机现状与缺口 + +- 本机：1× RTX 4050 Laptop 6GB（空闲）、16 CPU、19GB RAM、磁盘充足，nvcc 13.4，gcc 16，cmake 4.4。 +- 缺：`build/` 未建、submodule 未初始化、torch 未装、无 MNIST 数据、无 NCCL（标准路径下未找到 `nccl.h` / `libnccl*`）。 +- 唯一算力缺口：**双卡 DDP 等价验证**（1GPU bs=128 vs 2×64）。本地单卡只能做单卡全量 + DDP 流程联调，最终等价性需在双卡机或 CI 跑一次并留 log 作报告证据。 + +### 6.3 环境准备清单 + +```bash +# 1. 子模块（BUILD_TEST 需要 googletest） +git submodule update --init third_party/googletest + +# 2. MNIST 数据（官方 TorchVision 镜像，约 50MB） +./scripts/assets/prepare-infinitrain-assets.sh mnist +# 产物：data/mnist/{train-images-idx3-ubyte,train-labels-idx1-ubyte,t10k-images-idx3-ubyte,t10k-labels-idx1-ubyte} + +# 3. 构建（本地单卡；4050 为 Ada sm_89，仓库默认 CUDA_ARCHITECTURES=75;80;90 +# 不含 89，本地建议覆盖为 native；缺 NCCL 时先关 DDP 构建） +cmake -S . -B build -DUSE_CUDA=ON -DUSE_NCCL=OFF -DBUILD_TEST=ON \ + -DCMAKE_CUDA_ARCHITECTURES=native -DCMAKE_POLICY_VERSION_MINIMUM=3.5 +cmake --build build -j +# 有 NCCL 的双卡机：-DUSE_NCCL=ON，并用 ./build/infini_run 启动多进程 DDP + +# 4. PyTorch 对标（仅 CPU 数值 oracle 即可，不需要 GPU torch） +uv venv .venv && uv pip install "torch>=2.4 --index-url https://download.pytorch.org/whl/cpu" numpy +# 对齐固定：seed、输入、初始权重（.bin 导出供 InfiniTrain 读取）、lr/batch +# 对比项：logits、loss、关键参数 grad、单步更新后参数、多步 loss 趋势 + +# 5. 端到端（默认 MLP 行为不变，CNN 显式指定） +./build/mnist --dataset data/mnist --device cpu +./build/mnist --dataset data/mnist --device cuda --model=cnn --num_epoch=3 --lr=0.1 +``` + +MNIST 四个 IDX 文件可按字节数 + 首 16 字节 magic 校验（images 0x00000803，labels 0x00000801），损坏时 `FORCE=1` 重跑准备脚本。 + +--- + +## 七、结果摘要 + +超参（端到端统一）：batch=64，lr=0.01，SGD，CNN 即 §6.1 参考网络。 + +| 任务 | 结论 | +| --- | --- | +| P0 构建 | `USE_CUDA=ON / USE_NCCL=OFF / arch native / POLICY_VERSION_MINIMUM 3.5` 全量通过；MNIST 数据 + torch 2.14+cpu 就绪 | +| P1 dataset | `test_mnist_dataset` 3/3；stride bug 实证：旧值 784B vs 正确 3136B | +| P2 dataloader | DataLoader 8/8（含新增 shape-preserving Stack 回归），MNISTDataset 3/3；MLP smoke acc 0.60 | +| P3 ReLU | CPU 3/3 | +| P4 conv CPU | 8/8 exact-value；全量 Autograd 144/144 | +| P5 gradcheck+golden | CPU finite-difference + torch golden 通过；Autograd 增至 145 pass | +| P6 im2col | RTX 4050 上通过；发现 Dispatcher const& vs by-value 坑（已在代码注释记录） | +| P7 CUDA forward | 96 配置 grid ≤1e-5；per-image 1903ms → strided-batched 849ms | +| P8/P9/P10 CUDA dW/dX/db | 三个 grid 全部 ≤1e-5；Conv 相关 20/20 | +| P11 e2e smoke（1 epoch） | mlp acc 0.60；cnn cpu acc 0.90、loss 2.3→0.06；cnn cuda 与 CPU bit-identical，2.3s/epoch | +| P12 网络级 parity | 140/140 @1e-5（worst 9.5e-7）；10-step loss 双边严格单调递减 | +| P13a DDP（main.cc only） | 单进程 cuda acc 0.9028；NCCL-ON 可编译（需手动 `-I/-L/-rpath`，顶层 CMake 未加 NCCL include 路径）；2 进程止步于 `SetDevice(cuda:1)`——单卡机上限 | +| P13b 双卡等价 | BLOCKED：无双卡机，已留诚实 blocker 记录 | +| 全量 ctest（NCCL-OFF） | 354/355 通过；唯一失败 `stacktrace`（glog 符号化环境问题，本分支未触及该文件，属 HEAD 预存失败）；CUDA/CPU 交叉实例化 SKIP 系 `ONLY_*` 宏按设计跳过 | + +HEAD 预存 bug 修复（保留，与上游 #220 一致）：`main.cc` 改用 `shared_ptr` 持有 Module（Module 是 `enable_shared_from_this`，栈实例触发 `bad_weak_ptr` 崩溃）。 + +ENV-ONLY 说明：gcc16 本地构建曾需在 `tensor.cc` / `profiler.cc` 各加一行 `#include `；上游 gcc13 不受影响，该两行已 revert，不在 PR 内。 + +> Sharded-eval 注记：`DistributedDataLoader` 用 floor 除法（`num_batches_ = Size / global_batch`，丢尾）， +> 单进程 eval 用 ceil 除法覆盖全集。因此 P13b 1GPU-vs-2GPU 等价比较必须满足二者之一： +> (a) 只在 EVALUATED 样本上比较全局 loss/accuracy（reduce 上来的三元组本身即已是在 evaluated 样本上的精确值）； +> (b) 取数据集大小能整除 global batch 的配置。 diff --git a/example/mnist/dataset.cc b/example/mnist/dataset.cc index ee683f6d4..d408a6b99 100644 --- a/example/mnist/dataset.cc +++ b/example/mnist/dataset.cc @@ -87,9 +87,18 @@ MNISTDataset::MNISTDataset(const std::string &dataset, bool train) ReadSN3PascalVincentFile(std::format("{}/{}-images-idx3-ubyte", dataset, train ? kTrainPrefix : kTestPrefix))), label_file_(ReadSN3PascalVincentFile( std::format("{}/{}-labels-idx1-ubyte", dataset, train ? kTrainPrefix : kTestPrefix))), - image_dims_(image_file_.dims.begin() + 1, image_file_.dims.end()), + // Insert a leading channel dim so samples are [1, 28, 28] (NCHW-ready); the + // underlying buffer is unchanged, only the per-sample view dims gain the channel. + // Built inline so image_size_in_bytes_ below (init-list order) already sees [1, 28, 28]. + image_dims_([&] { + std::vector dims = {1}; + dims.insert(dims.end(), image_file_.dims.begin() + 1, image_file_.dims.end()); + return dims; + }()), label_dims_(label_file_.dims.begin() + 1, label_file_.dims.end()), - image_size_in_bytes_(kSN3TypeToSize.at(image_file_.type) + // NOTE: image_file_.tensor is converted to FLOAT32 below, so the per-sample stride must count FLOAT32 bytes, + // not the on-disk UINT8 bytes from kSN3TypeToSize. + 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/example/mnist/main.cc b/example/mnist/main.cc index 7744e0947..2014b434c 100644 --- a/example/mnist/main.cc +++ b/example/mnist/main.cc @@ -12,6 +12,12 @@ #include "infini_train/include/dataloader.h" #include "infini_train/include/device.h" #include "infini_train/include/nn/modules/loss.h" +#include "infini_train/include/nn/parallel/ddp/distributed_data_parallel.h" +#include "infini_train/include/nn/parallel/global.h" +#include "infini_train/include/nn/parallel/process_group.h" +#include "infini_train/include/nn/parallel/rank.h" +#include "infini_train/include/nn/parallel/reduce_op_type.h" +#include "infini_train/include/nn/parallel/utils.h" #include "infini_train/include/optimizer.h" #include "example/mnist/dataset.h" @@ -22,6 +28,7 @@ DEFINE_int32(bs, 64, "batch size"); DEFINE_int32(num_epoch, 1, "num epochs"); DEFINE_double(lr, 0.01, "learning rate"); DEFINE_string(device, "cpu", "device type (cpu/cuda)"); +DEFINE_string(model, "mlp", "model type (mlp/cnn)"); using namespace infini_train; @@ -31,30 +38,78 @@ constexpr int kNumClasses = 10; constexpr char kDeviceCPU[] = "cpu"; constexpr char kDeviceCUDA[] = "cuda"; +constexpr char kModelMLP[] = "mlp"; +constexpr char kModelCNN[] = "cnn"; }; // namespace DEFINE_validator(device, [](const char *, const std::string &value) { return value == kDeviceCPU || value == kDeviceCUDA; }); +DEFINE_validator(model, + [](const char *, const std::string &value) { return value == kModelMLP || value == kModelCNN; }); int main(int argc, char *argv[]) { gflags::ParseCommandLineFlags(&argc, &argv, true); google::InitGoogleLogging(argv[0]); + // Distributed init from env (WORLD_SIZE/RANK/LOCAL_RANK set by infini_run). + // Same helper pattern as example/gpt2: pure data-parallel layout (TP=PP=1). + nn::parallel::global::InitAllEnv(/*nthread_per_process=*/1, /*tensor_parallel_size=*/1, + /*sequence_parallel_enabled=*/false, /*pipeline_parallel_size=*/1, + /*virtual_pipeline_parallel=*/1); + const int ddp_world_size = nn::parallel::global::GetDataParallelSize(); + const bool distributed = ddp_world_size > 1; + nn::parallel::Rank rank(nn::parallel::global::GetGlobalProcRank(), 0, + nn::parallel::global::GetNprocPerNode(), /*threads_per_process=*/1); + nn::parallel::global::thread_global_rank = rank.GlobalRank(); + const int ddp_rank = distributed ? rank.GlobalRank() : 0; + const bool is_main_rank = rank.IsMainRank(); + + const nn::parallel::ProcessGroup *ddp_pg = nullptr; + if (distributed) { + const auto device_type = Device::DeviceType::kCUDA; + auto *pg_factory = nn::parallel::ProcessGroupFactory::Instance(device_type); + ddp_pg = pg_factory->GetOrCreate(nn::parallel::GetDataParallelProcessGroupName(rank.GlobalRank()), + nn::parallel::GetDataParallelGroupRanks(rank.GlobalRank())); + } + auto train_dataset = std::make_shared(FLAGS_dataset, true); - DataLoader train_dataloader(train_dataset, FLAGS_bs); + std::unique_ptr train_dataloader + = distributed ? std::unique_ptr( + std::make_unique(train_dataset, FLAGS_bs, ddp_rank, ddp_world_size)) + : std::unique_ptr(std::make_unique(train_dataset, FLAGS_bs)); // TODO(dcj): Add sampler & eval dataloader later. auto test_dataset = std::make_shared(FLAGS_dataset, false); - DataLoader test_dataloader(test_dataset, FLAGS_bs); - - auto network = MNIST(); + std::unique_ptr test_dataloader + = distributed ? std::unique_ptr( + std::make_unique(test_dataset, FLAGS_bs, ddp_rank, ddp_world_size)) + : std::unique_ptr(std::make_unique(test_dataset, FLAGS_bs)); + + std::shared_ptr network; + if (FLAGS_model == kModelCNN) { + network = std::make_shared(); + } else { + network = std::make_shared(); + } Device device = FLAGS_device == kDeviceCPU ? Device() : Device(Device::DeviceType::kCUDA, 0); + if (distributed) { + device = Device(Device::DeviceType::kCUDA, nn::parallel::global::GetLocalProcRank()); + } Device cpu_device = Device(); - network.To(device); + network->To(device); + + if (distributed) { + // Sync initial params from root so all ranks start from identical weights, + // then wrap with DDP (grad-only sync during training, no per-step loss sync). + // NOTE: complete all .To(device) conversions before wrapping (same rule as gpt2). + ddp_pg->Broadcast(network->Parameters(), /*root_rank_in_group=*/0); + network = std::make_shared(network, rank, + nn::parallel::DistributedDataParallelConfig{}); + } - auto loss_fn = nn::CrossEntropyLoss(); - loss_fn.To(device); - auto optimizer = optimizers::SGD(network.Parameters(), FLAGS_lr); + auto loss_fn = std::make_shared(); + loss_fn->To(device); + auto optimizer = optimizers::SGD(network->Parameters(), FLAGS_lr); for (int epoch = 0; epoch < FLAGS_num_epoch; ++epoch) { int train_idx = 0; @@ -62,14 +117,14 @@ int main(int argc, char *argv[]) { const auto epoch_start = std::chrono::high_resolution_clock::now(); - for (const auto &[image, label] : train_dataloader) { + for (const auto &[image, label] : *train_dataloader) { auto new_image = std::make_shared(image->To(device)); auto new_label = std::make_shared(label->To(device)); - auto outputs = network.Forward({new_image}); + auto outputs = network->Forward({new_image}); optimizer.ZeroGrad(); - auto loss = loss_fn.Forward({outputs[0], new_label}); + auto loss = loss_fn->Forward({outputs[0], new_label}); loss[0]->Backward(); // Defer the loss D2H copy until after backward; reading it earlier would synchronize CUDA @@ -77,7 +132,8 @@ int main(int argc, char *argv[]) { auto loss_cpu = loss[0]->To(cpu_device); float current_loss = static_cast(loss_cpu.DataPtr())[0]; total_loss += current_loss; - if (train_idx % kNumItersOfOutputDuration == 0) { + // Distributed: gradients sync via DDP; loss is rank-local only (no per-step AllReduce). + if (is_main_rank && train_idx % kNumItersOfOutputDuration == 0) { LOG(ERROR) << "epoch: " << epoch << ", [" << train_idx * FLAGS_bs << "/" << train_dataset->Size() << "] " << " loss: " << current_loss; @@ -90,23 +146,28 @@ int main(int argc, char *argv[]) { const auto epoch_end = std::chrono::high_resolution_clock::now(); const double duration_us = std::chrono::duration(epoch_end - epoch_start).count(); - LOG(ERROR) << std::format("epoch {:2d}/{} | train loss {:.6f} | lr {:.2e} | ({:.2f} ms | {:.0f} samples/s)", - epoch, FLAGS_num_epoch - 1, total_loss / train_idx, FLAGS_lr, duration_us / 1e3f, - train_dataset->Size() / (duration_us / 1e6)); + if (is_main_rank) { + LOG(ERROR) << std::format( + "epoch {:2d}/{} | train loss {:.6f} | lr {:.2e} | ({:.2f} ms | {:.0f} samples/s)", epoch, + FLAGS_num_epoch - 1, total_loss / train_idx, FLAGS_lr, duration_us / 1e3f, + train_dataset->Size() / (duration_us / 1e6)); + } } // TODO(dcj): Add no_grad() context manager later. std::vector test_losses; int correct = 0; int total = 0; - for (const auto &[image, label] : test_dataloader) { + // Weighted loss accumulator for the distributed path (sharded eval + epoch-level AllReduce). + double local_loss_sum = 0.0; + for (const auto &[image, label] : *test_dataloader) { auto new_image = std::make_shared(image->To(device)); auto new_label = std::make_shared(label->To(device)); auto label_cpu = label->To(cpu_device); - auto outputs = network.Forward({new_image}); + auto outputs = network->Forward({new_image}); auto output_cpu = outputs[0]->To(cpu_device); - auto loss = loss_fn.Forward({outputs[0], new_label}); + auto loss = loss_fn->Forward({outputs[0], new_label}); auto loss_cpu = loss[0]->To(cpu_device); const int batch_size = output_cpu.Dims()[0]; @@ -119,11 +180,28 @@ int main(int argc, char *argv[]) { } } total += batch_size; - test_losses.push_back(static_cast(loss_cpu.DataPtr())[0]); + const float batch_loss = static_cast(loss_cpu.DataPtr())[0]; + test_losses.push_back(batch_loss); + local_loss_sum += static_cast(batch_loss) * batch_size; + } + if (distributed) { + // Each rank evaluated its own shard; reduce (loss_sum, correct, samples) once per epoch. + const float stats[3] + = {static_cast(local_loss_sum), static_cast(correct), static_cast(total)}; + auto stats_tensor + = std::make_shared(stats, std::vector{3}, DataType::kFLOAT32, device); + ddp_pg->AllReduce(stats_tensor, nn::parallel::function::ReduceOpType::kSum); + auto stats_cpu = stats_tensor->To(cpu_device); + const auto *reduced = static_cast(stats_cpu.DataPtr()); + if (is_main_rank) { + LOG(ERROR) << "Total: " << static_cast(reduced[2]) << ", Correct: " << static_cast(reduced[1]) + << ", Accuracy: " << reduced[1] / reduced[2] << ", AverageLoss: " << reduced[0] / reduced[2]; + } + } else { + const auto avg_loss = std::accumulate(test_losses.begin(), test_losses.end(), 0.0) / test_losses.size(); + LOG(ERROR) << "Total: " << total << ", Correct: " << correct + << ", Accuracy: " << static_cast(correct) / total << ", AverageLoss: " << avg_loss; } - const auto avg_loss = std::accumulate(test_losses.begin(), test_losses.end(), 0.0) / test_losses.size(); - LOG(ERROR) << "Total: " << total << ", Correct: " << correct - << ", Accuracy: " << static_cast(correct) / total << ", AverageLoss: " << avg_loss; gflags::ShutDownCommandLineFlags(); google::ShutdownGoogleLogging(); diff --git a/example/mnist/net.cc b/example/mnist/net.cc index 501fee7ef..c12ad7bd7 100644 --- a/example/mnist/net.cc +++ b/example/mnist/net.cc @@ -8,6 +8,7 @@ #include "infini_train/include/nn/modules/activations.h" #include "infini_train/include/nn/modules/container.h" +#include "infini_train/include/nn/modules/conv.h" #include "infini_train/include/nn/modules/linear.h" #include "infini_train/include/nn/modules/module.h" #include "infini_train/include/tensor.h" @@ -25,7 +26,32 @@ MNIST::MNIST() { std::vector> MNIST::Forward(const std::vector> &x) { CHECK_EQ(x.size(), 1); - auto x1 = (*modules_["sequential"])(x); + // Batches arrive as [B, 1, 28, 28] (NCHW); flatten trailing dims so the MLP sees [B, 784] exactly as before. + auto x0 = x[0]->Flatten(1); + auto x1 = (*modules_["sequential"])({x0}); auto x2 = (*modules_["linear2"])(x1); return x2; } + +MnistCnn::MnistCnn() { + modules_["conv1"] = std::make_shared(1, 16, 3); + modules_["relu1"] = std::make_shared(); + modules_["conv2"] = std::make_shared(16, 32, 3); + modules_["relu2"] = std::make_shared(); + // Shape math (kernel 3, stride 1, padding 0): 28x28 -> 26x26 (16ch) -> 24x24 (32ch), + // so the classifier sees 32 * 24 * 24 = 18432 features. + modules_["fc"] = std::make_shared(18432, 10); +} + +std::vector> +MnistCnn::Forward(const std::vector> &x) { + CHECK_EQ(x.size(), 1); + // Input is [B, 1, 28, 28] (NCHW); conv stack keeps NCHW, output is [B, 32, 24, 24]. + auto h1 = (*modules_["conv1"])(x); + auto h2 = (*modules_["relu1"])(h1); + auto h3 = (*modules_["conv2"])(h2); + auto h4 = (*modules_["relu2"])(h3); + auto flat = h4[0]->Flatten(1); + CHECK_EQ(flat->Dims()[1], 32 * 24 * 24) << "MnistCnn feature size mismatch: expected 18432"; + return (*modules_["fc"])({flat}); +} diff --git a/example/mnist/net.h b/example/mnist/net.h index 5f4cfa33b..1e06825d5 100644 --- a/example/mnist/net.h +++ b/example/mnist/net.h @@ -16,3 +16,13 @@ class MNIST : public infini_train::nn::Module { std::vector> Forward(const std::vector> &x) override; }; + +// Small CNN for MNIST: Conv2d(1,16,3)->ReLU->Conv2d(16,32,3)->ReLU->Flatten(1)->Linear(18432,10). +// Input arrives as [B, 1, 28, 28] (NCHW) from the dataloader pipeline; no reshape hacks needed. +class MnistCnn : public infini_train::nn::Module { +public: + MnistCnn(); + + std::vector> + Forward(const std::vector> &x) override; +}; diff --git a/infini_train/include/autograd/activations.h b/infini_train/include/autograd/activations.h index a63977263..809a7b67f 100644 --- a/infini_train/include/autograd/activations.h +++ b/infini_train/include/autograd/activations.h @@ -21,4 +21,16 @@ class Sigmoid : public Function { const std::vector> &output_tensors) override; std::vector> Backward(const std::vector> &grad_outputs) override; }; + +class ReLU : public Function { +public: + static constexpr char kType[] = "ReLUFunction"; + + ReLU() : 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; +}; } // namespace infini_train::autograd diff --git a/infini_train/include/autograd/conv.h b/infini_train/include/autograd/conv.h new file mode 100644 index 000000000..62367f618 --- /dev/null +++ b/infini_train/include/autograd/conv.h @@ -0,0 +1,34 @@ +#pragma once + +#include +#include +#include + +#include "infini_train/include/autograd/function.h" + +namespace infini_train { +class Tensor; +} + +namespace infini_train::autograd { + +class Conv2d : public Function { +public: + static constexpr char kType[] = "Conv2dFunction"; + + Conv2d() : Function(kType) {} + Conv2d(int64_t stride, int64_t padding) : Function(kType), stride_(stride), padding_(padding) {} + + 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: + int64_t stride_ = 1; + int64_t padding_ = 0; + bool has_bias_ = false; + std::vector input_dims_; + std::vector weight_dims_; +}; +} // namespace infini_train::autograd diff --git a/infini_train/include/nn/modules/activations.h b/infini_train/include/nn/modules/activations.h index deb029576..ded703c51 100644 --- a/infini_train/include/nn/modules/activations.h +++ b/infini_train/include/nn/modules/activations.h @@ -32,4 +32,11 @@ class SwiGLU : public CloneableModule { std::vector> Forward(const std::vector> &x) override; }; + +class ReLU : public CloneableModule { +public: + static constexpr char kType[] = "ReLU"; + ReLU() : CloneableModule(kType) {} + std::vector> Forward(const std::vector> &input_tensors) override; +}; } // namespace infini_train::nn diff --git a/infini_train/include/nn/modules/conv.h b/infini_train/include/nn/modules/conv.h new file mode 100644 index 000000000..deff4b9a1 --- /dev/null +++ b/infini_train/include/nn/modules/conv.h @@ -0,0 +1,35 @@ +#pragma once + +#include +#include +#include + +#include "infini_train/include/device.h" +#include "infini_train/include/nn/modules/module.h" + +namespace infini_train { +class Tensor; +class Device; +} // namespace infini_train + +namespace infini_train::nn { +class Conv2d : public CloneableModule { +public: + static constexpr char kType[] = "Conv2d"; + + static constexpr char kParamWeightName[] = "weight"; + static constexpr char kParamBiasName[] = "bias"; + + Conv2d(int64_t in_channels, int64_t out_channels, int64_t kernel_size, int64_t stride = 1, int64_t padding = 0, + bool bias = true, Device device = Device()); + std::vector> Forward(const std::vector> &input_tensors) override; + + bool has_bias() const { return bias_; } + +private: + void ResetParameters(); + int64_t stride_ = 1; + int64_t padding_ = 0; + bool bias_ = true; +}; +} // namespace infini_train::nn diff --git a/infini_train/src/autograd/activations.cc b/infini_train/src/autograd/activations.cc index bb8b8e5ea..b7aa1e6db 100644 --- a/infini_train/src/autograd/activations.cc +++ b/infini_train/src/autograd/activations.cc @@ -30,4 +30,30 @@ std::vector> Sigmoid::Backward(const std::vectorGetDevice().type(); return {Dispatcher::Instance().Call>({device, "SigmoidBackward"}, output, grad_output)}; } + +std::vector> ReLU::Forward(const std::vector> &input_tensors) { + CHECK_EQ(input_tensors.size(), 1); + const auto &input = input_tensors[0]; + + auto device = input->GetDevice().type(); + return {Dispatcher::Instance().Call>({device, "ReLUForward"}, input)}; +} + +void ReLU::SetupContext(const std::vector> &input_tensors, + const std::vector> &) { + CHECK_EQ(input_tensors.size(), 1); + const auto &input = input_tensors[0]; + ctx_.SaveForBackward({input}); +} + +std::vector> ReLU::Backward(const std::vector> &grad_outputs) { + auto saved_tensors = ctx_.GetSavedTensors(); + CHECK_EQ(saved_tensors.size(), 1); + const auto &input = saved_tensors[0]; + CHECK_EQ(grad_outputs.size(), 1); + const auto &grad_output = grad_outputs[0]; + + auto device = input->GetDevice().type(); + return {Dispatcher::Instance().Call>({device, "ReLUBackward"}, input, grad_output)}; +} } // namespace infini_train::autograd diff --git a/infini_train/src/autograd/conv.cc b/infini_train/src/autograd/conv.cc new file mode 100644 index 000000000..e1c526d82 --- /dev/null +++ b/infini_train/src/autograd/conv.cc @@ -0,0 +1,73 @@ +#include "infini_train/include/autograd/conv.h" + +#include "glog/logging.h" + +#include "infini_train/include/dispatcher.h" +#include "infini_train/include/tensor.h" + +namespace infini_train::autograd { +std::vector> Conv2d::Forward(const std::vector> &input_tensors) { + CHECK_GE(input_tensors.size(), 2); + CHECK_LE(input_tensors.size(), 3); + const auto &input = input_tensors[0]; + const auto &weight = input_tensors[1]; + const std::shared_ptr &bias = input_tensors.size() == 3 ? input_tensors[2] : nullptr; + + auto device = input->GetDevice().type(); + return {Dispatcher::Instance().Call>({device, "Conv2dForward"}, input, weight, bias, + stride_, padding_)}; +} + +void Conv2d::SetupContext(const std::vector> &input_tensors, + const std::vector> &) { + const auto &input = input_tensors[0]; + const auto &weight = input_tensors[1]; + bool need_input = ctx_.needs_input_grad().size() > 0 && ctx_.needs_input_grad()[0]; + bool need_weight = ctx_.needs_input_grad().size() > 1 && ctx_.needs_input_grad()[1]; + + // grad_input needs weight, grad_weight needs input + ctx_.SaveForBackward({need_weight ? input : nullptr, need_input ? weight : nullptr}); + + has_bias_ = input_tensors.size() == 3; + input_dims_ = input->Dims(); + weight_dims_ = weight->Dims(); +} + +std::vector> Conv2d::Backward(const std::vector> &grad_outputs) { + auto saved_tensors = ctx_.GetSavedTensors(); + CHECK_EQ(saved_tensors.size(), 2); + const auto &input = saved_tensors[0]; + const auto &weight = saved_tensors[1]; + CHECK_EQ(grad_outputs.size(), 1); + const auto &grad_output = grad_outputs[0]; + + CHECK(!ctx_.needs_input_grad().empty()) << "needs_input_grad not populated in Conv2d::Backward"; + bool need_grad_input = ctx_.needs_input_grad()[0]; + bool need_grad_weight = ctx_.needs_input_grad().size() > 1 && ctx_.needs_input_grad()[1]; + bool need_grad_bias = has_bias_ && ctx_.needs_input_grad().size() > 2 && ctx_.needs_input_grad()[2]; + + auto device = grad_output->GetDevice().type(); + + std::shared_ptr grad_input = nullptr; + std::shared_ptr grad_weight = nullptr; + std::shared_ptr grad_bias = nullptr; + + if (need_grad_input) { + grad_input = Dispatcher::Instance().Call>({device, "Conv2dBackwardInput"}, weight, + grad_output, stride_, padding_, input_dims_); + } + if (need_grad_weight) { + grad_weight = Dispatcher::Instance().Call>( + {device, "Conv2dBackwardWeight"}, input, grad_output, stride_, padding_, weight_dims_); + } + if (need_grad_bias) { + grad_bias = Dispatcher::Instance().Call>({device, "Conv2dBackwardBias"}, grad_output); + } + + if (has_bias_) { + return {grad_input, grad_weight, grad_bias}; + } else { + return {grad_input, grad_weight}; + } +} +} // namespace infini_train::autograd diff --git a/infini_train/src/dataloader.cc b/infini_train/src/dataloader.cc index 198536779..f546f80ff 100644 --- a/infini_train/src/dataloader.cc +++ b/infini_train/src/dataloader.cc @@ -22,15 +22,16 @@ size_t CheckedCeilDiv(size_t numerator, size_t denominator) { // TODO(dcj): Use official stack implementation later. std::shared_ptr Stack(const std::vector> &tensors) { CHECK(!tensors.empty()) << "Cannot stack an empty batch. Check DataLoader iterator end handling."; - const int batch_size = tensors.size(); - const auto &dims = tensors[0]->Dims(); - const int stacked_dim = std::accumulate(dims.begin(), dims.end(), 1, std::multiplies()); - auto stacked_tensor = std::make_shared(std::vector{batch_size, stacked_dim}, tensors[0]->Dtype()); + const int64_t batch_size = static_cast(tensors.size()); + const auto &sample_dims = tensors[0]->Dims(); for (const auto &tensor : tensors) { - CHECK_EQ(static_cast(tensors[0]->Dtype()), static_cast(tensor->Dtype())); - const auto &dims = tensor->Dims(); - CHECK_EQ(stacked_dim, std::accumulate(dims.begin(), dims.end(), 1, std::multiplies())); + CHECK_EQ(static_cast(tensors[0]->Dtype()), static_cast(tensor->Dtype())) + << "Cannot stack tensors with different dtypes."; + CHECK(tensor->Dims() == sample_dims) << "Cannot stack tensors with different shapes."; } + std::vector stacked_dims = {batch_size}; + stacked_dims.insert(stacked_dims.end(), sample_dims.begin(), sample_dims.end()); + auto stacked_tensor = std::make_shared(stacked_dims, tensors[0]->Dtype()); size_t offset = 0; for (const auto &tensor : tensors) { diff --git a/infini_train/src/kernels/common/conv.h b/infini_train/src/kernels/common/conv.h new file mode 100644 index 000000000..e3fadc78c --- /dev/null +++ b/infini_train/src/kernels/common/conv.h @@ -0,0 +1,69 @@ +#pragma once + +#include +#include +#include + +#include "glog/logging.h" + +#include "infini_train/include/tensor.h" + +namespace infini_train::kernels { + +// Shared shape arithmetic for Conv2d (CPU/CUDA must not redefine it). +// Layout: input NCHW, weight OIHW. Kh/Kw are read independently from +// weight dims [2]/[3]; no Kh == Kw assumption. +struct Conv2dMeta { + int64_t batch = 0; + int64_t in_channels = 0; + int64_t out_channels = 0; + int64_t input_h = 0; + int64_t input_w = 0; + int64_t kernel_h = 0; + int64_t kernel_w = 0; + int64_t stride = 1; + int64_t padding = 0; + int64_t output_h = 0; + int64_t output_w = 0; + int64_t patches = 0; // Hout * Wout + int64_t kernel_elems = 0; // Cin * Kh * Kw +}; + +inline Conv2dMeta MakeConv2dMeta(const std::shared_ptr &input, const std::shared_ptr &weight, + int64_t stride, int64_t padding) { + CHECK(input != nullptr) << "Conv2d input must not be null"; + CHECK(weight != nullptr) << "Conv2d weight must not be null"; + CHECK_GT(stride, 0) << "Conv2d stride must be positive"; + CHECK_GE(padding, 0) << "Conv2d padding must be non-negative"; + + const auto &in_dims = input->Dims(); + const auto &w_dims = weight->Dims(); + CHECK_EQ(in_dims.size(), 4) << "Conv2d input must be NCHW"; + CHECK_EQ(w_dims.size(), 4) << "Conv2d weight must be OIHW"; + CHECK_EQ(in_dims[1], w_dims[1]) << "Conv2d input channels must match weight in-channels"; + + Conv2dMeta meta; + meta.batch = in_dims[0]; + meta.in_channels = in_dims[1]; + meta.input_h = in_dims[2]; + meta.input_w = in_dims[3]; + meta.out_channels = w_dims[0]; + meta.kernel_h = w_dims[2]; + meta.kernel_w = w_dims[3]; + meta.stride = stride; + meta.padding = padding; + + CHECK_GE(meta.input_h + 2 * padding, meta.kernel_h) << "Conv2d kernel taller than padded input"; + CHECK_GE(meta.input_w + 2 * padding, meta.kernel_w) << "Conv2d kernel wider than padded input"; + meta.output_h = (meta.input_h + 2 * padding - meta.kernel_h) / stride + 1; + meta.output_w = (meta.input_w + 2 * padding - meta.kernel_w) / stride + 1; + meta.patches = meta.output_h * meta.output_w; + meta.kernel_elems = meta.in_channels * meta.kernel_h * meta.kernel_w; + return meta; +} + +inline std::vector Conv2dOutputDims(const Conv2dMeta &meta) { + return {meta.batch, meta.out_channels, meta.output_h, meta.output_w}; +} + +} // namespace infini_train::kernels diff --git a/infini_train/src/kernels/cpu/conv.cc b/infini_train/src/kernels/cpu/conv.cc new file mode 100644 index 000000000..eddb11767 --- /dev/null +++ b/infini_train/src/kernels/cpu/conv.cc @@ -0,0 +1,240 @@ +#include +#include +#include + +#include "glog/logging.h" + +#include "infini_train/include/dispatcher.h" +#include "infini_train/include/tensor.h" +#include "infini_train/src/kernels/common/conv.h" + +namespace infini_train::kernels::cpu { +namespace { + +inline int64_t NCHWOffset(int64_t n, int64_t c, int64_t h, int64_t w, int64_t channels, int64_t height, int64_t width) { + return ((n * channels + c) * height + h) * width + w; +} + +inline int64_t OIHWOffset(int64_t o, int64_t i, int64_t kh, int64_t kw, int64_t in_channels, int64_t kernel_h, + int64_t kernel_w) { + return ((o * in_channels + i) * kernel_h + kh) * kernel_w + kw; +} + +void CheckFP32(const std::shared_ptr &tensor, const char *name) { + CHECK(tensor->Dtype() == DataType::kFLOAT32) << name << " currently supports float32 only"; +} + +} // namespace + +// Direct cross-correlation (no kernel flip), scalar stride/padding. +std::shared_ptr Conv2dForward(const std::shared_ptr &input, const std::shared_ptr &weight, + const std::shared_ptr &bias, int64_t stride, int64_t padding) { + CheckFP32(input, "CPU Conv2dForward"); + CheckFP32(weight, "CPU Conv2dForward"); + if (bias) { + CheckFP32(bias, "CPU Conv2dForward"); + } + + const Conv2dMeta meta = MakeConv2dMeta(input, weight, stride, padding); + if (bias) { + CHECK_EQ(bias->Dims().size(), 1); + CHECK_EQ(bias->Dims()[0], meta.out_channels); + } + + auto output = std::make_shared(Conv2dOutputDims(meta), DataType::kFLOAT32); + const float *input_ptr = static_cast(input->DataPtr()); + const float *weight_ptr = static_cast(weight->DataPtr()); + const float *bias_ptr = bias ? static_cast(bias->DataPtr()) : nullptr; + float *output_ptr = static_cast(output->DataPtr()); + + for (int64_t n = 0; n < meta.batch; ++n) { + for (int64_t oc = 0; oc < meta.out_channels; ++oc) { + for (int64_t oh = 0; oh < meta.output_h; ++oh) { + for (int64_t ow = 0; ow < meta.output_w; ++ow) { + float acc = bias_ptr ? bias_ptr[oc] : 0.0f; + for (int64_t ic = 0; ic < meta.in_channels; ++ic) { + for (int64_t kh = 0; kh < meta.kernel_h; ++kh) { + const int64_t ih = oh * stride + kh - padding; + if (ih < 0 || ih >= meta.input_h) { + continue; + } + for (int64_t kw = 0; kw < meta.kernel_w; ++kw) { + const int64_t iw = ow * stride + kw - padding; + if (iw < 0 || iw >= meta.input_w) { + continue; + } + acc += input_ptr[NCHWOffset(n, ic, ih, iw, meta.in_channels, meta.input_h, + meta.input_w)] + * weight_ptr[OIHWOffset(oc, ic, kh, kw, meta.in_channels, meta.kernel_h, + meta.kernel_w)]; + } + } + } + output_ptr[NCHWOffset(n, oc, oh, ow, meta.out_channels, meta.output_h, meta.output_w)] = acc; + } + } + } + } + return output; +} + +std::shared_ptr Conv2dBackwardInput(const std::shared_ptr &weight, + const std::shared_ptr &grad_output, int64_t stride, int64_t padding, + const std::vector &input_dims) { + CheckFP32(weight, "CPU Conv2dBackwardInput"); + CheckFP32(grad_output, "CPU Conv2dBackwardInput"); + CHECK_EQ(input_dims.size(), 4); + + const auto &go_dims = grad_output->Dims(); + CHECK_EQ(go_dims.size(), 4); + const int64_t batch = input_dims[0]; + const int64_t in_channels = input_dims[1]; + const int64_t input_h = input_dims[2]; + const int64_t input_w = input_dims[3]; + const int64_t out_channels = go_dims[1]; + const int64_t output_h = go_dims[2]; + const int64_t output_w = go_dims[3]; + + const auto &w_dims = weight->Dims(); + CHECK_EQ(w_dims.size(), 4); + CHECK_EQ(w_dims[0], out_channels); + CHECK_EQ(w_dims[1], in_channels); + const int64_t kernel_h = w_dims[2]; + const int64_t kernel_w = w_dims[3]; + CHECK_EQ(output_h, (input_h + 2 * padding - kernel_h) / stride + 1); + CHECK_EQ(output_w, (input_w + 2 * padding - kernel_w) / stride + 1); + + auto grad_input = std::make_shared(input_dims, DataType::kFLOAT32); + grad_input->Fill(0.0f); + const float *weight_ptr = static_cast(weight->DataPtr()); + const float *grad_output_ptr = static_cast(grad_output->DataPtr()); + float *grad_input_ptr = static_cast(grad_input->DataPtr()); + + for (int64_t n = 0; n < batch; ++n) { + for (int64_t oc = 0; oc < out_channels; ++oc) { + for (int64_t oh = 0; oh < output_h; ++oh) { + for (int64_t ow = 0; ow < output_w; ++ow) { + const float go = grad_output_ptr[NCHWOffset(n, oc, oh, ow, out_channels, output_h, output_w)]; + for (int64_t ic = 0; ic < in_channels; ++ic) { + for (int64_t kh = 0; kh < kernel_h; ++kh) { + const int64_t ih = oh * stride + kh - padding; + if (ih < 0 || ih >= input_h) { + continue; + } + for (int64_t kw = 0; kw < kernel_w; ++kw) { + const int64_t iw = ow * stride + kw - padding; + if (iw < 0 || iw >= input_w) { + continue; + } + grad_input_ptr[NCHWOffset(n, ic, ih, iw, in_channels, input_h, input_w)] + += go * weight_ptr[OIHWOffset(oc, ic, kh, kw, in_channels, kernel_h, kernel_w)]; + } + } + } + } + } + } + } + return grad_input; +} + +std::shared_ptr Conv2dBackwardWeight(const std::shared_ptr &input, + const std::shared_ptr &grad_output, int64_t stride, + int64_t padding, const std::vector &weight_dims) { + CheckFP32(input, "CPU Conv2dBackwardWeight"); + CheckFP32(grad_output, "CPU Conv2dBackwardWeight"); + CHECK_EQ(weight_dims.size(), 4); + + const auto &in_dims = input->Dims(); + CHECK_EQ(in_dims.size(), 4); + const int64_t batch = in_dims[0]; + const int64_t in_channels = in_dims[1]; + const int64_t input_h = in_dims[2]; + const int64_t input_w = in_dims[3]; + + const auto &go_dims = grad_output->Dims(); + CHECK_EQ(go_dims.size(), 4); + CHECK_EQ(go_dims[0], batch); + const int64_t out_channels = go_dims[1]; + const int64_t output_h = go_dims[2]; + const int64_t output_w = go_dims[3]; + + CHECK_EQ(weight_dims[0], out_channels); + CHECK_EQ(weight_dims[1], in_channels); + const int64_t kernel_h = weight_dims[2]; + const int64_t kernel_w = weight_dims[3]; + CHECK_EQ(output_h, (input_h + 2 * padding - kernel_h) / stride + 1); + CHECK_EQ(output_w, (input_w + 2 * padding - kernel_w) / stride + 1); + + auto grad_weight = std::make_shared(weight_dims, DataType::kFLOAT32); + grad_weight->Fill(0.0f); + const float *input_ptr = static_cast(input->DataPtr()); + const float *grad_output_ptr = static_cast(grad_output->DataPtr()); + float *grad_weight_ptr = static_cast(grad_weight->DataPtr()); + + for (int64_t oc = 0; oc < out_channels; ++oc) { + for (int64_t ic = 0; ic < in_channels; ++ic) { + for (int64_t kh = 0; kh < kernel_h; ++kh) { + for (int64_t kw = 0; kw < kernel_w; ++kw) { + float acc = 0.0f; + for (int64_t n = 0; n < batch; ++n) { + for (int64_t oh = 0; oh < output_h; ++oh) { + const int64_t ih = oh * stride + kh - padding; + if (ih < 0 || ih >= input_h) { + continue; + } + for (int64_t ow = 0; ow < output_w; ++ow) { + const int64_t iw = ow * stride + kw - padding; + if (iw < 0 || iw >= input_w) { + continue; + } + acc += input_ptr[NCHWOffset(n, ic, ih, iw, in_channels, input_h, input_w)] + * grad_output_ptr[NCHWOffset(n, oc, oh, ow, out_channels, output_h, output_w)]; + } + } + } + grad_weight_ptr[OIHWOffset(oc, ic, kh, kw, in_channels, kernel_h, kernel_w)] = acc; + } + } + } + } + return grad_weight; +} + +std::shared_ptr Conv2dBackwardBias(const std::shared_ptr &grad_output) { + CheckFP32(grad_output, "CPU Conv2dBackwardBias"); + const auto &go_dims = grad_output->Dims(); + CHECK_EQ(go_dims.size(), 4); + const int64_t batch = go_dims[0]; + const int64_t out_channels = go_dims[1]; + const int64_t output_h = go_dims[2]; + const int64_t output_w = go_dims[3]; + + auto grad_bias = std::make_shared(std::vector{out_channels}, DataType::kFLOAT32); + grad_bias->Fill(0.0f); + const float *grad_output_ptr = static_cast(grad_output->DataPtr()); + float *grad_bias_ptr = static_cast(grad_bias->DataPtr()); + + for (int64_t n = 0; n < batch; ++n) { + for (int64_t oc = 0; oc < out_channels; ++oc) { + for (int64_t oh = 0; oh < output_h; ++oh) { + for (int64_t ow = 0; ow < output_w; ++ow) { + grad_bias_ptr[oc] += grad_output_ptr[NCHWOffset(n, oc, oh, ow, out_channels, output_h, output_w)]; + } + } + } + } + return grad_bias; +} + +} // namespace infini_train::kernels::cpu + +#define REGISTER_CPU_CONV_KERNEL(kernel_name) \ + REGISTER_KERNEL(infini_train::Device::DeviceType::kCPU, kernel_name, infini_train::kernels::cpu::kernel_name) + +REGISTER_CPU_CONV_KERNEL(Conv2dForward) +REGISTER_CPU_CONV_KERNEL(Conv2dBackwardInput) +REGISTER_CPU_CONV_KERNEL(Conv2dBackwardWeight) +REGISTER_CPU_CONV_KERNEL(Conv2dBackwardBias) + +#undef REGISTER_CPU_CONV_KERNEL diff --git a/infini_train/src/kernels/cpu/relu.cc b/infini_train/src/kernels/cpu/relu.cc new file mode 100644 index 000000000..63d31cff6 --- /dev/null +++ b/infini_train/src/kernels/cpu/relu.cc @@ -0,0 +1,43 @@ +#include + +#include "glog/logging.h" + +#include "infini_train/include/dispatcher.h" +#include "infini_train/include/tensor.h" + +namespace infini_train::kernels::cpu { +std::shared_ptr ReLUForward(const std::shared_ptr &input) { + CHECK(input->Dtype() == DataType::kFLOAT32) << "CPU ReLUForward currently supports float32 only"; + auto output = std::make_shared(input->Dims(), DataType::kFLOAT32); + const float *input_ptr = static_cast(input->DataPtr()); + float *output_ptr = static_cast(output->DataPtr()); + + const int64_t numel = input->NumElements(); + for (int64_t idx = 0; idx < numel; ++idx) { output_ptr[idx] = input_ptr[idx] > 0.0f ? input_ptr[idx] : 0.0f; } + + return output; +} + +std::shared_ptr ReLUBackward(const std::shared_ptr &input, const std::shared_ptr &grad_output) { + CHECK(input->Dtype() == DataType::kFLOAT32) << "CPU ReLUBackward currently supports float32 only"; + CHECK(grad_output->Dtype() == DataType::kFLOAT32) << "CPU ReLUBackward currently supports float32 only"; + auto grad_input = std::make_shared(input->Dims(), DataType::kFLOAT32); + const float *input_ptr = static_cast(input->DataPtr()); + const float *grad_output_ptr = static_cast(grad_output->DataPtr()); + float *grad_input_ptr = static_cast(grad_input->DataPtr()); + + const int64_t numel = input->NumElements(); + for (int64_t idx = 0; idx < numel; ++idx) { + grad_input_ptr[idx] = input_ptr[idx] > 0.0f ? grad_output_ptr[idx] : 0.0f; + } + return grad_input; +} +} // namespace infini_train::kernels::cpu + +#define REGISTER_CPU_RELU_KERNEL(kernel_name) \ + REGISTER_KERNEL(infini_train::Device::DeviceType::kCPU, kernel_name, infini_train::kernels::cpu::kernel_name) + +REGISTER_CPU_RELU_KERNEL(ReLUForward) +REGISTER_CPU_RELU_KERNEL(ReLUBackward) + +#undef REGISTER_CPU_RELU_KERNEL diff --git a/infini_train/src/kernels/cuda/conv.cu b/infini_train/src/kernels/cuda/conv.cu new file mode 100644 index 000000000..f21f7a2b1 --- /dev/null +++ b/infini_train/src/kernels/cuda/conv.cu @@ -0,0 +1,495 @@ +#include +#include +#include + +#include + +#include "glog/logging.h" + +#include "infini_train/include/core/runtime/device_guard.h" +#include "infini_train/include/dispatcher.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_runtime_common.h" +#include "infini_train/src/kernels/common/conv.h" +#include "infini_train/src/kernels/cuda/common/gemm.cuh" + +namespace infini_train::kernels::cuda { +namespace { + +// Column layout: per image [K, P] row-major with K = Cin * Kh * Kw (kernel row +// k = (ic * Kh + kh) * Kw + kw) and P = Hout * Wout (patch column +// p = oh * Wout + ow). Out-of-bounds taps (padding) read as zero. +// Batched output is [N, K, P]. All shape arithmetic comes from Conv2dMeta. +__global__ void Im2colKernel(const float *input, float *columns, Conv2dMeta meta) { + const int64_t total = meta.batch * meta.kernel_elems * meta.patches; + const int64_t idx = static_cast(blockIdx.x) * blockDim.x + threadIdx.x; + if (idx >= total) { + return; + } + const int64_t p = idx % meta.patches; + const int64_t k = (idx / meta.patches) % meta.kernel_elems; + const int64_t n = idx / (meta.patches * meta.kernel_elems); + + const int64_t kw = k % meta.kernel_w; + const int64_t kh = (k / meta.kernel_w) % meta.kernel_h; + const int64_t ic = k / (meta.kernel_w * meta.kernel_h); + + const int64_t oh = p / meta.output_w; + const int64_t ow = p % meta.output_w; + const int64_t ih = oh * meta.stride + kh - meta.padding; + const int64_t iw = ow * meta.stride + kw - meta.padding; + + float value = 0.0f; + if (ih >= 0 && ih < meta.input_h && iw >= 0 && iw < meta.input_w) { + value = input[((n * meta.in_channels + ic) * meta.input_h + ih) * meta.input_w + iw]; + } + columns[idx] = value; +} + +// Broadcast per-channel bias over an [N, Cout, P] row-major buffer: +// out[(n * Cout + oc) * P + p] = bias[oc]. +__global__ void ConvBiasCopyKernel(float *output, const float *bias, int64_t out_channels, int64_t patches, + int64_t total) { + const int64_t idx = static_cast(blockIdx.x) * blockDim.x + threadIdx.x; + if (idx >= total) { + return; + } + output[idx] = bias[(idx / patches) % out_channels]; +} + +// dB tiles are reduced over N*P on-device: grad_output holds row-major +// [N, Cout, P] (P = Hout * Wout); channel oc's taps are strided as +// [(n * Cout + oc) * P + p]. One block per channel (same pattern as +// cuda::LinearBackwardBias ReduceRowsKernel), each thread striding over +// rows = N * P with a cub BlockReduce sum. +template +__global__ void ConvBackwardBiasKernel(const float *__restrict__ grad_output, float *__restrict__ grad_bias, + int64_t rows, int64_t out_channels, int64_t patches) { + using BlockReduce = cub::BlockReduce; + __shared__ typename BlockReduce::TempStorage temp_storage; + + const int64_t oc = blockIdx.x; + float sum = 0.0f; + for (int64_t r = threadIdx.x; r < rows; r += blockDim.x) { + const int64_t n = r / patches; + const int64_t p = r % patches; + sum += grad_output[(n * out_channels + oc) * patches + p]; + } + + const float reduced = BlockReduce(temp_storage).Sum(sum); + if (threadIdx.x == 0) { + grad_bias[oc] = reduced; + } +} + +// dW tiles are reduced over the batch on-device: per_image holds one +// row-major [Cout, K] tile per image (tile = oc * K + k); grad_weight is +// the flat [Cout * K] sum (= OIHW element count, same memory order). +__global__ void ConvBackwardWeightSumKernel(const float *per_image, float *grad_weight, int64_t batch, int64_t elems) { + const int64_t idx = static_cast(blockIdx.x) * blockDim.x + threadIdx.x; + if (idx >= elems) { + return; + } + float sum = 0.0f; + for (int64_t n = 0; n < batch; ++n) { sum += per_image[n * elems + idx]; } + grad_weight[idx] = sum; +} + +} // namespace + +// dXcol tiles hold one row-major [K, P] tile per image (tile = k * P + p, +// K = Cin * Kh * Kw matches Im2colKernel's row order, P = Hout * Wout); +// Col2imKernel gathers each input element's covering patches into grad_input. +__global__ void Col2imKernel(const float *columns, float *grad_input, Conv2dMeta meta) { + const int64_t hw = meta.input_h * meta.input_w; + const int64_t chw = meta.in_channels * hw; + const int64_t total = meta.batch * chw; + const int64_t idx = static_cast(blockIdx.x) * blockDim.x + threadIdx.x; + if (idx >= total) { + return; + } + const int64_t w = idx % meta.input_w; + const int64_t h = (idx / meta.input_w) % meta.input_h; + const int64_t c = (idx / hw) % meta.in_channels; + const int64_t n = idx / chw; + + // Gather: input tap (h, w) is covered by patches whose receptive field + // contains it, i.e. h == oh * stride + kh - padding (same for w). + // Padding taps contribute 0 (skipped via the divisibility/bounds checks). + float sum = 0.0f; + for (int64_t kh = 0; kh < meta.kernel_h; ++kh) { + const int64_t t = h + meta.padding - kh; + if (t < 0 || t % meta.stride != 0) { + continue; + } + const int64_t oh = t / meta.stride; + if (oh >= meta.output_h) { + continue; + } + for (int64_t kw = 0; kw < meta.kernel_w; ++kw) { + const int64_t u = w + meta.padding - kw; + if (u < 0 || u % meta.stride != 0) { + continue; + } + const int64_t ow = u / meta.stride; + if (ow >= meta.output_w) { + continue; + } + const int64_t k = (c * meta.kernel_h + kh) * meta.kernel_w + kw; + const int64_t p = oh * meta.output_w + ow; + sum += columns[(n * meta.kernel_elems + k) * meta.patches + p]; + } + } + grad_input[idx] = sum; +} + +// NOTE: meta is taken by value (not const-ref) on purpose. Dispatcher::Call +// deduces argument types by value, and a trivially-copyable struct of this +// size is passed on the stack by value but as a pointer when bound to const&, +// so const-ref here would read garbage through the type-erased call. +std::shared_ptr Im2colForward(const std::shared_ptr &input, Conv2dMeta meta) { + CHECK(input->Dtype() == DataType::kFLOAT32) << "CUDA Im2colForward currently supports float32 only"; + const auto &dims = input->Dims(); + CHECK_EQ(dims.size(), 4) << "Im2colForward input must be NCHW"; + CHECK_EQ(dims[0], meta.batch); + CHECK_EQ(dims[1], meta.in_channels); + CHECK_EQ(dims[2], meta.input_h); + CHECK_EQ(dims[3], meta.input_w); + + auto columns = std::make_shared(std::vector{meta.batch, meta.kernel_elems, meta.patches}, + DataType::kFLOAT32, input->GetDevice()); + const int64_t total = meta.batch * meta.kernel_elems * meta.patches; + if (total == 0) { + return columns; + } + + const int threads_per_block = 256; + const int num_blocks = static_cast((total + threads_per_block - 1) / threads_per_block); + auto device = input->GetDevice(); + const auto &cuda_stream = dynamic_cast( + infini_train::core::GetDeviceGuardImpl(device.type())->GetStream(device)) + ->cuda_stream(); + Im2colKernel<<>>(static_cast(input->DataPtr()), + static_cast(columns->DataPtr()), meta); + return columns; +} + +// Forward via im2col + framework GEMM (FP32 only). Weight OIHW is viewed as a +// row-major [Cout, K] matrix (K = Cin * Kh * Kw matches Im2colKernel's row +// order); per image Y[Cout, P] = W * Xcol[K, P]. +// NOTE: this entry point mirrors the CPU signature (shared_ptr const&, int64 +// by value) so the dispatcher resolves the identical "Conv2dForward" overload +// on CUDA. +std::shared_ptr Conv2dForward(const std::shared_ptr &input, const std::shared_ptr &weight, + const std::shared_ptr &bias, int64_t stride, int64_t padding) { + CHECK(input->Dtype() == DataType::kFLOAT32) << "CUDA Conv2dForward currently supports float32 only"; + CHECK(weight->Dtype() == DataType::kFLOAT32) << "CUDA Conv2dForward currently supports float32 only"; + if (bias) { + CHECK(bias->Dtype() == DataType::kFLOAT32) << "CUDA Conv2dForward currently supports float32 only"; + } + + const Conv2dMeta meta = MakeConv2dMeta(input, weight, stride, padding); + if (bias) { + CHECK_EQ(bias->Dims().size(), 1); + CHECK_EQ(bias->Dims()[0], meta.out_channels); + } + + auto device = input->GetDevice(); + auto output = std::make_shared(Conv2dOutputDims(meta), DataType::kFLOAT32, device); + if (meta.batch == 0) { + return output; + } + + auto columns = Im2colForward(input, meta); + + const auto &cuda_stream = dynamic_cast( + infini_train::core::GetDeviceGuardImpl(device.type())->GetStream(device)) + ->cuda_stream(); + + // Bias is pre-staged into the output and accumulated with beta=1 (same + // pattern as cuda::LinearForward); without bias beta=0 overwrites output. + float beta = 0.0f; + if (bias) { + const int64_t total = meta.batch * meta.out_channels * meta.patches; + const int threads_per_block = 256; + const int num_blocks = static_cast((total + threads_per_block - 1) / threads_per_block); + ConvBiasCopyKernel<<>>( + static_cast(output->DataPtr()), static_cast(bias->DataPtr()), meta.out_channels, + meta.patches, total); + beta = 1.0f; + } + + // cuBLAS is column-major: row-major Y[Cout, P] = W[Cout, K] * X[K, P] is + // computed as Y^T[P, Cout] = X^T[P, K] * W^T[K, Cout]. + // C = Y^T[P, Cout], A = X^T[P, K], B = W^T[K, Cout]. + // One strided-batched GEMM over the batch: weight is shared across images + // (stride_b = 0); batch_count == 1 uses zero strides per Gemm convention. + const int batch_count = static_cast(meta.batch); + const bool batched = meta.batch > 1; + Dispatcher::Instance().Call({device.type(), "Gemm"}, device, + GemmParams{ + .trans_a = GemmTranspose::kNoTranspose, + .trans_b = GemmTranspose::kNoTranspose, + .m = static_cast(meta.patches), + .n = static_cast(meta.out_channels), + .k = static_cast(meta.kernel_elems), + .A = columns->DataPtr(), + .lda = static_cast(meta.patches), + .B = weight->DataPtr(), + .ldb = static_cast(meta.kernel_elems), + .C = output->DataPtr(), + .ldc = static_cast(meta.patches), + .alpha = 1.0f, + .beta = beta, + .batch_count = batch_count, + .stride_a = batched ? meta.kernel_elems * meta.patches : 0, + .stride_b = 0, + .stride_c = batched ? meta.out_channels * meta.patches : 0, + .input_dtype = DataType::kFLOAT32, + .output_dtype = DataType::kFLOAT32, + }); + return output; +} + +// Weight grad via im2col + framework GEMM (FP32 only). Per image, +// row-major dW_n[Cout, K] = dY_n[Cout, P] * Xcol_n[K, P]^T (K = Cin*Kh*Kw, +// P = Hout*Wout matches Im2colKernel's row order), summed over N. +// NOTE: this entry point mirrors the CPU signature (shared_ptr const&, +// int64 by value, vector const&) so the dispatcher resolves the +// identical "Conv2dBackwardWeight" overload on CUDA. +std::shared_ptr Conv2dBackwardWeight(const std::shared_ptr &input, + const std::shared_ptr &grad_output, int64_t stride, + int64_t padding, const std::vector &weight_dims) { + CHECK(input->Dtype() == DataType::kFLOAT32) << "CUDA Conv2dBackwardWeight currently supports float32 only"; + CHECK(grad_output->Dtype() == DataType::kFLOAT32) << "CUDA Conv2dBackwardWeight currently supports float32 only"; + CHECK_EQ(weight_dims.size(), 4); + + const auto &in_dims = input->Dims(); + CHECK_EQ(in_dims.size(), 4); + const auto &go_dims = grad_output->Dims(); + CHECK_EQ(go_dims.size(), 4); + CHECK_EQ(go_dims[0], in_dims[0]); + CHECK_EQ(weight_dims[0], go_dims[1]); + CHECK_EQ(weight_dims[1], in_dims[1]); + CHECK_EQ(go_dims[2], (in_dims[2] + 2 * padding - weight_dims[2]) / stride + 1); + CHECK_EQ(go_dims[3], (in_dims[3] + 2 * padding - weight_dims[3]) / stride + 1); + + Conv2dMeta meta; + meta.batch = in_dims[0]; + meta.in_channels = in_dims[1]; + meta.input_h = in_dims[2]; + meta.input_w = in_dims[3]; + meta.out_channels = go_dims[1]; + meta.kernel_h = weight_dims[2]; + meta.kernel_w = weight_dims[3]; + meta.stride = stride; + meta.padding = padding; + meta.output_h = go_dims[2]; + meta.output_w = go_dims[3]; + meta.patches = meta.output_h * meta.output_w; + meta.kernel_elems = meta.in_channels * meta.kernel_h * meta.kernel_w; + + auto device = input->GetDevice(); + auto grad_weight = std::make_shared(weight_dims, DataType::kFLOAT32, device); + if (meta.batch == 0) { + return grad_weight; + } + + auto columns = Im2colForward(input, meta); + + const auto &cuda_stream = dynamic_cast( + infini_train::core::GetDeviceGuardImpl(device.type())->GetStream(device)) + ->cuda_stream(); + + // cuBLAS is column-major: row-major dW_n[Cout, K] = dY_n[Cout, P] * + // Xcol_n[K, P]^T is computed as dW_n^T[K, Cout] = op(Xcol)[K, P] * + // dY_n^T[P, Cout], with Xcol viewed as the column-major [P, K] + // transpose (trans_a=T) and dY viewed as column-major [P, Cout] + // (trans_b=N). + // C = dW^T[K, Cout], A = Xcol[P, K]^T, B = dY^T[P, Cout]. + // Strides are per-image tiles: Xcol K*P, dY Cout*P, dW Cout*K; + // batch_count == 1 uses zero strides per Gemm convention. + const int batch_count = static_cast(meta.batch); + const bool batched = meta.batch > 1; + void *gemm_c = grad_weight->DataPtr(); + long long stride_c = 0; + std::shared_ptr per_image; + if (batched) { + per_image = std::make_shared(std::vector{meta.batch, meta.out_channels, meta.kernel_elems}, + DataType::kFLOAT32, device); + gemm_c = per_image->DataPtr(); + stride_c = meta.out_channels * meta.kernel_elems; + } + Dispatcher::Instance().Call({device.type(), "Gemm"}, device, + GemmParams{ + .trans_a = GemmTranspose::kTranspose, + .trans_b = GemmTranspose::kNoTranspose, + .m = static_cast(meta.kernel_elems), + .n = static_cast(meta.out_channels), + .k = static_cast(meta.patches), + .A = columns->DataPtr(), + .lda = static_cast(meta.patches), + .B = grad_output->DataPtr(), + .ldb = static_cast(meta.patches), + .C = gemm_c, + .ldc = static_cast(meta.kernel_elems), + .alpha = 1.0f, + .beta = 0.0f, + .batch_count = batch_count, + .stride_a = batched ? meta.kernel_elems * meta.patches : 0, + .stride_b = batched ? meta.out_channels * meta.patches : 0, + .stride_c = stride_c, + .input_dtype = DataType::kFLOAT32, + .output_dtype = DataType::kFLOAT32, + }); + if (batched) { + // N == 1 fast path above overwrote grad_weight directly (beta=0). + const int64_t elems = meta.out_channels * meta.kernel_elems; + const int threads_per_block = 256; + const int num_blocks = static_cast((elems + threads_per_block - 1) / threads_per_block); + ConvBackwardWeightSumKernel<<>>( + static_cast(per_image->DataPtr()), static_cast(grad_weight->DataPtr()), meta.batch, + elems); + } + return grad_weight; +} + +// Input grad via GEMM + gather col2im (FP32 only). Per image, +// row-major dXcol[K, P] = W^T[K, Cout] * dY[Cout, P] (K = Cin*Kh*Kw, +// P = Hout*Wout matches Im2colKernel's row order), then Col2imKernel +// gathers each input element's covering patches (no atomics, deterministic). +// NOTE: this entry point mirrors the CPU signature (shared_ptr const&, +// int64 by value, vector const&) so the dispatcher resolves the +// identical "Conv2dBackwardInput" overload on CUDA. +std::shared_ptr Conv2dBackwardInput(const std::shared_ptr &weight, + const std::shared_ptr &grad_output, int64_t stride, int64_t padding, + const std::vector &input_dims) { + CHECK(weight->Dtype() == DataType::kFLOAT32) << "CUDA Conv2dBackwardInput currently supports float32 only"; + CHECK(grad_output->Dtype() == DataType::kFLOAT32) << "CUDA Conv2dBackwardInput currently supports float32 only"; + CHECK_EQ(input_dims.size(), 4); + + const auto &go_dims = grad_output->Dims(); + CHECK_EQ(go_dims.size(), 4); + const auto &w_dims = weight->Dims(); + CHECK_EQ(w_dims.size(), 4); + CHECK_EQ(w_dims[0], go_dims[1]); + CHECK_EQ(w_dims[1], input_dims[1]); + CHECK_EQ(go_dims[2], (input_dims[2] + 2 * padding - w_dims[2]) / stride + 1); + CHECK_EQ(go_dims[3], (input_dims[3] + 2 * padding - w_dims[3]) / stride + 1); + + Conv2dMeta meta; + meta.batch = input_dims[0]; + meta.in_channels = input_dims[1]; + meta.input_h = input_dims[2]; + meta.input_w = input_dims[3]; + meta.out_channels = go_dims[1]; + meta.kernel_h = w_dims[2]; + meta.kernel_w = w_dims[3]; + meta.stride = stride; + meta.padding = padding; + meta.output_h = go_dims[2]; + meta.output_w = go_dims[3]; + meta.patches = meta.output_h * meta.output_w; + meta.kernel_elems = meta.in_channels * meta.kernel_h * meta.kernel_w; + + auto device = grad_output->GetDevice(); + auto grad_input = std::make_shared(input_dims, DataType::kFLOAT32, device); + if (meta.batch == 0) { + return grad_input; + } + + auto columns = std::make_shared(std::vector{meta.batch, meta.kernel_elems, meta.patches}, + DataType::kFLOAT32, device); + + // cuBLAS is column-major: row-major dXcol[K, P] = W^T[K, Cout] * + // dY[Cout, P] is computed as dXcol^T[P, K] = dY^T[P, Cout] * W[Cout, K], + // with dY viewed as the column-major [P, Cout] transpose (trans_a=N) + // and W viewed as column-major [K, Cout] transposed (trans_b=T). + // C = dXcol^T[P, K], A = dY^T[P, Cout], B = W[K, Cout]^T. + // Weight is shared across images (stride_b = 0); batch_count == 1 uses + // zero strides per Gemm convention. + const int batch_count = static_cast(meta.batch); + const bool batched = meta.batch > 1; + Dispatcher::Instance().Call({device.type(), "Gemm"}, device, + GemmParams{ + .trans_a = GemmTranspose::kNoTranspose, + .trans_b = GemmTranspose::kTranspose, + .m = static_cast(meta.patches), + .n = static_cast(meta.kernel_elems), + .k = static_cast(meta.out_channels), + .A = grad_output->DataPtr(), + .lda = static_cast(meta.patches), + .B = weight->DataPtr(), + .ldb = static_cast(meta.kernel_elems), + .C = columns->DataPtr(), + .ldc = static_cast(meta.patches), + .alpha = 1.0f, + .beta = 0.0f, + .batch_count = batch_count, + .stride_a = batched ? meta.out_channels * meta.patches : 0, + .stride_b = 0, + .stride_c = batched ? meta.kernel_elems * meta.patches : 0, + .input_dtype = DataType::kFLOAT32, + .output_dtype = DataType::kFLOAT32, + }); + + const int64_t total = meta.batch * meta.in_channels * meta.input_h * meta.input_w; + if (total == 0) { + return grad_input; + } + const int threads_per_block = 256; + const int num_blocks = static_cast((total + threads_per_block - 1) / threads_per_block); + const auto &cuda_stream = dynamic_cast( + infini_train::core::GetDeviceGuardImpl(device.type())->GetStream(device)) + ->cuda_stream(); + Col2imKernel<<>>(static_cast(columns->DataPtr()), + static_cast(grad_input->DataPtr()), meta); + return grad_input; +} + +// Bias grad via N*P reduction (FP32 only). db[oc] = sum over N*Hout*Wout +// of dY (grad_output is row-major [N, Cout, P], P = Hout * Wout). +// NOTE: this entry point mirrors the CPU signature (shared_ptr const&) so +// the dispatcher resolves the identical "Conv2dBackwardBias" overload on +// CUDA. +std::shared_ptr Conv2dBackwardBias(const std::shared_ptr &grad_output) { + CHECK(grad_output->Dtype() == DataType::kFLOAT32) << "CUDA Conv2dBackwardBias currently supports float32 only"; + const auto &go_dims = grad_output->Dims(); + CHECK_EQ(go_dims.size(), 4); + + const int64_t out_channels = go_dims[1]; + const int64_t patches = go_dims[2] * go_dims[3]; + const int64_t rows = go_dims[0] * patches; + + auto device = grad_output->GetDevice(); + auto grad_bias = std::make_shared(std::vector{out_channels}, DataType::kFLOAT32, device); + if (out_channels == 0) { + return grad_bias; + } + + const auto &cuda_stream = dynamic_cast( + infini_train::core::GetDeviceGuardImpl(device.type())->GetStream(device)) + ->cuda_stream(); + // One block per channel (same pattern as cuda::LinearBackwardBias); + // rows == 0 writes zeros via the empty BlockReduce sum. + constexpr int kBlockSize = 256; + ConvBackwardBiasKernel<<>>( + static_cast(grad_output->DataPtr()), static_cast(grad_bias->DataPtr()), rows, + out_channels, patches); + return grad_bias; +} + +} // namespace infini_train::kernels::cuda + +#define REGISTER_CUDA_CONV_KERNEL(kernel_name) \ + REGISTER_KERNEL(infini_train::Device::DeviceType::kCUDA, kernel_name, infini_train::kernels::cuda::kernel_name) + +REGISTER_CUDA_CONV_KERNEL(Im2colForward) +REGISTER_CUDA_CONV_KERNEL(Conv2dForward) +REGISTER_CUDA_CONV_KERNEL(Conv2dBackwardInput) +REGISTER_CUDA_CONV_KERNEL(Conv2dBackwardWeight) +REGISTER_CUDA_CONV_KERNEL(Conv2dBackwardBias) + +#undef REGISTER_CUDA_CONV_KERNEL diff --git a/infini_train/src/kernels/cuda/relu.cu b/infini_train/src/kernels/cuda/relu.cu new file mode 100644 index 000000000..2acce9020 --- /dev/null +++ b/infini_train/src/kernels/cuda/relu.cu @@ -0,0 +1,86 @@ +#include +#include + +#include "glog/logging.h" + +#include "infini_train/include/core/runtime/device_guard.h" +#include "infini_train/include/dispatcher.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_runtime_common.h" + +namespace infini_train::kernels::cuda { +namespace { +__global__ void ReLUForwardKernel(float *output, const float *input, size_t num_elements) { + size_t idx = blockIdx.x * blockDim.x + threadIdx.x; + if (idx < num_elements) { + float x = input[idx]; + output[idx] = x > 0.0f ? x : 0.0f; + } +} + +__global__ void ReLUBackwardKernel(float *grad_input, const float *input, const float *grad_output, + size_t num_elements) { + size_t idx = blockIdx.x * blockDim.x + threadIdx.x; + if (idx < num_elements) { + grad_input[idx] = input[idx] > 0.0f ? grad_output[idx] : 0.0f; + } +} + +void LaunchForward(const std::shared_ptr &output, const std::shared_ptr &input) { + const size_t num_elements = static_cast(output->NumElements()); + if (num_elements == 0) { + return; + } + const int threads_per_block = 256; + const int num_blocks = static_cast((num_elements + threads_per_block - 1) / threads_per_block); + auto device = output->GetDevice(); + const auto &cuda_stream = dynamic_cast( + infini_train::core::GetDeviceGuardImpl(device.type())->GetStream(device)) + ->cuda_stream(); + ReLUForwardKernel<<>>( + static_cast(output->DataPtr()), static_cast(input->DataPtr()), num_elements); +} + +void LaunchBackward(const std::shared_ptr &grad_input, const std::shared_ptr &input, + const std::shared_ptr &grad_output) { + const size_t num_elements = static_cast(grad_input->NumElements()); + if (num_elements == 0) { + return; + } + const int threads_per_block = 256; + const int num_blocks = static_cast((num_elements + threads_per_block - 1) / threads_per_block); + auto device = grad_input->GetDevice(); + const auto &cuda_stream = dynamic_cast( + infini_train::core::GetDeviceGuardImpl(device.type())->GetStream(device)) + ->cuda_stream(); + ReLUBackwardKernel<<>>( + static_cast(grad_input->DataPtr()), static_cast(input->DataPtr()), + static_cast(grad_output->DataPtr()), num_elements); +} +} // namespace + +std::shared_ptr ReLUForward(const std::shared_ptr &input) { + CHECK(input->Dtype() == DataType::kFLOAT32) << "CUDA ReLUForward currently supports float32 only"; + auto output = std::make_shared(input->Dims(), DataType::kFLOAT32, input->GetDevice()); + LaunchForward(output, input); + return output; +} + +std::shared_ptr ReLUBackward(const std::shared_ptr &input, const std::shared_ptr &grad_output) { + CHECK(input->Dtype() == DataType::kFLOAT32) << "CUDA ReLUBackward currently supports float32 only"; + CHECK(grad_output->Dtype() == DataType::kFLOAT32) << "CUDA ReLUBackward currently supports float32 only"; + auto grad_input = std::make_shared(input->Dims(), DataType::kFLOAT32, input->GetDevice()); + LaunchBackward(grad_input, input, grad_output); + return grad_input; +} +} // namespace infini_train::kernels::cuda + +#define REGISTER_CUDA_RELU_KERNEL(kernel_name) \ + REGISTER_KERNEL(infini_train::Device::DeviceType::kCUDA, kernel_name, infini_train::kernels::cuda::kernel_name) + +REGISTER_CUDA_RELU_KERNEL(ReLUForward) +REGISTER_CUDA_RELU_KERNEL(ReLUBackward) + +#undef REGISTER_CUDA_RELU_KERNEL diff --git a/infini_train/src/nn/modules/activations.cc b/infini_train/src/nn/modules/activations.cc index d1bbc9da8..d41baf64f 100644 --- a/infini_train/src/nn/modules/activations.cc +++ b/infini_train/src/nn/modules/activations.cc @@ -21,4 +21,8 @@ std::vector> NewGELU::Forward(const std::vector> SwiGLU::Forward(const std::vector> &x) { return {x[0] * function::Sigmoid(x[0])}; } + +std::vector> ReLU::Forward(const std::vector> &input_tensors) { + return std::make_shared()->Apply(input_tensors); +} } // namespace infini_train::nn diff --git a/infini_train/src/nn/modules/conv.cc b/infini_train/src/nn/modules/conv.cc new file mode 100644 index 000000000..673893e30 --- /dev/null +++ b/infini_train/src/nn/modules/conv.cc @@ -0,0 +1,44 @@ +#include "infini_train/include/nn/modules/conv.h" + +#include +#include +#include + +#include "infini_train/include/autograd/conv.h" +#include "infini_train/include/device.h" +#include "infini_train/include/nn/init.h" +#include "infini_train/include/tensor.h" + +namespace infini_train::nn { +Conv2d::Conv2d(int64_t in_channels, int64_t out_channels, int64_t kernel_size, int64_t stride, int64_t padding, + bool bias, Device device) + : CloneableModule(kType), stride_(stride), padding_(padding), bias_(bias) { + device_ = device; + + parameters_[kParamWeightName] + = std::make_shared(std::vector{out_channels, in_channels, kernel_size, kernel_size}, + DataType::kFLOAT32, device_) + ->RequiresGrad(); + if (bias) { + parameters_[kParamBiasName] + = std::make_shared(std::vector{out_channels}, DataType::kFLOAT32, device_)->RequiresGrad(); + } + ResetParameters(); +} + +std::vector> Conv2d::Forward(const std::vector> &input_tensors) { + return std::make_shared(stride_, padding_) + ->Apply(bias_ ? std::vector>{input_tensors[0], parameters_[kParamWeightName], + parameters_[kParamBiasName]} + : std::vector>{input_tensors[0], parameters_[kParamWeightName]}); +} + +void Conv2d::ResetParameters() { + init::KaimingUniform(parameters_[kParamWeightName], sqrt(5.0f)); + if (bias_) { + const auto [fan_in, _] = init::CalculateFanInAndFanOut(parameters_[kParamWeightName]); + const float bound = fan_in > 0 ? 1.0 / sqrt(fan_in) : 0.0; + init::Uniform(parameters_[kParamBiasName], -bound, bound); + } +} +} // namespace infini_train::nn diff --git a/tests/CMakeLists.txt b/tests/CMakeLists.txt index 43bef143c..61f7c6bd0 100644 --- a/tests/CMakeLists.txt +++ b/tests/CMakeLists.txt @@ -38,3 +38,6 @@ add_subdirectory(transformer) # Checkpoint tests add_subdirectory(checkpoint) + +# Example tests +add_subdirectory(example) diff --git a/tests/autograd/test_autograd_conv_cpu.cc b/tests/autograd/test_autograd_conv_cpu.cc new file mode 100644 index 000000000..3392d4189 --- /dev/null +++ b/tests/autograd/test_autograd_conv_cpu.cc @@ -0,0 +1,150 @@ +#include + +#include "gtest/gtest.h" + +#include "infini_train/include/autograd/conv.h" +#include "infini_train/include/nn/modules/conv.h" +#include "infini_train/include/nn/parallel/global.h" +#include "infini_train/include/tensor.h" + +#include "tests/common/test_utils.h" + +using namespace infini_train; + +class AutogradConvCPUTest : public infini_train::test::InfiniTrainTest {}; + +TEST_P(AutogradConvCPUTest, Conv2dForwardNoBias) { + ONLY_CPU(); + // input 3x3 = 1..9, weight [[1, 0], [0, -1]], s=1, p=0 -> 2x2 of -4 + const float input_values[] = {1.0f, 2.0f, 3.0f, 4.0f, 5.0f, 6.0f, 7.0f, 8.0f, 9.0f}; + const float weight_values[] = {1.0f, 0.0f, 0.0f, -1.0f}; + auto input + = std::make_shared(input_values, std::vector{1, 1, 3, 3}, DataType::kFLOAT32, GetDevice()); + auto weight + = std::make_shared(weight_values, std::vector{1, 1, 2, 2}, DataType::kFLOAT32, GetDevice()); + auto conv_fn = std::make_shared(1, 0); + auto result = conv_fn->Apply({input, weight}); + ASSERT_EQ(result.size(), 1); + EXPECT_EQ(result[0]->Dims(), (std::vector{1, 1, 2, 2})); + test::ExpectTensorFloatEqual(result[0], {-4.0f, -4.0f, -4.0f, -4.0f}); +} + +TEST_P(AutogradConvCPUTest, Conv2dForwardWithBias) { + ONLY_CPU(); + const float input_values[] = {1.0f, 2.0f, 3.0f, 4.0f, 5.0f, 6.0f, 7.0f, 8.0f, 9.0f}; + const float weight_values[] = {1.0f, 0.0f, 0.0f, -1.0f}; + const float bias_values[] = {1.0f}; + auto input + = std::make_shared(input_values, std::vector{1, 1, 3, 3}, DataType::kFLOAT32, GetDevice()); + auto weight + = std::make_shared(weight_values, std::vector{1, 1, 2, 2}, DataType::kFLOAT32, GetDevice()); + auto bias = std::make_shared(bias_values, std::vector{1}, DataType::kFLOAT32, GetDevice()); + auto conv_fn = std::make_shared(1, 0); + auto result = conv_fn->Apply({input, weight, bias}); + ASSERT_EQ(result.size(), 1); + EXPECT_EQ(result[0]->Dims(), (std::vector{1, 1, 2, 2})); + test::ExpectTensorFloatEqual(result[0], {-3.0f, -3.0f, -3.0f, -3.0f}); +} + +TEST_P(AutogradConvCPUTest, Conv2dForwardStridePadding) { + ONLY_CPU(); + // input 4x4 = 1..16, 3x3 ones, s=2, p=1 -> 2x2 {14, 30, 57, 99} + std::vector input_values(16); + for (int i = 0; i < 16; ++i) { input_values[i] = static_cast(i + 1); } + auto input = std::make_shared(input_values.data(), std::vector{1, 1, 4, 4}, DataType::kFLOAT32, + GetDevice()); + auto weight = std::make_shared(std::vector{1, 1, 3, 3}, DataType::kFLOAT32, GetDevice()); + weight->Fill(1.0f); + auto conv_fn = std::make_shared(2, 1); + auto result = conv_fn->Apply({input, weight}); + ASSERT_EQ(result.size(), 1); + EXPECT_EQ(result[0]->Dims(), (std::vector{1, 1, 2, 2})); + test::ExpectTensorFloatEqual(result[0], {14.0f, 30.0f, 57.0f, 99.0f}); +} + +TEST_P(AutogradConvCPUTest, Conv2dForwardNonSquareKernel) { + ONLY_CPU(); + // input 3x4 = 1..12, 2x3 ones (Kh != Kw), s=1, p=0 -> 2x2 {24, 30, 48, 54} + std::vector input_values(12); + for (int i = 0; i < 12; ++i) { input_values[i] = static_cast(i + 1); } + auto input = std::make_shared(input_values.data(), std::vector{1, 1, 3, 4}, DataType::kFLOAT32, + GetDevice()); + auto weight = std::make_shared(std::vector{1, 1, 2, 3}, DataType::kFLOAT32, GetDevice()); + weight->Fill(1.0f); + auto conv_fn = std::make_shared(1, 0); + auto result = conv_fn->Apply({input, weight}); + ASSERT_EQ(result.size(), 1); + EXPECT_EQ(result[0]->Dims(), (std::vector{1, 1, 2, 2})); + test::ExpectTensorFloatEqual(result[0], {24.0f, 30.0f, 48.0f, 54.0f}); +} + +TEST_P(AutogradConvCPUTest, Conv2dBackward) { + ONLY_CPU(); + const float input_values[] = {1.0f, 2.0f, 3.0f, 4.0f, 5.0f, 6.0f, 7.0f, 8.0f, 9.0f}; + const float weight_values[] = {1.0f, 0.0f, 0.0f, -1.0f}; + auto input + = std::make_shared(input_values, std::vector{1, 1, 3, 3}, DataType::kFLOAT32, GetDevice()) + ->RequiresGrad(); + auto weight + = std::make_shared(weight_values, std::vector{1, 1, 2, 2}, DataType::kFLOAT32, GetDevice()) + ->RequiresGrad(); + auto bias = std::make_shared(std::vector{1}, DataType::kFLOAT32, GetDevice(), true); + bias->Fill(0.0f); + auto conv_fn = std::make_shared(1, 0); + auto result = conv_fn->Apply({input, weight, bias}); + auto grad = std::make_shared(std::vector{1, 1, 2, 2}, DataType::kFLOAT32, GetDevice()); + grad->Fill(1.0f); + auto grad_inputs = conv_fn->Backward({grad}); + ASSERT_EQ(grad_inputs.size(), 3); + ASSERT_NE(grad_inputs[0], nullptr); + ASSERT_NE(grad_inputs[1], nullptr); + ASSERT_NE(grad_inputs[2], nullptr); + EXPECT_EQ(grad_inputs[0]->Dims(), (std::vector{1, 1, 3, 3})); + EXPECT_EQ(grad_inputs[1]->Dims(), (std::vector{1, 1, 2, 2})); + EXPECT_EQ(grad_inputs[2]->Dims(), (std::vector{1})); + test::ExpectTensorFloatEqual(grad_inputs[0], {1.0f, 1.0f, 0.0f, 1.0f, 0.0f, -1.0f, 0.0f, -1.0f, -1.0f}); + test::ExpectTensorFloatEqual(grad_inputs[1], {12.0f, 16.0f, 24.0f, 28.0f}); + test::ExpectTensorFloatEqual(grad_inputs[2], {4.0f}); +} + +TEST_P(AutogradConvCPUTest, Conv2dBackwardNoBias) { + ONLY_CPU(); + auto input = std::make_shared(std::vector{1, 1, 3, 3}, DataType::kFLOAT32, GetDevice(), true); + input->Fill(1.0f); + auto weight = std::make_shared(std::vector{1, 1, 2, 2}, DataType::kFLOAT32, GetDevice(), true); + weight->Fill(1.0f); + auto conv_fn = std::make_shared(1, 0); + auto result = conv_fn->Apply({input, weight}); + auto grad = std::make_shared(std::vector{1, 1, 2, 2}, DataType::kFLOAT32, GetDevice()); + grad->Fill(1.0f); + auto grad_inputs = conv_fn->Backward({grad}); + ASSERT_EQ(grad_inputs.size(), 2); + ASSERT_NE(grad_inputs[0], nullptr); + ASSERT_NE(grad_inputs[1], nullptr); + test::ExpectTensorFloatEqual(grad_inputs[0], {1.0f, 2.0f, 1.0f, 2.0f, 4.0f, 2.0f, 1.0f, 2.0f, 1.0f}); + test::ExpectTensorFloatEqual(grad_inputs[1], {4.0f, 4.0f, 4.0f, 4.0f}); +} + +TEST_P(AutogradConvCPUTest, Conv2dModuleForward) { + ONLY_CPU(); + nn::Conv2d conv(1, 2, 3, 1, 1, true, GetDevice()); + EXPECT_TRUE(conv.has_bias()); + auto input = std::make_shared(std::vector{1, 1, 4, 4}, DataType::kFLOAT32, GetDevice()); + input->Fill(1.0f); + auto result = conv.Forward({input}); + ASSERT_EQ(result.size(), 1); + EXPECT_EQ(result[0]->Dims(), (std::vector{1, 2, 4, 4})); +} + +TEST_P(AutogradConvCPUTest, Conv2dModuleNoBias) { + ONLY_CPU(); + nn::Conv2d conv(1, 1, 2, 2, 0, false, GetDevice()); + EXPECT_FALSE(conv.has_bias()); + auto input = std::make_shared(std::vector{1, 1, 4, 4}, DataType::kFLOAT32, GetDevice()); + input->Fill(1.0f); + auto result = conv.Forward({input}); + ASSERT_EQ(result.size(), 1); + EXPECT_EQ(result[0]->Dims(), (std::vector{1, 1, 2, 2})); +} + +INFINI_TRAIN_REGISTER_TEST(AutogradConvCPUTest); diff --git a/tests/autograd/test_autograd_conv_cuda.cc b/tests/autograd/test_autograd_conv_cuda.cc new file mode 100644 index 000000000..5e7707de7 --- /dev/null +++ b/tests/autograd/test_autograd_conv_cuda.cc @@ -0,0 +1,283 @@ +// CUDA Conv2d forward parity vs CPU direct kernel (im2col + GEMM path). +// Grid: B{1,2} x Cin{1,3} x Cout{2,3} x H!=W x K{1,3} x s{1,2} x p{0,1}, +// with and without bias. Threshold: abs err <= 1e-5. +// Plus weight-grad parity (im2col + batched GEMM + N-reduce) over a smaller +// grid: B{1,2} x Cin{1,3} x Cout{2,3} x K{1,3} x s{1,2} x p{0,1}. Threshold: abs err <= 1e-5. +// Plus input-grad parity (W^T * dY GEMM + gather col2im) over the same +// smaller grid. Threshold: abs err <= 1e-5. +// Plus bias-grad parity (sum over N*Hout*Wout) over B{1,2} x Cin{1,3} x Cout{2,3}. +// Threshold: abs err <= 1e-5. +#include +#include + +#include "gtest/gtest.h" + +#include "infini_train/include/autograd/conv.h" +#include "infini_train/include/tensor.h" + +#include "tests/common/test_utils.h" + +using namespace infini_train; + +class AutogradConvCudaTest : public infini_train::test::InfiniTrainTest {}; + +namespace { + +std::vector DeterministicVals(int64_t n, int seed) { + std::vector vals(static_cast(n)); + uint32_t state = static_cast(seed); + for (int64_t i = 0; i < n; ++i) { + state = state * 1664525u + 1013904223u; + vals[static_cast(i)] = static_cast(state % 1000) / 1000.0f - 0.5f; + } + return vals; +} + +void CheckParity(const Device &cuda_device, int64_t batch, int64_t cin, int64_t cout, int64_t h, int64_t w, int64_t k, + int64_t stride, int64_t padding, bool with_bias) { + const Device cpu_device = Device(); + const std::vector in_dims = {batch, cin, h, w}; + const std::vector w_dims = {cout, cin, k, k}; + const std::vector b_dims = {cout}; + + int64_t in_n = batch * cin * h * w; + int64_t w_n = cout * cin * k * k; + auto in_vals = DeterministicVals(in_n, 11); + auto w_vals = DeterministicVals(w_n, 23); + auto b_vals = DeterministicVals(cout, 37); + + auto cpu_in = std::make_shared(in_vals.data(), in_dims, DataType::kFLOAT32, cpu_device); + auto cpu_w = std::make_shared(w_vals.data(), w_dims, DataType::kFLOAT32, cpu_device); + auto cpu_b = with_bias ? std::make_shared(b_vals.data(), b_dims, DataType::kFLOAT32, cpu_device) : nullptr; + + auto cuda_in = std::make_shared(cpu_in->To(cuda_device)); + auto cuda_w = std::make_shared(cpu_w->To(cuda_device)); + std::shared_ptr cuda_b = nullptr; + if (with_bias) { + cuda_b = std::make_shared(cpu_b->To(cuda_device)); + } + + auto cpu_fn = std::make_shared(stride, padding); + auto cuda_fn = std::make_shared(stride, padding); + std::vector> cpu_args = with_bias + ? std::vector>{cpu_in, cpu_w, cpu_b} + : std::vector>{cpu_in, cpu_w}; + std::vector> cuda_args = with_bias + ? std::vector>{cuda_in, cuda_w, cuda_b} + : std::vector>{cuda_in, cuda_w}; + auto cpu_out = cpu_fn->Apply(cpu_args); + auto cuda_out = cuda_fn->Apply(cuda_args); + ASSERT_EQ(cpu_out.size(), 1); + ASSERT_EQ(cuda_out.size(), 1); + EXPECT_EQ(cuda_out[0]->Dims(), cpu_out[0]->Dims()); + + auto expected_cpu = cpu_out[0]->To(Device()); + const float *expected = static_cast(expected_cpu.DataPtr()); + test::ExpectTensorNear(cuda_out[0], std::vector(expected, expected + expected_cpu.NumElements()), 1e-5f); +} + +} // namespace + +void CheckWeightGradParity(const Device &cuda_device, int64_t batch, int64_t cin, int64_t cout, int64_t h, int64_t w, + int64_t k, int64_t stride, int64_t padding) { + const Device cpu_device = Device(); + const std::vector in_dims = {batch, cin, h, w}; + const std::vector w_dims = {cout, cin, k, k}; + const int64_t oh = (h + 2 * padding - k) / stride + 1; + const int64_t ow = (w + 2 * padding - k) / stride + 1; + const std::vector go_dims = {batch, cout, oh, ow}; + + auto in_vals = DeterministicVals(batch * cin * h * w, 11); + auto w_vals = DeterministicVals(cout * cin * k * k, 23); + auto go_vals = DeterministicVals(batch * cout * oh * ow, 41); + + auto cpu_in = std::make_shared(in_vals.data(), in_dims, DataType::kFLOAT32, cpu_device); + auto cpu_w = std::make_shared(w_vals.data(), w_dims, DataType::kFLOAT32, cpu_device)->RequiresGrad(); + auto cpu_go = std::make_shared(go_vals.data(), go_dims, DataType::kFLOAT32, cpu_device); + + // NOTE: input deliberately does NOT require grad (weight-grad path only; + // CUDA Conv2dBackwardInput is a different task), so grads[0] is null. + auto cuda_in = std::make_shared(cpu_in->To(cuda_device)); + auto cuda_w = std::make_shared(cpu_w->To(cuda_device))->RequiresGrad(); + auto cuda_go = std::make_shared(cpu_go->To(cuda_device)); + + auto cpu_fn = std::make_shared(stride, padding); + auto cuda_fn = std::make_shared(stride, padding); + cpu_fn->Apply({cpu_in, cpu_w}); + cuda_fn->Apply({cuda_in, cuda_w}); + + auto cpu_grads = cpu_fn->Backward({cpu_go}); + auto cuda_grads = cuda_fn->Backward({cuda_go}); + ASSERT_EQ(cpu_grads.size(), 2); + ASSERT_EQ(cuda_grads.size(), 2); + EXPECT_EQ(cuda_grads[0], nullptr); + ASSERT_NE(cuda_grads[1], nullptr); + EXPECT_EQ(cuda_grads[1]->Dims(), cpu_grads[1]->Dims()); + + auto expected_cpu = cpu_grads[1]->To(Device()); + const float *expected = static_cast(expected_cpu.DataPtr()); + test::ExpectTensorNear(cuda_grads[1], std::vector(expected, expected + expected_cpu.NumElements()), 1e-5f); +} + +void CheckBiasGradParity(const Device &cuda_device, int64_t batch, int64_t cin, int64_t cout) { + const Device cpu_device = Device(); + const int64_t h = 5; + const int64_t w = 4; + const int64_t k = 3; + const int64_t stride = 1; + const int64_t padding = 1; + const std::vector in_dims = {batch, cin, h, w}; + const std::vector w_dims = {cout, cin, k, k}; + const std::vector b_dims = {cout}; + const int64_t oh = (h + 2 * padding - k) / stride + 1; + const int64_t ow = (w + 2 * padding - k) / stride + 1; + const std::vector go_dims = {batch, cout, oh, ow}; + + auto in_vals = DeterministicVals(batch * cin * h * w, 11); + auto w_vals = DeterministicVals(cout * cin * k * k, 23); + auto b_vals = DeterministicVals(cout, 37); + auto go_vals = DeterministicVals(batch * cout * oh * ow, 41); + + auto cpu_in = std::make_shared(in_vals.data(), in_dims, DataType::kFLOAT32, cpu_device); + auto cpu_w = std::make_shared(w_vals.data(), w_dims, DataType::kFLOAT32, cpu_device); + auto cpu_b = std::make_shared(b_vals.data(), b_dims, DataType::kFLOAT32, cpu_device)->RequiresGrad(); + auto cpu_go = std::make_shared(go_vals.data(), go_dims, DataType::kFLOAT32, cpu_device); + + // NOTE: input/weight deliberately do NOT require grad (bias-grad path + // only), so grads[0] and grads[1] are null. + auto cuda_in = std::make_shared(cpu_in->To(cuda_device)); + auto cuda_w = std::make_shared(cpu_w->To(cuda_device)); + auto cuda_b = std::make_shared(cpu_b->To(cuda_device))->RequiresGrad(); + auto cuda_go = std::make_shared(cpu_go->To(cuda_device)); + + auto cpu_fn = std::make_shared(stride, padding); + auto cuda_fn = std::make_shared(stride, padding); + cpu_fn->Apply({cpu_in, cpu_w, cpu_b}); + cuda_fn->Apply({cuda_in, cuda_w, cuda_b}); + + auto cpu_grads = cpu_fn->Backward({cpu_go}); + auto cuda_grads = cuda_fn->Backward({cuda_go}); + ASSERT_EQ(cpu_grads.size(), 3); + ASSERT_EQ(cuda_grads.size(), 3); + EXPECT_EQ(cuda_grads[0], nullptr); + EXPECT_EQ(cuda_grads[1], nullptr); + ASSERT_NE(cuda_grads[2], nullptr); + EXPECT_EQ(cuda_grads[2]->Dims(), cpu_grads[2]->Dims()); + + auto expected_cpu = cpu_grads[2]->To(Device()); + const float *expected = static_cast(expected_cpu.DataPtr()); + test::ExpectTensorNear(cuda_grads[2], std::vector(expected, expected + expected_cpu.NumElements()), 1e-5f); +} + +void CheckInputGradParity(const Device &cuda_device, int64_t batch, int64_t cin, int64_t cout, int64_t h, int64_t w, + int64_t k, int64_t stride, int64_t padding) { + const Device cpu_device = Device(); + const std::vector in_dims = {batch, cin, h, w}; + const std::vector w_dims = {cout, cin, k, k}; + const int64_t oh = (h + 2 * padding - k) / stride + 1; + const int64_t ow = (w + 2 * padding - k) / stride + 1; + const std::vector go_dims = {batch, cout, oh, ow}; + + auto in_vals = DeterministicVals(batch * cin * h * w, 11); + auto w_vals = DeterministicVals(cout * cin * k * k, 23); + auto go_vals = DeterministicVals(batch * cout * oh * ow, 41); + + auto cpu_in = std::make_shared(in_vals.data(), in_dims, DataType::kFLOAT32, cpu_device)->RequiresGrad(); + auto cpu_w = std::make_shared(w_vals.data(), w_dims, DataType::kFLOAT32, cpu_device); + auto cpu_go = std::make_shared(go_vals.data(), go_dims, DataType::kFLOAT32, cpu_device); + + // NOTE: weight deliberately does NOT require grad (input-grad path only), + // so grads[1] is null. + auto cuda_in = std::make_shared(cpu_in->To(cuda_device))->RequiresGrad(); + auto cuda_w = std::make_shared(cpu_w->To(cuda_device)); + auto cuda_go = std::make_shared(cpu_go->To(cuda_device)); + + auto cpu_fn = std::make_shared(stride, padding); + auto cuda_fn = std::make_shared(stride, padding); + cpu_fn->Apply({cpu_in, cpu_w}); + cuda_fn->Apply({cuda_in, cuda_w}); + + auto cpu_grads = cpu_fn->Backward({cpu_go}); + auto cuda_grads = cuda_fn->Backward({cuda_go}); + ASSERT_EQ(cpu_grads.size(), 2); + ASSERT_EQ(cuda_grads.size(), 2); + ASSERT_NE(cuda_grads[0], nullptr); + EXPECT_EQ(cuda_grads[1], nullptr); + EXPECT_EQ(cuda_grads[0]->Dims(), cpu_grads[0]->Dims()); + + auto expected_cpu = cpu_grads[0]->To(Device()); + const float *expected = static_cast(expected_cpu.DataPtr()); + test::ExpectTensorNear(cuda_grads[0], std::vector(expected, expected + expected_cpu.NumElements()), 1e-5f); +} + +TEST_P(AutogradConvCudaTest, Conv2dBackwardBiasParityGrid) { + ONLY_CUDA(); + const Device cuda_device = GetDevice(); + for (int64_t batch : {1, 2}) { + for (int64_t cin : {1, 3}) { + for (int64_t cout : {2, 3}) { CheckBiasGradParity(cuda_device, batch, cin, cout); } + } + } +} + +TEST_P(AutogradConvCudaTest, Conv2dBackwardInputParityGrid) { + ONLY_CUDA(); + const Device cuda_device = GetDevice(); + for (int64_t batch : {1, 2}) { + for (int64_t cin : {1, 3}) { + for (int64_t cout : {2, 3}) { + for (int64_t k : {1, 3}) { + for (int64_t stride : {1, 2}) { + for (int64_t padding : {0, 1}) { + // H != W; H + 2p >= K and W + 2p >= K always hold here. + CheckInputGradParity(cuda_device, batch, cin, cout, /*h=*/5, /*w=*/4, k, stride, padding); + } + } + } + } + } + } +} + +TEST_P(AutogradConvCudaTest, Conv2dBackwardWeightParityGrid) { + ONLY_CUDA(); + const Device cuda_device = GetDevice(); + for (int64_t batch : {1, 2}) { + for (int64_t cin : {1, 3}) { + for (int64_t cout : {2, 3}) { + for (int64_t k : {1, 3}) { + for (int64_t stride : {1, 2}) { + for (int64_t padding : {0, 1}) { + // H != W; H + 2p >= K and W + 2p >= K always hold here. + CheckWeightGradParity(cuda_device, batch, cin, cout, /*h=*/5, /*w=*/4, k, stride, padding); + } + } + } + } + } + } +} + +TEST_P(AutogradConvCudaTest, Conv2dForwardParityGrid) { + ONLY_CUDA(); + const Device cuda_device = GetDevice(); + for (int64_t batch : {1, 2}) { + for (int64_t cin : {1, 3}) { + for (int64_t cout : {2, 3}) { + for (int64_t k : {1, 3}) { + for (int64_t stride : {1, 2}) { + for (int64_t padding : {0, 1}) { + for (bool with_bias : {false, true}) { + // H != W; H + 2p >= K and W + 2p >= K always hold here. + CheckParity(cuda_device, batch, cin, cout, /*h=*/5, /*w=*/4, k, stride, padding, + with_bias); + } + } + } + } + } + } + } +} + +INFINI_TRAIN_REGISTER_TEST(AutogradConvCudaTest); diff --git a/tests/autograd/test_autograd_conv_gradcheck.cc b/tests/autograd/test_autograd_conv_gradcheck.cc new file mode 100644 index 000000000..c735d2f03 --- /dev/null +++ b/tests/autograd/test_autograd_conv_gradcheck.cc @@ -0,0 +1,128 @@ +// Finite-difference gradcheck for Conv2d autograd (CPU). +// +// Config: N=1, Cin=2, Cout=2, H=4, W=5, K=3, stride=2, padding=1. +// Compares analytic Backward() grads (dInput/dWeight/dBias) against central +// differences of loss = dot(forward, grad_output). Self-contained: no torch +// dependency. Threshold: relative err < 1e-3 per element. +#include +#include +#include +#include + +#include "gtest/gtest.h" + +#include "infini_train/include/autograd/conv.h" +#include "infini_train/include/device.h" +#include "infini_train/include/tensor.h" + +#include "tests/common/test_utils.h" + +using namespace infini_train; + +namespace { +// N=1, Cin=2, Cout=2, H=4, W=5, K=3, stride=2, padding=1 -> out (1,2,2,3) +constexpr int64_t kStride = 2, kPadding = 1; +const std::vector kInputDims = {1, 2, 4, 5}; +const std::vector kWeightDims = {2, 2, 3, 3}; +const std::vector kBiasDims = {2}; +const std::vector kOutDims = {1, 2, 2, 3}; +constexpr float kEps = 1e-2f; +constexpr float kRelTol = 1e-3f; + +// Deterministic pseudo-random values in roughly [-0.5, 0.5]; distinct stream per salt. +float Pattern(size_t idx, int salt) { + int v = static_cast((idx * 37u + static_cast(salt) * 17u) % 11u) - 5; // [-5, 5] + return static_cast(v) * 0.1f; +} +} // namespace + +class AutogradConvGradcheckTest : public infini_train::test::InfiniTrainTest {}; + +TEST_P(AutogradConvGradcheckTest, Conv2dCentralDifference) { + ONLY_CPU(); + const Device device = GetDevice(); + + const size_t nx = 40, nw = 36, nb = 2, ny = 12; + ASSERT_EQ(nx, static_cast(1 * 2 * 4 * 5)); + ASSERT_EQ(nw, static_cast(2 * 2 * 3 * 3)); + ASSERT_EQ(ny, static_cast(1 * 2 * 2 * 3)); + + std::vector x0(nx), w0(nw), b0(nb), go(ny); + for (size_t i = 0; i < nx; ++i) { x0[i] = Pattern(i, 1); } + for (size_t i = 0; i < nw; ++i) { w0[i] = Pattern(i, 2); } + for (size_t i = 0; i < nb; ++i) { b0[i] = Pattern(i, 3); } + for (size_t i = 0; i < ny; ++i) { go[i] = 0.5f + Pattern(i, 4); } + + // Scalar loss for a given parameter setting: dot(forward, go). + // NOTE: the loss is exactly linear in each single perturbed parameter + // (conv is multilinear), so central differences have no truncation error + // and a larger step (1e-2) minimises fp32 cancellation in the difference. + auto loss_fn = [&](const std::vector &x, const std::vector &w, + const std::vector &b) { + auto input = std::make_shared(x.data(), kInputDims, DataType::kFLOAT32, device); + auto weight = std::make_shared(w.data(), kWeightDims, DataType::kFLOAT32, device); + auto bias = std::make_shared(b.data(), kBiasDims, DataType::kFLOAT32, device); + auto fn = std::make_shared(kStride, kPadding); + auto out = fn->Apply({input, weight, bias}); + EXPECT_EQ(out.size(), 1); + auto out_cpu = out[0]->To(Device()); + const auto *p = static_cast(out_cpu.DataPtr()); + double acc = 0.0; + for (size_t i = 0; i < ny; ++i) { acc += static_cast(p[i]) * static_cast(go[i]); } + return acc; + }; + + // Analytic grads via autograd Backward. + auto input = std::make_shared(x0.data(), kInputDims, DataType::kFLOAT32, device)->RequiresGrad(); + auto weight = std::make_shared(w0.data(), kWeightDims, DataType::kFLOAT32, device)->RequiresGrad(); + auto bias = std::make_shared(b0.data(), kBiasDims, DataType::kFLOAT32, device)->RequiresGrad(); + auto fn = std::make_shared(kStride, kPadding); + auto out = fn->Apply({input, weight, bias}); + ASSERT_EQ(out.size(), 1); + EXPECT_EQ(out[0]->Dims(), kOutDims); + auto grad_out = std::make_shared(go.data(), kOutDims, DataType::kFLOAT32, device); + auto grads = fn->Backward({grad_out}); + ASSERT_EQ(grads.size(), 3); + ASSERT_NE(grads[0], nullptr); + ASSERT_NE(grads[1], nullptr); + ASSERT_NE(grads[2], nullptr); + + const std::vector *> analytic = {&grads[0], &grads[1], &grads[2]}; + const std::vector sizes = {nx, nw, nb}; + std::vector> bases = {x0, w0, b0}; + + for (int p = 0; p < 3; ++p) { + auto g_cpu = (*analytic[p])->To(Device()); + const auto *g = static_cast(g_cpu.DataPtr()); + ASSERT_EQ(g_cpu.NumElements(), sizes[p]); + double max_rel = 0.0; + for (size_t i = 0; i < sizes[p]; ++i) { + auto plus = bases[p], minus = bases[p]; + plus[i] += kEps; + minus[i] -= kEps; + // Select which parameter vector is perturbed. + const std::vector &xp = (p == 0 ? plus : bases[0]); + const std::vector &xm = (p == 0 ? minus : bases[0]); + const std::vector &wp = (p == 1 ? plus : bases[1]); + const std::vector &wm = (p == 1 ? minus : bases[1]); + const std::vector &bp = (p == 2 ? plus : bases[2]); + const std::vector &bm = (p == 2 ? minus : bases[2]); + const double numeric = (loss_fn(xp, wp, bp) - loss_fn(xm, wm, bm)) / (2.0 * kEps); + const double a = g[i]; + const double denom = std::max({std::fabs(a), std::fabs(numeric), 1e-8}); + double rel; + if (std::max(std::fabs(a), std::fabs(numeric)) < 1e-5) { + // Both ~0: fall back to absolute error so tiny noise cannot blow up the ratio. + rel = std::fabs(a - numeric) / 1e-5; + } else { + rel = std::fabs(a - numeric) / denom; + } + max_rel = std::max(max_rel, rel); + EXPECT_LT(rel, kRelTol) << "param " << p << " index " << i << " analytic=" << a + << " numeric=" << numeric; + } + EXPECT_LT(max_rel, kRelTol) << "param " << p << " max relative error too large"; + } +} + +INFINI_TRAIN_REGISTER_TEST(AutogradConvGradcheckTest); diff --git a/tests/autograd/test_autograd_conv_im2col.cc b/tests/autograd/test_autograd_conv_im2col.cc new file mode 100644 index 000000000..dd83dc1bd --- /dev/null +++ b/tests/autograd/test_autograd_conv_im2col.cc @@ -0,0 +1,62 @@ +#include +#include + +#include "gtest/gtest.h" + +#include "infini_train/include/dispatcher.h" +#include "infini_train/include/nn/parallel/global.h" +#include "infini_train/include/tensor.h" + +#include "infini_train/src/kernels/common/conv.h" +#include "tests/common/test_utils.h" + +using namespace infini_train; + +class AutogradConvIm2colTest : public infini_train::test::InfiniTrainTest {}; + +namespace { + +std::shared_ptr CallIm2col(const std::shared_ptr &input, const kernels::Conv2dMeta &meta, + Device::DeviceType device_type) { + return Dispatcher::Instance().Call>({device_type, "Im2colForward"}, input, meta); +} + +} // namespace + +TEST_P(AutogradConvIm2colTest, Im2colBasic) { + ONLY_CUDA(); + // input 3x3 = 1..9, kernel 2x2, stride 1, padding 0 -> [N=1, K=4, P=4], + // row k = (kh * Kw + kw), column p = (oh * Wout + ow). + const float input_values[] = {1.0f, 2.0f, 3.0f, 4.0f, 5.0f, 6.0f, 7.0f, 8.0f, 9.0f}; + auto input + = std::make_shared(input_values, std::vector{1, 1, 3, 3}, DataType::kFLOAT32, GetDevice()); + auto weight = std::make_shared(std::vector{1, 1, 2, 2}, DataType::kFLOAT32, GetDevice()); + weight->Fill(0.0f); + const auto meta = kernels::MakeConv2dMeta(input, weight, /*stride=*/1, /*padding=*/0); + auto columns = CallIm2col(input, meta, GetDevice().type()); + ASSERT_NE(columns, nullptr); + EXPECT_EQ(columns->Dims(), (std::vector{1, 4, 4})); + test::ExpectTensorFloatEqual( + columns, {1.0f, 2.0f, 4.0f, 5.0f, 2.0f, 3.0f, 5.0f, 6.0f, 4.0f, 5.0f, 7.0f, 8.0f, 5.0f, 6.0f, 8.0f, 9.0f}); +} + +TEST_P(AutogradConvIm2colTest, Im2colStridePadding) { + ONLY_CUDA(); + // input 4x4 = 1..16, kernel 3x3, stride 2, padding 1 -> [N=1, K=9, P=4]. + std::vector input_values(16); + for (int i = 0; i < 16; ++i) { input_values[i] = static_cast(i + 1); } + auto input = std::make_shared(input_values.data(), std::vector{1, 1, 4, 4}, DataType::kFLOAT32, + GetDevice()); + auto weight = std::make_shared(std::vector{1, 1, 3, 3}, DataType::kFLOAT32, GetDevice()); + weight->Fill(0.0f); + const auto meta = kernels::MakeConv2dMeta(input, weight, /*stride=*/2, /*padding=*/1); + auto columns = CallIm2col(input, meta, GetDevice().type()); + ASSERT_NE(columns, nullptr); + EXPECT_EQ(columns->Dims(), (std::vector{1, 9, 4})); + test::ExpectTensorFloatEqual(columns, + {0.0f, 0.0f, 0.0f, 6.0f, 0.0f, 0.0f, 5.0f, 7.0f, 0.0f, 0.0f, 6.0f, 8.0f, + 0.0f, 2.0f, 0.0f, 10.0f, 1.0f, 3.0f, 9.0f, 11.0f, 2.0f, 4.0f, 10.0f, 12.0f, + 0.0f, 6.0f, 0.0f, 14.0f, 5.0f, 7.0f, 13.0f, 15.0f, 6.0f, 8.0f, 14.0f, 16.0f}); +} + +INFINI_TRAIN_REGISTER_TEST(AutogradConvIm2colTest); diff --git a/tests/autograd/test_autograd_conv_parity.cc b/tests/autograd/test_autograd_conv_parity.cc new file mode 100644 index 000000000..f02069e47 --- /dev/null +++ b/tests/autograd/test_autograd_conv_parity.cc @@ -0,0 +1,105 @@ +// Golden-value parity for Conv2d forward/backward vs PyTorch (CPU). +// +// Config: N=1, Cin=2, Cout=2, H=4, W=5, K=3, stride=2, padding=1, seed=0. +// Goldens generated ONCE from torch (float32, torch 2.14+cpu) via +// /tmp/opencode/conv2d_golden.py (NOT committed; see OSpec P5 notes). +// Threshold: abs err <= 1e-5 (fp32 vs-PyTorch). +#include + +#include "gtest/gtest.h" + +#include "infini_train/include/autograd/conv.h" +#include "infini_train/include/tensor.h" + +#include "tests/common/test_utils.h" + +using namespace infini_train; + +class AutogradConvParityTest : public infini_train::test::InfiniTrainTest {}; + +TEST_P(AutogradConvParityTest, Conv2dGoldenParity) { + ONLY_CPU(); + const Device device = GetDevice(); + + // torch.manual_seed(0) randn inputs. + const std::vector input_vals = { + -1.1258398294f, -1.1523602009f, -0.2505785823f, -0.4338788092f, 0.8487103581f, 0.6920091510f, + -0.3160127699f, -2.1152193546f, 0.3222749233f, -1.2633347511f, 0.3499831855f, 0.3081339300f, + 0.1198415086f, 1.2376579046f, 1.1167771816f, -0.2472781539f, -1.3526537418f, -1.6959311962f, + 0.5666506290f, 0.7935083508f, 0.5988394618f, -1.5550950766f, -0.3413603902f, 1.8530061245f, + -0.2158632576f, -0.7425481677f, 0.5627213717f, 0.2596274018f, -0.1739609987f, -0.6787462234f, + 0.9382607341f, 0.4888698161f, 1.2032237053f, 0.0845346972f, -1.2001394033f, -0.0047857380f, + -0.5180748105f, -0.3067041934f, -1.5809938908f, 1.7066433430f, + }; + const std::vector weight_vals = { + 0.2055256665f, -0.4503297508f, -0.5730770826f, -0.5553584099f, 0.5943230391f, 1.5419425964f, + 0.5073344111f, -0.5910331607f, -1.3253259659f, 0.1885535717f, -0.0690726861f, -0.4949253500f, + -1.4959149361f, -0.1938371211f, 0.4455121756f, 1.3252748251f, 1.5091218948f, 2.0819554329f, + 1.7067116499f, 2.3803675175f, 1.9414620399f, 0.7914980650f, -0.0202518273f, -0.4371695518f, + 1.6458669901f, -1.3601689339f, 0.3445654213f, 0.5198677182f, -0.3656187952f, -1.3024404049f, + 0.0994034633f, 0.4418220222f, 0.2469263971f, 0.0768870041f, 0.3380058110f, 0.4544017613f, + }; + const std::vector bias_vals = {0.1752833277f, -0.9315211177f}; + const std::vector grad_out_vals = { + -1.5054897070f, -0.6609825492f, 1.3232016563f, 0.0371143036f, -0.2849093080f, -0.1334417462f, + 1.8929104805f, 3.1110441685f, -0.4583958089f, -0.3359880745f, -1.5699861050f, 1.2315003872f, + }; + + // PyTorch goldens (F.conv2d + autograd, 10 decimals). + const std::vector expected_forward = { + -3.0188088417f, 4.6534223557f, -2.1541304588f, 1.3896170855f, -2.9418153763f, 1.2058488131f, + -1.5697779655f, 1.0232539177f, 0.8026736379f, -0.3325381279f, -5.7407088280f, -2.4872004986f, + }; + const std::vector expected_dinput = { + -0.9330821037f, -0.3194339275f, -0.4558414817f, -3.4769215584f, 0.7956926227f, -2.5013723373f, + 4.0208745003f, -7.4497237206f, 1.0544195175f, 2.8329558372f, 0.0288622584f, -0.8803023696f, + -0.1375330836f, 1.2958744764f, -0.1042476371f, 0.4350647628f, -2.8934910297f, 2.3038370609f, + 1.7958209515f, -1.5961800814f, 1.1281492710f, 1.0947177410f, 1.5026507378f, -1.5512402058f, + -0.4590149522f, -1.5118727684f, -3.3616752625f, 0.6477436423f, 4.5567674637f, 1.4008895159f, + -0.1556410640f, 0.2037085742f, -0.6384283900f, -0.1925686896f, 0.5699699521f, -0.0575559139f, + -0.5736979246f, -0.9606273174f, -1.3887335062f, 0.2148744166f, + }; + const std::vector expected_dweight = { + 0.0470300466f, 0.7969107032f, -0.1035477147f, -0.0653646588f, 2.8134040833f, 1.6804685593f, + 0.9450824261f, -0.9472144842f, 0.0510890186f, -0.1371109039f, -0.0109563395f, 0.0704481155f, + 3.3292274475f, -1.1093821526f, 1.1104340553f, -0.2435595840f, -0.0923609436f, -0.3009741902f, + 0.8930173516f, 1.5325608253f, -0.3997906148f, -2.3457450867f, -2.2301483154f, -5.5777654648f, + 1.6906189919f, -0.9686450958f, -0.0307304859f, -1.0976977348f, -0.9940003157f, 0.0840486735f, + -6.3507938385f, -3.5117480755f, 2.5241556168f, 0.6967696548f, 2.2981307507f, 3.1801860332f, + }; + const std::vector expected_dbias = {-1.2245074511f, 3.8710851669f}; + constexpr float kAbsTol = 1e-5f; + + auto input = std::make_shared(input_vals.data(), std::vector{1, 2, 4, 5}, + DataType::kFLOAT32, device) + ->RequiresGrad(); + auto weight = std::make_shared(weight_vals.data(), std::vector{2, 2, 3, 3}, + DataType::kFLOAT32, device) + ->RequiresGrad(); + auto bias = std::make_shared(bias_vals.data(), std::vector{2}, DataType::kFLOAT32, + device) + ->RequiresGrad(); + + auto fn = std::make_shared(2, 1); + auto out = fn->Apply({input, weight, bias}); + ASSERT_EQ(out.size(), 1); + EXPECT_EQ(out[0]->Dims(), (std::vector{1, 2, 2, 3})); + test::ExpectTensorNear(out[0], expected_forward, kAbsTol); + + auto grad_out = std::make_shared(grad_out_vals.data(), std::vector{1, 2, 2, 3}, + DataType::kFLOAT32, device); + auto grads = fn->Backward({grad_out}); + ASSERT_EQ(grads.size(), 3); + ASSERT_NE(grads[0], nullptr); + ASSERT_NE(grads[1], nullptr); + ASSERT_NE(grads[2], nullptr); + EXPECT_EQ(grads[0]->Dims(), (std::vector{1, 2, 4, 5})); + EXPECT_EQ(grads[1]->Dims(), (std::vector{2, 2, 3, 3})); + EXPECT_EQ(grads[2]->Dims(), (std::vector{2})); + // Forward logits + full backward grads (dInput spot-check superset, dWeight, dBias). + test::ExpectTensorNear(grads[0], expected_dinput, kAbsTol); + test::ExpectTensorNear(grads[1], expected_dweight, kAbsTol); + test::ExpectTensorNear(grads[2], expected_dbias, kAbsTol); +} + +INFINI_TRAIN_REGISTER_TEST(AutogradConvParityTest); diff --git a/tests/autograd/test_autograd_relu.cc b/tests/autograd/test_autograd_relu.cc new file mode 100644 index 000000000..87a23ac06 --- /dev/null +++ b/tests/autograd/test_autograd_relu.cc @@ -0,0 +1,49 @@ +#include + +#include "gtest/gtest.h" + +#include "infini_train/include/autograd/activations.h" +#include "infini_train/include/nn/modules/activations.h" +#include "infini_train/include/nn/parallel/global.h" +#include "infini_train/include/tensor.h" + +#include "tests/common/test_utils.h" + +using namespace infini_train; + +class AutogradReLUTest : public infini_train::test::InfiniTrainTest {}; + +TEST_P(AutogradReLUTest, ReLUForward) { + std::vector input_values{-2.0f, -0.0f, 0.0f, 0.5f, 1.0f, 3.0f}; + auto a = std::make_shared(input_values.data(), std::vector{2, 3}, DataType::kFLOAT32, GetDevice()); + auto relu_fn = std::make_shared(); + auto result = relu_fn->Apply({a}); + EXPECT_EQ(result.size(), 1); + EXPECT_EQ(result[0]->Dims(), (std::vector{2, 3})); + test::ExpectTensorFloatEqual(result[0], {0.0f, 0.0f, 0.0f, 0.5f, 1.0f, 3.0f}); +} + +TEST_P(AutogradReLUTest, ReLUBackward) { + std::vector input_values{-2.0f, -0.0f, 0.0f, 0.5f, 1.0f, 3.0f}; + auto a = std::make_shared(input_values.data(), std::vector{2, 3}, DataType::kFLOAT32, GetDevice()); + auto relu_fn = std::make_shared(); + auto result = relu_fn->Apply({a}); + std::vector grad_values{1.0f, 2.0f, 3.0f, 4.0f, 5.0f, 6.0f}; + auto grad + = std::make_shared(grad_values.data(), std::vector{2, 3}, DataType::kFLOAT32, GetDevice()); + auto grad_inputs = relu_fn->Backward({grad}); + EXPECT_EQ(grad_inputs.size(), 1); + test::ExpectTensorFloatEqual(grad_inputs[0], {0.0f, 0.0f, 0.0f, 4.0f, 5.0f, 6.0f}); +} + +TEST_P(AutogradReLUTest, ReLUModuleForward) { + std::vector input_values{-1.0f, 0.0f, 2.0f}; + auto a = std::make_shared(input_values.data(), std::vector{3}, DataType::kFLOAT32, GetDevice()); + nn::ReLU relu; + auto result = relu.Forward({a}); + EXPECT_EQ(result.size(), 1); + EXPECT_EQ(result[0]->Dims(), (std::vector{3})); + test::ExpectTensorFloatEqual(result[0], {0.0f, 0.0f, 2.0f}); +} + +INFINI_TRAIN_REGISTER_TEST(AutogradReLUTest); diff --git a/tests/dataloader/CMakeLists.txt b/tests/dataloader/CMakeLists.txt index 7dcf5c887..2d432878d 100644 --- a/tests/dataloader/CMakeLists.txt +++ b/tests/dataloader/CMakeLists.txt @@ -6,3 +6,8 @@ infini_train_add_test(test_dataloader SOURCES test_dataloader.cc LABELS cpu ) + +infini_train_add_test(test_dataloader_stack + SOURCES test_dataloader_stack.cc + LABELS cpu +) diff --git a/tests/dataloader/test_dataloader_stack.cc b/tests/dataloader/test_dataloader_stack.cc new file mode 100644 index 000000000..2ba77bab6 --- /dev/null +++ b/tests/dataloader/test_dataloader_stack.cc @@ -0,0 +1,64 @@ +#include +#include +#include +#include + +#include "gtest/gtest.h" + +#include "infini_train/include/dataloader.h" +#include "infini_train/include/dataset.h" +#include "infini_train/include/tensor.h" + +using namespace infini_train; + +namespace { +// Dataset yielding [1, 2, 2] FLOAT32 samples with values {idx*4, ..., idx*4+3}. +class ChannelDataset : public Dataset { +public: + explicit ChannelDataset(size_t size) : size_(size) {} + + std::pair, std::shared_ptr> operator[](size_t idx) const override { + auto data = std::make_shared(std::vector{1, 2, 2}, DataType::kFLOAT32); + auto label = std::make_shared(std::vector{1}, DataType::kINT64); + auto *data_ptr = static_cast(data->DataPtr()); + for (int i = 0; i < 4; ++i) { data_ptr[i] = static_cast(idx * 4 + i); } + *static_cast(label->DataPtr()) = static_cast(idx); + return {data, label}; + } + + size_t Size() const override { return size_; } + +private: + size_t size_ = 0; +}; + +// Dataset whose sample 1 has a different shape ([1, 2, 3]) than sample 0 ([1, 2, 2]). +class RaggedDataset : public Dataset { +public: + std::pair, std::shared_ptr> operator[](size_t idx) const override { + auto data = std::make_shared(idx == 0 ? std::vector{1, 2, 2} : std::vector{1, 2, 3}, + DataType::kFLOAT32); + auto label = std::make_shared(std::vector{1}, DataType::kINT64); + return {data, label}; + } + + size_t Size() const override { return 2; } +}; +} // namespace + +TEST(DataLoaderStackTest, PreservesSampleDims) { + DataLoader loader(std::make_shared(2), 2); + auto it = loader.begin(); + const auto &[x, y] = *it; + EXPECT_EQ(x->Dims(), (std::vector{2, 1, 2, 2})); + EXPECT_EQ(x->Dtype(), DataType::kFLOAT32); + const auto *data = static_cast(x->DataPtr()); + for (int i = 0; i < 8; ++i) { EXPECT_FLOAT_EQ(data[i], static_cast(i)); } + EXPECT_EQ(y->Dims(), (std::vector{2, 1})); +} + +TEST(DataLoaderStackTest, MismatchedShapesFail) { + DataLoader loader(std::make_shared(), 2); + auto it = loader.begin(); + EXPECT_DEATH({ const auto batch = *it; }, "different shapes"); +} diff --git a/tests/example/CMakeLists.txt b/tests/example/CMakeLists.txt new file mode 100644 index 000000000..a589a8ff7 --- /dev/null +++ b/tests/example/CMakeLists.txt @@ -0,0 +1,17 @@ +# ========================================================================== +# Example tests (MNIST dataset, ...) +# ========================================================================== + +infini_train_add_test(test_mnist_dataset + SOURCES test_mnist_dataset.cc ${CMAKE_SOURCE_DIR}/example/mnist/dataset.cc + LABELS cpu + TEST_TIMEOUT 60 +) + +# MnistCnn <-> torch parity fixture (OSpec P12): dumps a deterministic 10-step +# SGD trajectory to $MNIST_PARITY_OUT_DIR (or ./mnist_parity_dump). +infini_train_add_test(test_mnist_parity + SOURCES mnist_parity.cc ${CMAKE_SOURCE_DIR}/example/mnist/net.cc + LABELS cpu + TEST_TIMEOUT 300 +) diff --git a/tests/example/mnist_parity.cc b/tests/example/mnist_parity.cc new file mode 100644 index 000000000..82d179acb --- /dev/null +++ b/tests/example/mnist_parity.cc @@ -0,0 +1,262 @@ +// MnistCnn <-> torch parity fixture (OSpec P12). +// +// Dumps a fully deterministic 10-step SGD trajectory of MnistCnn +// (Conv2d(1,16,3)->ReLU->Conv2d(16,32,3)->ReLU->Flatten->Linear(18432,10)) +// on a fixed synthetic batch so an out-of-repo torch script can compare +// logits / loss / grads / updated params at 1e-5. +// +// Layout (all little-endian, row-major / C-contiguous flat order): +// meta.txt human-readable shapes, dtypes, hyperparams, file map +// input.bin float32 [B,1,28,28] +// labels.bin uint8 [B] +// init__.bin float32, initial param `` (e.g. conv1.weight) +// step__logits.bin float32 [B,10], NNN = 000..009 +// step__loss.bin float32 scalar (1 element, mean CE loss) +// step__grad__.bin float32, grad of `` BEFORE the SGD step +// step__param__.bin float32, value of `` AFTER the SGD step +// +// Determinism notes: +// * Params are NOT the constructor RNG values: they are refilled here with a +// single-threaded std::mt19937(42), U(-bound,bound), bound = 1/sqrt(fan_in), +// consumed in sorted NamedParameters order. This matches the torch default +// init math (kaiming_uniform_(a=sqrt(5)) == U(-1/sqrt(fan_in),1/sqrt(fan_in))) +// while staying independent of the framework global RNG / OMP thread count. +// * Input uses std::mt19937(1234) U[0,1); labels use std::mt19937(5678) % 10. +// * Output dir: $MNIST_PARITY_OUT_DIR or ./mnist_parity_dump. +// +// Run: MNIST_PARITY_OUT_DIR=/tmp/opencode/mnist_parity \ +// ./build/tests/example/test_mnist_parity + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#include "gtest/gtest.h" + +#include "Eigen/Dense" + +#ifdef USE_OMP +#include +#endif + +#include "example/mnist/net.h" +#include "infini_train/include/datatype.h" +#include "infini_train/include/device.h" +#include "infini_train/include/nn/modules/loss.h" +#include "infini_train/include/optimizer.h" +#include "infini_train/include/tensor.h" + +namespace { + +constexpr int kBatch = 4; +constexpr int kSteps = 10; +constexpr float kLr = 0.01f; +constexpr uint32_t kParamSeed = 42; +constexpr uint32_t kInputSeed = 1234; +constexpr uint32_t kLabelSeed = 5678; + +std::string OutDir() { + if (const char *env = std::getenv("MNIST_PARITY_OUT_DIR")) { + if (env[0] != '\0') { + return env; + } + } + return "mnist_parity_dump"; +} + +std::string StepTag(int step) { + std::ostringstream oss; + oss << "step" << std::setw(3) << std::setfill('0') << step; + return oss.str(); +} + +void WriteFloats(const std::filesystem::path &path, const float *data, size_t n) { + std::ofstream ofs(path, std::ios::binary); + ASSERT_TRUE(ofs.good()) << path; + ofs.write(reinterpret_cast(data), n * sizeof(float)); + ASSERT_TRUE(ofs.good()) << path; +} + +void WriteBytes(const std::filesystem::path &path, const uint8_t *data, size_t n) { + std::ofstream ofs(path, std::ios::binary); + ASSERT_TRUE(ofs.good()) << path; + ofs.write(reinterpret_cast(data), n * sizeof(uint8_t)); + ASSERT_TRUE(ofs.good()) << path; +} + +void DumpTensorFp32(const std::filesystem::path &path, const std::shared_ptr &t) { + ASSERT_NE(t, nullptr) << path; + ASSERT_EQ(t->Dtype(), infini_train::DataType::kFLOAT32) << path; + // Fixture runs on CPU; tensors are CPU already. + WriteFloats(path, static_cast(t->DataPtr()), t->NumElements()); +} + +std::string DimsStr(const std::vector &dims) { + std::ostringstream oss; + oss << "["; + for (size_t i = 0; i < dims.size(); ++i) { + oss << (i ? "," : "") << dims[i]; + } + oss << "]"; + return oss.str(); +} + +// fan_in for bound = 1/sqrt(fan_in): weight [O,I,...] -> I*prod(rest); +// bias "m.bias" reuses sibling "m.weight" fan_in. +int64_t FanIn(const std::string &name, const std::vector &dims, + const std::vector>> &named) { + if (dims.size() >= 2) { + int64_t fan = dims[1]; + for (size_t i = 2; i < dims.size(); ++i) { fan *= dims[i]; } + return fan; + } + const auto pos = name.rfind('.'); + if (pos != std::string::npos) { + const std::string wname = name.substr(0, pos) + ".weight"; + for (const auto &[n, t] : named) { + if (n == wname) { return FanIn(wname, t->Dims(), named); } + } + } + return dims.empty() ? 1 : dims[0]; +} + +bool AllFinite(const float *data, size_t n) { + for (size_t i = 0; i < n; ++i) { + if (!std::isfinite(data[i])) { + return false; + } + } + return true; +} + +TEST(MnistParity, DumpTrajectory) { + namespace it = infini_train; + // Pin single-threaded execution: Eigen/OpenMP parallel reductions change FP + // summation order with the thread count (~2e-8 jitter), which would make the + // dumped trajectory environment-dependent. +#ifdef USE_OMP + omp_set_num_threads(1); +#endif + Eigen::setNbThreads(1); + + const std::filesystem::path out = OutDir(); + std::filesystem::create_directories(out); + + auto net = std::make_shared(); + auto named = net->NamedParameters(); + std::sort(named.begin(), named.end(), + [](const auto &a, const auto &b) { return a.first < b.first; }); + ASSERT_EQ(named.size(), 6); + EXPECT_EQ(named[0].first, "conv1.bias"); + EXPECT_EQ(named[1].first, "conv1.weight"); + EXPECT_EQ(named[2].first, "conv2.bias"); + EXPECT_EQ(named[3].first, "conv2.weight"); + EXPECT_EQ(named[4].first, "fc.bias"); + EXPECT_EQ(named[5].first, "fc.weight"); + + // Deterministic init refill (single-threaded mt19937, sorted order). + { + std::mt19937 gen(kParamSeed); + for (const auto &[name, param] : named) { + ASSERT_EQ(param->Dtype(), it::DataType::kFLOAT32) << name; + const int64_t fan_in = FanIn(name, param->Dims(), named); + ASSERT_GT(fan_in, 0) << name; + const float bound = 1.0f / std::sqrt(static_cast(fan_in)); + std::uniform_real_distribution dis(-bound, bound); + float *dst = static_cast(param->DataPtr()); + for (size_t i = 0; i < param->NumElements(); ++i) { dst[i] = dis(gen); } + } + } + + // Fixed synthetic input: U[0,1) like normalized MNIST pixels. + std::vector input_data(kBatch * 1 * 28 * 28); + { + std::mt19937 gen(kInputSeed); + std::uniform_real_distribution dis(0.0f, 1.0f); + for (auto &v : input_data) { v = dis(gen); } + } + auto input = std::make_shared(std::vector{kBatch, 1, 28, 28}, it::DataType::kFLOAT32, + it::Device()); + std::copy(input_data.begin(), input_data.end(), static_cast(input->DataPtr())); + + // Fixed labels. + std::vector label_data(kBatch); + { + std::mt19937 gen(kLabelSeed); + for (auto &v : label_data) { v = static_cast(gen() % 10); } + } + auto labels = std::make_shared(std::vector{kBatch}, it::DataType::kUINT8, it::Device()); + std::copy(label_data.begin(), label_data.end(), static_cast(labels->DataPtr())); + + auto loss_fn = std::make_shared(); + auto optimizer = it::optimizers::SGD(net->Parameters(), kLr); + + // meta.txt first (shapes/dtypes/hparams), then blobs. + { + std::ofstream meta(out / "meta.txt"); + ASSERT_TRUE(meta.good()); + meta << "# MnistCnn parity fixture (OSpec P12)\n"; + meta << "model: Conv2d(1,16,3,s1,p0)+ReLU+Conv2d(16,32,3,s1,p0)+ReLU+Flatten(1)+Linear(18432,10)\n"; + meta << "batch: " << kBatch << "\nsteps: " << kSteps << "\nloss: CrossEntropy(mean)\noptimizer: SGD lr=" + << kLr << "\n"; + meta << "seeds: param=" << kParamSeed << " input=" << kInputSeed << " label=" << kLabelSeed << "\n"; + meta << "input: float32 " << DimsStr(input->Dims()) << " input.bin\n"; + meta << "labels: uint8 " << DimsStr(labels->Dims()) << " labels.bin values="; + for (int i = 0; i < kBatch; ++i) { meta << (i ? "," : "") << static_cast(label_data[i]); } + meta << "\n"; + for (const auto &[name, param] : named) { + meta << "param " << name << ": float32 " << DimsStr(param->Dims()) << " init__" << name + << ".bin step__grad__" << name << ".bin step__param__" << name << ".bin\n"; + } + meta << "logits: float32 [B,10] step__logits.bin\nloss: float32 scalar step__loss.bin\n"; + ASSERT_TRUE(meta.good()); + } + + WriteFloats(out / "input.bin", input_data.data(), input_data.size()); + WriteBytes(out / "labels.bin", label_data.data(), label_data.size()); + for (const auto &[name, param] : named) { DumpTensorFp32(out / ("init__" + name + ".bin"), param); } + + std::vector losses; + for (int step = 0; step < kSteps; ++step) { + auto outputs = net->Forward({input}); + ASSERT_EQ(outputs.size(), 1); + ASSERT_EQ(outputs[0]->Dims(), (std::vector{kBatch, 10})) << "step " << step; + optimizer.ZeroGrad(); + auto loss = loss_fn->Forward({outputs[0], labels}); + ASSERT_EQ(loss.size(), 1); + loss[0]->Backward(); + + const float *logit_ptr = static_cast(outputs[0]->DataPtr()); + const float loss_val = static_cast(loss[0]->DataPtr())[0]; + ASSERT_TRUE(AllFinite(logit_ptr, outputs[0]->NumElements())) << "step " << step; + ASSERT_TRUE(std::isfinite(loss_val)) << "step " << step; + losses.push_back(loss_val); + + const std::string tag = StepTag(step); + DumpTensorFp32(out / (tag + "__logits.bin"), outputs[0]); + WriteFloats(out / (tag + "__loss.bin"), &loss_val, 1); + for (const auto &[name, param] : named) { + ASSERT_NE(param->grad(), nullptr) << "step " << step << " " << name; + DumpTensorFp32(out / (tag + "__grad__" + name + ".bin"), param->grad()); + } + + optimizer.Step(); + for (const auto &[name, param] : named) { DumpTensorFp32(out / (tag + "__param__" + name + ".bin"), param); } + + std::cout << "[parity] " << tag << " loss=" << loss_val << "\n"; + } + std::cout << "[parity] dumped " << (3 + named.size()) << " init files + " << kSteps * (2 + 2 * named.size()) + << " step files to " << out << "\n"; +} + +} // namespace diff --git a/tests/example/test_mnist_dataset.cc b/tests/example/test_mnist_dataset.cc new file mode 100644 index 000000000..2f242d09f --- /dev/null +++ b/tests/example/test_mnist_dataset.cc @@ -0,0 +1,113 @@ +#include +#include +#include +#include +#include +#include + +#include "gtest/gtest.h" + +#include "example/mnist/dataset.h" +#include "infini_train/include/datatype.h" +#include "infini_train/include/tensor.h" + +namespace { + +std::string ResolveMnistDir() { + if (const char *env = std::getenv("MNIST_DATA_DIR")) { + if (std::filesystem::exists(env)) { + return env; + } + } + // Read-only shared copy (see task description); never copied into the worktree. + constexpr char kSharedPath[] = "/home/Sota/infra/InfiniTrain/data/mnist"; + if (std::filesystem::exists(kSharedPath)) { + return kSharedPath; + } + return ""; +} + +uint32_t ReadBigEndianU32(std::ifstream &ifs) { + uint8_t bytes[4]; + ifs.read(reinterpret_cast(bytes), 4); + return (static_cast(bytes[0]) << 24) | (static_cast(bytes[1]) << 16) + | (static_cast(bytes[2]) << 8) | static_cast(bytes[3]); +} + +// Raw IDX pixel bytes for sample `idx`, read straight from disk (no framework code involved). +std::vector ReadRawImageBytes(const std::string &dir, size_t idx) { + std::ifstream ifs(dir + "/train-images-idx3-ubyte", std::ios::binary); + EXPECT_TRUE(ifs.good()); + ifs.seekg(0); + EXPECT_EQ(ReadBigEndianU32(ifs), 0x00000803u); + const uint32_t num_images = ReadBigEndianU32(ifs); + const uint32_t rows = ReadBigEndianU32(ifs); + const uint32_t cols = ReadBigEndianU32(ifs); + EXPECT_LT(idx, num_images); + EXPECT_EQ(rows, 28); + EXPECT_EQ(cols, 28); + std::vector raw(rows * cols); + ifs.seekg(16 + idx * rows * cols); + ifs.read(reinterpret_cast(raw.data()), raw.size()); + EXPECT_TRUE(ifs.good()); + return raw; +} + +uint8_t ReadRawLabelByte(const std::string &dir, size_t idx) { + std::ifstream ifs(dir + "/train-labels-idx1-ubyte", std::ios::binary); + EXPECT_TRUE(ifs.good()); + EXPECT_EQ(ReadBigEndianU32(ifs), 0x00000801u); + const uint32_t num_labels = ReadBigEndianU32(ifs); + EXPECT_LT(idx, num_labels); + uint8_t label = 0; + ifs.seekg(8 + idx); + ifs.read(reinterpret_cast(&label), 1); + EXPECT_TRUE(ifs.good()); + return label; +} + +void ExpectSampleMatchesRaw(const MNISTDataset &dataset, const std::string &dir, size_t idx) { + const auto [image, label] = dataset[idx]; + const std::vector raw = ReadRawImageBytes(dir, idx); + ASSERT_EQ(image->NumElements(), raw.size()); + const auto *data = static_cast(image->DataPtr()); + for (size_t i = 0; i < raw.size(); ++i) { + EXPECT_FLOAT_EQ(data[i], raw[i] / 255.0f) << "idx=" << idx << " pixel=" << i; + } + ASSERT_EQ(label->NumElements(), 1); + EXPECT_EQ(*static_cast(label->DataPtr()), ReadRawLabelByte(dir, idx)) << "idx=" << idx; +} + +TEST(MNISTDatasetTest, SampleDimsAndDtypes) { + const std::string dir = ResolveMnistDir(); + if (dir.empty()) { + GTEST_SKIP() << "MNIST data not found (set MNIST_DATA_DIR)"; + } + MNISTDataset dataset(dir, true); + const auto [image, label] = dataset[0]; + EXPECT_EQ(image->Dims(), (std::vector{1, 28, 28})); + EXPECT_EQ(image->Dtype(), infini_train::DataType::kFLOAT32); + EXPECT_EQ(label->Dtype(), infini_train::DataType::kUINT8); + EXPECT_EQ(label->NumElements(), 1); +} + +TEST(MNISTDatasetTest, NormalizationMatchesRawBytes) { + const std::string dir = ResolveMnistDir(); + if (dir.empty()) { + GTEST_SKIP() << "MNIST data not found (set MNIST_DATA_DIR)"; + } + MNISTDataset dataset(dir, true); + ExpectSampleMatchesRaw(dataset, dir, 0); +} + +TEST(MNISTDatasetTest, NonZeroIndicesAreCorrectlyAligned) { + const std::string dir = ResolveMnistDir(); + if (dir.empty()) { + GTEST_SKIP() << "MNIST data not found (set MNIST_DATA_DIR)"; + } + MNISTDataset dataset(dir, true); + // idx > 0 regresses the byte-stride bug (UINT8 stride used on a FLOAT32 buffer). + ExpectSampleMatchesRaw(dataset, dir, 1); + ExpectSampleMatchesRaw(dataset, dir, 100); +} +} // namespace From ed8d51a70583c2664cecaaaf5309bbfc4d465e7e Mon Sep 17 00:00:00 2001 From: opencode Date: Sat, 19 Sep 2026 11:03:13 +0800 Subject: [PATCH 2/8] fix: move ZeroGrad before DDP Forward in MNIST loop Reducer::PrepareForBackward (inside DDP Forward, gradient_as_bucket_view=true) binds param.grad to the bucket view; ZeroGrad(set_to_none=true) after Forward reset that binding, so backward accumulated into a standalone grad the reducer never all-reduced (silent no-sync, cross-rank weight fork). Same order as gpt2/mixtral/llama3 demos. --- example/mnist/main.cc | 7 ++++++- 1 file changed, 6 insertions(+), 1 deletion(-) diff --git a/example/mnist/main.cc b/example/mnist/main.cc index 2014b434c..01cf3fc1f 100644 --- a/example/mnist/main.cc +++ b/example/mnist/main.cc @@ -121,8 +121,13 @@ int main(int argc, char *argv[]) { auto new_image = std::make_shared(image->To(device)); auto new_label = std::make_shared(label->To(device)); - auto outputs = network->Forward({new_image}); + // NOTE: ZeroGrad must run BEFORE DDP Forward. Reducer::PrepareForBackward() + // (inside DDP Forward, gradient_as_bucket_view=true) binds param.grad to the + // bucket view; ZeroGrad(set_to_none=true) after Forward would reset that + // binding, so backward would accumulate into a standalone grad that the + // reducer never all-reduces (silent no-sync, cross-rank weight fork). optimizer.ZeroGrad(); + auto outputs = network->Forward({new_image}); auto loss = loss_fn->Forward({outputs[0], new_label}); loss[0]->Backward(); From b0307ece3aaf6d3a309c01f2c43f639c6189bbd5 Mon Sep 17 00:00:00 2001 From: opencode Date: Sat, 19 Sep 2026 11:24:08 +0800 Subject: [PATCH 3/8] docs: record P13b dual-GPU equivalence PASS and ZeroGrad fix --- docs/small_model_patch.md | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/docs/small_model_patch.md b/docs/small_model_patch.md index b52f9ebd1..8313df3cb 100644 --- a/docs/small_model_patch.md +++ b/docs/small_model_patch.md @@ -138,7 +138,7 @@ MNIST 四个 IDX 文件可按字节数 + 首 16 字节 magic 校验（images 0x0 | P11 e2e smoke（1 epoch） | mlp acc 0.60；cnn cpu acc 0.90、loss 2.3→0.06；cnn cuda 与 CPU bit-identical，2.3s/epoch | | P12 网络级 parity | 140/140 @1e-5（worst 9.5e-7）；10-step loss 双边严格单调递减 | | P13a DDP（main.cc only） | 单进程 cuda acc 0.9028；NCCL-ON 可编译（需手动 `-I/-L/-rpath`，顶层 CMake 未加 NCCL include 路径）；2 进程止步于 `SetDevice(cuda:1)`——单卡机上限 | -| P13b 双卡等价 | BLOCKED：无双卡机，已留诚实 blocker 记录 | +| P13b 双卡等价 | PASS（口径：`Step` 后 cross-rank 权重一致）：RTX 4090D ×2，CNN，SGD lr=0.01，3 epoch，`1×128` vs `2×64`，`post-bcast w0` 两 rank 均为 `-0.0836399`，`post-epoch w0` 两 rank 均为 `-0.001326`（修前 `-0.00206 vs -0.00215` 分叉）。loss：单卡 `0.596128/0.325636/0.301542`（acc 0.9123），双卡 `0.594806/0.322919/0.298881`（acc 0.9018，eval 9984 样本）；残差来自 floor 掉尾（训练少 96/epoch、eval 少 16）与 rank0-local 打印口径，非同步问题。根因：`example/mnist/main.cc` 曾用 `Forward → ZeroGrad(true)`，清掉 `PrepareForBackward` 建立的 bucket-view 绑定致静默不同步；fix（`ed8d51a`）改为 `ZeroGrad → Forward`，与 gpt2/mixtral/llama3 对齐。限定：单机双卡；多机/更大 world size/grad-accum 未测；train loss 打印仍为 rank0-local | | 全量 ctest（NCCL-OFF） | 354/355 通过；唯一失败 `stacktrace`（glog 符号化环境问题，本分支未触及该文件，属 HEAD 预存失败）；CUDA/CPU 交叉实例化 SKIP 系 `ONLY_*` 宏按设计跳过 | HEAD 预存 bug 修复（保留，与上游 #220 一致）：`main.cc` 改用 `shared_ptr` 持有 Module（Module 是 `enable_shared_from_this`，栈实例触发 `bad_weak_ptr` 崩溃）。 From 8770716583f3615f32265f74d9a9135bc7ccbf51 Mon Sep 17 00:00:00 2001 From: opencode Date: Sat, 19 Sep 2026 11:32:13 +0800 Subject: [PATCH 4/8] fix: route MNIST loops through Module operator() and guard eval with NoGradGuard Call modules via operator() (the hook entry, same as gpt2/llama3 and PRs 219/220) instead of Forward() directly. Wrap evaluation in NoGradGuard so forward-only graphs never prime grad accumulators (PR 220 acfbe87 hazard); resolves TODO(dcj) no_grad(). --- example/mnist/main.cc | 15 ++++++++++----- 1 file changed, 10 insertions(+), 5 deletions(-) diff --git a/example/mnist/main.cc b/example/mnist/main.cc index 01cf3fc1f..6fcaa2788 100644 --- a/example/mnist/main.cc +++ b/example/mnist/main.cc @@ -9,6 +9,8 @@ #include "gflags/gflags.h" #include "glog/logging.h" +#include "infini_train/include/autograd/grad_mode.h" + #include "infini_train/include/dataloader.h" #include "infini_train/include/device.h" #include "infini_train/include/nn/modules/loss.h" @@ -127,9 +129,9 @@ int main(int argc, char *argv[]) { // binding, so backward would accumulate into a standalone grad that the // reducer never all-reduces (silent no-sync, cross-rank weight fork). optimizer.ZeroGrad(); - auto outputs = network->Forward({new_image}); + auto outputs = (*network)({new_image}); - auto loss = loss_fn->Forward({outputs[0], new_label}); + auto loss = (*loss_fn)({outputs[0], new_label}); loss[0]->Backward(); // Defer the loss D2H copy until after backward; reading it earlier would synchronize CUDA @@ -159,7 +161,10 @@ int main(int argc, char *argv[]) { } } - // TODO(dcj): Add no_grad() context manager later. + // Evaluation builds forward-only graphs; keep it under NoGradGuard so it never + // primes grad accumulators with a dependency count the next backward cannot satisfy + // (which would silently stop gradient accumulation). Resolves TODO(dcj) no_grad(). + autograd::NoGradGuard no_grad; std::vector test_losses; int correct = 0; int total = 0; @@ -170,9 +175,9 @@ int main(int argc, char *argv[]) { auto new_label = std::make_shared(label->To(device)); auto label_cpu = label->To(cpu_device); - auto outputs = network->Forward({new_image}); + auto outputs = (*network)({new_image}); auto output_cpu = outputs[0]->To(cpu_device); - auto loss = loss_fn->Forward({outputs[0], new_label}); + auto loss = (*loss_fn)({outputs[0], new_label}); auto loss_cpu = loss[0]->To(cpu_device); const int batch_size = output_cpu.Dims()[0]; From b010284528eca8cb7db102a23232d7c74f932381 Mon Sep 17 00:00:00 2001 From: opencode Date: Sat, 19 Sep 2026 11:41:30 +0800 Subject: [PATCH 5/8] docs: restructure P13b result into tables with final dual-GPU confirmation --- docs/small_model_patch.md | 31 ++++++++++++++++++++++++++++++- 1 file changed, 30 insertions(+), 1 deletion(-) diff --git a/docs/small_model_patch.md b/docs/small_model_patch.md index 8313df3cb..6604b4117 100644 --- a/docs/small_model_patch.md +++ b/docs/small_model_patch.md @@ -138,9 +138,38 @@ MNIST 四个 IDX 文件可按字节数 + 首 16 字节 magic 校验（images 0x0 | P11 e2e smoke（1 epoch） | mlp acc 0.60；cnn cpu acc 0.90、loss 2.3→0.06；cnn cuda 与 CPU bit-identical，2.3s/epoch | | P12 网络级 parity | 140/140 @1e-5（worst 9.5e-7）；10-step loss 双边严格单调递减 | | P13a DDP（main.cc only） | 单进程 cuda acc 0.9028；NCCL-ON 可编译（需手动 `-I/-L/-rpath`，顶层 CMake 未加 NCCL include 路径）；2 进程止步于 `SetDevice(cuda:1)`——单卡机上限 | -| P13b 双卡等价 | PASS（口径：`Step` 后 cross-rank 权重一致）：RTX 4090D ×2，CNN，SGD lr=0.01，3 epoch，`1×128` vs `2×64`，`post-bcast w0` 两 rank 均为 `-0.0836399`，`post-epoch w0` 两 rank 均为 `-0.001326`（修前 `-0.00206 vs -0.00215` 分叉）。loss：单卡 `0.596128/0.325636/0.301542`（acc 0.9123），双卡 `0.594806/0.322919/0.298881`（acc 0.9018，eval 9984 样本）；残差来自 floor 掉尾（训练少 96/epoch、eval 少 16）与 rank0-local 打印口径，非同步问题。根因：`example/mnist/main.cc` 曾用 `Forward → ZeroGrad(true)`，清掉 `PrepareForBackward` 建立的 bucket-view 绑定致静默不同步；fix（`ed8d51a`）改为 `ZeroGrad → Forward`，与 gpt2/mixtral/llama3 对齐。限定：单机双卡；多机/更大 world size/grad-accum 未测；train loss 打印仍为 rank0-local | +| P13b 双卡等价 | PASS。判据为 `Step` 后 cross-rank 权重一致（详见下文复测结论小节） | | 全量 ctest（NCCL-OFF） | 354/355 通过；唯一失败 `stacktrace`（glog 符号化环境问题，本分支未触及该文件，属 HEAD 预存失败）；CUDA/CPU 交叉实例化 SKIP 系 `ONLY_*` 宏按设计跳过 | +### P13b 复测结论（2026-09-19，RTX 4090D ×2） + +配置：CNN，SGD lr=0.01，3 epoch，`1×128` vs `2×64`。 + +**主判据（权重一致性）：** + +| 探针 | rank0 | rank1 | 结论 | +| --- | --- | --- | --- | +| `post-bcast w0` | `-0.0836399` | `-0.0836399` | 初值一致 | +| `post-epoch w0`（修前） | `-0.00206` | `-0.00215` | 分叉 | +| `post-epoch w0`（修后） | `-0.001326` | `-0.001326` | 一致 ✅ | + +**Loss 对照（参考，非主判据）：** + +| 配置 | epoch0 / 1 / 2 loss | test acc | +| --- | --- | --- | +| 1卡 bs=128 | 0.596128 / 0.325636 / 0.301542 | 0.9123（10000） | +| 2卡 2×64 | 0.594806 / 0.322919 / 0.298881 | 0.9018（9984） | + +残差来自 floor 掉尾（训练每 epoch 少 96 样本、eval 少 16 样本）与 rank0-local 打印口径，非同步问题。 + +**根因与 fix：** + +- 根因：`example/mnist/main.cc` 曾用 `Forward → ZeroGrad(true)`，清掉 `PrepareForBackward` 建立的 bucket-view 绑定，致静默不同步。 +- `ed8d51a`：改为 `ZeroGrad → Forward`，与 gpt2/mixtral/llama3 对齐。 +- `8770716`：训练/eval 改走 `Module::operator()`，eval 加 `NoGradGuard`（抄 PR 219/220）；双卡复测行为零变化（上表即含该改动的结果）。 + +**限定：** 单机双卡；多机、更大 world size、grad-accum 未测；train loss 打印仍为 rank0-local。 + HEAD 预存 bug 修复（保留，与上游 #220 一致）：`main.cc` 改用 `shared_ptr` 持有 Module（Module 是 `enable_shared_from_this`，栈实例触发 `bad_weak_ptr` 崩溃）。 ENV-ONLY 说明：gcc16 本地构建曾需在 `tensor.cc` / `profiler.cc` 各加一行 `#include `；上游 gcc13 不受影响，该两行已 revert，不在 PR 内。 From b0f64ec51d110f3e71950c2dfa3b39f9fe704720 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E7=8B=AC=E8=A7=92=E4=BB=99=E4=B8=8E=E6=98=9F?= Date: Mon, 21 Sep 2026 15:53:00 +0800 Subject: [PATCH 6/8] =?UTF-8?q?docs:=20=E4=B8=AD=E6=96=87=E9=A1=B9?= =?UTF-8?q?=E7=9B=AE=E6=8A=A5=E5=91=8A(a/b/c)+PyTorch=E5=AF=B9=E9=BD=90?= =?UTF-8?q?=E8=84=9A=E6=9C=AC?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- docs/small_model_patch.md | 310 ++++++++++++++++++++++------------ scripts/mnist_parity_torch.py | 212 +++++++++++++++++++++++ 2 files changed, 415 insertions(+), 107 deletions(-) create mode 100644 scripts/mnist_parity_torch.py diff --git a/docs/small_model_patch.md b/docs/small_model_patch.md index 6604b4117..bcb10c4d4 100644 --- a/docs/small_model_patch.md +++ b/docs/small_model_patch.md @@ -1,128 +1,123 @@ -# 小模型训练支持 — 技术方案与结果摘要（分支 `work/small-model-cnn-impl`） +# 小模型训练支持 — 项目报告（分支 `work/small-model-cnn-impl`） -> 来源：本文件 §5–§6 逐字移植自上游提案文档 `/docs/small_model_patch.md` -> （只读来源，未修改）；§7 为本分支实测结果摘要。 +本报告对应 PR「【训练营】小模型训练支持」：功能代码 + CMake 集成 + 单元测试 + MNIST CNN 演示。 +PyTorch 对齐脚本见 `scripts/mnist_parity_torch.py`，C++ dump 端见 `tests/example/mnist_parity.cc`。 --- -## 五、技术方案（v1 锁定） +## 一、设计与实现 -结合上游 PR #220（kernel / 验证 / PR discipline）与 #219（数据形状语义 / DDP metric 设计）的对比结论，本项目做第三版：**用 #219 的语义设计 + #220 的 kernel 与验证方法**，并去掉两边共有的 demo-level special case。 +### 1.1 整体方案 -### 5.1 Conv2d v1 scope +目标是在 InfiniTrain 上跑通小模型（MNIST CNN）端到端训练，同时保证数学正确性可审计。 +方案是第三版取舍：采用 PR #219 的语义设计（数据形状语义、DDP metric 设计）+ PR #220 的 kernel 与验证方法（kernel 实现、验证 discipline），去掉两边共有的 demo-level special case。 + +整体训练链路： ```text -dtype: FP32 only（入口显式 CHECK，不静默 fallback） -layout: input NCHW，weight OIHW -API: scalar kernel_size / stride / padding + optional bias -内部语义: Kh/Kw 独立（读 weight.Dims()[2/3]），不加 Kh == Kw 限制 -不支持: tuple kernel/stride/padding、dilation、groups、FP16/BF16 +MNISTDataset（单样本 [1,28,28]） + → DataLoader Stack（shape-preserving，batch 后 [B,1,28,28]） + → MnistCnn 前向（Conv→ReLU→Conv→ReLU→Flatten→Linear） + → CrossEntropyLoss（mean） + → Backward（autograd selective-save） + → SGD 更新（lr=0.01） + → 可选 DDP（Broadcast → wrap，训练只同步梯度；eval 分片 + epoch 级 AllReduce） ``` -即：**public API 收敛到任务 scope，internal semantics 不做无意义限制**。非 square 权重只在 kernel / autograd 层单测覆盖，Module 层不暴露 tuple API。 - -### 5.2 双后端职责划分 +Oracle 策略（双独立 oracle，不依赖 CUDA direct 中间实现）： ```text -CPU: direct reference implementation（forward + dX + dW + db），暂不用 Eigen / im2col -CUDA: im2col + framework GEMM（Dispatcher "Gemm"），gather 风格 col2im，bias reduction + PyTorch + │ + ▼ +CPU direct ─────── CUDA GEMM + │ │ + └── finite diff ────┘ ``` -CPU 与现有 `cpu/matmul.cc` 朴素循环风格一致，职责是可审计的数学基准；CUDA 复用 `cuda/common/gemm.cu` 的 strided-batched GEMM，与 `cuda/matmul.cu`、`cuda/linear.cu` 的调用方式一致。两条独立实现路径互为 cross-check，避免 im2col layout 错误在两边同时成立。 - -`Conv2dMeta`（common 层）集中保存 stride / padding / batch / 通道 / `kernel_h,w` / `output_h,w` / `patches` / `kernel_elems`，shape arithmetic 全局只有一份，CPU / CUDA 不再重复定义。 +即 CPU numerical gradcheck（tiny tensor，dX/dW/db）+ PyTorch 对标（logits / loss / grad / 单步更新 / 多步趋势）。 -Autograd 沿用两边已验证的 selective-save（需 dW 才存 input，需 dX 才存 weight），`Conv2dBackwardInput / Weight / Bias` 保持三个独立 backend op，不合一。 +### 1.2 CNN 相关模块 / 算子设计 -### 5.3 实现次序 +**Conv2d v1 scope**（`infini_train/include/nn/modules/conv.h`，`infini_train/include/autograd/conv.h`）： ```text -latest upstream - → MNIST stride bug 修复（image_size_in_bytes_ 按 FLOAT32 计） - → shape-preserving Stack（[B,*sample_dims]）+ 回归测试，不碰 DDP sharding - → ReLU - → Conv2dMeta + CPU direct（forward/dW/dX/db，stride/padding） - → CPU finite-difference gradcheck + PyTorch parity（数学 correctness 结束） - → im2col 单测（tiny 人工可算例子，隔离 layout/transpose bug） - → CUDA forward（先 per-image GEMM，再 strided-batched 优化，语义不变） - → CUDA dW → dX（gather col2im）→ db - → CNN demo（--model 默认 mlp 不变，cnn 显式 opt-in） - → PyTorch 网络级 parity - → infini_run 多进程 DDP + sharded eval + epoch 级 metric reduce - → 1GPU vs 2GPU 数值等价 + 性能清理 +dtype: FP32 only（入口显式 CHECK，不静默 fallback） +layout: input NCHW，weight OIHW +API: scalar kernel_size / stride / padding + optional bias +内部语义: Kh/Kw 独立（读 weight.Dims()[2/3]），不加 Kh == Kw 限制 +不支持: tuple kernel/stride/padding、dilation、groups、FP16/BF16 ``` -### 5.4 DataLoader 与 Demo - -Dataset 单样本为 `[1,28,28]`，`Stack` 保持形状得 `[B,1,28,28]`，模型内不再 `View` 恢复形状；MLP 自行 `Flatten(1)`。DDP 采用 `To(device) → Broadcast → DDP wrap` 顺序，训练 step 只同步梯度，logging 不引入额外 collective。 +public API 收敛到任务 scope，internal semantics 不做无意义限制。非 square 权重只在 kernel / autograd 层单测覆盖，Module 层不暴露 tuple API。 -### 5.5 Oracle 策略 +**双后端职责划分**： ```text - PyTorch - │ - ▼ -CPU direct ─────── CUDA GEMM - │ │ - └── finite diff ────┘ +CPU: direct reference 实现（forward + dX + dW + db），暂不用 Eigen / im2col +CUDA: im2col + framework GEMM（Dispatcher "Gemm"），gather 风格 col2im，bias reduction ``` -CPU numerical gradcheck（tiny tensor，dX/dW/db）+ PyTorch 对标（logits / loss / grad / 单步更新 / 多步趋势）两套独立 oracle，不依赖 CUDA direct 中间实现。 +CPU 与现有 `cpu/matmul.cc` 朴素循环风格一致，职责是可审计的数学基准；CUDA 复用 `cuda/common/gemm.cu` 的 strided-batched GEMM，与 `cuda/matmul.cu`、`cuda/linear.cu` 调用方式一致。两条独立实现路径互为 cross-check，避免 im2col layout 错误在两边同时成立。 ---- +`Conv2dMeta`（`infini_train/src/kernels/common/conv.h`）集中保存 stride / padding / batch / 通道 / `kernel_h,w` / `output_h,w` / `patches` / `kernel_elems`，shape arithmetic 全局只有一份，CPU / CUDA 不再重复定义。 -## 六、环境与算力需求 +**ReLU**（`infini_train/include/autograd/activations.h`，`infini_train/include/nn/modules/activations.h`）：CPU direct + CUDA kernel，与 Conv 相同的 Dispatcher 注册方式，前向存 mask（或输出符号）供反向使用。 -### 6.1 开销估算 +**Autograd**：沿用 selective-save（需 dW 才存 input，需 dX 才存 weight），`Conv2dBackwardInput / Weight / Bias` 保持三个独立 backend op，不合一，便于逐项 gradcheck。 -参考网络 `Conv2d(1,16,3) → ReLU → Conv2d(16,32,3) → ReLU → Flatten → Linear(18432,10)`（s=1/p=0）： +### 1.3 模型构建 -- 参数约 189K（conv 4.8K + fc 184K），权重 < 1MB； -- 约 2.9M MAC/sample，单 epoch（60K）forward 约 0.35 TFLOP，fwd+bwd 约 1 TFLOP； -- batch=64 时激活约 7.7MB，整网 VRAM < 100MB。 +`example/mnist/net.h` / `net.cc` 中的 `MnistCnn`： -结论：单卡开发全够，瓶颈在 kernel 启动与 DataLoader，不在 FLOPs；CPU direct 跑 60K 为分钟级，可接受。 +```text +Conv2d(1,16,3) → ReLU → Conv2d(16,32,3) → ReLU → Flatten(1) → Linear(18432,10) +输入 [B,1,28,28]（NCHW），conv 栈保持 NCHW，输出 [B,32,24,24]，Flatten 后 18432 特征。 +``` -### 6.2 本机现状与缺口 +形状推导（kernel 3 / stride 1 / padding 0）：28×28 → 26×26（16ch）→ 24×24（32ch），classifier 输入 32×24×24 = 18432。`Forward` 内有 `CHECK_EQ(flat->Dims()[1], 18432)` 防形状漂移。MLP（`MNIST` 类）保留原行为，自行 `Flatten(1)`，与 CNN 共用同一 dataloader 输出。 -- 本机：1× RTX 4050 Laptop 6GB（空闲）、16 CPU、19GB RAM、磁盘充足，nvcc 13.4，gcc 16，cmake 4.4。 -- 缺：`build/` 未建、submodule 未初始化、torch 未装、无 MNIST 数据、无 NCCL（标准路径下未找到 `nccl.h` / `libnccl*`）。 -- 唯一算力缺口：**双卡 DDP 等价验证**（1GPU bs=128 vs 2×64）。本地单卡只能做单卡全量 + DDP 流程联调，最终等价性需在双卡机或 CI 跑一次并留 log 作报告证据。 +参数约 189K（conv 4.8K + fc 184K），权重 < 1MB；约 2.9M MAC/sample；batch=64 时激活约 7.7MB，整网显存 < 100MB。 -### 6.3 环境准备清单 +### 1.4 数据加载 -```bash -# 1. 子模块（BUILD_TEST 需要 googletest） -git submodule update --init third_party/googletest +- `MNISTDataset` 单样本 `[1,28,28]`；修复 stride bug（`image_size_in_bytes_` 按 FLOAT32 计，旧值 784B vs 正确 3136B）。 +- `Stack` 改为 shape-preserving（`[B,*sample_dims]`），MNIST batch 为 `[B,1,28,28]`，模型内不再 `View` 恢复形状；不碰 DDP sharding。 +- DDP 顺序：`To(device) → Broadcast → DDP wrap`（与 gpt2/mixtral/llama3 一致），训练 step 只同步梯度，logging 不引入额外 collective。 +- `DistributedDataLoader` 用 floor 除法（丢尾），单进程 eval 用 ceil 除法覆盖全集；等价比较时注意口径（见 2.4 注记）。 -# 2. MNIST 数据（官方 TorchVision 镜像，约 50MB） -./scripts/assets/prepare-infinitrain-assets.sh mnist -# 产物：data/mnist/{train-images-idx3-ubyte,train-labels-idx1-ubyte,t10k-images-idx3-ubyte,t10k-labels-idx1-ubyte} +### 1.5 前向 / 反向 / 优化器更新 -# 3. 构建（本地单卡；4050 为 Ada sm_89，仓库默认 CUDA_ARCHITECTURES=75;80;90 -# 不含 89，本地建议覆盖为 native；缺 NCCL 时先关 DDP 构建） -cmake -S . -B build -DUSE_CUDA=ON -DUSE_NCCL=OFF -DBUILD_TEST=ON \ - -DCMAKE_CUDA_ARCHITECTURES=native -DCMAKE_POLICY_VERSION_MINIMUM=3.5 -cmake --build build -j -# 有 NCCL 的双卡机：-DUSE_NCCL=ON，并用 ./build/infini_run 启动多进程 DDP - -# 4. PyTorch 对标（仅 CPU 数值 oracle 即可，不需要 GPU torch） -uv venv .venv && uv pip install "torch>=2.4 --index-url https://download.pytorch.org/whl/cpu" numpy -# 对齐固定：seed、输入、初始权重（.bin 导出供 InfiniTrain 读取）、lr/batch -# 对比项：logits、loss、关键参数 grad、单步更新后参数、多步 loss 趋势 +训练循环（`example/mnist/main.cc`）： -# 5. 端到端（默认 MLP 行为不变，CNN 显式指定） -./build/mnist --dataset data/mnist --device cpu -./build/mnist --dataset data/mnist --device cuda --model=cnn --num_epoch=3 --lr=0.1 +```text +optimizer.ZeroGrad() # 必须在 DDP Forward 之前（见 1.6） + → (*network)({image}) # 走 Module::operator()，保持 DDP hook 生效 + → (*loss_fn)({logits, label}) + → loss.Backward() + → optimizer.Step() # SGD，lr 由 --lr 指定 ``` -MNIST 四个 IDX 文件可按字节数 + 首 16 字节 magic 校验（images 0x00000803，labels 0x00000801），损坏时 `FORCE=1` 重跑准备脚本。 +eval 全程加 `autograd::NoGradGuard`，避免 forward-only 图污染 grad 累加器的依赖计数（曾导致静默停梯）。 + +### 1.6 关键实现取舍 + +| 取舍 | 决策 | 理由 | +| --- | --- | --- | +| CPU 不用 Eigen/im2col | direct 朴素循环 | 与 `cpu/matmul.cc` 风格一致，可审计的数学基准 | +| CUDA 用 im2col + 框架 GEMM | 复用 `cuda/common/gemm.cu` | 与 linear/matmul 一致，per-image GEMM → strided-batched 优化语义不变（1903ms → 849ms） | +| `ZeroGrad → Forward` 顺序 | 与 gpt2/mixtral/llama3 对齐 | `Forward → ZeroGrad(true)` 会清掉 `PrepareForBackward` 建立的 bucket-view 绑定，导致静默不同步、cross-rank 权重分叉（P13b 根因，`ed8d51a` 修复） | +| 训练/eval 走 `Module::operator()` | 不直接调 `Forward` | 保证 DDP hook 生效；eval 加 `NoGradGuard`（`8770716`） | +| `--model` 默认 `mlp` | CNN 显式 opt-in | 默认行为不变，老 demo 不受影响 | +| loss 打印 rank0-local | 不做 per-step AllReduce | 减少 collective；等价判据用 `Step` 后权重一致而非 loss 相等 | --- -## 七、结果摘要 +## 二、结果展示 -超参（端到端统一）：batch=64，lr=0.01，SGD，CNN 即 §6.1 参考网络。 +端到端统一超参：batch=64，lr=0.01，SGD，CNN 即 1.3 参考网络。 + +### 2.1 单元测试 | 任务 | 结论 | | --- | --- | @@ -133,19 +128,52 @@ MNIST 四个 IDX 文件可按字节数 + 首 16 字节 magic 校验（images 0x0 | P4 conv CPU | 8/8 exact-value；全量 Autograd 144/144 | | P5 gradcheck+golden | CPU finite-difference + torch golden 通过；Autograd 增至 145 pass | | P6 im2col | RTX 4050 上通过；发现 Dispatcher const& vs by-value 坑（已在代码注释记录） | -| P7 CUDA forward | 96 配置 grid ≤1e-5；per-image 1903ms → strided-batched 849ms | -| P8/P9/P10 CUDA dW/dX/db | 三个 grid 全部 ≤1e-5；Conv 相关 20/20 | -| P11 e2e smoke（1 epoch） | mlp acc 0.60；cnn cpu acc 0.90、loss 2.3→0.06；cnn cuda 与 CPU bit-identical，2.3s/epoch | -| P12 网络级 parity | 140/140 @1e-5（worst 9.5e-7）；10-step loss 双边严格单调递减 | -| P13a DDP（main.cc only） | 单进程 cuda acc 0.9028；NCCL-ON 可编译（需手动 `-I/-L/-rpath`，顶层 CMake 未加 NCCL include 路径）；2 进程止步于 `SetDevice(cuda:1)`——单卡机上限 | -| P13b 双卡等价 | PASS。判据为 `Step` 后 cross-rank 权重一致（详见下文复测结论小节） | -| 全量 ctest（NCCL-OFF） | 354/355 通过；唯一失败 `stacktrace`（glog 符号化环境问题，本分支未触及该文件，属 HEAD 预存失败）；CUDA/CPU 交叉实例化 SKIP 系 `ONLY_*` 宏按设计跳过 | +| P7 CUDA forward | 96 配置 grid ≤ 1e-5；per-image 1903ms → strided-batched 849ms | +| P8/P9/P10 CUDA dW/dX/db | 三个 grid 全部 ≤ 1e-5；Conv 相关 20/20 | +| 全量 ctest（NCCL-OFF） | 354/355 通过；唯一失败 `stacktrace`（glog 符号化环境问题，本分支未触及，属 HEAD 预存失败）；CUDA/CPU 交叉实例化 SKIP 系 `ONLY_*` 宏按设计跳过 | + +### 2.2 InfiniTrain 与 PyTorch 对齐（前向 / 损失 / 梯度 / 参数更新） -### P13b 复测结论（2026-09-19，RTX 4090D ×2） +对齐口径：`test_mnist_parity` dump 确定性 10 步 SGD 轨迹（batch=4，lr=0.01，seed param=42/input=1234/label=5678， +参数按 `NamedParameters` 排序后 `U(-1/sqrt(fan_in), 1/sqrt(fan_in))` 回填，与 torch `kaiming_uniform_(a=sqrt(5))` 同数学）， +`scripts/mnist_parity_torch.py --dump-dir ` 读取同一份 dump 重跑对比，阈值 1e-5。 +每步 14 项（logits、loss、6×grad、6×param），10 步共 140 项。 -配置：CNN，SGD lr=0.01，3 epoch，`1×128` vs `2×64`。 +| 对比项 | 规模 | 结论 | +| --- | --- | --- | +| 前向 logits `[B,10]` | 10 步 | 全部 ≤ 1e-5 | +| 损失（mean CE 标量） | 10 步 | 全部 ≤ 1e-5 | +| 梯度（6 参数，step 前） | 60 项 | 全部 ≤ 1e-5 | +| 参数更新（6 参数，SGD step 后） | 60 项 | 全部 ≤ 1e-5 | +| 合计 | 140/140 通过，最坏误差 9.5e-7 | 通过 | + +10 步 loss 双边严格单调递减（趋势一致）。复现命令： + +```bash +MNIST_PARITY_OUT_DIR=/tmp/mnist_parity ./build/tests/example/test_mnist_parity +python3 scripts/mnist_parity_torch.py --dump-dir /tmp/mnist_parity +``` + +### 2.3 端到端训练（loss / 准确率随迭代变化） + +1 epoch smoke（batch=64，lr=0.01，SGD）： + +| 模型 / 设备 | 训练 loss | 测试 acc | 备注 | +| --- | --- | --- | --- | +| mlp / cpu | — | 0.60 | 基线，行为不变 | +| cnn / cpu | 2.3 → 0.06 | 0.90 | 1 epoch | +| cnn / cuda | 与 CPU bit-identical | 0.9028（单进程） | 2.3s/epoch | + +P13b 双卡等价（RTX 4090D ×2，CNN，SGD lr=0.01，3 epoch，`1×128` vs `2×64`）： + +| 配置 | epoch0 / 1 / 2 train loss | test acc | +| --- | --- | --- | +| 1 卡 bs=128 | 0.596128 / 0.325636 / 0.301542 | 0.9123（10000 样本） | +| 2 卡 2×64 | 0.594806 / 0.322919 / 0.298881 | 0.9018（9984 样本） | -**主判据（权重一致性）：** +loss 随 epoch 单调下降；1 卡 vs 2 卡残差来自 floor 掉尾（训练每 epoch 少 96 样本、eval 少 16 样本）与 rank0-local 打印口径，非同步问题。 + +### 2.4 DDP 等价（主判据：`Step` 后 cross-rank 权重一致） | 探针 | rank0 | rank1 | 结论 | | --- | --- | --- | --- | @@ -153,28 +181,96 @@ MNIST 四个 IDX 文件可按字节数 + 首 16 字节 magic 校验（images 0x0 | `post-epoch w0`（修前） | `-0.00206` | `-0.00215` | 分叉 | | `post-epoch w0`（修后） | `-0.001326` | `-0.001326` | 一致 ✅ | -**Loss 对照（参考，非主判据）：** +根因：`example/mnist/main.cc` 曾用 `Forward → ZeroGrad(true)`，清掉 bucket-view 绑定致静默不同步。 +`ed8d51a` 改为 `ZeroGrad → Forward`；`8770716` 训练/eval 改走 `Module::operator()` + eval 加 `NoGradGuard`；双卡复测行为零变化（上表即含该改动的结果）。 +限定：单机双卡；多机、更大 world size、grad-accum 未测；train loss 打印仍为 rank0-local。 + +> Sharded-eval 注记：`DistributedDataLoader` 用 floor 除法（`num_batches_ = Size / global_batch`，丢尾）， +> 单进程 eval 用 ceil 除法覆盖全集。因此 1GPU-vs-2GPU 等价比较必须满足二者之一： +> (a) 只在 EVALUATED 样本上比较全局 loss/accuracy（reduce 上来的三元组本身即已是在 evaluated 样本上的精确值）； +> (b) 取数据集大小能整除 global batch 的配置。 + +--- + +## 三、使用说明 + +### 3.1 数据集与模型权重准备 + +```bash +# 子模块（BUILD_TEST 需要 googletest） +git submodule update --init third_party/googletest -| 配置 | epoch0 / 1 / 2 loss | test acc | +# MNIST 数据（约 50MB） +./scripts/assets/prepare-infinitrain-assets.sh mnist +# 产物：data/mnist/{train-images-idx3-ubyte,train-labels-idx1-ubyte,t10k-images-idx3-ubyte,t10k-labels-idx1-ubyte} +``` + +IDX 校验：按字节数 + 首 16 字节 magic（images `0x00000803`，labels `0x00000801`）；损坏时 `FORCE=1` 重跑准备脚本。 +模型权重无需预下载：`MnistCnn` / `MNIST` 构造时 `KaimingUniform` 初始化；对齐场景用 `test_mnist_parity` 的 `init__*.bin` 作为双方共同起点（见 2.2）。 + +### 3.2 公共 C++ 接口 + +```cpp +// 卷积 / 激活（infini_train/include/nn/modules/{conv,activations}.h） +nn::Conv2d(in_channels, out_channels, kernel_size, stride = 1, padding = 0, bias = true, device); +nn::ReLU(); + +// 模型（example/mnist/net.h） +MnistCnn(); // Conv2d(1,16,3)->ReLU->Conv2d(16,32,3)->ReLU->Flatten(1)->Linear(18432,10) +MNIST(); // MLP 基线：Linear(784,30)->Sigmoid->Linear(30,10) + +// 底层（按需）：autograd::Conv2d / autograd::ReLU（Function 接口）， +// kernels::MakeConv2dMeta(input, weight, stride, padding)（shape 唯一来源）。 +``` + +### 3.3 gflags 参数 + +| 参数 | 默认 | 说明 | | --- | --- | --- | -| 1卡 bs=128 | 0.596128 / 0.325636 / 0.301542 | 0.9123（10000） | -| 2卡 2×64 | 0.594806 / 0.322919 / 0.298881 | 0.9018（9984） | +| `--dataset` | `""` | MNIST 数据目录（必填，指向含 4 个 IDX 文件的目录） | +| `--bs` | `64` | batch size | +| `--num_epoch` | `1` | 训练 epoch 数 | +| `--lr` | `0.01` | SGD 学习率 | +| `--device` | `cpu` | `cpu` / `cuda` | +| `--model` | `mlp` | `mlp` / `cnn`（CNN 需显式 `--model=cnn`） | -残差来自 floor 掉尾（训练每 epoch 少 96 样本、eval 少 16 样本）与 rank0-local 打印口径，非同步问题。 +### 3.4 模型启动方式 -**根因与 fix：** +```bash +# 构建（单卡示例；RTX 4050 为 Ada sm_89，仓库默认 75;80;90 不含 89，本地建议 native； +# 缺 NCCL 时先关 DDP） +cmake -S . -B build -DUSE_CUDA=ON -DUSE_NCCL=OFF -DBUILD_TEST=ON \ + -DCMAKE_CUDA_ARCHITECTURES=native -DCMAKE_POLICY_VERSION_MINIMUM=3.5 +cmake --build build -j +# 双卡机：-DUSE_NCCL=ON，并用 ./build/infini_run 启动多进程 DDP -- 根因：`example/mnist/main.cc` 曾用 `Forward → ZeroGrad(true)`，清掉 `PrepareForBackward` 建立的 bucket-view 绑定，致静默不同步。 -- `ed8d51a`：改为 `ZeroGrad → Forward`，与 gpt2/mixtral/llama3 对齐。 -- `8770716`：训练/eval 改走 `Module::operator()`，eval 加 `NoGradGuard`（抄 PR 219/220）；双卡复测行为零变化（上表即含该改动的结果）。 +# 端到端（默认 MLP 行为不变，CNN 显式指定） +./build/mnist --dataset data/mnist --device cpu +./build/mnist --dataset data/mnist --device cuda --model=cnn --num_epoch=3 --lr=0.1 -**限定：** 单机双卡；多机、更大 world size、grad-accum 未测；train loss 打印仍为 rank0-local。 +# 对齐（C++ dump + torch 对比） +MNIST_PARITY_OUT_DIR=/tmp/mnist_parity ./build/tests/example/test_mnist_parity +python3 scripts/mnist_parity_torch.py --dump-dir /tmp/mnist_parity +``` -HEAD 预存 bug 修复（保留，与上游 #220 一致）：`main.cc` 改用 `shared_ptr` 持有 Module（Module 是 `enable_shared_from_this`，栈实例触发 `bad_weak_ptr` 崩溃）。 +### 3.5 默认行为 -ENV-ONLY 说明：gcc16 本地构建曾需在 `tensor.cc` / `profiler.cc` 各加一行 `#include `；上游 gcc13 不受影响，该两行已 revert，不在 PR 内。 +- `--model` 缺省为 `mlp`，老行为不变；CNN 必须显式 opt-in。 +- 分布式通过环境变量（`WORLD_SIZE/RANK/LOCAL_RANK`，由 `infini_run` 设置）自动 detection；单进程即非分布式。 +- 分布式下 device 自动取 `LOCAL_RANK` 对应 CUDA 卡；初参 rank0 Broadcast；训练只同步梯度，loss 为 rank0-local 打印。 -> Sharded-eval 注记：`DistributedDataLoader` 用 floor 除法（`num_batches_ = Size / global_batch`，丢尾）， -> 单进程 eval 用 ceil 除法覆盖全集。因此 P13b 1GPU-vs-2GPU 等价比较必须满足二者之一： -> (a) 只在 EVALUATED 样本上比较全局 loss/accuracy（reduce 上来的三元组本身即已是在 evaluated 样本上的精确值）； -> (b) 取数据集大小能整除 global batch 的配置。 +### 3.6 常见错误处理 + +| 现象 | 原因 | 处理 | +| --- | --- | --- | +| `--device/--model` 非法值被拒 | `DEFINE_validator` 校验 | 取 `cpu/cuda`、`mlp/cnn` | +| `--dataset` 为空或 IDX 缺失 | 路径错误 | 指向 `prepare-infinitrain-assets.sh mnist` 产物目录 | +| `SetDevice(cuda:1)` 失败 | 单卡机跑 2 进程 DDP | 到双卡机或 CI 跑 P13b | +| NCCL 头/库找不到 | 顶层 CMake 未加 NCCL include 路径 | 手动 `-I/-L/-rpath` 或 `-DUSE_NCCL=OFF` 先编单卡 | +| `bad_weak_ptr` 崩溃 | Module 是 `enable_shared_from_this`，栈实例触发 | 用 `shared_ptr` 持有 Module（本分支已修） | +| 双卡权重分叉但 loss 正常 | `Forward → ZeroGrad(true)` 清 bucket-view 绑定 | 用 `ZeroGrad → Forward`（本分支已修，切勿改回） | +| 1GPU vs 2GPU loss/acc 差小量 | floor 丢尾 + rank0-local 口径 | 按 2.4 注记 (a)/(b) 口径比较；主判据看 `Step` 后权重一致 | +| `ctest` 唯一失败 `stacktrace` | glog 符号化环境问题，HEAD 预存 | 本分支未触及，可忽略 | + +ENV-ONLY 说明：gcc16 本地构建曾需在 `tensor.cc` / `profiler.cc` 各加一行 `#include `； +上游 gcc13 不受影响，该两行已 revert，不在 PR 内。 diff --git a/scripts/mnist_parity_torch.py b/scripts/mnist_parity_torch.py new file mode 100644 index 000000000..9ccaa318e --- /dev/null +++ b/scripts/mnist_parity_torch.py @@ -0,0 +1,212 @@ +#!/usr/bin/env python3 +"""MnistCnn 对齐用 PyTorch 脚本（与 tests/example/mnist_parity.cc 配对）。 + +C++ 端（fixture）用固定的 mt19937 seed 生成输入/标签/初始参数， +跑 10 步 SGD 并把每步的 logits / loss / grad / 更新后参数 dump 成 .bin； +本脚本读取同一份 dump 目录，用完全相同的初始权重和输入在 PyTorch 里 +重跑 10 步 SGD，逐项对比，阈值 1e-5。 + +文件布局（C++ 端生成，见 mnist_parity.cc 头注释）： + meta.txt, input.bin(float32 [B,1,28,28]), labels.bin(uint8 [B]), + init__.bin, + step__logits.bin, step__loss.bin, + step__grad__.bin, step__param__.bin + +参数名与形状（与 MnistCnn::NamedParameters 排序后一致）： + conv1.weight [16,1,3,3], conv1.bias [16], + conv2.weight [32,16,3,3], conv2.bias [32], + fc.weight [10,18432], fc.bias [10] +注意：InfiniTrain Linear 权重形状为 [out_features, in_features]， +与 torch.nn.Linear.weight 一致，可直接按行主序拷贝。 + +用法： + pip install "torch>=2.4 --index-url https://download.pytorch.org/whl/cpu" numpy + MNIST_PARITY_OUT_DIR=/tmp/mnist_parity ./build/tests/example/test_mnist_parity + python3 scripts/mnist_parity_torch.py --dump-dir /tmp/mnist_parity + +退出码：全部通过返回 0，有任一项超差返回 1。 +""" + +from __future__ import annotations + +import argparse +import struct +import sys +from pathlib import Path + +import numpy as np + +try: + import torch + import torch.nn.functional as F +except ImportError: + print("需要先安装 torch（CPU 版即可）：", file=sys.stderr) + print(' pip install "torch>=2.4 --index-url https://download.pytorch.org/whl/cpu" numpy', file=sys.stderr) + sys.exit(2) + +TOL = 1e-5 +STEPS = 10 +LR = 0.01 + +# 与 mnist_parity.cc 中排序后的 NamedParameters 顺序保持一致 +PARAM_SPECS = [ + ("conv1.bias", (16,)), + ("conv1.weight", (16, 1, 3, 3)), + ("conv2.bias", (32,)), + ("conv2.weight", (32, 16, 3, 3)), + ("fc.bias", (10,)), + ("fc.weight", (10, 18432)), +] + + +class MnistCnn(torch.nn.Module): + """Conv2d(1,16,3)->ReLU->Conv2d(16,32,3)->ReLU->Flatten(1)->Linear(18432,10)。""" + + def __init__(self) -> None: + super().__init__() + self.conv1 = torch.nn.Conv2d(1, 16, kernel_size=3, stride=1, padding=0, bias=True) + self.relu1 = torch.nn.ReLU() + self.conv2 = torch.nn.Conv2d(16, 32, kernel_size=3, stride=1, padding=0, bias=True) + self.relu2 = torch.nn.ReLU() + self.fc = torch.nn.Linear(18432, 10, bias=True) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + x = self.relu1(self.conv1(x)) + x = self.relu2(self.conv2(x)) + x = torch.flatten(x, 1) + return self.fc(x) + + +def read_f32(path: Path, shape: tuple) -> np.ndarray: + raw = path.read_bytes() + expect = int(np.prod(shape)) * 4 + assert len(raw) == expect, f"{path}: 字节数 {len(raw)} != 期望 {expect}（形状 {shape}）" + return np.frombuffer(raw, dtype=" np.ndarray: + raw = path.read_bytes() + assert len(raw) == n, f"{path}: 字节数 {len(raw)} != 期望 {n}" + return np.frombuffer(raw, dtype=np.uint8).copy() + + +def max_abs_diff(a: np.ndarray, b: np.ndarray) -> float: + return float(np.max(np.abs(a.astype(np.float64) - b.astype(np.float64)))) + + +def main() -> int: + ap = argparse.ArgumentParser(description="MnistCnn PyTorch 对齐对比（阈值 1e-5）") + ap.add_argument("--dump-dir", required=True, help="C++ fixture 的输出目录（含 meta.txt/input.bin/...）") + ap.add_argument("--tol", type=float, default=TOL, help="最大允许绝对误差（默认 1e-5）") + ap.add_argument("--steps", type=int, default=STEPS, help="对比步数（默认 10，需 <= dump 步数）") + ap.add_argument("--lr", type=float, default=LR, help="SGD 学习率（默认 0.01，需与 fixture 一致）") + args = ap.parse_args() + + dump = Path(args.dump_dir) + assert (dump / "meta.txt").exists(), f"找不到 {dump / 'meta.txt'}，请先运行 test_mnist_parity 生成 dump" + torch.manual_seed(0) + torch.set_num_threads(1) + + batch = None + # 从 init 文件推断 batch：读 input.bin 头部无法得知 B，先用 labels.bin 长度 + labels_np = read_u8(dump / "labels.bin", len((dump / "labels.bin").read_bytes())) + batch = labels_np.shape[0] + input_np = read_f32(dump / "input.bin", (batch, 1, 28, 28)) + + model = MnistCnn() + # 用 C++ dump 的初始权重覆盖 torch 默认初始化（逐元素精确对齐起点） + with torch.no_grad(): + state = model.state_dict() + name_map = { + "conv1.bias": "conv1.bias", + "conv1.weight": "conv1.weight", + "conv2.bias": "conv2.bias", + "conv2.weight": "conv2.weight", + "fc.bias": "fc.bias", + "fc.weight": "fc.weight", + } + for name, shape in PARAM_SPECS: + arr = read_f32(dump / f"init__{name}.bin", shape) + state[name_map[name]].copy_(torch.from_numpy(arr)) + + x = torch.from_numpy(input_np) + y = torch.from_numpy(labels_np.astype(np.int64)) + opt = torch.optim.SGD(model.parameters(), lr=args.lr) + + n_fail = 0 + total_checks = 0 + worst = 0.0 + print(f"{'step':>8} {'check':>22} {'max_abs_diff':>14} 结论") + for step in range(args.steps): + tag = f"step{step:03d}" + model.zero_grad() + logits = model(x) + loss = F.cross_entropy(logits, y, reduction="mean") + loss.backward() + + # 1) logits + ref_logits = read_f32(dump / f"{tag}__logits.bin", (batch, 10)) + d = max_abs_diff(logits.detach().numpy(), ref_logits) + ok = d <= args.tol + print(f"{tag:>8} {'logits':>22} {d:14.3e} {'通过' if ok else '超差'}") + total_checks += 1 + worst = max(worst, d) + n_fail += not ok + + # 2) loss + ref_loss = read_f32(dump / f"{tag}__loss.bin", (1,))[0] + d = abs(float(loss.item()) - float(ref_loss)) + ok = d <= args.tol + print(f"{tag:>8} {'loss':>22} {d:14.3e} {'通过' if ok else '超差'}") + total_checks += 1 + worst = max(worst, d) + n_fail += not ok + + # 3) 每个参数的 grad（SGD step 之前） + grad_map = { + "conv1.bias": model.conv1.bias.grad, + "conv1.weight": model.conv1.weight.grad, + "conv2.bias": model.conv2.bias.grad, + "conv2.weight": model.conv2.weight.grad, + "fc.bias": model.fc.bias.grad, + "fc.weight": model.fc.weight.grad, + } + for name, shape in PARAM_SPECS: + ref_g = read_f32(dump / f"{tag}__grad__{name}.bin", shape) + d = max_abs_diff(grad_map[name].detach().numpy(), ref_g) + ok = d <= args.tol + print(f"{tag:>8} {('grad/' + name):>22} {d:14.3e} {'通过' if ok else '超差'}") + total_checks += 1 + worst = max(worst, d) + n_fail += not ok + + opt.step() + + # 4) 每个参数 SGD 更新后的值 + param_map = { + "conv1.bias": model.conv1.bias.detach(), + "conv1.weight": model.conv1.weight.detach(), + "conv2.bias": model.conv2.bias.detach(), + "conv2.weight": model.conv2.weight.detach(), + "fc.bias": model.fc.bias.detach(), + "fc.weight": model.fc.weight.detach(), + } + for name, shape in PARAM_SPECS: + ref_p = read_f32(dump / f"{tag}__param__{name}.bin", shape) + d = max_abs_diff(param_map[name].numpy(), ref_p) + ok = d <= args.tol + print(f"{tag:>8} {('param/' + name):>22} {d:14.3e} {'通过' if ok else '超差'}") + total_checks += 1 + worst = max(worst, d) + n_fail += not ok + + print(f"\n共 {total_checks} 项对比，超差 {n_fail} 项，最坏误差 {worst:.3e}（阈值 {args.tol:.0e}）") + if n_fail == 0: + print("结论：通过（140/140 口径：每步 14 项 × 10 步）") + return 0 + print("结论：未通过，请检查 dump 目录是否与当前代码版本匹配", file=sys.stderr) + return 1 + + +if __name__ == "__main__": + sys.exit(main()) From fdb7de61b2bdc39bd5c4a0d72e28424a80e8f57a Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E7=8B=AC=E8=A7=92=E4=BB=99=E4=B8=8E=E6=98=9F?= Date: Mon, 21 Sep 2026 16:14:55 +0800 Subject: [PATCH 7/8] =?UTF-8?q?revert:=20=E6=8A=A5=E5=91=8A=E8=B5=B0?= =?UTF-8?q?=E9=82=AE=E4=BB=B6=EF=BC=8C=E4=B8=8D=E8=BF=9BPR=EF=BC=88?= =?UTF-8?q?=E4=BF=9D=E7=95=99=E5=AF=B9=E9=BD=90=E8=84=9A=E6=9C=AC=EF=BC=89?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- docs/small_model_patch.md | 276 -------------------------------------- 1 file changed, 276 deletions(-) delete mode 100644 docs/small_model_patch.md diff --git a/docs/small_model_patch.md b/docs/small_model_patch.md deleted file mode 100644 index bcb10c4d4..000000000 --- a/docs/small_model_patch.md +++ /dev/null @@ -1,276 +0,0 @@ -# 小模型训练支持 — 项目报告（分支 `work/small-model-cnn-impl`） - -本报告对应 PR「【训练营】小模型训练支持」：功能代码 + CMake 集成 + 单元测试 + MNIST CNN 演示。 -PyTorch 对齐脚本见 `scripts/mnist_parity_torch.py`，C++ dump 端见 `tests/example/mnist_parity.cc`。 - ---- - -## 一、设计与实现 - -### 1.1 整体方案 - -目标是在 InfiniTrain 上跑通小模型（MNIST CNN）端到端训练，同时保证数学正确性可审计。 -方案是第三版取舍：采用 PR #219 的语义设计（数据形状语义、DDP metric 设计）+ PR #220 的 kernel 与验证方法（kernel 实现、验证 discipline），去掉两边共有的 demo-level special case。 - -整体训练链路： - -```text -MNISTDataset（单样本 [1,28,28]） - → DataLoader Stack（shape-preserving，batch 后 [B,1,28,28]） - → MnistCnn 前向（Conv→ReLU→Conv→ReLU→Flatten→Linear） - → CrossEntropyLoss（mean） - → Backward（autograd selective-save） - → SGD 更新（lr=0.01） - → 可选 DDP（Broadcast → wrap，训练只同步梯度；eval 分片 + epoch 级 AllReduce） -``` - -Oracle 策略（双独立 oracle，不依赖 CUDA direct 中间实现）： - -```text - PyTorch - │ - ▼ -CPU direct ─────── CUDA GEMM - │ │ - └── finite diff ────┘ -``` - -即 CPU numerical gradcheck（tiny tensor，dX/dW/db）+ PyTorch 对标（logits / loss / grad / 单步更新 / 多步趋势）。 - -### 1.2 CNN 相关模块 / 算子设计 - -**Conv2d v1 scope**（`infini_train/include/nn/modules/conv.h`，`infini_train/include/autograd/conv.h`）： - -```text -dtype: FP32 only（入口显式 CHECK，不静默 fallback） -layout: input NCHW，weight OIHW -API: scalar kernel_size / stride / padding + optional bias -内部语义: Kh/Kw 独立（读 weight.Dims()[2/3]），不加 Kh == Kw 限制 -不支持: tuple kernel/stride/padding、dilation、groups、FP16/BF16 -``` - -public API 收敛到任务 scope，internal semantics 不做无意义限制。非 square 权重只在 kernel / autograd 层单测覆盖，Module 层不暴露 tuple API。 - -**双后端职责划分**： - -```text -CPU: direct reference 实现（forward + dX + dW + db），暂不用 Eigen / im2col -CUDA: im2col + framework GEMM（Dispatcher "Gemm"），gather 风格 col2im，bias reduction -``` - -CPU 与现有 `cpu/matmul.cc` 朴素循环风格一致，职责是可审计的数学基准；CUDA 复用 `cuda/common/gemm.cu` 的 strided-batched GEMM，与 `cuda/matmul.cu`、`cuda/linear.cu` 调用方式一致。两条独立实现路径互为 cross-check，避免 im2col layout 错误在两边同时成立。 - -`Conv2dMeta`（`infini_train/src/kernels/common/conv.h`）集中保存 stride / padding / batch / 通道 / `kernel_h,w` / `output_h,w` / `patches` / `kernel_elems`，shape arithmetic 全局只有一份，CPU / CUDA 不再重复定义。 - -**ReLU**（`infini_train/include/autograd/activations.h`，`infini_train/include/nn/modules/activations.h`）：CPU direct + CUDA kernel，与 Conv 相同的 Dispatcher 注册方式，前向存 mask（或输出符号）供反向使用。 - -**Autograd**：沿用 selective-save（需 dW 才存 input，需 dX 才存 weight），`Conv2dBackwardInput / Weight / Bias` 保持三个独立 backend op，不合一，便于逐项 gradcheck。 - -### 1.3 模型构建 - -`example/mnist/net.h` / `net.cc` 中的 `MnistCnn`： - -```text -Conv2d(1,16,3) → ReLU → Conv2d(16,32,3) → ReLU → Flatten(1) → Linear(18432,10) -输入 [B,1,28,28]（NCHW），conv 栈保持 NCHW，输出 [B,32,24,24]，Flatten 后 18432 特征。 -``` - -形状推导（kernel 3 / stride 1 / padding 0）：28×28 → 26×26（16ch）→ 24×24（32ch），classifier 输入 32×24×24 = 18432。`Forward` 内有 `CHECK_EQ(flat->Dims()[1], 18432)` 防形状漂移。MLP（`MNIST` 类）保留原行为，自行 `Flatten(1)`，与 CNN 共用同一 dataloader 输出。 - -参数约 189K（conv 4.8K + fc 184K），权重 < 1MB；约 2.9M MAC/sample；batch=64 时激活约 7.7MB，整网显存 < 100MB。 - -### 1.4 数据加载 - -- `MNISTDataset` 单样本 `[1,28,28]`；修复 stride bug（`image_size_in_bytes_` 按 FLOAT32 计，旧值 784B vs 正确 3136B）。 -- `Stack` 改为 shape-preserving（`[B,*sample_dims]`），MNIST batch 为 `[B,1,28,28]`，模型内不再 `View` 恢复形状；不碰 DDP sharding。 -- DDP 顺序：`To(device) → Broadcast → DDP wrap`（与 gpt2/mixtral/llama3 一致），训练 step 只同步梯度，logging 不引入额外 collective。 -- `DistributedDataLoader` 用 floor 除法（丢尾），单进程 eval 用 ceil 除法覆盖全集；等价比较时注意口径（见 2.4 注记）。 - -### 1.5 前向 / 反向 / 优化器更新 - -训练循环（`example/mnist/main.cc`）： - -```text -optimizer.ZeroGrad() # 必须在 DDP Forward 之前（见 1.6） - → (*network)({image}) # 走 Module::operator()，保持 DDP hook 生效 - → (*loss_fn)({logits, label}) - → loss.Backward() - → optimizer.Step() # SGD，lr 由 --lr 指定 -``` - -eval 全程加 `autograd::NoGradGuard`，避免 forward-only 图污染 grad 累加器的依赖计数（曾导致静默停梯）。 - -### 1.6 关键实现取舍 - -| 取舍 | 决策 | 理由 | -| --- | --- | --- | -| CPU 不用 Eigen/im2col | direct 朴素循环 | 与 `cpu/matmul.cc` 风格一致，可审计的数学基准 | -| CUDA 用 im2col + 框架 GEMM | 复用 `cuda/common/gemm.cu` | 与 linear/matmul 一致，per-image GEMM → strided-batched 优化语义不变（1903ms → 849ms） | -| `ZeroGrad → Forward` 顺序 | 与 gpt2/mixtral/llama3 对齐 | `Forward → ZeroGrad(true)` 会清掉 `PrepareForBackward` 建立的 bucket-view 绑定，导致静默不同步、cross-rank 权重分叉（P13b 根因，`ed8d51a` 修复） | -| 训练/eval 走 `Module::operator()` | 不直接调 `Forward` | 保证 DDP hook 生效；eval 加 `NoGradGuard`（`8770716`） | -| `--model` 默认 `mlp` | CNN 显式 opt-in | 默认行为不变，老 demo 不受影响 | -| loss 打印 rank0-local | 不做 per-step AllReduce | 减少 collective；等价判据用 `Step` 后权重一致而非 loss 相等 | - ---- - -## 二、结果展示 - -端到端统一超参：batch=64，lr=0.01，SGD，CNN 即 1.3 参考网络。 - -### 2.1 单元测试 - -| 任务 | 结论 | -| --- | --- | -| P0 构建 | `USE_CUDA=ON / USE_NCCL=OFF / arch native / POLICY_VERSION_MINIMUM 3.5` 全量通过；MNIST 数据 + torch 2.14+cpu 就绪 | -| P1 dataset | `test_mnist_dataset` 3/3；stride bug 实证：旧值 784B vs 正确 3136B | -| P2 dataloader | DataLoader 8/8（含新增 shape-preserving Stack 回归），MNISTDataset 3/3；MLP smoke acc 0.60 | -| P3 ReLU | CPU 3/3 | -| P4 conv CPU | 8/8 exact-value；全量 Autograd 144/144 | -| P5 gradcheck+golden | CPU finite-difference + torch golden 通过；Autograd 增至 145 pass | -| P6 im2col | RTX 4050 上通过；发现 Dispatcher const& vs by-value 坑（已在代码注释记录） | -| P7 CUDA forward | 96 配置 grid ≤ 1e-5；per-image 1903ms → strided-batched 849ms | -| P8/P9/P10 CUDA dW/dX/db | 三个 grid 全部 ≤ 1e-5；Conv 相关 20/20 | -| 全量 ctest（NCCL-OFF） | 354/355 通过；唯一失败 `stacktrace`（glog 符号化环境问题，本分支未触及，属 HEAD 预存失败）；CUDA/CPU 交叉实例化 SKIP 系 `ONLY_*` 宏按设计跳过 | - -### 2.2 InfiniTrain 与 PyTorch 对齐（前向 / 损失 / 梯度 / 参数更新） - -对齐口径：`test_mnist_parity` dump 确定性 10 步 SGD 轨迹（batch=4，lr=0.01，seed param=42/input=1234/label=5678， -参数按 `NamedParameters` 排序后 `U(-1/sqrt(fan_in), 1/sqrt(fan_in))` 回填，与 torch `kaiming_uniform_(a=sqrt(5))` 同数学）， -`scripts/mnist_parity_torch.py --dump-dir ` 读取同一份 dump 重跑对比，阈值 1e-5。 -每步 14 项（logits、loss、6×grad、6×param），10 步共 140 项。 - -| 对比项 | 规模 | 结论 | -| --- | --- | --- | -| 前向 logits `[B,10]` | 10 步 | 全部 ≤ 1e-5 | -| 损失（mean CE 标量） | 10 步 | 全部 ≤ 1e-5 | -| 梯度（6 参数，step 前） | 60 项 | 全部 ≤ 1e-5 | -| 参数更新（6 参数，SGD step 后） | 60 项 | 全部 ≤ 1e-5 | -| 合计 | 140/140 通过，最坏误差 9.5e-7 | 通过 | - -10 步 loss 双边严格单调递减（趋势一致）。复现命令： - -```bash -MNIST_PARITY_OUT_DIR=/tmp/mnist_parity ./build/tests/example/test_mnist_parity -python3 scripts/mnist_parity_torch.py --dump-dir /tmp/mnist_parity -``` - -### 2.3 端到端训练（loss / 准确率随迭代变化） - -1 epoch smoke（batch=64，lr=0.01，SGD）： - -| 模型 / 设备 | 训练 loss | 测试 acc | 备注 | -| --- | --- | --- | --- | -| mlp / cpu | — | 0.60 | 基线，行为不变 | -| cnn / cpu | 2.3 → 0.06 | 0.90 | 1 epoch | -| cnn / cuda | 与 CPU bit-identical | 0.9028（单进程） | 2.3s/epoch | - -P13b 双卡等价（RTX 4090D ×2，CNN，SGD lr=0.01，3 epoch，`1×128` vs `2×64`）： - -| 配置 | epoch0 / 1 / 2 train loss | test acc | -| --- | --- | --- | -| 1 卡 bs=128 | 0.596128 / 0.325636 / 0.301542 | 0.9123（10000 样本） | -| 2 卡 2×64 | 0.594806 / 0.322919 / 0.298881 | 0.9018（9984 样本） | - -loss 随 epoch 单调下降；1 卡 vs 2 卡残差来自 floor 掉尾（训练每 epoch 少 96 样本、eval 少 16 样本）与 rank0-local 打印口径，非同步问题。 - -### 2.4 DDP 等价（主判据：`Step` 后 cross-rank 权重一致） - -| 探针 | rank0 | rank1 | 结论 | -| --- | --- | --- | --- | -| `post-bcast w0` | `-0.0836399` | `-0.0836399` | 初值一致 | -| `post-epoch w0`（修前） | `-0.00206` | `-0.00215` | 分叉 | -| `post-epoch w0`（修后） | `-0.001326` | `-0.001326` | 一致 ✅ | - -根因：`example/mnist/main.cc` 曾用 `Forward → ZeroGrad(true)`，清掉 bucket-view 绑定致静默不同步。 -`ed8d51a` 改为 `ZeroGrad → Forward`；`8770716` 训练/eval 改走 `Module::operator()` + eval 加 `NoGradGuard`；双卡复测行为零变化（上表即含该改动的结果）。 -限定：单机双卡；多机、更大 world size、grad-accum 未测；train loss 打印仍为 rank0-local。 - -> Sharded-eval 注记：`DistributedDataLoader` 用 floor 除法（`num_batches_ = Size / global_batch`，丢尾）， -> 单进程 eval 用 ceil 除法覆盖全集。因此 1GPU-vs-2GPU 等价比较必须满足二者之一： -> (a) 只在 EVALUATED 样本上比较全局 loss/accuracy（reduce 上来的三元组本身即已是在 evaluated 样本上的精确值）； -> (b) 取数据集大小能整除 global batch 的配置。 - ---- - -## 三、使用说明 - -### 3.1 数据集与模型权重准备 - -```bash -# 子模块（BUILD_TEST 需要 googletest） -git submodule update --init third_party/googletest - -# MNIST 数据（约 50MB） -./scripts/assets/prepare-infinitrain-assets.sh mnist -# 产物：data/mnist/{train-images-idx3-ubyte,train-labels-idx1-ubyte,t10k-images-idx3-ubyte,t10k-labels-idx1-ubyte} -``` - -IDX 校验：按字节数 + 首 16 字节 magic（images `0x00000803`，labels `0x00000801`）；损坏时 `FORCE=1` 重跑准备脚本。 -模型权重无需预下载：`MnistCnn` / `MNIST` 构造时 `KaimingUniform` 初始化；对齐场景用 `test_mnist_parity` 的 `init__*.bin` 作为双方共同起点（见 2.2）。 - -### 3.2 公共 C++ 接口 - -```cpp -// 卷积 / 激活（infini_train/include/nn/modules/{conv,activations}.h） -nn::Conv2d(in_channels, out_channels, kernel_size, stride = 1, padding = 0, bias = true, device); -nn::ReLU(); - -// 模型（example/mnist/net.h） -MnistCnn(); // Conv2d(1,16,3)->ReLU->Conv2d(16,32,3)->ReLU->Flatten(1)->Linear(18432,10) -MNIST(); // MLP 基线：Linear(784,30)->Sigmoid->Linear(30,10) - -// 底层（按需）：autograd::Conv2d / autograd::ReLU（Function 接口）， -// kernels::MakeConv2dMeta(input, weight, stride, padding)（shape 唯一来源）。 -``` - -### 3.3 gflags 参数 - -| 参数 | 默认 | 说明 | -| --- | --- | --- | -| `--dataset` | `""` | MNIST 数据目录（必填，指向含 4 个 IDX 文件的目录） | -| `--bs` | `64` | batch size | -| `--num_epoch` | `1` | 训练 epoch 数 | -| `--lr` | `0.01` | SGD 学习率 | -| `--device` | `cpu` | `cpu` / `cuda` | -| `--model` | `mlp` | `mlp` / `cnn`（CNN 需显式 `--model=cnn`） | - -### 3.4 模型启动方式 - -```bash -# 构建（单卡示例；RTX 4050 为 Ada sm_89，仓库默认 75;80;90 不含 89，本地建议 native； -# 缺 NCCL 时先关 DDP） -cmake -S . -B build -DUSE_CUDA=ON -DUSE_NCCL=OFF -DBUILD_TEST=ON \ - -DCMAKE_CUDA_ARCHITECTURES=native -DCMAKE_POLICY_VERSION_MINIMUM=3.5 -cmake --build build -j -# 双卡机：-DUSE_NCCL=ON，并用 ./build/infini_run 启动多进程 DDP - -# 端到端（默认 MLP 行为不变，CNN 显式指定） -./build/mnist --dataset data/mnist --device cpu -./build/mnist --dataset data/mnist --device cuda --model=cnn --num_epoch=3 --lr=0.1 - -# 对齐（C++ dump + torch 对比） -MNIST_PARITY_OUT_DIR=/tmp/mnist_parity ./build/tests/example/test_mnist_parity -python3 scripts/mnist_parity_torch.py --dump-dir /tmp/mnist_parity -``` - -### 3.5 默认行为 - -- `--model` 缺省为 `mlp`，老行为不变；CNN 必须显式 opt-in。 -- 分布式通过环境变量（`WORLD_SIZE/RANK/LOCAL_RANK`，由 `infini_run` 设置）自动 detection；单进程即非分布式。 -- 分布式下 device 自动取 `LOCAL_RANK` 对应 CUDA 卡；初参 rank0 Broadcast；训练只同步梯度，loss 为 rank0-local 打印。 - -### 3.6 常见错误处理 - -| 现象 | 原因 | 处理 | -| --- | --- | --- | -| `--device/--model` 非法值被拒 | `DEFINE_validator` 校验 | 取 `cpu/cuda`、`mlp/cnn` | -| `--dataset` 为空或 IDX 缺失 | 路径错误 | 指向 `prepare-infinitrain-assets.sh mnist` 产物目录 | -| `SetDevice(cuda:1)` 失败 | 单卡机跑 2 进程 DDP | 到双卡机或 CI 跑 P13b | -| NCCL 头/库找不到 | 顶层 CMake 未加 NCCL include 路径 | 手动 `-I/-L/-rpath` 或 `-DUSE_NCCL=OFF` 先编单卡 | -| `bad_weak_ptr` 崩溃 | Module 是 `enable_shared_from_this`，栈实例触发 | 用 `shared_ptr` 持有 Module（本分支已修） | -| 双卡权重分叉但 loss 正常 | `Forward → ZeroGrad(true)` 清 bucket-view 绑定 | 用 `ZeroGrad → Forward`（本分支已修，切勿改回） | -| 1GPU vs 2GPU loss/acc 差小量 | floor 丢尾 + rank0-local 口径 | 按 2.4 注记 (a)/(b) 口径比较；主判据看 `Step` 后权重一致 | -| `ctest` 唯一失败 `stacktrace` | glog 符号化环境问题，HEAD 预存 | 本分支未触及，可忽略 | - -ENV-ONLY 说明：gcc16 本地构建曾需在 `tensor.cc` / `profiler.cc` 各加一行 `#include `； -上游 gcc13 不受影响，该两行已 revert，不在 PR 内。 From 67752bd51f8bc4329a7d8cfdc86066cdfb58f449 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E7=8B=AC=E8=A7=92=E4=BB=99=E4=B8=8E=E6=98=9F?= Date: Mon, 21 Sep 2026 17:03:37 +0800 Subject: [PATCH 8/8] =?UTF-8?q?chore:=20ignore=20mail/=EF=BC=88=E9=82=AE?= =?UTF-8?q?=E4=BB=B6=E6=9D=90=E6=96=99=E7=9B=AE=E5=BD=95=EF=BC=8C=E4=B8=8D?= =?UTF-8?q?=E8=BF=9B=E7=89=88=E6=9C=AC=EF=BC=89?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- .gitignore | 3 +++ 1 file changed, 3 insertions(+) diff --git a/.gitignore b/.gitignore index 4ad6f92ff..4595f3b02 100644 --- a/.gitignore +++ b/.gitignore @@ -8,3 +8,6 @@ build/ __pycache__/ /data/ + +# 邮件材料（项目报告+对齐脚本副本）：只走邮件，不进 PR +/mail/