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;
+
     // 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/layernorm.cu

@@ -101,11 +101,9 @@ LayerNormForward(const std::shared_ptr<Tensor> &input, const std::shared_ptr<Ten
 }
 
 template <int BLOCK_SIZE, typename T>
-__global__ void LayerNormBackwardKernel(const T *__restrict__ input, const T *__restrict__ grad_output,
-                                        const float *__restrict__ mean, const float *__restrict__ rstd,
-                                        const T *__restrict__ weight, T *__restrict__ grad_input,
-                                        T *__restrict__ grad_weight, T *__restrict__ grad_bias, int embed_dim,
-                                        size_t weight_num_elements, size_t bias_num_elements) {
+__global__ void LayerNormInputGradKernel(const T *__restrict__ input, const T *__restrict__ grad_output,
+                                         const float *__restrict__ mean, const float *__restrict__ rstd,
+                                         const T *__restrict__ weight, T *__restrict__ grad_input, int embed_dim) {
     using BlockReduce = cub::BlockReduce<float, BLOCK_SIZE>;
     __shared__ typename BlockReduce::TempStorage temp_storage_mean;
     __shared__ typename BlockReduce::TempStorage temp_storage_norm;
@@ -126,17 +124,18 @@ __global__ void LayerNormBackwardKernel(const T *__restrict__ input, const T *__
     float dnorm_norm_mean = 0.f;
 
     for (int i = tid; i < embed_dim; i += BLOCK_SIZE) {
-        float dnorm = common::cuda::Cast<float>(common::cuda::Mul(weight[i], grad_output_ptr[i]));
+        float dnorm = common::cuda::Cast<float>(weight[i]) * common::cuda::Cast<float>(grad_output_ptr[i]);
+        float norm = (common::cuda::Cast<float>(input_ptr[i]) - mean_val) * rstd_val;
         dnorm_mean += dnorm;
-        dnorm_norm_mean += dnorm * (common::cuda::Cast<float>(input_ptr[i]) - mean_val);
+        dnorm_norm_mean += dnorm * norm;
     }
 
     dnorm_mean = BlockReduce(temp_storage_mean).Sum(dnorm_mean);
     dnorm_norm_mean = BlockReduce(temp_storage_norm).Sum(dnorm_norm_mean);
 
     if (tid == 0) {
         float mean_d = dnorm_mean / embed_dim;
-        float norm_d = (dnorm_norm_mean / embed_dim) * rstd_val - mean_d * mean_val * rstd_val;
+        float norm_d = dnorm_norm_mean / embed_dim;
         shared_mean = mean_d;
         shared_norm = norm_d;
     }
@@ -148,20 +147,51 @@ __global__ void LayerNormBackwardKernel(const T *__restrict__ input, const T *__
 
         grad_input_ptr[i] = common::cuda::Cast<T>(
             (common::cuda::Cast<float>(weight[i]) * grad_output_val - shared_mean - norm * shared_norm) * rstd_val);
+    }
+}
 
-        common::cuda::fastAtomicAdd<T, size_t>(grad_weight, i, weight_num_elements,
-                                               common::cuda::Cast<T>(grad_output_val * norm), true);
-        common::cuda::fastAtomicAdd<T, size_t>(grad_bias, i, bias_num_elements, grad_output_ptr[i], true);
+template <int BLOCK_SIZE, typename T>
+__global__ void LayerNormParameterGradKernel(const T *__restrict__ input, const T *__restrict__ grad_output,
+                                             const float *__restrict__ mean, const float *__restrict__ rstd,
+                                             T *__restrict__ grad_weight, T *__restrict__ grad_bias, int64_t num_tokens,
+                                             int embed_dim) {
+    using BlockReduce = cub::BlockReduce<float, BLOCK_SIZE>;
+    __shared__ typename BlockReduce::TempStorage temp_storage_weight;
+    __shared__ typename BlockReduce::TempStorage temp_storage_bias;
+
+    int feature_idx = blockIdx.x;
+    float grad_weight_sum = 0.0f;
+    float grad_bias_sum = 0.0f;
+
+    for (int64_t token_idx = threadIdx.x; token_idx < num_tokens; token_idx += BLOCK_SIZE) {
+        int64_t offset = token_idx * embed_dim + feature_idx;
+        float grad_output_val = common::cuda::Cast<float>(grad_output[offset]);
+        float norm = (common::cuda::Cast<float>(input[offset]) - mean[token_idx]) * rstd[token_idx];
+        grad_weight_sum += grad_output_val * norm;
+        grad_bias_sum += grad_output_val;
+    }
+
+    float grad_weight_reduced = BlockReduce(temp_storage_weight).Sum(grad_weight_sum);
+    float grad_bias_reduced = BlockReduce(temp_storage_bias).Sum(grad_bias_sum);
+
+    if (threadIdx.x == 0) {
+        grad_weight[feature_idx] = common::cuda::Cast<T>(grad_weight_reduced);
+        grad_bias[feature_idx] = common::cuda::Cast<T>(grad_bias_reduced);
     }
 }
 
 std::tuple<std::shared_ptr<Tensor>, std::shared_ptr<Tensor>, std::shared_ptr<Tensor>>
 LayerNormBackward(const std::shared_ptr<Tensor> &input, const std::shared_ptr<Tensor> &weight,
                   const std::shared_ptr<Tensor> &bias, const std::shared_ptr<Tensor> &mean,
                   const std::shared_ptr<Tensor> &rstd, const std::shared_ptr<Tensor> &grad_output) {
+    CHECK_EQ(input->Dims().size(), 3);
+    CHECK_LE(input->Dims()[2], weight->Dims()[0]);
+    CHECK_LE(input->Dims()[2], bias->Dims()[0]);
+
     const int batch_size = input->Dims()[0];
     const int max_seqlen = input->Dims()[1];
     const int embed_dim = input->Dims()[2];
+    const int64_t num_tokens = static_cast<int64_t>(batch_size) * max_seqlen;
 
     auto dtype = input->Dtype();
     CHECK(dtype == weight->Dtype() && dtype == bias->Dtype() && dtype == grad_output->Dtype()
@@ -182,15 +212,17 @@ LayerNormBackward(const std::shared_ptr<Tensor> &input, const std::shared_ptr<Te
     core::cuda::DispatchCudaFunc<INFINI_ALL_FLOATING_TYPES>(
         dtype,
         [=]<typename T>() {
-            // Each token block writes its complete grad_input slice; no Fill is needed.
-            grad_weight->Fill(0.0);
-            grad_bias->Fill(0.0);
-            LayerNormBackwardKernel<BLOCK_SIZE><<<num_blocks, threads_per_block, 0, cuda_stream>>>(
+            grad_weight->Fill(0);
+            grad_bias->Fill(0);
+            LayerNormInputGradKernel<BLOCK_SIZE><<<num_blocks, threads_per_block, 0, cuda_stream>>>(
+                static_cast<const T *>(input->DataPtr()), static_cast<const T *>(grad_output->DataPtr()),
+                static_cast<const float *>(mean->DataPtr()), static_cast<const float *>(rstd->DataPtr()),
+                static_cast<const T *>(weight->DataPtr()), static_cast<T *>(grad_input->DataPtr()), embed_dim);
+            LayerNormParameterGradKernel<BLOCK_SIZE><<<embed_dim, threads_per_block, 0, cuda_stream>>>(
                 static_cast<const T *>(input->DataPtr()), static_cast<const T *>(grad_output->DataPtr()),
                 static_cast<const float *>(mean->DataPtr()), static_cast<const float *>(rstd->DataPtr()),
-                static_cast<const T *>(weight->DataPtr()), static_cast<T *>(grad_input->DataPtr()),
-                static_cast<T *>(grad_weight->DataPtr()), static_cast<T *>(grad_bias->DataPtr()), embed_dim,
-                grad_weight->NumElements(), grad_bias->NumElements());
+                static_cast<T *>(grad_weight->DataPtr()), static_cast<T *>(grad_bias->DataPtr()), num_tokens,
+                embed_dim);
         },
         "CUDA LayerNormBackward");
 

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