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

README.md

@@ -141,10 +141,26 @@ The generated files can be passed directly to the corresponding executables:
 
 ```bash
 ./build/mnist \
+  --model cnn \
   --device cpu \
   --dataset data/mnist
 ```
 
+`--model` selects the network, `cnn` (default) or `mlp`. Launching through
+`infini_run` distributes the training across processes with DDP; each process
+picks its GPU from `LOCAL_RANK` automatically:
+
+```bash
+./build/infini_run \
+  --nnodes=1 \
+  --nproc_per_node=2 \
+  ./build/mnist \
+    --model cnn \
+    --device cuda \
+    --dataset data/mnist \
+    --num_epoch 3
+```
+
 ##### GPT-2 124M
 
 ```bash

example/mnist/cnn_net.h

@@ -0,0 +1,44 @@
+#pragma once
+
+#include <memory>
+#include <utility>
+#include <vector>
+
+#include "glog/logging.h"
+
+#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"
+
+// Small CNN classifier for the MNIST demo. Structure matches the reference:
+// Conv2d(1,16,3) -> ReLU -> Conv2d(16,32,3) -> ReLU -> Flatten -> Linear(18432,10).
+// The DataLoader::Stack helper flattens each image into a [N, 784] matrix, so the
+// Forward entry restores the (N, 1, 28, 28) spatial layout before the first conv.
+class MnistCnn : public infini_train::nn::Module {
+public:
+    MnistCnn() {
+        std::vector<std::shared_ptr<infini_train::nn::Module>> layers;
+        // Two 3x3 valid convs shrink 28 -> 26 -> 24; no pooling in the reference net.
+        layers.push_back(std::make_shared<infini_train::nn::Conv2d>(1, 16, 3));
+        layers.push_back(std::make_shared<infini_train::nn::ReLU>());
+        layers.push_back(std::make_shared<infini_train::nn::Conv2d>(16, 32, 3));
+        layers.push_back(std::make_shared<infini_train::nn::ReLU>());
+        modules_["sequential"] = std::make_shared<infini_train::nn::Sequential>(std::move(layers));
+        // 32 * 24 * 24 = 18432 flattened features into the 10-class head.
+        modules_["linear"] = std::make_shared<infini_train::nn::Linear>(32 * 24 * 24, 10);
+    }
+
+    std::vector<std::shared_ptr<infini_train::Tensor>>
+    Forward(const std::vector<std::shared_ptr<infini_train::Tensor>> &x) override {
+        CHECK_EQ(x.size(), 1);
+        // Restore the batch dimension from the runtime shape, then reshape the flattened
+        // [N, 784] input to (N, 1, 28, 28). View is element-order preserving.
+        const auto batch = x[0]->Dims()[0];
+        auto x_view = x[0]->View({batch, 1, 28, 28});
+        auto x_feat = (*modules_["sequential"])({x_view})[0]->Flatten(1, -1);
+        return (*modules_["linear"])({x_feat});
+    }
+};

example/mnist/dataset.cc

@@ -89,7 +89,12 @@ MNISTDataset::MNISTDataset(const std::string &dataset, bool train)
           std::format("{}/{}-labels-idx1-ubyte", dataset, train ? kTrainPrefix : kTestPrefix))),
       image_dims_(image_file_.dims.begin() + 1, image_file_.dims.end()),
       label_dims_(label_file_.dims.begin() + 1, label_file_.dims.end()),
-      image_size_in_bytes_(kSN3TypeToSize.at(image_file_.type)
+      // The image tensor is normalized to FLOAT32 in the constructor body, so the per-sample
+      // byte stride must use the float element size, not the on-disk UINT8 size. After the
+      // first sample, the UINT8 stride reads the wrong bytes: most views are a misaligned
+      // mix of real pixels, and some coincide with another sample's exact copy but are still
+      // paired with an unrelated label.
+      image_size_in_bytes_(infini_train::kDataTypeToSize.at(DataType::kFLOAT32)
                            * 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

@@ -1,26 +1,41 @@
+#include <algorithm>
 #include <chrono>
 #include <cstdlib>
 #include <format>
 #include <iostream>
 #include <memory>
 #include <numeric>
+#include <optional>
 #include <vector>
 
 #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/ddp/distributed_data_parallel_config.h"
+#include "infini_train/include/nn/parallel/global.h"
+#include "infini_train/include/nn/parallel/parallel_functional.h"
+#include "infini_train/include/nn/parallel/process_group.h"
+#include "infini_train/include/nn/parallel/rank.h"
+#include "infini_train/include/nn/parallel/utils.h"
 #include "infini_train/include/optimizer.h"
+#include "infini_train/include/tensor.h"
 
+#include "example/mnist/cnn_net.h"
 #include "example/mnist/dataset.h"
 #include "example/mnist/net.h"
 
 DEFINE_string(dataset, "", "mnist dataset path");
+DEFINE_string(model, "cnn", "model type (mlp/cnn)");
 DEFINE_int32(bs, 64, "batch size");
-DEFINE_int32(num_epoch, 1, "num epochs");
-DEFINE_double(lr, 0.01, "learning rate");
+// Defaults are tuned for the CNN demo (default --model=cnn) to reach ~97.8% test accuracy.
+// The MLP reaches ~92% at these defaults; pass more epochs (e.g. --num_epoch=20) to reach ~95%.
+DEFINE_int32(num_epoch, 3, "num epochs");
+DEFINE_double(lr, 0.1, "learning rate");
 DEFINE_string(device, "cpu", "device type (cpu/cuda)");
 
 using namespace infini_train;
@@ -31,56 +46,171 @@ 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]);
 
+    CHECK_GT(FLAGS_bs, 0) << "--bs must be a positive batch size (got " << FLAGS_bs << ")";
+
+    // Consume the WORLD_SIZE/RANK/LOCAL_RANK env injected by infini_run; without the launcher
+    // every size defaults to 1 and the single-process path below is unchanged.
+    nn::parallel::global::InitAllEnv(/*nthread_per_process=*/1, /*tensor_parallel_size=*/1,
+                                     /*sequence_parallel_enabled=*/false, /*pipeline_parallel_size=*/1,
+                                     /*virtual_pipeline_parallel_size=*/1);
+    nn::parallel::Rank rank(nn::parallel::global::GetGlobalProcRank(), /*thread_rank=*/0,
+                            nn::parallel::global::GetNprocPerNode(), nn::parallel::global::GetNthreadPerProc());
+    const int ddp_world_size = nn::parallel::global::GetDataParallelSize();
+
     auto train_dataset = std::make_shared<MNISTDataset>(FLAGS_dataset, true);
-    DataLoader train_dataloader(train_dataset, FLAGS_bs);
+
+    Device device;
+    const nn::parallel::ProcessGroup *ddp_pg = nullptr;
+    int ddp_rank = 0;
+    if (rank.IsParallel()) {
+        CHECK_EQ(FLAGS_device, kDeviceCUDA) << "Distributed training requires --device=cuda";
+        // One GPU per process, taken from the LOCAL_RANK assigned by the launcher.
+        device = Device(Device::DeviceType::kCUDA, nn::parallel::global::GetDeviceIndex(rank.thread_rank()));
+        auto *pg_factory = nn::parallel::ProcessGroupFactory::Instance(device.type());
+        ddp_pg = pg_factory->GetOrCreate(nn::parallel::GetDataParallelProcessGroupName(rank.GlobalRank()),
+                                         nn::parallel::GetDataParallelGroupRanks(rank.GlobalRank()));
+        ddp_rank = ddp_pg->GetGroupRank(rank.GlobalRank());
+    } else {
+        device = FLAGS_device == kDeviceCPU ? Device() : Device(Device::DeviceType::kCUDA, 0);
+    }
+
+    // Sharded batches per rank under DDP; the legacy loader keeps the single-process path.
+    std::optional<DistributedDataLoader> ddp_train_dataloader;
+    std::optional<DataLoader> plain_train_dataloader;
+    if (ddp_pg != nullptr) {
+        // The loss AllReduce runs step-for-step on every rank, so all ranks must see the same
+        // batch count or a collective would hang out of step. DistributedDataLoader derives
+        // that count from the global batch size, which keeps every rank aligned.
+        ddp_train_dataloader.emplace(train_dataset, FLAGS_bs, ddp_rank, ddp_world_size);
+    } else {
+        plain_train_dataloader.emplace(train_dataset, FLAGS_bs);
+    }
+    const DataLoader &train_loader = ddp_pg != nullptr ? *ddp_train_dataloader : *plain_train_dataloader;
 
     // TODO(dcj): Add sampler & eval dataloader later.
+    // The test loader stays unsharded so every rank reports metrics over the full test set.
     auto test_dataset = std::make_shared<MNISTDataset>(FLAGS_dataset, false);
     DataLoader test_dataloader(test_dataset, FLAGS_bs);
 
-    auto network = MNIST();
-    Device device = FLAGS_device == kDeviceCPU ? Device() : Device(Device::DeviceType::kCUDA, 0);
+    std::shared_ptr<nn::Module> network;
+    if (FLAGS_model == kModelCNN) {
+        network = std::make_shared<MnistCnn>();
+    } else {
+        network = std::make_shared<MNIST>();
+    }
     Device cpu_device = Device();
-    network.To(device);
+    network->To(device);
+
+    auto loss_fn = std::make_shared<nn::CrossEntropyLoss>();
+    loss_fn->To(device);
+
+    // Wrap with DDP only after all device conversions: a later .To() recreates parameter
+    // tensors and would leave the gradient hooks registered at wrap time dangling.
+    if (ddp_pg != nullptr) {
+        network = std::make_shared<nn::parallel::DistributedDataParallel>(
+            network, rank, nn::parallel::DistributedDataParallelConfig{});
+        // Keep every replica starting from the same state; parameter broadcast before training
+        // is part of the DDP contract and left to the caller by the wrapper. Parameters only:
+        // the demo network carries no buffers, which PyTorch would sync alongside them.
+        ddp_pg->Broadcast(network->Parameters(), /*root_rank_in_group=*/0);
+    }
+
+    auto optimizer = optimizers::SGD(network->Parameters(), FLAGS_lr);
+
+    // Steps per epoch, needed by the progress line's "[step k/total]" field. Counted through the
+    // loader's own iterators so the total follows its batching and sharding rules.
+    int num_train_iters = 0;
+    for (auto it = train_loader.begin(); it != train_loader.end(); ++it) { ++num_train_iters; }
+
+    // Runs the full test set on every rank, once per epoch, so a single run shows how the test
+    // metrics evolve instead of reporting them only after the last epoch. no_grad: evaluation
+    // never backpropagates, and a forward-only graph would leave the parameters' grad accumulators
+    // primed with a dependency count the next training step can never satisfy, which would silently
+    // stop their gradient accumulation.
+    auto evaluate = [&](int epoch) {
+        autograd::NoGradGuard no_grad;
+        std::vector<float> test_losses;
+        int correct = 0;
+        int total = 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)({new_image});
+            auto output_cpu = outputs[0]->To(cpu_device);
+            auto loss = (*loss_fn)({outputs[0], new_label});
+            auto loss_cpu = loss[0]->To(cpu_device);
 
-    auto loss_fn = nn::CrossEntropyLoss();
-    loss_fn.To(device);
-    auto optimizer = optimizers::SGD(network.Parameters(), FLAGS_lr);
+            const int batch_size = output_cpu.Dims()[0];
+            for (int batch_idx = 0; batch_idx < batch_size; ++batch_idx) {
+                auto label_index = reinterpret_cast<uint8_t *>(label_cpu.DataPtr())[batch_idx];
+                const auto *output_values = static_cast<float *>(output_cpu.DataPtr()) + batch_idx * kNumClasses;
+                const int output_index = std::max_element(output_values, output_values + kNumClasses) - output_values;
+                if (output_index == label_index) {
+                    ++correct;
+                }
+            }
+            total += batch_size;
+            test_losses.push_back(static_cast<float *>(loss_cpu.DataPtr())[0]);
+        }
+        const auto avg_loss = std::accumulate(test_losses.begin(), test_losses.end(), 0.0) / test_losses.size();
+        LOG(ERROR) << "epoch " << epoch << " | Total: " << total << ", Correct: " << correct
+                   << ", Accuracy: " << static_cast<float>(correct) / total << ", AverageLoss: " << avg_loss;
+    };
 
     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_loader) {
             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});
+            // Zero grads before forward: DDP rebinds param.grad to its bucket view during forward.
             optimizer.ZeroGrad();
 
-            auto loss = loss_fn.Forward({outputs[0], new_label});
+            auto outputs = (*network)({new_image});
+
+            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];
+            if (ddp_pg != nullptr) {
+                // Average the per-rank loss so the logged value matches the global batch. With an
+                // uneven trailing batch the equal-weight average is a bounded per-step approximation
+                // until a sampler lands.
+                auto loss_stat
+                    = std::make_shared<Tensor>(&current_loss, std::vector<int64_t>{}, DataType::kFLOAT32, device);
+                nn::parallel::function::AllReduce(loss_stat, nn::parallel::function::ReduceOpType::kAvg, ddp_pg);
+                auto loss_stat_cpu = loss_stat->To(cpu_device);
+                current_loss = static_cast<const float *>(loss_stat_cpu.DataPtr())[0];
+            }
             total_loss += current_loss;
             if (train_idx % kNumItersOfOutputDuration == 0) {
-                LOG(ERROR) << "epoch: " << epoch << ", [" << train_idx * FLAGS_bs << "/" << train_dataset->Size()
-                           << "] "
-                           << " loss: " << current_loss;
+                // step is 1-based while samples is the count consumed before it, so both fields
+                // describe the same instant: the start of this step.
+                LOG(ERROR) << std::format("epoch {:2d}, [step {:4d}/{}] [samples {}/{}] loss: {:.6f}", epoch,
+                                          train_idx + 1, num_train_iters, train_idx * FLAGS_bs * ddp_world_size,
+                                          train_dataset->Size(), current_loss);
             }
 
             optimizer.Step();
@@ -93,37 +223,9 @@ int main(int argc, char *argv[]) {
         LOG(ERROR) << std::format("epoch {:2d}/{} | train loss {:.6f} | lr {:.2e} | ({:.2f} ms | {:.0f} samples/s)",
                                   epoch, FLAGS_num_epoch - 1, total_loss / train_idx, FLAGS_lr, duration_us / 1e3f,
                                   train_dataset->Size() / (duration_us / 1e6));
-    }
 
-    // TODO(dcj): Add no_grad() context manager later.
-    std::vector<float> test_losses;
-    int correct = 0;
-    int total = 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 output_cpu = outputs[0]->To(cpu_device);
-        auto loss = loss_fn.Forward({outputs[0], new_label});
-        auto loss_cpu = loss[0]->To(cpu_device);
-
-        const int batch_size = output_cpu.Dims()[0];
-        for (int batch_idx = 0; batch_idx < batch_size; ++batch_idx) {
-            auto label_index = reinterpret_cast<uint8_t *>(label_cpu.DataPtr())[batch_idx];
-            const auto *output_values = static_cast<float *>(output_cpu.DataPtr()) + batch_idx * kNumClasses;
-            const int output_index = std::max_element(output_values, output_values + kNumClasses) - output_values;
-            if (output_index == label_index) {
-                ++correct;
-            }
-        }
-        total += batch_size;
-        test_losses.push_back(static_cast<float *>(loss_cpu.DataPtr())[0]);
+        evaluate(epoch);
     }
-    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();

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