fix: correct DDP handling and backward gradients

CMakeLists.txt

@@ -3,7 +3,7 @@ cmake_minimum_required(VERSION 3.28)
 option(USE_CUDA "Support NVIDIA CUDA" OFF)
 option(PROFILE_MODE "ENABLE PROFILE MODE" OFF)
 option(USE_OMP "Use OpenMP as backend for Eigen" ON)
-option(USE_NCCL "Build project for distributed running" ON)
+option(USE_NCCL "Build project for distributed running on CUDA using NCCL" ON)
 option(BUILD_TEST "Build InfiniTrain tests" OFF)
 
 project(infini_train VERSION 0.6.0 LANGUAGES CXX)
@@ -64,12 +64,15 @@ endif()
 # Framework core sources (*.cc), excluding cpu kernels (they are built separately)
 file(GLOB_RECURSE SRC ${PROJECT_SOURCE_DIR}/infini_train/src/*.cc)
 list(FILTER SRC EXCLUDE REGEX ".*kernels/cpu/.*")
+
+# Exclude backend-specific runtime/ccl translation units when the corresponding
+# backend is disabled. This keeps each build self-contained and avoids pulling
+# in headers (e.g. <cuda_runtime.h> / <mcr/mc_runtime.h>) that are not on the
+# include path.
 if(NOT USE_CUDA)
-  list(FILTER SRC EXCLUDE REGEX ".*runtime/cuda/.*")
-  list(FILTER SRC EXCLUDE REGEX ".*ccl/cuda/.*")
-endif()
-if(NOT USE_NCCL)
-  list(FILTER SRC EXCLUDE REGEX ".*infini_train/src/core/ccl/cuda/.*")
+  list(FILTER SRC EXCLUDE REGEX ".*/(ccl|runtime)/cuda/.*")
+elseif(NOT USE_NCCL)
+  list(FILTER SRC EXCLUDE REGEX ".*/ccl/cuda/.*")
 endif()
 
 # CPU kernels (*.cc)

docs/test_infrastructure_design.md

@@ -8,15 +8,16 @@
 tests/
 ├── CMakeLists.txt              # 顶层：include 宏 + add_subdirectory
 ├── common/
-│   ├── CMakeLists.txt          # header-only interface library
+│   ├── CMakeLists.txt          # 公共 test_main target
 │   └── test_utils.h            # C++ 基类、skip 宏、填充工具函数
 ├── tensor/                     # Tensor 创建 / 拷贝 / 销毁 / 算子
 ├── optimizer/                  # Optimizer 创建 / step
 ├── autograd/                   # 各 autograd op 的 forward / backward
 ├── hook/                       # Module hook + precision check
 ├── lora/                       # LoRA 相关
 ├── dtype/                      # Scalar / dtype dispatch + 编译期负面测试
-└── transformer/                # Transformer 架构测试
+├── transformer/                # Transformer 架构测试
+└── checkpoint/                 # Checkpoint 序列化测试
 
 cmake/
 └── test_macros.cmake           # CMake 宏：infini_train_add_test / infini_train_add_test_suite
@@ -88,11 +89,11 @@ ctest -L cpu --output-on-failure
 ctest -L cuda --output-on-failure
 
 # 运行单个测试二进制（看完整 GTest 输出）
-./test_tensor_cpu
-./test_autograd_cuda
+./tests/tensor/test_tensor_cpu
+./tests/autograd/test_autograd_cuda
 
 # GTest filter 过滤特定用例
-./test_tensor_cpu --gtest_filter="CPU/TensorCreateTest.*"
+./tests/tensor/test_tensor_cpu --gtest_filter="CPU/TensorCreateTest.*"
 ```
 
 无 GPU 机器上 `cmake -DBUILD_TEST=ON -DUSE_CUDA=OFF ..` 即可，CUDA 测试实例不会注册。

docs/test_usage_guide.md

@@ -36,8 +36,8 @@ ctest -L cuda --output-on-failure
 ctest -R tensor --output-on-failure
 
 # 直接运行测试二进制，使用 GTest 过滤器
-./tests/tensor/test_tensor_create_cpu --gtest_filter="CPU/TensorCreateTest.*"
-./tests/tensor/test_tensor_create_cuda --gtest_filter="CUDA/TensorCreateTest.*"
+./tests/tensor/test_tensor_cpu --gtest_filter="CPU/TensorCreateTest.*"
+./tests/tensor/test_tensor_cuda --gtest_filter="CUDA/TensorCreateTest.*"
 ```
 
 ---
@@ -75,7 +75,9 @@ INFINI_TRAIN_REGISTER_TEST(TensorCopyTest);
 在子目录的 `CMakeLists.txt`（例如 `tests/tensor/CMakeLists.txt`）中添加：
 
 ```cmake
-infini_train_add_test_suite(test_tensor_copy test_tensor_copy.cc)
+infini_train_add_test_suite(test_tensor_copy
+  SOURCES test_tensor_copy.cc
+)
 ```
 
 这会生成两个 CTest 目标：`test_tensor_copy_cpu`（标签 `cpu`）和 `test_tensor_copy_cuda`（标签 `cuda`）。

example/gpt2/main.cc

@@ -4,6 +4,7 @@
 #include <format>
 #include <memory>
 #include <optional>
+#include <thread>
 #include <unordered_map>
 #include <unordered_set>
 

example/llama3/main.cc

@@ -3,6 +3,7 @@
 #include <format>
 #include <memory>
 #include <optional>
+#include <thread>
 #include <unordered_set>
 
 #include "gflags/gflags.h"

infini_train/include/autograd/elementwise.h

@@ -104,6 +104,8 @@ class Exp : public Function {
     explicit Exp() : 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;
 };
 

infini_train/include/nn/parallel/process_group.h

@@ -39,6 +39,8 @@ class ProcessGroup {
 
     virtual int GetGroupRank(int global_rank) const;
 
+    Device::DeviceType backend() const { return backend_; }
+
     // Asynchronous communication APIs (Compute / Communication stream decoupled)
     virtual std::shared_ptr<Work> AllReduce(const std::shared_ptr<Tensor> &tensor,
                                             function::ReduceOpType reduce_op = function::ReduceOpType::kSum,

infini_train/src/autograd/elementwise.cc

@@ -182,12 +182,21 @@ std::vector<std::shared_ptr<Tensor>> Exp::Forward(const std::vector<std::shared_
     return {Dispatcher::Instance().Call<std::shared_ptr<Tensor>>({device, "ExpForward"}, input)};
 }
 
+void Exp::SetupContext(const std::vector<std::shared_ptr<Tensor>> &,
+                       const std::vector<std::shared_ptr<Tensor>> &output_tensors) {
+    const auto &output = output_tensors[0];
+    ctx_.SaveForBackward({output});
+}
+
 std::vector<std::shared_ptr<Tensor>> Exp::Backward(const std::vector<std::shared_ptr<Tensor>> &grad_outputs) {
+    auto saved_tensors = ctx_.GetSavedTensors();
+    CHECK_EQ(saved_tensors.size(), 1);
+    const auto &output = saved_tensors[0];
     CHECK_EQ(grad_outputs.size(), 1);
     const auto &grad_output = grad_outputs[0];
 
-    auto device = grad_output->GetDevice().type();
-    return {Dispatcher::Instance().Call<std::shared_ptr<Tensor>>({device, "ExpBackward"}, grad_output)};
+    auto device = output->GetDevice().type();
+    return {Dispatcher::Instance().Call<std::shared_ptr<Tensor>>({device, "ExpBackward"}, grad_output, output)};
 }
 
 std::vector<std::shared_ptr<Tensor>> Log::Forward(const std::vector<std::shared_ptr<Tensor>> &input_tensors) {
@@ -397,11 +406,6 @@ std::vector<std::shared_ptr<Tensor>> Add::Backward(const std::vector<std::shared
     CHECK_EQ(grad_outputs.size(), 1);
     const auto &grad_output = grad_outputs[0];
 
-    // Fast path: no broadcast — grad_a and grad_b are both just grad_output
-    if (a_dims_ == b_dims_) {
-        return {grad_output, grad_output};
-    }
-
     auto device = grad_output->GetDevice().type();
     auto [grad_a, grad_b] = Dispatcher::Instance().Call<std::pair<std::shared_ptr<Tensor>, std::shared_ptr<Tensor>>>(
         {device, "AddBackward"}, grad_output, a_dims_, b_dims_);

infini_train/src/core/runtime/device_guard.cc

@@ -140,7 +140,7 @@ void DeviceGuardImplRegistry::Register(Device::DeviceType type, std::unique_ptr<
     }
 
     if (impls_.contains(type)) {
-        LOG(FATAL) << std::format("DeviceGuardImpl for type {} already registrered", static_cast<int>(type));
+        LOG(FATAL) << std::format("DeviceGuardImpl for type {} already registered", static_cast<int>(type));
     }
 
     if (!impls_.empty()) {

infini_train/src/kernels/cuda/linear.cu

@@ -136,22 +136,22 @@ std::shared_ptr<Tensor> LinearForward(const std::shared_ptr<Tensor> &input, cons
 }
 
 template <int BLOCK_SIZE, typename TIn, typename TOut>
-__global__ void ReduceColumnsKernel(const TIn *__restrict__ input, TOut *__restrict__ output, int num_rows,
-                                    int num_cols) {
+__global__ void ReduceRowsKernel(const TIn *__restrict__ input, TOut *__restrict__ output, int64_t num_rows,
+                                 int64_t num_cols) {
     using BlockReduce = cub::BlockReduce<float, BLOCK_SIZE>;
     __shared__ typename BlockReduce::TempStorage temp_storage;
 
-    int row = blockIdx.x;
+    const int64_t col = blockIdx.x;
     float sum = 0.0f;
 
-    for (int col = threadIdx.x; col < num_cols; col += blockDim.x) {
+    for (int64_t row = threadIdx.x; row < num_rows; row += blockDim.x) {
         sum += common::cuda::Cast<float>(input[row * num_cols + col]);
     }
 
     float reduced = BlockReduce(temp_storage).Sum(sum);
 
     if (threadIdx.x == 0) {
-        output[row] = reduced;
+        output[col] = common::cuda::Cast<TOut>(reduced);
     }
 }
 
@@ -289,7 +289,8 @@ std::shared_ptr<Tensor> LinearBackwardWeight(const std::shared_ptr<Tensor> &inpu
 std::shared_ptr<Tensor> LinearBackwardBias(const std::shared_ptr<Tensor> &grad_output, int64_t out_features) {
     const auto &dims = grad_output->Dims();
     CHECK_GE(dims.size(), 2);
-    const int64_t bs = std::accumulate(dims.rbegin() + 1, dims.rend(), 1, std::multiplies<int64_t>{});
+    CHECK_EQ(dims.back(), out_features);
+    const int64_t bs = std::accumulate(dims.rbegin() + 1, dims.rend(), int64_t{1}, std::multiplies<int64_t>{});
 
     auto compute_dtype = grad_output->Dtype();
     // FIXME(cx): output dtype promotion is a temporary hack; revisit when autograd/autocast is fixed.
@@ -307,15 +308,15 @@ std::shared_ptr<Tensor> LinearBackwardBias(const std::shared_ptr<Tensor> &grad_o
     constexpr int BLOCK_SIZE = 256;
     switch (compute_dtype) {
         DISPATCH_CASE(WRAP({
-                          ReduceColumnsKernel<BLOCK_SIZE><<<out_features, BLOCK_SIZE, 0, cuda_stream>>>(
+                          ReduceRowsKernel<BLOCK_SIZE><<<out_features, BLOCK_SIZE, 0, cuda_stream>>>(
                               static_cast<const float *>(grad_output->DataPtr()),
-                              static_cast<float *>(grad_bias->DataPtr()), out_features, bs);
+                              static_cast<float *>(grad_bias->DataPtr()), bs, out_features);
                       }),
                       DataType::kFLOAT32)
         DISPATCH_CASE(WRAP({
-                          ReduceColumnsKernel<BLOCK_SIZE><<<out_features, BLOCK_SIZE, 0, cuda_stream>>>(
+                          ReduceRowsKernel<BLOCK_SIZE><<<out_features, BLOCK_SIZE, 0, cuda_stream>>>(
                               static_cast<const nv_bfloat16 *>(grad_output->DataPtr()),
-                              static_cast<float *>(grad_bias->DataPtr()), out_features, bs);
+                              static_cast<float *>(grad_bias->DataPtr()), bs, out_features);
                       }),
                       DataType::kBFLOAT16)
     }

infini_train/src/nn/parallel/ddp/distributed_data_parallel.cc

@@ -28,23 +28,30 @@ DistributedDataParallel::DistributedDataParallel(std::shared_ptr<nn::Module> mod
     if (ddp_config_.zero_stage == 3) {
         LOG(FATAL) << "DistributedDataParallel: ZeRO-3 is not implemented yet.";
     }
+    CHECK_NOTNULL(ddp_pg_);
+    const auto expected_backend = ddp_pg_->backend();
+    const int expected_device_index = global::GetDeviceIndex(rank.thread_rank());
+    const auto validate_device = [expected_backend, expected_device_index](Device device, const char *kind) {
+        CHECK_EQ(static_cast<int>(device.type()), static_cast<int>(expected_backend))
+            << "DistributedDataParallel " << kind << " backend must match the process group backend";
+        CHECK_EQ(device.index(), expected_device_index)
+            << "DistributedDataParallel " << kind << " must use the device assigned to this rank";
+    };
+
     for (auto &param : module->Parameters()) {
+        auto device = param->GetDevice();
+        validate_device(device, "parameter");
         if (!param->requires_grad()) {
             continue;
         }
-        auto device = param->GetDevice();
-        CHECK_EQ(device.index(), global::GetDeviceIndex(rank.thread_rank()))
-            << "All parameters must be on the same device as the module";
         if (!ddp_config.gradient_bucketing_enabled && ddp_config.zero_stage < 1) {
-            auto hook = std::make_unique<infini_train::autograd::AllReducePostAccumulateHook>(
-                function::ReduceOpType::kAvg, ddp_pg_);
+            const auto reduce_op
+                = ddp_config.average_in_collective ? function::ReduceOpType::kAvg : function::ReduceOpType::kSum;
+            auto hook = std::make_unique<infini_train::autograd::AllReducePostAccumulateHook>(reduce_op, ddp_pg_);
             param->RegisterPostAccumulateGradHook(std::move(hook));
         }
     }
-    for (auto &buffer : module->Buffers()) {
-        CHECK_EQ(buffer->GetDevice().index(), global::GetDeviceIndex(rank.thread_rank()))
-            << "All buffers must be on the same device as the module";
-    }
+    for (auto &buffer : module->Buffers()) { validate_device(buffer->GetDevice(), "buffer"); }
     modules_[kModuleName] = std::move(module);
 
     if (ddp_config.zero_stage >= 1) {

infini_train/src/nn/parallel/ddp/reducer.cc

@@ -401,7 +401,9 @@ void Reducer::FinalizeBucketDense(size_t bucket_index) {
         // FIXME(zbl): support custom hook later
         LOG(FATAL) << "Custom hook is not supported now";
     } else {
-        bucket.work = ddp_pg->AllReduce(bucket.contents, function::ReduceOpType::kAvg, true);
+        const auto reduce_op
+            = ddp_config_.average_in_collective ? function::ReduceOpType::kAvg : function::ReduceOpType::kSum;
+        bucket.work = ddp_pg->AllReduce(bucket.contents, reduce_op, true);
     }
 }
 

infini_train/src/tensor.cc

@@ -157,6 +157,8 @@ Tensor Tensor::To(Device device) {
         //  1. D2H
         Tensor cpu_tensor = To(Device());
         //  2. H2D
+        // FIXME: Use the destination device for the guard, runtime implementation, and stream
+        // when cross-backend copies are supported.
         core::DeviceGuard guard(buffer_device);
         auto *impl = core::GetDeviceGuardImpl(buffer_device.type());
         impl->MemcpyAsync(new_tensor.DataPtr(), cpu_tensor.DataPtr(), SizeInBytes(), core::MemcpyKind::kH2D,

tests/autograd/test_autograd_elementwise_backward.cc

@@ -4,13 +4,35 @@
 #include "gtest/gtest.h"
 
 #include "infini_train/include/autograd/elementwise.h"
+#include "infini_train/include/core/runtime/device_guard.h"
 #include "infini_train/include/nn/parallel/global.h"
 #include "infini_train/include/tensor.h"
 
 #include "tests/common/test_utils.h"
 
 using namespace infini_train;
 
+namespace {
+
+void ExpectExpGradient(const std::shared_ptr<Tensor> &actual, const std::vector<float> &input_values,
+                       const std::vector<float> &grad_values, Device expected_device) {
+    ASSERT_NE(actual, nullptr);
+    ASSERT_EQ(input_values.size(), grad_values.size());
+    ASSERT_EQ(actual->NumElements(), input_values.size());
+    EXPECT_EQ(actual->Dtype(), DataType::kFLOAT32);
+    EXPECT_EQ(actual->GetDevice(), expected_device);
+
+    auto actual_cpu = actual->To(Device());
+    core::GetDeviceGuardImpl(actual->GetDevice().type())->SynchronizeDevice(actual->GetDevice());
+    const auto *actual_data = static_cast<const float *>(actual_cpu.DataPtr());
+    for (size_t idx = 0; idx < input_values.size(); ++idx) {
+        const float expected = grad_values[idx] * std::exp(input_values[idx]);
+        EXPECT_NEAR(actual_data[idx], expected, 1e-5f) << "Mismatch at index " << idx;
+    }
+}
+
+} // namespace
+
 class AutogradElementwiseBackwardTest : public infini_train::test::InfiniTrainTest {};
 
 TEST_P(AutogradElementwiseBackwardTest, AddBackward) {
@@ -23,7 +45,9 @@ TEST_P(AutogradElementwiseBackwardTest, AddBackward) {
     auto grad = std::make_shared<Tensor>(std::vector<int64_t>{2, 3}, DataType::kFLOAT32, GetDevice(), true);
     grad->Fill(1.0f);
     auto grad_inputs = add_fn->Backward({grad});
-    EXPECT_EQ(grad_inputs.size(), 2);
+    ASSERT_EQ(grad_inputs.size(), 2);
+    EXPECT_NE(grad_inputs[0].get(), grad_inputs[1].get());
+    EXPECT_NE(grad_inputs[0]->DataPtr(), grad_inputs[1]->DataPtr());
 }
 
 TEST_P(AutogradElementwiseBackwardTest, SubBackward) {
@@ -126,14 +150,33 @@ TEST_P(AutogradElementwiseBackwardTest, TanhBackward) {
 }
 
 TEST_P(AutogradElementwiseBackwardTest, ExpBackward) {
-    auto a = std::make_shared<Tensor>(std::vector<int64_t>{2, 3}, DataType::kFLOAT32, GetDevice(), true);
-    a->Fill(1.0f);
+    const std::vector<float> input_values = {-1.0f, -0.5f, 0.0f, 0.5f, 1.0f, 2.0f};
+    const std::vector<float> grad_values = {0.25f, -0.5f, 1.0f, 1.5f, -2.0f, 0.125f};
+    auto a = std::make_shared<Tensor>(input_values.data(), std::vector<int64_t>{2, 3}, DataType::kFLOAT32, GetDevice());
+    a->RequiresGrad();
     auto exp_fn = std::make_shared<autograd::Exp>();
     auto result = exp_fn->Apply({a});
-    auto grad = std::make_shared<Tensor>(std::vector<int64_t>{2, 3}, DataType::kFLOAT32, GetDevice(), true);
-    grad->Fill(1.0f);
+    ASSERT_EQ(result.size(), 1);
+    auto grad
+        = std::make_shared<Tensor>(grad_values.data(), std::vector<int64_t>{2, 3}, DataType::kFLOAT32, GetDevice());
     auto grad_inputs = exp_fn->Backward({grad});
-    EXPECT_EQ(grad_inputs.size(), 1);
+    ASSERT_EQ(grad_inputs.size(), 1);
+    ExpectExpGradient(grad_inputs[0], input_values, grad_values, GetDevice());
+}
+
+TEST_P(AutogradElementwiseBackwardTest, ExpBackwardAccumulatesIntoLeaf) {
+    const std::vector<float> input_values = {-1.0f, -0.5f, 0.0f, 0.5f, 1.0f, 2.0f};
+    const std::vector<float> grad_values = {0.25f, -0.5f, 1.0f, 1.5f, -2.0f, 0.125f};
+    auto input
+        = std::make_shared<Tensor>(input_values.data(), std::vector<int64_t>{2, 3}, DataType::kFLOAT32, GetDevice());
+    input->RequiresGrad();
+    auto output = input->Exp();
+    auto grad
+        = std::make_shared<Tensor>(grad_values.data(), std::vector<int64_t>{2, 3}, DataType::kFLOAT32, GetDevice());
+
+    output->Backward(grad);
+
+    ExpectExpGradient(input->grad(), input_values, grad_values, GetDevice());
 }
 
 TEST_P(AutogradElementwiseBackwardTest, LogBackward) {

tests/autograd/test_autograd_linear_backward.cc

@@ -3,6 +3,7 @@
 #include "gtest/gtest.h"
 
 #include "infini_train/include/autograd/linear.h"
+#include "infini_train/include/core/runtime/device_guard.h"
 #include "infini_train/include/nn/parallel/global.h"
 #include "infini_train/include/tensor.h"
 
@@ -21,10 +22,17 @@ TEST_P(AutogradLinearBackwardTest, LinearBackward) {
     bias->Fill(0.0f);
     auto linear_fn = std::make_shared<autograd::Linear>();
     auto result = linear_fn->Apply({input, weight, bias});
-    auto grad = std::make_shared<Tensor>(std::vector<int64_t>{2, 4}, DataType::kFLOAT32, GetDevice(), true);
-    grad->Fill(1.0f);
+    const float grad_values[] = {1.0f, 2.0f, 3.0f, 4.0f, 5.0f, 6.0f, 7.0f, 8.0f};
+    auto grad = std::make_shared<Tensor>(grad_values, std::vector<int64_t>{2, 4}, DataType::kFLOAT32, GetDevice());
     auto grad_inputs = linear_fn->Backward({grad});
-    EXPECT_EQ(grad_inputs.size(), 3);
+    ASSERT_EQ(grad_inputs.size(), 3);
+    ASSERT_NE(grad_inputs[2], nullptr);
+
+    auto bias_grad_cpu = grad_inputs[2]->To(Device());
+    core::GetDeviceGuardImpl(GetDevice().type())->SynchronizeDevice(GetDevice());
+    const auto *bias_grad = static_cast<const float *>(bias_grad_cpu.DataPtr());
+    const float expected_bias_grad[] = {6.0f, 8.0f, 10.0f, 12.0f};
+    for (int idx = 0; idx < 4; ++idx) { EXPECT_FLOAT_EQ(bias_grad[idx], expected_bias_grad[idx]); }
 }
 
 TEST_P(AutogradLinearBackwardTest, LinearBackwardNoBias) {