From c1c90288edebcadaef246be01076c5209054ef6b Mon Sep 17 00:00:00 2001 From: accelerator-llc Date: Fri, 4 Sep 2026 21:17:08 +0800 Subject: [PATCH 01/15] fix: use shared_ptr for MNIST network to avoid bad_weak_ptr --- example/mnist/main.cc | 10 +++++----- 1 file changed, 5 insertions(+), 5 deletions(-) diff --git a/example/mnist/main.cc b/example/mnist/main.cc index 7744e094..f9581dd3 100644 --- a/example/mnist/main.cc +++ b/example/mnist/main.cc @@ -47,14 +47,14 @@ int main(int argc, char *argv[]) { auto test_dataset = std::make_shared(FLAGS_dataset, false); DataLoader test_dataloader(test_dataset, FLAGS_bs); - auto network = MNIST(); + auto network = std::make_shared(); Device device = FLAGS_device == kDeviceCPU ? Device() : Device(Device::DeviceType::kCUDA, 0); Device cpu_device = Device(); - network.To(device); + network->To(device); auto loss_fn = nn::CrossEntropyLoss(); loss_fn.To(device); - auto optimizer = optimizers::SGD(network.Parameters(), FLAGS_lr); + auto optimizer = optimizers::SGD(network->Parameters(), FLAGS_lr); for (int epoch = 0; epoch < FLAGS_num_epoch; ++epoch) { int train_idx = 0; @@ -66,7 +66,7 @@ int main(int argc, char *argv[]) { auto new_image = std::make_shared(image->To(device)); auto new_label = std::make_shared(label->To(device)); - auto outputs = network.Forward({new_image}); + auto outputs = network->Forward({new_image}); optimizer.ZeroGrad(); auto loss = loss_fn.Forward({outputs[0], new_label}); @@ -104,7 +104,7 @@ int main(int argc, char *argv[]) { auto new_label = std::make_shared(label->To(device)); auto label_cpu = label->To(cpu_device); - auto outputs = network.Forward({new_image}); + 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); From 440b5228f69dcd1b9c3857ecc5fe33a1e87a8a1c Mon Sep 17 00:00:00 2001 From: accelerator-llc Date: Sat, 5 Sep 2026 13:42:34 +0800 Subject: [PATCH 02/15] feat(conv): add Conv2d operator for CPU with autograd and unit tests Implemented inline (im2col + GEMM) rather than exposing a public operator since the MNIST CNN is the only consumer. Scoped to square kernel, stride 1, padding 0, optional bias, and FP32 to meet the project need without covering the full Conv2d parameter space. Uses PyTorch cross-correlation semantics (no kernel flip) to keep numerical alignment with torch.nn.functional.conv2d. --- infini_train/include/autograd/conv.h | 31 ++ infini_train/include/nn/modules/conv.h | 31 ++ infini_train/src/autograd/conv.cc | 72 +++++ infini_train/src/kernels/cpu/conv.cc | 269 ++++++++++++++++++ infini_train/src/nn/modules/conv.cc | 43 +++ tests/autograd/test_autograd_conv_backward.cc | 223 +++++++++++++++ tests/autograd/test_autograd_conv_forward.cc | 206 ++++++++++++++ tests/autograd/test_autograd_conv_train.cc | 89 ++++++ 8 files changed, 964 insertions(+) create mode 100644 infini_train/include/autograd/conv.h create mode 100644 infini_train/include/nn/modules/conv.h create mode 100644 infini_train/src/autograd/conv.cc create mode 100644 infini_train/src/kernels/cpu/conv.cc create mode 100644 infini_train/src/nn/modules/conv.cc create mode 100644 tests/autograd/test_autograd_conv_backward.cc create mode 100644 tests/autograd/test_autograd_conv_forward.cc create mode 100644 tests/autograd/test_autograd_conv_train.cc diff --git a/infini_train/include/autograd/conv.h b/infini_train/include/autograd/conv.h new file mode 100644 index 00000000..db3655f7 --- /dev/null +++ b/infini_train/include/autograd/conv.h @@ -0,0 +1,31 @@ +#pragma once + +#include +#include +#include + +#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) {} + + std::vector> Forward(const std::vector> &input_tensors) override; + void SetupContext(const std::vector> &input_tensors, + const std::vector> &output_tensors) override; + std::vector> Backward(const std::vector> &grad_outputs) override; + +private: + bool bias_ = false; + int64_t out_channels_ = 0; + std::vector input_dims_; +}; +} // namespace infini_train::autograd diff --git a/infini_train/include/nn/modules/conv.h b/infini_train/include/nn/modules/conv.h new file mode 100644 index 00000000..e0e83853 --- /dev/null +++ b/infini_train/include/nn/modules/conv.h @@ -0,0 +1,31 @@ +#pragma once + +#include +#include + +#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 { +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, bool bias = true, Device device = Device()); + std::vector> Forward(const std::vector> &input_tensors) override; + + bool has_bias() const { return bias_; } + +private: + void ResetParameters(); + bool bias_ = true; +}; +} // namespace infini_train::nn diff --git a/infini_train/src/autograd/conv.cc b/infini_train/src/autograd/conv.cc new file mode 100644 index 00000000..b10894e0 --- /dev/null +++ b/infini_train/src/autograd/conv.cc @@ -0,0 +1,72 @@ +#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> Conv2d::Forward(const std::vector> &input_tensors) { + CHECK_GE(input_tensors.size(), 2); + const auto &input = input_tensors[0]; + const auto &weight = input_tensors[1]; + const auto &bias = input_tensors.size() == 3 ? input_tensors[2] : nullptr; + + auto device = input->GetDevice().type(); + return {Dispatcher::Instance().Call>({device, "Conv2dForward"}, input, weight, bias)}; +} + +void Conv2d::SetupContext(const std::vector> &input_tensors, + const std::vector> &) { + 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}); + + bias_ = input_tensors.size() == 3; + out_channels_ = weight->Dims()[0]; + input_dims_ = input->Dims(); +} + +std::vector> Conv2d::Backward(const std::vector> &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 = bias_ && ctx_.needs_input_grad().size() > 2 && ctx_.needs_input_grad()[2]; + + auto device = grad_output->GetDevice().type(); + + std::shared_ptr grad_input = nullptr; + std::shared_ptr grad_weight = nullptr; + std::shared_ptr grad_bias = nullptr; + + if (need_grad_input) { + grad_input = Dispatcher::Instance().Call>({device, "Conv2dBackwardInput"}, weight, + grad_output, input_dims_); + } + if (need_grad_weight) { + grad_weight = Dispatcher::Instance().Call>({device, "Conv2dBackwardWeight"}, input, + grad_output); + } + if (need_grad_bias) { + grad_bias = Dispatcher::Instance().Call>({device, "Conv2dBackwardBias"}, grad_output, + out_channels_); + } + + if (bias_) { + return {grad_input, grad_weight, grad_bias}; + } else { + return {grad_input, grad_weight}; + } +} +} // namespace infini_train::autograd diff --git a/infini_train/src/kernels/cpu/conv.cc b/infini_train/src/kernels/cpu/conv.cc new file mode 100644 index 00000000..268590d4 --- /dev/null +++ b/infini_train/src/kernels/cpu/conv.cc @@ -0,0 +1,269 @@ +#include +#include +#include + +#include "Eigen/Dense" +#include "glog/logging.h" + +#include "infini_train/include/dispatcher.h" +#include "infini_train/include/tensor.h" + +namespace infini_train::kernels::cpu { +namespace { + +using RowMatrix = Eigen::Matrix; + +// im2col for stride=1 / padding=0 cross-correlation: extracts every (C, kH, kW) sliding window of a +// (C, H, W) image into a (out_height*out_width, C*kH*kW) row-major matrix. The k-th column of a row +// corresponds to the (c, u, v) tap ordered like the row-major flattening of an (O, C, kH, kW) +// weight, so a conv step is a plain GEMM against the weight viewed as (O, C*kH*kW). +void Im2col(const float *image, int64_t channels, int64_t height, int64_t width, int64_t kernel_h, int64_t kernel_w, + int64_t out_height, int64_t out_width, float *col) { + const int64_t flat_kernel = channels * kernel_h * kernel_w; + for (int64_t i = 0; i < out_height; ++i) { + for (int64_t j = 0; j < out_width; ++j) { + float *row = col + (i * out_width + j) * flat_kernel; + for (int64_t c = 0; c < channels; ++c) { + const float *window = image + (c * height + i) * width + j; + for (int64_t u = 0; u < kernel_h; ++u) { + for (int64_t v = 0; v < kernel_w; ++v) { + row[(c * kernel_h + u) * kernel_w + v] = window[u * width + v]; + } + } + } + } + } +} + +// col2im: inverse of Im2col for gradients. An input element is covered by multiple output windows, +// so contributions from all patches touching it are accumulated. +void Col2im(const float *col, int64_t channels, int64_t height, int64_t width, int64_t kernel_h, int64_t kernel_w, + int64_t out_height, int64_t out_width, float *grad_image) { + const int64_t flat_kernel = channels * kernel_h * kernel_w; + for (int64_t i = 0; i < out_height; ++i) { + for (int64_t j = 0; j < out_width; ++j) { + const float *row = col + (i * out_width + j) * flat_kernel; + for (int64_t c = 0; c < channels; ++c) { + float *window = grad_image + (c * height + i) * width + j; + for (int64_t u = 0; u < kernel_h; ++u) { + for (int64_t v = 0; v < kernel_w; ++v) { + window[u * width + v] += row[(c * kernel_h + u) * kernel_w + v]; + } + } + } + } + } +} + +} // namespace + +std::shared_ptr Conv2dForward(const std::shared_ptr &input, const std::shared_ptr &weight, + const std::shared_ptr &bias) { + /* + Cross-correlation (PyTorch conv2d semantics, kernel not flipped), stride=1, padding=0: + output(n, o, i, j) = bias(o) + sum_{c,u,v} input(n, c, i+u, j+v) * weight(o, c, u, v) + Computed per image as output(n) = weight(o, C*kH*kW) * im2col(input(n))^T + bias + */ + + const auto &input_dims = input->Dims(); + CHECK_EQ(input_dims.size(), 4); + const int64_t batch = input_dims[0]; + const int64_t channels = input_dims[1]; + const int64_t height = input_dims[2]; + const int64_t width = input_dims[3]; + + const auto &weight_dims = weight->Dims(); + CHECK_EQ(weight_dims.size(), 4); + const int64_t out_channels = weight_dims[0]; + const int64_t kernel_h = weight_dims[2]; + const int64_t kernel_w = weight_dims[3]; + CHECK_EQ(weight_dims[1], channels); + CHECK_GE(height, kernel_h); + CHECK_GE(width, kernel_w); + + if (bias) { + CHECK_EQ(bias->Dims().size(), 1); + CHECK_EQ(bias->Dims()[0], out_channels); + } + + const int64_t out_height = height - kernel_h + 1; + const int64_t out_width = width - kernel_w + 1; + const int64_t patches = out_height * out_width; + const int64_t flat_kernel = channels * kernel_h * kernel_w; + + auto output = std::make_shared(std::vector{batch, out_channels, out_height, out_width}, + DataType::kFLOAT32, input->GetDevice()); + + const float *input_data = static_cast(input->DataPtr()); + const float *weight_data = static_cast(weight->DataPtr()); + float *output_data = static_cast(output->DataPtr()); + + // (batch, patches, C*kH*kW) scratch: the im2col expansion of every image in the batch. + const int64_t image_size = channels * height * width; + std::vector col_buffer(batch * patches * flat_kernel); + for (int64_t n = 0; n < batch; ++n) { + Im2col(input_data + n * image_size, channels, height, width, kernel_h, kernel_w, out_height, out_width, + col_buffer.data() + n * patches * flat_kernel); + } + + Eigen::Map weight_mat(weight_data, out_channels, flat_kernel); + for (int64_t n = 0; n < batch; ++n) { + Eigen::Map col_n(col_buffer.data() + n * patches * flat_kernel, patches, flat_kernel); + Eigen::Map out_n(output_data + n * out_channels * patches, out_channels, patches); + out_n.noalias() = weight_mat * col_n.transpose(); + if (bias) { + out_n.colwise() + += Eigen::Map(static_cast(bias->DataPtr()), out_channels); + } + } + + return output; +} + +std::shared_ptr Conv2dBackwardInput(const std::shared_ptr &weight, + const std::shared_ptr &grad_output, + const std::vector &input_dims) { + /* + grad_input(n, c, x, y) + = sum_{o} sum_{0<=x-uDims(); + CHECK_EQ(weight_dims.size(), 4); + const int64_t out_channels = weight_dims[0]; + const int64_t kernel_h = weight_dims[2]; + const int64_t kernel_w = weight_dims[3]; + CHECK_EQ(weight_dims[1], channels); + CHECK_GE(height, kernel_h); + CHECK_GE(width, kernel_w); + + const int64_t out_height = height - kernel_h + 1; + const int64_t out_width = width - kernel_w + 1; + const int64_t patches = out_height * out_width; + const int64_t flat_kernel = channels * kernel_h * kernel_w; + + const auto &grad_dims = grad_output->Dims(); + CHECK_EQ(grad_dims.size(), 4); + CHECK_EQ(grad_dims[0], batch); + CHECK_EQ(grad_dims[1], out_channels); + CHECK_EQ(grad_dims[2], out_height); + CHECK_EQ(grad_dims[3], out_width); + + auto grad_input = std::make_shared(input_dims, DataType::kFLOAT32, grad_output->GetDevice()); + grad_input->Fill(0.0f); + + const float *weight_data = static_cast(weight->DataPtr()); + const float *grad_output_data = static_cast(grad_output->DataPtr()); + float *grad_input_data = static_cast(grad_input->DataPtr()); + + Eigen::Map weight_mat(weight_data, out_channels, flat_kernel); + const int64_t image_size = channels * height * width; + std::vector col_buffer(patches * flat_kernel); + for (int64_t n = 0; n < batch; ++n) { + Eigen::Map grad_output_n(grad_output_data + n * out_channels * patches, out_channels, patches); + Eigen::Map col_n(col_buffer.data(), patches, flat_kernel); + col_n.noalias() = grad_output_n.transpose() * weight_mat; + Col2im(col_buffer.data(), channels, height, width, kernel_h, kernel_w, out_height, out_width, + grad_input_data + n * image_size); + } + + return grad_input; +} + +std::shared_ptr Conv2dBackwardWeight(const std::shared_ptr &input, + const std::shared_ptr &grad_output) { + /* + grad_weight(o, c, u, v) = sum_{n, i, j} grad_output(n, o, i, j) * input(n, c, i+u, j+v) + Computed per image as grad_output(n) * im2col(input(n)) accumulated over the batch, the + transpose of the forward GEMM. + */ + + const auto &input_dims = input->Dims(); + CHECK_EQ(input_dims.size(), 4); + const int64_t batch = input_dims[0]; + const int64_t channels = input_dims[1]; + const int64_t height = input_dims[2]; + const int64_t width = input_dims[3]; + + const auto &grad_dims = grad_output->Dims(); + CHECK_EQ(grad_dims.size(), 4); + CHECK_EQ(grad_dims[0], batch); + + // The kernel extent is the difference between the input and output spatial extents. + const int64_t kernel_h = height - grad_dims[2] + 1; + const int64_t kernel_w = width - grad_dims[3] + 1; + CHECK_GT(kernel_h, 0); + CHECK_GT(kernel_w, 0); + + const int64_t out_channels = grad_dims[1]; + const int64_t out_height = grad_dims[2]; + const int64_t out_width = grad_dims[3]; + const int64_t patches = out_height * out_width; + const int64_t flat_kernel = channels * kernel_h * kernel_w; + + auto grad_weight = std::make_shared(std::vector{out_channels, channels, kernel_h, kernel_w}, + DataType::kFLOAT32, input->GetDevice()); + grad_weight->Fill(0.0f); + + const float *input_data = static_cast(input->DataPtr()); + const float *grad_output_data = static_cast(grad_output->DataPtr()); + float *grad_weight_data = static_cast(grad_weight->DataPtr()); + + const int64_t image_size = channels * height * width; + std::vector col_buffer(batch * patches * flat_kernel); + for (int64_t n = 0; n < batch; ++n) { + Im2col(input_data + n * image_size, channels, height, width, kernel_h, kernel_w, out_height, out_width, + col_buffer.data() + n * patches * flat_kernel); + } + + Eigen::Map grad_weight_mat(grad_weight_data, out_channels, flat_kernel); + for (int64_t n = 0; n < batch; ++n) { + Eigen::Map grad_output_n(grad_output_data + n * out_channels * patches, out_channels, patches); + Eigen::Map col_n(col_buffer.data() + n * patches * flat_kernel, patches, flat_kernel); + grad_weight_mat.noalias() += grad_output_n * col_n; + } + + return grad_weight; +} + +std::shared_ptr Conv2dBackwardBias(const std::shared_ptr &grad_output, int64_t out_channels) { + /* + grad_bias(o) = sum_{n, i, j} grad_output(n, o, i, j) + */ + const auto &grad_dims = grad_output->Dims(); + CHECK_EQ(grad_dims.size(), 4); + CHECK_EQ(grad_dims[1], out_channels); + + auto grad_bias + = std::make_shared(std::vector{out_channels}, DataType::kFLOAT32, grad_output->GetDevice()); + grad_bias->Fill(0.0f); + Eigen::Map grad_bias_vec(static_cast(grad_bias->DataPtr()), out_channels); + + const int64_t batch = grad_dims[0]; + const int64_t patches = grad_dims[2] * grad_dims[3]; + const float *grad_output_data = static_cast(grad_output->DataPtr()); + for (int64_t n = 0; n < batch; ++n) { + Eigen::Map grad_output_n(grad_output_data + n * out_channels * patches, out_channels, patches); + grad_bias_vec += grad_output_n.rowwise().sum(); + } + return grad_bias; +} + +} // namespace infini_train::kernels::cpu + +#define REGISTER_CPU_CONV2D_KERNEL(kernel_name) \ + REGISTER_KERNEL(infini_train::Device::DeviceType::kCPU, kernel_name, infini_train::kernels::cpu::kernel_name) + +REGISTER_CPU_CONV2D_KERNEL(Conv2dForward) +REGISTER_CPU_CONV2D_KERNEL(Conv2dBackwardInput) +REGISTER_CPU_CONV2D_KERNEL(Conv2dBackwardWeight) +REGISTER_CPU_CONV2D_KERNEL(Conv2dBackwardBias) + +#undef REGISTER_CPU_CONV2D_KERNEL diff --git a/infini_train/src/nn/modules/conv.cc b/infini_train/src/nn/modules/conv.cc new file mode 100644 index 00000000..7ad8a4e4 --- /dev/null +++ b/infini_train/src/nn/modules/conv.cc @@ -0,0 +1,43 @@ +#include "infini_train/include/nn/modules/conv.h" + +#include +#include +#include + +#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, bool bias, Device device) + : CloneableModule(kType), bias_(bias) { + device_ = device; + + parameters_[kParamWeightName] + = std::make_shared(std::vector{out_channels, in_channels, kernel_size, kernel_size}, + DataType::kFLOAT32, device_) + ->RequiresGrad(); + if (bias) { + parameters_[kParamBiasName] + = std::make_shared(std::vector{out_channels}, DataType::kFLOAT32, device_)->RequiresGrad(); + } + ResetParameters(); +} + +std::vector> Conv2d::Forward(const std::vector> &input_tensors) { + return std::make_shared()->Apply( + bias_ ? std::vector>{input_tensors[0], parameters_[kParamWeightName], + parameters_[kParamBiasName]} + : std::vector>{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/tests/autograd/test_autograd_conv_backward.cc b/tests/autograd/test_autograd_conv_backward.cc new file mode 100644 index 00000000..153f7668 --- /dev/null +++ b/tests/autograd/test_autograd_conv_backward.cc @@ -0,0 +1,223 @@ +#include +#include +#include + +#include "gtest/gtest.h" + +#include "infini_train/include/autograd/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; + +namespace { +// Same golden case as the forward test (torch 2.14, seed 20260905, autograd.grad reference). +constexpr float kGoldenAbsError = 2e-6f; +} // namespace + +class AutogradConvBackwardTest : public infini_train::test::InfiniTrainTest {}; + +// Hand-checkable case: x = arange(9) reshaped (1,1,3,3), w = [[1,2],[3,4]], bias = 0.5, +// grad_output = ones(1,1,2,2). bias does not influence grad_input / grad_weight. +TEST_P(AutogradConvBackwardTest, ConvBackwardGradients) { + ONLY_CPU(); + std::vector input_values; + for (int idx = 0; idx < 9; ++idx) { input_values.push_back(static_cast(idx)); } + auto input = std::make_shared(input_values.data(), std::vector{1, 1, 3, 3}, DataType::kFLOAT32, + GetDevice()) + ->RequiresGrad(); + const std::vector weight_values{1.0f, 2.0f, 3.0f, 4.0f}; + auto weight = std::make_shared(weight_values.data(), std::vector{1, 1, 2, 2}, DataType::kFLOAT32, + GetDevice()) + ->RequiresGrad(); + auto bias = std::make_shared(std::vector{1}, DataType::kFLOAT32, GetDevice(), true); + bias->Fill(0.5f); + + auto conv_fn = std::make_shared(); + auto result = conv_fn->Apply({input, weight, bias}); + auto grad_output + = std::make_shared(std::vector{1, 1, 2, 2}, DataType::kFLOAT32, GetDevice(), true); + grad_output->Fill(1.0f); + auto grad_inputs = conv_fn->Backward({grad_output}); + ASSERT_EQ(grad_inputs.size(), 3); + ASSERT_NE(grad_inputs[0], nullptr); + ASSERT_NE(grad_inputs[1], nullptr); + ASSERT_NE(grad_inputs[2], nullptr); + + // grad_input(x, y) sums the taps of every window covering (x, y). + const std::vector expected_grad_input{1.0f, 3.0f, 2.0f, 4.0f, 10.0f, 6.0f, 3.0f, 7.0f, 4.0f}; + EXPECT_EQ(grad_inputs[0]->Dims(), (std::vector{1, 1, 3, 3})); + test::ExpectTensorFloatEqual(grad_inputs[0], expected_grad_input); + + // grad_weight(u, v) sums the input elements seen through tap (u, v) over all windows. + const std::vector expected_grad_weight{8.0f, 12.0f, 20.0f, 24.0f}; + EXPECT_EQ(grad_inputs[1]->Dims(), (std::vector{1, 1, 2, 2})); + test::ExpectTensorFloatEqual(grad_inputs[1], expected_grad_weight); + + test::ExpectTensorFloatEqual(grad_inputs[2], std::vector{4.0f}); +} + +TEST_P(AutogradConvBackwardTest, ConvBackwardNoBias) { + ONLY_CPU(); + auto input = std::make_shared(std::vector{2, 2, 4, 4}, DataType::kFLOAT32, GetDevice(), true); + input->Fill(1.0f); + auto weight = std::make_shared(std::vector{3, 2, 3, 3}, DataType::kFLOAT32, GetDevice(), true); + weight->Fill(1.0f); + auto conv_fn = std::make_shared(); + auto result = conv_fn->Apply({input, weight}); + auto grad = std::make_shared(std::vector{2, 3, 2, 2}, DataType::kFLOAT32, GetDevice(), true); + grad->Fill(1.0f); + auto grad_inputs = conv_fn->Backward({grad}); + EXPECT_EQ(grad_inputs.size(), 2); + + // No-bias path returns {grad_input, grad_weight}: both must be produced and correct. + ASSERT_NE(grad_inputs[0], nullptr); + EXPECT_EQ(grad_inputs[0]->Dims(), (std::vector{2, 2, 4, 4})); + ASSERT_NE(grad_inputs[1], nullptr); + EXPECT_EQ(grad_inputs[1]->Dims(), (std::vector{3, 2, 3, 3})); + // With all-ones input/weight/grad_output, grad_weight(o,c,u,v) = sum_{n,i,j} 1 = 2*2*2. + test::ExpectTensorFloatEqual(grad_inputs[1], 8.0f); +} + +// H = W = kernel: a single window, so grad_input mirrors the weight and grad_weight the input. +TEST_P(AutogradConvBackwardTest, ConvBackwardKernelEqualsSpatial) { + ONLY_CPU(); + const std::vector input_values{1.0f, 2.0f, 3.0f, 4.0f}; + auto input = std::make_shared(input_values.data(), std::vector{1, 1, 2, 2}, DataType::kFLOAT32, + GetDevice()) + ->RequiresGrad(); + const std::vector weight_values{5.0f, 6.0f, 7.0f, 8.0f}; + auto weight = std::make_shared(weight_values.data(), std::vector{1, 1, 2, 2}, DataType::kFLOAT32, + GetDevice()) + ->RequiresGrad(); + auto bias = std::make_shared(std::vector{1}, DataType::kFLOAT32, GetDevice(), true); + bias->Fill(0.25f); + + auto conv_fn = std::make_shared(); + auto result = conv_fn->Apply({input, weight, bias}); + auto grad_output + = std::make_shared(std::vector{1, 1, 1, 1}, DataType::kFLOAT32, GetDevice(), true); + grad_output->Fill(1.0f); + auto grad_inputs = conv_fn->Backward({grad_output}); + ASSERT_EQ(grad_inputs.size(), 3); + + EXPECT_EQ(grad_inputs[0]->Dims(), (std::vector{1, 1, 2, 2})); + test::ExpectTensorFloatEqual(grad_inputs[0], weight_values); + EXPECT_EQ(grad_inputs[1]->Dims(), (std::vector{1, 1, 2, 2})); + test::ExpectTensorFloatEqual(grad_inputs[1], input_values); + test::ExpectTensorFloatEqual(grad_inputs[2], std::vector{1.0f}); +} + +// Training-shaped call: the activation is a leaf without requires_grad, so only the parameter +// gradients must be produced. +TEST_P(AutogradConvBackwardTest, ConvBackwardInputNotRequired) { + ONLY_CPU(); + auto input = std::make_shared(std::vector{2, 2, 4, 4}, DataType::kFLOAT32, GetDevice()); + input->Fill(1.0f); + auto weight = std::make_shared(std::vector{3, 2, 3, 3}, DataType::kFLOAT32, GetDevice(), true); + weight->Fill(1.0f); + auto bias = std::make_shared(std::vector{3}, DataType::kFLOAT32, GetDevice(), true); + bias->Fill(1.0f); + + auto conv_fn = std::make_shared(); + auto result = conv_fn->Apply({input, weight, bias}); + auto grad = std::make_shared(std::vector{2, 3, 2, 2}, DataType::kFLOAT32, GetDevice(), true); + grad->Fill(1.0f); + auto grad_inputs = conv_fn->Backward({grad}); + ASSERT_EQ(grad_inputs.size(), 3); + EXPECT_EQ(grad_inputs[0], nullptr); + EXPECT_NE(grad_inputs[1], nullptr); + EXPECT_NE(grad_inputs[2], nullptr); + // 2 images * 4 windows * 9 taps of 1.0 + test::ExpectTensorFloatEqual(grad_inputs[1], 8.0f); + // 2 images * 4 spatial positions + test::ExpectTensorFloatEqual(grad_inputs[2], 8.0f); +} + +// Random 4D case checked against PyTorch autograd (fixed seed, see file-level comment). +TEST_P(AutogradConvBackwardTest, ConvBackwardTorchGolden) { + ONLY_CPU(); + const std::vector input_values{ + 2.014464855f, -0.671311080f, -0.945236862f, -0.096128546f, 0.889558852f, 1.726294398f, -0.078931913f, + 0.205890238f, 0.094608001f, 0.173616216f, 0.437049866f, -0.557070732f, 0.455583423f, -0.736482263f, + -0.718647420f, 0.926798999f, 1.732108712f, -0.138735890f, -1.508729339f, 2.127460241f, 0.079739936f, + 0.217489868f, 0.589455068f, -0.052687794f, -1.659490585f, -0.380473137f, -0.166608080f, 0.350353301f, + -1.411424637f, 0.626776755f, 1.336345553f, -0.203511983f, -0.757243156f, 1.381434083f, -0.417290032f, + -0.668979943f, -0.569557011f, -1.035447836f, -0.555948615f, -1.279154539f, 0.056448594f, -0.388204783f, + 2.071142673f, -0.803012371f, 1.956272721f, 0.610922813f, -1.668144941f, -0.628582716f, -0.756579638f, + 0.987797141f, -0.849777400f, -1.654165864f, -0.326211363f, -0.040599853f, 0.281366080f, 0.499955416f, + -0.518889010f, -0.409607649f, 0.689743340f, 1.129151702f, 0.798748374f, -1.004696369f, -0.104189672f, + 1.257721663f, + }; + auto input = std::make_shared(input_values.data(), std::vector{2, 2, 4, 4}, DataType::kFLOAT32, + GetDevice()) + ->RequiresGrad(); + const std::vector weight_values{ + -2.448281050f, -0.608355999f, -0.958948731f, 1.009868145f, 0.031469870f, -1.101765275f, 0.571945250f, + -1.224039674f, -0.426795214f, -0.612229824f, -0.744571507f, -1.164491534f, -0.726353824f, -1.124407530f, + -0.498482078f, -0.423681349f, -1.231436968f, -0.689355612f, -0.170030892f, -0.120427191f, 1.032825112f, + 1.369398594f, 1.237134099f, 0.458183825f, -0.396747291f, 0.548917472f, 0.043999992f, 0.891736746f, + -1.292207360f, 1.713299513f, 0.982937992f, 0.432634443f, 0.293460459f, 0.318505853f, -0.146791920f, + 0.983720183f, -0.162205622f, 0.616937041f, 0.929942310f, -0.411319345f, 0.301661789f, -1.444623590f, + -0.730446517f, 2.391211271f, 0.400584310f, -0.909495890f, -0.702420592f, 0.372547686f, 0.994188488f, + 0.768918931f, 0.740486383f, -1.176032424f, 0.868664801f, 0.422459871f, + }; + auto weight = std::make_shared(weight_values.data(), std::vector{3, 2, 3, 3}, DataType::kFLOAT32, + GetDevice()) + ->RequiresGrad(); + const std::vector bias_values{-0.640517414f, 0.982376575f, 0.729396820f}; + auto bias = std::make_shared(bias_values.data(), std::vector{3}, DataType::kFLOAT32, GetDevice()) + ->RequiresGrad(); + + auto conv_fn = std::make_shared(); + auto result = conv_fn->Apply({input, weight, bias}); + const std::vector grad_output_values{ + -1.277466655f, -1.339194655f, -1.159473658f, 0.086730979f, 0.438649237f, 0.602657557f, + -0.738640666f, -0.410562515f, 0.201109841f, 0.716394842f, 0.003985748f, 0.676588118f, + 0.412507862f, -0.146394536f, -0.880041420f, -1.247093678f, -1.212666154f, -0.393535852f, + 0.564200640f, 0.279093444f, -2.798053741f, -1.615590930f, -0.518746257f, -0.283052236f, + }; + auto grad_output = std::make_shared(grad_output_values.data(), std::vector{2, 3, 2, 2}, + DataType::kFLOAT32, GetDevice()); + + auto grad_inputs = conv_fn->Backward({grad_output}); + ASSERT_EQ(grad_inputs.size(), 3); + ASSERT_NE(grad_inputs[0], nullptr); + ASSERT_NE(grad_inputs[1], nullptr); + ASSERT_NE(grad_inputs[2], nullptr); + + const std::vector expected_grad_input{ + 3.020392179f, 3.908452511f, 3.049194813f, 2.572864771f, 2.191555023f, 0.285838872f, 3.004242659f, + 0.838663280f, -3.235622168f, -0.945002913f, 4.960352898f, -0.376030058f, -0.373013139f, 0.741603315f, + 1.750292301f, 0.215949476f, 0.990354240f, 0.948824525f, 2.029216766f, 2.858904839f, 1.606565118f, + 4.838562012f, 3.308694601f, 0.822548687f, 0.564552128f, 2.793590069f, 4.189336777f, 2.155962467f, + 0.251298606f, 0.576505065f, 0.615549624f, -0.177835792f, -0.349884391f, -1.143750548f, -5.110340118f, + -1.768475175f, 2.049649715f, 0.845555902f, 3.547870159f, 3.535832882f, 2.858142614f, -6.855783939f, + -3.052171230f, 1.308768988f, -0.348264217f, -0.470771074f, 1.195474505f, 0.431147456f, 1.210889816f, + 4.433355331f, -1.848075271f, -1.105654240f, -2.759691477f, -3.466245174f, -1.316413283f, 0.586114824f, + 3.407652855f, 0.809629261f, -2.299243927f, -0.474773228f, 1.162620664f, 1.500420690f, 2.191398382f, + 1.014662623f, + }; + const std::vector expected_grad_weight{ + -1.278161168f, 2.350493193f, 3.455903053f, -3.191775799f, -4.849991322f, -1.593288898f, -3.348839760f, + 1.063873172f, 4.325201988f, -2.219501734f, 2.213379383f, -2.589243412f, 2.336852074f, -1.198173761f, + -2.431171656f, 4.716279507f, 0.842501462f, -3.181046009f, -1.122435212f, -4.357190609f, -0.400157630f, + 2.311078310f, 2.235447884f, 2.117345333f, 1.470544577f, 0.713402152f, -3.305890560f, 0.861350715f, + -2.180587769f, 2.186314106f, 1.552511692f, 0.699985623f, 0.372371852f, 0.788659871f, -1.649568439f, + -1.754016399f, 1.571256876f, -3.355353594f, 2.778887510f, 4.880561829f, 3.997585535f, 2.531520605f, + -1.071573257f, -2.245790720f, -3.141553879f, 1.098210573f, -2.158659220f, 5.950089931f, 1.271268964f, + 0.028057545f, -1.955230832f, 1.795806646f, 1.293214321f, -3.971021652f, + }; + const std::vector expected_grad_bias{-5.550425529f, -0.870804369f, -3.617365122f}; + + EXPECT_EQ(grad_inputs[0]->Dims(), (std::vector{2, 2, 4, 4})); + test::ExpectTensorNear(grad_inputs[0], expected_grad_input, kGoldenAbsError); + EXPECT_EQ(grad_inputs[1]->Dims(), (std::vector{3, 2, 3, 3})); + test::ExpectTensorNear(grad_inputs[1], expected_grad_weight, kGoldenAbsError); + EXPECT_EQ(grad_inputs[2]->Dims(), (std::vector{3})); + test::ExpectTensorNear(grad_inputs[2], expected_grad_bias, kGoldenAbsError); +} + +INFINI_TRAIN_REGISTER_TEST(AutogradConvBackwardTest); diff --git a/tests/autograd/test_autograd_conv_forward.cc b/tests/autograd/test_autograd_conv_forward.cc new file mode 100644 index 00000000..26c8ea8d --- /dev/null +++ b/tests/autograd/test_autograd_conv_forward.cc @@ -0,0 +1,206 @@ +#include +#include +#include + +#include "gtest/gtest.h" + +#include "infini_train/include/autograd/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; + +namespace { +// Deterministic pseudo-random golden data, generated with torch 2.14 under seed 20260905 +// (conv2d / autograd.grad reference). +constexpr float kGoldenAbsError = 2e-6f; +} // namespace + +class AutogradConvForwardTest : public infini_train::test::InfiniTrainTest {}; + +// An asymmetric kernel pins the cross-correlation semantics: flipping the kernel would swap +// the roles of the taps and change every output value. +TEST_P(AutogradConvForwardTest, ConvForwardAsymmetricKernel) { + ONLY_CPU(); + std::vector input_values; + for (int idx = 0; idx < 16; ++idx) { input_values.push_back(static_cast(idx)); } + auto input = std::make_shared(input_values.data(), std::vector{1, 1, 4, 4}, DataType::kFLOAT32, + GetDevice()); + const std::vector weight_values{1.0f, 2.0f, 3.0f, 4.0f}; + auto weight = std::make_shared(weight_values.data(), std::vector{1, 1, 2, 2}, DataType::kFLOAT32, + GetDevice()); + // out[i][j] = 1*x[i][j] + 2*x[i][j+1] + 3*x[i+1][j] + 4*x[i+1][j+1] + const std::vector expected{34.0f, 44.0f, 54.0f, 74.0f, 84.0f, 94.0f, 114.0f, 124.0f, 134.0f}; + + auto conv_fn = std::make_shared(); + auto result = conv_fn->Apply({input, weight}); + ASSERT_EQ(result.size(), 1); + EXPECT_EQ(result[0]->Dims(), (std::vector{1, 1, 3, 3})); + test::ExpectTensorFloatEqual(result[0], expected); +} + +TEST_P(AutogradConvForwardTest, ConvForwardMultiChannelBatchBias) { + ONLY_CPU(); + // input(n, c, h, w) = 10*n + 5*c + 3*h + w with H = W = 3. + std::vector input_values; + for (int n = 0; n < 2; ++n) { + for (int c = 0; c < 2; ++c) { + for (int h = 0; h < 3; ++h) { + for (int w = 0; w < 3; ++w) { input_values.push_back(10.0f * n + 5.0f * c + 3.0f * h + w); } + } + } + } + auto input = std::make_shared(input_values.data(), std::vector{2, 2, 3, 3}, DataType::kFLOAT32, + GetDevice()); + // out channel 0 sums both channel patches (all-ones kernel), out channel 1 mixes + // x(n, 0, i, j) + 2 * x(n, 1, i+1, j+1) to exercise per-tap / per-channel weighting. + const std::vector weight_values{1.0f, 1.0f, 1.0f, 1.0f, 1.0f, 1.0f, 1.0f, 1.0f, + 1.0f, 0.0f, 0.0f, 0.0f, 0.0f, 0.0f, 0.0f, 2.0f}; + auto weight = std::make_shared(weight_values.data(), std::vector{2, 2, 2, 2}, DataType::kFLOAT32, + GetDevice()); + const std::vector bias_values{0.5f, -0.25f}; + auto bias = std::make_shared(bias_values.data(), std::vector{2}, DataType::kFLOAT32, GetDevice()); + const std::vector expected{36.5f, 44.5f, 60.5f, 68.5f, 17.75f, 20.75f, 26.75f, 29.75f, + 116.5f, 124.5f, 140.5f, 148.5f, 47.75f, 50.75f, 56.75f, 59.75f}; + + auto conv_fn = std::make_shared(); + auto result = conv_fn->Apply({input, weight, bias}); + ASSERT_EQ(result.size(), 1); + EXPECT_EQ(result[0]->Dims(), (std::vector{2, 2, 2, 2})); + test::ExpectTensorFloatEqual(result[0], expected); +} + +TEST_P(AutogradConvForwardTest, ConvForwardKernelOne) { + ONLY_CPU(); + const std::vector input_values{1.0f, 2.0f, 3.0f, 4.0f, 5.0f, 6.0f, 7.0f, 8.0f, 9.0f, + 10.0f, 20.0f, 30.0f, 40.0f, 50.0f, 60.0f, 70.0f, 80.0f, 90.0f}; + auto input = std::make_shared(input_values.data(), std::vector{1, 2, 3, 3}, DataType::kFLOAT32, + GetDevice()); + // 1x1 kernels reduce conv to a per-position channel mix: out0 = 2*c0 + 3*c1, out1 = c0 - c1. + const std::vector weight_values{2.0f, 3.0f, 1.0f, -1.0f}; + auto weight = std::make_shared(weight_values.data(), std::vector{2, 2, 1, 1}, DataType::kFLOAT32, + GetDevice()); + const std::vector expected{32.0f, 64.0f, 96.0f, 128.0f, 160.0f, 192.0f, 224.0f, 256.0f, 288.0f, + -9.0f, -18.0f, -27.0f, -36.0f, -45.0f, -54.0f, -63.0f, -72.0f, -81.0f}; + + auto conv_fn = std::make_shared(); + auto result = conv_fn->Apply({input, weight}); + ASSERT_EQ(result.size(), 1); + EXPECT_EQ(result[0]->Dims(), (std::vector{1, 2, 3, 3})); + test::ExpectTensorFloatEqual(result[0], expected); +} + +// H = W = kernel collapses the output to 1x1: the only patch covers the whole image. +TEST_P(AutogradConvForwardTest, ConvForwardKernelEqualsSpatial) { + ONLY_CPU(); + std::vector input_values; + for (int idx = 0; idx < 9; ++idx) { input_values.push_back(static_cast(idx)); } + for (int idx = 0; idx < 9; ++idx) { input_values.push_back(2.0f * idx); } + auto input = std::make_shared(input_values.data(), std::vector{2, 1, 3, 3}, DataType::kFLOAT32, + GetDevice()); + auto weight = std::make_shared(std::vector{1, 1, 3, 3}, DataType::kFLOAT32, GetDevice(), true); + weight->Fill(1.0f); + auto bias = std::make_shared(std::vector{1}, DataType::kFLOAT32, GetDevice(), true); + bias->Fill(4.0f); + + auto conv_fn = std::make_shared(); + auto result = conv_fn->Apply({input, weight, bias}); + ASSERT_EQ(result.size(), 1); + EXPECT_EQ(result[0]->Dims(), (std::vector{2, 1, 1, 1})); + // batch 0: sum(0..8) + 4, batch 1: sum(0, 2, ..., 16) + 4 + test::ExpectTensorFloatEqual(result[0], std::vector{40.0f, 76.0f}); +} + +// Magnitudes near the fp32 integer-exact range: every product and partial sum stays exact, so +// this guards against spurious overflow / precision loss on the accumulation path. +TEST_P(AutogradConvForwardTest, ConvForwardExtremeValues) { + ONLY_CPU(); + auto input = std::make_shared(std::vector{1, 1, 4, 4}, DataType::kFLOAT32, GetDevice(), true); + input->Fill(static_cast(1 << 30)); + auto weight = std::make_shared(std::vector{1, 1, 3, 3}, DataType::kFLOAT32, GetDevice(), true); + weight->Fill(4.0f); + + auto conv_fn = std::make_shared(); + auto result = conv_fn->Apply({input, weight}); + ASSERT_EQ(result.size(), 1); + EXPECT_EQ(result[0]->Dims(), (std::vector{1, 1, 2, 2})); + // 9 taps * 2^30 * 4 = 9 * 2^32, exactly representable in fp32. + test::ExpectTensorFloatEqual(result[0], static_cast(9.0 * (1ULL << 32))); + + auto neg_input = std::make_shared(std::vector{1, 1, 4, 4}, DataType::kFLOAT32, GetDevice(), true); + neg_input->Fill(static_cast(-(1 << 30))); + auto neg_result = conv_fn->Apply({neg_input, weight}); + test::ExpectTensorFloatEqual(neg_result[0], static_cast(-9.0 * (1ULL << 32))); +} + +// Random 4D case checked against PyTorch conv2d (fixed seed, see file-level comment). +TEST_P(AutogradConvForwardTest, ConvForwardTorchGolden) { + ONLY_CPU(); + const std::vector input_values{ + 2.014464855f, -0.671311080f, -0.945236862f, -0.096128546f, 0.889558852f, 1.726294398f, -0.078931913f, + 0.205890238f, 0.094608001f, 0.173616216f, 0.437049866f, -0.557070732f, 0.455583423f, -0.736482263f, + -0.718647420f, 0.926798999f, 1.732108712f, -0.138735890f, -1.508729339f, 2.127460241f, 0.079739936f, + 0.217489868f, 0.589455068f, -0.052687794f, -1.659490585f, -0.380473137f, -0.166608080f, 0.350353301f, + -1.411424637f, 0.626776755f, 1.336345553f, -0.203511983f, -0.757243156f, 1.381434083f, -0.417290032f, + -0.668979943f, -0.569557011f, -1.035447836f, -0.555948615f, -1.279154539f, 0.056448594f, -0.388204783f, + 2.071142673f, -0.803012371f, 1.956272721f, 0.610922813f, -1.668144941f, -0.628582716f, -0.756579638f, + 0.987797141f, -0.849777400f, -1.654165864f, -0.326211363f, -0.040599853f, 0.281366080f, 0.499955416f, + -0.518889010f, -0.409607649f, 0.689743340f, 1.129151702f, 0.798748374f, -1.004696369f, -0.104189672f, + 1.257721663f, + }; + auto input = std::make_shared(input_values.data(), std::vector{2, 2, 4, 4}, DataType::kFLOAT32, + GetDevice()); + const std::vector weight_values{ + -2.448281050f, -0.608355999f, -0.958948731f, 1.009868145f, 0.031469870f, -1.101765275f, 0.571945250f, + -1.224039674f, -0.426795214f, -0.612229824f, -0.744571507f, -1.164491534f, -0.726353824f, -1.124407530f, + -0.498482078f, -0.423681349f, -1.231436968f, -0.689355612f, -0.170030892f, -0.120427191f, 1.032825112f, + 1.369398594f, 1.237134099f, 0.458183825f, -0.396747291f, 0.548917472f, 0.043999992f, 0.891736746f, + -1.292207360f, 1.713299513f, 0.982937992f, 0.432634443f, 0.293460459f, 0.318505853f, -0.146791920f, + 0.983720183f, -0.162205622f, 0.616937041f, 0.929942310f, -0.411319345f, 0.301661789f, -1.444623590f, + -0.730446517f, 2.391211271f, 0.400584310f, -0.909495890f, -0.702420592f, 0.372547686f, 0.994188488f, + 0.768918931f, 0.740486383f, -1.176032424f, 0.868664801f, 0.422459871f, + }; + auto weight = std::make_shared(weight_values.data(), std::vector{3, 2, 3, 3}, DataType::kFLOAT32, + GetDevice()); + const std::vector bias_values{-0.640517414f, 0.982376575f, 0.729396820f}; + auto bias = std::make_shared(bias_values.data(), std::vector{3}, DataType::kFLOAT32, GetDevice()); + const std::vector expected{ + -2.073040247f, 1.047590375f, -2.980665922f, -6.143202305f, 1.986891150f, 9.790361404f, + 0.198000669f, 0.174701050f, 0.093229778f, 2.889255762f, -0.725508392f, 0.107455865f, + 1.509838820f, -5.231677532f, 2.129449368f, 4.556076050f, -5.035542011f, -1.076976299f, + 0.719715834f, 1.997749925f, 1.977930188f, 7.832613468f, -5.848879814f, -0.770412445f, + }; + + auto conv_fn = std::make_shared(); + auto result = conv_fn->Apply({input, weight, bias}); + ASSERT_EQ(result.size(), 1); + EXPECT_EQ(result[0]->Dims(), (std::vector{2, 3, 2, 2})); + test::ExpectTensorNear(result[0], expected, kGoldenAbsError); +} + +TEST_P(AutogradConvForwardTest, ConvForwardRejectsInvalidShapes) { + ONLY_CPU(); + EXPECT_DEATH( + { + auto input = std::make_shared(std::vector{1, 2, 3, 3}, DataType::kFLOAT32, GetDevice()); + auto weight = std::make_shared(std::vector{1, 3, 2, 2}, DataType::kFLOAT32, GetDevice()); + auto conv_fn = std::make_shared(); + auto result = conv_fn->Apply({input, weight}); + (void)result; + }, + ""); + + EXPECT_DEATH( + { + auto input = std::make_shared(std::vector{1, 1, 2, 2}, DataType::kFLOAT32, GetDevice()); + auto weight = std::make_shared(std::vector{1, 1, 3, 3}, DataType::kFLOAT32, GetDevice()); + auto conv_fn = std::make_shared(); + auto result = conv_fn->Apply({input, weight}); + (void)result; + }, + ""); +} + +INFINI_TRAIN_REGISTER_TEST(AutogradConvForwardTest); diff --git a/tests/autograd/test_autograd_conv_train.cc b/tests/autograd/test_autograd_conv_train.cc new file mode 100644 index 00000000..613c0a78 --- /dev/null +++ b/tests/autograd/test_autograd_conv_train.cc @@ -0,0 +1,89 @@ +#include +#include +#include +#include + +#include "gtest/gtest.h" + +#include "infini_train/include/nn/modules/conv.h" +#include "infini_train/include/nn/modules/linear.h" +#include "infini_train/include/nn/modules/loss.h" +#include "infini_train/include/nn/parallel/global.h" +#include "infini_train/include/optimizer.h" +#include "infini_train/include/tensor.h" + +#include "tests/common/test_utils.h" + +using namespace infini_train; + +// Minimal end-to-end training step: Conv2d -> Flatten -> Linear -> CrossEntropy, one +// loss->Backward() through the autograd graph and one SGD step, on CPU. +class AutogradConvTrainTest : public infini_train::test::InfiniTrainTest {}; + +TEST_P(AutogradConvTrainTest, ConvTrainStepUpdatesParams) { + ONLY_CPU(); + const Device device = GetDevice(); + constexpr float kLearningRate = 0.1f; + + std::vector input_values; + for (int idx = 0; idx < 2 * 1 * 8 * 8; ++idx) { input_values.push_back((idx * 37) % 101 / 50.0f - 1.0f); } + auto input + = std::make_shared(input_values.data(), std::vector{2, 1, 8, 8}, DataType::kFLOAT32, device); + auto target = std::make_shared(std::vector{2}, DataType::kINT64, device); + target->Fill(0); + + auto conv = std::make_shared(1, 2, 3, true, device); + auto fc = std::make_shared(2 * 6 * 6, 3, true, device); + auto loss_fn = std::make_shared(); + + auto conv_out = (*conv)({input})[0]; + ASSERT_EQ(conv_out->Dims(), (std::vector{2, 2, 6, 6})); + auto flat = conv_out->Flatten(1); + ASSERT_EQ(flat->Dims(), (std::vector{2, 72})); + auto logits = (*fc)({flat})[0]; + ASSERT_EQ(logits->Dims(), (std::vector{2, 3})); + auto loss = (*loss_fn)({logits, target})[0]; + ASSERT_TRUE(loss->Dims().empty()); + EXPECT_TRUE(std::isfinite(*static_cast(loss->DataPtr()))); + + loss->Backward(); + + std::vector> params = conv->Parameters(); + for (auto ¶m : fc->Parameters()) { params.push_back(param); } + ASSERT_EQ(params.size(), 4); + + // Every parameter must have received a finite, non-trivial gradient through the graph. + for (const auto ¶m : params) { + const auto &grad = param->grad(); + ASSERT_NE(grad, nullptr); + ASSERT_EQ(grad->Dims(), param->Dims()); + const auto *grad_data = static_cast(grad->DataPtr()); + float max_abs = 0.0f; + for (size_t idx = 0; idx < grad->NumElements(); ++idx) { + ASSERT_TRUE(std::isfinite(grad_data[idx])); + max_abs = std::max(max_abs, std::fabs(grad_data[idx])); + } + EXPECT_GT(max_abs, 0.0f); + } + + std::vector> old_values; + for (const auto ¶m : params) { + auto old_value = std::make_shared(param->Dims(), param->Dtype(), device); + old_value->CopyFrom(*param); + old_values.push_back(old_value); + } + + auto optimizer = std::make_shared(params, kLearningRate); + optimizer->Step(); + + for (size_t p = 0; p < params.size(); ++p) { + const auto *old_data = static_cast(old_values[p]->DataPtr()); + const auto *grad_data = static_cast(params[p]->grad()->DataPtr()); + const auto *new_data = static_cast(params[p]->DataPtr()); + for (size_t idx = 0; idx < params[p]->NumElements(); ++idx) { + EXPECT_FLOAT_EQ(new_data[idx], old_data[idx] - kLearningRate * grad_data[idx]) << "param " << p; + } + } +} + +INFINI_TRAIN_REGISTER_TEST(AutogradConvTrainTest); From 51a34901312d051c1f0d0ce2170497a65d8917db Mon Sep 17 00:00:00 2001 From: accelerator-llc Date: Sat, 5 Sep 2026 15:49:18 +0800 Subject: [PATCH 03/15] feat(conv): add CUDA Conv2d kernels with autograd and unit tests --- infini_train/src/kernels/cuda/conv.cu | 410 ++++++++++++++++++ .../test_autograd_conv_cuda_backward.cc | 247 +++++++++++ .../test_autograd_conv_cuda_forward.cc | 189 ++++++++ .../autograd/test_autograd_conv_cuda_train.cc | 89 ++++ 4 files changed, 935 insertions(+) create mode 100644 infini_train/src/kernels/cuda/conv.cu create mode 100644 tests/autograd/test_autograd_conv_cuda_backward.cc create mode 100644 tests/autograd/test_autograd_conv_cuda_forward.cc create mode 100644 tests/autograd/test_autograd_conv_cuda_train.cc diff --git a/infini_train/src/kernels/cuda/conv.cu b/infini_train/src/kernels/cuda/conv.cu new file mode 100644 index 00000000..10499847 --- /dev/null +++ b/infini_train/src/kernels/cuda/conv.cu @@ -0,0 +1,410 @@ +#include +#include +#include + +#include + +#include "glog/logging.h" + +#include "infini_train/include/common/cuda/common_cuda.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_runtime_common.h" +#include "infini_train/src/kernels/common/gemm.h" +#include "infini_train/src/kernels/cuda/common/gemm.cuh" + +namespace infini_train::kernels::cuda { +namespace { + +// im2col for stride=1 / padding=0 cross-correlation. Each thread writes one +// element of the (flat_kernel, patches) column-major scratch buffer, so a single +// image's buffer is addressable as col[kk * patches + p]. This column-major +// order is what the cuBLAS GEMM geometry below expects. The (c, u, v) tap index +// is the same row-major flattening as an (O, C, kH, kW) weight. +__global__ void Im2colKernel(const float *__restrict__ input, float *__restrict__ col, int64_t total, int64_t channels, + int64_t height, int64_t width, int64_t kernel_h, int64_t kernel_w, int64_t out_width, + int64_t patches, int64_t flat_kernel) { + for (int64_t idx = blockIdx.x * blockDim.x + threadIdx.x; idx < total; idx += blockDim.x * gridDim.x) { + const int64_t p = idx % patches; + const int64_t kk = (idx / patches) % flat_kernel; + const int64_t n = idx / (flat_kernel * patches); + + const int64_t i = p / out_width; + const int64_t j = p % out_width; + const int64_t c = kk / (kernel_h * kernel_w); + const int64_t rem = kk % (kernel_h * kernel_w); + const int64_t u = rem / kernel_w; + const int64_t v = rem % kernel_w; + + const int64_t in_idx = n * channels * height * width + c * height * width + (i + u) * width + (j + v); + col[idx] = input[in_idx]; + } +} + +// col2im: inverse of Im2col for the input gradient. An input pixel is covered by +// multiple output windows, so all patches that touch it are gathered. Each thread +// owns one (c, x, y) output cell and writes it once, so no atomics are needed. +__global__ void Col2imKernel(const float *__restrict__ col, float *__restrict__ grad_input, int64_t total, + int64_t channels, int64_t height, int64_t width, int64_t kernel_h, int64_t kernel_w, + int64_t out_height, int64_t out_width, int64_t patches, int64_t flat_kernel) { + for (int64_t idx = blockIdx.x * blockDim.x + threadIdx.x; idx < total; idx += blockDim.x * gridDim.x) { + const int64_t y = idx % width; + const int64_t x = (idx / width) % height; + const int64_t c = (idx / (width * height)) % channels; + const int64_t n = idx / (width * height * channels); + + float sum = 0.0f; + for (int64_t u = 0; u < kernel_h; ++u) { + const int64_t i = x - u; + if (i < 0 || i >= out_height) { + continue; + } + for (int64_t v = 0; v < kernel_w; ++v) { + const int64_t j = y - v; + if (j < 0 || j >= out_width) { + continue; + } + const int64_t kk = c * kernel_h * kernel_w + u * kernel_w + v; + const int64_t col_idx = n * (flat_kernel * patches) + kk * patches + i * out_width + j; + sum += col[col_idx]; + } + } + grad_input[idx] = sum; + } +} + +// Broadcast bias over every spatial position of the (N, O, H, W) output, so a +// subsequent GEMM with beta=1 accumulates the convolution onto it. +__global__ void BroadcastBiasKernel(float *__restrict__ output, const float *__restrict__ bias, int64_t total, + int64_t patches, int64_t out_channels) { + for (int64_t idx = blockIdx.x * blockDim.x + threadIdx.x; idx < total; idx += blockDim.x * gridDim.x) { + const int64_t o = (idx / patches) % out_channels; + output[idx] = bias[o]; + } +} + +// grad_bias(o) = sum_{n, i, j} grad_output(n, o, i, j). A parallel block reduction keeps the +// accumulation error at the same order as the cuBLAS reference rather than a single-thread sum. +template +__global__ void GradBiasKernel(const float *__restrict__ grad_output, float *__restrict__ grad_bias, int64_t batch, + int64_t patches, int64_t out_channels) { + using BlockReduce = cub::BlockReduce; + __shared__ typename BlockReduce::TempStorage temp_storage; + + const int64_t o = blockIdx.x; + if (o >= out_channels) { + return; + } + + const int64_t total = batch * patches; + float sum = 0.0f; + for (int64_t idx = threadIdx.x; idx < total; idx += blockDim.x) { + const int64_t n = idx / patches; + const int64_t p = idx % patches; + sum += grad_output[n * out_channels * patches + o * patches + p]; + } + + const float reduced = BlockReduce(temp_storage).Sum(sum); + if (threadIdx.x == 0) { + grad_bias[o] = reduced; + } +} + +constexpr int kThreads = 256; +constexpr int kMaxBlocks = 1024; + +int64_t NumBlocks(int64_t total) { + int64_t blocks = (total + kThreads - 1) / kThreads; + return blocks > kMaxBlocks ? kMaxBlocks : blocks; +} + +cudaStream_t CurrentCudaStream(const Device &device) { + return dynamic_cast( + infini_train::core::GetDeviceGuardImpl(device.type())->GetStream(device)) + ->cuda_stream(); +} + +} // namespace + +std::shared_ptr Conv2dForward(const std::shared_ptr &input, const std::shared_ptr &weight, + const std::shared_ptr &bias) { + const auto &input_dims = input->Dims(); + CHECK_EQ(input_dims.size(), 4); + const int64_t batch = input_dims[0]; + const int64_t channels = input_dims[1]; + const int64_t height = input_dims[2]; + const int64_t width = input_dims[3]; + + const auto &weight_dims = weight->Dims(); + CHECK_EQ(weight_dims.size(), 4); + const int64_t out_channels = weight_dims[0]; + const int64_t kernel_h = weight_dims[2]; + const int64_t kernel_w = weight_dims[3]; + CHECK_EQ(weight_dims[1], channels); + CHECK_GE(height, kernel_h); + CHECK_GE(width, kernel_w); + + CHECK(input->Dtype() == DataType::kFLOAT32) << "Conv2d requires FP32 tensors"; + CHECK(weight->Dtype() == DataType::kFLOAT32) << "Conv2d requires FP32 tensors"; + CHECK_EQ(input->GetDevice(), weight->GetDevice()); + + if (bias) { + CHECK_EQ(bias->Dims().size(), 1); + CHECK_EQ(bias->Dims()[0], out_channels); + CHECK(bias->Dtype() == DataType::kFLOAT32) << "Conv2d requires FP32 tensors"; + CHECK_EQ(input->GetDevice(), bias->GetDevice()); + } + + const int64_t out_height = height - kernel_h + 1; + const int64_t out_width = width - kernel_w + 1; + const int64_t patches = out_height * out_width; + const int64_t flat_kernel = channels * kernel_h * kernel_w; + + auto output = std::make_shared(std::vector{batch, out_channels, out_height, out_width}, + DataType::kFLOAT32, input->GetDevice()); + if (batch == 0) { + return output; + } + + auto device = input->GetDevice(); + const cudaStream_t stream = CurrentCudaStream(device); + + // Scratch im2col buffer, indexed column-major (patches, flat_kernel) per image. + auto col = std::make_shared(std::vector{batch, flat_kernel, patches}, DataType::kFLOAT32, device); + + Im2colKernel<<>>( + static_cast(input->DataPtr()), static_cast(col->DataPtr()), + batch * patches * flat_kernel, channels, height, width, kernel_h, kernel_w, out_width, patches, flat_kernel); + CUDA_CHECK(cudaGetLastError()); + + if (bias) { + // Prefill output with the bias, then the GEMM below accumulates onto it. + BroadcastBiasKernel<<>>( + static_cast(output->DataPtr()), static_cast(bias->DataPtr()), + batch * out_channels * patches, patches, out_channels); + CUDA_CHECK(cudaGetLastError()); + } + + const float *col_data = static_cast(col->DataPtr()); + const float *weight_data = static_cast(weight->DataPtr()); + float *output_data = static_cast(output->DataPtr()); + const float beta = bias ? 1.0f : 0.0f; + for (int64_t n = 0; n < batch; ++n) { + Dispatcher::Instance().Call({device.type(), "Gemm"}, device, + GemmParams{ + .trans_a = GemmTranspose::kNoTranspose, + .trans_b = GemmTranspose::kNoTranspose, + .m = static_cast(patches), + .n = static_cast(out_channels), + .k = static_cast(flat_kernel), + .A = col_data + n * patches * flat_kernel, + .lda = static_cast(patches), + .B = weight_data, + .ldb = static_cast(flat_kernel), + .C = output_data + n * out_channels * patches, + .ldc = static_cast(patches), + .alpha = 1.0f, + .beta = beta, + .batch_count = 1, + .input_dtype = DataType::kFLOAT32, + .output_dtype = DataType::kFLOAT32, + }); + } + + return output; +} + +std::shared_ptr Conv2dBackwardInput(const std::shared_ptr &weight, + const std::shared_ptr &grad_output, + const std::vector &input_dims) { + CHECK_EQ(input_dims.size(), 4); + const int64_t batch = input_dims[0]; + const int64_t channels = input_dims[1]; + const int64_t height = input_dims[2]; + const int64_t width = input_dims[3]; + + const auto &weight_dims = weight->Dims(); + CHECK_EQ(weight_dims.size(), 4); + const int64_t out_channels = weight_dims[0]; + const int64_t kernel_h = weight_dims[2]; + const int64_t kernel_w = weight_dims[3]; + CHECK_EQ(weight_dims[1], channels); + CHECK_GE(height, kernel_h); + CHECK_GE(width, kernel_w); + + CHECK(weight->Dtype() == DataType::kFLOAT32) << "Conv2d requires FP32 tensors"; + CHECK(grad_output->Dtype() == DataType::kFLOAT32) << "Conv2d requires FP32 tensors"; + CHECK(weight->GetDevice() == grad_output->GetDevice()); + + const int64_t out_height = height - kernel_h + 1; + const int64_t out_width = width - kernel_w + 1; + const int64_t patches = out_height * out_width; + const int64_t flat_kernel = channels * kernel_h * kernel_w; + + const auto &grad_dims = grad_output->Dims(); + CHECK_EQ(grad_dims.size(), 4); + CHECK_EQ(grad_dims[0], batch); + CHECK_EQ(grad_dims[1], out_channels); + CHECK_EQ(grad_dims[2], out_height); + CHECK_EQ(grad_dims[3], out_width); + + auto grad_input = std::make_shared(input_dims, DataType::kFLOAT32, grad_output->GetDevice()); + if (batch == 0) { + return grad_input; + } + + auto device = grad_output->GetDevice(); + const cudaStream_t stream = CurrentCudaStream(device); + + // Per-image scratch: grad_output(n)^T * weight, stored column-major (patches, flat_kernel). + auto col = std::make_shared(std::vector{flat_kernel, patches}, DataType::kFLOAT32, device); + + const float *weight_data = static_cast(weight->DataPtr()); + const float *grad_output_data = static_cast(grad_output->DataPtr()); + float *grad_input_data = static_cast(grad_input->DataPtr()); + float *col_data = static_cast(col->DataPtr()); + + for (int64_t n = 0; n < batch; ++n) { + // col = grad_output(n)^T * weight, the transpose of the forward GEMM. + Dispatcher::Instance().Call({device.type(), "Gemm"}, device, + GemmParams{ + .trans_a = GemmTranspose::kNoTranspose, + .trans_b = GemmTranspose::kTranspose, + .m = static_cast(patches), + .n = static_cast(flat_kernel), + .k = static_cast(out_channels), + .A = grad_output_data + n * out_channels * patches, + .lda = static_cast(patches), + .B = weight_data, + .ldb = static_cast(flat_kernel), + .C = col_data, + .ldc = static_cast(patches), + .alpha = 1.0f, + .beta = 0.0f, + .batch_count = 1, + .input_dtype = DataType::kFLOAT32, + .output_dtype = DataType::kFLOAT32, + }); + + Col2imKernel<<>>( + col_data, grad_input_data + n * channels * height * width, channels * height * width, channels, height, + width, kernel_h, kernel_w, out_height, out_width, patches, flat_kernel); + CUDA_CHECK(cudaGetLastError()); + } + + return grad_input; +} + +std::shared_ptr Conv2dBackwardWeight(const std::shared_ptr &input, + const std::shared_ptr &grad_output) { + const auto &input_dims = input->Dims(); + CHECK_EQ(input_dims.size(), 4); + const int64_t batch = input_dims[0]; + const int64_t channels = input_dims[1]; + const int64_t height = input_dims[2]; + const int64_t width = input_dims[3]; + + const auto &grad_dims = grad_output->Dims(); + CHECK_EQ(grad_dims.size(), 4); + CHECK_EQ(grad_dims[0], batch); + + CHECK(input->Dtype() == DataType::kFLOAT32) << "Conv2d requires FP32 tensors"; + CHECK(grad_output->Dtype() == DataType::kFLOAT32) << "Conv2d requires FP32 tensors"; + CHECK(input->GetDevice() == grad_output->GetDevice()); + + const int64_t kernel_h = height - grad_dims[2] + 1; + const int64_t kernel_w = width - grad_dims[3] + 1; + CHECK_GT(kernel_h, 0); + CHECK_GT(kernel_w, 0); + + const int64_t out_channels = grad_dims[1]; + const int64_t out_height = grad_dims[2]; + const int64_t out_width = grad_dims[3]; + const int64_t patches = out_height * out_width; + const int64_t flat_kernel = channels * kernel_h * kernel_w; + + auto grad_weight = std::make_shared(std::vector{out_channels, channels, kernel_h, kernel_w}, + DataType::kFLOAT32, input->GetDevice()); + grad_weight->Fill(0.0f); + if (batch == 0) { + return grad_weight; + } + + auto device = input->GetDevice(); + const cudaStream_t stream = CurrentCudaStream(device); + + auto col = std::make_shared(std::vector{batch, flat_kernel, patches}, DataType::kFLOAT32, device); + + Im2colKernel<<>>( + static_cast(input->DataPtr()), static_cast(col->DataPtr()), + batch * patches * flat_kernel, channels, height, width, kernel_h, kernel_w, out_width, patches, flat_kernel); + CUDA_CHECK(cudaGetLastError()); + + const float *col_data = static_cast(col->DataPtr()); + const float *grad_output_data = static_cast(grad_output->DataPtr()); + float *grad_weight_data = static_cast(grad_weight->DataPtr()); + + // grad_weight(o, kk) = sum_n grad_output(n, o, p) * col(n, kk, p). Computed as + // the transpose of the forward GEMM and accumulated over the batch. + for (int64_t n = 0; n < batch; ++n) { + Dispatcher::Instance().Call({device.type(), "Gemm"}, device, + GemmParams{ + .trans_a = GemmTranspose::kTranspose, + .trans_b = GemmTranspose::kNoTranspose, + .m = static_cast(flat_kernel), + .n = static_cast(out_channels), + .k = static_cast(patches), + .A = col_data + n * patches * flat_kernel, + .lda = static_cast(patches), + .B = grad_output_data + n * out_channels * patches, + .ldb = static_cast(patches), + .C = grad_weight_data, + .ldc = static_cast(flat_kernel), + .alpha = 1.0f, + .beta = 1.0f, + .batch_count = 1, + .input_dtype = DataType::kFLOAT32, + .output_dtype = DataType::kFLOAT32, + }); + } + + return grad_weight; +} + +std::shared_ptr Conv2dBackwardBias(const std::shared_ptr &grad_output, int64_t out_channels) { + const auto &grad_dims = grad_output->Dims(); + CHECK_EQ(grad_dims.size(), 4); + CHECK_EQ(grad_dims[1], out_channels); + CHECK(grad_output->Dtype() == DataType::kFLOAT32) << "Conv2d requires FP32 tensors"; + + const int64_t batch = grad_dims[0]; + const int64_t patches = grad_dims[2] * grad_dims[3]; + + auto grad_bias + = std::make_shared(std::vector{out_channels}, DataType::kFLOAT32, grad_output->GetDevice()); + + auto device = grad_output->GetDevice(); + const cudaStream_t stream = CurrentCudaStream(device); + + constexpr int kGradBiasBlock = 256; + GradBiasKernel<<>>( + static_cast(grad_output->DataPtr()), static_cast(grad_bias->DataPtr()), batch, patches, + out_channels); + CUDA_CHECK(cudaGetLastError()); + + return grad_bias; +} + +} // namespace infini_train::kernels::cuda + +#define REGISTER_CUDA_CONV2D_KERNEL(kernel_name) \ + REGISTER_KERNEL(infini_train::Device::DeviceType::kCUDA, kernel_name, infini_train::kernels::cuda::kernel_name) + +REGISTER_CUDA_CONV2D_KERNEL(Conv2dForward) +REGISTER_CUDA_CONV2D_KERNEL(Conv2dBackwardInput) +REGISTER_CUDA_CONV2D_KERNEL(Conv2dBackwardWeight) +REGISTER_CUDA_CONV2D_KERNEL(Conv2dBackwardBias) + +#undef REGISTER_CUDA_CONV2D_KERNEL diff --git a/tests/autograd/test_autograd_conv_cuda_backward.cc b/tests/autograd/test_autograd_conv_cuda_backward.cc new file mode 100644 index 00000000..b7b050ea --- /dev/null +++ b/tests/autograd/test_autograd_conv_cuda_backward.cc @@ -0,0 +1,247 @@ +#include +#include +#include + +#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; + +namespace { +// Same golden case as the forward test (torch 2.14, seed 20260905, autograd.grad reference). +constexpr float kGoldenAbsError = 2e-6f; +} // namespace + +class AutogradConvCudaBackwardTest : public infini_train::test::InfiniTrainTest {}; + +// Hand-checkable case: x = arange(9) reshaped (1,1,3,3), w = [[1,2],[3,4]], bias = 0.5, +// grad_output = ones(1,1,2,2). +TEST_P(AutogradConvCudaBackwardTest, ConvBackwardGradients) { + ONLY_CUDA(); + std::vector input_values; + for (int idx = 0; idx < 9; ++idx) { input_values.push_back(static_cast(idx)); } + auto input = std::make_shared(input_values.data(), std::vector{1, 1, 3, 3}, DataType::kFLOAT32, + GetDevice()) + ->RequiresGrad(); + const std::vector weight_values{1.0f, 2.0f, 3.0f, 4.0f}; + auto weight = std::make_shared(weight_values.data(), std::vector{1, 1, 2, 2}, DataType::kFLOAT32, + GetDevice()) + ->RequiresGrad(); + auto bias = std::make_shared(std::vector{1}, DataType::kFLOAT32, GetDevice(), true); + bias->Fill(0.5f); + + auto conv_fn = std::make_shared(); + auto result = conv_fn->Apply({input, weight, bias}); + auto grad_output + = std::make_shared(std::vector{1, 1, 2, 2}, DataType::kFLOAT32, GetDevice(), true); + grad_output->Fill(1.0f); + auto grad_inputs = conv_fn->Backward({grad_output}); + ASSERT_EQ(grad_inputs.size(), 3); + ASSERT_NE(grad_inputs[0], nullptr); + ASSERT_NE(grad_inputs[1], nullptr); + ASSERT_NE(grad_inputs[2], nullptr); + + const std::vector expected_grad_input{1.0f, 3.0f, 2.0f, 4.0f, 10.0f, 6.0f, 3.0f, 7.0f, 4.0f}; + EXPECT_EQ(grad_inputs[0]->Dims(), (std::vector{1, 1, 3, 3})); + test::ExpectTensorFloatEqual(grad_inputs[0], expected_grad_input); + + const std::vector expected_grad_weight{8.0f, 12.0f, 20.0f, 24.0f}; + EXPECT_EQ(grad_inputs[1]->Dims(), (std::vector{1, 1, 2, 2})); + test::ExpectTensorFloatEqual(grad_inputs[1], expected_grad_weight); + + test::ExpectTensorFloatEqual(grad_inputs[2], std::vector{4.0f}); +} + +// No-bias path returns {grad_input, grad_weight} and both must be correct. +TEST_P(AutogradConvCudaBackwardTest, ConvBackwardNoBias) { + ONLY_CUDA(); + auto input = std::make_shared(std::vector{2, 2, 4, 4}, DataType::kFLOAT32, GetDevice(), true); + input->Fill(1.0f); + auto weight = std::make_shared(std::vector{3, 2, 3, 3}, DataType::kFLOAT32, GetDevice(), true); + weight->Fill(1.0f); + auto conv_fn = std::make_shared(); + auto result = conv_fn->Apply({input, weight}); + auto grad = std::make_shared(std::vector{2, 3, 2, 2}, DataType::kFLOAT32, GetDevice(), true); + grad->Fill(1.0f); + auto grad_inputs = conv_fn->Backward({grad}); + EXPECT_EQ(grad_inputs.size(), 2); + + ASSERT_NE(grad_inputs[0], nullptr); + EXPECT_EQ(grad_inputs[0]->Dims(), (std::vector{2, 2, 4, 4})); + ASSERT_NE(grad_inputs[1], nullptr); + EXPECT_EQ(grad_inputs[1]->Dims(), (std::vector{3, 2, 3, 3})); + // With all-ones input/weight/grad_output, grad_weight(o,c,u,v) = sum_{n,i,j} 1 = 2*2*2. + test::ExpectTensorFloatEqual(grad_inputs[1], 8.0f); +} + +// H = W = kernel: a single window, so grad_input mirrors the weight and grad_weight the input. +TEST_P(AutogradConvCudaBackwardTest, ConvBackwardKernelEqualsSpatial) { + ONLY_CUDA(); + const std::vector input_values{1.0f, 2.0f, 3.0f, 4.0f}; + auto input = std::make_shared(input_values.data(), std::vector{1, 1, 2, 2}, DataType::kFLOAT32, + GetDevice()) + ->RequiresGrad(); + const std::vector weight_values{5.0f, 6.0f, 7.0f, 8.0f}; + auto weight = std::make_shared(weight_values.data(), std::vector{1, 1, 2, 2}, DataType::kFLOAT32, + GetDevice()) + ->RequiresGrad(); + auto bias = std::make_shared(std::vector{1}, DataType::kFLOAT32, GetDevice(), true); + bias->Fill(0.25f); + + auto conv_fn = std::make_shared(); + auto result = conv_fn->Apply({input, weight, bias}); + auto grad_output + = std::make_shared(std::vector{1, 1, 1, 1}, DataType::kFLOAT32, GetDevice(), true); + grad_output->Fill(1.0f); + auto grad_inputs = conv_fn->Backward({grad_output}); + ASSERT_EQ(grad_inputs.size(), 3); + + EXPECT_EQ(grad_inputs[0]->Dims(), (std::vector{1, 1, 2, 2})); + test::ExpectTensorFloatEqual(grad_inputs[0], weight_values); + EXPECT_EQ(grad_inputs[1]->Dims(), (std::vector{1, 1, 2, 2})); + test::ExpectTensorFloatEqual(grad_inputs[1], input_values); + test::ExpectTensorFloatEqual(grad_inputs[2], std::vector{1.0f}); +} + +// Training-shaped call: the activation is a leaf without requires_grad, so only the parameter +// gradients must be produced. +TEST_P(AutogradConvCudaBackwardTest, ConvBackwardInputNotRequired) { + ONLY_CUDA(); + auto input = std::make_shared(std::vector{2, 2, 4, 4}, DataType::kFLOAT32, GetDevice()); + input->Fill(1.0f); + auto weight = std::make_shared(std::vector{3, 2, 3, 3}, DataType::kFLOAT32, GetDevice(), true); + weight->Fill(1.0f); + auto bias = std::make_shared(std::vector{3}, DataType::kFLOAT32, GetDevice(), true); + bias->Fill(1.0f); + + auto conv_fn = std::make_shared(); + auto result = conv_fn->Apply({input, weight, bias}); + auto grad = std::make_shared(std::vector{2, 3, 2, 2}, DataType::kFLOAT32, GetDevice(), true); + grad->Fill(1.0f); + auto grad_inputs = conv_fn->Backward({grad}); + ASSERT_EQ(grad_inputs.size(), 3); + EXPECT_EQ(grad_inputs[0], nullptr); + EXPECT_NE(grad_inputs[1], nullptr); + EXPECT_NE(grad_inputs[2], nullptr); + test::ExpectTensorFloatEqual(grad_inputs[1], 8.0f); + test::ExpectTensorFloatEqual(grad_inputs[2], 8.0f); +} + +// Random 4D case checked against PyTorch autograd (fixed seed, see file-level comment). +TEST_P(AutogradConvCudaBackwardTest, ConvBackwardTorchGolden) { + ONLY_CUDA(); + const std::vector input_values{ + 2.014464855f, -0.671311080f, -0.945236862f, -0.096128546f, 0.889558852f, 1.726294398f, -0.078931913f, + 0.205890238f, 0.094608001f, 0.173616216f, 0.437049866f, -0.557070732f, 0.455583423f, -0.736482263f, + -0.718647420f, 0.926798999f, 1.732108712f, -0.138735890f, -1.508729339f, 2.127460241f, 0.079739936f, + 0.217489868f, 0.589455068f, -0.052687794f, -1.659490585f, -0.380473137f, -0.166608080f, 0.350353301f, + -1.411424637f, 0.626776755f, 1.336345553f, -0.203511983f, -0.757243156f, 1.381434083f, -0.417290032f, + -0.668979943f, -0.569557011f, -1.035447836f, -0.555948615f, -1.279154539f, 0.056448594f, -0.388204783f, + 2.071142673f, -0.803012371f, 1.956272721f, 0.610922813f, -1.668144941f, -0.628582716f, -0.756579638f, + 0.987797141f, -0.849777400f, -1.654165864f, -0.326211363f, -0.040599853f, 0.281366080f, 0.499955416f, + -0.518889010f, -0.409607649f, 0.689743340f, 1.129151702f, 0.798748374f, -1.004696369f, -0.104189672f, + 1.257721663f, + }; + auto input = std::make_shared(input_values.data(), std::vector{2, 2, 4, 4}, DataType::kFLOAT32, + GetDevice()) + ->RequiresGrad(); + const std::vector weight_values{ + -2.448281050f, -0.608355999f, -0.958948731f, 1.009868145f, 0.031469870f, -1.101765275f, 0.571945250f, + -1.224039674f, -0.426795214f, -0.612229824f, -0.744571507f, -1.164491534f, -0.726353824f, -1.124407530f, + -0.498482078f, -0.423681349f, -1.231436968f, -0.689355612f, -0.170030892f, -0.120427191f, 1.032825112f, + 1.369398594f, 1.237134099f, 0.458183825f, -0.396747291f, 0.548917472f, 0.043999992f, 0.891736746f, + -1.292207360f, 1.713299513f, 0.982937992f, 0.432634443f, 0.293460459f, 0.318505853f, -0.146791920f, + 0.983720183f, -0.162205622f, 0.616937041f, 0.929942310f, -0.411319345f, 0.301661789f, -1.444623590f, + -0.730446517f, 2.391211271f, 0.400584310f, -0.909495890f, -0.702420592f, 0.372547686f, 0.994188488f, + 0.768918931f, 0.740486383f, -1.176032424f, 0.868664801f, 0.422459871f, + }; + auto weight = std::make_shared(weight_values.data(), std::vector{3, 2, 3, 3}, DataType::kFLOAT32, + GetDevice()) + ->RequiresGrad(); + const std::vector bias_values{-0.640517414f, 0.982376575f, 0.729396820f}; + auto bias = std::make_shared(bias_values.data(), std::vector{3}, DataType::kFLOAT32, GetDevice()) + ->RequiresGrad(); + + auto conv_fn = std::make_shared(); + auto result = conv_fn->Apply({input, weight, bias}); + const std::vector grad_output_values{ + -1.277466655f, -1.339194655f, -1.159473658f, 0.086730979f, 0.438649237f, 0.602657557f, + -0.738640666f, -0.410562515f, 0.201109841f, 0.716394842f, 0.003985748f, 0.676588118f, + 0.412507862f, -0.146394536f, -0.880041420f, -1.247093678f, -1.212666154f, -0.393535852f, + 0.564200640f, 0.279093444f, -2.798053741f, -1.615590930f, -0.518746257f, -0.283052236f, + }; + auto grad_output = std::make_shared(grad_output_values.data(), std::vector{2, 3, 2, 2}, + DataType::kFLOAT32, GetDevice()); + + auto grad_inputs = conv_fn->Backward({grad_output}); + ASSERT_EQ(grad_inputs.size(), 3); + ASSERT_NE(grad_inputs[0], nullptr); + ASSERT_NE(grad_inputs[1], nullptr); + ASSERT_NE(grad_inputs[2], nullptr); + + const std::vector expected_grad_input{ + 3.020392179f, 3.908452511f, 3.049194813f, 2.572864771f, 2.191555023f, 0.285838872f, 3.004242659f, + 0.838663280f, -3.235622168f, -0.945002913f, 4.960352898f, -0.376030058f, -0.373013139f, 0.741603315f, + 1.750292301f, 0.215949476f, 0.990354240f, 0.948824525f, 2.029216766f, 2.858904839f, 1.606565118f, + 4.838562012f, 3.308694601f, 0.822548687f, 0.564552128f, 2.793590069f, 4.189336777f, 2.155962467f, + 0.251298606f, 0.576505065f, 0.615549624f, -0.177835792f, -0.349884391f, -1.143750548f, -5.110340118f, + -1.768475175f, 2.049649715f, 0.845555902f, 3.547870159f, 3.535832882f, 2.858142614f, -6.855783939f, + -3.052171230f, 1.308768988f, -0.348264217f, -0.470771074f, 1.195474505f, 0.431147456f, 1.210889816f, + 4.433355331f, -1.848075271f, -1.105654240f, -2.759691477f, -3.466245174f, -1.316413283f, 0.586114824f, + 3.407652855f, 0.809629261f, -2.299243927f, -0.474773228f, 1.162620664f, 1.500420690f, 2.191398382f, + 1.014662623f, + }; + const std::vector expected_grad_weight{ + -1.278161168f, 2.350493193f, 3.455903053f, -3.191775799f, -4.849991322f, -1.593288898f, -3.348839760f, + 1.063873172f, 4.325201988f, -2.219501734f, 2.213379383f, -2.589243412f, 2.336852074f, -1.198173761f, + -2.431171656f, 4.716279507f, 0.842501462f, -3.181046009f, -1.122435212f, -4.357190609f, -0.400157630f, + 2.311078310f, 2.235447884f, 2.117345333f, 1.470544577f, 0.713402152f, -3.305890560f, 0.861350715f, + -2.180587769f, 2.186314106f, 1.552511692f, 0.699985623f, 0.372371852f, 0.788659871f, -1.649568439f, + -1.754016399f, 1.571256876f, -3.355353594f, 2.778887510f, 4.880561829f, 3.997585535f, 2.531520605f, + -1.071573257f, -2.245790720f, -3.141553879f, 1.098210573f, -2.158659220f, 5.950089931f, 1.271268964f, + 0.028057545f, -1.955230832f, 1.795806646f, 1.293214321f, -3.971021652f, + }; + const std::vector expected_grad_bias{-5.550425529f, -0.870804369f, -3.617365122f}; + + EXPECT_EQ(grad_inputs[0]->Dims(), (std::vector{2, 2, 4, 4})); + test::ExpectTensorNear(grad_inputs[0], expected_grad_input, kGoldenAbsError); + EXPECT_EQ(grad_inputs[1]->Dims(), (std::vector{3, 2, 3, 3})); + test::ExpectTensorNear(grad_inputs[1], expected_grad_weight, kGoldenAbsError); + EXPECT_EQ(grad_inputs[2]->Dims(), (std::vector{3})); + test::ExpectTensorNear(grad_inputs[2], expected_grad_bias, kGoldenAbsError); +} + +// An empty batch is degenerate but must not crash nor return uninitialized gradients: the +// parameter gradients (weight/bias) are sums over zero images, so they are exact zeros. +TEST_P(AutogradConvCudaBackwardTest, ConvBackwardEmptyBatch) { + ONLY_CUDA(); + auto input + = std::make_shared(std::vector{0, 2, 4, 4}, DataType::kFLOAT32, GetDevice())->RequiresGrad(); + auto weight = std::make_shared(std::vector{3, 2, 3, 3}, DataType::kFLOAT32, GetDevice(), true) + ->RequiresGrad(); + auto bias + = std::make_shared(std::vector{3}, DataType::kFLOAT32, GetDevice(), true)->RequiresGrad(); + + auto conv_fn = std::make_shared(); + auto result = conv_fn->Apply({input, weight, bias}); + auto grad_output = std::make_shared(std::vector{0, 3, 2, 2}, DataType::kFLOAT32, GetDevice()); + auto grads = conv_fn->Backward({grad_output}); + ASSERT_EQ(grads.size(), 3); + ASSERT_NE(grads[0], nullptr); + ASSERT_NE(grads[1], nullptr); + ASSERT_NE(grads[2], nullptr); + + // grad_input has an empty batch, grad_weight and grad_bias are exact zeros over the empty batch. + EXPECT_EQ(grads[0]->Dims(), (std::vector{0, 2, 4, 4})); + EXPECT_EQ(grads[0]->NumElements(), 0); + EXPECT_EQ(grads[1]->Dims(), (std::vector{3, 2, 3, 3})); + test::ExpectTensorFloatEqual(grads[1], 0.0f); + EXPECT_EQ(grads[2]->Dims(), (std::vector{3})); + test::ExpectTensorFloatEqual(grads[2], 0.0f); +} + +INFINI_TRAIN_REGISTER_TEST(AutogradConvCudaBackwardTest); diff --git a/tests/autograd/test_autograd_conv_cuda_forward.cc b/tests/autograd/test_autograd_conv_cuda_forward.cc new file mode 100644 index 00000000..40ecacdc --- /dev/null +++ b/tests/autograd/test_autograd_conv_cuda_forward.cc @@ -0,0 +1,189 @@ +#include +#include +#include + +#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; + +namespace { +// Deterministic golden data, generated with torch 2.14 under seed 20260905 +// (conv2d / autograd.grad reference). +constexpr float kGoldenAbsError = 2e-6f; +} // namespace + +class AutogradConvCudaForwardTest : public infini_train::test::InfiniTrainTest {}; + +// An asymmetric kernel pins the cross-correlation semantics: flipping the kernel would swap +// the roles of the taps and change every output value. The values are small exact integers, +// so the CUDA path reproduces them bit-exactly. +TEST_P(AutogradConvCudaForwardTest, ConvForwardAsymmetricKernel) { + ONLY_CUDA(); + std::vector input_values; + for (int idx = 0; idx < 16; ++idx) { input_values.push_back(static_cast(idx)); } + auto input = std::make_shared(input_values.data(), std::vector{1, 1, 4, 4}, DataType::kFLOAT32, + GetDevice()); + const std::vector weight_values{1.0f, 2.0f, 3.0f, 4.0f}; + auto weight = std::make_shared(weight_values.data(), std::vector{1, 1, 2, 2}, DataType::kFLOAT32, + GetDevice()); + // out[i][j] = 1*x[i][j] + 2*x[i][j+1] + 3*x[i+1][j] + 4*x[i+1][j+1] + const std::vector expected{34.0f, 44.0f, 54.0f, 74.0f, 84.0f, 94.0f, 114.0f, 124.0f, 134.0f}; + + auto conv_fn = std::make_shared(); + auto result = conv_fn->Apply({input, weight}); + ASSERT_EQ(result.size(), 1); + EXPECT_EQ(result[0]->Dims(), (std::vector{1, 1, 3, 3})); + test::ExpectTensorFloatEqual(result[0], expected); +} + +// Cross-channel / batch weighting with a per-channel bias. +TEST_P(AutogradConvCudaForwardTest, ConvForwardMultiChannelBatchBias) { + ONLY_CUDA(); + std::vector input_values; + for (int n = 0; n < 2; ++n) { + for (int c = 0; c < 2; ++c) { + for (int h = 0; h < 3; ++h) { + for (int w = 0; w < 3; ++w) { input_values.push_back(10.0f * n + 5.0f * c + 3.0f * h + w); } + } + } + } + auto input = std::make_shared(input_values.data(), std::vector{2, 2, 3, 3}, DataType::kFLOAT32, + GetDevice()); + const std::vector weight_values{1.0f, 1.0f, 1.0f, 1.0f, 1.0f, 1.0f, 1.0f, 1.0f, + 1.0f, 0.0f, 0.0f, 0.0f, 0.0f, 0.0f, 0.0f, 2.0f}; + auto weight = std::make_shared(weight_values.data(), std::vector{2, 2, 2, 2}, DataType::kFLOAT32, + GetDevice()); + const std::vector bias_values{0.5f, -0.25f}; + auto bias = std::make_shared(bias_values.data(), std::vector{2}, DataType::kFLOAT32, GetDevice()); + const std::vector expected{36.5f, 44.5f, 60.5f, 68.5f, 17.75f, 20.75f, 26.75f, 29.75f, + 116.5f, 124.5f, 140.5f, 148.5f, 47.75f, 50.75f, 56.75f, 59.75f}; + + auto conv_fn = std::make_shared(); + auto result = conv_fn->Apply({input, weight, bias}); + ASSERT_EQ(result.size(), 1); + EXPECT_EQ(result[0]->Dims(), (std::vector{2, 2, 2, 2})); + test::ExpectTensorFloatEqual(result[0], expected); +} + +// 1x1 kernels reduce conv to a per-position channel mix. +TEST_P(AutogradConvCudaForwardTest, ConvForwardKernelOne) { + ONLY_CUDA(); + const std::vector input_values{1.0f, 2.0f, 3.0f, 4.0f, 5.0f, 6.0f, 7.0f, 8.0f, 9.0f, + 10.0f, 20.0f, 30.0f, 40.0f, 50.0f, 60.0f, 70.0f, 80.0f, 90.0f}; + auto input = std::make_shared(input_values.data(), std::vector{1, 2, 3, 3}, DataType::kFLOAT32, + GetDevice()); + const std::vector weight_values{2.0f, 3.0f, 1.0f, -1.0f}; + auto weight = std::make_shared(weight_values.data(), std::vector{2, 2, 1, 1}, DataType::kFLOAT32, + GetDevice()); + const std::vector expected{32.0f, 64.0f, 96.0f, 128.0f, 160.0f, 192.0f, 224.0f, 256.0f, 288.0f, + -9.0f, -18.0f, -27.0f, -36.0f, -45.0f, -54.0f, -63.0f, -72.0f, -81.0f}; + + auto conv_fn = std::make_shared(); + auto result = conv_fn->Apply({input, weight}); + ASSERT_EQ(result.size(), 1); + EXPECT_EQ(result[0]->Dims(), (std::vector{1, 2, 3, 3})); + test::ExpectTensorFloatEqual(result[0], expected); +} + +// H = W = kernel collapses the output to 1x1. +TEST_P(AutogradConvCudaForwardTest, ConvForwardKernelEqualsSpatial) { + ONLY_CUDA(); + std::vector input_values; + for (int idx = 0; idx < 9; ++idx) { input_values.push_back(static_cast(idx)); } + for (int idx = 0; idx < 9; ++idx) { input_values.push_back(2.0f * idx); } + auto input = std::make_shared(input_values.data(), std::vector{2, 1, 3, 3}, DataType::kFLOAT32, + GetDevice()); + auto weight = std::make_shared(std::vector{1, 1, 3, 3}, DataType::kFLOAT32, GetDevice(), true); + weight->Fill(1.0f); + auto bias = std::make_shared(std::vector{1}, DataType::kFLOAT32, GetDevice(), true); + bias->Fill(4.0f); + + auto conv_fn = std::make_shared(); + auto result = conv_fn->Apply({input, weight, bias}); + ASSERT_EQ(result.size(), 1); + EXPECT_EQ(result[0]->Dims(), (std::vector{2, 1, 1, 1})); + test::ExpectTensorFloatEqual(result[0], std::vector{40.0f, 76.0f}); +} + +// Magnitudes near the fp32 integer-exact range guard against overflow on accumulation. +TEST_P(AutogradConvCudaForwardTest, ConvForwardExtremeValues) { + ONLY_CUDA(); + auto input = std::make_shared(std::vector{1, 1, 4, 4}, DataType::kFLOAT32, GetDevice(), true); + input->Fill(static_cast(1 << 30)); + auto weight = std::make_shared(std::vector{1, 1, 3, 3}, DataType::kFLOAT32, GetDevice(), true); + weight->Fill(4.0f); + + auto conv_fn = std::make_shared(); + auto result = conv_fn->Apply({input, weight}); + ASSERT_EQ(result.size(), 1); + EXPECT_EQ(result[0]->Dims(), (std::vector{1, 1, 2, 2})); + // 9 taps * 2^30 * 4 = 9 * 2^32, exactly representable in fp32. + test::ExpectTensorFloatEqual(result[0], static_cast(9.0 * (1ULL << 32))); +} + +// Random 4D case checked against PyTorch conv2d (fixed seed, see file-level comment). +TEST_P(AutogradConvCudaForwardTest, ConvForwardTorchGolden) { + ONLY_CUDA(); + const std::vector input_values{ + 2.014464855f, -0.671311080f, -0.945236862f, -0.096128546f, 0.889558852f, 1.726294398f, -0.078931913f, + 0.205890238f, 0.094608001f, 0.173616216f, 0.437049866f, -0.557070732f, 0.455583423f, -0.736482263f, + -0.718647420f, 0.926798999f, 1.732108712f, -0.138735890f, -1.508729339f, 2.127460241f, 0.079739936f, + 0.217489868f, 0.589455068f, -0.052687794f, -1.659490585f, -0.380473137f, -0.166608080f, 0.350353301f, + -1.411424637f, 0.626776755f, 1.336345553f, -0.203511983f, -0.757243156f, 1.381434083f, -0.417290032f, + -0.668979943f, -0.569557011f, -1.035447836f, -0.555948615f, -1.279154539f, 0.056448594f, -0.388204783f, + 2.071142673f, -0.803012371f, 1.956272721f, 0.610922813f, -1.668144941f, -0.628582716f, -0.756579638f, + 0.987797141f, -0.849777400f, -1.654165864f, -0.326211363f, -0.040599853f, 0.281366080f, 0.499955416f, + -0.518889010f, -0.409607649f, 0.689743340f, 1.129151702f, 0.798748374f, -1.004696369f, -0.104189672f, + 1.257721663f, + }; + auto input = std::make_shared(input_values.data(), std::vector{2, 2, 4, 4}, DataType::kFLOAT32, + GetDevice()); + const std::vector weight_values{ + -2.448281050f, -0.608355999f, -0.958948731f, 1.009868145f, 0.031469870f, -1.101765275f, 0.571945250f, + -1.224039674f, -0.426795214f, -0.612229824f, -0.744571507f, -1.164491534f, -0.726353824f, -1.124407530f, + -0.498482078f, -0.423681349f, -1.231436968f, -0.689355612f, -0.170030892f, -0.120427191f, 1.032825112f, + 1.369398594f, 1.237134099f, 0.458183825f, -0.396747291f, 0.548917472f, 0.043999992f, 0.891736746f, + -1.292207360f, 1.713299513f, 0.982937992f, 0.432634443f, 0.293460459f, 0.318505853f, -0.146791920f, + 0.983720183f, -0.162205622f, 0.616937041f, 0.929942310f, -0.411319345f, 0.301661789f, -1.444623590f, + -0.730446517f, 2.391211271f, 0.400584310f, -0.909495890f, -0.702420592f, 0.372547686f, 0.994188488f, + 0.768918931f, 0.740486383f, -1.176032424f, 0.868664801f, 0.422459871f, + }; + auto weight = std::make_shared(weight_values.data(), std::vector{3, 2, 3, 3}, DataType::kFLOAT32, + GetDevice()); + const std::vector bias_values{-0.640517414f, 0.982376575f, 0.729396820f}; + auto bias = std::make_shared(bias_values.data(), std::vector{3}, DataType::kFLOAT32, GetDevice()); + const std::vector expected{ + -2.073040247f, 1.047590375f, -2.980665922f, -6.143202305f, 1.986891150f, 9.790361404f, + 0.198000669f, 0.174701050f, 0.093229778f, 2.889255762f, -0.725508392f, 0.107455865f, + 1.509838820f, -5.231677532f, 2.129449368f, 4.556076050f, -5.035542011f, -1.076976299f, + 0.719715834f, 1.997749925f, 1.977930188f, 7.832613468f, -5.848879814f, -0.770412445f, + }; + + auto conv_fn = std::make_shared(); + auto result = conv_fn->Apply({input, weight, bias}); + ASSERT_EQ(result.size(), 1); + EXPECT_EQ(result[0]->Dims(), (std::vector{2, 3, 2, 2})); + test::ExpectTensorNear(result[0], expected, kGoldenAbsError); +} + +// An empty batch is a degenerate but valid call: the kernel must not launch a zero-block +// configuration; the output tensor keeps the (N=0, ...) shape and is returned untouched. +TEST_P(AutogradConvCudaForwardTest, ConvForwardEmptyBatch) { + ONLY_CUDA(); + auto input = std::make_shared(std::vector{0, 2, 4, 4}, DataType::kFLOAT32, GetDevice()); + auto weight = std::make_shared(std::vector{3, 2, 3, 3}, DataType::kFLOAT32, GetDevice(), true); + weight->Fill(1.0f); + + auto conv_fn = std::make_shared(); + auto result = conv_fn->Apply({input, weight}); + ASSERT_EQ(result.size(), 1); + EXPECT_EQ(result[0]->Dims(), (std::vector{0, 3, 2, 2})); + EXPECT_EQ(result[0]->NumElements(), 0); +} + +INFINI_TRAIN_REGISTER_TEST(AutogradConvCudaForwardTest); diff --git a/tests/autograd/test_autograd_conv_cuda_train.cc b/tests/autograd/test_autograd_conv_cuda_train.cc new file mode 100644 index 00000000..39c0777b --- /dev/null +++ b/tests/autograd/test_autograd_conv_cuda_train.cc @@ -0,0 +1,89 @@ +#include +#include +#include +#include + +#include "gtest/gtest.h" + +#include "infini_train/include/nn/modules/conv.h" +#include "infini_train/include/nn/modules/linear.h" +#include "infini_train/include/nn/modules/loss.h" +#include "infini_train/include/optimizer.h" +#include "infini_train/include/tensor.h" + +#include "tests/common/test_utils.h" + +using namespace infini_train; + +// Minimal end-to-end training step on CUDA: Conv2d -> Flatten -> Linear -> CrossEntropy, +// one loss->Backward() through the autograd graph and one SGD step. +class AutogradConvCudaTrainTest : public infini_train::test::InfiniTrainTest {}; + +TEST_P(AutogradConvCudaTrainTest, ConvTrainStepUpdatesParams) { + ONLY_CUDA(); + const Device device = GetDevice(); + const Device host = Device(); + constexpr float kLearningRate = 0.1f; + + std::vector input_values; + for (int idx = 0; idx < 2 * 1 * 8 * 8; ++idx) { input_values.push_back((idx * 37) % 101 / 50.0f - 1.0f); } + auto input + = std::make_shared(input_values.data(), std::vector{2, 1, 8, 8}, DataType::kFLOAT32, device); + auto target = std::make_shared(std::vector{2}, DataType::kINT64, device); + target->Fill(0); + + auto conv = std::make_shared(1, 2, 3, true, device); + auto fc = std::make_shared(2 * 6 * 6, 3, true, device); + auto loss_fn = std::make_shared(); + + auto conv_out = (*conv)({input})[0]; + ASSERT_EQ(conv_out->Dims(), (std::vector{2, 2, 6, 6})); + auto flat = conv_out->Flatten(1); + ASSERT_EQ(flat->Dims(), (std::vector{2, 72})); + auto logits = (*fc)({flat})[0]; + ASSERT_EQ(logits->Dims(), (std::vector{2, 3})); + auto loss = (*loss_fn)({logits, target})[0]; + ASSERT_TRUE(loss->Dims().empty()); + const auto loss_cpu = loss->To(host); + EXPECT_TRUE(std::isfinite(*static_cast(loss_cpu.DataPtr()))); + + loss->Backward(); + + std::vector> params = conv->Parameters(); + for (auto ¶m : fc->Parameters()) { params.push_back(param); } + ASSERT_EQ(params.size(), 4); + + // Every parameter must have received a finite, non-trivial gradient through the graph. + for (const auto ¶m : params) { + const auto &grad = param->grad(); + ASSERT_NE(grad, nullptr); + ASSERT_EQ(grad->Dims(), param->Dims()); + const auto grad_cpu = grad->To(host); + const auto *grad_data = static_cast(grad_cpu.DataPtr()); + float max_abs = 0.0f; + for (size_t idx = 0; idx < grad_cpu.NumElements(); ++idx) { + ASSERT_TRUE(std::isfinite(grad_data[idx])); + max_abs = std::max(max_abs, std::fabs(grad_data[idx])); + } + EXPECT_GT(max_abs, 0.0f); + } + + std::vector> old_values_cpu; + for (const auto ¶m : params) { old_values_cpu.push_back(std::make_shared(param->To(host))); } + + auto optimizer = std::make_shared(params, kLearningRate); + optimizer->Step(); + + for (size_t p = 0; p < params.size(); ++p) { + const auto new_cpu = params[p]->To(host); + const auto grad_cpu = params[p]->grad()->To(host); + const auto *old_data = static_cast(old_values_cpu[p]->DataPtr()); + const auto *grad_data = static_cast(grad_cpu.DataPtr()); + const auto *new_data = static_cast(new_cpu.DataPtr()); + for (size_t idx = 0; idx < params[p]->NumElements(); ++idx) { + EXPECT_FLOAT_EQ(new_data[idx], old_data[idx] - kLearningRate * grad_data[idx]) << "param " << p; + } + } +} + +INFINI_TRAIN_REGISTER_TEST(AutogradConvCudaTrainTest); From 57e95fe5bc54d60912dee922a778fba98ca7159c Mon Sep 17 00:00:00 2001 From: accelerator-llc Date: Sat, 5 Sep 2026 15:49:18 +0800 Subject: [PATCH 04/15] fix(conv): guard Conv2d CPU kernels against non-FP32 inputs --- infini_train/src/kernels/cpu/conv.cc | 8 ++++++++ 1 file changed, 8 insertions(+) diff --git a/infini_train/src/kernels/cpu/conv.cc b/infini_train/src/kernels/cpu/conv.cc index 268590d4..bd5c8312 100644 --- a/infini_train/src/kernels/cpu/conv.cc +++ b/infini_train/src/kernels/cpu/conv.cc @@ -80,10 +80,13 @@ std::shared_ptr Conv2dForward(const std::shared_ptr &input, cons CHECK_EQ(weight_dims[1], channels); CHECK_GE(height, kernel_h); CHECK_GE(width, kernel_w); + CHECK(input->Dtype() == DataType::kFLOAT32) << "Conv2d requires FP32 tensors"; + CHECK(weight->Dtype() == DataType::kFLOAT32) << "Conv2d requires FP32 tensors"; if (bias) { CHECK_EQ(bias->Dims().size(), 1); CHECK_EQ(bias->Dims()[0], out_channels); + CHECK(bias->Dtype() == DataType::kFLOAT32) << "Conv2d requires FP32 tensors"; } const int64_t out_height = height - kernel_h + 1; @@ -143,6 +146,8 @@ std::shared_ptr Conv2dBackwardInput(const std::shared_ptr &weigh CHECK_EQ(weight_dims[1], channels); CHECK_GE(height, kernel_h); CHECK_GE(width, kernel_w); + CHECK(weight->Dtype() == DataType::kFLOAT32) << "Conv2d requires FP32 tensors"; + CHECK(grad_output->Dtype() == DataType::kFLOAT32) << "Conv2d requires FP32 tensors"; const int64_t out_height = height - kernel_h + 1; const int64_t out_width = width - kernel_w + 1; @@ -195,6 +200,8 @@ std::shared_ptr Conv2dBackwardWeight(const std::shared_ptr &inpu const auto &grad_dims = grad_output->Dims(); CHECK_EQ(grad_dims.size(), 4); CHECK_EQ(grad_dims[0], batch); + CHECK(input->Dtype() == DataType::kFLOAT32) << "Conv2d requires FP32 tensors"; + CHECK(grad_output->Dtype() == DataType::kFLOAT32) << "Conv2d requires FP32 tensors"; // The kernel extent is the difference between the input and output spatial extents. const int64_t kernel_h = height - grad_dims[2] + 1; @@ -240,6 +247,7 @@ std::shared_ptr Conv2dBackwardBias(const std::shared_ptr &grad_o const auto &grad_dims = grad_output->Dims(); CHECK_EQ(grad_dims.size(), 4); CHECK_EQ(grad_dims[1], out_channels); + CHECK(grad_output->Dtype() == DataType::kFLOAT32) << "Conv2d requires FP32 tensors"; auto grad_bias = std::make_shared(std::vector{out_channels}, DataType::kFLOAT32, grad_output->GetDevice()); From 4885d9dd941278549c49da3fc5a40e8f00202182 Mon Sep 17 00:00:00 2001 From: accelerator-llc Date: Sat, 5 Sep 2026 22:13:02 +0800 Subject: [PATCH 05/15] feat(relu): add ReLU activation with autograd for CPU and CUDA Add autograd::ReLU function (forward with clamp_min semantics, backward with threshold_backward semantics) and nn::ReLU module, following the existing activation pattern. CPU and CUDA elementwise kernels preserve the exact NaN and negative-zero behavior of PyTorch, verified bit-exact against PyTorch (fixed seed). Includes forward, backward, end-to-end training chain and Flatten gradient tests. --- infini_train/include/autograd/activations.h | 12 ++ infini_train/include/nn/modules/activations.h | 7 + infini_train/src/autograd/activations.cc | 25 ++++ infini_train/src/kernels/cpu/relu.cc | 47 +++++++ infini_train/src/kernels/cuda/relu.cu | 116 ++++++++++++++++ infini_train/src/nn/modules/activations.cc | 4 + tests/autograd/test_autograd_flatten_grad.cc | 65 +++++++++ tests/autograd/test_autograd_relu_backward.cc | 111 ++++++++++++++++ tests/autograd/test_autograd_relu_forward.cc | 124 ++++++++++++++++++ tests/autograd/test_autograd_relu_train.cc | 120 +++++++++++++++++ 10 files changed, 631 insertions(+) create mode 100644 infini_train/src/kernels/cpu/relu.cc create mode 100644 infini_train/src/kernels/cuda/relu.cu create mode 100644 tests/autograd/test_autograd_flatten_grad.cc create mode 100644 tests/autograd/test_autograd_relu_backward.cc create mode 100644 tests/autograd/test_autograd_relu_forward.cc create mode 100644 tests/autograd/test_autograd_relu_train.cc diff --git a/infini_train/include/autograd/activations.h b/infini_train/include/autograd/activations.h index a6397726..809a7b67 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> &output_tensors) override; std::vector> Backward(const std::vector> &grad_outputs) override; }; + +class ReLU : public Function { +public: + static constexpr char kType[] = "ReLUFunction"; + + ReLU() : Function(kType) {} + + std::vector> Forward(const std::vector> &input_tensors) override; + void SetupContext(const std::vector> &input_tensors, + const std::vector> &output_tensors) override; + std::vector> Backward(const std::vector> &grad_outputs) override; +}; } // namespace infini_train::autograd diff --git a/infini_train/include/nn/modules/activations.h b/infini_train/include/nn/modules/activations.h index deb02957..4b69ef18 100644 --- a/infini_train/include/nn/modules/activations.h +++ b/infini_train/include/nn/modules/activations.h @@ -17,6 +17,13 @@ class Sigmoid : public CloneableModule { std::vector> Forward(const std::vector> &input_tensors) override; }; +class ReLU : public CloneableModule { +public: + static constexpr char kType[] = "ReLU"; + ReLU() : CloneableModule(kType) {} + std::vector> Forward(const std::vector> &input_tensors) override; +}; + class NewGELU : public CloneableModule { public: static constexpr char kType[] = "NewGELU"; diff --git a/infini_train/src/autograd/activations.cc b/infini_train/src/autograd/activations.cc index bb8b8e5e..4b2ae37a 100644 --- a/infini_train/src/autograd/activations.cc +++ b/infini_train/src/autograd/activations.cc @@ -30,4 +30,29 @@ std::vector> Sigmoid::Backward(const std::vectorGetDevice().type(); return {Dispatcher::Instance().Call>({device, "SigmoidBackward"}, output, grad_output)}; } + +std::vector> ReLU::Forward(const std::vector> &input_tensors) { + CHECK_EQ(input_tensors.size(), 1); + const auto &input = input_tensors[0]; + + auto device = input->GetDevice().type(); + return {Dispatcher::Instance().Call>({device, "ReLUForward"}, input)}; +} + +void ReLU::SetupContext(const std::vector> &input_tensors, + const std::vector> &) { + // ReLU backward needs the input activations to decide where the gradient flows. + ctx_.SaveForBackward({input_tensors[0]}); +} + +std::vector> ReLU::Backward(const std::vector> &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>({device, "ReLUBackward"}, input, grad_output)}; +} } // namespace infini_train::autograd diff --git a/infini_train/src/kernels/cpu/relu.cc b/infini_train/src/kernels/cpu/relu.cc new file mode 100644 index 00000000..7d4086a0 --- /dev/null +++ b/infini_train/src/kernels/cpu/relu.cc @@ -0,0 +1,47 @@ +#include + +#include "glog/logging.h" + +#include "infini_train/include/dispatcher.h" +#include "infini_train/include/tensor.h" + +namespace infini_train::kernels::cpu { +std::shared_ptr ReLUForward(const std::shared_ptr &input) { + CHECK(input->Dtype() == DataType::kFLOAT32) << "ReLU requires FP32 tensors"; + + auto output = std::make_shared(input->Dims(), DataType::kFLOAT32); + const float *input_ptr = static_cast(input->DataPtr()); + float *output_ptr = static_cast(output->DataPtr()); + + const int64_t numel = input->NumElements(); + // Strict `<` keeps NaN and -0 (both compare false), matching clamp_min(x, 0). + for (int64_t idx = 0; idx < numel; ++idx) { output_ptr[idx] = input_ptr[idx] < 0.0f ? 0.0f : input_ptr[idx]; } + + return output; +} + +std::shared_ptr ReLUBackward(const std::shared_ptr &input, const std::shared_ptr &grad_output) { + CHECK(input->Dtype() == DataType::kFLOAT32) << "ReLU requires FP32 tensors"; + CHECK(grad_output->Dtype() == DataType::kFLOAT32) << "ReLU requires FP32 tensors"; + + auto grad_input = std::make_shared(input->Dims(), DataType::kFLOAT32); + const float *input_ptr = static_cast(input->DataPtr()); + const float *grad_output_ptr = static_cast(grad_output->DataPtr()); + float *grad_input_ptr = static_cast(grad_input->DataPtr()); + + const int64_t numel = input->NumElements(); + // `<=` matches threshold_backward: NaN (false) passes grad through, 0/-0/negative mask to 0. + for (int64_t idx = 0; idx < numel; ++idx) { + grad_input_ptr[idx] = input_ptr[idx] <= 0.0f ? 0.0f : grad_output_ptr[idx]; + } + 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/relu.cu b/infini_train/src/kernels/cuda/relu.cu new file mode 100644 index 00000000..7471edbf --- /dev/null +++ b/infini_train/src/kernels/cuda/relu.cu @@ -0,0 +1,116 @@ +#include +#include + +#include "infini_train/include/common/common.h" +#include "infini_train/include/core/runtime/device_guard.h" +#include "infini_train/include/dispatcher.h" +#include "infini_train/include/dtype_dispatch.h" +#include "infini_train/include/tensor.h" + +#include "infini_train/src/core/runtime/cuda/cuda_runtime_common.h" + +namespace infini_train::kernels::cuda { +namespace { + +template __global__ void ReLUForwardKernel(T *output, const T *input, size_t num_elements, size_t offset) { + size_t idx = blockIdx.x * blockDim.x + threadIdx.x + offset; + if (idx < num_elements) { + // Strict `<` keeps NaN and -0 (both compare false), matching clamp_min(x, 0). + output[idx] = input[idx] < T(0) ? T(0) : input[idx]; + } +} + +template +__global__ void ReLUBackwardKernel(T *grad_input, const T *input, const T *grad_output, size_t num_elements, + size_t offset) { + size_t idx = blockIdx.x * blockDim.x + threadIdx.x + offset; + if (idx < num_elements) { + // `<=` matches threshold_backward: NaN (false) passes grad through, 0/-0/negative mask to 0. + grad_input[idx] = input[idx] <= T(0) ? T(0) : grad_output[idx]; + } +} + +// Mirrors the unary launch pattern in kernels/cuda/elementwise.cu: same block-size +// heuristic, stream lookup and offset-loop emission. That file's helpers live in an +// anonymous namespace and cannot be shared from here; keep this in sync with it. +inline size_t ChooseBlockSize(size_t num_elements) { + if (num_elements < 1024) { + return 64; + } + if (num_elements < 65536) { + return 128; + } + if (num_elements < 1048576) { + return 256; + } + return 512; +} + +cudaStream_t CurrentStream(const std::shared_ptr &tensor) { + auto device = tensor->GetDevice(); + return dynamic_cast( + infini_train::core::GetDeviceGuardImpl(device.type())->GetStream(device)) + ->cuda_stream(); +} + +template +void ReLUForwardImpl(const std::shared_ptr &input, const std::shared_ptr &output) { + const size_t num_elements = output->NumElements(); + if (num_elements == 0) { + return; + } + cudaStream_t stream = CurrentStream(output); + T *out_ptr = static_cast(output->DataPtr()); + const T *in_ptr = static_cast(input->DataPtr()); + + dim3 block(ChooseBlockSize(num_elements)); + dim3 grid(CEIL_DIV(num_elements, block.x)); + const size_t step = grid.x * block.x; + for (size_t offset = 0; offset < num_elements; offset += step) { + ReLUForwardKernel<<>>(out_ptr, in_ptr, num_elements, offset); + } +} + +template +void ReLUBackwardImpl(const std::shared_ptr &input, const std::shared_ptr &grad_output, + const std::shared_ptr &grad_input) { + const size_t num_elements = grad_input->NumElements(); + if (num_elements == 0) { + return; + } + cudaStream_t stream = CurrentStream(grad_input); + T *grad_input_ptr = static_cast(grad_input->DataPtr()); + const T *input_ptr = static_cast(input->DataPtr()); + const T *grad_output_ptr = static_cast(grad_output->DataPtr()); + + dim3 block(ChooseBlockSize(num_elements)); + dim3 grid(CEIL_DIV(num_elements, block.x)); + const size_t step = grid.x * block.x; + for (size_t offset = 0; offset < num_elements; offset += step) { + ReLUBackwardKernel + <<>>(grad_input_ptr, input_ptr, grad_output_ptr, num_elements, offset); + } +} + +} // namespace + +std::shared_ptr ReLUForward(const std::shared_ptr &input) { + auto output = std::make_shared(input->Dims(), input->Dtype(), input->GetDevice()); + DISPATCH(input->Dtype(), ReLUForwardImpl(input, output);, DataType::kFLOAT32) + return output; +} + +std::shared_ptr ReLUBackward(const std::shared_ptr &input, const std::shared_ptr &grad_output) { + auto grad_input = std::make_shared(input->Dims(), input->Dtype(), input->GetDevice()); + DISPATCH(input->Dtype(), ReLUBackwardImpl(input, grad_output, grad_input);, DataType::kFLOAT32) + 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 d1bbc9da..a5373847 100644 --- a/infini_train/src/nn/modules/activations.cc +++ b/infini_train/src/nn/modules/activations.cc @@ -12,6 +12,10 @@ std::vector> Sigmoid::Forward(const std::vector()->Apply(input_tensors); } +std::vector> ReLU::Forward(const std::vector> &input_tensors) { + return std::make_shared()->Apply(input_tensors); +} + std::vector> NewGELU::Forward(const std::vector> &x) { auto &input = x[0]; return {0.5 * input diff --git a/tests/autograd/test_autograd_flatten_grad.cc b/tests/autograd/test_autograd_flatten_grad.cc new file mode 100644 index 00000000..d9917dcb --- /dev/null +++ b/tests/autograd/test_autograd_flatten_grad.cc @@ -0,0 +1,65 @@ +#include +#include +#include + +#include "gtest/gtest.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; + +// Gradient contract of Tensor::Flatten: it is a pure reshape (via the no-op view), +// so d(loss)/d(x) must be the reshape-back of the downstream gradient. +// The composition x -> flatten -> x*x -> sum makes the expected grad = 2*x. +class AutogradFlattenGradTest : public infini_train::test::InfiniTrainTest {}; + +namespace { +void CheckFlattenGrad(const Device &device, const std::vector &input_dims, int64_t start, int64_t end, + const std::vector &expected_flat_dims) { + const Device host = Device(); + + std::vector values; + for (size_t idx = 0; idx < static_cast( + std::accumulate(input_dims.begin(), input_dims.end(), 1, std::multiplies())); + ++idx) { + values.push_back(((idx % 17) - 8) / 4.0f); + } + auto x = std::make_shared(input_dims, DataType::kFLOAT32, device, true); + auto storage = std::make_shared(values.data(), input_dims, DataType::kFLOAT32, device); + x->CopyFrom(*storage); + + auto flat = x->Flatten(start, end); + ASSERT_EQ(flat->Dims(), expected_flat_dims); + + auto sq = flat->Mul(flat); + auto reduced = sq; + while (reduced->Dims().size() > 0) { reduced = reduced->Sum(0); } + reduced->Backward(); + + const auto &grad = x->grad(); + ASSERT_NE(grad, nullptr); + ASSERT_EQ(grad->Dims(), input_dims); + const auto grad_cpu = grad->To(host); + const float *grad_data = static_cast(grad_cpu.DataPtr()); + for (size_t idx = 0; idx < values.size(); ++idx) { + EXPECT_FLOAT_EQ(grad_data[idx], 2.0f * values[idx]) << "gradient mismatch at flat index " << idx; + } +} +} // namespace + +// 4-D conv-output layout: (N, C, H, W) flattening the channel dim with start=1, +// the same call pattern as the MNIST demo. +TEST_P(AutogradFlattenGradTest, FlattenGradFourDConvLayout) { + CheckFlattenGrad(GetDevice(), {2, 1, 4, 4}, 1, -1, {2, 16}); +} + +TEST_P(AutogradFlattenGradTest, FlattenGradThreeDStartEnd) { CheckFlattenGrad(GetDevice(), {2, 3, 4}, 0, 1, {6, 4}); } + +TEST_P(AutogradFlattenGradTest, FlattenGradTwoDStartEnd) { CheckFlattenGrad(GetDevice(), {3, 4}, 0, 1, {12}); } + +TEST_P(AutogradFlattenGradTest, FlattenGradOneD) { CheckFlattenGrad(GetDevice(), {5}, 0, 0, {5}); } + +INFINI_TRAIN_REGISTER_TEST(AutogradFlattenGradTest); diff --git a/tests/autograd/test_autograd_relu_backward.cc b/tests/autograd/test_autograd_relu_backward.cc new file mode 100644 index 00000000..e76c7b78 --- /dev/null +++ b/tests/autograd/test_autograd_relu_backward.cc @@ -0,0 +1,111 @@ +#include +#include +#include +#include + +#include "gtest/gtest.h" + +#include "infini_train/include/autograd/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; + +namespace { +uint32_t FloatBits(float value) { + uint32_t bits; + std::memcpy(&bits, &value, sizeof(bits)); + return bits; +} + +// grad respect to input: NaN passes grad through, 0/-0/negative mask to 0, positive keeps grad. +float ExpectedReluGrad(float x, float grad) { return x <= 0.0f ? 0.0f : grad; } +} // namespace + +class AutogradReLUBackwardTest : public infini_train::test::InfiniTrainTest {}; + +TEST_P(AutogradReLUBackwardTest, ReLUBackwardNaNTransparentAndZeroMask) { + const std::vector input_values{std::nanf(""), + -0.0f, + 0.0f, + 1.5f, + -2.5f, + std::numeric_limits::infinity(), + -std::numeric_limits::infinity(), + 1e-30f, + -1e-30f, + std::numeric_limits::max(), + -std::numeric_limits::max()}; + const std::vector grad_values{1.0f, 2.0f, 3.0f, 4.0f, 5.0f, 6.0f, 7.0f, 8.0f, 9.0f, 10.0f, 11.0f}; + + auto input + = std::make_shared(input_values.data(), std::vector{static_cast(input_values.size())}, + DataType::kFLOAT32, GetDevice()); + auto grad + = std::make_shared(grad_values.data(), std::vector{static_cast(grad_values.size())}, + DataType::kFLOAT32, GetDevice()); + + auto relu_fn = std::make_shared(); + auto result = relu_fn->Apply({input}); + ASSERT_EQ(result.size(), 1); + auto grad_inputs = relu_fn->Backward({grad}); + ASSERT_EQ(grad_inputs.size(), 1); + ASSERT_EQ(grad_inputs[0]->Dims(), input->Dims()); + + const auto host_grad_input = grad_inputs[0]->To(Device()); + const float *out = static_cast(host_grad_input.DataPtr()); + + for (size_t idx = 0; idx < input_values.size(); ++idx) { + const float expected = ExpectedReluGrad(input_values[idx], grad_values[idx]); + if (std::isnan(input_values[idx])) { + // The gradient must be passed through bit-for-bit, not zeroed. + EXPECT_EQ(FloatBits(out[idx]), FloatBits(grad_values[idx])) << "NaN gradient must pass through at " << idx; + } else { + EXPECT_EQ(FloatBits(out[idx]), FloatBits(expected)) << "gradient bit mismatch at position " << idx; + } + } +} + +// Two-dimensional case with a mixed-sign input, checking mask and passthrough jointly. +TEST_P(AutogradReLUBackwardTest, ReLUBackwardTwoD) { + const std::vector input_values{-1.0f, 2.0f, -0.0f, 4.0f, -5.0f, 6.0f}; + const std::vector grad_values{1.0f, 2.0f, 3.0f, 4.0f, 5.0f, 6.0f}; + auto input + = std::make_shared(input_values.data(), std::vector{2, 3}, DataType::kFLOAT32, GetDevice()); + auto grad + = std::make_shared(grad_values.data(), std::vector{2, 3}, DataType::kFLOAT32, GetDevice()); + + auto relu_fn = std::make_shared(); + auto result = relu_fn->Apply({input}); + ASSERT_EQ(result.size(), 1); + auto grad_inputs = relu_fn->Backward({grad}); + ASSERT_EQ(grad_inputs.size(), 1); + ASSERT_EQ(grad_inputs[0]->Dims(), (std::vector{2, 3})); + test::ExpectTensorFloatEqual(grad_inputs[0], {0.0f, 2.0f, 0.0f, 4.0f, 0.0f, 6.0f}); +} + +// Gradients entering at the masked (zero) positions must stay 0 regardless of the +// incoming grad value, and vice versa for the active positions. +TEST_P(AutogradReLUBackwardTest, ReLUBackwardGradValuesIndependentOfMask) { + const std::vector input_values{-3.0f, -0.5f, -0.0f, 0.0f, 1.0f, 7.0f}; + const std::vector grad_values{-2.0f, 100.0f, -5.0f, 3.0f, -1.0f, 0.25f}; + auto input + = std::make_shared(input_values.data(), std::vector{6}, DataType::kFLOAT32, GetDevice()); + auto grad = std::make_shared(grad_values.data(), std::vector{6}, DataType::kFLOAT32, GetDevice()); + + auto relu_fn = std::make_shared(); + auto result = relu_fn->Apply({input}); + ASSERT_EQ(result.size(), 1); + auto grad_inputs = relu_fn->Backward({grad}); + ASSERT_EQ(grad_inputs.size(), 1); + + const auto host_grad_input = grad_inputs[0]->To(Device()); + const float *out = static_cast(host_grad_input.DataPtr()); + for (size_t idx = 0; idx < input_values.size(); ++idx) { + EXPECT_EQ(out[idx], ExpectedReluGrad(input_values[idx], grad_values[idx])) << "position " << idx; + } +} + +INFINI_TRAIN_REGISTER_TEST(AutogradReLUBackwardTest); diff --git a/tests/autograd/test_autograd_relu_forward.cc b/tests/autograd/test_autograd_relu_forward.cc new file mode 100644 index 00000000..bc7d5458 --- /dev/null +++ b/tests/autograd/test_autograd_relu_forward.cc @@ -0,0 +1,124 @@ +#include +#include +#include +#include + +#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; + +namespace { +uint32_t FloatBits(float value) { + uint32_t bits; + std::memcpy(&bits, &value, sizeof(bits)); + return bits; +} + +// Reference for the finite-valued part: NaN and -0 must survive untouched, +// negatives collapse to +0. +float ExpectedRelu(float x) { return x < 0.0f ? 0.0f : x; } +} // namespace + +class AutogradReLUForwardTest : public infini_train::test::InfiniTrainTest {}; + +// A single-case battery of NaN, -0, 0, +/-values, infinities and the fp32 extremes. +// NaN and -0 are asserted bit- or positionally because they have no ordering. +TEST_P(AutogradReLUForwardTest, ReLUForwardNaNZeroExtremes) { + const std::vector input_values{std::nanf(""), + -0.0f, + 0.0f, + 1.5f, + -2.5f, + std::numeric_limits::infinity(), + -std::numeric_limits::infinity(), + 1e-30f, + -1e-30f, + std::numeric_limits::max(), + -std::numeric_limits::max()}; + auto input + = std::make_shared(input_values.data(), std::vector{static_cast(input_values.size())}, + DataType::kFLOAT32, GetDevice()); + + auto relu_fn = std::make_shared(); + auto result = relu_fn->Apply({input}); + ASSERT_EQ(result.size(), 1); + ASSERT_EQ(result[0]->Dims(), input->Dims()); + + const auto host_result = result[0]->To(Device()); + const float *out = static_cast(host_result.DataPtr()); + + for (size_t idx = 0; idx < input_values.size(); ++idx) { + if (std::isnan(input_values[idx])) { + EXPECT_TRUE(std::isnan(out[idx])) << "NaN position " << idx << " must stay NaN"; + } else if (std::signbit(input_values[idx]) && input_values[idx] == 0.0f) { + EXPECT_EQ(out[idx], 0.0f) << "-0 position " << idx; + EXPECT_TRUE(std::signbit(out[idx])) << "-0 signbit lost at position " << idx; + } else { + EXPECT_EQ(FloatBits(out[idx]), FloatBits(ExpectedRelu(input_values[idx]))) + << "finite bit mismatch at position " << idx; + } + } +} + +TEST_P(AutogradReLUForwardTest, ReLUForwardModuleTwoD) { + const std::vector input_values{-1.0f, 0.0f, 2.0f, 3.5f, -4.25f, 0.5f}; + auto input + = std::make_shared(input_values.data(), std::vector{2, 3}, DataType::kFLOAT32, GetDevice()); + + auto relu = std::make_shared(); + auto result = (*relu)({input}); + ASSERT_EQ(result.size(), 1); + ASSERT_EQ(result[0]->Dims(), (std::vector{2, 3})); + test::ExpectTensorFloatEqual(result[0], {0.0f, 0.0f, 2.0f, 3.5f, 0.0f, 0.5f}); +} + +// ReLU is elementwise; the 4-D conv-output layout must keep its shape. +TEST_P(AutogradReLUForwardTest, ReLUForwardFourD) { + std::vector input_values; + for (int idx = 0; idx < 2 * 2 * 4 * 4; ++idx) { input_values.push_back((idx * 7) % 13 / 3.0f - 2.0f); } + auto input = std::make_shared(input_values.data(), std::vector{2, 2, 4, 4}, DataType::kFLOAT32, + GetDevice()); + + auto relu_fn = std::make_shared(); + auto result = relu_fn->Apply({input}); + ASSERT_EQ(result.size(), 1); + ASSERT_EQ(result[0]->Dims(), (std::vector{2, 2, 4, 4})); + + const auto host_result = result[0]->To(Device()); + const float *out = static_cast(host_result.DataPtr()); + for (size_t idx = 0; idx < input_values.size(); ++idx) { + EXPECT_EQ(out[idx], ExpectedRelu(input_values[idx])) << "position " << idx; + } +} + +// An empty batch is a valid 0-element tensor: the kernel must not touch memory. +TEST_P(AutogradReLUForwardTest, ReLUForwardEmptyBatch) { + auto input = std::make_shared(std::vector{0, 3}, DataType::kFLOAT32, GetDevice()); + auto relu_fn = std::make_shared(); + auto result = relu_fn->Apply({input}); + ASSERT_EQ(result.size(), 1); + EXPECT_EQ(result[0]->Dims(), (std::vector{0, 3})); + EXPECT_EQ(result[0]->NumElements(), 0); +} + +// The op is declared FP32-only; other dtypes must be rejected loudly. +TEST_P(AutogradReLUForwardTest, ReLUForwardRejectsNonFloat) { + ONLY_CPU(); + auto input = std::make_shared(std::vector{2, 3}, DataType::kFLOAT16, GetDevice()); + EXPECT_DEATH( + { + auto relu_fn = std::make_shared(); + auto result = relu_fn->Apply({input}); + (void)result; + }, + ""); +} + +INFINI_TRAIN_REGISTER_TEST(AutogradReLUForwardTest); diff --git a/tests/autograd/test_autograd_relu_train.cc b/tests/autograd/test_autograd_relu_train.cc new file mode 100644 index 00000000..1d0d130f --- /dev/null +++ b/tests/autograd/test_autograd_relu_train.cc @@ -0,0 +1,120 @@ +#include +#include +#include +#include + +#include "gtest/gtest.h" + +#include "infini_train/include/nn/modules/activations.h" +#include "infini_train/include/nn/modules/conv.h" +#include "infini_train/include/nn/modules/linear.h" +#include "infini_train/include/nn/modules/loss.h" +#include "infini_train/include/nn/parallel/global.h" +#include "infini_train/include/optimizer.h" +#include "infini_train/include/tensor.h" + +#include "tests/common/test_utils.h" + +using namespace infini_train; + +// Minimal end-to-end training step with ReLU in the middle: +// Conv2d -> ReLU -> Flatten -> Linear -> CrossEntropy, one loss->Backward() through the +// autograd graph and one SGD step, on CPU and CUDA. +class AutogradReLUTrainTest : public infini_train::test::InfiniTrainTest {}; + +TEST_P(AutogradReLUTrainTest, ReluConvFlattenLinearCrossEntropySgdStep) { + const Device device = GetDevice(); + const Device host = Device(); + constexpr float kLearningRate = 0.1f; + + std::vector input_values; + for (int idx = 0; idx < 2 * 1 * 8 * 8; ++idx) { input_values.push_back((idx * 37) % 101 / 50.0f - 1.0f); } + auto input + = std::make_shared(input_values.data(), std::vector{2, 1, 8, 8}, DataType::kFLOAT32, device); + auto target = std::make_shared(std::vector{2}, DataType::kINT64, device); + target->Fill(0); + + auto conv = std::make_shared(1, 2, 3, true, device); + auto relu = std::make_shared(); + auto fc = std::make_shared(2 * 6 * 6, 3, true, device); + auto loss_fn = std::make_shared(); + + auto conv_out = (*conv)({input})[0]; + ASSERT_EQ(conv_out->Dims(), (std::vector{2, 2, 6, 6})); + auto relu_out = (*relu)({conv_out})[0]; + ASSERT_EQ(relu_out->Dims(), (std::vector{2, 2, 6, 6})); + + // The pre-activations must contain negatives so the ReLU mask actually fires. + bool saw_negative = false; + bool saw_masked_zero = false; + const auto conv_cpu = conv_out->To(host); + const auto relu_cpu = relu_out->To(host); + const float *conv_data = static_cast(conv_cpu.DataPtr()); + const float *relu_data = static_cast(relu_cpu.DataPtr()); + for (size_t idx = 0; idx < conv_cpu.NumElements(); ++idx) { + ASSERT_TRUE(std::isfinite(relu_data[idx])); + if (conv_data[idx] < 0.0f) { + saw_negative = true; + } + if (relu_data[idx] == 0.0f) { + saw_masked_zero = true; + } + } + ASSERT_TRUE(saw_negative); + ASSERT_TRUE(saw_masked_zero); + + auto flat = relu_out->Flatten(1); + ASSERT_EQ(flat->Dims(), (std::vector{2, 72})); + auto logits = (*fc)({flat})[0]; + ASSERT_EQ(logits->Dims(), (std::vector{2, 3})); + auto loss = (*loss_fn)({logits, target})[0]; + ASSERT_TRUE(loss->Dims().empty()); + const auto loss_cpu = loss->To(host); + EXPECT_TRUE(std::isfinite(*static_cast(loss_cpu.DataPtr()))); + + loss->Backward(); + + std::vector> params = conv->Parameters(); + for (auto ¶m : fc->Parameters()) { params.push_back(param); } + ASSERT_EQ(params.size(), 4); + + // Every parameter must have received a finite, non-trivial gradient through the graph. + for (const auto ¶m : params) { + const auto &grad = param->grad(); + ASSERT_NE(grad, nullptr); + ASSERT_EQ(grad->Dims(), param->Dims()); + const auto grad_cpu = grad->To(host); + const auto *grad_data = static_cast(grad_cpu.DataPtr()); + float max_abs = 0.0f; + for (size_t idx = 0; idx < grad_cpu.NumElements(); ++idx) { + ASSERT_TRUE(std::isfinite(grad_data[idx])); + max_abs = std::max(max_abs, std::fabs(grad_data[idx])); + } + EXPECT_GT(max_abs, 0.0f); + } + + std::vector> old_values_cpu; + for (const auto ¶m : params) { + // CopyFrom (not To) so the snapshot owns its buffer: on CPU To(home device) + // returns a view, which would alias the in-place update below. + auto old_value = std::make_shared(param->Dims(), param->Dtype(), host); + old_value->CopyFrom(*param); + old_values_cpu.push_back(old_value); + } + + auto optimizer = std::make_shared(params, kLearningRate); + optimizer->Step(); + + for (size_t p = 0; p < params.size(); ++p) { + const auto new_cpu = params[p]->To(host); + const auto grad_cpu = params[p]->grad()->To(host); + const auto *old_data = static_cast(old_values_cpu[p]->DataPtr()); + const auto *grad_data = static_cast(grad_cpu.DataPtr()); + const auto *new_data = static_cast(new_cpu.DataPtr()); + for (size_t idx = 0; idx < params[p]->NumElements(); ++idx) { + EXPECT_FLOAT_EQ(new_data[idx], old_data[idx] - kLearningRate * grad_data[idx]) << "param " << p; + } + } +} + +INFINI_TRAIN_REGISTER_TEST(AutogradReLUTrainTest); From 3dec1e53beafe90f53d7fd85fc9cf51056ffc008 Mon Sep 17 00:00:00 2001 From: accelerator-llc Date: Sun, 6 Sep 2026 20:37:32 +0800 Subject: [PATCH 06/15] fix(mnist): compute dataset sample stride from the normalized dtype The MNIST dataset loader computed the per-sample byte stride from the on-disk UINT8 element size, but the image tensor is normalized to FLOAT32 in the constructor. Sample views after the first therefore read from a wrong, overlapping byte range, so most training images are misaligned mixes of real pixels and the network cannot learn. Use the FLOAT32 element size when computing the stride. Validated: MNIST MLP and CNN training reach >=95% test accuracy. --- example/mnist/dataset.cc | 7 ++++++- 1 file changed, 6 insertions(+), 1 deletion(-) diff --git a/example/mnist/dataset.cc b/example/mnist/dataset.cc index ee683f6d..fc798831 100644 --- a/example/mnist/dataset.cc +++ b/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())), label_size_in_bytes_(kSN3TypeToSize.at(label_file_.type) * std::accumulate(label_dims_.begin(), label_dims_.end(), 1, std::multiplies())) { From 2d50f4117a1ec61b2d6ccae9c16c03d709ef6b02 Mon Sep 17 00:00:00 2001 From: accelerator-llc Date: Sun, 6 Sep 2026 22:24:27 +0800 Subject: [PATCH 07/15] feat(mnist): add CNN demo selectable via --model flag Add MnistCnn (Conv2d(1,16,3)->ReLU->Conv2d(16,32,3)->ReLU->Flatten->Linear(18432,10)) alongside the existing MLP. A --model=mlp|cnn flag (default cnn) selects the network; the CNN restores the (N,1,28,28) spatial layout from the flattened [N,784] input with Tensor::View. Hold the loss in a shared_ptr to match the other examples. Tune the default --num_epoch/--lr (3 / 0.1) so the default CNN demo reaches ~97.8% test accuracy. Training converges after the MNIST dataset loader's per-sample stride is fixed to use the normalized FLOAT32 element size (older stride misaligned the image views). --- example/mnist/cnn_net.h | 44 +++++++++++++++++++++++++++++++++++++++++ example/mnist/main.cc | 27 ++++++++++++++++++------- 2 files changed, 64 insertions(+), 7 deletions(-) create mode 100644 example/mnist/cnn_net.h diff --git a/example/mnist/cnn_net.h b/example/mnist/cnn_net.h new file mode 100644 index 00000000..02e92d7d --- /dev/null +++ b/example/mnist/cnn_net.h @@ -0,0 +1,44 @@ +#pragma once + +#include +#include +#include + +#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> layers; + // Two 3x3 valid convs shrink 28 -> 26 -> 24; no pooling in the reference net. + layers.push_back(std::make_shared(1, 16, 3)); + layers.push_back(std::make_shared()); + layers.push_back(std::make_shared(16, 32, 3)); + layers.push_back(std::make_shared()); + modules_["sequential"] = std::make_shared(std::move(layers)); + // 32 * 24 * 24 = 18432 flattened features into the 10-class head. + modules_["linear"] = std::make_shared(32 * 24 * 24, 10); + } + + std::vector> + Forward(const std::vector> &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}); + } +}; diff --git a/example/mnist/main.cc b/example/mnist/main.cc index f9581dd3..7cb006aa 100644 --- a/example/mnist/main.cc +++ b/example/mnist/main.cc @@ -14,13 +14,17 @@ #include "infini_train/include/nn/modules/loss.h" #include "infini_train/include/optimizer.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,10 +35,14 @@ 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); @@ -47,13 +55,18 @@ int main(int argc, char *argv[]) { auto test_dataset = std::make_shared(FLAGS_dataset, false); DataLoader test_dataloader(test_dataset, FLAGS_bs); - auto network = std::make_shared(); + std::shared_ptr network; + if (FLAGS_model == kModelCNN) { + network = std::make_shared(); + } else { + network = std::make_shared(); + } Device device = FLAGS_device == kDeviceCPU ? Device() : Device(Device::DeviceType::kCUDA, 0); Device cpu_device = Device(); network->To(device); - auto loss_fn = nn::CrossEntropyLoss(); - loss_fn.To(device); + auto loss_fn = std::make_shared(); + loss_fn->To(device); auto optimizer = optimizers::SGD(network->Parameters(), FLAGS_lr); for (int epoch = 0; epoch < FLAGS_num_epoch; ++epoch) { @@ -69,7 +82,7 @@ int main(int argc, char *argv[]) { auto outputs = network->Forward({new_image}); optimizer.ZeroGrad(); - auto loss = loss_fn.Forward({outputs[0], new_label}); + auto loss = loss_fn->Forward({outputs[0], new_label}); loss[0]->Backward(); // Defer the loss D2H copy until after backward; reading it earlier would synchronize CUDA @@ -106,7 +119,7 @@ int main(int argc, char *argv[]) { 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 = loss_fn->Forward({outputs[0], new_label}); auto loss_cpu = loss[0]->To(cpu_device); const int batch_size = output_cpu.Dims()[0]; From 4d91fa1cc87509e64915b14836601ab71b7aaa1e Mon Sep 17 00:00:00 2001 From: accelerator-llc Date: Tue, 8 Sep 2026 12:32:13 +0800 Subject: [PATCH 08/15] feat(mnist): add DDP distributed training via infini_run Launch with infini_run to shard MNIST batches across processes: each process selects its GPU from LOCAL_RANK, wraps the network with DistributedDataParallel, broadcasts rank-0 parameters before training, and averages the per-step loss for global logging. Single-process behavior is unchanged; document --model and the multi-process launch in the README. --- README.md | 16 +++++++++ example/mnist/main.cc | 83 +++++++++++++++++++++++++++++++++++++++---- 2 files changed, 93 insertions(+), 6 deletions(-) diff --git a/README.md b/README.md index abd8070b..66f0e40e 100644 --- a/README.md +++ b/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 diff --git a/example/mnist/main.cc b/example/mnist/main.cc index 7cb006aa..dbfec2b2 100644 --- a/example/mnist/main.cc +++ b/example/mnist/main.cc @@ -4,6 +4,7 @@ #include #include #include +#include #include #include "gflags/gflags.h" @@ -12,7 +13,15 @@ #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" @@ -48,10 +57,49 @@ 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(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 ddp_train_dataloader; + std::optional 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(FLAGS_dataset, false); DataLoader test_dataloader(test_dataset, FLAGS_bs); @@ -61,12 +109,23 @@ int main(int argc, char *argv[]) { } else { network = std::make_shared(); } - Device device = FLAGS_device == kDeviceCPU ? Device() : Device(Device::DeviceType::kCUDA, 0); Device cpu_device = Device(); network->To(device); auto loss_fn = std::make_shared(); 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( + 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); for (int epoch = 0; epoch < FLAGS_num_epoch; ++epoch) { @@ -75,13 +134,15 @@ 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_loader) { auto new_image = std::make_shared(image->To(device)); auto new_label = std::make_shared(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 outputs = network->Forward({new_image}); + auto loss = loss_fn->Forward({outputs[0], new_label}); loss[0]->Backward(); @@ -89,10 +150,20 @@ int main(int argc, char *argv[]) { // between forward and backward. auto loss_cpu = loss[0]->To(cpu_device); float current_loss = static_cast(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(¤t_loss, std::vector{}, 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(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() - << "] " + LOG(ERROR) << "epoch: " << epoch << ", [" << train_idx * FLAGS_bs * ddp_world_size << "/" + << train_dataset->Size() << "] " << " loss: " << current_loss; } From 10923eb25a1257f9c997d6b5f9257fc849245727 Mon Sep 17 00:00:00 2001 From: accelerator-llc Date: Tue, 8 Sep 2026 20:56:33 +0800 Subject: [PATCH 09/15] test(linear): add bias-gradient regression tests Cover non-square grad_output shapes whose per-element values separate the sample-dimension sum from a row sum, and the bf16 branch whose grad_bias is promoted to fp32. --- .../autograd/test_autograd_linear_backward.cc | 104 ++++++++++++++++++ 1 file changed, 104 insertions(+) diff --git a/tests/autograd/test_autograd_linear_backward.cc b/tests/autograd/test_autograd_linear_backward.cc index 9ce88eee..8d5f9336 100644 --- a/tests/autograd/test_autograd_linear_backward.cc +++ b/tests/autograd/test_autograd_linear_backward.cc @@ -48,4 +48,108 @@ TEST_P(AutogradLinearBackwardTest, LinearBackwardNoBias) { EXPECT_EQ(grad_inputs.size(), 2); } +// Non-square (bs=4, out_features=3) with distinct values per element: the bias +// gradient must be the column sum of grad_output ([18, 22, 26] below), which a +// row-sum misread would report as [6, 22, 38]. Constant-valued grad_output +// cannot distinguish the two reductions, so values must vary. +TEST_P(AutogradLinearBackwardTest, LinearBackwardValues) { + const std::vector input_values{1.0f, 2.0f, 3.0f, 4.0f, 5.0f, 6.0f, 7.0f, 8.0f}; + const std::vector weight_values{1.0f, 0.0f, 0.0f, 1.0f, 1.0f, 1.0f}; + const std::vector grad_values{0.0f, 1.0f, 2.0f, 3.0f, 4.0f, 5.0f, 6.0f, 7.0f, 8.0f, 9.0f, 10.0f, 11.0f}; + + auto input + = std::make_shared(input_values.data(), std::vector{4, 2}, DataType::kFLOAT32, GetDevice()); + input->set_requires_grad(true); + auto weight + = std::make_shared(weight_values.data(), std::vector{3, 2}, DataType::kFLOAT32, GetDevice()); + weight->set_requires_grad(true); + auto bias = std::make_shared(std::vector{3}, DataType::kFLOAT32, GetDevice(), true); + bias->Fill(0.0f); + auto grad + = std::make_shared(grad_values.data(), std::vector{4, 3}, DataType::kFLOAT32, GetDevice()); + + auto linear_fn = std::make_shared(); + linear_fn->Apply({input, weight, bias}); + auto grad_inputs = linear_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); + + // grad_bias = grad_output.sum(dim=0) = [0+3+6+9, 1+4+7+10, 2+5+8+11] + EXPECT_EQ(grad_inputs[2]->Dims(), (std::vector{3})); + test::ExpectTensorFloatEqual(grad_inputs[2], {18.0f, 22.0f, 26.0f}); + // grad_input = grad_output * weight + EXPECT_EQ(grad_inputs[0]->Dims(), (std::vector{4, 2})); + test::ExpectTensorFloatEqual(grad_inputs[0], {2.0f, 3.0f, 8.0f, 9.0f, 14.0f, 15.0f, 20.0f, 21.0f}); + // grad_weight = grad_output^T * input + EXPECT_EQ(grad_inputs[1]->Dims(), (std::vector{3, 2})); + test::ExpectTensorFloatEqual(grad_inputs[1], {102.0f, 120.0f, 118.0f, 140.0f, 134.0f, 160.0f}); +} + +// Larger non-square shape with a deterministic pseudo-random pattern, checked +// against a double-precision host reference of the sample-dimension sum. +TEST_P(AutogradLinearBackwardTest, LinearBackwardBiasValues) { + constexpr int64_t bs = 16; + constexpr int64_t out_features = 7; + constexpr int64_t in_features = 3; + + std::vector grad_values(bs * out_features); + for (size_t i = 0; i < grad_values.size(); ++i) { grad_values[i] = static_cast((i * 13 + 7) % 23) - 11.0f; } + + auto input = std::make_shared(std::vector{bs, in_features}, DataType::kFLOAT32, GetDevice(), true); + input->Fill(1.0f); + auto weight = std::make_shared(std::vector{out_features, in_features}, DataType::kFLOAT32, + GetDevice(), true); + weight->Fill(1.0f); + auto bias = std::make_shared(std::vector{out_features}, DataType::kFLOAT32, GetDevice(), true); + bias->Fill(0.0f); + auto grad = std::make_shared(grad_values.data(), std::vector{bs, out_features}, DataType::kFLOAT32, + GetDevice()); + + auto linear_fn = std::make_shared(); + linear_fn->Apply({input, weight, bias}); + auto grad_inputs = linear_fn->Backward({grad}); + + ASSERT_EQ(grad_inputs.size(), 3); + std::vector expected_bias(out_features, 0.0); + for (int64_t i = 0; i < bs; ++i) { + for (int64_t j = 0; j < out_features; ++j) { expected_bias[j] += grad_values[i * out_features + j]; } + } + const std::vector expected_bias_f32(expected_bias.begin(), expected_bias.end()); + test::ExpectTensorNear(grad_inputs[2], expected_bias_f32, 1e-4f); +} + +// The bf16 branch accumulates in fp32 and returns a promoted fp32 grad_bias; +// bf16 kernels are CUDA-only (the CPU Linear path is fp32-only). +TEST_P(AutogradLinearBackwardTest, LinearBackwardBiasBFloat16) { + ONLY_CUDA(); + const std::vector grad_values{0.0f, 1.0f, 2.0f, 3.0f, 4.0f, 5.0f, 6.0f, 7.0f, 8.0f, 9.0f, 10.0f, 11.0f}; + + auto grad_f32 + = std::make_shared(grad_values.data(), std::vector{4, 3}, DataType::kFLOAT32, GetDevice()); + auto grad = std::make_shared(grad_f32->To(DataType::kBFLOAT16)); + + auto input_f32 = std::make_shared(std::vector{4, 2}, DataType::kFLOAT32, GetDevice(), true); + input_f32->Fill(1.0f); + auto input = std::make_shared(input_f32->To(DataType::kBFLOAT16)); + input->set_requires_grad(true); + auto weight_f32 = std::make_shared(std::vector{3, 2}, DataType::kFLOAT32, GetDevice(), true); + weight_f32->Fill(1.0f); + auto weight = std::make_shared(weight_f32->To(DataType::kBFLOAT16)); + weight->set_requires_grad(true); + auto bias = std::make_shared(std::vector{3}, DataType::kBFLOAT16, GetDevice(), true); + bias->Fill(0.0f); + + auto linear_fn = std::make_shared(); + linear_fn->Apply({input, weight, bias}); + auto grad_inputs = linear_fn->Backward({grad}); + + ASSERT_EQ(grad_inputs.size(), 3); + // bf16 inputs 0..11 are exact; the fp32 accumulation must reproduce [18, 22, 26]. + EXPECT_EQ(grad_inputs[2]->Dtype(), DataType::kFLOAT32); + test::ExpectTensorFloatEqual(grad_inputs[2], {18.0f, 22.0f, 26.0f}); +} + INFINI_TRAIN_REGISTER_TEST(AutogradLinearBackwardTest); From b07d1eeef0f3d7d52bf77f733bea44b0ba0f5b2e Mon Sep 17 00:00:00 2001 From: accelerator-llc Date: Tue, 8 Sep 2026 22:09:59 +0800 Subject: [PATCH 10/15] refactor(mnist): invoke modules through operator() Route the training and test loops through Module::operator() instead of calling Forward directly, matching the gpt2 example and the hook design doc: operator() is the entry that runs module hooks. No hook consumers exist yet, so outputs are unchanged. --- example/mnist/main.cc | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/example/mnist/main.cc b/example/mnist/main.cc index dbfec2b2..ef3c3063 100644 --- a/example/mnist/main.cc +++ b/example/mnist/main.cc @@ -141,9 +141,9 @@ int main(int argc, char *argv[]) { // Zero grads before forward: DDP rebinds param.grad to its bucket view during forward. optimizer.ZeroGrad(); - auto outputs = network->Forward({new_image}); + 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 @@ -188,9 +188,9 @@ int main(int argc, char *argv[]) { auto new_label = std::make_shared(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]; From 13f57d466811399ff447b6596e33c9323ac7ef78 Mon Sep 17 00:00:00 2001 From: accelerator-llc Date: Tue, 8 Sep 2026 22:11:05 +0800 Subject: [PATCH 11/15] test(conv): merge device-split conv tests into parameterized files Fold the *_cuda_* conv test files into their CPU counterparts so each test definition instantiates on every available device, matching the test infrastructure design and leaving no ONLY_ macros in the merged files. The CUDA-only empty-batch cases are device-independent and carry over, and the negative-value check in the extreme-values forward case now also runs on CUDA. --- tests/autograd/test_autograd_conv_backward.cc | 34 ++- .../test_autograd_conv_cuda_backward.cc | 247 ------------------ .../test_autograd_conv_cuda_forward.cc | 189 -------------- tests/autograd/test_autograd_conv_forward.cc | 22 +- 4 files changed, 42 insertions(+), 450 deletions(-) delete mode 100644 tests/autograd/test_autograd_conv_cuda_backward.cc delete mode 100644 tests/autograd/test_autograd_conv_cuda_forward.cc diff --git a/tests/autograd/test_autograd_conv_backward.cc b/tests/autograd/test_autograd_conv_backward.cc index 153f7668..8390ce58 100644 --- a/tests/autograd/test_autograd_conv_backward.cc +++ b/tests/autograd/test_autograd_conv_backward.cc @@ -5,7 +5,6 @@ #include "gtest/gtest.h" #include "infini_train/include/autograd/conv.h" -#include "infini_train/include/nn/parallel/global.h" #include "infini_train/include/tensor.h" #include "tests/common/test_utils.h" @@ -22,7 +21,6 @@ class AutogradConvBackwardTest : public infini_train::test::InfiniTrainTest {}; // Hand-checkable case: x = arange(9) reshaped (1,1,3,3), w = [[1,2],[3,4]], bias = 0.5, // grad_output = ones(1,1,2,2). bias does not influence grad_input / grad_weight. TEST_P(AutogradConvBackwardTest, ConvBackwardGradients) { - ONLY_CPU(); std::vector input_values; for (int idx = 0; idx < 9; ++idx) { input_values.push_back(static_cast(idx)); } auto input = std::make_shared(input_values.data(), std::vector{1, 1, 3, 3}, DataType::kFLOAT32, @@ -60,7 +58,6 @@ TEST_P(AutogradConvBackwardTest, ConvBackwardGradients) { } TEST_P(AutogradConvBackwardTest, ConvBackwardNoBias) { - ONLY_CPU(); auto input = std::make_shared(std::vector{2, 2, 4, 4}, DataType::kFLOAT32, GetDevice(), true); input->Fill(1.0f); auto weight = std::make_shared(std::vector{3, 2, 3, 3}, DataType::kFLOAT32, GetDevice(), true); @@ -83,7 +80,6 @@ TEST_P(AutogradConvBackwardTest, ConvBackwardNoBias) { // H = W = kernel: a single window, so grad_input mirrors the weight and grad_weight the input. TEST_P(AutogradConvBackwardTest, ConvBackwardKernelEqualsSpatial) { - ONLY_CPU(); const std::vector input_values{1.0f, 2.0f, 3.0f, 4.0f}; auto input = std::make_shared(input_values.data(), std::vector{1, 1, 2, 2}, DataType::kFLOAT32, GetDevice()) @@ -113,7 +109,6 @@ TEST_P(AutogradConvBackwardTest, ConvBackwardKernelEqualsSpatial) { // Training-shaped call: the activation is a leaf without requires_grad, so only the parameter // gradients must be produced. TEST_P(AutogradConvBackwardTest, ConvBackwardInputNotRequired) { - ONLY_CPU(); auto input = std::make_shared(std::vector{2, 2, 4, 4}, DataType::kFLOAT32, GetDevice()); input->Fill(1.0f); auto weight = std::make_shared(std::vector{3, 2, 3, 3}, DataType::kFLOAT32, GetDevice(), true); @@ -138,7 +133,6 @@ TEST_P(AutogradConvBackwardTest, ConvBackwardInputNotRequired) { // Random 4D case checked against PyTorch autograd (fixed seed, see file-level comment). TEST_P(AutogradConvBackwardTest, ConvBackwardTorchGolden) { - ONLY_CPU(); const std::vector input_values{ 2.014464855f, -0.671311080f, -0.945236862f, -0.096128546f, 0.889558852f, 1.726294398f, -0.078931913f, 0.205890238f, 0.094608001f, 0.173616216f, 0.437049866f, -0.557070732f, 0.455583423f, -0.736482263f, @@ -220,4 +214,32 @@ TEST_P(AutogradConvBackwardTest, ConvBackwardTorchGolden) { test::ExpectTensorNear(grad_inputs[2], expected_grad_bias, kGoldenAbsError); } +// An empty batch is degenerate but must not crash nor return uninitialized gradients: the +// parameter gradients (weight/bias) are sums over zero images, so they are exact zeros. +TEST_P(AutogradConvBackwardTest, ConvBackwardEmptyBatch) { + auto input + = std::make_shared(std::vector{0, 2, 4, 4}, DataType::kFLOAT32, GetDevice())->RequiresGrad(); + auto weight = std::make_shared(std::vector{3, 2, 3, 3}, DataType::kFLOAT32, GetDevice(), true) + ->RequiresGrad(); + auto bias + = std::make_shared(std::vector{3}, DataType::kFLOAT32, GetDevice(), true)->RequiresGrad(); + + auto conv_fn = std::make_shared(); + auto result = conv_fn->Apply({input, weight, bias}); + auto grad_output = std::make_shared(std::vector{0, 3, 2, 2}, DataType::kFLOAT32, GetDevice()); + auto grads = conv_fn->Backward({grad_output}); + ASSERT_EQ(grads.size(), 3); + ASSERT_NE(grads[0], nullptr); + ASSERT_NE(grads[1], nullptr); + ASSERT_NE(grads[2], nullptr); + + // grad_input has an empty batch, grad_weight and grad_bias are exact zeros over the empty batch. + EXPECT_EQ(grads[0]->Dims(), (std::vector{0, 2, 4, 4})); + EXPECT_EQ(grads[0]->NumElements(), 0); + EXPECT_EQ(grads[1]->Dims(), (std::vector{3, 2, 3, 3})); + test::ExpectTensorFloatEqual(grads[1], 0.0f); + EXPECT_EQ(grads[2]->Dims(), (std::vector{3})); + test::ExpectTensorFloatEqual(grads[2], 0.0f); +} + INFINI_TRAIN_REGISTER_TEST(AutogradConvBackwardTest); diff --git a/tests/autograd/test_autograd_conv_cuda_backward.cc b/tests/autograd/test_autograd_conv_cuda_backward.cc deleted file mode 100644 index b7b050ea..00000000 --- a/tests/autograd/test_autograd_conv_cuda_backward.cc +++ /dev/null @@ -1,247 +0,0 @@ -#include -#include -#include - -#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; - -namespace { -// Same golden case as the forward test (torch 2.14, seed 20260905, autograd.grad reference). -constexpr float kGoldenAbsError = 2e-6f; -} // namespace - -class AutogradConvCudaBackwardTest : public infini_train::test::InfiniTrainTest {}; - -// Hand-checkable case: x = arange(9) reshaped (1,1,3,3), w = [[1,2],[3,4]], bias = 0.5, -// grad_output = ones(1,1,2,2). -TEST_P(AutogradConvCudaBackwardTest, ConvBackwardGradients) { - ONLY_CUDA(); - std::vector input_values; - for (int idx = 0; idx < 9; ++idx) { input_values.push_back(static_cast(idx)); } - auto input = std::make_shared(input_values.data(), std::vector{1, 1, 3, 3}, DataType::kFLOAT32, - GetDevice()) - ->RequiresGrad(); - const std::vector weight_values{1.0f, 2.0f, 3.0f, 4.0f}; - auto weight = std::make_shared(weight_values.data(), std::vector{1, 1, 2, 2}, DataType::kFLOAT32, - GetDevice()) - ->RequiresGrad(); - auto bias = std::make_shared(std::vector{1}, DataType::kFLOAT32, GetDevice(), true); - bias->Fill(0.5f); - - auto conv_fn = std::make_shared(); - auto result = conv_fn->Apply({input, weight, bias}); - auto grad_output - = std::make_shared(std::vector{1, 1, 2, 2}, DataType::kFLOAT32, GetDevice(), true); - grad_output->Fill(1.0f); - auto grad_inputs = conv_fn->Backward({grad_output}); - ASSERT_EQ(grad_inputs.size(), 3); - ASSERT_NE(grad_inputs[0], nullptr); - ASSERT_NE(grad_inputs[1], nullptr); - ASSERT_NE(grad_inputs[2], nullptr); - - const std::vector expected_grad_input{1.0f, 3.0f, 2.0f, 4.0f, 10.0f, 6.0f, 3.0f, 7.0f, 4.0f}; - EXPECT_EQ(grad_inputs[0]->Dims(), (std::vector{1, 1, 3, 3})); - test::ExpectTensorFloatEqual(grad_inputs[0], expected_grad_input); - - const std::vector expected_grad_weight{8.0f, 12.0f, 20.0f, 24.0f}; - EXPECT_EQ(grad_inputs[1]->Dims(), (std::vector{1, 1, 2, 2})); - test::ExpectTensorFloatEqual(grad_inputs[1], expected_grad_weight); - - test::ExpectTensorFloatEqual(grad_inputs[2], std::vector{4.0f}); -} - -// No-bias path returns {grad_input, grad_weight} and both must be correct. -TEST_P(AutogradConvCudaBackwardTest, ConvBackwardNoBias) { - ONLY_CUDA(); - auto input = std::make_shared(std::vector{2, 2, 4, 4}, DataType::kFLOAT32, GetDevice(), true); - input->Fill(1.0f); - auto weight = std::make_shared(std::vector{3, 2, 3, 3}, DataType::kFLOAT32, GetDevice(), true); - weight->Fill(1.0f); - auto conv_fn = std::make_shared(); - auto result = conv_fn->Apply({input, weight}); - auto grad = std::make_shared(std::vector{2, 3, 2, 2}, DataType::kFLOAT32, GetDevice(), true); - grad->Fill(1.0f); - auto grad_inputs = conv_fn->Backward({grad}); - EXPECT_EQ(grad_inputs.size(), 2); - - ASSERT_NE(grad_inputs[0], nullptr); - EXPECT_EQ(grad_inputs[0]->Dims(), (std::vector{2, 2, 4, 4})); - ASSERT_NE(grad_inputs[1], nullptr); - EXPECT_EQ(grad_inputs[1]->Dims(), (std::vector{3, 2, 3, 3})); - // With all-ones input/weight/grad_output, grad_weight(o,c,u,v) = sum_{n,i,j} 1 = 2*2*2. - test::ExpectTensorFloatEqual(grad_inputs[1], 8.0f); -} - -// H = W = kernel: a single window, so grad_input mirrors the weight and grad_weight the input. -TEST_P(AutogradConvCudaBackwardTest, ConvBackwardKernelEqualsSpatial) { - ONLY_CUDA(); - const std::vector input_values{1.0f, 2.0f, 3.0f, 4.0f}; - auto input = std::make_shared(input_values.data(), std::vector{1, 1, 2, 2}, DataType::kFLOAT32, - GetDevice()) - ->RequiresGrad(); - const std::vector weight_values{5.0f, 6.0f, 7.0f, 8.0f}; - auto weight = std::make_shared(weight_values.data(), std::vector{1, 1, 2, 2}, DataType::kFLOAT32, - GetDevice()) - ->RequiresGrad(); - auto bias = std::make_shared(std::vector{1}, DataType::kFLOAT32, GetDevice(), true); - bias->Fill(0.25f); - - auto conv_fn = std::make_shared(); - auto result = conv_fn->Apply({input, weight, bias}); - auto grad_output - = std::make_shared(std::vector{1, 1, 1, 1}, DataType::kFLOAT32, GetDevice(), true); - grad_output->Fill(1.0f); - auto grad_inputs = conv_fn->Backward({grad_output}); - ASSERT_EQ(grad_inputs.size(), 3); - - EXPECT_EQ(grad_inputs[0]->Dims(), (std::vector{1, 1, 2, 2})); - test::ExpectTensorFloatEqual(grad_inputs[0], weight_values); - EXPECT_EQ(grad_inputs[1]->Dims(), (std::vector{1, 1, 2, 2})); - test::ExpectTensorFloatEqual(grad_inputs[1], input_values); - test::ExpectTensorFloatEqual(grad_inputs[2], std::vector{1.0f}); -} - -// Training-shaped call: the activation is a leaf without requires_grad, so only the parameter -// gradients must be produced. -TEST_P(AutogradConvCudaBackwardTest, ConvBackwardInputNotRequired) { - ONLY_CUDA(); - auto input = std::make_shared(std::vector{2, 2, 4, 4}, DataType::kFLOAT32, GetDevice()); - input->Fill(1.0f); - auto weight = std::make_shared(std::vector{3, 2, 3, 3}, DataType::kFLOAT32, GetDevice(), true); - weight->Fill(1.0f); - auto bias = std::make_shared(std::vector{3}, DataType::kFLOAT32, GetDevice(), true); - bias->Fill(1.0f); - - auto conv_fn = std::make_shared(); - auto result = conv_fn->Apply({input, weight, bias}); - auto grad = std::make_shared(std::vector{2, 3, 2, 2}, DataType::kFLOAT32, GetDevice(), true); - grad->Fill(1.0f); - auto grad_inputs = conv_fn->Backward({grad}); - ASSERT_EQ(grad_inputs.size(), 3); - EXPECT_EQ(grad_inputs[0], nullptr); - EXPECT_NE(grad_inputs[1], nullptr); - EXPECT_NE(grad_inputs[2], nullptr); - test::ExpectTensorFloatEqual(grad_inputs[1], 8.0f); - test::ExpectTensorFloatEqual(grad_inputs[2], 8.0f); -} - -// Random 4D case checked against PyTorch autograd (fixed seed, see file-level comment). -TEST_P(AutogradConvCudaBackwardTest, ConvBackwardTorchGolden) { - ONLY_CUDA(); - const std::vector input_values{ - 2.014464855f, -0.671311080f, -0.945236862f, -0.096128546f, 0.889558852f, 1.726294398f, -0.078931913f, - 0.205890238f, 0.094608001f, 0.173616216f, 0.437049866f, -0.557070732f, 0.455583423f, -0.736482263f, - -0.718647420f, 0.926798999f, 1.732108712f, -0.138735890f, -1.508729339f, 2.127460241f, 0.079739936f, - 0.217489868f, 0.589455068f, -0.052687794f, -1.659490585f, -0.380473137f, -0.166608080f, 0.350353301f, - -1.411424637f, 0.626776755f, 1.336345553f, -0.203511983f, -0.757243156f, 1.381434083f, -0.417290032f, - -0.668979943f, -0.569557011f, -1.035447836f, -0.555948615f, -1.279154539f, 0.056448594f, -0.388204783f, - 2.071142673f, -0.803012371f, 1.956272721f, 0.610922813f, -1.668144941f, -0.628582716f, -0.756579638f, - 0.987797141f, -0.849777400f, -1.654165864f, -0.326211363f, -0.040599853f, 0.281366080f, 0.499955416f, - -0.518889010f, -0.409607649f, 0.689743340f, 1.129151702f, 0.798748374f, -1.004696369f, -0.104189672f, - 1.257721663f, - }; - auto input = std::make_shared(input_values.data(), std::vector{2, 2, 4, 4}, DataType::kFLOAT32, - GetDevice()) - ->RequiresGrad(); - const std::vector weight_values{ - -2.448281050f, -0.608355999f, -0.958948731f, 1.009868145f, 0.031469870f, -1.101765275f, 0.571945250f, - -1.224039674f, -0.426795214f, -0.612229824f, -0.744571507f, -1.164491534f, -0.726353824f, -1.124407530f, - -0.498482078f, -0.423681349f, -1.231436968f, -0.689355612f, -0.170030892f, -0.120427191f, 1.032825112f, - 1.369398594f, 1.237134099f, 0.458183825f, -0.396747291f, 0.548917472f, 0.043999992f, 0.891736746f, - -1.292207360f, 1.713299513f, 0.982937992f, 0.432634443f, 0.293460459f, 0.318505853f, -0.146791920f, - 0.983720183f, -0.162205622f, 0.616937041f, 0.929942310f, -0.411319345f, 0.301661789f, -1.444623590f, - -0.730446517f, 2.391211271f, 0.400584310f, -0.909495890f, -0.702420592f, 0.372547686f, 0.994188488f, - 0.768918931f, 0.740486383f, -1.176032424f, 0.868664801f, 0.422459871f, - }; - auto weight = std::make_shared(weight_values.data(), std::vector{3, 2, 3, 3}, DataType::kFLOAT32, - GetDevice()) - ->RequiresGrad(); - const std::vector bias_values{-0.640517414f, 0.982376575f, 0.729396820f}; - auto bias = std::make_shared(bias_values.data(), std::vector{3}, DataType::kFLOAT32, GetDevice()) - ->RequiresGrad(); - - auto conv_fn = std::make_shared(); - auto result = conv_fn->Apply({input, weight, bias}); - const std::vector grad_output_values{ - -1.277466655f, -1.339194655f, -1.159473658f, 0.086730979f, 0.438649237f, 0.602657557f, - -0.738640666f, -0.410562515f, 0.201109841f, 0.716394842f, 0.003985748f, 0.676588118f, - 0.412507862f, -0.146394536f, -0.880041420f, -1.247093678f, -1.212666154f, -0.393535852f, - 0.564200640f, 0.279093444f, -2.798053741f, -1.615590930f, -0.518746257f, -0.283052236f, - }; - auto grad_output = std::make_shared(grad_output_values.data(), std::vector{2, 3, 2, 2}, - DataType::kFLOAT32, GetDevice()); - - auto grad_inputs = conv_fn->Backward({grad_output}); - ASSERT_EQ(grad_inputs.size(), 3); - ASSERT_NE(grad_inputs[0], nullptr); - ASSERT_NE(grad_inputs[1], nullptr); - ASSERT_NE(grad_inputs[2], nullptr); - - const std::vector expected_grad_input{ - 3.020392179f, 3.908452511f, 3.049194813f, 2.572864771f, 2.191555023f, 0.285838872f, 3.004242659f, - 0.838663280f, -3.235622168f, -0.945002913f, 4.960352898f, -0.376030058f, -0.373013139f, 0.741603315f, - 1.750292301f, 0.215949476f, 0.990354240f, 0.948824525f, 2.029216766f, 2.858904839f, 1.606565118f, - 4.838562012f, 3.308694601f, 0.822548687f, 0.564552128f, 2.793590069f, 4.189336777f, 2.155962467f, - 0.251298606f, 0.576505065f, 0.615549624f, -0.177835792f, -0.349884391f, -1.143750548f, -5.110340118f, - -1.768475175f, 2.049649715f, 0.845555902f, 3.547870159f, 3.535832882f, 2.858142614f, -6.855783939f, - -3.052171230f, 1.308768988f, -0.348264217f, -0.470771074f, 1.195474505f, 0.431147456f, 1.210889816f, - 4.433355331f, -1.848075271f, -1.105654240f, -2.759691477f, -3.466245174f, -1.316413283f, 0.586114824f, - 3.407652855f, 0.809629261f, -2.299243927f, -0.474773228f, 1.162620664f, 1.500420690f, 2.191398382f, - 1.014662623f, - }; - const std::vector expected_grad_weight{ - -1.278161168f, 2.350493193f, 3.455903053f, -3.191775799f, -4.849991322f, -1.593288898f, -3.348839760f, - 1.063873172f, 4.325201988f, -2.219501734f, 2.213379383f, -2.589243412f, 2.336852074f, -1.198173761f, - -2.431171656f, 4.716279507f, 0.842501462f, -3.181046009f, -1.122435212f, -4.357190609f, -0.400157630f, - 2.311078310f, 2.235447884f, 2.117345333f, 1.470544577f, 0.713402152f, -3.305890560f, 0.861350715f, - -2.180587769f, 2.186314106f, 1.552511692f, 0.699985623f, 0.372371852f, 0.788659871f, -1.649568439f, - -1.754016399f, 1.571256876f, -3.355353594f, 2.778887510f, 4.880561829f, 3.997585535f, 2.531520605f, - -1.071573257f, -2.245790720f, -3.141553879f, 1.098210573f, -2.158659220f, 5.950089931f, 1.271268964f, - 0.028057545f, -1.955230832f, 1.795806646f, 1.293214321f, -3.971021652f, - }; - const std::vector expected_grad_bias{-5.550425529f, -0.870804369f, -3.617365122f}; - - EXPECT_EQ(grad_inputs[0]->Dims(), (std::vector{2, 2, 4, 4})); - test::ExpectTensorNear(grad_inputs[0], expected_grad_input, kGoldenAbsError); - EXPECT_EQ(grad_inputs[1]->Dims(), (std::vector{3, 2, 3, 3})); - test::ExpectTensorNear(grad_inputs[1], expected_grad_weight, kGoldenAbsError); - EXPECT_EQ(grad_inputs[2]->Dims(), (std::vector{3})); - test::ExpectTensorNear(grad_inputs[2], expected_grad_bias, kGoldenAbsError); -} - -// An empty batch is degenerate but must not crash nor return uninitialized gradients: the -// parameter gradients (weight/bias) are sums over zero images, so they are exact zeros. -TEST_P(AutogradConvCudaBackwardTest, ConvBackwardEmptyBatch) { - ONLY_CUDA(); - auto input - = std::make_shared(std::vector{0, 2, 4, 4}, DataType::kFLOAT32, GetDevice())->RequiresGrad(); - auto weight = std::make_shared(std::vector{3, 2, 3, 3}, DataType::kFLOAT32, GetDevice(), true) - ->RequiresGrad(); - auto bias - = std::make_shared(std::vector{3}, DataType::kFLOAT32, GetDevice(), true)->RequiresGrad(); - - auto conv_fn = std::make_shared(); - auto result = conv_fn->Apply({input, weight, bias}); - auto grad_output = std::make_shared(std::vector{0, 3, 2, 2}, DataType::kFLOAT32, GetDevice()); - auto grads = conv_fn->Backward({grad_output}); - ASSERT_EQ(grads.size(), 3); - ASSERT_NE(grads[0], nullptr); - ASSERT_NE(grads[1], nullptr); - ASSERT_NE(grads[2], nullptr); - - // grad_input has an empty batch, grad_weight and grad_bias are exact zeros over the empty batch. - EXPECT_EQ(grads[0]->Dims(), (std::vector{0, 2, 4, 4})); - EXPECT_EQ(grads[0]->NumElements(), 0); - EXPECT_EQ(grads[1]->Dims(), (std::vector{3, 2, 3, 3})); - test::ExpectTensorFloatEqual(grads[1], 0.0f); - EXPECT_EQ(grads[2]->Dims(), (std::vector{3})); - test::ExpectTensorFloatEqual(grads[2], 0.0f); -} - -INFINI_TRAIN_REGISTER_TEST(AutogradConvCudaBackwardTest); diff --git a/tests/autograd/test_autograd_conv_cuda_forward.cc b/tests/autograd/test_autograd_conv_cuda_forward.cc deleted file mode 100644 index 40ecacdc..00000000 --- a/tests/autograd/test_autograd_conv_cuda_forward.cc +++ /dev/null @@ -1,189 +0,0 @@ -#include -#include -#include - -#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; - -namespace { -// Deterministic golden data, generated with torch 2.14 under seed 20260905 -// (conv2d / autograd.grad reference). -constexpr float kGoldenAbsError = 2e-6f; -} // namespace - -class AutogradConvCudaForwardTest : public infini_train::test::InfiniTrainTest {}; - -// An asymmetric kernel pins the cross-correlation semantics: flipping the kernel would swap -// the roles of the taps and change every output value. The values are small exact integers, -// so the CUDA path reproduces them bit-exactly. -TEST_P(AutogradConvCudaForwardTest, ConvForwardAsymmetricKernel) { - ONLY_CUDA(); - std::vector input_values; - for (int idx = 0; idx < 16; ++idx) { input_values.push_back(static_cast(idx)); } - auto input = std::make_shared(input_values.data(), std::vector{1, 1, 4, 4}, DataType::kFLOAT32, - GetDevice()); - const std::vector weight_values{1.0f, 2.0f, 3.0f, 4.0f}; - auto weight = std::make_shared(weight_values.data(), std::vector{1, 1, 2, 2}, DataType::kFLOAT32, - GetDevice()); - // out[i][j] = 1*x[i][j] + 2*x[i][j+1] + 3*x[i+1][j] + 4*x[i+1][j+1] - const std::vector expected{34.0f, 44.0f, 54.0f, 74.0f, 84.0f, 94.0f, 114.0f, 124.0f, 134.0f}; - - auto conv_fn = std::make_shared(); - auto result = conv_fn->Apply({input, weight}); - ASSERT_EQ(result.size(), 1); - EXPECT_EQ(result[0]->Dims(), (std::vector{1, 1, 3, 3})); - test::ExpectTensorFloatEqual(result[0], expected); -} - -// Cross-channel / batch weighting with a per-channel bias. -TEST_P(AutogradConvCudaForwardTest, ConvForwardMultiChannelBatchBias) { - ONLY_CUDA(); - std::vector input_values; - for (int n = 0; n < 2; ++n) { - for (int c = 0; c < 2; ++c) { - for (int h = 0; h < 3; ++h) { - for (int w = 0; w < 3; ++w) { input_values.push_back(10.0f * n + 5.0f * c + 3.0f * h + w); } - } - } - } - auto input = std::make_shared(input_values.data(), std::vector{2, 2, 3, 3}, DataType::kFLOAT32, - GetDevice()); - const std::vector weight_values{1.0f, 1.0f, 1.0f, 1.0f, 1.0f, 1.0f, 1.0f, 1.0f, - 1.0f, 0.0f, 0.0f, 0.0f, 0.0f, 0.0f, 0.0f, 2.0f}; - auto weight = std::make_shared(weight_values.data(), std::vector{2, 2, 2, 2}, DataType::kFLOAT32, - GetDevice()); - const std::vector bias_values{0.5f, -0.25f}; - auto bias = std::make_shared(bias_values.data(), std::vector{2}, DataType::kFLOAT32, GetDevice()); - const std::vector expected{36.5f, 44.5f, 60.5f, 68.5f, 17.75f, 20.75f, 26.75f, 29.75f, - 116.5f, 124.5f, 140.5f, 148.5f, 47.75f, 50.75f, 56.75f, 59.75f}; - - auto conv_fn = std::make_shared(); - auto result = conv_fn->Apply({input, weight, bias}); - ASSERT_EQ(result.size(), 1); - EXPECT_EQ(result[0]->Dims(), (std::vector{2, 2, 2, 2})); - test::ExpectTensorFloatEqual(result[0], expected); -} - -// 1x1 kernels reduce conv to a per-position channel mix. -TEST_P(AutogradConvCudaForwardTest, ConvForwardKernelOne) { - ONLY_CUDA(); - const std::vector input_values{1.0f, 2.0f, 3.0f, 4.0f, 5.0f, 6.0f, 7.0f, 8.0f, 9.0f, - 10.0f, 20.0f, 30.0f, 40.0f, 50.0f, 60.0f, 70.0f, 80.0f, 90.0f}; - auto input = std::make_shared(input_values.data(), std::vector{1, 2, 3, 3}, DataType::kFLOAT32, - GetDevice()); - const std::vector weight_values{2.0f, 3.0f, 1.0f, -1.0f}; - auto weight = std::make_shared(weight_values.data(), std::vector{2, 2, 1, 1}, DataType::kFLOAT32, - GetDevice()); - const std::vector expected{32.0f, 64.0f, 96.0f, 128.0f, 160.0f, 192.0f, 224.0f, 256.0f, 288.0f, - -9.0f, -18.0f, -27.0f, -36.0f, -45.0f, -54.0f, -63.0f, -72.0f, -81.0f}; - - auto conv_fn = std::make_shared(); - auto result = conv_fn->Apply({input, weight}); - ASSERT_EQ(result.size(), 1); - EXPECT_EQ(result[0]->Dims(), (std::vector{1, 2, 3, 3})); - test::ExpectTensorFloatEqual(result[0], expected); -} - -// H = W = kernel collapses the output to 1x1. -TEST_P(AutogradConvCudaForwardTest, ConvForwardKernelEqualsSpatial) { - ONLY_CUDA(); - std::vector input_values; - for (int idx = 0; idx < 9; ++idx) { input_values.push_back(static_cast(idx)); } - for (int idx = 0; idx < 9; ++idx) { input_values.push_back(2.0f * idx); } - auto input = std::make_shared(input_values.data(), std::vector{2, 1, 3, 3}, DataType::kFLOAT32, - GetDevice()); - auto weight = std::make_shared(std::vector{1, 1, 3, 3}, DataType::kFLOAT32, GetDevice(), true); - weight->Fill(1.0f); - auto bias = std::make_shared(std::vector{1}, DataType::kFLOAT32, GetDevice(), true); - bias->Fill(4.0f); - - auto conv_fn = std::make_shared(); - auto result = conv_fn->Apply({input, weight, bias}); - ASSERT_EQ(result.size(), 1); - EXPECT_EQ(result[0]->Dims(), (std::vector{2, 1, 1, 1})); - test::ExpectTensorFloatEqual(result[0], std::vector{40.0f, 76.0f}); -} - -// Magnitudes near the fp32 integer-exact range guard against overflow on accumulation. -TEST_P(AutogradConvCudaForwardTest, ConvForwardExtremeValues) { - ONLY_CUDA(); - auto input = std::make_shared(std::vector{1, 1, 4, 4}, DataType::kFLOAT32, GetDevice(), true); - input->Fill(static_cast(1 << 30)); - auto weight = std::make_shared(std::vector{1, 1, 3, 3}, DataType::kFLOAT32, GetDevice(), true); - weight->Fill(4.0f); - - auto conv_fn = std::make_shared(); - auto result = conv_fn->Apply({input, weight}); - ASSERT_EQ(result.size(), 1); - EXPECT_EQ(result[0]->Dims(), (std::vector{1, 1, 2, 2})); - // 9 taps * 2^30 * 4 = 9 * 2^32, exactly representable in fp32. - test::ExpectTensorFloatEqual(result[0], static_cast(9.0 * (1ULL << 32))); -} - -// Random 4D case checked against PyTorch conv2d (fixed seed, see file-level comment). -TEST_P(AutogradConvCudaForwardTest, ConvForwardTorchGolden) { - ONLY_CUDA(); - const std::vector input_values{ - 2.014464855f, -0.671311080f, -0.945236862f, -0.096128546f, 0.889558852f, 1.726294398f, -0.078931913f, - 0.205890238f, 0.094608001f, 0.173616216f, 0.437049866f, -0.557070732f, 0.455583423f, -0.736482263f, - -0.718647420f, 0.926798999f, 1.732108712f, -0.138735890f, -1.508729339f, 2.127460241f, 0.079739936f, - 0.217489868f, 0.589455068f, -0.052687794f, -1.659490585f, -0.380473137f, -0.166608080f, 0.350353301f, - -1.411424637f, 0.626776755f, 1.336345553f, -0.203511983f, -0.757243156f, 1.381434083f, -0.417290032f, - -0.668979943f, -0.569557011f, -1.035447836f, -0.555948615f, -1.279154539f, 0.056448594f, -0.388204783f, - 2.071142673f, -0.803012371f, 1.956272721f, 0.610922813f, -1.668144941f, -0.628582716f, -0.756579638f, - 0.987797141f, -0.849777400f, -1.654165864f, -0.326211363f, -0.040599853f, 0.281366080f, 0.499955416f, - -0.518889010f, -0.409607649f, 0.689743340f, 1.129151702f, 0.798748374f, -1.004696369f, -0.104189672f, - 1.257721663f, - }; - auto input = std::make_shared(input_values.data(), std::vector{2, 2, 4, 4}, DataType::kFLOAT32, - GetDevice()); - const std::vector weight_values{ - -2.448281050f, -0.608355999f, -0.958948731f, 1.009868145f, 0.031469870f, -1.101765275f, 0.571945250f, - -1.224039674f, -0.426795214f, -0.612229824f, -0.744571507f, -1.164491534f, -0.726353824f, -1.124407530f, - -0.498482078f, -0.423681349f, -1.231436968f, -0.689355612f, -0.170030892f, -0.120427191f, 1.032825112f, - 1.369398594f, 1.237134099f, 0.458183825f, -0.396747291f, 0.548917472f, 0.043999992f, 0.891736746f, - -1.292207360f, 1.713299513f, 0.982937992f, 0.432634443f, 0.293460459f, 0.318505853f, -0.146791920f, - 0.983720183f, -0.162205622f, 0.616937041f, 0.929942310f, -0.411319345f, 0.301661789f, -1.444623590f, - -0.730446517f, 2.391211271f, 0.400584310f, -0.909495890f, -0.702420592f, 0.372547686f, 0.994188488f, - 0.768918931f, 0.740486383f, -1.176032424f, 0.868664801f, 0.422459871f, - }; - auto weight = std::make_shared(weight_values.data(), std::vector{3, 2, 3, 3}, DataType::kFLOAT32, - GetDevice()); - const std::vector bias_values{-0.640517414f, 0.982376575f, 0.729396820f}; - auto bias = std::make_shared(bias_values.data(), std::vector{3}, DataType::kFLOAT32, GetDevice()); - const std::vector expected{ - -2.073040247f, 1.047590375f, -2.980665922f, -6.143202305f, 1.986891150f, 9.790361404f, - 0.198000669f, 0.174701050f, 0.093229778f, 2.889255762f, -0.725508392f, 0.107455865f, - 1.509838820f, -5.231677532f, 2.129449368f, 4.556076050f, -5.035542011f, -1.076976299f, - 0.719715834f, 1.997749925f, 1.977930188f, 7.832613468f, -5.848879814f, -0.770412445f, - }; - - auto conv_fn = std::make_shared(); - auto result = conv_fn->Apply({input, weight, bias}); - ASSERT_EQ(result.size(), 1); - EXPECT_EQ(result[0]->Dims(), (std::vector{2, 3, 2, 2})); - test::ExpectTensorNear(result[0], expected, kGoldenAbsError); -} - -// An empty batch is a degenerate but valid call: the kernel must not launch a zero-block -// configuration; the output tensor keeps the (N=0, ...) shape and is returned untouched. -TEST_P(AutogradConvCudaForwardTest, ConvForwardEmptyBatch) { - ONLY_CUDA(); - auto input = std::make_shared(std::vector{0, 2, 4, 4}, DataType::kFLOAT32, GetDevice()); - auto weight = std::make_shared(std::vector{3, 2, 3, 3}, DataType::kFLOAT32, GetDevice(), true); - weight->Fill(1.0f); - - auto conv_fn = std::make_shared(); - auto result = conv_fn->Apply({input, weight}); - ASSERT_EQ(result.size(), 1); - EXPECT_EQ(result[0]->Dims(), (std::vector{0, 3, 2, 2})); - EXPECT_EQ(result[0]->NumElements(), 0); -} - -INFINI_TRAIN_REGISTER_TEST(AutogradConvCudaForwardTest); diff --git a/tests/autograd/test_autograd_conv_forward.cc b/tests/autograd/test_autograd_conv_forward.cc index 26c8ea8d..85d45800 100644 --- a/tests/autograd/test_autograd_conv_forward.cc +++ b/tests/autograd/test_autograd_conv_forward.cc @@ -5,7 +5,6 @@ #include "gtest/gtest.h" #include "infini_train/include/autograd/conv.h" -#include "infini_train/include/nn/parallel/global.h" #include "infini_train/include/tensor.h" #include "tests/common/test_utils.h" @@ -23,7 +22,6 @@ class AutogradConvForwardTest : public infini_train::test::InfiniTrainTest {}; // An asymmetric kernel pins the cross-correlation semantics: flipping the kernel would swap // the roles of the taps and change every output value. TEST_P(AutogradConvForwardTest, ConvForwardAsymmetricKernel) { - ONLY_CPU(); std::vector input_values; for (int idx = 0; idx < 16; ++idx) { input_values.push_back(static_cast(idx)); } auto input = std::make_shared(input_values.data(), std::vector{1, 1, 4, 4}, DataType::kFLOAT32, @@ -42,7 +40,6 @@ TEST_P(AutogradConvForwardTest, ConvForwardAsymmetricKernel) { } TEST_P(AutogradConvForwardTest, ConvForwardMultiChannelBatchBias) { - ONLY_CPU(); // input(n, c, h, w) = 10*n + 5*c + 3*h + w with H = W = 3. std::vector input_values; for (int n = 0; n < 2; ++n) { @@ -73,7 +70,6 @@ TEST_P(AutogradConvForwardTest, ConvForwardMultiChannelBatchBias) { } TEST_P(AutogradConvForwardTest, ConvForwardKernelOne) { - ONLY_CPU(); const std::vector input_values{1.0f, 2.0f, 3.0f, 4.0f, 5.0f, 6.0f, 7.0f, 8.0f, 9.0f, 10.0f, 20.0f, 30.0f, 40.0f, 50.0f, 60.0f, 70.0f, 80.0f, 90.0f}; auto input = std::make_shared(input_values.data(), std::vector{1, 2, 3, 3}, DataType::kFLOAT32, @@ -94,7 +90,6 @@ TEST_P(AutogradConvForwardTest, ConvForwardKernelOne) { // H = W = kernel collapses the output to 1x1: the only patch covers the whole image. TEST_P(AutogradConvForwardTest, ConvForwardKernelEqualsSpatial) { - ONLY_CPU(); std::vector input_values; for (int idx = 0; idx < 9; ++idx) { input_values.push_back(static_cast(idx)); } for (int idx = 0; idx < 9; ++idx) { input_values.push_back(2.0f * idx); } @@ -116,7 +111,6 @@ TEST_P(AutogradConvForwardTest, ConvForwardKernelEqualsSpatial) { // Magnitudes near the fp32 integer-exact range: every product and partial sum stays exact, so // this guards against spurious overflow / precision loss on the accumulation path. TEST_P(AutogradConvForwardTest, ConvForwardExtremeValues) { - ONLY_CPU(); auto input = std::make_shared(std::vector{1, 1, 4, 4}, DataType::kFLOAT32, GetDevice(), true); input->Fill(static_cast(1 << 30)); auto weight = std::make_shared(std::vector{1, 1, 3, 3}, DataType::kFLOAT32, GetDevice(), true); @@ -137,7 +131,6 @@ TEST_P(AutogradConvForwardTest, ConvForwardExtremeValues) { // Random 4D case checked against PyTorch conv2d (fixed seed, see file-level comment). TEST_P(AutogradConvForwardTest, ConvForwardTorchGolden) { - ONLY_CPU(); const std::vector input_values{ 2.014464855f, -0.671311080f, -0.945236862f, -0.096128546f, 0.889558852f, 1.726294398f, -0.078931913f, 0.205890238f, 0.094608001f, 0.173616216f, 0.437049866f, -0.557070732f, 0.455583423f, -0.736482263f, @@ -181,7 +174,6 @@ TEST_P(AutogradConvForwardTest, ConvForwardTorchGolden) { } TEST_P(AutogradConvForwardTest, ConvForwardRejectsInvalidShapes) { - ONLY_CPU(); EXPECT_DEATH( { auto input = std::make_shared(std::vector{1, 2, 3, 3}, DataType::kFLOAT32, GetDevice()); @@ -203,4 +195,18 @@ TEST_P(AutogradConvForwardTest, ConvForwardRejectsInvalidShapes) { ""); } +// An empty batch is a degenerate but valid call: the kernel must not launch a zero-block +// configuration; the output tensor keeps the (N=0, ...) shape and is returned untouched. +TEST_P(AutogradConvForwardTest, ConvForwardEmptyBatch) { + auto input = std::make_shared(std::vector{0, 2, 4, 4}, DataType::kFLOAT32, GetDevice()); + auto weight = std::make_shared(std::vector{3, 2, 3, 3}, DataType::kFLOAT32, GetDevice(), true); + weight->Fill(1.0f); + + auto conv_fn = std::make_shared(); + auto result = conv_fn->Apply({input, weight}); + ASSERT_EQ(result.size(), 1); + EXPECT_EQ(result[0]->Dims(), (std::vector{0, 3, 2, 2})); + EXPECT_EQ(result[0]->NumElements(), 0); +} + INFINI_TRAIN_REGISTER_TEST(AutogradConvForwardTest); From d69875189a3d4d1c217ed2470c721b819d32dd83 Mon Sep 17 00:00:00 2001 From: accelerator-llc Date: Wed, 9 Sep 2026 11:27:17 +0800 Subject: [PATCH 12/15] perf(conv): batch the conv GEMMs and drop redundant col buffer init Forward and backward-input dispatched one cuBLAS call per image with batch_count=1, while the Gemm interface already exposes a strided-batched path that Matmul uses. Both now issue a single strided-batched GEMM with the weight shared via stride 0, and backward-input runs one col2im launch for the whole batch. The CPU kernels zero-initialized their im2col scratch right before overwriting every element; allocate it uninitialized instead. Single-GPU MNIST CNN throughput: 16582 -> 32802 samples/s per epoch. --- infini_train/src/kernels/cpu/conv.cc | 29 ++++--- infini_train/src/kernels/cuda/conv.cu | 110 ++++++++++++++------------ 2 files changed, 79 insertions(+), 60 deletions(-) diff --git a/infini_train/src/kernels/cpu/conv.cc b/infini_train/src/kernels/cpu/conv.cc index bd5c8312..1fa8f634 100644 --- a/infini_train/src/kernels/cpu/conv.cc +++ b/infini_train/src/kernels/cpu/conv.cc @@ -102,16 +102,19 @@ std::shared_ptr Conv2dForward(const std::shared_ptr &input, cons float *output_data = static_cast(output->DataPtr()); // (batch, patches, C*kH*kW) scratch: the im2col expansion of every image in the batch. + // Left uninitialized: im2col overwrites every element before the first read. const int64_t image_size = channels * height * width; - std::vector col_buffer(batch * patches * flat_kernel); + auto col = std::make_shared(std::vector{batch, patches, flat_kernel}, DataType::kFLOAT32, + input->GetDevice()); + float *col_buffer = static_cast(col->DataPtr()); for (int64_t n = 0; n < batch; ++n) { Im2col(input_data + n * image_size, channels, height, width, kernel_h, kernel_w, out_height, out_width, - col_buffer.data() + n * patches * flat_kernel); + col_buffer + n * patches * flat_kernel); } Eigen::Map weight_mat(weight_data, out_channels, flat_kernel); for (int64_t n = 0; n < batch; ++n) { - Eigen::Map col_n(col_buffer.data() + n * patches * flat_kernel, patches, flat_kernel); + Eigen::Map col_n(col_buffer + n * patches * flat_kernel, patches, flat_kernel); Eigen::Map out_n(output_data + n * out_channels * patches, out_channels, patches); out_n.noalias() = weight_mat * col_n.transpose(); if (bias) { @@ -170,12 +173,16 @@ std::shared_ptr Conv2dBackwardInput(const std::shared_ptr &weigh Eigen::Map weight_mat(weight_data, out_channels, flat_kernel); const int64_t image_size = channels * height * width; - std::vector col_buffer(patches * flat_kernel); + // (patches, C*kH*kW) scratch reused across images. Left uninitialized: the GEMM below + // assigns every element before Col2im reads it. + auto col = std::make_shared(std::vector{patches, flat_kernel}, DataType::kFLOAT32, + grad_output->GetDevice()); + float *col_buffer = static_cast(col->DataPtr()); for (int64_t n = 0; n < batch; ++n) { Eigen::Map grad_output_n(grad_output_data + n * out_channels * patches, out_channels, patches); - Eigen::Map col_n(col_buffer.data(), patches, flat_kernel); + Eigen::Map col_n(col_buffer, patches, flat_kernel); col_n.noalias() = grad_output_n.transpose() * weight_mat; - Col2im(col_buffer.data(), channels, height, width, kernel_h, kernel_w, out_height, out_width, + Col2im(col_buffer, channels, height, width, kernel_h, kernel_w, out_height, out_width, grad_input_data + n * image_size); } @@ -224,16 +231,20 @@ std::shared_ptr Conv2dBackwardWeight(const std::shared_ptr &inpu float *grad_weight_data = static_cast(grad_weight->DataPtr()); const int64_t image_size = channels * height * width; - std::vector col_buffer(batch * patches * flat_kernel); + // (batch, patches, C*kH*kW) scratch: the im2col expansion of every image in the batch. + // Left uninitialized: im2col overwrites every element before the first read. + auto col = std::make_shared(std::vector{batch, patches, flat_kernel}, DataType::kFLOAT32, + input->GetDevice()); + float *col_buffer = static_cast(col->DataPtr()); for (int64_t n = 0; n < batch; ++n) { Im2col(input_data + n * image_size, channels, height, width, kernel_h, kernel_w, out_height, out_width, - col_buffer.data() + n * patches * flat_kernel); + col_buffer + n * patches * flat_kernel); } Eigen::Map grad_weight_mat(grad_weight_data, out_channels, flat_kernel); for (int64_t n = 0; n < batch; ++n) { Eigen::Map grad_output_n(grad_output_data + n * out_channels * patches, out_channels, patches); - Eigen::Map col_n(col_buffer.data() + n * patches * flat_kernel, patches, flat_kernel); + Eigen::Map col_n(col_buffer + n * patches * flat_kernel, patches, flat_kernel); grad_weight_mat.noalias() += grad_output_n * col_n; } diff --git a/infini_train/src/kernels/cuda/conv.cu b/infini_train/src/kernels/cuda/conv.cu index 10499847..5141b0d8 100644 --- a/infini_train/src/kernels/cuda/conv.cu +++ b/infini_train/src/kernels/cuda/conv.cu @@ -45,7 +45,7 @@ __global__ void Im2colKernel(const float *__restrict__ input, float *__restrict_ // col2im: inverse of Im2col for the input gradient. An input pixel is covered by // multiple output windows, so all patches that touch it are gathered. Each thread -// owns one (c, x, y) output cell and writes it once, so no atomics are needed. +// owns one (n, c, x, y) output cell and writes it once, so no atomics are needed. __global__ void Col2imKernel(const float *__restrict__ col, float *__restrict__ grad_input, int64_t total, int64_t channels, int64_t height, int64_t width, int64_t kernel_h, int64_t kernel_w, int64_t out_height, int64_t out_width, int64_t patches, int64_t flat_kernel) { @@ -191,27 +191,31 @@ std::shared_ptr Conv2dForward(const std::shared_ptr &input, cons const float *weight_data = static_cast(weight->DataPtr()); float *output_data = static_cast(output->DataPtr()); const float beta = bias ? 1.0f : 0.0f; - for (int64_t n = 0; n < batch; ++n) { - Dispatcher::Instance().Call({device.type(), "Gemm"}, device, - GemmParams{ - .trans_a = GemmTranspose::kNoTranspose, - .trans_b = GemmTranspose::kNoTranspose, - .m = static_cast(patches), - .n = static_cast(out_channels), - .k = static_cast(flat_kernel), - .A = col_data + n * patches * flat_kernel, - .lda = static_cast(patches), - .B = weight_data, - .ldb = static_cast(flat_kernel), - .C = output_data + n * out_channels * patches, - .ldc = static_cast(patches), - .alpha = 1.0f, - .beta = beta, - .batch_count = 1, - .input_dtype = DataType::kFLOAT32, - .output_dtype = DataType::kFLOAT32, - }); - } + // One strided-batched GEMM over the batch: A steps through the per-image col blocks, + // B (the weight) is shared by every image via stride 0, C steps through the output. + // batch == 1 falls back to the non-batched Gemm path, which expects zero strides. + Dispatcher::Instance().Call({device.type(), "Gemm"}, device, + GemmParams{ + .trans_a = GemmTranspose::kNoTranspose, + .trans_b = GemmTranspose::kNoTranspose, + .m = static_cast(patches), + .n = static_cast(out_channels), + .k = static_cast(flat_kernel), + .A = col_data, + .lda = static_cast(patches), + .B = weight_data, + .ldb = static_cast(flat_kernel), + .C = output_data, + .ldc = static_cast(patches), + .alpha = 1.0f, + .beta = beta, + .batch_count = static_cast(batch), + .stride_a = batch > 1 ? patches * flat_kernel : 0, + .stride_b = 0, + .stride_c = batch > 1 ? out_channels * patches : 0, + .input_dtype = DataType::kFLOAT32, + .output_dtype = DataType::kFLOAT32, + }); return output; } @@ -258,41 +262,45 @@ std::shared_ptr Conv2dBackwardInput(const std::shared_ptr &weigh auto device = grad_output->GetDevice(); const cudaStream_t stream = CurrentCudaStream(device); - // Per-image scratch: grad_output(n)^T * weight, stored column-major (patches, flat_kernel). - auto col = std::make_shared(std::vector{flat_kernel, patches}, DataType::kFLOAT32, device); + // Whole-batch scratch: col(n) = grad_output(n)^T * weight per image, stored column-major + // (patches, flat_kernel), so the col2im below can gather every image in a single launch. + auto col = std::make_shared(std::vector{batch, flat_kernel, patches}, DataType::kFLOAT32, device); const float *weight_data = static_cast(weight->DataPtr()); const float *grad_output_data = static_cast(grad_output->DataPtr()); float *grad_input_data = static_cast(grad_input->DataPtr()); float *col_data = static_cast(col->DataPtr()); - for (int64_t n = 0; n < batch; ++n) { - // col = grad_output(n)^T * weight, the transpose of the forward GEMM. - Dispatcher::Instance().Call({device.type(), "Gemm"}, device, - GemmParams{ - .trans_a = GemmTranspose::kNoTranspose, - .trans_b = GemmTranspose::kTranspose, - .m = static_cast(patches), - .n = static_cast(flat_kernel), - .k = static_cast(out_channels), - .A = grad_output_data + n * out_channels * patches, - .lda = static_cast(patches), - .B = weight_data, - .ldb = static_cast(flat_kernel), - .C = col_data, - .ldc = static_cast(patches), - .alpha = 1.0f, - .beta = 0.0f, - .batch_count = 1, - .input_dtype = DataType::kFLOAT32, - .output_dtype = DataType::kFLOAT32, - }); - - Col2imKernel<<>>( - col_data, grad_input_data + n * channels * height * width, channels * height * width, channels, height, - width, kernel_h, kernel_w, out_height, out_width, patches, flat_kernel); - CUDA_CHECK(cudaGetLastError()); - } + // The transpose of the forward GEMM for every image at once: A steps through grad_output, + // B (the weight) is shared by every image via stride 0, C steps through the col buffer. + // batch == 1 falls back to the non-batched Gemm path, which expects zero strides. + Dispatcher::Instance().Call({device.type(), "Gemm"}, device, + GemmParams{ + .trans_a = GemmTranspose::kNoTranspose, + .trans_b = GemmTranspose::kTranspose, + .m = static_cast(patches), + .n = static_cast(flat_kernel), + .k = static_cast(out_channels), + .A = grad_output_data, + .lda = static_cast(patches), + .B = weight_data, + .ldb = static_cast(flat_kernel), + .C = col_data, + .ldc = static_cast(patches), + .alpha = 1.0f, + .beta = 0.0f, + .batch_count = static_cast(batch), + .stride_a = batch > 1 ? out_channels * patches : 0, + .stride_b = 0, + .stride_c = batch > 1 ? flat_kernel * patches : 0, + .input_dtype = DataType::kFLOAT32, + .output_dtype = DataType::kFLOAT32, + }); + + Col2imKernel<<>>( + col_data, grad_input_data, batch * channels * height * width, channels, height, width, kernel_h, kernel_w, + out_height, out_width, patches, flat_kernel); + CUDA_CHECK(cudaGetLastError()); return grad_input; } From 23c4263ca19435de0fa799c8bdf91376967c9d88 Mon Sep 17 00:00:00 2001 From: accelerator-llc Date: Wed, 9 Sep 2026 15:19:12 +0800 Subject: [PATCH 13/15] test(conv): merge the device-split conv train tests The train pair duplicated the same end-to-end SGD step per device, leaving a *_cuda_* file behind after the forward/backward merge. Fold them into one parameterized body that inspects loss, gradients and parameters through host copies, with parameter snapshots taken into freshly allocated host buffers so they survive the in-place update. --- .../autograd/test_autograd_conv_cuda_train.cc | 89 ------------------- tests/autograd/test_autograd_conv_train.cc | 24 +++-- 2 files changed, 15 insertions(+), 98 deletions(-) delete mode 100644 tests/autograd/test_autograd_conv_cuda_train.cc diff --git a/tests/autograd/test_autograd_conv_cuda_train.cc b/tests/autograd/test_autograd_conv_cuda_train.cc deleted file mode 100644 index 39c0777b..00000000 --- a/tests/autograd/test_autograd_conv_cuda_train.cc +++ /dev/null @@ -1,89 +0,0 @@ -#include -#include -#include -#include - -#include "gtest/gtest.h" - -#include "infini_train/include/nn/modules/conv.h" -#include "infini_train/include/nn/modules/linear.h" -#include "infini_train/include/nn/modules/loss.h" -#include "infini_train/include/optimizer.h" -#include "infini_train/include/tensor.h" - -#include "tests/common/test_utils.h" - -using namespace infini_train; - -// Minimal end-to-end training step on CUDA: Conv2d -> Flatten -> Linear -> CrossEntropy, -// one loss->Backward() through the autograd graph and one SGD step. -class AutogradConvCudaTrainTest : public infini_train::test::InfiniTrainTest {}; - -TEST_P(AutogradConvCudaTrainTest, ConvTrainStepUpdatesParams) { - ONLY_CUDA(); - const Device device = GetDevice(); - const Device host = Device(); - constexpr float kLearningRate = 0.1f; - - std::vector input_values; - for (int idx = 0; idx < 2 * 1 * 8 * 8; ++idx) { input_values.push_back((idx * 37) % 101 / 50.0f - 1.0f); } - auto input - = std::make_shared(input_values.data(), std::vector{2, 1, 8, 8}, DataType::kFLOAT32, device); - auto target = std::make_shared(std::vector{2}, DataType::kINT64, device); - target->Fill(0); - - auto conv = std::make_shared(1, 2, 3, true, device); - auto fc = std::make_shared(2 * 6 * 6, 3, true, device); - auto loss_fn = std::make_shared(); - - auto conv_out = (*conv)({input})[0]; - ASSERT_EQ(conv_out->Dims(), (std::vector{2, 2, 6, 6})); - auto flat = conv_out->Flatten(1); - ASSERT_EQ(flat->Dims(), (std::vector{2, 72})); - auto logits = (*fc)({flat})[0]; - ASSERT_EQ(logits->Dims(), (std::vector{2, 3})); - auto loss = (*loss_fn)({logits, target})[0]; - ASSERT_TRUE(loss->Dims().empty()); - const auto loss_cpu = loss->To(host); - EXPECT_TRUE(std::isfinite(*static_cast(loss_cpu.DataPtr()))); - - loss->Backward(); - - std::vector> params = conv->Parameters(); - for (auto ¶m : fc->Parameters()) { params.push_back(param); } - ASSERT_EQ(params.size(), 4); - - // Every parameter must have received a finite, non-trivial gradient through the graph. - for (const auto ¶m : params) { - const auto &grad = param->grad(); - ASSERT_NE(grad, nullptr); - ASSERT_EQ(grad->Dims(), param->Dims()); - const auto grad_cpu = grad->To(host); - const auto *grad_data = static_cast(grad_cpu.DataPtr()); - float max_abs = 0.0f; - for (size_t idx = 0; idx < grad_cpu.NumElements(); ++idx) { - ASSERT_TRUE(std::isfinite(grad_data[idx])); - max_abs = std::max(max_abs, std::fabs(grad_data[idx])); - } - EXPECT_GT(max_abs, 0.0f); - } - - std::vector> old_values_cpu; - for (const auto ¶m : params) { old_values_cpu.push_back(std::make_shared(param->To(host))); } - - auto optimizer = std::make_shared(params, kLearningRate); - optimizer->Step(); - - for (size_t p = 0; p < params.size(); ++p) { - const auto new_cpu = params[p]->To(host); - const auto grad_cpu = params[p]->grad()->To(host); - const auto *old_data = static_cast(old_values_cpu[p]->DataPtr()); - const auto *grad_data = static_cast(grad_cpu.DataPtr()); - const auto *new_data = static_cast(new_cpu.DataPtr()); - for (size_t idx = 0; idx < params[p]->NumElements(); ++idx) { - EXPECT_FLOAT_EQ(new_data[idx], old_data[idx] - kLearningRate * grad_data[idx]) << "param " << p; - } - } -} - -INFINI_TRAIN_REGISTER_TEST(AutogradConvCudaTrainTest); diff --git a/tests/autograd/test_autograd_conv_train.cc b/tests/autograd/test_autograd_conv_train.cc index 613c0a78..d2bb6999 100644 --- a/tests/autograd/test_autograd_conv_train.cc +++ b/tests/autograd/test_autograd_conv_train.cc @@ -8,7 +8,6 @@ #include "infini_train/include/nn/modules/conv.h" #include "infini_train/include/nn/modules/linear.h" #include "infini_train/include/nn/modules/loss.h" -#include "infini_train/include/nn/parallel/global.h" #include "infini_train/include/optimizer.h" #include "infini_train/include/tensor.h" @@ -17,12 +16,13 @@ using namespace infini_train; // Minimal end-to-end training step: Conv2d -> Flatten -> Linear -> CrossEntropy, one -// loss->Backward() through the autograd graph and one SGD step, on CPU. +// loss->Backward() through the autograd graph and one SGD step. Loss, gradients and +// parameters are inspected through host copies, so one body serves both devices. class AutogradConvTrainTest : public infini_train::test::InfiniTrainTest {}; TEST_P(AutogradConvTrainTest, ConvTrainStepUpdatesParams) { - ONLY_CPU(); const Device device = GetDevice(); + const Device host = Device(); constexpr float kLearningRate = 0.1f; std::vector input_values; @@ -44,7 +44,8 @@ TEST_P(AutogradConvTrainTest, ConvTrainStepUpdatesParams) { ASSERT_EQ(logits->Dims(), (std::vector{2, 3})); auto loss = (*loss_fn)({logits, target})[0]; ASSERT_TRUE(loss->Dims().empty()); - EXPECT_TRUE(std::isfinite(*static_cast(loss->DataPtr()))); + const auto loss_cpu = loss->To(host); + EXPECT_TRUE(std::isfinite(*static_cast(loss_cpu.DataPtr()))); loss->Backward(); @@ -57,18 +58,21 @@ TEST_P(AutogradConvTrainTest, ConvTrainStepUpdatesParams) { const auto &grad = param->grad(); ASSERT_NE(grad, nullptr); ASSERT_EQ(grad->Dims(), param->Dims()); - const auto *grad_data = static_cast(grad->DataPtr()); + const auto grad_cpu = grad->To(host); + const auto *grad_data = static_cast(grad_cpu.DataPtr()); float max_abs = 0.0f; - for (size_t idx = 0; idx < grad->NumElements(); ++idx) { + for (size_t idx = 0; idx < grad_cpu.NumElements(); ++idx) { ASSERT_TRUE(std::isfinite(grad_data[idx])); max_abs = std::max(max_abs, std::fabs(grad_data[idx])); } EXPECT_GT(max_abs, 0.0f); } + // Snapshot into freshly allocated host buffers: a same-device To(host) view would alias + // the live storage and silently track the optimizer's in-place update. std::vector> old_values; for (const auto ¶m : params) { - auto old_value = std::make_shared(param->Dims(), param->Dtype(), device); + auto old_value = std::make_shared(param->Dims(), param->Dtype(), host); old_value->CopyFrom(*param); old_values.push_back(old_value); } @@ -77,9 +81,11 @@ TEST_P(AutogradConvTrainTest, ConvTrainStepUpdatesParams) { optimizer->Step(); for (size_t p = 0; p < params.size(); ++p) { + const auto new_cpu = params[p]->To(host); + const auto grad_cpu = params[p]->grad()->To(host); const auto *old_data = static_cast(old_values[p]->DataPtr()); - const auto *grad_data = static_cast(params[p]->grad()->DataPtr()); - const auto *new_data = static_cast(params[p]->DataPtr()); + const auto *grad_data = static_cast(grad_cpu.DataPtr()); + const auto *new_data = static_cast(new_cpu.DataPtr()); for (size_t idx = 0; idx < params[p]->NumElements(); ++idx) { EXPECT_FLOAT_EQ(new_data[idx], old_data[idx] - kLearningRate * grad_data[idx]) << "param " << p; } From badbb9ffa81415388316cfabe110c11aef596ea6 Mon Sep 17 00:00:00 2001 From: accelerator-llc Date: Wed, 9 Sep 2026 15:19:12 +0800 Subject: [PATCH 14/15] chore(kernels): fix include hygiene in conv and relu Drop the unused Gemm declaration header from the CUDA conv kernel, add the missing to the CPU ReLU kernel, and remove the copy-and- keep-in-sync note above the CUDA ReLU block-size helper. --- infini_train/src/kernels/cpu/relu.cc | 1 + infini_train/src/kernels/cuda/conv.cu | 1 - infini_train/src/kernels/cuda/relu.cu | 3 --- 3 files changed, 1 insertion(+), 4 deletions(-) diff --git a/infini_train/src/kernels/cpu/relu.cc b/infini_train/src/kernels/cpu/relu.cc index 7d4086a0..2a6bddbb 100644 --- a/infini_train/src/kernels/cpu/relu.cc +++ b/infini_train/src/kernels/cpu/relu.cc @@ -1,3 +1,4 @@ +#include #include #include "glog/logging.h" diff --git a/infini_train/src/kernels/cuda/conv.cu b/infini_train/src/kernels/cuda/conv.cu index 5141b0d8..19143609 100644 --- a/infini_train/src/kernels/cuda/conv.cu +++ b/infini_train/src/kernels/cuda/conv.cu @@ -13,7 +13,6 @@ #include "infini_train/src/core/runtime/cuda/cuda_runtime_common.h" #include "infini_train/src/kernels/common/gemm.h" -#include "infini_train/src/kernels/cuda/common/gemm.cuh" namespace infini_train::kernels::cuda { namespace { diff --git a/infini_train/src/kernels/cuda/relu.cu b/infini_train/src/kernels/cuda/relu.cu index 7471edbf..3f33d4f3 100644 --- a/infini_train/src/kernels/cuda/relu.cu +++ b/infini_train/src/kernels/cuda/relu.cu @@ -30,9 +30,6 @@ __global__ void ReLUBackwardKernel(T *grad_input, const T *input, const T *grad_ } } -// Mirrors the unary launch pattern in kernels/cuda/elementwise.cu: same block-size -// heuristic, stream lookup and offset-loop emission. That file's helpers live in an -// anonymous namespace and cannot be shared from here; keep this in sync with it. inline size_t ChooseBlockSize(size_t num_elements) { if (num_elements < 1024) { return 64; From 1c11b69bcf83e733085b677e1141c5f79e918117 Mon Sep 17 00:00:00 2001 From: accelerator-llc Date: Sat, 12 Sep 2026 14:17:01 +0800 Subject: [PATCH 15/15] feat(mnist): report test metrics every epoch and log the training step Run the full test set after each epoch instead of only once at the end of training, so a single run shows the test loss falling and the test accuracy rising. The evaluation runs under no_grad: a forward-only graph leaves each parameter's grad accumulator with a dependency count that the next backward pass cannot satisfy, silently stopping its gradient accumulation. The training progress line now carries an explicit 1-based [step k/total] field next to the samples consumed so far. --- example/mnist/main.cc | 82 ++++++++++++++++++++++++++----------------- 1 file changed, 50 insertions(+), 32 deletions(-) diff --git a/example/mnist/main.cc b/example/mnist/main.cc index ef3c3063..69011d7c 100644 --- a/example/mnist/main.cc +++ b/example/mnist/main.cc @@ -1,3 +1,4 @@ +#include #include #include #include @@ -10,6 +11,7 @@ #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" @@ -128,6 +130,48 @@ int main(int argc, char *argv[]) { 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 test_losses; + int correct = 0; + int total = 0; + for (const auto &[image, label] : test_dataloader) { + auto new_image = std::make_shared(image->To(device)); + auto new_label = std::make_shared(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); + + const int batch_size = output_cpu.Dims()[0]; + for (int batch_idx = 0; batch_idx < batch_size; ++batch_idx) { + auto label_index = reinterpret_cast(label_cpu.DataPtr())[batch_idx]; + const auto *output_values = static_cast(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(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(correct) / total << ", AverageLoss: " << avg_loss; + }; + for (int epoch = 0; epoch < FLAGS_num_epoch; ++epoch) { int train_idx = 0; float total_loss = 0.0; @@ -162,9 +206,11 @@ int main(int argc, char *argv[]) { } total_loss += current_loss; if (train_idx % kNumItersOfOutputDuration == 0) { - LOG(ERROR) << "epoch: " << epoch << ", [" << train_idx * FLAGS_bs * ddp_world_size << "/" - << 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(); @@ -177,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 test_losses; - int correct = 0; - int total = 0; - for (const auto &[image, label] : test_dataloader) { - auto new_image = std::make_shared(image->To(device)); - auto new_label = std::make_shared(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); - - const int batch_size = output_cpu.Dims()[0]; - for (int batch_idx = 0; batch_idx < batch_size; ++batch_idx) { - auto label_index = reinterpret_cast(label_cpu.DataPtr())[batch_idx]; - const auto *output_values = static_cast(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(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(correct) / total << ", AverageLoss: " << avg_loss; gflags::ShutDownCommandLineFlags(); google::ShutdownGoogleLogging();