【训练营】小模型训练支持

.gitignore

@@ -8,3 +8,6 @@ build/
 
 __pycache__/
 /data/
+
+# 邮件材料（项目报告+对齐脚本副本）：只走邮件，不进 PR
+/mail/

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>())) {

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,53 +40,107 @@ 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;
         float total_loss = 0.0;
 
         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
             // between forward and backward.
             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();

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});
+}

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;
+};

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

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

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