From 76e6c9824cb470dc83da8c61d581cf26ca350bba Mon Sep 17 00:00:00 2001 From: gavin-richie <2904793019@qq.com> Date: Mon, 7 Sep 2026 06:37:39 +0800 Subject: [PATCH 01/14] feat: add ReLU activation operator and module Add the Relu autograd function with forward/backward kernels for CPU and CUDA (fp32/bf16 via the elementwise dispatch helpers), a functional wrapper, and an nn::ReLU module, mirroring the existing Sigmoid operator. --- infini_train/include/autograd/activations.h | 12 ++++++ infini_train/include/nn/functional.h | 11 +++++ infini_train/include/nn/modules/activations.h | 7 ++++ infini_train/src/autograd/activations.cc | 25 +++++++++++ infini_train/src/kernels/cpu/relu.cc | 41 +++++++++++++++++++ infini_train/src/kernels/cuda/elementwise.cu | 16 ++++++++ infini_train/src/nn/functional.cc | 4 ++ infini_train/src/nn/modules/activations.cc | 4 ++ 8 files changed, 120 insertions(+) create mode 100644 infini_train/src/kernels/cpu/relu.cc diff --git a/infini_train/include/autograd/activations.h b/infini_train/include/autograd/activations.h index a63977263..caec954e3 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/functional.h b/infini_train/include/nn/functional.h index e4354fd10..82b64dc7a 100644 --- a/infini_train/include/nn/functional.h +++ b/infini_train/include/nn/functional.h @@ -137,6 +137,17 @@ std::shared_ptr Max(const std::shared_ptr &input, int64_t dim, b // A tensor containing sigmoid applied element-wise to the input. std::shared_ptr Sigmoid(const std::shared_ptr &input); +// Applies the rectified linear unit function element-wise. +// +// ReLU(x) = max(0, x). +// +// Args: +// input: The input tensor. +// +// Returns: +// A tensor containing relu applied element-wise to the input. +std::shared_ptr Relu(const std::shared_ptr &input); + // Applies the softmax function along the specified dimension. // // The softmax function maps input values to the range [0, 1] and ensures they sum to 1. diff --git a/infini_train/include/nn/modules/activations.h b/infini_train/include/nn/modules/activations.h index deb029576..663b90368 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 bb8b8e5ea..2d8398a69 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> &, + const std::vector> &output_tensors) { + const auto &output = output_tensors[0]; + ctx_.SaveForBackward({output}); +} + +std::vector> Relu::Backward(const std::vector> &grad_outputs) { + auto saved_tensors = ctx_.GetSavedTensors(); + CHECK_EQ(saved_tensors.size(), 1); + const auto &output = saved_tensors[0]; + CHECK_EQ(grad_outputs.size(), 1); + const auto &grad_output = grad_outputs[0]; + + auto device = output->GetDevice().type(); + return {Dispatcher::Instance().Call>({device, "ReluBackward"}, output, 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 000000000..3fb874df8 --- /dev/null +++ b/infini_train/src/kernels/cpu/relu.cc @@ -0,0 +1,41 @@ +#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) { + 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(); + for (int64_t idx = 0; idx < numel; ++idx) { output_ptr[idx] = input_ptr[idx] > 0.0f ? input_ptr[idx] : 0.0f; } + + return output; +} + +std::shared_ptr ReluBackward(const std::shared_ptr &output, + const std::shared_ptr &grad_output) { + auto grad_input = std::make_shared(output->Dims(), DataType::kFLOAT32); + const float *output_ptr = static_cast(output->DataPtr()); + const float *grad_output_ptr = static_cast(grad_output->DataPtr()); + float *grad_input_ptr = static_cast(grad_input->DataPtr()); + + const int64_t numel = output->NumElements(); + for (int64_t idx = 0; idx < numel; ++idx) { + grad_input_ptr[idx] = output_ptr[idx] > 0.0f ? grad_output_ptr[idx] : 0.0f; + } + 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/elementwise.cu b/infini_train/src/kernels/cuda/elementwise.cu index fc423b35f..9fc73491c 100644 --- a/infini_train/src/kernels/cuda/elementwise.cu +++ b/infini_train/src/kernels/cuda/elementwise.cu @@ -1213,6 +1213,20 @@ std::shared_ptr SigmoidBackward(const std::shared_ptr &output, return UnaryBackward(grad_output, output, [] __device__(auto x) { return Mul(x, Sub(decltype(x){1}, x)); }); , INFINI_ALL_FLOATING_TYPES) } + +std::shared_ptr ReluForward(const std::shared_ptr &input) { + DISPATCH(input->Dtype(), return UnaryForward(input, [] __device__(auto x) { return Max(x, decltype(x){0}); }); + , INFINI_ALL_FLOATING_TYPES) +} + +std::shared_ptr ReluBackward(const std::shared_ptr &output, + const std::shared_ptr &grad_output) { + DISPATCH(grad_output->Dtype(), return UnaryBackward(grad_output, output, + [] __device__(auto x) { + return x > decltype(x){0} ? decltype(x){1} : decltype(x){0}; + }); + , INFINI_ALL_FLOATING_TYPES) +} } // namespace infini_train::kernels::cuda #define REGISTER_CUDA_ELEMENTWISE_KERNEL(kernel_name) \ @@ -1262,5 +1276,7 @@ REGISTER_CUDA_ELEMENTWISE_KERNEL(DivForward) REGISTER_CUDA_ELEMENTWISE_KERNEL(DivBackward) REGISTER_CUDA_ELEMENTWISE_KERNEL(SigmoidForward) REGISTER_CUDA_ELEMENTWISE_KERNEL(SigmoidBackward) +REGISTER_CUDA_ELEMENTWISE_KERNEL(ReluForward) +REGISTER_CUDA_ELEMENTWISE_KERNEL(ReluBackward) #undef REGISTER_CUDA_ELEMENTWISE_KERNEL diff --git a/infini_train/src/nn/functional.cc b/infini_train/src/nn/functional.cc index c33e23684..01cc79917 100644 --- a/infini_train/src/nn/functional.cc +++ b/infini_train/src/nn/functional.cc @@ -78,4 +78,8 @@ std::shared_ptr Softmax(const std::shared_ptr &input, int64_t di std::shared_ptr Sigmoid(const std::shared_ptr &input) { return std::make_shared()->Apply({input})[0]; } + +std::shared_ptr Relu(const std::shared_ptr &input) { + return std::make_shared()->Apply({input})[0]; +} } // namespace infini_train::nn::function diff --git a/infini_train/src/nn/modules/activations.cc b/infini_train/src/nn/modules/activations.cc index d1bbc9da8..e7596c2c9 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 From 0e768109500d584ae31f6a4ec26198080df15ee6 Mon Sep 17 00:00:00 2001 From: gavin-richie <2904793019@qq.com> Date: Mon, 7 Sep 2026 06:37:50 +0800 Subject: [PATCH 02/14] feat: add Conv2d operator with im2col and GEMM kernels Add the Conv2d autograd function supporting symmetric stride and padding with input, weight, and bias gradients, plus an nn::Conv2d module using the same default initialization as torch.nn.Conv2d (kaiming_uniform with a=sqrt(5)). CPU: im2col unfolding with row-major Eigen GEMMs and a col2im scatter for grad_input. CUDA: an im2col kernel, strided-batched cuBLAS GEMMs through the existing Gemm wrapper (weight broadcast via stride_b=0), an atomic col2im kernel, and per-channel bias reductions. Kernels cover the fp32 path required by the MNIST CNN training use case and reject other dtypes explicitly. --- infini_train/include/autograd/conv.h | 31 ++ infini_train/include/nn/modules/conv.h | 35 +++ infini_train/src/autograd/conv.cc | 72 +++++ infini_train/src/kernels/cpu/conv2d.cc | 255 ++++++++++++++++ infini_train/src/kernels/cuda/conv2d.cu | 380 ++++++++++++++++++++++++ infini_train/src/nn/modules/conv.cc | 53 ++++ 6 files changed, 826 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/conv2d.cc create mode 100644 infini_train/src/kernels/cuda/conv2d.cu create mode 100644 infini_train/src/nn/modules/conv.cc diff --git a/infini_train/include/autograd/conv.h b/infini_train/include/autograd/conv.h new file mode 100644 index 000000000..97bd6a677 --- /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(int64_t stride, int64_t padding) : Function(kType), stride_(stride), padding_(padding) {} + + 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: + int64_t stride_ = 1; + int64_t padding_ = 0; + std::vector input_dims_; + std::vector weight_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 000000000..deff4b9a1 --- /dev/null +++ b/infini_train/include/nn/modules/conv.h @@ -0,0 +1,35 @@ +#pragma once + +#include +#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, int64_t stride = 1, int64_t padding = 0, + bool bias = true, Device device = Device()); + std::vector> Forward(const std::vector> &input_tensors) override; + + bool has_bias() const { return bias_; } + +private: + void ResetParameters(); + int64_t stride_ = 1; + int64_t padding_ = 0; + 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 000000000..8a4294121 --- /dev/null +++ b/infini_train/src/autograd/conv.cc @@ -0,0 +1,72 @@ +#include "infini_train/include/autograd/conv.h" + +#include +#include +#include + +#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, + stride_, padding_)}; +} + +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}); + input_dims_ = input->Dims(); + weight_dims_ = weight->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 = 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) { + CHECK_NE(weight, nullptr) << "Conv2d::Backward requires the saved weight to compute grad_input"; + grad_input = Dispatcher::Instance().Call>({device, "Conv2dBackwardInput"}, weight, + grad_output, stride_, padding_, input_dims_); + } + if (need_grad_weight) { + CHECK_NE(input, nullptr) << "Conv2d::Backward requires the saved input to compute grad_weight"; + grad_weight = Dispatcher::Instance().Call>( + {device, "Conv2dBackwardWeight"}, input, grad_output, stride_, padding_, weight_dims_); + } + if (need_grad_bias) { + grad_bias = Dispatcher::Instance().Call>({device, "Conv2dBackwardBias"}, grad_output); + } + + return {grad_input, grad_weight, grad_bias}; +} +} // namespace infini_train::autograd diff --git a/infini_train/src/kernels/cpu/conv2d.cc b/infini_train/src/kernels/cpu/conv2d.cc new file mode 100644 index 000000000..5dc3170e4 --- /dev/null +++ b/infini_train/src/kernels/cpu/conv2d.cc @@ -0,0 +1,255 @@ +#include +#include +#include +#include + +#include + +#include "glog/logging.h" + +#include "infini_train/include/dispatcher.h" +#include "infini_train/include/tensor.h" + +namespace infini_train::kernels::cpu { +namespace { +using RowMajorMatrix = Eigen::Matrix; +using RowMajorMatrixMap = Eigen::Map; +using ConstRowMajorMatrixMap = Eigen::Map; + +// Conv kernels only support fp32, which covers the MNIST CNN training use case. +void CheckFloat32(const std::shared_ptr &tensor, const char *name) { + CHECK_EQ(static_cast(tensor->Dtype()), static_cast(DataType::kFLOAT32)) + << "Conv2d kernel expects fp32 " << name; + CHECK_EQ(static_cast(tensor->GetDevice().type()), static_cast(Device::DeviceType::kCPU)) + << "Conv2d cpu kernel expects a cpu tensor: " << name; +} + +int64_t ConvOutSize(int64_t in_size, int64_t kernel_size, int64_t stride, int64_t padding) { + CHECK_GT(stride, 0); + CHECK_GE(padding, 0); + const int64_t out_size = (in_size + 2 * padding - kernel_size) / stride + 1; + CHECK_GT(out_size, 0) << "Non-positive convolution output size"; + return out_size; +} + +struct Conv2dDims { + int64_t N; + int64_t C_in; + int64_t H; + int64_t W; + int64_t C_out; + int64_t kH; + int64_t kW; + int64_t H_out; + int64_t W_out; +}; + +Conv2dDims ResolveConv2dDims(const std::vector &input_dims, const std::vector &weight_dims, + int64_t stride, int64_t padding) { + CHECK_EQ(input_dims.size(), 4); + CHECK_EQ(weight_dims.size(), 4); + Conv2dDims d{}; + d.N = input_dims[0]; + d.C_in = input_dims[1]; + d.H = input_dims[2]; + d.W = input_dims[3]; + d.C_out = weight_dims[0]; + CHECK_EQ(d.C_in, weight_dims[1]) << "Conv2d input channel mismatch with weight"; + d.kH = weight_dims[2]; + d.kW = weight_dims[3]; + d.H_out = ConvOutSize(d.H, d.kH, stride, padding); + d.W_out = ConvOutSize(d.W, d.kW, stride, padding); + return d; +} + +// Unfolds input (C_in, H, W) into col (C_in*kH*kW, H_out*W_out); padded regions contribute zeros. +// col row (c*kH+kh)*kW+kw holds the input values under the kernel tap (kh, kw) for every output position. +void Im2Col(const float *input, int64_t C_in, int64_t H, int64_t W, int64_t kH, int64_t kW, int64_t stride, + int64_t padding, int64_t H_out, int64_t W_out, float *col) { + const int64_t P = H_out * W_out; + for (int64_t c = 0; c < C_in; ++c) { + for (int64_t kh = 0; kh < kH; ++kh) { + for (int64_t kw = 0; kw < kW; ++kw) { + float *col_row = col + ((c * kH + kh) * kW + kw) * P; + for (int64_t oh = 0; oh < H_out; ++oh) { + const int64_t ih = oh * stride + kh - padding; + for (int64_t ow = 0; ow < W_out; ++ow) { + const int64_t iw = ow * stride + kw - padding; + const bool in_bounds = ih >= 0 && ih < H && iw >= 0 && iw < W; + col_row[oh * W_out + ow] = in_bounds ? input[(c * H + ih) * W + iw] : 0.0f; + } + } + } + } + } +} + +// Scatter-adds grad_col (C_in*kH*kW, H_out*W_out) into grad_input (C_in, H, W) at the source positions. +void Col2ImAccumulate(const float *grad_col, int64_t C_in, int64_t H, int64_t W, int64_t kH, int64_t kW, int64_t stride, + int64_t padding, int64_t H_out, int64_t W_out, float *grad_input) { + const int64_t P = H_out * W_out; + for (int64_t c = 0; c < C_in; ++c) { + for (int64_t kh = 0; kh < kH; ++kh) { + for (int64_t kw = 0; kw < kW; ++kw) { + const float *grad_col_row = grad_col + ((c * kH + kh) * kW + kw) * P; + for (int64_t oh = 0; oh < H_out; ++oh) { + const int64_t ih = oh * stride + kh - padding; + for (int64_t ow = 0; ow < W_out; ++ow) { + const int64_t iw = ow * stride + kw - padding; + if (ih >= 0 && ih < H && iw >= 0 && iw < W) { + grad_input[(c * H + ih) * W + iw] += grad_col_row[oh * W_out + ow]; + } + } + } + } + } + } +} +} // namespace + +std::shared_ptr Conv2dForward(const std::shared_ptr &input, const std::shared_ptr &weight, + const std::shared_ptr &bias, int64_t stride, int64_t padding) { + /* + input: (N, C_in, H, W), weight: (C_out, C_in, kH, kW), bias: (C_out) + output: (N, C_out, H_out, W_out), where H_out = (H + 2*padding - kH) / stride + 1 + */ + CheckFloat32(input, "input"); + CheckFloat32(weight, "weight"); + if (bias) { + CheckFloat32(bias, "bias"); + CHECK_EQ(bias->Dims().size(), 1); + CHECK_EQ(bias->Dims()[0], weight->Dims()[0]); + } + + const Conv2dDims d = ResolveConv2dDims(input->Dims(), weight->Dims(), stride, padding); + const int64_t K = d.C_in * d.kH * d.kW; + const int64_t P = d.H_out * d.W_out; + + auto output = std::make_shared(std::vector{d.N, d.C_out, d.H_out, d.W_out}, DataType::kFLOAT32); + const float *input_ptr = static_cast(input->DataPtr()); + const float *weight_ptr = static_cast(weight->DataPtr()); + float *output_ptr = static_cast(output->DataPtr()); + + const ConstRowMajorMatrixMap weight_mat(weight_ptr, d.C_out, K); + RowMajorMatrix col(K, P); + for (int64_t n = 0; n < d.N; ++n) { + Im2Col(input_ptr + n * d.C_in * d.H * d.W, d.C_in, d.H, d.W, d.kH, d.kW, stride, padding, d.H_out, d.W_out, + col.data()); + RowMajorMatrixMap out_mat(output_ptr + n * d.C_out * P, d.C_out, P); + out_mat = weight_mat * col; + if (bias) { + const Eigen::Map bias_vec(static_cast(bias->DataPtr()), d.C_out); + out_mat.colwise() += bias_vec; + } + } + return output; +} + +std::shared_ptr Conv2dBackwardInput(const std::shared_ptr &weight, + const std::shared_ptr &grad_output, int64_t stride, int64_t padding, + const std::vector &input_dims) { + /* + grad_input = col2im(weight^T * grad_output) + weight: (C_out, C_in, kH, kW), grad_output: (N, C_out, H_out, W_out), grad_input: (N, C_in, H, W) + */ + CheckFloat32(weight, "weight"); + CheckFloat32(grad_output, "grad_output"); + + const Conv2dDims d = ResolveConv2dDims(input_dims, weight->Dims(), stride, padding); + const int64_t P = d.H_out * d.W_out; + const auto &grad_output_dims = grad_output->Dims(); + CHECK_EQ(grad_output_dims.size(), 4); + CHECK_EQ(grad_output_dims[0], d.N); + CHECK_EQ(grad_output_dims[1], d.C_out); + CHECK_EQ(grad_output_dims[2], d.H_out); + CHECK_EQ(grad_output_dims[3], d.W_out); + + auto grad_input = std::make_shared(input_dims, DataType::kFLOAT32); + float *grad_input_ptr = static_cast(grad_input->DataPtr()); + std::fill_n(grad_input_ptr, grad_input->NumElements(), 0.0f); + + const int64_t K = d.C_in * d.kH * d.kW; + const ConstRowMajorMatrixMap weight_mat(static_cast(weight->DataPtr()), d.C_out, K); + RowMajorMatrix grad_col(K, P); + for (int64_t n = 0; n < d.N; ++n) { + const ConstRowMajorMatrixMap grad_out_mat(static_cast(grad_output->DataPtr()) + n * d.C_out * P, + d.C_out, P); + grad_col = weight_mat.transpose() * grad_out_mat; + Col2ImAccumulate(grad_col.data(), d.C_in, d.H, d.W, d.kH, d.kW, stride, padding, d.H_out, d.W_out, + grad_input_ptr + n * d.C_in * d.H * d.W); + } + return grad_input; +} + +std::shared_ptr Conv2dBackwardWeight(const std::shared_ptr &input, + const std::shared_ptr &grad_output, int64_t stride, + int64_t padding, const std::vector &weight_dims) { + /* + grad_weight = sum_n grad_output[n] * col(input[n])^T + input: (N, C_in, H, W), grad_output: (N, C_out, H_out, W_out), grad_weight: (C_out, C_in, kH, kW) + */ + CheckFloat32(input, "input"); + CheckFloat32(grad_output, "grad_output"); + CHECK_EQ(weight_dims.size(), 4); + + const Conv2dDims d = ResolveConv2dDims(input->Dims(), weight_dims, stride, padding); + const auto &grad_output_dims = grad_output->Dims(); + CHECK_EQ(grad_output_dims.size(), 4); + CHECK_EQ(grad_output_dims[0], d.N); + CHECK_EQ(grad_output_dims[1], d.C_out); + CHECK_EQ(grad_output_dims[2], d.H_out); + CHECK_EQ(grad_output_dims[3], d.W_out); + + auto grad_weight = std::make_shared(weight_dims, DataType::kFLOAT32); + float *grad_weight_ptr = static_cast(grad_weight->DataPtr()); + std::fill_n(grad_weight_ptr, grad_weight->NumElements(), 0.0f); + + const int64_t K = d.C_in * d.kH * d.kW; + const int64_t P = d.H_out * d.W_out; + RowMajorMatrixMap grad_weight_mat(grad_weight_ptr, d.C_out, K); + RowMajorMatrix col(K, P); + for (int64_t n = 0; n < d.N; ++n) { + Im2Col(static_cast(input->DataPtr()) + n * d.C_in * d.H * d.W, d.C_in, d.H, d.W, d.kH, d.kW, + stride, padding, d.H_out, d.W_out, col.data()); + const ConstRowMajorMatrixMap grad_out_mat(static_cast(grad_output->DataPtr()) + n * d.C_out * P, + d.C_out, P); + grad_weight_mat.noalias() += grad_out_mat * col.transpose(); + } + return grad_weight; +} + +std::shared_ptr Conv2dBackwardBias(const std::shared_ptr &grad_output) { + /* + grad_bias = sum over batch and spatial positions of grad_output + grad_output: (N, C_out, H_out, W_out), grad_bias: (C_out) + */ + CheckFloat32(grad_output, "grad_output"); + const auto &grad_output_dims = grad_output->Dims(); + CHECK_EQ(grad_output_dims.size(), 4); + const int64_t N = grad_output_dims[0]; + const int64_t C_out = grad_output_dims[1]; + const int64_t P = grad_output_dims[2] * grad_output_dims[3]; + + auto grad_bias = std::make_shared(std::vector{C_out}, DataType::kFLOAT32); + float *grad_bias_ptr = static_cast(grad_bias->DataPtr()); + std::fill_n(grad_bias_ptr, C_out, 0.0f); + Eigen::Map> grad_bias_vec(grad_bias_ptr, C_out); + + for (int64_t n = 0; n < N; ++n) { + const ConstRowMajorMatrixMap grad_out_mat(static_cast(grad_output->DataPtr()) + n * C_out * P, + C_out, P); + grad_bias_vec += grad_out_mat.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/kernels/cuda/conv2d.cu b/infini_train/src/kernels/cuda/conv2d.cu new file mode 100644 index 000000000..e8a38f3ca --- /dev/null +++ b/infini_train/src/kernels/cuda/conv2d.cu @@ -0,0 +1,380 @@ +#include +#include +#include + +#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 { + +constexpr int kThreadsPerBlock = 256; + +// Conv kernels only support fp32, which covers the MNIST CNN training use case. +void CheckFloat32(const std::shared_ptr &tensor, const char *name) { + CHECK_EQ(static_cast(tensor->Dtype()), static_cast(DataType::kFLOAT32)) + << "Conv2d kernel expects fp32 " << name; + CHECK_EQ(static_cast(tensor->GetDevice().type()), static_cast(Device::DeviceType::kCUDA)) + << "Conv2d cuda kernel expects a cuda tensor: " << name; +} + +int64_t ConvOutSize(int64_t in_size, int64_t kernel_size, int64_t stride, int64_t padding) { + CHECK_GT(stride, 0); + CHECK_GE(padding, 0); + const int64_t out_size = (in_size + 2 * padding - kernel_size) / stride + 1; + CHECK_GT(out_size, 0) << "Non-positive convolution output size"; + return out_size; +} + +struct Conv2dDims { + int64_t N; + int64_t C_in; + int64_t H; + int64_t W; + int64_t C_out; + int64_t kH; + int64_t kW; + int64_t H_out; + int64_t W_out; +}; + +Conv2dDims ResolveConv2dDims(const std::vector &input_dims, const std::vector &weight_dims, + int64_t stride, int64_t padding) { + CHECK_EQ(input_dims.size(), 4); + CHECK_EQ(weight_dims.size(), 4); + Conv2dDims d{}; + d.N = input_dims[0]; + d.C_in = input_dims[1]; + d.H = input_dims[2]; + d.W = input_dims[3]; + d.C_out = weight_dims[0]; + CHECK_EQ(d.C_in, weight_dims[1]) << "Conv2d input channel mismatch with weight"; + d.kH = weight_dims[2]; + d.kW = weight_dims[3]; + d.H_out = ConvOutSize(d.H, d.kH, stride, padding); + d.W_out = ConvOutSize(d.W, d.kW, stride, padding); + return d; +} + +cudaStream_t GetCudaStream(const Device &device) { + return dynamic_cast( + infini_train::core::GetDeviceGuardImpl(device.type())->GetStream(device)) + ->cuda_stream(); +} + +int NumBlocks(int64_t total) { return static_cast((total + kThreadsPerBlock - 1) / kThreadsPerBlock); } + +// cuBLAS ignores strides for batchCount==1, but the Gemm wrapper expects them zeroed in that case. +void CallConvGemm(const Device &device, GemmParams params) { + if (params.batch_count == 1) { + params.stride_a = 0; + params.stride_b = 0; + params.stride_c = 0; + } + Gemm(device, params); +} + +// One thread per (n, k, p): gathers input[n] (C_in, H, W) into col[n] (K, P); padded regions contribute zeros. +// col row k = (c*kH+kh)*kW+kw holds the input values under the kernel tap (kh, kw) for every output position. +__global__ void Im2ColKernel(const float *input, float *col, int64_t total, int64_t K, int64_t P, int64_t C_in, + int64_t H, int64_t W, int64_t kH, int64_t kW, int64_t stride, int64_t padding, + int64_t H_out, int64_t W_out) { + const int64_t idx = static_cast(blockIdx.x) * blockDim.x + threadIdx.x; + if (idx >= total) { + return; + } + const int64_t p = idx % P; + const int64_t k = (idx / P) % K; + const int64_t n = idx / (P * K); + const int64_t c = k / (kH * kW); + const int64_t kh = (k / kW) % kH; + const int64_t kw = k % kW; + const int64_t oh = p / W_out; + const int64_t ow = p % W_out; + const int64_t ih = oh * stride + kh - padding; + const int64_t iw = ow * stride + kw - padding; + const bool in_bounds = ih >= 0 && ih < H && iw >= 0 && iw < W; + col[idx] = in_bounds ? input[(n * C_in + c) * H * W + ih * W + iw] : 0.0f; +} + +// Inverse of Im2ColKernel: scatter-adds grad_col[n] (K, P) into grad_input[n] (C_in, H, W). +__global__ void Col2ImKernel(const float *grad_col, float *grad_input, int64_t total, int64_t K, int64_t P, + int64_t C_in, int64_t H, int64_t W, int64_t kH, int64_t kW, int64_t stride, + int64_t padding, int64_t H_out, int64_t W_out) { + const int64_t idx = static_cast(blockIdx.x) * blockDim.x + threadIdx.x; + if (idx >= total) { + return; + } + const int64_t p = idx % P; + const int64_t k = (idx / P) % K; + const int64_t n = idx / (P * K); + const int64_t c = k / (kH * kW); + const int64_t kh = (k / kW) % kH; + const int64_t kw = k % kW; + const int64_t oh = p / W_out; + const int64_t ow = p % W_out; + const int64_t ih = oh * stride + kh - padding; + const int64_t iw = ow * stride + kw - padding; + if (ih >= 0 && ih < H && iw >= 0 && iw < W) { + atomicAdd(&grad_input[(n * C_in + c) * H * W + ih * W + iw], grad_col[idx]); + } +} + +// output (N, C_out, P): adds bias[c] to every spatial position of channel c. +__global__ void BiasAddKernel(float *output, const float *bias, int64_t total, int64_t C_out, int64_t P) { + const int64_t idx = static_cast(blockIdx.x) * blockDim.x + threadIdx.x; + if (idx >= total) { + return; + } + const int64_t c = (idx / P) % C_out; + output[idx] += bias[c]; +} + +// grad_bias (C_out): one block per output channel, sums over batch and spatial positions of +// grad_output (N, C_out, P). +__global__ void BiasBackwardKernel(const float *grad_output, float *grad_bias, int64_t N, int64_t C_out, int64_t P) { + const int64_t c = blockIdx.x; + if (c >= C_out) { + return; + } + float sum = 0.0f; + for (int64_t i = threadIdx.x; i < N * P; i += blockDim.x) { + sum += grad_output[(i / P) * C_out * P + c * P + (i % P)]; + } + atomicAdd(&grad_bias[c], sum); +} + +// out(idx) = sum over the leading batch dimension of batched(b, idx). +__global__ void BatchSumKernel(const float *batched, float *out, int64_t batch, int64_t num_elements) { + const int64_t idx = static_cast(blockIdx.x) * blockDim.x + threadIdx.x; + if (idx >= num_elements) { + return; + } + float sum = 0.0f; + for (int64_t b = 0; b < batch; ++b) { sum += batched[b * num_elements + idx]; } + out[idx] = sum; +} +} // namespace + +std::shared_ptr Conv2dForward(const std::shared_ptr &input, const std::shared_ptr &weight, + const std::shared_ptr &bias, int64_t stride, int64_t padding) { + /* + input: (N, C_in, H, W), weight: (C_out, C_in, kH, kW), bias: (C_out) + output: (N, C_out, H_out, W_out), where H_out = (H + 2*padding - kH) / stride + 1 + */ + CheckFloat32(input, "input"); + CheckFloat32(weight, "weight"); + if (bias) { + CheckFloat32(bias, "bias"); + CHECK_EQ(bias->Dims().size(), 1); + CHECK_EQ(bias->Dims()[0], weight->Dims()[0]); + } + + const Conv2dDims d = ResolveConv2dDims(input->Dims(), weight->Dims(), stride, padding); + const int64_t K = d.C_in * d.kH * d.kW; + const int64_t P = d.H_out * d.W_out; + + auto device = input->GetDevice(); + const auto cuda_stream = GetCudaStream(device); + + auto output + = std::make_shared(std::vector{d.N, d.C_out, d.H_out, d.W_out}, DataType::kFLOAT32, device); + auto col = std::make_shared(std::vector{d.N * K * P}, DataType::kFLOAT32, device); + + const int64_t im2col_total = d.N * K * P; + Im2ColKernel<<>>( + static_cast(input->DataPtr()), static_cast(col->DataPtr()), im2col_total, K, P, d.C_in, + d.H, d.W, d.kH, d.kW, stride, padding, d.H_out, d.W_out); + + // out_n (C_out, P) = weight (C_out, K) * col_n (K, P) per sample; cuBLAS is column-major, so a row-major GEMM + // C(M, N) = L(M, K) * R(K, N) maps to op(A)=R, op(B)=L with m=N, n=M, k=K (see linear.cu). + CallConvGemm(device, GemmParams{ + .trans_a = GemmTranspose::kNoTranspose, + .trans_b = GemmTranspose::kNoTranspose, + .m = static_cast(P), + .n = static_cast(d.C_out), + .k = static_cast(K), + .A = col->DataPtr(), + .lda = static_cast(P), + .B = weight->DataPtr(), + .ldb = static_cast(K), + .C = output->DataPtr(), + .ldc = static_cast(P), + .alpha = 1.0f, + .beta = 0.0f, + .batch_count = static_cast(d.N), + .stride_a = K * P, + .stride_b = 0, + .stride_c = d.C_out * P, + .input_dtype = DataType::kFLOAT32, + .output_dtype = DataType::kFLOAT32, + }); + + if (bias) { + const int64_t total = d.N * d.C_out * P; + BiasAddKernel<<>>( + static_cast(output->DataPtr()), static_cast(bias->DataPtr()), total, d.C_out, P); + } + return output; +} + +std::shared_ptr Conv2dBackwardInput(const std::shared_ptr &weight, + const std::shared_ptr &grad_output, int64_t stride, int64_t padding, + const std::vector &input_dims) { + /* + grad_input = col2im(weight^T * grad_output) + weight: (C_out, C_in, kH, kW), grad_output: (N, C_out, H_out, W_out), grad_input: (N, C_in, H, W) + */ + CheckFloat32(weight, "weight"); + CheckFloat32(grad_output, "grad_output"); + + const Conv2dDims d = ResolveConv2dDims(input_dims, weight->Dims(), stride, padding); + const auto &grad_output_dims = grad_output->Dims(); + CHECK_EQ(grad_output_dims.size(), 4); + CHECK_EQ(grad_output_dims[0], d.N); + CHECK_EQ(grad_output_dims[1], d.C_out); + CHECK_EQ(grad_output_dims[2], d.H_out); + CHECK_EQ(grad_output_dims[3], d.W_out); + + const int64_t K = d.C_in * d.kH * d.kW; + const int64_t P = d.H_out * d.W_out; + + auto device = grad_output->GetDevice(); + const auto cuda_stream = GetCudaStream(device); + + auto grad_input = std::make_shared(input_dims, DataType::kFLOAT32, device); + CUDA_CHECK(cudaMemsetAsync(grad_input->DataPtr(), 0, grad_input->SizeInBytes(), cuda_stream)); + auto grad_col = std::make_shared(std::vector{d.N * K * P}, DataType::kFLOAT32, device); + + // grad_col_n (K, P) = weight^T (K, C_out) * grad_out_n (C_out, P): op(B) transposes the weight. + CallConvGemm(device, GemmParams{ + .trans_a = GemmTranspose::kNoTranspose, + .trans_b = GemmTranspose::kTranspose, + .m = static_cast(P), + .n = static_cast(K), + .k = static_cast(d.C_out), + .A = grad_output->DataPtr(), + .lda = static_cast(P), + .B = weight->DataPtr(), + .ldb = static_cast(K), + .C = grad_col->DataPtr(), + .ldc = static_cast(P), + .alpha = 1.0f, + .beta = 0.0f, + .batch_count = static_cast(d.N), + .stride_a = d.C_out * P, + .stride_b = 0, + .stride_c = K * P, + .input_dtype = DataType::kFLOAT32, + .output_dtype = DataType::kFLOAT32, + }); + + const int64_t col2im_total = d.N * K * P; + Col2ImKernel<<>>( + static_cast(grad_col->DataPtr()), static_cast(grad_input->DataPtr()), col2im_total, K, + P, d.C_in, d.H, d.W, d.kH, d.kW, stride, padding, d.H_out, d.W_out); + return grad_input; +} + +std::shared_ptr Conv2dBackwardWeight(const std::shared_ptr &input, + const std::shared_ptr &grad_output, int64_t stride, + int64_t padding, const std::vector &weight_dims) { + /* + grad_weight = sum_n grad_out_n (C_out, P) * col(input[n])^T (P, K) + input: (N, C_in, H, W), grad_output: (N, C_out, H_out, W_out), grad_weight: (C_out, C_in, kH, kW) + */ + CheckFloat32(input, "input"); + CheckFloat32(grad_output, "grad_output"); + CHECK_EQ(weight_dims.size(), 4); + + const Conv2dDims d = ResolveConv2dDims(input->Dims(), weight_dims, stride, padding); + const auto &grad_output_dims = grad_output->Dims(); + CHECK_EQ(grad_output_dims.size(), 4); + CHECK_EQ(grad_output_dims[0], d.N); + CHECK_EQ(grad_output_dims[1], d.C_out); + CHECK_EQ(grad_output_dims[2], d.H_out); + CHECK_EQ(grad_output_dims[3], d.W_out); + + const int64_t K = d.C_in * d.kH * d.kW; + const int64_t P = d.H_out * d.W_out; + + auto device = grad_output->GetDevice(); + const auto cuda_stream = GetCudaStream(device); + + auto grad_weight = std::make_shared(weight_dims, DataType::kFLOAT32, device); + auto grad_weight_batched + = std::make_shared(std::vector{d.N * d.C_out * K}, DataType::kFLOAT32, device); + auto col = std::make_shared(std::vector{d.N * K * P}, DataType::kFLOAT32, device); + + const int64_t im2col_total = d.N * K * P; + Im2ColKernel<<>>( + static_cast(input->DataPtr()), static_cast(col->DataPtr()), im2col_total, K, P, d.C_in, + d.H, d.W, d.kH, d.kW, stride, padding, d.H_out, d.W_out); + + // Per-sample grad_w_n (C_out, K) = grad_out_n (C_out, P) * col_n^T (P, K): op(A) transposes col. + CallConvGemm(device, GemmParams{ + .trans_a = GemmTranspose::kTranspose, + .trans_b = GemmTranspose::kNoTranspose, + .m = static_cast(K), + .n = static_cast(d.C_out), + .k = static_cast(P), + .A = col->DataPtr(), + .lda = static_cast(P), + .B = grad_output->DataPtr(), + .ldb = static_cast(P), + .C = grad_weight_batched->DataPtr(), + .ldc = static_cast(K), + .alpha = 1.0f, + .beta = 0.0f, + .batch_count = static_cast(d.N), + .stride_a = K * P, + .stride_b = d.C_out * P, + .stride_c = d.C_out * K, + .input_dtype = DataType::kFLOAT32, + .output_dtype = DataType::kFLOAT32, + }); + + BatchSumKernel<<>>( + static_cast(grad_weight_batched->DataPtr()), static_cast(grad_weight->DataPtr()), d.N, + d.C_out * K); + return grad_weight; +} + +std::shared_ptr Conv2dBackwardBias(const std::shared_ptr &grad_output) { + /* + grad_bias = sum over batch and spatial positions of grad_output + grad_output: (N, C_out, H_out, W_out), grad_bias: (C_out) + */ + CheckFloat32(grad_output, "grad_output"); + const auto &grad_output_dims = grad_output->Dims(); + CHECK_EQ(grad_output_dims.size(), 4); + const int64_t N = grad_output_dims[0]; + const int64_t C_out = grad_output_dims[1]; + const int64_t P = grad_output_dims[2] * grad_output_dims[3]; + + auto device = grad_output->GetDevice(); + const auto cuda_stream = GetCudaStream(device); + + auto grad_bias = std::make_shared(std::vector{C_out}, DataType::kFLOAT32, device); + CUDA_CHECK(cudaMemsetAsync(grad_bias->DataPtr(), 0, grad_bias->SizeInBytes(), cuda_stream)); + BiasBackwardKernel<<(C_out), kThreadsPerBlock, 0, cuda_stream>>>( + static_cast(grad_output->DataPtr()), static_cast(grad_bias->DataPtr()), N, C_out, P); + 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/infini_train/src/nn/modules/conv.cc b/infini_train/src/nn/modules/conv.cc new file mode 100644 index 000000000..6db8ba36c --- /dev/null +++ b/infini_train/src/nn/modules/conv.cc @@ -0,0 +1,53 @@ +#include "infini_train/include/nn/modules/conv.h" + +#include +#include +#include + +#include "glog/logging.h" + +#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, int64_t stride, int64_t padding, + bool bias, Device device) + : CloneableModule(kType), stride_(stride), padding_(padding), bias_(bias) { + CHECK_GT(in_channels, 0); + CHECK_GT(out_channels, 0); + CHECK_GT(kernel_size, 0); + CHECK_GT(stride, 0); + CHECK_GE(padding, 0); + 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(stride_, padding_) + ->Apply(bias_ ? std::vector>{input_tensors[0], parameters_[kParamWeightName], + parameters_[kParamBiasName]} + : std::vector>{input_tensors[0], parameters_[kParamWeightName]}); +} + +void Conv2d::ResetParameters() { + // Same scheme as torch.nn.Conv2d: kaiming_uniform_ on the weight with a=sqrt(5), then + // uniform_ on the bias with bound 1/sqrt(fan_in). + 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 From 5f4742f838f8b93569aba253d9e80198c9fe04f4 Mon Sep 17 00:00:00 2001 From: gavin-richie <2904793019@qq.com> Date: Mon, 7 Sep 2026 06:38:10 +0800 Subject: [PATCH 03/14] feat: add nn::Flatten module --- infini_train/include/nn/modules/flatten.h | 27 +++++++++++++++++++++++ infini_train/src/nn/modules/flatten.cc | 15 +++++++++++++ 2 files changed, 42 insertions(+) create mode 100644 infini_train/include/nn/modules/flatten.h create mode 100644 infini_train/src/nn/modules/flatten.cc diff --git a/infini_train/include/nn/modules/flatten.h b/infini_train/include/nn/modules/flatten.h new file mode 100644 index 000000000..9c64887e8 --- /dev/null +++ b/infini_train/include/nn/modules/flatten.h @@ -0,0 +1,27 @@ +#pragma once + +#include +#include +#include + +#include "infini_train/include/nn/modules/module.h" + +namespace infini_train { +class Tensor; +} + +namespace infini_train::nn { +class Flatten : public CloneableModule { +public: + static constexpr char kType[] = "Flatten"; + + Flatten(int64_t start_dim = 1, int64_t end_dim = -1) + : CloneableModule(kType), start_dim_(start_dim), end_dim_(end_dim) {} + + std::vector> Forward(const std::vector> &input_tensors) override; + +private: + int64_t start_dim_; + int64_t end_dim_; +}; +} // namespace infini_train::nn diff --git a/infini_train/src/nn/modules/flatten.cc b/infini_train/src/nn/modules/flatten.cc new file mode 100644 index 000000000..e1673ffab --- /dev/null +++ b/infini_train/src/nn/modules/flatten.cc @@ -0,0 +1,15 @@ +#include "infini_train/include/nn/modules/flatten.h" + +#include +#include + +#include "glog/logging.h" + +#include "infini_train/include/tensor.h" + +namespace infini_train::nn { +std::vector> Flatten::Forward(const std::vector> &input_tensors) { + CHECK_EQ(input_tensors.size(), 1); + return {input_tensors[0]->Flatten(start_dim_, end_dim_)}; +} +} // namespace infini_train::nn From 5f887c6531ecd88579722256c81a5271febfe442 Mon Sep 17 00:00:00 2001 From: gavin-richie <2904793019@qq.com> Date: Mon, 7 Sep 2026 06:38:10 +0800 Subject: [PATCH 04/14] fix: recompute MNIST image stride after float32 conversion image_size_in_bytes_ was computed in the member initializer list against the original uint8 IDX payload, but the constructor then replaces the image tensor with a float32 copy. Every sample view therefore advanced by 784 bytes instead of 3136, so images were misaligned with their labels and training silently stagnated around chance accuracy. Recompute the per-sample stride once the conversion to float32 is done. --- example/mnist/dataset.cc | 5 ++++- example/mnist/dataset.h | 2 +- 2 files changed, 5 insertions(+), 2 deletions(-) diff --git a/example/mnist/dataset.cc b/example/mnist/dataset.cc index ee683f6d4..ae6648dbc 100644 --- a/example/mnist/dataset.cc +++ b/example/mnist/dataset.cc @@ -103,13 +103,16 @@ MNISTDataset::MNISTDataset(const std::string &dataset, bool train) const auto bs = image_dims[0]; infini_train::Tensor transposed_tensor(image_dims, DataType::kFLOAT32); for (int idx = 0; idx < bs; ++idx) { - const auto *image_data = reinterpret_cast(image_file_.tensor.DataPtr()) + idx * 28 * 28; + const auto *image_data = reinterpret_cast(image_file_.tensor.DataPtr()) + idx * 28 * 28; auto *transposed_data = static_cast(transposed_tensor.DataPtr()) + idx * 28 * 28; for (int i = 0; i < 28; ++i) { for (int j = 0; j < 28; ++j) { transposed_data[i * 28 + j] = image_data[i * 28 + j] / 255.0f; } } } image_file_.tensor = std::move(transposed_tensor); + // The images are now float32; recompute the per-sample stride, which was computed against the + // original uint8 file in the member initializer list. + image_size_in_bytes_ = image_file_.tensor.SizeInBytes() / bs; } std::pair, std::shared_ptr> diff --git a/example/mnist/dataset.h b/example/mnist/dataset.h index 8d6193fb7..0c7dfd2fe 100644 --- a/example/mnist/dataset.h +++ b/example/mnist/dataset.h @@ -40,6 +40,6 @@ class MNISTDataset : public infini_train::Dataset { SN3PascalVincentFile label_file_; std::vector image_dims_; std::vector label_dims_; - const size_t image_size_in_bytes_ = 0; + size_t image_size_in_bytes_ = 0; const size_t label_size_in_bytes_ = 0; }; From 12aef4031f60377e11f49466d2871a4c51527f82 Mon Sep 17 00:00:00 2001 From: gavin-richie <2904793019@qq.com> Date: Mon, 7 Sep 2026 06:38:32 +0800 Subject: [PATCH 05/14] feat: add CNN training demo to the MNIST example Add a --model flag (mlp|cnn, defaulting to the existing MLP) with a CNN classifier following Conv2d(1,16,3) -> ReLU -> Conv2d(16,32,3) -> ReLU -> Flatten -> Linear(18432, 10). The evaluation loop now runs under no_grad, per-step losses are reported, and the final test loss and accuracy are logged explicitly. A new --init_weights flag loads an InfiniTrain checkpoint as initial weights, enabling PyTorch-alignment runs. --- example/mnist/main.cc | 77 +++++++++++++++++++++++++------------------ example/mnist/net.cc | 42 +++++++++++++++++++++++ example/mnist/net.h | 16 +++++++++ 3 files changed, 103 insertions(+), 32 deletions(-) diff --git a/example/mnist/main.cc b/example/mnist/main.cc index 7744e0947..95536050a 100644 --- a/example/mnist/main.cc +++ b/example/mnist/main.cc @@ -9,6 +9,8 @@ #include "gflags/gflags.h" #include "glog/logging.h" +#include "infini_train/include/autograd/grad_mode.h" +#include "infini_train/include/checkpoint/checkpoint.h" #include "infini_train/include/dataloader.h" #include "infini_train/include/device.h" #include "infini_train/include/nn/modules/loss.h" @@ -18,10 +20,12 @@ #include "example/mnist/net.h" DEFINE_string(dataset, "", "mnist dataset path"); +DEFINE_string(model, "mlp", "model type (mlp/cnn)"); DEFINE_int32(bs, 64, "batch size"); DEFINE_int32(num_epoch, 1, "num epochs"); DEFINE_double(lr, 0.01, "learning rate"); DEFINE_string(device, "cpu", "device type (cpu/cuda)"); +DEFINE_string(init_weights, "", "checkpoint dir to load initial weights from (e.g. exported by PyTorch)"); using namespace infini_train; @@ -36,6 +40,8 @@ constexpr char kDeviceCUDA[] = "cuda"; 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 == "mlp" || value == "cnn"; }); + int main(int argc, char *argv[]) { gflags::ParseCommandLineFlags(&argc, &argv, true); google::InitGoogleLogging(argv[0]); @@ -47,14 +53,19 @@ 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 = CreateMNISTNetwork(FLAGS_model); Device device = FLAGS_device == kDeviceCPU ? Device() : Device(Device::DeviceType::kCUDA, 0); Device cpu_device = Device(); - network.To(device); + network->To(device); + + if (!FLAGS_init_weights.empty()) { + TrainerState state; + Checkpoint::Load(FLAGS_init_weights, *network, /*optimizer=*/nullptr, state, /*lr_scheduler=*/nullptr); + } 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 +77,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}); @@ -78,8 +89,8 @@ int main(int argc, char *argv[]) { float current_loss = static_cast(loss_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 << ", step: " << train_idx << ", [" << train_idx * FLAGS_bs << "/" + << train_dataset->Size() << "] " << " loss: " << current_loss; } @@ -95,35 +106,37 @@ int main(int argc, char *argv[]) { 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.Forward({new_image}); - auto output_cpu = outputs[0]->To(cpu_device); - auto loss = loss_fn.Forward({outputs[0], new_label}); - auto loss_cpu = loss[0]->To(cpu_device); - - const int batch_size = output_cpu.Dims()[0]; - for (int batch_idx = 0; batch_idx < batch_size; ++batch_idx) { - auto label_index = reinterpret_cast(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; + { + 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->Forward({new_image}); + auto output_cpu = outputs[0]->To(cpu_device); + auto loss = loss_fn.Forward({outputs[0], new_label}); + auto loss_cpu = loss[0]->To(cpu_device); + + const int batch_size = output_cpu.Dims()[0]; + for (int batch_idx = 0; batch_idx < batch_size; ++batch_idx) { + auto label_index = reinterpret_cast(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]); } - 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) << std::format("test | test loss {:.6f} | test accuracy {:.4f} ({}/{})", avg_loss, + static_cast(correct) / total, correct, total); } - 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(); diff --git a/example/mnist/net.cc b/example/mnist/net.cc index 501fee7ef..951ac6715 100644 --- a/example/mnist/net.cc +++ b/example/mnist/net.cc @@ -8,12 +8,19 @@ #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/flatten.h" #include "infini_train/include/nn/modules/linear.h" #include "infini_train/include/nn/modules/module.h" #include "infini_train/include/tensor.h" namespace nn = infini_train::nn; +namespace { +constexpr char kModelMLP[] = "mlp"; +constexpr char kModelCNN[] = "cnn"; +} // namespace + MNIST::MNIST() { std::vector> layers; layers.push_back(std::make_shared(784, 30)); @@ -29,3 +36,38 @@ MNIST::Forward(const std::vector> &x) { auto x2 = (*modules_["linear2"])(x1); return x2; } + +MnistCNN::MnistCNN() { + std::vector> features; + features.push_back(std::make_shared(1, 16, 3)); + features.push_back(std::make_shared()); + features.push_back(std::make_shared(16, 32, 3)); + features.push_back(std::make_shared()); + modules_["features"] = std::make_shared(std::move(features)); + modules_["flatten"] = std::make_shared(); + modules_["classifier"] = std::make_shared(32 * 24 * 24, kNumClasses); +} + +std::vector> +MnistCNN::Forward(const std::vector> &x) { + CHECK_EQ(x.size(), 1); + // The MNIST DataLoader stacks flattened images as (N, 784); restore the (N, 1, H, W) layout for conv. + const auto &input = x[0]; + CHECK_EQ(input->Dims().size(), 2); + auto image = input->View({input->Dims()[0], 1, kImageSize, kImageSize}); + + auto features = (*modules_["features"])({image}); + auto flattened = (*modules_["flatten"])(features); + return (*modules_["classifier"])(flattened); +} + +std::shared_ptr CreateMNISTNetwork(const std::string &model) { + if (model == kModelMLP) { + return std::make_shared(); + } + if (model == kModelCNN) { + return std::make_shared(); + } + LOG(FATAL) << "Unknown model type: " << model << " (expected '" << kModelMLP << "' or '" << kModelCNN << "')"; + return nullptr; +} diff --git a/example/mnist/net.h b/example/mnist/net.h index 5f4cfa33b..dad29f077 100644 --- a/example/mnist/net.h +++ b/example/mnist/net.h @@ -2,6 +2,7 @@ #include #include +#include #include #include "glog/logging.h" @@ -16,3 +17,18 @@ class MNIST : public infini_train::nn::Module { std::vector> Forward(const std::vector> &x) override; }; + +// CNN classifier for MNIST: +// Conv2d(1, 16, 3) -> ReLU -> Conv2d(16, 32, 3) -> ReLU -> Flatten -> Linear(32*24*24, 10) +class MnistCNN : public infini_train::nn::Module { +public: + static constexpr int kImageSize = 28; + static constexpr int kNumClasses = 10; + + MnistCNN(); + + std::vector> + Forward(const std::vector> &x) override; +}; + +std::shared_ptr CreateMNISTNetwork(const std::string &model); From e1cf0d6f33d000c91f450732304ea6543d318e06 Mon Sep 17 00:00:00 2001 From: gavin-richie <2904793019@qq.com> Date: Mon, 7 Sep 2026 06:38:32 +0800 Subject: [PATCH 06/14] test: cover Conv2d and ReLU forward and backward Device-parameterized gtest coverage with PyTorch-derived reference vectors: ReLU elementwise forward/backward, Conv2d forward with and without bias across stride/padding combinations, Conv2d input, weight, and bias gradients, and module-level parameter, shape, and conv-relu-flatten-linear-cross-entropy training integration checks. --- .../autograd/test_autograd_conv2d_backward.cc | 152 ++++++++++++++++++ .../autograd/test_autograd_conv2d_forward.cc | 147 +++++++++++++++++ tests/autograd/test_autograd_relu.cc | 47 ++++++ tests/module/test_module_conv2d.cc | 106 ++++++++++++ 4 files changed, 452 insertions(+) create mode 100644 tests/autograd/test_autograd_conv2d_backward.cc create mode 100644 tests/autograd/test_autograd_conv2d_forward.cc create mode 100644 tests/autograd/test_autograd_relu.cc create mode 100644 tests/module/test_module_conv2d.cc diff --git a/tests/autograd/test_autograd_conv2d_backward.cc b/tests/autograd/test_autograd_conv2d_backward.cc new file mode 100644 index 000000000..7a39c9f42 --- /dev/null +++ b/tests/autograd/test_autograd_conv2d_backward.cc @@ -0,0 +1,152 @@ +#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 { +// a_input: shape [1, 2, 4, 4] +const std::vector a_input + = {1.926915f, 1.487284f, 0.900717f, -2.105521f, 0.678418f, -1.234545f, -0.043067f, -1.604667f, + -0.752135f, 1.648723f, -0.392479f, -1.403607f, -0.727881f, -0.559430f, -0.768839f, 0.762445f, + 1.642317f, -0.159597f, -0.497398f, 0.439589f, -0.758131f, 1.078318f, 0.800801f, 1.680621f, + 1.279124f, 1.296423f, 0.610466f, 1.334738f, -0.231624f, 0.041759f, -0.251575f, 0.859859f}; +// a_weight: shape [3, 2, 3, 3] +const std::vector a_weight + = {-1.384674f, -0.871236f, -0.223366f, 1.717361f, 0.318880f, -0.424519f, 0.305721f, -0.774593f, -1.557572f, + 0.995636f, -0.879786f, -0.601142f, -1.274151f, 2.122785f, -1.234653f, -0.487914f, -0.913823f, -0.658137f, + 0.078024f, 0.525809f, -0.487992f, 1.191369f, -0.814008f, -0.735993f, -1.403248f, 0.036004f, -0.063477f, + 0.675615f, -0.097807f, 1.844594f, -1.184537f, 1.383549f, 1.445134f, 0.856413f, 2.218076f, 0.523166f, + 0.346647f, -0.197331f, 1.141202f, 0.051644f, 0.728110f, -0.710642f, -0.602068f, 0.960449f, 0.404814f, + -1.354343f, 1.334703f, 0.483539f, -0.197562f, 1.268311f, 1.224263f, 0.098117f, 1.742253f, -1.352674f}; +// a_bias: shape [3] +const std::vector a_bias = {-0.437065f, 0.077618f, -0.922759f}; +// a_output: shape [1, 3, 2, 2] +const std::vector a_output = {-2.577903f, -5.072696f, -1.058225f, 1.205943f, 11.600315f, 5.375094f, + 0.386035f, 9.683105f, 2.902302f, -1.870026f, 5.919574f, -1.493631f}; +// a_grad_output: shape [1, 3, 2, 2] +const std::vector a_grad_output = {-0.728986f, 0.727698f, -0.007153f, -0.977891f, 1.003980f, -0.024754f, + 0.147247f, -0.130241f, -0.755130f, -0.636322f, -0.684569f, -1.285823f}; +// a_grad_input: shape [1, 2, 4, 4] +const std::vector a_grad_input + = {0.825979f, 0.081896f, -1.710306f, -0.876634f, -0.310730f, 0.704758f, 0.081797f, -1.023905f, + -1.049271f, -2.005676f, -1.170570f, 0.035285f, 0.203346f, 0.011301f, -0.757518f, 1.010884f, + 0.975205f, 1.104877f, 0.437927f, -0.790798f, 0.908273f, -2.130304f, 1.231674f, -1.987390f, + 1.111350f, 2.113331f, -4.305401f, -0.186185f, 0.062427f, -0.620124f, -0.627753f, 2.314747f}; +// a_grad_weight: shape [3, 2, 3, 3] +const std::vector a_grad_weight + = {0.879997f, -0.377813f, -0.619297f, -2.999825f, 1.240634f, 0.239064f, 2.300337f, -0.731660f, -1.475380f, + -2.362419f, -1.036420f, -0.966707f, 0.060448f, -0.809580f, -0.670382f, -0.028238f, -0.255124f, -0.312783f, + 2.158452f, 1.294732f, 1.159075f, 0.386198f, -0.944505f, 0.121499f, -0.830260f, 1.682760f, -0.571807f, + 1.400730f, -0.093438f, -0.611229f, -0.768341f, 1.174173f, 0.678437f, 1.212578f, 1.325385f, 0.430823f, + -1.278481f, -0.795731f, 2.752429f, -1.331804f, 0.335641f, 3.127074f, 0.736455f, 0.376301f, 0.735471f, + -2.006141f, -1.330848f, -2.613303f, -2.656287f, -2.996279f, -3.808266f, -1.685980f, -1.072527f, -2.243709f}; +// a_grad_bias: shape [3] +const std::vector a_grad_bias = {-0.986331f, 0.996232f, -3.361843f}; +// b_input: shape [2, 1, 5, 5] +const std::vector b_input + = {0.654741f, 0.576005f, -0.360910f, -0.060590f, 0.073255f, 0.818653f, 1.480474f, 0.344929f, -0.686651f, + 0.636814f, 0.217553f, -0.046655f, -1.433521f, -0.566527f, -0.425283f, 0.262519f, -0.732803f, 0.104298f, + 1.041401f, -0.399731f, -2.293334f, 0.497563f, -0.425723f, -1.337147f, -1.195454f, 0.812337f, -0.306278f, + -0.330158f, -0.980803f, 0.194734f, -1.653521f, 0.681419f, 0.174820f, -1.093929f, 0.716537f, 1.533467f, + -1.450979f, -0.786135f, -0.956316f, -1.247602f, -0.749942f, -0.592190f, -1.532617f, -0.725131f, 0.466404f, + 0.666725f, -0.043871f, 0.236808f, -0.706068f, -0.716906f}; +// b_weight: shape [2, 1, 2, 2] +const std::vector b_weight + = {0.494319f, -0.642960f, 0.711319f, 0.399978f, -1.203922f, -0.419752f, -1.192891f, -0.935063f}; +// b_bias: shape [2] +const std::vector b_bias = {0.213803f, -1.284212f}; +// b_output: shape [2, 2, 3, 3] +const std::vector b_output + = {0.475685f, 0.479170f, 0.200004f, -0.225542f, 0.117290f, -1.108152f, -0.872269f, -0.031849f, -0.443695f, + -1.896435f, -1.633849f, -1.280432f, -1.831269f, -1.815284f, 0.348629f, 0.750008f, -0.641214f, 0.342706f, + 0.538720f, -0.136115f, -0.405972f, 1.890304f, -0.908306f, -1.966910f, 0.962661f, 0.969996f, -1.233508f, + -2.043798f, -0.610136f, -0.296308f, -2.024031f, 0.287977f, 2.039392f, -1.592851f, -0.097037f, 0.905630f}; +// b_grad_output: shape [2, 2, 3, 3] +const std::vector b_grad_output + = {0.361835f, 1.999347f, 0.663007f, 0.704733f, -0.004505f, 1.666792f, 0.153920f, -1.060253f, 0.507098f, + 0.082078f, 0.443975f, -0.724034f, -0.071988f, -0.906094f, -2.048712f, -1.081056f, -0.982707f, 0.301771f, + 0.178692f, -0.129309f, -0.685482f, 0.563559f, -1.507175f, -1.610666f, -1.479047f, 0.432274f, -0.125025f, + 0.782118f, -1.598768f, -0.109130f, 0.715199f, 0.039139f, 1.305860f, 0.246593f, -1.977591f, 0.017896f}; +// b_grad_input: shape [2, 1, 5, 5] +const std::vector b_grad_input + = {0.067978f, 0.892559f, 0.384550f, 1.335303f, 0.942206f, -0.422898f, 1.088640f, 0.383232f, 3.290416f, + -0.211729f, 0.349190f, 1.077667f, 0.845453f, 3.629510f, 2.582355f, 0.354811f, 0.658999f, 1.094194f, + -0.112641f, -0.452713f, 1.072420f, 0.418085f, 0.494815f, 0.000728f, -0.079347f, -0.659857f, 1.815175f, + 1.443228f, -0.357416f, -0.172134f, -0.662553f, -0.792145f, 0.952625f, -2.368335f, 0.487456f, -0.443345f, + -1.118770f, -0.639434f, -2.703445f, -1.865292f, 0.847460f, 2.594546f, 0.552163f, -0.083348f, 0.072874f, + -0.822166f, 2.666535f, 2.022072f, -0.110281f, -0.066741f}; +// b_grad_weight: shape [2, 1, 2, 2] +const std::vector b_grad_weight + = {0.723482f, -0.597131f, 3.470205f, 1.825737f, 0.855981f, 0.431211f, -0.024724f, 4.793805f}; +// b_grad_bias: shape [2] +const std::vector b_grad_bias = {0.629796f, -5.565448f}; + +std::shared_ptr MakeTensor(const std::vector &dims, const std::vector &values, + const Device &device, bool requires_grad = false) { + auto cpu_tensor = std::make_shared(dims, DataType::kFLOAT32, Device()); + std::copy(values.begin(), values.end(), static_cast(cpu_tensor->DataPtr())); + auto tensor = std::make_shared(cpu_tensor->To(device)); + return requires_grad ? tensor->RequiresGrad() : tensor; +} +} // namespace + +class AutogradConv2dBackwardTest : public infini_train::test::InfiniTrainTest {}; + +TEST_P(AutogradConv2dBackwardTest, Conv2dBackwardStride1NoPadding) { + auto input = MakeTensor({1, 2, 4, 4}, a_input, GetDevice(), /*requires_grad=*/true); + auto weight = MakeTensor({3, 2, 3, 3}, a_weight, GetDevice(), /*requires_grad=*/true); + auto bias = MakeTensor({3}, a_bias, GetDevice(), /*requires_grad=*/true); + auto grad_output = MakeTensor({1, 3, 2, 2}, a_grad_output, GetDevice()); + + auto conv_fn = std::make_shared(/*stride=*/1, /*padding=*/0); + conv_fn->Apply({input, weight, bias}); + + const auto grads = conv_fn->Backward({grad_output}); + EXPECT_EQ(grads.size(), 3); + test::ExpectTensorNear(grads[0], a_grad_input, 1e-4f); + test::ExpectTensorNear(grads[1], a_grad_weight, 1e-4f); + test::ExpectTensorNear(grads[2], a_grad_bias, 1e-4f); +} + +TEST_P(AutogradConv2dBackwardTest, Conv2dBackwardStride2Padding1Batch) { + auto input = MakeTensor({2, 1, 5, 5}, b_input, GetDevice(), /*requires_grad=*/true); + auto weight = MakeTensor({2, 1, 2, 2}, b_weight, GetDevice(), /*requires_grad=*/true); + auto bias = MakeTensor({2}, b_bias, GetDevice(), /*requires_grad=*/true); + auto grad_output = MakeTensor({2, 2, 3, 3}, b_grad_output, GetDevice()); + + auto conv_fn = std::make_shared(/*stride=*/2, /*padding=*/1); + conv_fn->Apply({input, weight, bias}); + + const auto grads = conv_fn->Backward({grad_output}); + EXPECT_EQ(grads.size(), 3); + test::ExpectTensorNear(grads[0], b_grad_input, 1e-4f); + test::ExpectTensorNear(grads[1], b_grad_weight, 1e-4f); + test::ExpectTensorNear(grads[2], b_grad_bias, 1e-4f); +} + +TEST_P(AutogradConv2dBackwardTest, Conv2dBackwardSkipsUnusedInputs) { + // Without bias and with a non-differentiable input, only the weight gradient is produced. + auto input = std::make_shared(std::vector{1, 2, 4, 4}, DataType::kFLOAT32, GetDevice(), false); + input->Fill(0.5f); + auto weight = std::make_shared(std::vector{3, 2, 3, 3}, DataType::kFLOAT32, GetDevice(), true); + weight->Fill(0.1f); + auto grad_output = MakeTensor({1, 3, 2, 2}, a_grad_output, GetDevice()); + + auto conv_fn = std::make_shared(/*stride=*/1, /*padding=*/0); + conv_fn->Apply({input, weight}); + + const auto grads = conv_fn->Backward({grad_output}); + EXPECT_EQ(grads.size(), 3); + EXPECT_EQ(grads[0], nullptr); + EXPECT_NE(grads[1], nullptr); + EXPECT_EQ(grads[2], nullptr); +} + +INFINI_TRAIN_REGISTER_TEST(AutogradConv2dBackwardTest); diff --git a/tests/autograd/test_autograd_conv2d_forward.cc b/tests/autograd/test_autograd_conv2d_forward.cc new file mode 100644 index 000000000..01cd6a2ea --- /dev/null +++ b/tests/autograd/test_autograd_conv2d_forward.cc @@ -0,0 +1,147 @@ +#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 { +// a_input: shape [1, 2, 4, 4] +const std::vector a_input + = {1.926915f, 1.487284f, 0.900717f, -2.105521f, 0.678418f, -1.234545f, -0.043067f, -1.604667f, + -0.752135f, 1.648723f, -0.392479f, -1.403607f, -0.727881f, -0.559430f, -0.768839f, 0.762445f, + 1.642317f, -0.159597f, -0.497398f, 0.439589f, -0.758131f, 1.078318f, 0.800801f, 1.680621f, + 1.279124f, 1.296423f, 0.610466f, 1.334738f, -0.231624f, 0.041759f, -0.251575f, 0.859859f}; +// a_weight: shape [3, 2, 3, 3] +const std::vector a_weight + = {-1.384674f, -0.871236f, -0.223366f, 1.717361f, 0.318880f, -0.424519f, 0.305721f, -0.774593f, -1.557572f, + 0.995636f, -0.879786f, -0.601142f, -1.274151f, 2.122785f, -1.234653f, -0.487914f, -0.913823f, -0.658137f, + 0.078024f, 0.525809f, -0.487992f, 1.191369f, -0.814008f, -0.735993f, -1.403248f, 0.036004f, -0.063477f, + 0.675615f, -0.097807f, 1.844594f, -1.184537f, 1.383549f, 1.445134f, 0.856413f, 2.218076f, 0.523166f, + 0.346647f, -0.197331f, 1.141202f, 0.051644f, 0.728110f, -0.710642f, -0.602068f, 0.960449f, 0.404814f, + -1.354343f, 1.334703f, 0.483539f, -0.197562f, 1.268311f, 1.224263f, 0.098117f, 1.742253f, -1.352674f}; +// a_bias: shape [3] +const std::vector a_bias = {-0.437065f, 0.077618f, -0.922759f}; +// a_output: shape [1, 3, 2, 2] +const std::vector a_output = {-2.577903f, -5.072696f, -1.058225f, 1.205943f, 11.600315f, 5.375094f, + 0.386035f, 9.683105f, 2.902302f, -1.870026f, 5.919574f, -1.493631f}; +// a_grad_output: shape [1, 3, 2, 2] +const std::vector a_grad_output = {-0.728986f, 0.727698f, -0.007153f, -0.977891f, 1.003980f, -0.024754f, + 0.147247f, -0.130241f, -0.755130f, -0.636322f, -0.684569f, -1.285823f}; +// a_grad_input: shape [1, 2, 4, 4] +const std::vector a_grad_input + = {0.825979f, 0.081896f, -1.710306f, -0.876634f, -0.310730f, 0.704758f, 0.081797f, -1.023905f, + -1.049271f, -2.005676f, -1.170570f, 0.035285f, 0.203346f, 0.011301f, -0.757518f, 1.010884f, + 0.975205f, 1.104877f, 0.437927f, -0.790798f, 0.908273f, -2.130304f, 1.231674f, -1.987390f, + 1.111350f, 2.113331f, -4.305401f, -0.186185f, 0.062427f, -0.620124f, -0.627753f, 2.314747f}; +// a_grad_weight: shape [3, 2, 3, 3] +const std::vector a_grad_weight + = {0.879997f, -0.377813f, -0.619297f, -2.999825f, 1.240634f, 0.239064f, 2.300337f, -0.731660f, -1.475380f, + -2.362419f, -1.036420f, -0.966707f, 0.060448f, -0.809580f, -0.670382f, -0.028238f, -0.255124f, -0.312783f, + 2.158452f, 1.294732f, 1.159075f, 0.386198f, -0.944505f, 0.121499f, -0.830260f, 1.682760f, -0.571807f, + 1.400730f, -0.093438f, -0.611229f, -0.768341f, 1.174173f, 0.678437f, 1.212578f, 1.325385f, 0.430823f, + -1.278481f, -0.795731f, 2.752429f, -1.331804f, 0.335641f, 3.127074f, 0.736455f, 0.376301f, 0.735471f, + -2.006141f, -1.330848f, -2.613303f, -2.656287f, -2.996279f, -3.808266f, -1.685980f, -1.072527f, -2.243709f}; +// a_grad_bias: shape [3] +const std::vector a_grad_bias = {-0.986331f, 0.996232f, -3.361843f}; +// b_input: shape [2, 1, 5, 5] +const std::vector b_input + = {0.654741f, 0.576005f, -0.360910f, -0.060590f, 0.073255f, 0.818653f, 1.480474f, 0.344929f, -0.686651f, + 0.636814f, 0.217553f, -0.046655f, -1.433521f, -0.566527f, -0.425283f, 0.262519f, -0.732803f, 0.104298f, + 1.041401f, -0.399731f, -2.293334f, 0.497563f, -0.425723f, -1.337147f, -1.195454f, 0.812337f, -0.306278f, + -0.330158f, -0.980803f, 0.194734f, -1.653521f, 0.681419f, 0.174820f, -1.093929f, 0.716537f, 1.533467f, + -1.450979f, -0.786135f, -0.956316f, -1.247602f, -0.749942f, -0.592190f, -1.532617f, -0.725131f, 0.466404f, + 0.666725f, -0.043871f, 0.236808f, -0.706068f, -0.716906f}; +// b_weight: shape [2, 1, 2, 2] +const std::vector b_weight + = {0.494319f, -0.642960f, 0.711319f, 0.399978f, -1.203922f, -0.419752f, -1.192891f, -0.935063f}; +// b_bias: shape [2] +const std::vector b_bias = {0.213803f, -1.284212f}; +// b_output: shape [2, 2, 3, 3] +const std::vector b_output + = {0.475685f, 0.479170f, 0.200004f, -0.225542f, 0.117290f, -1.108152f, -0.872269f, -0.031849f, -0.443695f, + -1.896435f, -1.633849f, -1.280432f, -1.831269f, -1.815284f, 0.348629f, 0.750008f, -0.641214f, 0.342706f, + 0.538720f, -0.136115f, -0.405972f, 1.890304f, -0.908306f, -1.966910f, 0.962661f, 0.969996f, -1.233508f, + -2.043798f, -0.610136f, -0.296308f, -2.024031f, 0.287977f, 2.039392f, -1.592851f, -0.097037f, 0.905630f}; +// b_grad_output: shape [2, 2, 3, 3] +const std::vector b_grad_output + = {0.361835f, 1.999347f, 0.663007f, 0.704733f, -0.004505f, 1.666792f, 0.153920f, -1.060253f, 0.507098f, + 0.082078f, 0.443975f, -0.724034f, -0.071988f, -0.906094f, -2.048712f, -1.081056f, -0.982707f, 0.301771f, + 0.178692f, -0.129309f, -0.685482f, 0.563559f, -1.507175f, -1.610666f, -1.479047f, 0.432274f, -0.125025f, + 0.782118f, -1.598768f, -0.109130f, 0.715199f, 0.039139f, 1.305860f, 0.246593f, -1.977591f, 0.017896f}; +// b_grad_input: shape [2, 1, 5, 5] +const std::vector b_grad_input + = {0.067978f, 0.892559f, 0.384550f, 1.335303f, 0.942206f, -0.422898f, 1.088640f, 0.383232f, 3.290416f, + -0.211729f, 0.349190f, 1.077667f, 0.845453f, 3.629510f, 2.582355f, 0.354811f, 0.658999f, 1.094194f, + -0.112641f, -0.452713f, 1.072420f, 0.418085f, 0.494815f, 0.000728f, -0.079347f, -0.659857f, 1.815175f, + 1.443228f, -0.357416f, -0.172134f, -0.662553f, -0.792145f, 0.952625f, -2.368335f, 0.487456f, -0.443345f, + -1.118770f, -0.639434f, -2.703445f, -1.865292f, 0.847460f, 2.594546f, 0.552163f, -0.083348f, 0.072874f, + -0.822166f, 2.666535f, 2.022072f, -0.110281f, -0.066741f}; +// b_grad_weight: shape [2, 1, 2, 2] +const std::vector b_grad_weight + = {0.723482f, -0.597131f, 3.470205f, 1.825737f, 0.855981f, 0.431211f, -0.024724f, 4.793805f}; +// b_grad_bias: shape [2] +const std::vector b_grad_bias = {0.629796f, -5.565448f}; + +std::shared_ptr MakeTensor(const std::vector &dims, const std::vector &values, + const Device &device, bool requires_grad = false) { + auto cpu_tensor = std::make_shared(dims, DataType::kFLOAT32, Device()); + std::copy(values.begin(), values.end(), static_cast(cpu_tensor->DataPtr())); + auto tensor = std::make_shared(cpu_tensor->To(device)); + return requires_grad ? tensor->RequiresGrad() : tensor; +} +} // namespace + +class AutogradConv2dForwardTest : public infini_train::test::InfiniTrainTest {}; + +TEST_P(AutogradConv2dForwardTest, Conv2dForwardStride1NoPadding) { + const std::vector input_dims = {1, 2, 4, 4}; + const std::vector output_dims = {1, 3, 2, 2}; + auto input = MakeTensor(input_dims, a_input, GetDevice()); + auto weight = MakeTensor({3, 2, 3, 3}, a_weight, GetDevice()); + auto bias = MakeTensor({3}, a_bias, GetDevice()); + + auto conv_fn = std::make_shared(/*stride=*/1, /*padding=*/0); + auto result = conv_fn->Apply({input, weight, bias}); + EXPECT_EQ(result.size(), 1); + EXPECT_EQ(result[0]->Dims(), output_dims); + test::ExpectTensorNear(result[0], a_output, 1e-4f); +} + +TEST_P(AutogradConv2dForwardTest, Conv2dForwardStride2Padding1Batch) { + const std::vector input_dims = {2, 1, 5, 5}; + const std::vector output_dims = {2, 2, 3, 3}; + auto input = MakeTensor(input_dims, b_input, GetDevice()); + auto weight = MakeTensor({2, 1, 2, 2}, b_weight, GetDevice()); + auto bias = MakeTensor({2}, b_bias, GetDevice()); + + auto conv_fn = std::make_shared(/*stride=*/2, /*padding=*/1); + auto result = conv_fn->Apply({input, weight, bias}); + EXPECT_EQ(result.size(), 1); + EXPECT_EQ(result[0]->Dims(), output_dims); + test::ExpectTensorNear(result[0], b_output, 1e-4f); +} + +TEST_P(AutogradConv2dForwardTest, Conv2dForwardNoBias) { + const std::vector input_dims = {1, 2, 4, 4}; + auto input = MakeTensor(input_dims, a_input, GetDevice()); + auto weight = MakeTensor({3, 2, 3, 3}, a_weight, GetDevice()); + + auto conv_fn = std::make_shared(/*stride=*/1, /*padding=*/0); + auto result = conv_fn->Apply({input, weight}); + EXPECT_EQ(result.size(), 1); + EXPECT_EQ(result[0]->Dims(), (std::vector{1, 3, 2, 2})); + // Bias-free output is the biased reference minus the bias per channel; layout is (C, H_out, W_out). + std::vector expected(a_output); + constexpr int64_t kP = 2 * 2; + for (int64_t c = 0; c < 3; ++c) { + for (int64_t i = 0; i < kP; ++i) { expected[c * kP + i] -= a_bias[c]; } + } + test::ExpectTensorNear(result[0], expected, 1e-4f); +} + +INFINI_TRAIN_REGISTER_TEST(AutogradConv2dForwardTest); diff --git a/tests/autograd/test_autograd_relu.cc b/tests/autograd/test_autograd_relu.cc new file mode 100644 index 000000000..7f39615f0 --- /dev/null +++ b/tests/autograd/test_autograd_relu.cc @@ -0,0 +1,47 @@ +#include +#include + +#include "gtest/gtest.h" + +#include "infini_train/include/autograd/activations.h" +#include "infini_train/include/tensor.h" + +#include "tests/common/test_utils.h" + +using namespace infini_train; + +namespace { +std::shared_ptr MakeTensor(const std::vector &dims, const std::vector &values, + const Device &device, bool requires_grad = false) { + auto cpu_tensor = std::make_shared(dims, DataType::kFLOAT32, Device()); + std::copy(values.begin(), values.end(), static_cast(cpu_tensor->DataPtr())); + auto tensor = std::make_shared(cpu_tensor->To(device)); + return requires_grad ? tensor->RequiresGrad() : tensor; +} +} // namespace + +class AutogradReluTest : public infini_train::test::InfiniTrainTest {}; + +TEST_P(AutogradReluTest, ReluForward) { + auto input = MakeTensor({2, 3}, {-1.0f, 0.0f, 2.5f, -3.0f, 1.5f, -0.25f}, GetDevice()); + auto relu_fn = std::make_shared(); + auto result = relu_fn->Apply({input}); + EXPECT_EQ(result.size(), 1); + EXPECT_EQ(result[0]->Dims(), (std::vector{2, 3})); + test::ExpectTensorFloatEqual(result[0], {0.0f, 0.0f, 2.5f, 0.0f, 1.5f, 0.0f}); +} + +TEST_P(AutogradReluTest, ReluBackward) { + auto output = MakeTensor({2, 3}, {0.0f, 0.0f, 2.5f, 0.0f, 1.5f, 0.0f}, GetDevice()); + auto grad_output = MakeTensor({2, 3}, {1.0f, -2.0f, 0.5f, -1.0f, 3.0f, 7.0f}, GetDevice()); + + auto relu_fn = std::make_shared(); + auto input = MakeTensor({2, 3}, {-1.0f, 0.0f, 2.5f, -3.0f, 1.5f, -0.25f}, GetDevice()); + relu_fn->Apply({input}); + + const auto grads = relu_fn->Backward({grad_output}); + EXPECT_EQ(grads.size(), 1); + test::ExpectTensorFloatEqual(grads[0], {0.0f, 0.0f, 0.5f, 0.0f, 3.0f, 0.0f}); +} + +INFINI_TRAIN_REGISTER_TEST(AutogradReluTest); diff --git a/tests/module/test_module_conv2d.cc b/tests/module/test_module_conv2d.cc new file mode 100644 index 000000000..2154d3ac1 --- /dev/null +++ b/tests/module/test_module_conv2d.cc @@ -0,0 +1,106 @@ +#include +#include + +#include "gtest/gtest.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/flatten.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; + +class ModuleConv2dTest : public infini_train::test::InfiniTrainTest {}; + +TEST_P(ModuleConv2dTest, Conv2dParameterShapes) { + const int64_t in_channels = 2; + const int64_t out_channels = 4; + const int64_t kernel_size = 3; + nn::Conv2d conv(in_channels, out_channels, kernel_size); + + const auto state_dict = conv.StateDict(); + EXPECT_EQ(state_dict.size(), 2); + ASSERT_TRUE(state_dict.contains("weight")); + EXPECT_EQ(state_dict.at("weight")->Dims(), + (std::vector{out_channels, in_channels, kernel_size, kernel_size})); + ASSERT_TRUE(state_dict.contains("bias")); + EXPECT_EQ(state_dict.at("bias")->Dims(), (std::vector{out_channels})); + EXPECT_TRUE(state_dict.at("weight")->requires_grad()); + EXPECT_TRUE(state_dict.at("bias")->requires_grad()); +} + +TEST_P(ModuleConv2dTest, Conv2dNoBiasParameterShapes) { + nn::Conv2d conv(2, 4, 3, /*stride=*/1, /*padding=*/0, /*bias=*/false); + const auto state_dict = conv.StateDict(); + EXPECT_EQ(state_dict.size(), 1); + EXPECT_TRUE(state_dict.contains("weight")); +} + +TEST_P(ModuleConv2dTest, Conv2dForwardShape) { + nn::Conv2d conv(2, 4, 3, /*stride=*/1, /*padding=*/1, /*bias=*/true, GetDevice()); + auto input = std::make_shared(std::vector{2, 2, 5, 5}, DataType::kFLOAT32, GetDevice(), false); + input->Fill(1.0f); + + const auto output = conv.Forward({input}); + EXPECT_EQ(output.size(), 1); + // With padding 1, the spatial size is preserved: (5 + 2 * 1 - 3) / 1 + 1 = 5. + EXPECT_EQ(output[0]->Dims(), (std::vector{2, 4, 5, 5})); +} + +TEST_P(ModuleConv2dTest, FlattenModule) { + auto input = std::make_shared(std::vector{2, 3, 4, 5}, DataType::kFLOAT32, GetDevice(), true); + input->Fill(1.0f); + + nn::Flatten flatten; + const auto output = flatten.Forward({input}); + EXPECT_EQ(output.size(), 1); + EXPECT_EQ(output[0]->Dims(), (std::vector{2, 60})); +} + +TEST_P(ModuleConv2dTest, ConvReluFlattenLinearChain) { + // End-to-end wiring test: conv -> relu -> flatten -> linear -> loss trains with SGD. + const int64_t batch_size = 2; + const int64_t num_classes = 10; + + const Device device = GetDevice(); + nn::Conv2d conv(1, 2, 3, /*stride=*/1, /*padding=*/0, /*bias=*/true, device); + nn::Relu relu; + nn::Flatten flatten; + nn::Linear linear(2 * 26 * 26, num_classes, /*bias=*/true, device); + nn::CrossEntropyLoss loss_fn; + loss_fn.To(device); + + auto input + = std::make_shared(std::vector{batch_size, 1, 28, 28}, DataType::kFLOAT32, device, false); + input->Fill(0.5f); + auto label_cpu = std::make_shared(std::vector{batch_size}, DataType::kUINT8, Device(), false); + static_cast(label_cpu->DataPtr())[0] = 1; + static_cast(label_cpu->DataPtr())[1] = 2; + auto label = std::make_shared(label_cpu->To(device)); + + auto outputs = conv.Forward({input}); + outputs = relu.Forward(outputs); + outputs = flatten.Forward(outputs); + outputs = linear.Forward(outputs); + const auto loss = loss_fn.Forward({outputs[0], label}); + ASSERT_EQ(loss.size(), 1); + loss[0]->Backward(); + + std::vector> params{conv.StateDict().at("weight"), conv.StateDict().at("bias"), + linear.StateDict().at("weight"), linear.StateDict().at("bias")}; + for (const auto ¶m : params) { + ASSERT_NE(param->grad(), nullptr); + EXPECT_EQ(param->grad()->Dims(), param->Dims()); + } + + optimizers::SGD optimizer(params, 0.01); + optimizer.Step(); +} + +INFINI_TRAIN_REGISTER_TEST(ModuleConv2dTest); From 580579439a084fa88e4e11a1edc79370163c4c82 Mon Sep 17 00:00:00 2001 From: gavin-richie <2904793019@qq.com> Date: Mon, 7 Sep 2026 06:38:32 +0800 Subject: [PATCH 07/14] build: add sm_89 to the CUDA architecture list --- CMakeLists.txt | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/CMakeLists.txt b/CMakeLists.txt index 6bd8069d4..4d90c3490 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -104,7 +104,7 @@ if(USE_CUDA) file(GLOB_RECURSE CUDA_KERNELS ${PROJECT_SOURCE_DIR}/infini_train/src/*.cu) add_library(infini_train_cuda_kernels STATIC ${CUDA_KERNELS}) - set_target_properties(infini_train_cuda_kernels PROPERTIES CUDA_ARCHITECTURES "75;80;90") + set_target_properties(infini_train_cuda_kernels PROPERTIES CUDA_ARCHITECTURES "75;80;89;90") target_link_libraries(infini_train_cuda_kernels PUBLIC From 16c92e7e518d986386d70fa8c7dd808bdbfba79b Mon Sep 17 00:00:00 2001 From: gavin-richie <2904793019@qq.com> Date: Mon, 7 Sep 2026 06:38:32 +0800 Subject: [PATCH 08/14] docs: document the MNIST CNN demo options --- README.md | 326 +++++++++++++++++++++++++++++------------------------- 1 file changed, 175 insertions(+), 151 deletions(-) diff --git a/README.md b/README.md index abd8070b2..8fa4958b6 100644 --- a/README.md +++ b/README.md @@ -50,7 +50,7 @@ Build Options: > Both options are optional and can be disabled for CPU-only builds. -## ✨ InfiniTrain Overview +## ✨ InfiniTrain Overview ### ✔ Support Matrix @@ -96,160 +96,184 @@ For example, the `llama3` example produces a binary named `llama3`. To view available runtime options: -```bash -./build/llama3 --help +```bash +./build/llama3 --help +``` + +### Getting Started + +#### Prepare Datasets and Weights + +Run the asset preparation script from the repository root. Prepared files are +written to `data/` by default. + +```bash +# MNIST dataset +./scripts/assets/prepare-infinitrain-assets.sh mnist + +# GPT-2 124M weights, tokenizer, and tokenized TinyShakespeare data +./scripts/assets/prepare-infinitrain-assets.sh gpt2 + +# LLaMA 3.2 1B weights and tokenized TinyShakespeare data +HF_TOKEN=hf_xxx ./scripts/assets/prepare-infinitrain-assets.sh llama3 +``` + +Preparing LLaMA requires access to the gated +`meta-llama/Llama-3.2-1B` repository. Accept its license on Hugging Face and +provide `HF_TOKEN`, or authenticate with `hf auth login`, before running the +command. The complete LLaMA preparation requires approximately 8.5 GB of free +disk space, including the downloaded checkpoint and converted FP32 weights. + +Use `DATA_DIR` to write the assets elsewhere, or prepare all supported assets +in one invocation: + +```bash +DATA_DIR=/path/to/data \ +HF_TOKEN=hf_xxx \ +./scripts/assets/prepare-infinitrain-assets.sh all +``` + +#### Model Examples + +The generated files can be passed directly to the corresponding executables: + +##### MNIST + +Train an MLP classifier (the original example) or a CNN classifier on MNIST. +The CNN stacks `Conv2d(1, 16, 3) -> ReLU -> Conv2d(16, 32, 3) -> ReLU -> +Flatten -> Linear(18432, 10)` and reports test loss and accuracy after every +epoch. + +```bash +# MLP (default: --model mlp) +./build/mnist \ + --device cpu \ + --dataset data/mnist + +# CNN on CUDA +./build/mnist \ + --model cnn \ + --device cuda \ + --dataset data/mnist ``` -### Getting Started - -#### Prepare Datasets and Weights - -Run the asset preparation script from the repository root. Prepared files are -written to `data/` by default. - -```bash -# MNIST dataset -./scripts/assets/prepare-infinitrain-assets.sh mnist - -# GPT-2 124M weights, tokenizer, and tokenized TinyShakespeare data -./scripts/assets/prepare-infinitrain-assets.sh gpt2 - -# LLaMA 3.2 1B weights and tokenized TinyShakespeare data -HF_TOKEN=hf_xxx ./scripts/assets/prepare-infinitrain-assets.sh llama3 -``` - -Preparing LLaMA requires access to the gated -`meta-llama/Llama-3.2-1B` repository. Accept its license on Hugging Face and -provide `HF_TOKEN`, or authenticate with `hf auth login`, before running the -command. The complete LLaMA preparation requires approximately 8.5 GB of free -disk space, including the downloaded checkpoint and converted FP32 weights. - -Use `DATA_DIR` to write the assets elsewhere, or prepare all supported assets -in one invocation: - -```bash -DATA_DIR=/path/to/data \ -HF_TOKEN=hf_xxx \ -./scripts/assets/prepare-infinitrain-assets.sh all -``` - -#### Model Examples - -The generated files can be passed directly to the corresponding executables: - -##### MNIST - -```bash -./build/mnist \ - --device cpu \ - --dataset data/mnist -``` - -##### GPT-2 124M - -```bash -./build/gpt2 \ - --device cuda \ - --input_bin data/gpt2/tiny_shakespeare_train.bin \ - --input_val_bin data/gpt2/tiny_shakespeare_val.bin \ - --tokenizer_bin data/gpt2/gpt2_tokenizer.bin \ - --llmc_filepath data/gpt2/gpt2_124M.bin \ - --num_iteration 10 -``` - -##### LLaMA 3.2 1B - -```bash -./build/llama3 \ - --device cuda \ - --input_bin data/llama3/tiny_shakespeare_train.bin \ - --input_val_bin data/llama3/tiny_shakespeare_val.bin \ - --llmc_filepath data/llama3/llama3.2_1B_fp32.bin \ - --num_iteration 10 -``` - -### Launch Modes - -GPT-2 and LLaMA training support both thread-based and process-based launches. -The examples below use LLaMA, but the same launch modes also apply to GPT-2. - -#### Direct Launch - -Running a model executable directly uses one process and one device by default. -Set `--nthread_per_process` to use multiple execution threads and devices in the -same process: - -```bash -./build/llama3 \ - --device cuda \ - --input_bin data/llama3/tiny_shakespeare_train.bin \ - --llmc_filepath data/llama3/llama3.2_1B_fp32.bin \ - --nthread_per_process 8 \ - --num_iteration 10 -``` - -#### Single-node Multi-process Launch - -Use `infini_run` to start multiple training processes on one node. Each process -uses one execution thread by default: - -```bash -./build/infini_run \ - --nnodes=1 \ - --nproc_per_node=8 \ - ./build/llama3 \ - --device cuda \ - --input_bin data/llama3/tiny_shakespeare_train.bin \ - --llmc_filepath data/llama3/llama3.2_1B_fp32.bin \ - --num_iteration 10 -``` - -#### Multi-node Multi-process Launch - -Run the following command on every node with the same rendezvous settings and -a distinct `node_rank`: - -```bash -./build/infini_run \ - --nnodes=2 \ - --nproc_per_node=4 \ - --node_rank=[rank_id] \ - --rdzv_endpoint=[master_addr]:29500 \ - --rdzv_id=[job_id] \ - ./build/llama3 \ - --device cuda \ - --input_bin data/llama3/tiny_shakespeare_train.bin \ - --llmc_filepath data/llama3/llama3.2_1B_fp32.bin \ - --num_iteration 10 \ - --tensor_parallel 2 \ - --pipeline_parallel 2 \ - --sequence_parallel -``` - -`--nproc_per_node` and `--nthread_per_process` can be combined. The total -training world size is: - -```text -world_size = nnodes × nproc_per_node × nthread_per_process -``` +Pass `--init_weights ` to load initial weights from an +InfiniTrain checkpoint (`model.ckpt` inside the directory), for example one +exported by PyTorch for numerical-alignment runs. + +```bash +./build/mnist \ + --model cnn \ + --device cuda \ + --dataset data/mnist \ + --init_weights data/cnn_align +``` + +##### GPT-2 124M + +```bash +./build/gpt2 \ + --device cuda \ + --input_bin data/gpt2/tiny_shakespeare_train.bin \ + --input_val_bin data/gpt2/tiny_shakespeare_val.bin \ + --tokenizer_bin data/gpt2/gpt2_tokenizer.bin \ + --llmc_filepath data/gpt2/gpt2_124M.bin \ + --num_iteration 10 +``` + +##### LLaMA 3.2 1B + +```bash +./build/llama3 \ + --device cuda \ + --input_bin data/llama3/tiny_shakespeare_train.bin \ + --input_val_bin data/llama3/tiny_shakespeare_val.bin \ + --llmc_filepath data/llama3/llama3.2_1B_fp32.bin \ + --num_iteration 10 +``` + +### Launch Modes + +GPT-2 and LLaMA training support both thread-based and process-based launches. +The examples below use LLaMA, but the same launch modes also apply to GPT-2. + +#### Direct Launch + +Running a model executable directly uses one process and one device by default. +Set `--nthread_per_process` to use multiple execution threads and devices in the +same process: + +```bash +./build/llama3 \ + --device cuda \ + --input_bin data/llama3/tiny_shakespeare_train.bin \ + --llmc_filepath data/llama3/llama3.2_1B_fp32.bin \ + --nthread_per_process 8 \ + --num_iteration 10 +``` + +#### Single-node Multi-process Launch + +Use `infini_run` to start multiple training processes on one node. Each process +uses one execution thread by default: + +```bash +./build/infini_run \ + --nnodes=1 \ + --nproc_per_node=8 \ + ./build/llama3 \ + --device cuda \ + --input_bin data/llama3/tiny_shakespeare_train.bin \ + --llmc_filepath data/llama3/llama3.2_1B_fp32.bin \ + --num_iteration 10 +``` + +#### Multi-node Multi-process Launch + +Run the following command on every node with the same rendezvous settings and +a distinct `node_rank`: + +```bash +./build/infini_run \ + --nnodes=2 \ + --nproc_per_node=4 \ + --node_rank=[rank_id] \ + --rdzv_endpoint=[master_addr]:29500 \ + --rdzv_id=[job_id] \ + ./build/llama3 \ + --device cuda \ + --input_bin data/llama3/tiny_shakespeare_train.bin \ + --llmc_filepath data/llama3/llama3.2_1B_fp32.bin \ + --num_iteration 10 \ + --tensor_parallel 2 \ + --pipeline_parallel 2 \ + --sequence_parallel +``` + +`--nproc_per_node` and `--nthread_per_process` can be combined. The total +training world size is: + +```text +world_size = nnodes × nproc_per_node × nthread_per_process +``` ### Parallelism Strategies -#### Distributed Data Parallelism (DDP) - -For a direct launch with TP and PP disabled, the following starts eight -data-parallel workers in one process: - -```bash ---nthread_per_process 8 # 8-way DDP when TP=1 and PP=1 -``` - -For all launch modes, the data-parallel size is derived from the total world -size after accounting for tensor and pipeline parallelism: - -```text -data_parallel_size = world_size / (tensor_parallel × pipeline_parallel) -``` +#### Distributed Data Parallelism (DDP) + +For a direct launch with TP and PP disabled, the following starts eight +data-parallel workers in one process: + +```bash +--nthread_per_process 8 # 8-way DDP when TP=1 and PP=1 +``` + +For all launch modes, the data-parallel size is derived from the total world +size after accounting for tensor and pipeline parallelism: + +```text +data_parallel_size = world_size / (tensor_parallel × pipeline_parallel) +``` #### Tensor Parallelism (TP) @@ -316,4 +340,4 @@ Multiple parallelism strategies (DDP, TP, SP, PP) can be freely combined to scal optimizations. Integrated a CTest + GTest based testing infrastructure to strengthen the - framework's automated test workflow. + framework's automated test workflow. From f801f748c4132a8e6ee48f795351de86111dce37 Mon Sep 17 00:00:00 2001 From: gavin-richie <2904793019@qq.com> Date: Mon, 7 Sep 2026 09:31:52 +0800 Subject: [PATCH 09/14] fix: sum the correct axis in the CUDA LinearBackwardBias kernel ReduceColumnsKernel summed along the wrong axis: it indexed the (bs, out_features) grad_output as if it were (out_features, bs), which happens to be correct only for bs == 1. Bias gradients on CUDA were therefore transposed sums and diverged from PyTorch for larger batches. The kernel now reduces rows within each column block for a row-major (num_rows, num_cols) input, and LinearBackwardBias passes (bs, out_features). Extend the linear backward test to check the bias gradient values, which previously only asserted the result size. --- infini_train/src/kernels/cuda/linear.cu | 11 ++++++----- tests/autograd/test_autograd_linear_backward.cc | 6 ++++++ 2 files changed, 12 insertions(+), 5 deletions(-) diff --git a/infini_train/src/kernels/cuda/linear.cu b/infini_train/src/kernels/cuda/linear.cu index 1b4c18190..34ea8c30a 100644 --- a/infini_train/src/kernels/cuda/linear.cu +++ b/infini_train/src/kernels/cuda/linear.cu @@ -135,23 +135,24 @@ std::shared_ptr LinearForward(const std::shared_ptr &input, cons return output; } +// Sums each column of the row-major (num_rows, num_cols) matrix in input; one block per column. template __global__ void ReduceColumnsKernel(const TIn *__restrict__ input, TOut *__restrict__ output, int num_rows, int num_cols) { using BlockReduce = cub::BlockReduce; __shared__ typename BlockReduce::TempStorage temp_storage; - int row = blockIdx.x; + int col = blockIdx.x; float sum = 0.0f; - for (int col = threadIdx.x; col < num_cols; col += blockDim.x) { + for (int row = threadIdx.x; row < num_rows; row += blockDim.x) { sum += common::cuda::Cast(input[row * num_cols + col]); } float reduced = BlockReduce(temp_storage).Sum(sum); if (threadIdx.x == 0) { - output[row] = reduced; + output[col] = reduced; } } @@ -309,13 +310,13 @@ std::shared_ptr LinearBackwardBias(const std::shared_ptr &grad_o DISPATCH_CASE(WRAP({ ReduceColumnsKernel<<>>( static_cast(grad_output->DataPtr()), - static_cast(grad_bias->DataPtr()), out_features, bs); + static_cast(grad_bias->DataPtr()), bs, out_features); }), DataType::kFLOAT32) DISPATCH_CASE(WRAP({ ReduceColumnsKernel<<>>( static_cast(grad_output->DataPtr()), - static_cast(grad_bias->DataPtr()), out_features, bs); + static_cast(grad_bias->DataPtr()), bs, out_features); }), DataType::kBFLOAT16) } diff --git a/tests/autograd/test_autograd_linear_backward.cc b/tests/autograd/test_autograd_linear_backward.cc index ba0f6fe1b..8e41c2eb3 100644 --- a/tests/autograd/test_autograd_linear_backward.cc +++ b/tests/autograd/test_autograd_linear_backward.cc @@ -25,6 +25,12 @@ TEST_P(AutogradLinearBackwardTest, LinearBackward) { grad->Fill(1.0f); auto grad_inputs = linear_fn->Backward({grad}); EXPECT_EQ(grad_inputs.size(), 3); + // With an all-ones grad, the bias gradient is the per-column sum of the (2, 4) grad output. + test::ExpectTensorFloatEqual(grad_inputs[2], {2.0f, 2.0f, 2.0f, 2.0f}); + // grad_input = grad * weight^T summed over output features; grad_weight = grad^T * input. + test::ExpectTensorFloatEqual(grad_inputs[0], {4.0f, 4.0f, 4.0f, 4.0f, 4.0f, 4.0f}); + test::ExpectTensorFloatEqual(grad_inputs[1], + {2.0f, 2.0f, 2.0f, 2.0f, 2.0f, 2.0f, 2.0f, 2.0f, 2.0f, 2.0f, 2.0f, 2.0f}); } TEST_P(AutogradLinearBackwardTest, LinearBackwardNoBias) { From ed7e9120dcc865bfb57f9b9c311623741ca9712e Mon Sep 17 00:00:00 2001 From: gavin-richie <2904793019@qq.com> Date: Mon, 7 Sep 2026 09:31:52 +0800 Subject: [PATCH 10/14] feat: add optional metrics logging to the mnist demo Add a --metrics_file flag that appends train and test metrics as JSON lines for visualization: a train_step line every 10 steps and an epoch_end line with train loss, test loss, and test accuracy after each epoch. Evaluation is extracted into an Evaluate() helper and now runs after every epoch under no_grad. --- example/mnist/main.cc | 100 ++++++++++++++++++++++++++++-------------- 1 file changed, 66 insertions(+), 34 deletions(-) diff --git a/example/mnist/main.cc b/example/mnist/main.cc index 95536050a..b86b4276a 100644 --- a/example/mnist/main.cc +++ b/example/mnist/main.cc @@ -1,9 +1,11 @@ #include #include #include +#include #include #include #include +#include #include #include "gflags/gflags.h" @@ -26,6 +28,7 @@ DEFINE_int32(num_epoch, 1, "num epochs"); DEFINE_double(lr, 0.01, "learning rate"); DEFINE_string(device, "cpu", "device type (cpu/cuda)"); DEFINE_string(init_weights, "", "checkpoint dir to load initial weights from (e.g. exported by PyTorch)"); +DEFINE_string(metrics_file, "", "append train/test metrics as JSON lines to this file for visualization"); using namespace infini_train; @@ -35,12 +38,59 @@ constexpr int kNumClasses = 10; constexpr char kDeviceCPU[] = "cpu"; constexpr char kDeviceCUDA[] = "cuda"; -}; // namespace +constexpr char kModelMLP[] = "mlp"; +constexpr char kModelCNN[] = "cnn"; + +// Appends one JSON line to the metrics file; a no-op when the flag is empty. +void AppendMetrics(const std::string &metrics_file, const std::string &json_line) { + if (metrics_file.empty()) { + return; + } + std::ofstream ofs(metrics_file, std::ios::app); + ofs << json_line << "\n"; +} +} // 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 == "mlp" || value == "cnn"; }); +DEFINE_validator(model, + [](const char *, const std::string &value) { return value == kModelMLP || value == kModelCNN; }); + +// Runs the test set under no_grad; returns {average test loss, accuracy}. +std::pair Evaluate(nn::Module &network, nn::CrossEntropyLoss &loss_fn, DataLoader &test_dataloader, + const Device &device) { + Device cpu_device = Device(); + std::vector test_losses; + int correct = 0; + int total = 0; + autograd::NoGradGuard no_grad; + 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.Forward({new_image}); + auto output_cpu = outputs[0]->To(cpu_device); + auto loss = loss_fn.Forward({outputs[0], new_label}); + auto loss_cpu = loss[0]->To(cpu_device); + + const int batch_size = output_cpu.Dims()[0]; + for (int batch_idx = 0; batch_idx < batch_size; ++batch_idx) { + auto label_index = reinterpret_cast(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 = static_cast(std::accumulate(test_losses.begin(), test_losses.end(), 0.0) + / std::max(test_losses.size(), 1)); + return {avg_loss, static_cast(correct) / total}; +} int main(int argc, char *argv[]) { gflags::ParseCommandLineFlags(&argc, &argv, true); @@ -67,6 +117,7 @@ int main(int argc, char *argv[]) { loss_fn.To(device); auto optimizer = optimizers::SGD(network->Parameters(), FLAGS_lr); + int64_t global_step = 0; for (int epoch = 0; epoch < FLAGS_num_epoch; ++epoch) { int train_idx = 0; float total_loss = 0.0; @@ -92,50 +143,31 @@ int main(int argc, char *argv[]) { LOG(ERROR) << "epoch: " << epoch << ", step: " << train_idx << ", [" << train_idx * FLAGS_bs << "/" << train_dataset->Size() << "] " << " loss: " << current_loss; + AppendMetrics(FLAGS_metrics_file, + std::format("{{\"type\": \"train_step\", \"epoch\": {}, \"step\": {}, \"loss\": {:.6f}}}", + epoch, global_step, current_loss)); } optimizer.Step(); train_idx += 1; + global_step += 1; } const auto epoch_end = std::chrono::high_resolution_clock::now(); const double duration_us = std::chrono::duration(epoch_end - epoch_start).count(); + const float train_loss = total_loss / train_idx; 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, + epoch, FLAGS_num_epoch - 1, train_loss, FLAGS_lr, duration_us / 1e3f, train_dataset->Size() / (duration_us / 1e6)); - } - { - 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->Forward({new_image}); - auto output_cpu = outputs[0]->To(cpu_device); - auto loss = loss_fn.Forward({outputs[0], new_label}); - auto loss_cpu = loss[0]->To(cpu_device); - - const int batch_size = output_cpu.Dims()[0]; - for (int batch_idx = 0; batch_idx < batch_size; ++batch_idx) { - auto label_index = reinterpret_cast(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) << std::format("test | test loss {:.6f} | test accuracy {:.4f} ({}/{})", avg_loss, - static_cast(correct) / total, correct, total); + const auto [test_loss, test_accuracy] = Evaluate(*network, loss_fn, test_dataloader, device); + LOG(ERROR) << std::format("epoch {:2d} | test loss {:.6f} | test accuracy {:.4f}", epoch, test_loss, + test_accuracy); + AppendMetrics(FLAGS_metrics_file, + std::format("{{\"type\": \"epoch_end\", \"epoch\": {}, \"train_loss\": {:.6f}, " + "\"test_loss\": {:.6f}, \"test_accuracy\": {:.6f}}}", + epoch, train_loss, test_loss, test_accuracy)); } gflags::ShutDownCommandLineFlags(); From 76c96001a44f276a205277835738ca6c2a3517e5 Mon Sep 17 00:00:00 2001 From: gavin-richie <2904793019@qq.com> Date: Mon, 7 Sep 2026 09:54:10 +0800 Subject: [PATCH 11/14] chore: add swanlab metrics upload script Upload the JSON lines written by the mnist demo's --metrics_file flag to SwanLab via the official swanlab package. The API key is read from the SWANLAB_API_KEY environment variable and is never stored in the repository. --- README.md | 10 +++++ scripts/swanlab_upload.py | 87 +++++++++++++++++++++++++++++++++++++++ 2 files changed, 97 insertions(+) create mode 100644 scripts/swanlab_upload.py diff --git a/README.md b/README.md index 8fa4958b6..91bda8172 100644 --- a/README.md +++ b/README.md @@ -169,6 +169,16 @@ exported by PyTorch for numerical-alignment runs. --init_weights data/cnn_align ``` +Pass `--metrics_file ` to append per-step training losses and per-epoch +test loss / accuracy as JSON lines, which can be uploaded to +[SwanLab](https://swanlab.cn) for training visualization: + +```bash +./build/mnist --model cnn --device cuda --dataset data/mnist --metrics_file metrics.jsonl +SWANLAB_API_KEY= python3 scripts/swanlab_upload.py \ + --metrics metrics.jsonl --name cnn-cuda-3epoch-lr0.05 --model cnn --device cuda --lr 0.05 +``` + ##### GPT-2 124M ```bash diff --git a/scripts/swanlab_upload.py b/scripts/swanlab_upload.py new file mode 100644 index 000000000..82834a52d --- /dev/null +++ b/scripts/swanlab_upload.py @@ -0,0 +1,87 @@ +"""Upload MNIST demo metrics (JSON lines from --metrics_file) to SwanLab. + +The demo writes one JSON object per line: + {"type": "train_step", "epoch": E, "step": S, "loss": L} + {"type": "epoch_end", "epoch": E, "train_loss": TL, "test_loss": XL, "test_accuracy": XA} + +Usage: + export SWANLAB_API_KEY= + python3 scripts/swanlab_upload.py --metrics /tmp/runs/cnn_cuda_3ep.jsonl \ + --name cnn-cuda-3epoch --model cnn --device cuda --lr 0.05 +""" + +import argparse +import json +import os + +import swanlab + + +def main(): + parser = argparse.ArgumentParser() + parser.add_argument( + "--metrics", required=True, help="JSON lines file written by the mnist demo" + ) + parser.add_argument("--project", default="infinitrain-mnist-cnn") + parser.add_argument("--name", required=True, help="experiment name") + parser.add_argument("--model", default="", help="recorded in the experiment config") + parser.add_argument( + "--device", default="", help="recorded in the experiment config" + ) + parser.add_argument("--lr", default="", help="recorded in the experiment config") + args = parser.parse_args() + + if not os.environ.get("SWANLAB_API_KEY"): + raise SystemExit("SWANLAB_API_KEY is not set") + + rows = [] + with open(args.metrics) as f: + for line in f: + line = line.strip() + if line: + rows.append(json.loads(line)) + + train_rows = [r for r in rows if r.get("type") == "train_step"] + epoch_rows = [r for r in rows if r.get("type") == "epoch_end"] + if not train_rows and not epoch_rows: + raise SystemExit("no metrics found") + + steps_per_epoch = ( + max((r["step"] for r in train_rows), default=0) + 1 if train_rows else 0 + ) + + run = swanlab.init( + project=args.project, + experiment_name=args.name, + api_key=os.environ["SWANLAB_API_KEY"], + config={ + "model": args.model, + "device": args.device, + "learning_rate": args.lr, + "epochs": len(epoch_rows), + "batch_size": 64, + "dataset": "MNIST", + }, + ) + + for r in train_rows: + swanlab.log({"train/loss": r["loss"]}, step=r["step"]) + for r in epoch_rows: + step = (r["epoch"] + 1) * steps_per_epoch - 1 + metrics = { + "train/epoch_loss": r["train_loss"], + "test/loss": r["test_loss"], + "test/accuracy": r["test_accuracy"], + } + swanlab.log(metrics, step=step) + + swanlab.finish() + print(f"uploaded {len(train_rows)} train points and {len(epoch_rows)} epoch points") + try: + print(f"run url: {run.url}") + except AttributeError: + pass + + +if __name__ == "__main__": + main() From 15a626567715c96f373184025db13a834c09befa Mon Sep 17 00:00:00 2001 From: gavin-richie <2904793019@qq.com> Date: Wed, 9 Sep 2026 02:24:01 +0800 Subject: [PATCH 12/14] build: add sm_86 to the CUDA architecture list --- CMakeLists.txt | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/CMakeLists.txt b/CMakeLists.txt index 4d90c3490..eed4f8fee 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -104,7 +104,7 @@ if(USE_CUDA) file(GLOB_RECURSE CUDA_KERNELS ${PROJECT_SOURCE_DIR}/infini_train/src/*.cu) add_library(infini_train_cuda_kernels STATIC ${CUDA_KERNELS}) - set_target_properties(infini_train_cuda_kernels PROPERTIES CUDA_ARCHITECTURES "75;80;89;90") + set_target_properties(infini_train_cuda_kernels PROPERTIES CUDA_ARCHITECTURES "75;80;86;89;90") target_link_libraries(infini_train_cuda_kernels PUBLIC From bfb0333721abb4e2a38d540694ed9275f2f052bb Mon Sep 17 00:00:00 2001 From: gavin-richie <2904793019@qq.com> Date: Wed, 9 Sep 2026 03:01:30 +0800 Subject: [PATCH 13/14] feat: support DDP training in the mnist demo Wire the mnist demo into the existing DDP stack: initialize the parallel environment from the RANK/LOCAL_RANK/WORLD_SIZE variables set by tools/infini_run, place each rank on its own CUDA device, create the data- parallel process group, wrap the model with DistributedDataParallel (zero_stage 0, bucketed gradient all-reduce), and shard the training data with DistributedDataLoader. Per-rank training losses are averaged with an AllReduce for logging, parameters are broadcast from the global-rank-0 process at startup, and rank-0-only logging/metrics keep single-process behavior unchanged when no parallel environment is set. The per-epoch step count is capped at the per-rank floor so every rank runs the same number of collectives; the strided DistributedDataLoader would otherwise hand one rank an extra batch and deadlock NCCL on the last all-reduce. Single-process 2-GPU verification (2x RTX 3090, CUDA 12.9/NCCL cu12): DDP with 2 ranks x batch 32 reproduces the single-process batch-64 loss trajectory exactly (max |diff| 2.6e-5 over 938 steps), and a 3-epoch DDP run reaches 97.08% test accuracy at ~236k samples/s. --- example/mnist/main.cc | 115 +++++++++++++++++++++++++++++++----------- 1 file changed, 86 insertions(+), 29 deletions(-) diff --git a/example/mnist/main.cc b/example/mnist/main.cc index b86b4276a..4765d5aeb 100644 --- a/example/mnist/main.cc +++ b/example/mnist/main.cc @@ -16,6 +16,14 @@ #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/reduce_op_type.h" +#include "infini_train/include/nn/parallel/utils.h" #include "infini_train/include/optimizer.h" #include "example/mnist/dataset.h" @@ -23,10 +31,10 @@ DEFINE_string(dataset, "", "mnist dataset path"); DEFINE_string(model, "mlp", "model type (mlp/cnn)"); -DEFINE_int32(bs, 64, "batch size"); +DEFINE_int32(bs, 64, "batch size per rank"); DEFINE_int32(num_epoch, 1, "num epochs"); DEFINE_double(lr, 0.01, "learning rate"); -DEFINE_string(device, "cpu", "device type (cpu/cuda)"); +DEFINE_string(device, "cpu", "device type (cpu/cuda) for single-process runs"); DEFINE_string(init_weights, "", "checkpoint dir to load initial weights from (e.g. exported by PyTorch)"); DEFINE_string(metrics_file, "", "append train/test metrics as JSON lines to this file for visualization"); @@ -49,13 +57,6 @@ void AppendMetrics(const std::string &metrics_file, const std::string &json_line std::ofstream ofs(metrics_file, std::ios::app); ofs << json_line << "\n"; } -} // 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; }); // Runs the test set under no_grad; returns {average test loss, accuracy}. std::pair Evaluate(nn::Module &network, nn::CrossEntropyLoss &loss_fn, DataLoader &test_dataloader, @@ -92,27 +93,53 @@ std::pair Evaluate(nn::Module &network, nn::CrossEntropyLoss &loss return {avg_loss, static_cast(correct) / total}; } -int main(int argc, char *argv[]) { - gflags::ParseCommandLineFlags(&argc, &argv, true); - google::InitGoogleLogging(argv[0]); +int Train(const nn::parallel::Rank &rank) { + const bool is_main_rank = rank.GlobalRank() == 0; + const int ddp_world_size = nn::parallel::global::GetDataParallelSize(); + Device cpu_device = Device(); auto train_dataset = std::make_shared(FLAGS_dataset, true); - DataLoader train_dataloader(train_dataset, FLAGS_bs); - - // TODO(dcj): Add sampler & eval dataloader later. + // Strided batch sharding across DDP ranks (a no-op when ddp_world_size == 1). + DistributedDataLoader train_dataloader(train_dataset, FLAGS_bs, rank.GlobalRank(), ddp_world_size); + // The strided sharding can hand one rank one more batch than the other when the batch count + // is not divisible by the world size; cap every rank at the same step count (dropping the + // incomplete trailing global batch), otherwise the extra collective deadlocks the peer. + const size_t train_samples_per_rank = train_dataset->Size() / ddp_world_size; + const size_t steps_per_epoch = train_samples_per_rank / FLAGS_bs; + + // TODO(dcj): Add sampler & eval dataloader later. Each rank evaluates the full test set. auto test_dataset = std::make_shared(FLAGS_dataset, false); DataLoader test_dataloader(test_dataset, FLAGS_bs); auto network = CreateMNISTNetwork(FLAGS_model); Device device = FLAGS_device == kDeviceCPU ? Device() : Device(Device::DeviceType::kCUDA, 0); - Device cpu_device = Device(); + const nn::parallel::ProcessGroup *ddp_pg = nullptr; + if (rank.IsParallel()) { + CHECK(device.IsCUDA()) << "DDP training requires --device cuda"; + device = Device(Device::DeviceType::kCUDA, nn::parallel::global::GetDeviceIndex(rank.thread_rank())); + } network->To(device); if (!FLAGS_init_weights.empty()) { + // Every rank loads the same file, so the initial weights are identical across ranks. TrainerState state; Checkpoint::Load(FLAGS_init_weights, *network, /*optimizer=*/nullptr, state, /*lr_scheduler=*/nullptr); } + // The DP process group must exist before constructing DistributedDataParallel, and all + // parameters must already live on the rank's device (see example/gpt2/main.cc). + if (ddp_world_size > 1) { + auto *pg_factory = nn::parallel::ProcessGroupFactory::Instance(device.type()); + ddp_pg = pg_factory->GetOrCreate(nn::parallel::GetDataParallelProcessGroupName(rank.GlobalRank()), + nn::parallel::GetDataParallelGroupRanks(rank.GlobalRank())); + // NOTE: Complete all device conversions before wrapping, otherwise gradient hooks are lost. + auto ddp_config = nn::parallel::DistributedDataParallelConfig{.zero_stage = 0}; + network = std::make_shared(network, rank, ddp_config); + // Match torch DDP semantics: broadcast the parameters of the global-rank-0 process so all + // ranks start from identical weights even if the initialization were rank-dependent. + ddp_pg->Broadcast(network->Parameters(), /*root_rank_in_group=*/0); + } + auto loss_fn = nn::CrossEntropyLoss(); loss_fn.To(device); auto optimizer = optimizers::SGD(network->Parameters(), FLAGS_lr); @@ -125,21 +152,31 @@ int main(int argc, char *argv[]) { const auto epoch_start = std::chrono::high_resolution_clock::now(); for (const auto &[image, label] : train_dataloader) { + if (static_cast(train_idx) >= steps_per_epoch) { + break; + } 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 before the forward pass: with DDP's gradient-as-bucket-view the forward's + // PrepareForBackward binds parameter gradients to the bucket slices, and zeroing + // afterwards would break that binding (matching example/gpt2/main.cc). optimizer.ZeroGrad(); + auto outputs = network->Forward({new_image}); + 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 - // between forward and backward. + // between forward and backward. With DDP the gradient all-reduce runs inside Backward(). + if (ddp_pg != nullptr) { + nn::parallel::function::AllReduce(loss[0], nn::parallel::function::ReduceOpType::kAvg, ddp_pg); + } auto loss_cpu = loss[0]->To(cpu_device); float current_loss = static_cast(loss_cpu.DataPtr())[0]; total_loss += current_loss; - if (train_idx % kNumItersOfOutputDuration == 0) { + if (is_main_rank && train_idx % kNumItersOfOutputDuration == 0) { LOG(ERROR) << "epoch: " << epoch << ", step: " << train_idx << ", [" << train_idx * FLAGS_bs << "/" << train_dataset->Size() << "] " << " loss: " << current_loss; @@ -157,19 +194,39 @@ int main(int argc, char *argv[]) { const double duration_us = std::chrono::duration(epoch_end - epoch_start).count(); const float train_loss = total_loss / train_idx; - LOG(ERROR) << std::format("epoch {:2d}/{} | train loss {:.6f} | lr {:.2e} | ({:.2f} ms | {:.0f} samples/s)", - epoch, FLAGS_num_epoch - 1, train_loss, FLAGS_lr, duration_us / 1e3f, - train_dataset->Size() / (duration_us / 1e6)); - const auto [test_loss, test_accuracy] = Evaluate(*network, loss_fn, test_dataloader, device); - LOG(ERROR) << std::format("epoch {:2d} | test loss {:.6f} | test accuracy {:.4f}", epoch, test_loss, - test_accuracy); - AppendMetrics(FLAGS_metrics_file, - std::format("{{\"type\": \"epoch_end\", \"epoch\": {}, \"train_loss\": {:.6f}, " - "\"test_loss\": {:.6f}, \"test_accuracy\": {:.6f}}}", - epoch, train_loss, test_loss, test_accuracy)); + if (is_main_rank) { + LOG(ERROR) << std::format("epoch {:2d}/{} | train loss {:.6f} | lr {:.2e} | ({:.2f} ms | {:.0f} samples/s)", + epoch, FLAGS_num_epoch - 1, train_loss, FLAGS_lr, duration_us / 1e3f, + train_dataset->Size() * ddp_world_size / (duration_us / 1e6)); + LOG(ERROR) << std::format("epoch {:2d} | test loss {:.6f} | test accuracy {:.4f}", epoch, test_loss, + test_accuracy); + AppendMetrics(FLAGS_metrics_file, + std::format("{{\"type\": \"epoch_end\", \"epoch\": {}, \"train_loss\": {:.6f}, " + "\"test_loss\": {:.6f}, \"test_accuracy\": {:.6f}}}", + epoch, train_loss, test_loss, test_accuracy)); + } } + return 0; +} +} // namespace + +int main(int argc, char *argv[]) { + gflags::ParseCommandLineFlags(&argc, &argv, true); + google::InitGoogleLogging(argv[0]); + + // Rank/world size come from the RANK / LOCAL_RANK / WORLD_SIZE / LOCAL_WORLD_SIZE environment + // variables, which tools/infini_run sets for each spawned process. Without them this is a + // plain single-process run. + 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(), /*threads_per_process=*/1); + + Train(rank); + gflags::ShutDownCommandLineFlags(); google::ShutdownGoogleLogging(); From a8674d7e9b7d6c2cf7262477b8810b83e2d9fdaf Mon Sep 17 00:00:00 2001 From: gavin-richie <2904793019@qq.com> Date: Wed, 9 Sep 2026 03:02:36 +0800 Subject: [PATCH 14/14] docs: document mnist DDP launch --- README.md | 16 ++++++++++++++++ 1 file changed, 16 insertions(+) diff --git a/README.md b/README.md index 91bda8172..1f5b32d5e 100644 --- a/README.md +++ b/README.md @@ -238,6 +238,22 @@ uses one execution thread by default: --num_iteration 10 ``` +The MNIST example supports the same DDP launch (`--device cuda` is required; +each rank trains on its own GPU and logs only on rank 0): + +```bash +./build/infini_run \ + --nnodes=1 \ + --nproc_per_node=2 \ + ./build/mnist \ + --model cnn \ + --device cuda \ + --dataset data/mnist \ + --num_epoch 3 \ + --lr 0.1 \ + --metrics_file metrics_ddp.jsonl +``` + #### Multi-node Multi-process Launch Run the following command on every node with the same rendezvous settings and