diff --git a/.gitignore b/.gitignore
index 4ad6f92ff..4595f3b02 100644
--- a/.gitignore
+++ b/.gitignore
@@ -8,3 +8,6 @@ build/
 
 __pycache__/
 /data/
+
+# 邮件材料（项目报告+对齐脚本副本）：只走邮件，不进 PR
+/mail/
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<int64_t> 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<int>())),
       label_size_in_bytes_(kSN3TypeToSize.at(label_file_.type)
                            * std::accumulate(label_dims_.begin(), label_dims_.end(), 1, std::multiplies<int>())) {
diff --git a/example/mnist/main.cc b/example/mnist/main.cc
index 7744e0947..6fcaa2788 100644
--- a/example/mnist/main.cc
+++ b/example/mnist/main.cc
@@ -9,9 +9,17 @@
 #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"
+#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 +30,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 +40,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<MNISTDataset>(FLAGS_dataset, true);
-    DataLoader train_dataloader(train_dataset, FLAGS_bs);
+    std::unique_ptr<DataLoader> train_dataloader
+        = distributed ? std::unique_ptr<DataLoader>(
+                            std::make_unique<DistributedDataLoader>(train_dataset, FLAGS_bs, ddp_rank, ddp_world_size))
+                      : std::unique_ptr<DataLoader>(std::make_unique<DataLoader>(train_dataset, FLAGS_bs));
 
     // TODO(dcj): Add sampler & eval dataloader later.
     auto test_dataset = std::make_shared<MNISTDataset>(FLAGS_dataset, false);
-    DataLoader test_dataloader(test_dataset, FLAGS_bs);
-
-    auto network = MNIST();
+    std::unique_ptr<DataLoader> test_dataloader
+        = distributed ? std::unique_ptr<DataLoader>(
+                            std::make_unique<DistributedDataLoader>(test_dataset, FLAGS_bs, ddp_rank, ddp_world_size))
+                      : std::unique_ptr<DataLoader>(std::make_unique<DataLoader>(test_dataset, FLAGS_bs));
+
+    std::shared_ptr<nn::Module> network;
+    if (FLAGS_model == kModelCNN) {
+        network = std::make_shared<MnistCnn>();
+    } else {
+        network = std::make_shared<MNIST>();
+    }
     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<nn::parallel::DistributedDataParallel>(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<nn::CrossEntropyLoss>();
+    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 +119,19 @@ 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<Tensor>(image->To(device));
             auto new_label = std::make_shared<Tensor>(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)({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
@@ -77,7 +139,8 @@ int main(int argc, char *argv[]) {
             auto loss_cpu = loss[0]->To(cpu_device);
             float current_loss = static_cast<float *>(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 +153,31 @@ int main(int argc, char *argv[]) {
         const auto epoch_end = std::chrono::high_resolution_clock::now();
         const double duration_us = std::chrono::duration<double, std::micro>(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.
+    // 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<float> 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<Tensor>(image->To(device));
         auto new_label = std::make_shared<Tensor>(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];
@@ -119,11 +190,28 @@ int main(int argc, char *argv[]) {
             }
         }
         total += batch_size;
-        test_losses.push_back(static_cast<float *>(loss_cpu.DataPtr())[0]);
+        const float batch_loss = static_cast<float *>(loss_cpu.DataPtr())[0];
+        test_losses.push_back(batch_loss);
+        local_loss_sum += static_cast<double>(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<float>(local_loss_sum), static_cast<float>(correct), static_cast<float>(total)};
+        auto stats_tensor
+            = std::make_shared<Tensor>(stats, std::vector<int64_t>{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<const float *>(stats_cpu.DataPtr());
+        if (is_main_rank) {
+            LOG(ERROR) << "Total: " << static_cast<int>(reduced[2]) << ", Correct: " << static_cast<int>(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<float>(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<float>(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<std::shared_ptr<infini_train::Tensor>>
 MNIST::Forward(const std::vector<std::shared_ptr<infini_train::Tensor>> &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<nn::Conv2d>(1, 16, 3);
+    modules_["relu1"] = std::make_shared<nn::ReLU>();
+    modules_["conv2"] = std::make_shared<nn::Conv2d>(16, 32, 3);
+    modules_["relu2"] = std::make_shared<nn::ReLU>();
+    // 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<nn::Linear>(18432, 10);
+}
+
+std::vector<std::shared_ptr<infini_train::Tensor>>
+MnistCnn::Forward(const std::vector<std::shared_ptr<infini_train::Tensor>> &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<std::shared_ptr<infini_train::Tensor>>
     Forward(const std::vector<std::shared_ptr<infini_train::Tensor>> &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<std::shared_ptr<infini_train::Tensor>>
+    Forward(const std::vector<std::shared_ptr<infini_train::Tensor>> &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<std::shared_ptr<Tensor>> &output_tensors) override;
     std::vector<std::shared_ptr<Tensor>> Backward(const std::vector<std::shared_ptr<Tensor>> &grad_outputs) override;
 };
+
+class ReLU : public Function {
+public:
+    static constexpr char kType[] = "ReLUFunction";
+
+    ReLU() : Function(kType) {}
+
+    std::vector<std::shared_ptr<Tensor>> Forward(const std::vector<std::shared_ptr<Tensor>> &input_tensors) override;
+    void SetupContext(const std::vector<std::shared_ptr<Tensor>> &input_tensors,
+                      const std::vector<std::shared_ptr<Tensor>> &output_tensors) override;
+    std::vector<std::shared_ptr<Tensor>> Backward(const std::vector<std::shared_ptr<Tensor>> &grad_outputs) override;
+};
 } // 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 <cstdint>
+#include <memory>
+#include <vector>
+
+#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<std::shared_ptr<Tensor>> Forward(const std::vector<std::shared_ptr<Tensor>> &input_tensors) override;
+    void SetupContext(const std::vector<std::shared_ptr<Tensor>> &input_tensors,
+                      const std::vector<std::shared_ptr<Tensor>> &output_tensors) override;
+    std::vector<std::shared_ptr<Tensor>> Backward(const std::vector<std::shared_ptr<Tensor>> &grad_outputs) override;
+
+private:
+    int64_t stride_ = 1;
+    int64_t padding_ = 0;
+    bool has_bias_ = false;
+    std::vector<int64_t> input_dims_;
+    std::vector<int64_t> 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<SwiGLU> {
 
     std::vector<std::shared_ptr<Tensor>> Forward(const std::vector<std::shared_ptr<Tensor>> &x) override;
 };
+
+class ReLU : public CloneableModule<ReLU> {
+public:
+    static constexpr char kType[] = "ReLU";
+    ReLU() : CloneableModule(kType) {}
+    std::vector<std::shared_ptr<Tensor>> Forward(const std::vector<std::shared_ptr<Tensor>> &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 <cstdint>
+#include <memory>
+#include <vector>
+
+#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<Conv2d> {
+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<std::shared_ptr<Tensor>> Forward(const std::vector<std::shared_ptr<Tensor>> &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<std::shared_ptr<Tensor>> Sigmoid::Backward(const std::vector<std::sh
     auto device = output->GetDevice().type();
     return {Dispatcher::Instance().Call<std::shared_ptr<Tensor>>({device, "SigmoidBackward"}, output, grad_output)};
 }
+
+std::vector<std::shared_ptr<Tensor>> ReLU::Forward(const std::vector<std::shared_ptr<Tensor>> &input_tensors) {
+    CHECK_EQ(input_tensors.size(), 1);
+    const auto &input = input_tensors[0];
+
+    auto device = input->GetDevice().type();
+    return {Dispatcher::Instance().Call<std::shared_ptr<Tensor>>({device, "ReLUForward"}, input)};
+}
+
+void ReLU::SetupContext(const std::vector<std::shared_ptr<Tensor>> &input_tensors,
+                        const std::vector<std::shared_ptr<Tensor>> &) {
+    CHECK_EQ(input_tensors.size(), 1);
+    const auto &input = input_tensors[0];
+    ctx_.SaveForBackward({input});
+}
+
+std::vector<std::shared_ptr<Tensor>> ReLU::Backward(const std::vector<std::shared_ptr<Tensor>> &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<std::shared_ptr<Tensor>>({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<std::shared_ptr<Tensor>> Conv2d::Forward(const std::vector<std::shared_ptr<Tensor>> &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<Tensor> &bias = input_tensors.size() == 3 ? input_tensors[2] : nullptr;
+
+    auto device = input->GetDevice().type();
+    return {Dispatcher::Instance().Call<std::shared_ptr<Tensor>>({device, "Conv2dForward"}, input, weight, bias,
+                                                                 stride_, padding_)};
+}
+
+void Conv2d::SetupContext(const std::vector<std::shared_ptr<Tensor>> &input_tensors,
+                          const std::vector<std::shared_ptr<Tensor>> &) {
+    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<std::shared_ptr<Tensor>> Conv2d::Backward(const std::vector<std::shared_ptr<Tensor>> &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<Tensor> grad_input = nullptr;
+    std::shared_ptr<Tensor> grad_weight = nullptr;
+    std::shared_ptr<Tensor> grad_bias = nullptr;
+
+    if (need_grad_input) {
+        grad_input = Dispatcher::Instance().Call<std::shared_ptr<Tensor>>({device, "Conv2dBackwardInput"}, weight,
+                                                                          grad_output, stride_, padding_, input_dims_);
+    }
+    if (need_grad_weight) {
+        grad_weight = Dispatcher::Instance().Call<std::shared_ptr<Tensor>>(
+            {device, "Conv2dBackwardWeight"}, input, grad_output, stride_, padding_, weight_dims_);
+    }
+    if (need_grad_bias) {
+        grad_bias = Dispatcher::Instance().Call<std::shared_ptr<Tensor>>({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<Tensor> Stack(const std::vector<std::shared_ptr<Tensor>> &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<int64_t>());
-    auto stacked_tensor = std::make_shared<Tensor>(std::vector<int64_t>{batch_size, stacked_dim}, tensors[0]->Dtype());
+    const int64_t batch_size = static_cast<int64_t>(tensors.size());
+    const auto &sample_dims = tensors[0]->Dims();
     for (const auto &tensor : tensors) {
-        CHECK_EQ(static_cast<int>(tensors[0]->Dtype()), static_cast<int>(tensor->Dtype()));
-        const auto &dims = tensor->Dims();
-        CHECK_EQ(stacked_dim, std::accumulate(dims.begin(), dims.end(), 1, std::multiplies<int64_t>()));
+        CHECK_EQ(static_cast<int>(tensors[0]->Dtype()), static_cast<int>(tensor->Dtype()))
+            << "Cannot stack tensors with different dtypes.";
+        CHECK(tensor->Dims() == sample_dims) << "Cannot stack tensors with different shapes.";
     }
+    std::vector<int64_t> stacked_dims = {batch_size};
+    stacked_dims.insert(stacked_dims.end(), sample_dims.begin(), sample_dims.end());
+    auto stacked_tensor = std::make_shared<Tensor>(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 <cstdint>
+#include <memory>
+#include <vector>
+
+#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<Tensor> &input, const std::shared_ptr<Tensor> &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<int64_t> 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 <cstdint>
+#include <memory>
+#include <vector>
+
+#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> &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<Tensor> Conv2dForward(const std::shared_ptr<Tensor> &input, const std::shared_ptr<Tensor> &weight,
+                                      const std::shared_ptr<Tensor> &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<Tensor>(Conv2dOutputDims(meta), DataType::kFLOAT32);
+    const float *input_ptr = static_cast<const float *>(input->DataPtr());
+    const float *weight_ptr = static_cast<const float *>(weight->DataPtr());
+    const float *bias_ptr = bias ? static_cast<const float *>(bias->DataPtr()) : nullptr;
+    float *output_ptr = static_cast<float *>(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<Tensor> Conv2dBackwardInput(const std::shared_ptr<Tensor> &weight,
+                                            const std::shared_ptr<Tensor> &grad_output, int64_t stride, int64_t padding,
+                                            const std::vector<int64_t> &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<Tensor>(input_dims, DataType::kFLOAT32);
+    grad_input->Fill(0.0f);
+    const float *weight_ptr = static_cast<const float *>(weight->DataPtr());
+    const float *grad_output_ptr = static_cast<const float *>(grad_output->DataPtr());
+    float *grad_input_ptr = static_cast<float *>(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<Tensor> Conv2dBackwardWeight(const std::shared_ptr<Tensor> &input,
+                                             const std::shared_ptr<Tensor> &grad_output, int64_t stride,
+                                             int64_t padding, const std::vector<int64_t> &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<Tensor>(weight_dims, DataType::kFLOAT32);
+    grad_weight->Fill(0.0f);
+    const float *input_ptr = static_cast<const float *>(input->DataPtr());
+    const float *grad_output_ptr = static_cast<const float *>(grad_output->DataPtr());
+    float *grad_weight_ptr = static_cast<float *>(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<Tensor> Conv2dBackwardBias(const std::shared_ptr<Tensor> &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<Tensor>(std::vector<int64_t>{out_channels}, DataType::kFLOAT32);
+    grad_bias->Fill(0.0f);
+    const float *grad_output_ptr = static_cast<const float *>(grad_output->DataPtr());
+    float *grad_bias_ptr = static_cast<float *>(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 <memory>
+
+#include "glog/logging.h"
+
+#include "infini_train/include/dispatcher.h"
+#include "infini_train/include/tensor.h"
+
+namespace infini_train::kernels::cpu {
+std::shared_ptr<Tensor> ReLUForward(const std::shared_ptr<Tensor> &input) {
+    CHECK(input->Dtype() == DataType::kFLOAT32) << "CPU ReLUForward currently supports float32 only";
+    auto output = std::make_shared<Tensor>(input->Dims(), DataType::kFLOAT32);
+    const float *input_ptr = static_cast<const float *>(input->DataPtr());
+    float *output_ptr = static_cast<float *>(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<Tensor> ReLUBackward(const std::shared_ptr<Tensor> &input, const std::shared_ptr<Tensor> &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<Tensor>(input->Dims(), DataType::kFLOAT32);
+    const float *input_ptr = static_cast<const float *>(input->DataPtr());
+    const float *grad_output_ptr = static_cast<const float *>(grad_output->DataPtr());
+    float *grad_input_ptr = static_cast<float *>(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 <cstdint>
+#include <memory>
+#include <vector>
+
+#include <cub/block/block_reduce.cuh>
+
+#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<int64_t>(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<int64_t>(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 <int BLOCK_SIZE>
+__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<float, BLOCK_SIZE>;
+    __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<int64_t>(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<int64_t>(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<Tensor> Im2colForward(const std::shared_ptr<Tensor> &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<Tensor>(std::vector<int64_t>{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<int>((total + threads_per_block - 1) / threads_per_block);
+    auto device = input->GetDevice();
+    const auto &cuda_stream = dynamic_cast<infini_train::core::cuda::CudaStream *>(
+                                  infini_train::core::GetDeviceGuardImpl(device.type())->GetStream(device))
+                                  ->cuda_stream();
+    Im2colKernel<<<num_blocks, threads_per_block, 0, cuda_stream>>>(static_cast<const float *>(input->DataPtr()),
+                                                                    static_cast<float *>(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<Tensor> Conv2dForward(const std::shared_ptr<Tensor> &input, const std::shared_ptr<Tensor> &weight,
+                                      const std::shared_ptr<Tensor> &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<Tensor>(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::cuda::CudaStream *>(
+                                  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<int>((total + threads_per_block - 1) / threads_per_block);
+        ConvBiasCopyKernel<<<num_blocks, threads_per_block, 0, cuda_stream>>>(
+            static_cast<float *>(output->DataPtr()), static_cast<const float *>(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<int>(meta.batch);
+    const bool batched = meta.batch > 1;
+    Dispatcher::Instance().Call<void>({device.type(), "Gemm"}, device,
+                                      GemmParams{
+                                          .trans_a = GemmTranspose::kNoTranspose,
+                                          .trans_b = GemmTranspose::kNoTranspose,
+                                          .m = static_cast<int>(meta.patches),
+                                          .n = static_cast<int>(meta.out_channels),
+                                          .k = static_cast<int>(meta.kernel_elems),
+                                          .A = columns->DataPtr(),
+                                          .lda = static_cast<int>(meta.patches),
+                                          .B = weight->DataPtr(),
+                                          .ldb = static_cast<int>(meta.kernel_elems),
+                                          .C = output->DataPtr(),
+                                          .ldc = static_cast<int>(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<int64_t> const&) so the dispatcher resolves the
+// identical "Conv2dBackwardWeight" overload on CUDA.
+std::shared_ptr<Tensor> Conv2dBackwardWeight(const std::shared_ptr<Tensor> &input,
+                                             const std::shared_ptr<Tensor> &grad_output, int64_t stride,
+                                             int64_t padding, const std::vector<int64_t> &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<Tensor>(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::cuda::CudaStream *>(
+                                  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<int>(meta.batch);
+    const bool batched = meta.batch > 1;
+    void *gemm_c = grad_weight->DataPtr();
+    long long stride_c = 0;
+    std::shared_ptr<Tensor> per_image;
+    if (batched) {
+        per_image = std::make_shared<Tensor>(std::vector<int64_t>{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<void>({device.type(), "Gemm"}, device,
+                                      GemmParams{
+                                          .trans_a = GemmTranspose::kTranspose,
+                                          .trans_b = GemmTranspose::kNoTranspose,
+                                          .m = static_cast<int>(meta.kernel_elems),
+                                          .n = static_cast<int>(meta.out_channels),
+                                          .k = static_cast<int>(meta.patches),
+                                          .A = columns->DataPtr(),
+                                          .lda = static_cast<int>(meta.patches),
+                                          .B = grad_output->DataPtr(),
+                                          .ldb = static_cast<int>(meta.patches),
+                                          .C = gemm_c,
+                                          .ldc = static_cast<int>(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<int>((elems + threads_per_block - 1) / threads_per_block);
+        ConvBackwardWeightSumKernel<<<num_blocks, threads_per_block, 0, cuda_stream>>>(
+            static_cast<const float *>(per_image->DataPtr()), static_cast<float *>(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<int64_t> const&) so the dispatcher resolves the
+// identical "Conv2dBackwardInput" overload on CUDA.
+std::shared_ptr<Tensor> Conv2dBackwardInput(const std::shared_ptr<Tensor> &weight,
+                                            const std::shared_ptr<Tensor> &grad_output, int64_t stride, int64_t padding,
+                                            const std::vector<int64_t> &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<Tensor>(input_dims, DataType::kFLOAT32, device);
+    if (meta.batch == 0) {
+        return grad_input;
+    }
+
+    auto columns = std::make_shared<Tensor>(std::vector<int64_t>{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<int>(meta.batch);
+    const bool batched = meta.batch > 1;
+    Dispatcher::Instance().Call<void>({device.type(), "Gemm"}, device,
+                                      GemmParams{
+                                          .trans_a = GemmTranspose::kNoTranspose,
+                                          .trans_b = GemmTranspose::kTranspose,
+                                          .m = static_cast<int>(meta.patches),
+                                          .n = static_cast<int>(meta.kernel_elems),
+                                          .k = static_cast<int>(meta.out_channels),
+                                          .A = grad_output->DataPtr(),
+                                          .lda = static_cast<int>(meta.patches),
+                                          .B = weight->DataPtr(),
+                                          .ldb = static_cast<int>(meta.kernel_elems),
+                                          .C = columns->DataPtr(),
+                                          .ldc = static_cast<int>(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<int>((total + threads_per_block - 1) / threads_per_block);
+    const auto &cuda_stream = dynamic_cast<infini_train::core::cuda::CudaStream *>(
+                                  infini_train::core::GetDeviceGuardImpl(device.type())->GetStream(device))
+                                  ->cuda_stream();
+    Col2imKernel<<<num_blocks, threads_per_block, 0, cuda_stream>>>(static_cast<const float *>(columns->DataPtr()),
+                                                                    static_cast<float *>(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<Tensor> Conv2dBackwardBias(const std::shared_ptr<Tensor> &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<Tensor>(std::vector<int64_t>{out_channels}, DataType::kFLOAT32, device);
+    if (out_channels == 0) {
+        return grad_bias;
+    }
+
+    const auto &cuda_stream = dynamic_cast<infini_train::core::cuda::CudaStream *>(
+                                  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<kBlockSize><<<out_channels, kBlockSize, 0, cuda_stream>>>(
+        static_cast<const float *>(grad_output->DataPtr()), static_cast<float *>(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 <cstddef>
+#include <memory>
+
+#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<Tensor> &output, const std::shared_ptr<Tensor> &input) {
+    const size_t num_elements = static_cast<size_t>(output->NumElements());
+    if (num_elements == 0) {
+        return;
+    }
+    const int threads_per_block = 256;
+    const int num_blocks = static_cast<int>((num_elements + threads_per_block - 1) / threads_per_block);
+    auto device = output->GetDevice();
+    const auto &cuda_stream = dynamic_cast<infini_train::core::cuda::CudaStream *>(
+                                  infini_train::core::GetDeviceGuardImpl(device.type())->GetStream(device))
+                                  ->cuda_stream();
+    ReLUForwardKernel<<<num_blocks, threads_per_block, 0, cuda_stream>>>(
+        static_cast<float *>(output->DataPtr()), static_cast<const float *>(input->DataPtr()), num_elements);
+}
+
+void LaunchBackward(const std::shared_ptr<Tensor> &grad_input, const std::shared_ptr<Tensor> &input,
+                    const std::shared_ptr<Tensor> &grad_output) {
+    const size_t num_elements = static_cast<size_t>(grad_input->NumElements());
+    if (num_elements == 0) {
+        return;
+    }
+    const int threads_per_block = 256;
+    const int num_blocks = static_cast<int>((num_elements + threads_per_block - 1) / threads_per_block);
+    auto device = grad_input->GetDevice();
+    const auto &cuda_stream = dynamic_cast<infini_train::core::cuda::CudaStream *>(
+                                  infini_train::core::GetDeviceGuardImpl(device.type())->GetStream(device))
+                                  ->cuda_stream();
+    ReLUBackwardKernel<<<num_blocks, threads_per_block, 0, cuda_stream>>>(
+        static_cast<float *>(grad_input->DataPtr()), static_cast<const float *>(input->DataPtr()),
+        static_cast<const float *>(grad_output->DataPtr()), num_elements);
+}
+} // namespace
+
+std::shared_ptr<Tensor> ReLUForward(const std::shared_ptr<Tensor> &input) {
+    CHECK(input->Dtype() == DataType::kFLOAT32) << "CUDA ReLUForward currently supports float32 only";
+    auto output = std::make_shared<Tensor>(input->Dims(), DataType::kFLOAT32, input->GetDevice());
+    LaunchForward(output, input);
+    return output;
+}
+
+std::shared_ptr<Tensor> ReLUBackward(const std::shared_ptr<Tensor> &input, const std::shared_ptr<Tensor> &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<Tensor>(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<std::shared_ptr<Tensor>> NewGELU::Forward(const std::vector<std::sha
 std::vector<std::shared_ptr<Tensor>> SwiGLU::Forward(const std::vector<std::shared_ptr<Tensor>> &x) {
     return {x[0] * function::Sigmoid(x[0])};
 }
+
+std::vector<std::shared_ptr<Tensor>> ReLU::Forward(const std::vector<std::shared_ptr<Tensor>> &input_tensors) {
+    return std::make_shared<autograd::ReLU>()->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 <cmath>
+#include <memory>
+#include <vector>
+
+#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<Tensor>(std::vector<int64_t>{out_channels, in_channels, kernel_size, kernel_size},
+                                   DataType::kFLOAT32, device_)
+              ->RequiresGrad();
+    if (bias) {
+        parameters_[kParamBiasName]
+            = std::make_shared<Tensor>(std::vector<int64_t>{out_channels}, DataType::kFLOAT32, device_)->RequiresGrad();
+    }
+    ResetParameters();
+}
+
+std::vector<std::shared_ptr<Tensor>> Conv2d::Forward(const std::vector<std::shared_ptr<Tensor>> &input_tensors) {
+    return std::make_shared<autograd::Conv2d>(stride_, padding_)
+        ->Apply(bias_ ? std::vector<std::shared_ptr<Tensor>>{input_tensors[0], parameters_[kParamWeightName],
+                                                             parameters_[kParamBiasName]}
+                      : std::vector<std::shared_ptr<Tensor>>{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/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__<name>.bin,
+  step<NNN>__logits.bin, step<NNN>__loss.bin,
+  step<NNN>__grad__<name>.bin, step<NNN>__param__<name>.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="<f4").reshape(shape).copy()
+
+
+def read_u8(path: Path, n: int) -> 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())
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 <vector>
+
+#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<Tensor>(input_values, std::vector<int64_t>{1, 1, 3, 3}, DataType::kFLOAT32, GetDevice());
+    auto weight
+        = std::make_shared<Tensor>(weight_values, std::vector<int64_t>{1, 1, 2, 2}, DataType::kFLOAT32, GetDevice());
+    auto conv_fn = std::make_shared<autograd::Conv2d>(1, 0);
+    auto result = conv_fn->Apply({input, weight});
+    ASSERT_EQ(result.size(), 1);
+    EXPECT_EQ(result[0]->Dims(), (std::vector<int64_t>{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<Tensor>(input_values, std::vector<int64_t>{1, 1, 3, 3}, DataType::kFLOAT32, GetDevice());
+    auto weight
+        = std::make_shared<Tensor>(weight_values, std::vector<int64_t>{1, 1, 2, 2}, DataType::kFLOAT32, GetDevice());
+    auto bias = std::make_shared<Tensor>(bias_values, std::vector<int64_t>{1}, DataType::kFLOAT32, GetDevice());
+    auto conv_fn = std::make_shared<autograd::Conv2d>(1, 0);
+    auto result = conv_fn->Apply({input, weight, bias});
+    ASSERT_EQ(result.size(), 1);
+    EXPECT_EQ(result[0]->Dims(), (std::vector<int64_t>{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<float> input_values(16);
+    for (int i = 0; i < 16; ++i) { input_values[i] = static_cast<float>(i + 1); }
+    auto input = std::make_shared<Tensor>(input_values.data(), std::vector<int64_t>{1, 1, 4, 4}, DataType::kFLOAT32,
+                                          GetDevice());
+    auto weight = std::make_shared<Tensor>(std::vector<int64_t>{1, 1, 3, 3}, DataType::kFLOAT32, GetDevice());
+    weight->Fill(1.0f);
+    auto conv_fn = std::make_shared<autograd::Conv2d>(2, 1);
+    auto result = conv_fn->Apply({input, weight});
+    ASSERT_EQ(result.size(), 1);
+    EXPECT_EQ(result[0]->Dims(), (std::vector<int64_t>{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<float> input_values(12);
+    for (int i = 0; i < 12; ++i) { input_values[i] = static_cast<float>(i + 1); }
+    auto input = std::make_shared<Tensor>(input_values.data(), std::vector<int64_t>{1, 1, 3, 4}, DataType::kFLOAT32,
+                                          GetDevice());
+    auto weight = std::make_shared<Tensor>(std::vector<int64_t>{1, 1, 2, 3}, DataType::kFLOAT32, GetDevice());
+    weight->Fill(1.0f);
+    auto conv_fn = std::make_shared<autograd::Conv2d>(1, 0);
+    auto result = conv_fn->Apply({input, weight});
+    ASSERT_EQ(result.size(), 1);
+    EXPECT_EQ(result[0]->Dims(), (std::vector<int64_t>{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<Tensor>(input_values, std::vector<int64_t>{1, 1, 3, 3}, DataType::kFLOAT32, GetDevice())
+              ->RequiresGrad();
+    auto weight
+        = std::make_shared<Tensor>(weight_values, std::vector<int64_t>{1, 1, 2, 2}, DataType::kFLOAT32, GetDevice())
+              ->RequiresGrad();
+    auto bias = std::make_shared<Tensor>(std::vector<int64_t>{1}, DataType::kFLOAT32, GetDevice(), true);
+    bias->Fill(0.0f);
+    auto conv_fn = std::make_shared<autograd::Conv2d>(1, 0);
+    auto result = conv_fn->Apply({input, weight, bias});
+    auto grad = std::make_shared<Tensor>(std::vector<int64_t>{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<int64_t>{1, 1, 3, 3}));
+    EXPECT_EQ(grad_inputs[1]->Dims(), (std::vector<int64_t>{1, 1, 2, 2}));
+    EXPECT_EQ(grad_inputs[2]->Dims(), (std::vector<int64_t>{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<Tensor>(std::vector<int64_t>{1, 1, 3, 3}, DataType::kFLOAT32, GetDevice(), true);
+    input->Fill(1.0f);
+    auto weight = std::make_shared<Tensor>(std::vector<int64_t>{1, 1, 2, 2}, DataType::kFLOAT32, GetDevice(), true);
+    weight->Fill(1.0f);
+    auto conv_fn = std::make_shared<autograd::Conv2d>(1, 0);
+    auto result = conv_fn->Apply({input, weight});
+    auto grad = std::make_shared<Tensor>(std::vector<int64_t>{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<Tensor>(std::vector<int64_t>{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<int64_t>{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<Tensor>(std::vector<int64_t>{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<int64_t>{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 <cstdint>
+#include <vector>
+
+#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<float> DeterministicVals(int64_t n, int seed) {
+    std::vector<float> vals(static_cast<size_t>(n));
+    uint32_t state = static_cast<uint32_t>(seed);
+    for (int64_t i = 0; i < n; ++i) {
+        state = state * 1664525u + 1013904223u;
+        vals[static_cast<size_t>(i)] = static_cast<float>(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<int64_t> in_dims = {batch, cin, h, w};
+    const std::vector<int64_t> w_dims = {cout, cin, k, k};
+    const std::vector<int64_t> 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<Tensor>(in_vals.data(), in_dims, DataType::kFLOAT32, cpu_device);
+    auto cpu_w = std::make_shared<Tensor>(w_vals.data(), w_dims, DataType::kFLOAT32, cpu_device);
+    auto cpu_b = with_bias ? std::make_shared<Tensor>(b_vals.data(), b_dims, DataType::kFLOAT32, cpu_device) : nullptr;
+
+    auto cuda_in = std::make_shared<Tensor>(cpu_in->To(cuda_device));
+    auto cuda_w = std::make_shared<Tensor>(cpu_w->To(cuda_device));
+    std::shared_ptr<Tensor> cuda_b = nullptr;
+    if (with_bias) {
+        cuda_b = std::make_shared<Tensor>(cpu_b->To(cuda_device));
+    }
+
+    auto cpu_fn = std::make_shared<autograd::Conv2d>(stride, padding);
+    auto cuda_fn = std::make_shared<autograd::Conv2d>(stride, padding);
+    std::vector<std::shared_ptr<Tensor>> cpu_args = with_bias
+                                                      ? std::vector<std::shared_ptr<Tensor>>{cpu_in, cpu_w, cpu_b}
+                                                      : std::vector<std::shared_ptr<Tensor>>{cpu_in, cpu_w};
+    std::vector<std::shared_ptr<Tensor>> cuda_args = with_bias
+                                                       ? std::vector<std::shared_ptr<Tensor>>{cuda_in, cuda_w, cuda_b}
+                                                       : std::vector<std::shared_ptr<Tensor>>{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<const float *>(expected_cpu.DataPtr());
+    test::ExpectTensorNear(cuda_out[0], std::vector<float>(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<int64_t> in_dims = {batch, cin, h, w};
+    const std::vector<int64_t> 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<int64_t> 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<Tensor>(in_vals.data(), in_dims, DataType::kFLOAT32, cpu_device);
+    auto cpu_w = std::make_shared<Tensor>(w_vals.data(), w_dims, DataType::kFLOAT32, cpu_device)->RequiresGrad();
+    auto cpu_go = std::make_shared<Tensor>(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<Tensor>(cpu_in->To(cuda_device));
+    auto cuda_w = std::make_shared<Tensor>(cpu_w->To(cuda_device))->RequiresGrad();
+    auto cuda_go = std::make_shared<Tensor>(cpu_go->To(cuda_device));
+
+    auto cpu_fn = std::make_shared<autograd::Conv2d>(stride, padding);
+    auto cuda_fn = std::make_shared<autograd::Conv2d>(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<const float *>(expected_cpu.DataPtr());
+    test::ExpectTensorNear(cuda_grads[1], std::vector<float>(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<int64_t> in_dims = {batch, cin, h, w};
+    const std::vector<int64_t> w_dims = {cout, cin, k, k};
+    const std::vector<int64_t> 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<int64_t> 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<Tensor>(in_vals.data(), in_dims, DataType::kFLOAT32, cpu_device);
+    auto cpu_w = std::make_shared<Tensor>(w_vals.data(), w_dims, DataType::kFLOAT32, cpu_device);
+    auto cpu_b = std::make_shared<Tensor>(b_vals.data(), b_dims, DataType::kFLOAT32, cpu_device)->RequiresGrad();
+    auto cpu_go = std::make_shared<Tensor>(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<Tensor>(cpu_in->To(cuda_device));
+    auto cuda_w = std::make_shared<Tensor>(cpu_w->To(cuda_device));
+    auto cuda_b = std::make_shared<Tensor>(cpu_b->To(cuda_device))->RequiresGrad();
+    auto cuda_go = std::make_shared<Tensor>(cpu_go->To(cuda_device));
+
+    auto cpu_fn = std::make_shared<autograd::Conv2d>(stride, padding);
+    auto cuda_fn = std::make_shared<autograd::Conv2d>(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<const float *>(expected_cpu.DataPtr());
+    test::ExpectTensorNear(cuda_grads[2], std::vector<float>(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<int64_t> in_dims = {batch, cin, h, w};
+    const std::vector<int64_t> 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<int64_t> 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<Tensor>(in_vals.data(), in_dims, DataType::kFLOAT32, cpu_device)->RequiresGrad();
+    auto cpu_w = std::make_shared<Tensor>(w_vals.data(), w_dims, DataType::kFLOAT32, cpu_device);
+    auto cpu_go = std::make_shared<Tensor>(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<Tensor>(cpu_in->To(cuda_device))->RequiresGrad();
+    auto cuda_w = std::make_shared<Tensor>(cpu_w->To(cuda_device));
+    auto cuda_go = std::make_shared<Tensor>(cpu_go->To(cuda_device));
+
+    auto cpu_fn = std::make_shared<autograd::Conv2d>(stride, padding);
+    auto cuda_fn = std::make_shared<autograd::Conv2d>(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<const float *>(expected_cpu.DataPtr());
+    test::ExpectTensorNear(cuda_grads[0], std::vector<float>(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 <cmath>
+#include <cstddef>
+#include <memory>
+#include <vector>
+
+#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<int64_t> kInputDims = {1, 2, 4, 5};
+const std::vector<int64_t> kWeightDims = {2, 2, 3, 3};
+const std::vector<int64_t> kBiasDims = {2};
+const std::vector<int64_t> 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<int>((idx * 37u + static_cast<unsigned>(salt) * 17u) % 11u) - 5; // [-5, 5]
+    return static_cast<float>(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<size_t>(1 * 2 * 4 * 5));
+    ASSERT_EQ(nw, static_cast<size_t>(2 * 2 * 3 * 3));
+    ASSERT_EQ(ny, static_cast<size_t>(1 * 2 * 2 * 3));
+
+    std::vector<float> 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<float> &x, const std::vector<float> &w,
+                       const std::vector<float> &b) {
+        auto input = std::make_shared<Tensor>(x.data(), kInputDims, DataType::kFLOAT32, device);
+        auto weight = std::make_shared<Tensor>(w.data(), kWeightDims, DataType::kFLOAT32, device);
+        auto bias = std::make_shared<Tensor>(b.data(), kBiasDims, DataType::kFLOAT32, device);
+        auto fn = std::make_shared<autograd::Conv2d>(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<const float *>(out_cpu.DataPtr());
+        double acc = 0.0;
+        for (size_t i = 0; i < ny; ++i) { acc += static_cast<double>(p[i]) * static_cast<double>(go[i]); }
+        return acc;
+    };
+
+    // Analytic grads via autograd Backward.
+    auto input = std::make_shared<Tensor>(x0.data(), kInputDims, DataType::kFLOAT32, device)->RequiresGrad();
+    auto weight = std::make_shared<Tensor>(w0.data(), kWeightDims, DataType::kFLOAT32, device)->RequiresGrad();
+    auto bias = std::make_shared<Tensor>(b0.data(), kBiasDims, DataType::kFLOAT32, device)->RequiresGrad();
+    auto fn = std::make_shared<autograd::Conv2d>(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<Tensor>(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<const std::shared_ptr<Tensor> *> analytic = {&grads[0], &grads[1], &grads[2]};
+    const std::vector<size_t> sizes = {nx, nw, nb};
+    std::vector<std::vector<float>> bases = {x0, w0, b0};
+
+    for (int p = 0; p < 3; ++p) {
+        auto g_cpu = (*analytic[p])->To(Device());
+        const auto *g = static_cast<const float *>(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<float> &xp = (p == 0 ? plus : bases[0]);
+            const std::vector<float> &xm = (p == 0 ? minus : bases[0]);
+            const std::vector<float> &wp = (p == 1 ? plus : bases[1]);
+            const std::vector<float> &wm = (p == 1 ? minus : bases[1]);
+            const std::vector<float> &bp = (p == 2 ? plus : bases[2]);
+            const std::vector<float> &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 <memory>
+#include <vector>
+
+#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<Tensor> CallIm2col(const std::shared_ptr<Tensor> &input, const kernels::Conv2dMeta &meta,
+                                   Device::DeviceType device_type) {
+    return Dispatcher::Instance().Call<std::shared_ptr<Tensor>>({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<Tensor>(input_values, std::vector<int64_t>{1, 1, 3, 3}, DataType::kFLOAT32, GetDevice());
+    auto weight = std::make_shared<Tensor>(std::vector<int64_t>{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<int64_t>{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<float> input_values(16);
+    for (int i = 0; i < 16; ++i) { input_values[i] = static_cast<float>(i + 1); }
+    auto input = std::make_shared<Tensor>(input_values.data(), std::vector<int64_t>{1, 1, 4, 4}, DataType::kFLOAT32,
+                                          GetDevice());
+    auto weight = std::make_shared<Tensor>(std::vector<int64_t>{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<int64_t>{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 <vector>
+
+#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<float> 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<float> 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<float> bias_vals = {0.1752833277f, -0.9315211177f};
+    const std::vector<float> 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<float> 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<float> 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<float> 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<float> expected_dbias = {-1.2245074511f, 3.8710851669f};
+    constexpr float kAbsTol = 1e-5f;
+
+    auto input = std::make_shared<Tensor>(input_vals.data(), std::vector<int64_t>{1, 2, 4, 5},
+                                          DataType::kFLOAT32, device)
+                     ->RequiresGrad();
+    auto weight = std::make_shared<Tensor>(weight_vals.data(), std::vector<int64_t>{2, 2, 3, 3},
+                                           DataType::kFLOAT32, device)
+                      ->RequiresGrad();
+    auto bias = std::make_shared<Tensor>(bias_vals.data(), std::vector<int64_t>{2}, DataType::kFLOAT32,
+                                         device)
+                    ->RequiresGrad();
+
+    auto fn = std::make_shared<autograd::Conv2d>(2, 1);
+    auto out = fn->Apply({input, weight, bias});
+    ASSERT_EQ(out.size(), 1);
+    EXPECT_EQ(out[0]->Dims(), (std::vector<int64_t>{1, 2, 2, 3}));
+    test::ExpectTensorNear(out[0], expected_forward, kAbsTol);
+
+    auto grad_out = std::make_shared<Tensor>(grad_out_vals.data(), std::vector<int64_t>{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<int64_t>{1, 2, 4, 5}));
+    EXPECT_EQ(grads[1]->Dims(), (std::vector<int64_t>{2, 2, 3, 3}));
+    EXPECT_EQ(grads[2]->Dims(), (std::vector<int64_t>{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 <vector>
+
+#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<float> input_values{-2.0f, -0.0f, 0.0f, 0.5f, 1.0f, 3.0f};
+    auto a = std::make_shared<Tensor>(input_values.data(), std::vector<int64_t>{2, 3}, DataType::kFLOAT32, GetDevice());
+    auto relu_fn = std::make_shared<autograd::ReLU>();
+    auto result = relu_fn->Apply({a});
+    EXPECT_EQ(result.size(), 1);
+    EXPECT_EQ(result[0]->Dims(), (std::vector<int64_t>{2, 3}));
+    test::ExpectTensorFloatEqual(result[0], {0.0f, 0.0f, 0.0f, 0.5f, 1.0f, 3.0f});
+}
+
+TEST_P(AutogradReLUTest, ReLUBackward) {
+    std::vector<float> input_values{-2.0f, -0.0f, 0.0f, 0.5f, 1.0f, 3.0f};
+    auto a = std::make_shared<Tensor>(input_values.data(), std::vector<int64_t>{2, 3}, DataType::kFLOAT32, GetDevice());
+    auto relu_fn = std::make_shared<autograd::ReLU>();
+    auto result = relu_fn->Apply({a});
+    std::vector<float> grad_values{1.0f, 2.0f, 3.0f, 4.0f, 5.0f, 6.0f};
+    auto grad
+        = std::make_shared<Tensor>(grad_values.data(), std::vector<int64_t>{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<float> input_values{-1.0f, 0.0f, 2.0f};
+    auto a = std::make_shared<Tensor>(input_values.data(), std::vector<int64_t>{3}, DataType::kFLOAT32, GetDevice());
+    nn::ReLU relu;
+    auto result = relu.Forward({a});
+    EXPECT_EQ(result.size(), 1);
+    EXPECT_EQ(result[0]->Dims(), (std::vector<int64_t>{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 <cstdint>
+#include <memory>
+#include <utility>
+#include <vector>
+
+#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<Tensor>, std::shared_ptr<Tensor>> operator[](size_t idx) const override {
+        auto data = std::make_shared<Tensor>(std::vector<int64_t>{1, 2, 2}, DataType::kFLOAT32);
+        auto label = std::make_shared<Tensor>(std::vector<int64_t>{1}, DataType::kINT64);
+        auto *data_ptr = static_cast<float *>(data->DataPtr());
+        for (int i = 0; i < 4; ++i) { data_ptr[i] = static_cast<float>(idx * 4 + i); }
+        *static_cast<int64_t *>(label->DataPtr()) = static_cast<int64_t>(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<Tensor>, std::shared_ptr<Tensor>> operator[](size_t idx) const override {
+        auto data = std::make_shared<Tensor>(idx == 0 ? std::vector<int64_t>{1, 2, 2} : std::vector<int64_t>{1, 2, 3},
+                                             DataType::kFLOAT32);
+        auto label = std::make_shared<Tensor>(std::vector<int64_t>{1}, DataType::kINT64);
+        return {data, label};
+    }
+
+    size_t Size() const override { return 2; }
+};
+} // namespace
+
+TEST(DataLoaderStackTest, PreservesSampleDims) {
+    DataLoader loader(std::make_shared<ChannelDataset>(2), 2);
+    auto it = loader.begin();
+    const auto &[x, y] = *it;
+    EXPECT_EQ(x->Dims(), (std::vector<int64_t>{2, 1, 2, 2}));
+    EXPECT_EQ(x->Dtype(), DataType::kFLOAT32);
+    const auto *data = static_cast<const float *>(x->DataPtr());
+    for (int i = 0; i < 8; ++i) { EXPECT_FLOAT_EQ(data[i], static_cast<float>(i)); }
+    EXPECT_EQ(y->Dims(), (std::vector<int64_t>{2, 1}));
+}
+
+TEST(DataLoaderStackTest, MismatchedShapesFail) {
+    DataLoader loader(std::make_shared<RaggedDataset>(), 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__<name>.bin              float32, initial param `<name>` (e.g. conv1.weight)
+//   step<NNN>__logits.bin        float32 [B,10], NNN = 000..009
+//   step<NNN>__loss.bin          float32 scalar (1 element, mean CE loss)
+//   step<NNN>__grad__<name>.bin  float32, grad of `<name>` BEFORE the SGD step
+//   step<NNN>__param__<name>.bin float32, value of `<name>` 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 <algorithm>
+#include <cmath>
+#include <cstdint>
+#include <cstdlib>
+#include <filesystem>
+#include <fstream>
+#include <iomanip>
+#include <iostream>
+#include <random>
+#include <sstream>
+#include <string>
+#include <utility>
+#include <vector>
+
+#include "gtest/gtest.h"
+
+#include "Eigen/Dense"
+
+#ifdef USE_OMP
+#include <omp.h>
+#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<const char *>(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<const char *>(data), n * sizeof(uint8_t));
+    ASSERT_TRUE(ofs.good()) << path;
+}
+
+void DumpTensorFp32(const std::filesystem::path &path, const std::shared_ptr<infini_train::Tensor> &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<const float *>(t->DataPtr()), t->NumElements());
+}
+
+std::string DimsStr(const std::vector<int64_t> &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<int64_t> &dims,
+              const std::vector<std::pair<std::string, std::shared_ptr<infini_train::Tensor>>> &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<MnistCnn>();
+    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<float>(fan_in));
+            std::uniform_real_distribution<float> dis(-bound, bound);
+            float *dst = static_cast<float *>(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<float> input_data(kBatch * 1 * 28 * 28);
+    {
+        std::mt19937 gen(kInputSeed);
+        std::uniform_real_distribution<float> dis(0.0f, 1.0f);
+        for (auto &v : input_data) { v = dis(gen); }
+    }
+    auto input = std::make_shared<it::Tensor>(std::vector<int64_t>{kBatch, 1, 28, 28}, it::DataType::kFLOAT32,
+                                              it::Device());
+    std::copy(input_data.begin(), input_data.end(), static_cast<float *>(input->DataPtr()));
+
+    // Fixed labels.
+    std::vector<uint8_t> label_data(kBatch);
+    {
+        std::mt19937 gen(kLabelSeed);
+        for (auto &v : label_data) { v = static_cast<uint8_t>(gen() % 10); }
+    }
+    auto labels = std::make_shared<it::Tensor>(std::vector<int64_t>{kBatch}, it::DataType::kUINT8, it::Device());
+    std::copy(label_data.begin(), label_data.end(), static_cast<uint8_t *>(labels->DataPtr()));
+
+    auto loss_fn = std::make_shared<it::nn::CrossEntropyLoss>();
+    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<int>(label_data[i]); }
+        meta << "\n";
+        for (const auto &[name, param] : named) {
+            meta << "param " << name << ": float32 " << DimsStr(param->Dims()) << " init__" << name
+                 << ".bin step<NNN>__grad__" << name << ".bin step<NNN>__param__" << name << ".bin\n";
+        }
+        meta << "logits: float32 [B,10] step<NNN>__logits.bin\nloss: float32 scalar step<NNN>__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<float> 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<int64_t>{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<const float *>(outputs[0]->DataPtr());
+        const float loss_val = static_cast<const float *>(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 <cstdint>
+#include <cstdlib>
+#include <filesystem>
+#include <fstream>
+#include <string>
+#include <vector>
+
+#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<char *>(bytes), 4);
+    return (static_cast<uint32_t>(bytes[0]) << 24) | (static_cast<uint32_t>(bytes[1]) << 16)
+         | (static_cast<uint32_t>(bytes[2]) << 8) | static_cast<uint32_t>(bytes[3]);
+}
+
+// Raw IDX pixel bytes for sample `idx`, read straight from disk (no framework code involved).
+std::vector<uint8_t> 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<uint8_t> raw(rows * cols);
+    ifs.seekg(16 + idx * rows * cols);
+    ifs.read(reinterpret_cast<char *>(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<char *>(&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<uint8_t> raw = ReadRawImageBytes(dir, idx);
+    ASSERT_EQ(image->NumElements(), raw.size());
+    const auto *data = static_cast<const float *>(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<const uint8_t *>(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<int64_t>{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