diff --git a/README.md b/README.md index abd8070b2..66f0e40e9 100644 --- a/README.md +++ b/README.md @@ -141,10 +141,26 @@ The generated files can be passed directly to the corresponding executables: ```bash ./build/mnist \ + --model cnn \ --device cpu \ --dataset data/mnist ``` +`--model` selects the network, `cnn` (default) or `mlp`. Launching through +`infini_run` distributes the training across processes with DDP; each process +picks its GPU from `LOCAL_RANK` automatically: + +```bash +./build/infini_run \ + --nnodes=1 \ + --nproc_per_node=2 \ + ./build/mnist \ + --model cnn \ + --device cuda \ + --dataset data/mnist \ + --num_epoch 3 +``` + ##### GPT-2 124M ```bash diff --git a/example/mnist/cnn_net.h b/example/mnist/cnn_net.h new file mode 100644 index 000000000..02e92d7d0 --- /dev/null +++ b/example/mnist/cnn_net.h @@ -0,0 +1,44 @@ +#pragma once + +#include +#include +#include + +#include "glog/logging.h" + +#include "infini_train/include/nn/modules/activations.h" +#include "infini_train/include/nn/modules/container.h" +#include "infini_train/include/nn/modules/conv.h" +#include "infini_train/include/nn/modules/linear.h" +#include "infini_train/include/nn/modules/module.h" +#include "infini_train/include/tensor.h" + +// Small CNN classifier for the MNIST demo. Structure matches the reference: +// Conv2d(1,16,3) -> ReLU -> Conv2d(16,32,3) -> ReLU -> Flatten -> Linear(18432,10). +// The DataLoader::Stack helper flattens each image into a [N, 784] matrix, so the +// Forward entry restores the (N, 1, 28, 28) spatial layout before the first conv. +class MnistCnn : public infini_train::nn::Module { +public: + MnistCnn() { + std::vector> layers; + // Two 3x3 valid convs shrink 28 -> 26 -> 24; no pooling in the reference net. + layers.push_back(std::make_shared(1, 16, 3)); + layers.push_back(std::make_shared()); + layers.push_back(std::make_shared(16, 32, 3)); + layers.push_back(std::make_shared()); + modules_["sequential"] = std::make_shared(std::move(layers)); + // 32 * 24 * 24 = 18432 flattened features into the 10-class head. + modules_["linear"] = std::make_shared(32 * 24 * 24, 10); + } + + std::vector> + Forward(const std::vector> &x) override { + CHECK_EQ(x.size(), 1); + // Restore the batch dimension from the runtime shape, then reshape the flattened + // [N, 784] input to (N, 1, 28, 28). View is element-order preserving. + const auto batch = x[0]->Dims()[0]; + auto x_view = x[0]->View({batch, 1, 28, 28}); + auto x_feat = (*modules_["sequential"])({x_view})[0]->Flatten(1, -1); + return (*modules_["linear"])({x_feat}); + } +}; diff --git a/example/mnist/dataset.cc b/example/mnist/dataset.cc index ee683f6d4..fc7988311 100644 --- a/example/mnist/dataset.cc +++ b/example/mnist/dataset.cc @@ -89,7 +89,12 @@ MNISTDataset::MNISTDataset(const std::string &dataset, bool train) std::format("{}/{}-labels-idx1-ubyte", dataset, train ? kTrainPrefix : kTestPrefix))), image_dims_(image_file_.dims.begin() + 1, image_file_.dims.end()), label_dims_(label_file_.dims.begin() + 1, label_file_.dims.end()), - image_size_in_bytes_(kSN3TypeToSize.at(image_file_.type) + // The image tensor is normalized to FLOAT32 in the constructor body, so the per-sample + // byte stride must use the float element size, not the on-disk UINT8 size. After the + // first sample, the UINT8 stride reads the wrong bytes: most views are a misaligned + // mix of real pixels, and some coincide with another sample's exact copy but are still + // paired with an unrelated label. + image_size_in_bytes_(infini_train::kDataTypeToSize.at(DataType::kFLOAT32) * std::accumulate(image_dims_.begin(), image_dims_.end(), 1, std::multiplies())), label_size_in_bytes_(kSN3TypeToSize.at(label_file_.type) * std::accumulate(label_dims_.begin(), label_dims_.end(), 1, std::multiplies())) { diff --git a/example/mnist/main.cc b/example/mnist/main.cc index 7744e0947..69011d7c0 100644 --- a/example/mnist/main.cc +++ b/example/mnist/main.cc @@ -1,26 +1,41 @@ +#include #include #include #include #include #include #include +#include #include #include "gflags/gflags.h" #include "glog/logging.h" +#include "infini_train/include/autograd/grad_mode.h" #include "infini_train/include/dataloader.h" #include "infini_train/include/device.h" #include "infini_train/include/nn/modules/loss.h" +#include "infini_train/include/nn/parallel/ddp/distributed_data_parallel.h" +#include "infini_train/include/nn/parallel/ddp/distributed_data_parallel_config.h" +#include "infini_train/include/nn/parallel/global.h" +#include "infini_train/include/nn/parallel/parallel_functional.h" +#include "infini_train/include/nn/parallel/process_group.h" +#include "infini_train/include/nn/parallel/rank.h" +#include "infini_train/include/nn/parallel/utils.h" #include "infini_train/include/optimizer.h" +#include "infini_train/include/tensor.h" +#include "example/mnist/cnn_net.h" #include "example/mnist/dataset.h" #include "example/mnist/net.h" DEFINE_string(dataset, "", "mnist dataset path"); +DEFINE_string(model, "cnn", "model type (mlp/cnn)"); DEFINE_int32(bs, 64, "batch size"); -DEFINE_int32(num_epoch, 1, "num epochs"); -DEFINE_double(lr, 0.01, "learning rate"); +// Defaults are tuned for the CNN demo (default --model=cnn) to reach ~97.8% test accuracy. +// The MLP reaches ~92% at these defaults; pass more epochs (e.g. --num_epoch=20) to reach ~95%. +DEFINE_int32(num_epoch, 3, "num epochs"); +DEFINE_double(lr, 0.1, "learning rate"); DEFINE_string(device, "cpu", "device type (cpu/cuda)"); using namespace infini_train; @@ -31,30 +46,131 @@ constexpr int kNumClasses = 10; constexpr char kDeviceCPU[] = "cpu"; constexpr char kDeviceCUDA[] = "cuda"; +constexpr char kModelMLP[] = "mlp"; +constexpr char kModelCNN[] = "cnn"; }; // namespace DEFINE_validator(device, [](const char *, const std::string &value) { return value == kDeviceCPU || value == kDeviceCUDA; }); +DEFINE_validator(model, + [](const char *, const std::string &value) { return value == kModelMLP || value == kModelCNN; }); int main(int argc, char *argv[]) { gflags::ParseCommandLineFlags(&argc, &argv, true); google::InitGoogleLogging(argv[0]); + CHECK_GT(FLAGS_bs, 0) << "--bs must be a positive batch size (got " << FLAGS_bs << ")"; + + // Consume the WORLD_SIZE/RANK/LOCAL_RANK env injected by infini_run; without the launcher + // every size defaults to 1 and the single-process path below is unchanged. + nn::parallel::global::InitAllEnv(/*nthread_per_process=*/1, /*tensor_parallel_size=*/1, + /*sequence_parallel_enabled=*/false, /*pipeline_parallel_size=*/1, + /*virtual_pipeline_parallel_size=*/1); + nn::parallel::Rank rank(nn::parallel::global::GetGlobalProcRank(), /*thread_rank=*/0, + nn::parallel::global::GetNprocPerNode(), nn::parallel::global::GetNthreadPerProc()); + const int ddp_world_size = nn::parallel::global::GetDataParallelSize(); + auto train_dataset = std::make_shared(FLAGS_dataset, true); - DataLoader train_dataloader(train_dataset, FLAGS_bs); + + Device device; + const nn::parallel::ProcessGroup *ddp_pg = nullptr; + int ddp_rank = 0; + if (rank.IsParallel()) { + CHECK_EQ(FLAGS_device, kDeviceCUDA) << "Distributed training requires --device=cuda"; + // One GPU per process, taken from the LOCAL_RANK assigned by the launcher. + device = Device(Device::DeviceType::kCUDA, nn::parallel::global::GetDeviceIndex(rank.thread_rank())); + auto *pg_factory = nn::parallel::ProcessGroupFactory::Instance(device.type()); + ddp_pg = pg_factory->GetOrCreate(nn::parallel::GetDataParallelProcessGroupName(rank.GlobalRank()), + nn::parallel::GetDataParallelGroupRanks(rank.GlobalRank())); + ddp_rank = ddp_pg->GetGroupRank(rank.GlobalRank()); + } else { + device = FLAGS_device == kDeviceCPU ? Device() : Device(Device::DeviceType::kCUDA, 0); + } + + // Sharded batches per rank under DDP; the legacy loader keeps the single-process path. + std::optional ddp_train_dataloader; + std::optional plain_train_dataloader; + if (ddp_pg != nullptr) { + // The loss AllReduce runs step-for-step on every rank, so all ranks must see the same + // batch count or a collective would hang out of step. DistributedDataLoader derives + // that count from the global batch size, which keeps every rank aligned. + ddp_train_dataloader.emplace(train_dataset, FLAGS_bs, ddp_rank, ddp_world_size); + } else { + plain_train_dataloader.emplace(train_dataset, FLAGS_bs); + } + const DataLoader &train_loader = ddp_pg != nullptr ? *ddp_train_dataloader : *plain_train_dataloader; // TODO(dcj): Add sampler & eval dataloader later. + // The test loader stays unsharded so every rank reports metrics over the full test set. auto test_dataset = std::make_shared(FLAGS_dataset, false); DataLoader test_dataloader(test_dataset, FLAGS_bs); - auto network = MNIST(); - Device device = FLAGS_device == kDeviceCPU ? Device() : Device(Device::DeviceType::kCUDA, 0); + std::shared_ptr network; + if (FLAGS_model == kModelCNN) { + network = std::make_shared(); + } else { + network = std::make_shared(); + } Device cpu_device = Device(); - network.To(device); + network->To(device); + + auto loss_fn = std::make_shared(); + loss_fn->To(device); + + // Wrap with DDP only after all device conversions: a later .To() recreates parameter + // tensors and would leave the gradient hooks registered at wrap time dangling. + if (ddp_pg != nullptr) { + network = std::make_shared( + network, rank, nn::parallel::DistributedDataParallelConfig{}); + // Keep every replica starting from the same state; parameter broadcast before training + // is part of the DDP contract and left to the caller by the wrapper. Parameters only: + // the demo network carries no buffers, which PyTorch would sync alongside them. + ddp_pg->Broadcast(network->Parameters(), /*root_rank_in_group=*/0); + } + + auto optimizer = optimizers::SGD(network->Parameters(), FLAGS_lr); + + // Steps per epoch, needed by the progress line's "[step k/total]" field. Counted through the + // loader's own iterators so the total follows its batching and sharding rules. + int num_train_iters = 0; + for (auto it = train_loader.begin(); it != train_loader.end(); ++it) { ++num_train_iters; } + + // Runs the full test set on every rank, once per epoch, so a single run shows how the test + // metrics evolve instead of reporting them only after the last epoch. no_grad: evaluation + // never backpropagates, and a forward-only graph would leave the parameters' grad accumulators + // primed with a dependency count the next training step can never satisfy, which would silently + // stop their gradient accumulation. + auto evaluate = [&](int epoch) { + autograd::NoGradGuard no_grad; + std::vector test_losses; + int correct = 0; + int total = 0; + for (const auto &[image, label] : test_dataloader) { + auto new_image = std::make_shared(image->To(device)); + auto new_label = std::make_shared(label->To(device)); + + auto label_cpu = label->To(cpu_device); + auto outputs = (*network)({new_image}); + auto output_cpu = outputs[0]->To(cpu_device); + auto loss = (*loss_fn)({outputs[0], new_label}); + auto loss_cpu = loss[0]->To(cpu_device); - auto loss_fn = nn::CrossEntropyLoss(); - loss_fn.To(device); - auto optimizer = optimizers::SGD(network.Parameters(), FLAGS_lr); + const int batch_size = output_cpu.Dims()[0]; + for (int batch_idx = 0; batch_idx < batch_size; ++batch_idx) { + auto label_index = reinterpret_cast(label_cpu.DataPtr())[batch_idx]; + const auto *output_values = static_cast(output_cpu.DataPtr()) + batch_idx * kNumClasses; + const int output_index = std::max_element(output_values, output_values + kNumClasses) - output_values; + if (output_index == label_index) { + ++correct; + } + } + total += batch_size; + test_losses.push_back(static_cast(loss_cpu.DataPtr())[0]); + } + const auto avg_loss = std::accumulate(test_losses.begin(), test_losses.end(), 0.0) / test_losses.size(); + LOG(ERROR) << "epoch " << epoch << " | Total: " << total << ", Correct: " << correct + << ", Accuracy: " << static_cast(correct) / total << ", AverageLoss: " << avg_loss; + }; for (int epoch = 0; epoch < FLAGS_num_epoch; ++epoch) { int train_idx = 0; @@ -62,25 +178,39 @@ int main(int argc, char *argv[]) { const auto epoch_start = std::chrono::high_resolution_clock::now(); - for (const auto &[image, label] : train_dataloader) { + for (const auto &[image, label] : train_loader) { auto new_image = std::make_shared(image->To(device)); auto new_label = std::make_shared(label->To(device)); - auto outputs = network.Forward({new_image}); + // Zero grads before forward: DDP rebinds param.grad to its bucket view during forward. optimizer.ZeroGrad(); - auto loss = loss_fn.Forward({outputs[0], new_label}); + auto outputs = (*network)({new_image}); + + auto loss = (*loss_fn)({outputs[0], new_label}); loss[0]->Backward(); // Defer the loss D2H copy until after backward; reading it earlier would synchronize CUDA // between forward and backward. auto loss_cpu = loss[0]->To(cpu_device); float current_loss = static_cast(loss_cpu.DataPtr())[0]; + if (ddp_pg != nullptr) { + // Average the per-rank loss so the logged value matches the global batch. With an + // uneven trailing batch the equal-weight average is a bounded per-step approximation + // until a sampler lands. + auto loss_stat + = std::make_shared(¤t_loss, std::vector{}, DataType::kFLOAT32, device); + nn::parallel::function::AllReduce(loss_stat, nn::parallel::function::ReduceOpType::kAvg, ddp_pg); + auto loss_stat_cpu = loss_stat->To(cpu_device); + current_loss = static_cast(loss_stat_cpu.DataPtr())[0]; + } total_loss += current_loss; if (train_idx % kNumItersOfOutputDuration == 0) { - LOG(ERROR) << "epoch: " << epoch << ", [" << train_idx * FLAGS_bs << "/" << train_dataset->Size() - << "] " - << " loss: " << current_loss; + // step is 1-based while samples is the count consumed before it, so both fields + // describe the same instant: the start of this step. + LOG(ERROR) << std::format("epoch {:2d}, [step {:4d}/{}] [samples {}/{}] loss: {:.6f}", epoch, + train_idx + 1, num_train_iters, train_idx * FLAGS_bs * ddp_world_size, + train_dataset->Size(), current_loss); } optimizer.Step(); @@ -93,37 +223,9 @@ int main(int argc, char *argv[]) { LOG(ERROR) << std::format("epoch {:2d}/{} | train loss {:.6f} | lr {:.2e} | ({:.2f} ms | {:.0f} samples/s)", epoch, FLAGS_num_epoch - 1, total_loss / train_idx, FLAGS_lr, duration_us / 1e3f, train_dataset->Size() / (duration_us / 1e6)); - } - // TODO(dcj): Add no_grad() context manager later. - std::vector 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]); + evaluate(epoch); } - const auto avg_loss = std::accumulate(test_losses.begin(), test_losses.end(), 0.0) / test_losses.size(); - LOG(ERROR) << "Total: " << total << ", Correct: " << correct - << ", Accuracy: " << static_cast(correct) / total << ", AverageLoss: " << avg_loss; gflags::ShutDownCommandLineFlags(); google::ShutdownGoogleLogging(); diff --git a/infini_train/include/autograd/activations.h b/infini_train/include/autograd/activations.h index a63977263..809a7b67f 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/autograd/conv.h b/infini_train/include/autograd/conv.h new file mode 100644 index 000000000..db3655f77 --- /dev/null +++ b/infini_train/include/autograd/conv.h @@ -0,0 +1,31 @@ +#pragma once + +#include +#include +#include + +#include "infini_train/include/autograd/function.h" + +namespace infini_train { +class Tensor; +} + +namespace infini_train::autograd { + +class Conv2d : public Function { +public: + static constexpr char kType[] = "Conv2dFunction"; + + Conv2d() : Function(kType) {} + + std::vector> Forward(const std::vector> &input_tensors) override; + void SetupContext(const std::vector> &input_tensors, + const std::vector> &output_tensors) override; + std::vector> Backward(const std::vector> &grad_outputs) override; + +private: + bool bias_ = false; + int64_t out_channels_ = 0; + std::vector input_dims_; +}; +} // namespace infini_train::autograd diff --git a/infini_train/include/nn/modules/activations.h b/infini_train/include/nn/modules/activations.h index deb029576..4b69ef180 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/include/nn/modules/conv.h b/infini_train/include/nn/modules/conv.h new file mode 100644 index 000000000..e0e83853b --- /dev/null +++ b/infini_train/include/nn/modules/conv.h @@ -0,0 +1,31 @@ +#pragma once + +#include +#include + +#include "infini_train/include/device.h" +#include "infini_train/include/nn/modules/module.h" + +namespace infini_train { +class Tensor; +class Device; +} // namespace infini_train + +namespace infini_train::nn { +class Conv2d : public CloneableModule { +public: + static constexpr char kType[] = "Conv2d"; + + static constexpr char kParamWeightName[] = "weight"; + static constexpr char kParamBiasName[] = "bias"; + + Conv2d(int64_t in_channels, int64_t out_channels, int64_t kernel_size, bool bias = true, Device device = Device()); + std::vector> Forward(const std::vector> &input_tensors) override; + + bool has_bias() const { return bias_; } + +private: + void ResetParameters(); + bool bias_ = true; +}; +} // namespace infini_train::nn diff --git a/infini_train/src/autograd/activations.cc b/infini_train/src/autograd/activations.cc index bb8b8e5ea..4b2ae37a5 100644 --- a/infini_train/src/autograd/activations.cc +++ b/infini_train/src/autograd/activations.cc @@ -30,4 +30,29 @@ std::vector> Sigmoid::Backward(const std::vectorGetDevice().type(); return {Dispatcher::Instance().Call>({device, "SigmoidBackward"}, output, grad_output)}; } + +std::vector> ReLU::Forward(const std::vector> &input_tensors) { + CHECK_EQ(input_tensors.size(), 1); + const auto &input = input_tensors[0]; + + auto device = input->GetDevice().type(); + return {Dispatcher::Instance().Call>({device, "ReLUForward"}, input)}; +} + +void ReLU::SetupContext(const std::vector> &input_tensors, + const std::vector> &) { + // ReLU backward needs the input activations to decide where the gradient flows. + ctx_.SaveForBackward({input_tensors[0]}); +} + +std::vector> ReLU::Backward(const std::vector> &grad_outputs) { + auto saved_tensors = ctx_.GetSavedTensors(); + CHECK_EQ(saved_tensors.size(), 1); + const auto &input = saved_tensors[0]; + CHECK_EQ(grad_outputs.size(), 1); + const auto &grad_output = grad_outputs[0]; + + auto device = input->GetDevice().type(); + return {Dispatcher::Instance().Call>({device, "ReLUBackward"}, input, grad_output)}; +} } // namespace infini_train::autograd diff --git a/infini_train/src/autograd/conv.cc b/infini_train/src/autograd/conv.cc new file mode 100644 index 000000000..b10894e08 --- /dev/null +++ b/infini_train/src/autograd/conv.cc @@ -0,0 +1,72 @@ +#include "infini_train/include/autograd/conv.h" + +#include "glog/logging.h" + +#include "infini_train/include/dispatcher.h" +#include "infini_train/include/tensor.h" + +namespace infini_train::autograd { +std::vector> Conv2d::Forward(const std::vector> &input_tensors) { + CHECK_GE(input_tensors.size(), 2); + const auto &input = input_tensors[0]; + const auto &weight = input_tensors[1]; + const auto &bias = input_tensors.size() == 3 ? input_tensors[2] : nullptr; + + auto device = input->GetDevice().type(); + return {Dispatcher::Instance().Call>({device, "Conv2dForward"}, input, weight, bias)}; +} + +void Conv2d::SetupContext(const std::vector> &input_tensors, + const std::vector> &) { + const auto &input = input_tensors[0]; + const auto &weight = input_tensors[1]; + bool need_input = ctx_.needs_input_grad().size() > 0 && ctx_.needs_input_grad()[0]; + bool need_weight = ctx_.needs_input_grad().size() > 1 && ctx_.needs_input_grad()[1]; + + // grad_input needs weight, grad_weight needs input + ctx_.SaveForBackward({need_weight ? input : nullptr, need_input ? weight : nullptr}); + + bias_ = input_tensors.size() == 3; + out_channels_ = weight->Dims()[0]; + input_dims_ = input->Dims(); +} + +std::vector> Conv2d::Backward(const std::vector> &grad_outputs) { + auto saved_tensors = ctx_.GetSavedTensors(); + CHECK_EQ(saved_tensors.size(), 2); + const auto &input = saved_tensors[0]; + const auto &weight = saved_tensors[1]; + CHECK_EQ(grad_outputs.size(), 1); + const auto &grad_output = grad_outputs[0]; + + CHECK(!ctx_.needs_input_grad().empty()) << "needs_input_grad not populated in Conv2d::Backward"; + bool need_grad_input = ctx_.needs_input_grad()[0]; + bool need_grad_weight = ctx_.needs_input_grad().size() > 1 && ctx_.needs_input_grad()[1]; + bool need_grad_bias = bias_ && ctx_.needs_input_grad().size() > 2 && ctx_.needs_input_grad()[2]; + + auto device = grad_output->GetDevice().type(); + + std::shared_ptr grad_input = nullptr; + std::shared_ptr grad_weight = nullptr; + std::shared_ptr grad_bias = nullptr; + + if (need_grad_input) { + grad_input = Dispatcher::Instance().Call>({device, "Conv2dBackwardInput"}, weight, + grad_output, input_dims_); + } + if (need_grad_weight) { + grad_weight = Dispatcher::Instance().Call>({device, "Conv2dBackwardWeight"}, input, + grad_output); + } + if (need_grad_bias) { + grad_bias = Dispatcher::Instance().Call>({device, "Conv2dBackwardBias"}, grad_output, + out_channels_); + } + + if (bias_) { + return {grad_input, grad_weight, grad_bias}; + } else { + return {grad_input, grad_weight}; + } +} +} // namespace infini_train::autograd diff --git a/infini_train/src/kernels/cpu/conv.cc b/infini_train/src/kernels/cpu/conv.cc new file mode 100644 index 000000000..1fa8f6349 --- /dev/null +++ b/infini_train/src/kernels/cpu/conv.cc @@ -0,0 +1,288 @@ +#include +#include +#include + +#include "Eigen/Dense" +#include "glog/logging.h" + +#include "infini_train/include/dispatcher.h" +#include "infini_train/include/tensor.h" + +namespace infini_train::kernels::cpu { +namespace { + +using RowMatrix = Eigen::Matrix; + +// im2col for stride=1 / padding=0 cross-correlation: extracts every (C, kH, kW) sliding window of a +// (C, H, W) image into a (out_height*out_width, C*kH*kW) row-major matrix. The k-th column of a row +// corresponds to the (c, u, v) tap ordered like the row-major flattening of an (O, C, kH, kW) +// weight, so a conv step is a plain GEMM against the weight viewed as (O, C*kH*kW). +void Im2col(const float *image, int64_t channels, int64_t height, int64_t width, int64_t kernel_h, int64_t kernel_w, + int64_t out_height, int64_t out_width, float *col) { + const int64_t flat_kernel = channels * kernel_h * kernel_w; + for (int64_t i = 0; i < out_height; ++i) { + for (int64_t j = 0; j < out_width; ++j) { + float *row = col + (i * out_width + j) * flat_kernel; + for (int64_t c = 0; c < channels; ++c) { + const float *window = image + (c * height + i) * width + j; + for (int64_t u = 0; u < kernel_h; ++u) { + for (int64_t v = 0; v < kernel_w; ++v) { + row[(c * kernel_h + u) * kernel_w + v] = window[u * width + v]; + } + } + } + } + } +} + +// col2im: inverse of Im2col for gradients. An input element is covered by multiple output windows, +// so contributions from all patches touching it are accumulated. +void Col2im(const float *col, int64_t channels, int64_t height, int64_t width, int64_t kernel_h, int64_t kernel_w, + int64_t out_height, int64_t out_width, float *grad_image) { + const int64_t flat_kernel = channels * kernel_h * kernel_w; + for (int64_t i = 0; i < out_height; ++i) { + for (int64_t j = 0; j < out_width; ++j) { + const float *row = col + (i * out_width + j) * flat_kernel; + for (int64_t c = 0; c < channels; ++c) { + float *window = grad_image + (c * height + i) * width + j; + for (int64_t u = 0; u < kernel_h; ++u) { + for (int64_t v = 0; v < kernel_w; ++v) { + window[u * width + v] += row[(c * kernel_h + u) * kernel_w + v]; + } + } + } + } + } +} + +} // namespace + +std::shared_ptr Conv2dForward(const std::shared_ptr &input, const std::shared_ptr &weight, + const std::shared_ptr &bias) { + /* + Cross-correlation (PyTorch conv2d semantics, kernel not flipped), stride=1, padding=0: + output(n, o, i, j) = bias(o) + sum_{c,u,v} input(n, c, i+u, j+v) * weight(o, c, u, v) + Computed per image as output(n) = weight(o, C*kH*kW) * im2col(input(n))^T + bias + */ + + const auto &input_dims = input->Dims(); + CHECK_EQ(input_dims.size(), 4); + const int64_t batch = input_dims[0]; + const int64_t channels = input_dims[1]; + const int64_t height = input_dims[2]; + const int64_t width = input_dims[3]; + + const auto &weight_dims = weight->Dims(); + CHECK_EQ(weight_dims.size(), 4); + const int64_t out_channels = weight_dims[0]; + const int64_t kernel_h = weight_dims[2]; + const int64_t kernel_w = weight_dims[3]; + CHECK_EQ(weight_dims[1], channels); + CHECK_GE(height, kernel_h); + CHECK_GE(width, kernel_w); + CHECK(input->Dtype() == DataType::kFLOAT32) << "Conv2d requires FP32 tensors"; + CHECK(weight->Dtype() == DataType::kFLOAT32) << "Conv2d requires FP32 tensors"; + + if (bias) { + CHECK_EQ(bias->Dims().size(), 1); + CHECK_EQ(bias->Dims()[0], out_channels); + CHECK(bias->Dtype() == DataType::kFLOAT32) << "Conv2d requires FP32 tensors"; + } + + const int64_t out_height = height - kernel_h + 1; + const int64_t out_width = width - kernel_w + 1; + const int64_t patches = out_height * out_width; + const int64_t flat_kernel = channels * kernel_h * kernel_w; + + auto output = std::make_shared(std::vector{batch, out_channels, out_height, out_width}, + DataType::kFLOAT32, input->GetDevice()); + + const float *input_data = static_cast(input->DataPtr()); + const float *weight_data = static_cast(weight->DataPtr()); + float *output_data = static_cast(output->DataPtr()); + + // (batch, patches, C*kH*kW) scratch: the im2col expansion of every image in the batch. + // Left uninitialized: im2col overwrites every element before the first read. + const int64_t image_size = channels * height * width; + auto col = std::make_shared(std::vector{batch, patches, flat_kernel}, DataType::kFLOAT32, + input->GetDevice()); + float *col_buffer = static_cast(col->DataPtr()); + for (int64_t n = 0; n < batch; ++n) { + Im2col(input_data + n * image_size, channels, height, width, kernel_h, kernel_w, out_height, out_width, + col_buffer + n * patches * flat_kernel); + } + + Eigen::Map weight_mat(weight_data, out_channels, flat_kernel); + for (int64_t n = 0; n < batch; ++n) { + Eigen::Map col_n(col_buffer + n * patches * flat_kernel, patches, flat_kernel); + Eigen::Map out_n(output_data + n * out_channels * patches, out_channels, patches); + out_n.noalias() = weight_mat * col_n.transpose(); + if (bias) { + out_n.colwise() + += Eigen::Map(static_cast(bias->DataPtr()), out_channels); + } + } + + return output; +} + +std::shared_ptr Conv2dBackwardInput(const std::shared_ptr &weight, + const std::shared_ptr &grad_output, + const std::vector &input_dims) { + /* + grad_input(n, c, x, y) + = sum_{o} sum_{0<=x-uDims(); + CHECK_EQ(weight_dims.size(), 4); + const int64_t out_channels = weight_dims[0]; + const int64_t kernel_h = weight_dims[2]; + const int64_t kernel_w = weight_dims[3]; + CHECK_EQ(weight_dims[1], channels); + CHECK_GE(height, kernel_h); + CHECK_GE(width, kernel_w); + CHECK(weight->Dtype() == DataType::kFLOAT32) << "Conv2d requires FP32 tensors"; + CHECK(grad_output->Dtype() == DataType::kFLOAT32) << "Conv2d requires FP32 tensors"; + + const int64_t out_height = height - kernel_h + 1; + const int64_t out_width = width - kernel_w + 1; + const int64_t patches = out_height * out_width; + const int64_t flat_kernel = channels * kernel_h * kernel_w; + + const auto &grad_dims = grad_output->Dims(); + CHECK_EQ(grad_dims.size(), 4); + CHECK_EQ(grad_dims[0], batch); + CHECK_EQ(grad_dims[1], out_channels); + CHECK_EQ(grad_dims[2], out_height); + CHECK_EQ(grad_dims[3], out_width); + + auto grad_input = std::make_shared(input_dims, DataType::kFLOAT32, grad_output->GetDevice()); + grad_input->Fill(0.0f); + + const float *weight_data = static_cast(weight->DataPtr()); + const float *grad_output_data = static_cast(grad_output->DataPtr()); + float *grad_input_data = static_cast(grad_input->DataPtr()); + + Eigen::Map weight_mat(weight_data, out_channels, flat_kernel); + const int64_t image_size = channels * height * width; + // (patches, C*kH*kW) scratch reused across images. Left uninitialized: the GEMM below + // assigns every element before Col2im reads it. + auto col = std::make_shared(std::vector{patches, flat_kernel}, DataType::kFLOAT32, + grad_output->GetDevice()); + float *col_buffer = static_cast(col->DataPtr()); + for (int64_t n = 0; n < batch; ++n) { + Eigen::Map grad_output_n(grad_output_data + n * out_channels * patches, out_channels, patches); + Eigen::Map col_n(col_buffer, patches, flat_kernel); + col_n.noalias() = grad_output_n.transpose() * weight_mat; + Col2im(col_buffer, channels, height, width, kernel_h, kernel_w, out_height, out_width, + grad_input_data + n * image_size); + } + + return grad_input; +} + +std::shared_ptr Conv2dBackwardWeight(const std::shared_ptr &input, + const std::shared_ptr &grad_output) { + /* + grad_weight(o, c, u, v) = sum_{n, i, j} grad_output(n, o, i, j) * input(n, c, i+u, j+v) + Computed per image as grad_output(n) * im2col(input(n)) accumulated over the batch, the + transpose of the forward GEMM. + */ + + const auto &input_dims = input->Dims(); + CHECK_EQ(input_dims.size(), 4); + const int64_t batch = input_dims[0]; + const int64_t channels = input_dims[1]; + const int64_t height = input_dims[2]; + const int64_t width = input_dims[3]; + + const auto &grad_dims = grad_output->Dims(); + CHECK_EQ(grad_dims.size(), 4); + CHECK_EQ(grad_dims[0], batch); + CHECK(input->Dtype() == DataType::kFLOAT32) << "Conv2d requires FP32 tensors"; + CHECK(grad_output->Dtype() == DataType::kFLOAT32) << "Conv2d requires FP32 tensors"; + + // The kernel extent is the difference between the input and output spatial extents. + const int64_t kernel_h = height - grad_dims[2] + 1; + const int64_t kernel_w = width - grad_dims[3] + 1; + CHECK_GT(kernel_h, 0); + CHECK_GT(kernel_w, 0); + + const int64_t out_channels = grad_dims[1]; + const int64_t out_height = grad_dims[2]; + const int64_t out_width = grad_dims[3]; + const int64_t patches = out_height * out_width; + const int64_t flat_kernel = channels * kernel_h * kernel_w; + + auto grad_weight = std::make_shared(std::vector{out_channels, channels, kernel_h, kernel_w}, + DataType::kFLOAT32, input->GetDevice()); + grad_weight->Fill(0.0f); + + const float *input_data = static_cast(input->DataPtr()); + const float *grad_output_data = static_cast(grad_output->DataPtr()); + float *grad_weight_data = static_cast(grad_weight->DataPtr()); + + const int64_t image_size = channels * height * width; + // (batch, patches, C*kH*kW) scratch: the im2col expansion of every image in the batch. + // Left uninitialized: im2col overwrites every element before the first read. + auto col = std::make_shared(std::vector{batch, patches, flat_kernel}, DataType::kFLOAT32, + input->GetDevice()); + float *col_buffer = static_cast(col->DataPtr()); + for (int64_t n = 0; n < batch; ++n) { + Im2col(input_data + n * image_size, channels, height, width, kernel_h, kernel_w, out_height, out_width, + col_buffer + n * patches * flat_kernel); + } + + Eigen::Map grad_weight_mat(grad_weight_data, out_channels, flat_kernel); + for (int64_t n = 0; n < batch; ++n) { + Eigen::Map grad_output_n(grad_output_data + n * out_channels * patches, out_channels, patches); + Eigen::Map col_n(col_buffer + n * patches * flat_kernel, patches, flat_kernel); + grad_weight_mat.noalias() += grad_output_n * col_n; + } + + return grad_weight; +} + +std::shared_ptr Conv2dBackwardBias(const std::shared_ptr &grad_output, int64_t out_channels) { + /* + grad_bias(o) = sum_{n, i, j} grad_output(n, o, i, j) + */ + const auto &grad_dims = grad_output->Dims(); + CHECK_EQ(grad_dims.size(), 4); + CHECK_EQ(grad_dims[1], out_channels); + CHECK(grad_output->Dtype() == DataType::kFLOAT32) << "Conv2d requires FP32 tensors"; + + auto grad_bias + = std::make_shared(std::vector{out_channels}, DataType::kFLOAT32, grad_output->GetDevice()); + grad_bias->Fill(0.0f); + Eigen::Map grad_bias_vec(static_cast(grad_bias->DataPtr()), out_channels); + + const int64_t batch = grad_dims[0]; + const int64_t patches = grad_dims[2] * grad_dims[3]; + const float *grad_output_data = static_cast(grad_output->DataPtr()); + for (int64_t n = 0; n < batch; ++n) { + Eigen::Map grad_output_n(grad_output_data + n * out_channels * patches, out_channels, patches); + grad_bias_vec += grad_output_n.rowwise().sum(); + } + return grad_bias; +} + +} // namespace infini_train::kernels::cpu + +#define REGISTER_CPU_CONV2D_KERNEL(kernel_name) \ + REGISTER_KERNEL(infini_train::Device::DeviceType::kCPU, kernel_name, infini_train::kernels::cpu::kernel_name) + +REGISTER_CPU_CONV2D_KERNEL(Conv2dForward) +REGISTER_CPU_CONV2D_KERNEL(Conv2dBackwardInput) +REGISTER_CPU_CONV2D_KERNEL(Conv2dBackwardWeight) +REGISTER_CPU_CONV2D_KERNEL(Conv2dBackwardBias) + +#undef REGISTER_CPU_CONV2D_KERNEL diff --git a/infini_train/src/kernels/cpu/relu.cc b/infini_train/src/kernels/cpu/relu.cc new file mode 100644 index 000000000..2a6bddbb6 --- /dev/null +++ b/infini_train/src/kernels/cpu/relu.cc @@ -0,0 +1,48 @@ +#include +#include + +#include "glog/logging.h" + +#include "infini_train/include/dispatcher.h" +#include "infini_train/include/tensor.h" + +namespace infini_train::kernels::cpu { +std::shared_ptr ReLUForward(const std::shared_ptr &input) { + CHECK(input->Dtype() == DataType::kFLOAT32) << "ReLU requires FP32 tensors"; + + auto output = std::make_shared(input->Dims(), DataType::kFLOAT32); + const float *input_ptr = static_cast(input->DataPtr()); + float *output_ptr = static_cast(output->DataPtr()); + + const int64_t numel = input->NumElements(); + // Strict `<` keeps NaN and -0 (both compare false), matching clamp_min(x, 0). + for (int64_t idx = 0; idx < numel; ++idx) { output_ptr[idx] = input_ptr[idx] < 0.0f ? 0.0f : input_ptr[idx]; } + + return output; +} + +std::shared_ptr ReLUBackward(const std::shared_ptr &input, const std::shared_ptr &grad_output) { + CHECK(input->Dtype() == DataType::kFLOAT32) << "ReLU requires FP32 tensors"; + CHECK(grad_output->Dtype() == DataType::kFLOAT32) << "ReLU requires FP32 tensors"; + + auto grad_input = std::make_shared(input->Dims(), DataType::kFLOAT32); + const float *input_ptr = static_cast(input->DataPtr()); + const float *grad_output_ptr = static_cast(grad_output->DataPtr()); + float *grad_input_ptr = static_cast(grad_input->DataPtr()); + + const int64_t numel = input->NumElements(); + // `<=` matches threshold_backward: NaN (false) passes grad through, 0/-0/negative mask to 0. + for (int64_t idx = 0; idx < numel; ++idx) { + grad_input_ptr[idx] = input_ptr[idx] <= 0.0f ? 0.0f : grad_output_ptr[idx]; + } + return grad_input; +} +} // namespace infini_train::kernels::cpu + +#define REGISTER_CPU_RELU_KERNEL(kernel_name) \ + REGISTER_KERNEL(infini_train::Device::DeviceType::kCPU, kernel_name, infini_train::kernels::cpu::kernel_name) + +REGISTER_CPU_RELU_KERNEL(ReLUForward) +REGISTER_CPU_RELU_KERNEL(ReLUBackward) + +#undef REGISTER_CPU_RELU_KERNEL diff --git a/infini_train/src/kernels/cuda/conv.cu b/infini_train/src/kernels/cuda/conv.cu new file mode 100644 index 000000000..191436094 --- /dev/null +++ b/infini_train/src/kernels/cuda/conv.cu @@ -0,0 +1,417 @@ +#include +#include +#include + +#include + +#include "glog/logging.h" + +#include "infini_train/include/common/cuda/common_cuda.h" +#include "infini_train/include/core/runtime/device_guard.h" +#include "infini_train/include/dispatcher.h" +#include "infini_train/include/tensor.h" + +#include "infini_train/src/core/runtime/cuda/cuda_runtime_common.h" +#include "infini_train/src/kernels/common/gemm.h" + +namespace infini_train::kernels::cuda { +namespace { + +// im2col for stride=1 / padding=0 cross-correlation. Each thread writes one +// element of the (flat_kernel, patches) column-major scratch buffer, so a single +// image's buffer is addressable as col[kk * patches + p]. This column-major +// order is what the cuBLAS GEMM geometry below expects. The (c, u, v) tap index +// is the same row-major flattening as an (O, C, kH, kW) weight. +__global__ void Im2colKernel(const float *__restrict__ input, float *__restrict__ col, int64_t total, int64_t channels, + int64_t height, int64_t width, int64_t kernel_h, int64_t kernel_w, int64_t out_width, + int64_t patches, int64_t flat_kernel) { + for (int64_t idx = blockIdx.x * blockDim.x + threadIdx.x; idx < total; idx += blockDim.x * gridDim.x) { + const int64_t p = idx % patches; + const int64_t kk = (idx / patches) % flat_kernel; + const int64_t n = idx / (flat_kernel * patches); + + const int64_t i = p / out_width; + const int64_t j = p % out_width; + const int64_t c = kk / (kernel_h * kernel_w); + const int64_t rem = kk % (kernel_h * kernel_w); + const int64_t u = rem / kernel_w; + const int64_t v = rem % kernel_w; + + const int64_t in_idx = n * channels * height * width + c * height * width + (i + u) * width + (j + v); + col[idx] = input[in_idx]; + } +} + +// col2im: inverse of Im2col for the input gradient. An input pixel is covered by +// multiple output windows, so all patches that touch it are gathered. Each thread +// owns one (n, c, x, y) output cell and writes it once, so no atomics are needed. +__global__ void Col2imKernel(const float *__restrict__ col, float *__restrict__ grad_input, int64_t total, + int64_t channels, int64_t height, int64_t width, int64_t kernel_h, int64_t kernel_w, + int64_t out_height, int64_t out_width, int64_t patches, int64_t flat_kernel) { + for (int64_t idx = blockIdx.x * blockDim.x + threadIdx.x; idx < total; idx += blockDim.x * gridDim.x) { + const int64_t y = idx % width; + const int64_t x = (idx / width) % height; + const int64_t c = (idx / (width * height)) % channels; + const int64_t n = idx / (width * height * channels); + + float sum = 0.0f; + for (int64_t u = 0; u < kernel_h; ++u) { + const int64_t i = x - u; + if (i < 0 || i >= out_height) { + continue; + } + for (int64_t v = 0; v < kernel_w; ++v) { + const int64_t j = y - v; + if (j < 0 || j >= out_width) { + continue; + } + const int64_t kk = c * kernel_h * kernel_w + u * kernel_w + v; + const int64_t col_idx = n * (flat_kernel * patches) + kk * patches + i * out_width + j; + sum += col[col_idx]; + } + } + grad_input[idx] = sum; + } +} + +// Broadcast bias over every spatial position of the (N, O, H, W) output, so a +// subsequent GEMM with beta=1 accumulates the convolution onto it. +__global__ void BroadcastBiasKernel(float *__restrict__ output, const float *__restrict__ bias, int64_t total, + int64_t patches, int64_t out_channels) { + for (int64_t idx = blockIdx.x * blockDim.x + threadIdx.x; idx < total; idx += blockDim.x * gridDim.x) { + const int64_t o = (idx / patches) % out_channels; + output[idx] = bias[o]; + } +} + +// grad_bias(o) = sum_{n, i, j} grad_output(n, o, i, j). A parallel block reduction keeps the +// accumulation error at the same order as the cuBLAS reference rather than a single-thread sum. +template +__global__ void GradBiasKernel(const float *__restrict__ grad_output, float *__restrict__ grad_bias, int64_t batch, + int64_t patches, int64_t out_channels) { + using BlockReduce = cub::BlockReduce; + __shared__ typename BlockReduce::TempStorage temp_storage; + + const int64_t o = blockIdx.x; + if (o >= out_channels) { + return; + } + + const int64_t total = batch * patches; + float sum = 0.0f; + for (int64_t idx = threadIdx.x; idx < total; idx += blockDim.x) { + const int64_t n = idx / patches; + const int64_t p = idx % patches; + sum += grad_output[n * out_channels * patches + o * patches + p]; + } + + const float reduced = BlockReduce(temp_storage).Sum(sum); + if (threadIdx.x == 0) { + grad_bias[o] = reduced; + } +} + +constexpr int kThreads = 256; +constexpr int kMaxBlocks = 1024; + +int64_t NumBlocks(int64_t total) { + int64_t blocks = (total + kThreads - 1) / kThreads; + return blocks > kMaxBlocks ? kMaxBlocks : blocks; +} + +cudaStream_t CurrentCudaStream(const Device &device) { + return dynamic_cast( + infini_train::core::GetDeviceGuardImpl(device.type())->GetStream(device)) + ->cuda_stream(); +} + +} // namespace + +std::shared_ptr Conv2dForward(const std::shared_ptr &input, const std::shared_ptr &weight, + const std::shared_ptr &bias) { + const auto &input_dims = input->Dims(); + CHECK_EQ(input_dims.size(), 4); + const int64_t batch = input_dims[0]; + const int64_t channels = input_dims[1]; + const int64_t height = input_dims[2]; + const int64_t width = input_dims[3]; + + const auto &weight_dims = weight->Dims(); + CHECK_EQ(weight_dims.size(), 4); + const int64_t out_channels = weight_dims[0]; + const int64_t kernel_h = weight_dims[2]; + const int64_t kernel_w = weight_dims[3]; + CHECK_EQ(weight_dims[1], channels); + CHECK_GE(height, kernel_h); + CHECK_GE(width, kernel_w); + + CHECK(input->Dtype() == DataType::kFLOAT32) << "Conv2d requires FP32 tensors"; + CHECK(weight->Dtype() == DataType::kFLOAT32) << "Conv2d requires FP32 tensors"; + CHECK_EQ(input->GetDevice(), weight->GetDevice()); + + if (bias) { + CHECK_EQ(bias->Dims().size(), 1); + CHECK_EQ(bias->Dims()[0], out_channels); + CHECK(bias->Dtype() == DataType::kFLOAT32) << "Conv2d requires FP32 tensors"; + CHECK_EQ(input->GetDevice(), bias->GetDevice()); + } + + const int64_t out_height = height - kernel_h + 1; + const int64_t out_width = width - kernel_w + 1; + const int64_t patches = out_height * out_width; + const int64_t flat_kernel = channels * kernel_h * kernel_w; + + auto output = std::make_shared(std::vector{batch, out_channels, out_height, out_width}, + DataType::kFLOAT32, input->GetDevice()); + if (batch == 0) { + return output; + } + + auto device = input->GetDevice(); + const cudaStream_t stream = CurrentCudaStream(device); + + // Scratch im2col buffer, indexed column-major (patches, flat_kernel) per image. + auto col = std::make_shared(std::vector{batch, flat_kernel, patches}, DataType::kFLOAT32, device); + + Im2colKernel<<>>( + static_cast(input->DataPtr()), static_cast(col->DataPtr()), + batch * patches * flat_kernel, channels, height, width, kernel_h, kernel_w, out_width, patches, flat_kernel); + CUDA_CHECK(cudaGetLastError()); + + if (bias) { + // Prefill output with the bias, then the GEMM below accumulates onto it. + BroadcastBiasKernel<<>>( + static_cast(output->DataPtr()), static_cast(bias->DataPtr()), + batch * out_channels * patches, patches, out_channels); + CUDA_CHECK(cudaGetLastError()); + } + + const float *col_data = static_cast(col->DataPtr()); + const float *weight_data = static_cast(weight->DataPtr()); + float *output_data = static_cast(output->DataPtr()); + const float beta = bias ? 1.0f : 0.0f; + // One strided-batched GEMM over the batch: A steps through the per-image col blocks, + // B (the weight) is shared by every image via stride 0, C steps through the output. + // batch == 1 falls back to the non-batched Gemm path, which expects zero strides. + Dispatcher::Instance().Call({device.type(), "Gemm"}, device, + GemmParams{ + .trans_a = GemmTranspose::kNoTranspose, + .trans_b = GemmTranspose::kNoTranspose, + .m = static_cast(patches), + .n = static_cast(out_channels), + .k = static_cast(flat_kernel), + .A = col_data, + .lda = static_cast(patches), + .B = weight_data, + .ldb = static_cast(flat_kernel), + .C = output_data, + .ldc = static_cast(patches), + .alpha = 1.0f, + .beta = beta, + .batch_count = static_cast(batch), + .stride_a = batch > 1 ? patches * flat_kernel : 0, + .stride_b = 0, + .stride_c = batch > 1 ? out_channels * patches : 0, + .input_dtype = DataType::kFLOAT32, + .output_dtype = DataType::kFLOAT32, + }); + + return output; +} + +std::shared_ptr Conv2dBackwardInput(const std::shared_ptr &weight, + const std::shared_ptr &grad_output, + const std::vector &input_dims) { + CHECK_EQ(input_dims.size(), 4); + const int64_t batch = input_dims[0]; + const int64_t channels = input_dims[1]; + const int64_t height = input_dims[2]; + const int64_t width = input_dims[3]; + + const auto &weight_dims = weight->Dims(); + CHECK_EQ(weight_dims.size(), 4); + const int64_t out_channels = weight_dims[0]; + const int64_t kernel_h = weight_dims[2]; + const int64_t kernel_w = weight_dims[3]; + CHECK_EQ(weight_dims[1], channels); + CHECK_GE(height, kernel_h); + CHECK_GE(width, kernel_w); + + CHECK(weight->Dtype() == DataType::kFLOAT32) << "Conv2d requires FP32 tensors"; + CHECK(grad_output->Dtype() == DataType::kFLOAT32) << "Conv2d requires FP32 tensors"; + CHECK(weight->GetDevice() == grad_output->GetDevice()); + + const int64_t out_height = height - kernel_h + 1; + const int64_t out_width = width - kernel_w + 1; + const int64_t patches = out_height * out_width; + const int64_t flat_kernel = channels * kernel_h * kernel_w; + + const auto &grad_dims = grad_output->Dims(); + CHECK_EQ(grad_dims.size(), 4); + CHECK_EQ(grad_dims[0], batch); + CHECK_EQ(grad_dims[1], out_channels); + CHECK_EQ(grad_dims[2], out_height); + CHECK_EQ(grad_dims[3], out_width); + + auto grad_input = std::make_shared(input_dims, DataType::kFLOAT32, grad_output->GetDevice()); + if (batch == 0) { + return grad_input; + } + + auto device = grad_output->GetDevice(); + const cudaStream_t stream = CurrentCudaStream(device); + + // Whole-batch scratch: col(n) = grad_output(n)^T * weight per image, stored column-major + // (patches, flat_kernel), so the col2im below can gather every image in a single launch. + auto col = std::make_shared(std::vector{batch, flat_kernel, patches}, DataType::kFLOAT32, device); + + const float *weight_data = static_cast(weight->DataPtr()); + const float *grad_output_data = static_cast(grad_output->DataPtr()); + float *grad_input_data = static_cast(grad_input->DataPtr()); + float *col_data = static_cast(col->DataPtr()); + + // The transpose of the forward GEMM for every image at once: A steps through grad_output, + // B (the weight) is shared by every image via stride 0, C steps through the col buffer. + // batch == 1 falls back to the non-batched Gemm path, which expects zero strides. + Dispatcher::Instance().Call({device.type(), "Gemm"}, device, + GemmParams{ + .trans_a = GemmTranspose::kNoTranspose, + .trans_b = GemmTranspose::kTranspose, + .m = static_cast(patches), + .n = static_cast(flat_kernel), + .k = static_cast(out_channels), + .A = grad_output_data, + .lda = static_cast(patches), + .B = weight_data, + .ldb = static_cast(flat_kernel), + .C = col_data, + .ldc = static_cast(patches), + .alpha = 1.0f, + .beta = 0.0f, + .batch_count = static_cast(batch), + .stride_a = batch > 1 ? out_channels * patches : 0, + .stride_b = 0, + .stride_c = batch > 1 ? flat_kernel * patches : 0, + .input_dtype = DataType::kFLOAT32, + .output_dtype = DataType::kFLOAT32, + }); + + Col2imKernel<<>>( + col_data, grad_input_data, batch * channels * height * width, channels, height, width, kernel_h, kernel_w, + out_height, out_width, patches, flat_kernel); + CUDA_CHECK(cudaGetLastError()); + + return grad_input; +} + +std::shared_ptr Conv2dBackwardWeight(const std::shared_ptr &input, + const std::shared_ptr &grad_output) { + const auto &input_dims = input->Dims(); + CHECK_EQ(input_dims.size(), 4); + const int64_t batch = input_dims[0]; + const int64_t channels = input_dims[1]; + const int64_t height = input_dims[2]; + const int64_t width = input_dims[3]; + + const auto &grad_dims = grad_output->Dims(); + CHECK_EQ(grad_dims.size(), 4); + CHECK_EQ(grad_dims[0], batch); + + CHECK(input->Dtype() == DataType::kFLOAT32) << "Conv2d requires FP32 tensors"; + CHECK(grad_output->Dtype() == DataType::kFLOAT32) << "Conv2d requires FP32 tensors"; + CHECK(input->GetDevice() == grad_output->GetDevice()); + + const int64_t kernel_h = height - grad_dims[2] + 1; + const int64_t kernel_w = width - grad_dims[3] + 1; + CHECK_GT(kernel_h, 0); + CHECK_GT(kernel_w, 0); + + const int64_t out_channels = grad_dims[1]; + const int64_t out_height = grad_dims[2]; + const int64_t out_width = grad_dims[3]; + const int64_t patches = out_height * out_width; + const int64_t flat_kernel = channels * kernel_h * kernel_w; + + auto grad_weight = std::make_shared(std::vector{out_channels, channels, kernel_h, kernel_w}, + DataType::kFLOAT32, input->GetDevice()); + grad_weight->Fill(0.0f); + if (batch == 0) { + return grad_weight; + } + + auto device = input->GetDevice(); + const cudaStream_t stream = CurrentCudaStream(device); + + auto col = std::make_shared(std::vector{batch, flat_kernel, patches}, DataType::kFLOAT32, device); + + Im2colKernel<<>>( + static_cast(input->DataPtr()), static_cast(col->DataPtr()), + batch * patches * flat_kernel, channels, height, width, kernel_h, kernel_w, out_width, patches, flat_kernel); + CUDA_CHECK(cudaGetLastError()); + + const float *col_data = static_cast(col->DataPtr()); + const float *grad_output_data = static_cast(grad_output->DataPtr()); + float *grad_weight_data = static_cast(grad_weight->DataPtr()); + + // grad_weight(o, kk) = sum_n grad_output(n, o, p) * col(n, kk, p). Computed as + // the transpose of the forward GEMM and accumulated over the batch. + for (int64_t n = 0; n < batch; ++n) { + Dispatcher::Instance().Call({device.type(), "Gemm"}, device, + GemmParams{ + .trans_a = GemmTranspose::kTranspose, + .trans_b = GemmTranspose::kNoTranspose, + .m = static_cast(flat_kernel), + .n = static_cast(out_channels), + .k = static_cast(patches), + .A = col_data + n * patches * flat_kernel, + .lda = static_cast(patches), + .B = grad_output_data + n * out_channels * patches, + .ldb = static_cast(patches), + .C = grad_weight_data, + .ldc = static_cast(flat_kernel), + .alpha = 1.0f, + .beta = 1.0f, + .batch_count = 1, + .input_dtype = DataType::kFLOAT32, + .output_dtype = DataType::kFLOAT32, + }); + } + + return grad_weight; +} + +std::shared_ptr Conv2dBackwardBias(const std::shared_ptr &grad_output, int64_t out_channels) { + const auto &grad_dims = grad_output->Dims(); + CHECK_EQ(grad_dims.size(), 4); + CHECK_EQ(grad_dims[1], out_channels); + CHECK(grad_output->Dtype() == DataType::kFLOAT32) << "Conv2d requires FP32 tensors"; + + const int64_t batch = grad_dims[0]; + const int64_t patches = grad_dims[2] * grad_dims[3]; + + auto grad_bias + = std::make_shared(std::vector{out_channels}, DataType::kFLOAT32, grad_output->GetDevice()); + + auto device = grad_output->GetDevice(); + const cudaStream_t stream = CurrentCudaStream(device); + + constexpr int kGradBiasBlock = 256; + GradBiasKernel<<>>( + static_cast(grad_output->DataPtr()), static_cast(grad_bias->DataPtr()), batch, patches, + out_channels); + CUDA_CHECK(cudaGetLastError()); + + return grad_bias; +} + +} // namespace infini_train::kernels::cuda + +#define REGISTER_CUDA_CONV2D_KERNEL(kernel_name) \ + REGISTER_KERNEL(infini_train::Device::DeviceType::kCUDA, kernel_name, infini_train::kernels::cuda::kernel_name) + +REGISTER_CUDA_CONV2D_KERNEL(Conv2dForward) +REGISTER_CUDA_CONV2D_KERNEL(Conv2dBackwardInput) +REGISTER_CUDA_CONV2D_KERNEL(Conv2dBackwardWeight) +REGISTER_CUDA_CONV2D_KERNEL(Conv2dBackwardBias) + +#undef REGISTER_CUDA_CONV2D_KERNEL diff --git a/infini_train/src/kernels/cuda/relu.cu b/infini_train/src/kernels/cuda/relu.cu new file mode 100644 index 000000000..3f33d4f38 --- /dev/null +++ b/infini_train/src/kernels/cuda/relu.cu @@ -0,0 +1,113 @@ +#include +#include + +#include "infini_train/include/common/common.h" +#include "infini_train/include/core/runtime/device_guard.h" +#include "infini_train/include/dispatcher.h" +#include "infini_train/include/dtype_dispatch.h" +#include "infini_train/include/tensor.h" + +#include "infini_train/src/core/runtime/cuda/cuda_runtime_common.h" + +namespace infini_train::kernels::cuda { +namespace { + +template __global__ void ReLUForwardKernel(T *output, const T *input, size_t num_elements, size_t offset) { + size_t idx = blockIdx.x * blockDim.x + threadIdx.x + offset; + if (idx < num_elements) { + // Strict `<` keeps NaN and -0 (both compare false), matching clamp_min(x, 0). + output[idx] = input[idx] < T(0) ? T(0) : input[idx]; + } +} + +template +__global__ void ReLUBackwardKernel(T *grad_input, const T *input, const T *grad_output, size_t num_elements, + size_t offset) { + size_t idx = blockIdx.x * blockDim.x + threadIdx.x + offset; + if (idx < num_elements) { + // `<=` matches threshold_backward: NaN (false) passes grad through, 0/-0/negative mask to 0. + grad_input[idx] = input[idx] <= T(0) ? T(0) : grad_output[idx]; + } +} + +inline size_t ChooseBlockSize(size_t num_elements) { + if (num_elements < 1024) { + return 64; + } + if (num_elements < 65536) { + return 128; + } + if (num_elements < 1048576) { + return 256; + } + return 512; +} + +cudaStream_t CurrentStream(const std::shared_ptr &tensor) { + auto device = tensor->GetDevice(); + return dynamic_cast( + infini_train::core::GetDeviceGuardImpl(device.type())->GetStream(device)) + ->cuda_stream(); +} + +template +void ReLUForwardImpl(const std::shared_ptr &input, const std::shared_ptr &output) { + const size_t num_elements = output->NumElements(); + if (num_elements == 0) { + return; + } + cudaStream_t stream = CurrentStream(output); + T *out_ptr = static_cast(output->DataPtr()); + const T *in_ptr = static_cast(input->DataPtr()); + + dim3 block(ChooseBlockSize(num_elements)); + dim3 grid(CEIL_DIV(num_elements, block.x)); + const size_t step = grid.x * block.x; + for (size_t offset = 0; offset < num_elements; offset += step) { + ReLUForwardKernel<<>>(out_ptr, in_ptr, num_elements, offset); + } +} + +template +void ReLUBackwardImpl(const std::shared_ptr &input, const std::shared_ptr &grad_output, + const std::shared_ptr &grad_input) { + const size_t num_elements = grad_input->NumElements(); + if (num_elements == 0) { + return; + } + cudaStream_t stream = CurrentStream(grad_input); + T *grad_input_ptr = static_cast(grad_input->DataPtr()); + const T *input_ptr = static_cast(input->DataPtr()); + const T *grad_output_ptr = static_cast(grad_output->DataPtr()); + + dim3 block(ChooseBlockSize(num_elements)); + dim3 grid(CEIL_DIV(num_elements, block.x)); + const size_t step = grid.x * block.x; + for (size_t offset = 0; offset < num_elements; offset += step) { + ReLUBackwardKernel + <<>>(grad_input_ptr, input_ptr, grad_output_ptr, num_elements, offset); + } +} + +} // namespace + +std::shared_ptr ReLUForward(const std::shared_ptr &input) { + auto output = std::make_shared(input->Dims(), input->Dtype(), input->GetDevice()); + DISPATCH(input->Dtype(), ReLUForwardImpl(input, output);, DataType::kFLOAT32) + return output; +} + +std::shared_ptr ReLUBackward(const std::shared_ptr &input, const std::shared_ptr &grad_output) { + auto grad_input = std::make_shared(input->Dims(), input->Dtype(), input->GetDevice()); + DISPATCH(input->Dtype(), ReLUBackwardImpl(input, grad_output, grad_input);, DataType::kFLOAT32) + return grad_input; +} +} // namespace infini_train::kernels::cuda + +#define REGISTER_CUDA_RELU_KERNEL(kernel_name) \ + REGISTER_KERNEL(infini_train::Device::DeviceType::kCUDA, kernel_name, infini_train::kernels::cuda::kernel_name) + +REGISTER_CUDA_RELU_KERNEL(ReLUForward) +REGISTER_CUDA_RELU_KERNEL(ReLUBackward) + +#undef REGISTER_CUDA_RELU_KERNEL diff --git a/infini_train/src/nn/modules/activations.cc b/infini_train/src/nn/modules/activations.cc index d1bbc9da8..a53738478 100644 --- a/infini_train/src/nn/modules/activations.cc +++ b/infini_train/src/nn/modules/activations.cc @@ -12,6 +12,10 @@ std::vector> Sigmoid::Forward(const std::vector()->Apply(input_tensors); } +std::vector> ReLU::Forward(const std::vector> &input_tensors) { + return std::make_shared()->Apply(input_tensors); +} + std::vector> NewGELU::Forward(const std::vector> &x) { auto &input = x[0]; return {0.5 * input diff --git a/infini_train/src/nn/modules/conv.cc b/infini_train/src/nn/modules/conv.cc new file mode 100644 index 000000000..7ad8a4e4e --- /dev/null +++ b/infini_train/src/nn/modules/conv.cc @@ -0,0 +1,43 @@ +#include "infini_train/include/nn/modules/conv.h" + +#include +#include +#include + +#include "infini_train/include/autograd/conv.h" +#include "infini_train/include/device.h" +#include "infini_train/include/nn/init.h" +#include "infini_train/include/tensor.h" + +namespace infini_train::nn { +Conv2d::Conv2d(int64_t in_channels, int64_t out_channels, int64_t kernel_size, bool bias, Device device) + : CloneableModule(kType), bias_(bias) { + device_ = device; + + parameters_[kParamWeightName] + = std::make_shared(std::vector{out_channels, in_channels, kernel_size, kernel_size}, + DataType::kFLOAT32, device_) + ->RequiresGrad(); + if (bias) { + parameters_[kParamBiasName] + = std::make_shared(std::vector{out_channels}, DataType::kFLOAT32, device_)->RequiresGrad(); + } + ResetParameters(); +} + +std::vector> Conv2d::Forward(const std::vector> &input_tensors) { + return std::make_shared()->Apply( + bias_ ? std::vector>{input_tensors[0], parameters_[kParamWeightName], + parameters_[kParamBiasName]} + : std::vector>{input_tensors[0], parameters_[kParamWeightName]}); +} + +void Conv2d::ResetParameters() { + init::KaimingUniform(parameters_[kParamWeightName], sqrt(5.0f)); + if (bias_) { + const auto [fan_in, _] = init::CalculateFanInAndFanOut(parameters_[kParamWeightName]); + const float bound = fan_in > 0 ? 1.0 / sqrt(fan_in) : 0.0; + init::Uniform(parameters_[kParamBiasName], -bound, bound); + } +} +} // namespace infini_train::nn diff --git a/tests/autograd/test_autograd_conv_backward.cc b/tests/autograd/test_autograd_conv_backward.cc new file mode 100644 index 000000000..8390ce586 --- /dev/null +++ b/tests/autograd/test_autograd_conv_backward.cc @@ -0,0 +1,245 @@ +#include +#include +#include + +#include "gtest/gtest.h" + +#include "infini_train/include/autograd/conv.h" +#include "infini_train/include/tensor.h" + +#include "tests/common/test_utils.h" + +using namespace infini_train; + +namespace { +// Same golden case as the forward test (torch 2.14, seed 20260905, autograd.grad reference). +constexpr float kGoldenAbsError = 2e-6f; +} // namespace + +class AutogradConvBackwardTest : public infini_train::test::InfiniTrainTest {}; + +// Hand-checkable case: x = arange(9) reshaped (1,1,3,3), w = [[1,2],[3,4]], bias = 0.5, +// grad_output = ones(1,1,2,2). bias does not influence grad_input / grad_weight. +TEST_P(AutogradConvBackwardTest, ConvBackwardGradients) { + std::vector input_values; + for (int idx = 0; idx < 9; ++idx) { input_values.push_back(static_cast(idx)); } + auto input = std::make_shared(input_values.data(), std::vector{1, 1, 3, 3}, DataType::kFLOAT32, + GetDevice()) + ->RequiresGrad(); + const std::vector weight_values{1.0f, 2.0f, 3.0f, 4.0f}; + auto weight = std::make_shared(weight_values.data(), std::vector{1, 1, 2, 2}, DataType::kFLOAT32, + GetDevice()) + ->RequiresGrad(); + auto bias = std::make_shared(std::vector{1}, DataType::kFLOAT32, GetDevice(), true); + bias->Fill(0.5f); + + auto conv_fn = std::make_shared(); + auto result = conv_fn->Apply({input, weight, bias}); + auto grad_output + = std::make_shared(std::vector{1, 1, 2, 2}, DataType::kFLOAT32, GetDevice(), true); + grad_output->Fill(1.0f); + auto grad_inputs = conv_fn->Backward({grad_output}); + ASSERT_EQ(grad_inputs.size(), 3); + ASSERT_NE(grad_inputs[0], nullptr); + ASSERT_NE(grad_inputs[1], nullptr); + ASSERT_NE(grad_inputs[2], nullptr); + + // grad_input(x, y) sums the taps of every window covering (x, y). + const std::vector expected_grad_input{1.0f, 3.0f, 2.0f, 4.0f, 10.0f, 6.0f, 3.0f, 7.0f, 4.0f}; + EXPECT_EQ(grad_inputs[0]->Dims(), (std::vector{1, 1, 3, 3})); + test::ExpectTensorFloatEqual(grad_inputs[0], expected_grad_input); + + // grad_weight(u, v) sums the input elements seen through tap (u, v) over all windows. + const std::vector expected_grad_weight{8.0f, 12.0f, 20.0f, 24.0f}; + EXPECT_EQ(grad_inputs[1]->Dims(), (std::vector{1, 1, 2, 2})); + test::ExpectTensorFloatEqual(grad_inputs[1], expected_grad_weight); + + test::ExpectTensorFloatEqual(grad_inputs[2], std::vector{4.0f}); +} + +TEST_P(AutogradConvBackwardTest, ConvBackwardNoBias) { + auto input = std::make_shared(std::vector{2, 2, 4, 4}, DataType::kFLOAT32, GetDevice(), true); + input->Fill(1.0f); + auto weight = std::make_shared(std::vector{3, 2, 3, 3}, DataType::kFLOAT32, GetDevice(), true); + weight->Fill(1.0f); + auto conv_fn = std::make_shared(); + auto result = conv_fn->Apply({input, weight}); + auto grad = std::make_shared(std::vector{2, 3, 2, 2}, DataType::kFLOAT32, GetDevice(), true); + grad->Fill(1.0f); + auto grad_inputs = conv_fn->Backward({grad}); + EXPECT_EQ(grad_inputs.size(), 2); + + // No-bias path returns {grad_input, grad_weight}: both must be produced and correct. + ASSERT_NE(grad_inputs[0], nullptr); + EXPECT_EQ(grad_inputs[0]->Dims(), (std::vector{2, 2, 4, 4})); + ASSERT_NE(grad_inputs[1], nullptr); + EXPECT_EQ(grad_inputs[1]->Dims(), (std::vector{3, 2, 3, 3})); + // With all-ones input/weight/grad_output, grad_weight(o,c,u,v) = sum_{n,i,j} 1 = 2*2*2. + test::ExpectTensorFloatEqual(grad_inputs[1], 8.0f); +} + +// H = W = kernel: a single window, so grad_input mirrors the weight and grad_weight the input. +TEST_P(AutogradConvBackwardTest, ConvBackwardKernelEqualsSpatial) { + const std::vector input_values{1.0f, 2.0f, 3.0f, 4.0f}; + auto input = std::make_shared(input_values.data(), std::vector{1, 1, 2, 2}, DataType::kFLOAT32, + GetDevice()) + ->RequiresGrad(); + const std::vector weight_values{5.0f, 6.0f, 7.0f, 8.0f}; + auto weight = std::make_shared(weight_values.data(), std::vector{1, 1, 2, 2}, DataType::kFLOAT32, + GetDevice()) + ->RequiresGrad(); + auto bias = std::make_shared(std::vector{1}, DataType::kFLOAT32, GetDevice(), true); + bias->Fill(0.25f); + + auto conv_fn = std::make_shared(); + auto result = conv_fn->Apply({input, weight, bias}); + auto grad_output + = std::make_shared(std::vector{1, 1, 1, 1}, DataType::kFLOAT32, GetDevice(), true); + grad_output->Fill(1.0f); + auto grad_inputs = conv_fn->Backward({grad_output}); + ASSERT_EQ(grad_inputs.size(), 3); + + EXPECT_EQ(grad_inputs[0]->Dims(), (std::vector{1, 1, 2, 2})); + test::ExpectTensorFloatEqual(grad_inputs[0], weight_values); + EXPECT_EQ(grad_inputs[1]->Dims(), (std::vector{1, 1, 2, 2})); + test::ExpectTensorFloatEqual(grad_inputs[1], input_values); + test::ExpectTensorFloatEqual(grad_inputs[2], std::vector{1.0f}); +} + +// Training-shaped call: the activation is a leaf without requires_grad, so only the parameter +// gradients must be produced. +TEST_P(AutogradConvBackwardTest, ConvBackwardInputNotRequired) { + auto input = std::make_shared(std::vector{2, 2, 4, 4}, DataType::kFLOAT32, GetDevice()); + input->Fill(1.0f); + auto weight = std::make_shared(std::vector{3, 2, 3, 3}, DataType::kFLOAT32, GetDevice(), true); + weight->Fill(1.0f); + auto bias = std::make_shared(std::vector{3}, DataType::kFLOAT32, GetDevice(), true); + bias->Fill(1.0f); + + auto conv_fn = std::make_shared(); + auto result = conv_fn->Apply({input, weight, bias}); + auto grad = std::make_shared(std::vector{2, 3, 2, 2}, DataType::kFLOAT32, GetDevice(), true); + grad->Fill(1.0f); + auto grad_inputs = conv_fn->Backward({grad}); + ASSERT_EQ(grad_inputs.size(), 3); + EXPECT_EQ(grad_inputs[0], nullptr); + EXPECT_NE(grad_inputs[1], nullptr); + EXPECT_NE(grad_inputs[2], nullptr); + // 2 images * 4 windows * 9 taps of 1.0 + test::ExpectTensorFloatEqual(grad_inputs[1], 8.0f); + // 2 images * 4 spatial positions + test::ExpectTensorFloatEqual(grad_inputs[2], 8.0f); +} + +// Random 4D case checked against PyTorch autograd (fixed seed, see file-level comment). +TEST_P(AutogradConvBackwardTest, ConvBackwardTorchGolden) { + const std::vector input_values{ + 2.014464855f, -0.671311080f, -0.945236862f, -0.096128546f, 0.889558852f, 1.726294398f, -0.078931913f, + 0.205890238f, 0.094608001f, 0.173616216f, 0.437049866f, -0.557070732f, 0.455583423f, -0.736482263f, + -0.718647420f, 0.926798999f, 1.732108712f, -0.138735890f, -1.508729339f, 2.127460241f, 0.079739936f, + 0.217489868f, 0.589455068f, -0.052687794f, -1.659490585f, -0.380473137f, -0.166608080f, 0.350353301f, + -1.411424637f, 0.626776755f, 1.336345553f, -0.203511983f, -0.757243156f, 1.381434083f, -0.417290032f, + -0.668979943f, -0.569557011f, -1.035447836f, -0.555948615f, -1.279154539f, 0.056448594f, -0.388204783f, + 2.071142673f, -0.803012371f, 1.956272721f, 0.610922813f, -1.668144941f, -0.628582716f, -0.756579638f, + 0.987797141f, -0.849777400f, -1.654165864f, -0.326211363f, -0.040599853f, 0.281366080f, 0.499955416f, + -0.518889010f, -0.409607649f, 0.689743340f, 1.129151702f, 0.798748374f, -1.004696369f, -0.104189672f, + 1.257721663f, + }; + auto input = std::make_shared(input_values.data(), std::vector{2, 2, 4, 4}, DataType::kFLOAT32, + GetDevice()) + ->RequiresGrad(); + const std::vector weight_values{ + -2.448281050f, -0.608355999f, -0.958948731f, 1.009868145f, 0.031469870f, -1.101765275f, 0.571945250f, + -1.224039674f, -0.426795214f, -0.612229824f, -0.744571507f, -1.164491534f, -0.726353824f, -1.124407530f, + -0.498482078f, -0.423681349f, -1.231436968f, -0.689355612f, -0.170030892f, -0.120427191f, 1.032825112f, + 1.369398594f, 1.237134099f, 0.458183825f, -0.396747291f, 0.548917472f, 0.043999992f, 0.891736746f, + -1.292207360f, 1.713299513f, 0.982937992f, 0.432634443f, 0.293460459f, 0.318505853f, -0.146791920f, + 0.983720183f, -0.162205622f, 0.616937041f, 0.929942310f, -0.411319345f, 0.301661789f, -1.444623590f, + -0.730446517f, 2.391211271f, 0.400584310f, -0.909495890f, -0.702420592f, 0.372547686f, 0.994188488f, + 0.768918931f, 0.740486383f, -1.176032424f, 0.868664801f, 0.422459871f, + }; + auto weight = std::make_shared(weight_values.data(), std::vector{3, 2, 3, 3}, DataType::kFLOAT32, + GetDevice()) + ->RequiresGrad(); + const std::vector bias_values{-0.640517414f, 0.982376575f, 0.729396820f}; + auto bias = std::make_shared(bias_values.data(), std::vector{3}, DataType::kFLOAT32, GetDevice()) + ->RequiresGrad(); + + auto conv_fn = std::make_shared(); + auto result = conv_fn->Apply({input, weight, bias}); + const std::vector grad_output_values{ + -1.277466655f, -1.339194655f, -1.159473658f, 0.086730979f, 0.438649237f, 0.602657557f, + -0.738640666f, -0.410562515f, 0.201109841f, 0.716394842f, 0.003985748f, 0.676588118f, + 0.412507862f, -0.146394536f, -0.880041420f, -1.247093678f, -1.212666154f, -0.393535852f, + 0.564200640f, 0.279093444f, -2.798053741f, -1.615590930f, -0.518746257f, -0.283052236f, + }; + auto grad_output = std::make_shared(grad_output_values.data(), std::vector{2, 3, 2, 2}, + DataType::kFLOAT32, GetDevice()); + + auto grad_inputs = conv_fn->Backward({grad_output}); + ASSERT_EQ(grad_inputs.size(), 3); + ASSERT_NE(grad_inputs[0], nullptr); + ASSERT_NE(grad_inputs[1], nullptr); + ASSERT_NE(grad_inputs[2], nullptr); + + const std::vector expected_grad_input{ + 3.020392179f, 3.908452511f, 3.049194813f, 2.572864771f, 2.191555023f, 0.285838872f, 3.004242659f, + 0.838663280f, -3.235622168f, -0.945002913f, 4.960352898f, -0.376030058f, -0.373013139f, 0.741603315f, + 1.750292301f, 0.215949476f, 0.990354240f, 0.948824525f, 2.029216766f, 2.858904839f, 1.606565118f, + 4.838562012f, 3.308694601f, 0.822548687f, 0.564552128f, 2.793590069f, 4.189336777f, 2.155962467f, + 0.251298606f, 0.576505065f, 0.615549624f, -0.177835792f, -0.349884391f, -1.143750548f, -5.110340118f, + -1.768475175f, 2.049649715f, 0.845555902f, 3.547870159f, 3.535832882f, 2.858142614f, -6.855783939f, + -3.052171230f, 1.308768988f, -0.348264217f, -0.470771074f, 1.195474505f, 0.431147456f, 1.210889816f, + 4.433355331f, -1.848075271f, -1.105654240f, -2.759691477f, -3.466245174f, -1.316413283f, 0.586114824f, + 3.407652855f, 0.809629261f, -2.299243927f, -0.474773228f, 1.162620664f, 1.500420690f, 2.191398382f, + 1.014662623f, + }; + const std::vector expected_grad_weight{ + -1.278161168f, 2.350493193f, 3.455903053f, -3.191775799f, -4.849991322f, -1.593288898f, -3.348839760f, + 1.063873172f, 4.325201988f, -2.219501734f, 2.213379383f, -2.589243412f, 2.336852074f, -1.198173761f, + -2.431171656f, 4.716279507f, 0.842501462f, -3.181046009f, -1.122435212f, -4.357190609f, -0.400157630f, + 2.311078310f, 2.235447884f, 2.117345333f, 1.470544577f, 0.713402152f, -3.305890560f, 0.861350715f, + -2.180587769f, 2.186314106f, 1.552511692f, 0.699985623f, 0.372371852f, 0.788659871f, -1.649568439f, + -1.754016399f, 1.571256876f, -3.355353594f, 2.778887510f, 4.880561829f, 3.997585535f, 2.531520605f, + -1.071573257f, -2.245790720f, -3.141553879f, 1.098210573f, -2.158659220f, 5.950089931f, 1.271268964f, + 0.028057545f, -1.955230832f, 1.795806646f, 1.293214321f, -3.971021652f, + }; + const std::vector expected_grad_bias{-5.550425529f, -0.870804369f, -3.617365122f}; + + EXPECT_EQ(grad_inputs[0]->Dims(), (std::vector{2, 2, 4, 4})); + test::ExpectTensorNear(grad_inputs[0], expected_grad_input, kGoldenAbsError); + EXPECT_EQ(grad_inputs[1]->Dims(), (std::vector{3, 2, 3, 3})); + test::ExpectTensorNear(grad_inputs[1], expected_grad_weight, kGoldenAbsError); + EXPECT_EQ(grad_inputs[2]->Dims(), (std::vector{3})); + test::ExpectTensorNear(grad_inputs[2], expected_grad_bias, kGoldenAbsError); +} + +// An empty batch is degenerate but must not crash nor return uninitialized gradients: the +// parameter gradients (weight/bias) are sums over zero images, so they are exact zeros. +TEST_P(AutogradConvBackwardTest, ConvBackwardEmptyBatch) { + auto input + = std::make_shared(std::vector{0, 2, 4, 4}, DataType::kFLOAT32, GetDevice())->RequiresGrad(); + auto weight = std::make_shared(std::vector{3, 2, 3, 3}, DataType::kFLOAT32, GetDevice(), true) + ->RequiresGrad(); + auto bias + = std::make_shared(std::vector{3}, DataType::kFLOAT32, GetDevice(), true)->RequiresGrad(); + + auto conv_fn = std::make_shared(); + auto result = conv_fn->Apply({input, weight, bias}); + auto grad_output = std::make_shared(std::vector{0, 3, 2, 2}, DataType::kFLOAT32, GetDevice()); + auto grads = conv_fn->Backward({grad_output}); + ASSERT_EQ(grads.size(), 3); + ASSERT_NE(grads[0], nullptr); + ASSERT_NE(grads[1], nullptr); + ASSERT_NE(grads[2], nullptr); + + // grad_input has an empty batch, grad_weight and grad_bias are exact zeros over the empty batch. + EXPECT_EQ(grads[0]->Dims(), (std::vector{0, 2, 4, 4})); + EXPECT_EQ(grads[0]->NumElements(), 0); + EXPECT_EQ(grads[1]->Dims(), (std::vector{3, 2, 3, 3})); + test::ExpectTensorFloatEqual(grads[1], 0.0f); + EXPECT_EQ(grads[2]->Dims(), (std::vector{3})); + test::ExpectTensorFloatEqual(grads[2], 0.0f); +} + +INFINI_TRAIN_REGISTER_TEST(AutogradConvBackwardTest); diff --git a/tests/autograd/test_autograd_conv_forward.cc b/tests/autograd/test_autograd_conv_forward.cc new file mode 100644 index 000000000..85d458008 --- /dev/null +++ b/tests/autograd/test_autograd_conv_forward.cc @@ -0,0 +1,212 @@ +#include +#include +#include + +#include "gtest/gtest.h" + +#include "infini_train/include/autograd/conv.h" +#include "infini_train/include/tensor.h" + +#include "tests/common/test_utils.h" + +using namespace infini_train; + +namespace { +// Deterministic pseudo-random golden data, generated with torch 2.14 under seed 20260905 +// (conv2d / autograd.grad reference). +constexpr float kGoldenAbsError = 2e-6f; +} // namespace + +class AutogradConvForwardTest : public infini_train::test::InfiniTrainTest {}; + +// An asymmetric kernel pins the cross-correlation semantics: flipping the kernel would swap +// the roles of the taps and change every output value. +TEST_P(AutogradConvForwardTest, ConvForwardAsymmetricKernel) { + std::vector input_values; + for (int idx = 0; idx < 16; ++idx) { input_values.push_back(static_cast(idx)); } + auto input = std::make_shared(input_values.data(), std::vector{1, 1, 4, 4}, DataType::kFLOAT32, + GetDevice()); + const std::vector weight_values{1.0f, 2.0f, 3.0f, 4.0f}; + auto weight = std::make_shared(weight_values.data(), std::vector{1, 1, 2, 2}, DataType::kFLOAT32, + GetDevice()); + // out[i][j] = 1*x[i][j] + 2*x[i][j+1] + 3*x[i+1][j] + 4*x[i+1][j+1] + const std::vector expected{34.0f, 44.0f, 54.0f, 74.0f, 84.0f, 94.0f, 114.0f, 124.0f, 134.0f}; + + auto conv_fn = std::make_shared(); + auto result = conv_fn->Apply({input, weight}); + ASSERT_EQ(result.size(), 1); + EXPECT_EQ(result[0]->Dims(), (std::vector{1, 1, 3, 3})); + test::ExpectTensorFloatEqual(result[0], expected); +} + +TEST_P(AutogradConvForwardTest, ConvForwardMultiChannelBatchBias) { + // input(n, c, h, w) = 10*n + 5*c + 3*h + w with H = W = 3. + std::vector input_values; + for (int n = 0; n < 2; ++n) { + for (int c = 0; c < 2; ++c) { + for (int h = 0; h < 3; ++h) { + for (int w = 0; w < 3; ++w) { input_values.push_back(10.0f * n + 5.0f * c + 3.0f * h + w); } + } + } + } + auto input = std::make_shared(input_values.data(), std::vector{2, 2, 3, 3}, DataType::kFLOAT32, + GetDevice()); + // out channel 0 sums both channel patches (all-ones kernel), out channel 1 mixes + // x(n, 0, i, j) + 2 * x(n, 1, i+1, j+1) to exercise per-tap / per-channel weighting. + const std::vector weight_values{1.0f, 1.0f, 1.0f, 1.0f, 1.0f, 1.0f, 1.0f, 1.0f, + 1.0f, 0.0f, 0.0f, 0.0f, 0.0f, 0.0f, 0.0f, 2.0f}; + auto weight = std::make_shared(weight_values.data(), std::vector{2, 2, 2, 2}, DataType::kFLOAT32, + GetDevice()); + const std::vector bias_values{0.5f, -0.25f}; + auto bias = std::make_shared(bias_values.data(), std::vector{2}, DataType::kFLOAT32, GetDevice()); + const std::vector expected{36.5f, 44.5f, 60.5f, 68.5f, 17.75f, 20.75f, 26.75f, 29.75f, + 116.5f, 124.5f, 140.5f, 148.5f, 47.75f, 50.75f, 56.75f, 59.75f}; + + auto conv_fn = std::make_shared(); + auto result = conv_fn->Apply({input, weight, bias}); + ASSERT_EQ(result.size(), 1); + EXPECT_EQ(result[0]->Dims(), (std::vector{2, 2, 2, 2})); + test::ExpectTensorFloatEqual(result[0], expected); +} + +TEST_P(AutogradConvForwardTest, ConvForwardKernelOne) { + const std::vector input_values{1.0f, 2.0f, 3.0f, 4.0f, 5.0f, 6.0f, 7.0f, 8.0f, 9.0f, + 10.0f, 20.0f, 30.0f, 40.0f, 50.0f, 60.0f, 70.0f, 80.0f, 90.0f}; + auto input = std::make_shared(input_values.data(), std::vector{1, 2, 3, 3}, DataType::kFLOAT32, + GetDevice()); + // 1x1 kernels reduce conv to a per-position channel mix: out0 = 2*c0 + 3*c1, out1 = c0 - c1. + const std::vector weight_values{2.0f, 3.0f, 1.0f, -1.0f}; + auto weight = std::make_shared(weight_values.data(), std::vector{2, 2, 1, 1}, DataType::kFLOAT32, + GetDevice()); + const std::vector expected{32.0f, 64.0f, 96.0f, 128.0f, 160.0f, 192.0f, 224.0f, 256.0f, 288.0f, + -9.0f, -18.0f, -27.0f, -36.0f, -45.0f, -54.0f, -63.0f, -72.0f, -81.0f}; + + auto conv_fn = std::make_shared(); + auto result = conv_fn->Apply({input, weight}); + ASSERT_EQ(result.size(), 1); + EXPECT_EQ(result[0]->Dims(), (std::vector{1, 2, 3, 3})); + test::ExpectTensorFloatEqual(result[0], expected); +} + +// H = W = kernel collapses the output to 1x1: the only patch covers the whole image. +TEST_P(AutogradConvForwardTest, ConvForwardKernelEqualsSpatial) { + std::vector input_values; + for (int idx = 0; idx < 9; ++idx) { input_values.push_back(static_cast(idx)); } + for (int idx = 0; idx < 9; ++idx) { input_values.push_back(2.0f * idx); } + auto input = std::make_shared(input_values.data(), std::vector{2, 1, 3, 3}, DataType::kFLOAT32, + GetDevice()); + auto weight = std::make_shared(std::vector{1, 1, 3, 3}, DataType::kFLOAT32, GetDevice(), true); + weight->Fill(1.0f); + auto bias = std::make_shared(std::vector{1}, DataType::kFLOAT32, GetDevice(), true); + bias->Fill(4.0f); + + auto conv_fn = std::make_shared(); + auto result = conv_fn->Apply({input, weight, bias}); + ASSERT_EQ(result.size(), 1); + EXPECT_EQ(result[0]->Dims(), (std::vector{2, 1, 1, 1})); + // batch 0: sum(0..8) + 4, batch 1: sum(0, 2, ..., 16) + 4 + test::ExpectTensorFloatEqual(result[0], std::vector{40.0f, 76.0f}); +} + +// Magnitudes near the fp32 integer-exact range: every product and partial sum stays exact, so +// this guards against spurious overflow / precision loss on the accumulation path. +TEST_P(AutogradConvForwardTest, ConvForwardExtremeValues) { + auto input = std::make_shared(std::vector{1, 1, 4, 4}, DataType::kFLOAT32, GetDevice(), true); + input->Fill(static_cast(1 << 30)); + auto weight = std::make_shared(std::vector{1, 1, 3, 3}, DataType::kFLOAT32, GetDevice(), true); + weight->Fill(4.0f); + + auto conv_fn = std::make_shared(); + auto result = conv_fn->Apply({input, weight}); + ASSERT_EQ(result.size(), 1); + EXPECT_EQ(result[0]->Dims(), (std::vector{1, 1, 2, 2})); + // 9 taps * 2^30 * 4 = 9 * 2^32, exactly representable in fp32. + test::ExpectTensorFloatEqual(result[0], static_cast(9.0 * (1ULL << 32))); + + auto neg_input = std::make_shared(std::vector{1, 1, 4, 4}, DataType::kFLOAT32, GetDevice(), true); + neg_input->Fill(static_cast(-(1 << 30))); + auto neg_result = conv_fn->Apply({neg_input, weight}); + test::ExpectTensorFloatEqual(neg_result[0], static_cast(-9.0 * (1ULL << 32))); +} + +// Random 4D case checked against PyTorch conv2d (fixed seed, see file-level comment). +TEST_P(AutogradConvForwardTest, ConvForwardTorchGolden) { + const std::vector input_values{ + 2.014464855f, -0.671311080f, -0.945236862f, -0.096128546f, 0.889558852f, 1.726294398f, -0.078931913f, + 0.205890238f, 0.094608001f, 0.173616216f, 0.437049866f, -0.557070732f, 0.455583423f, -0.736482263f, + -0.718647420f, 0.926798999f, 1.732108712f, -0.138735890f, -1.508729339f, 2.127460241f, 0.079739936f, + 0.217489868f, 0.589455068f, -0.052687794f, -1.659490585f, -0.380473137f, -0.166608080f, 0.350353301f, + -1.411424637f, 0.626776755f, 1.336345553f, -0.203511983f, -0.757243156f, 1.381434083f, -0.417290032f, + -0.668979943f, -0.569557011f, -1.035447836f, -0.555948615f, -1.279154539f, 0.056448594f, -0.388204783f, + 2.071142673f, -0.803012371f, 1.956272721f, 0.610922813f, -1.668144941f, -0.628582716f, -0.756579638f, + 0.987797141f, -0.849777400f, -1.654165864f, -0.326211363f, -0.040599853f, 0.281366080f, 0.499955416f, + -0.518889010f, -0.409607649f, 0.689743340f, 1.129151702f, 0.798748374f, -1.004696369f, -0.104189672f, + 1.257721663f, + }; + auto input = std::make_shared(input_values.data(), std::vector{2, 2, 4, 4}, DataType::kFLOAT32, + GetDevice()); + const std::vector weight_values{ + -2.448281050f, -0.608355999f, -0.958948731f, 1.009868145f, 0.031469870f, -1.101765275f, 0.571945250f, + -1.224039674f, -0.426795214f, -0.612229824f, -0.744571507f, -1.164491534f, -0.726353824f, -1.124407530f, + -0.498482078f, -0.423681349f, -1.231436968f, -0.689355612f, -0.170030892f, -0.120427191f, 1.032825112f, + 1.369398594f, 1.237134099f, 0.458183825f, -0.396747291f, 0.548917472f, 0.043999992f, 0.891736746f, + -1.292207360f, 1.713299513f, 0.982937992f, 0.432634443f, 0.293460459f, 0.318505853f, -0.146791920f, + 0.983720183f, -0.162205622f, 0.616937041f, 0.929942310f, -0.411319345f, 0.301661789f, -1.444623590f, + -0.730446517f, 2.391211271f, 0.400584310f, -0.909495890f, -0.702420592f, 0.372547686f, 0.994188488f, + 0.768918931f, 0.740486383f, -1.176032424f, 0.868664801f, 0.422459871f, + }; + auto weight = std::make_shared(weight_values.data(), std::vector{3, 2, 3, 3}, DataType::kFLOAT32, + GetDevice()); + const std::vector bias_values{-0.640517414f, 0.982376575f, 0.729396820f}; + auto bias = std::make_shared(bias_values.data(), std::vector{3}, DataType::kFLOAT32, GetDevice()); + const std::vector expected{ + -2.073040247f, 1.047590375f, -2.980665922f, -6.143202305f, 1.986891150f, 9.790361404f, + 0.198000669f, 0.174701050f, 0.093229778f, 2.889255762f, -0.725508392f, 0.107455865f, + 1.509838820f, -5.231677532f, 2.129449368f, 4.556076050f, -5.035542011f, -1.076976299f, + 0.719715834f, 1.997749925f, 1.977930188f, 7.832613468f, -5.848879814f, -0.770412445f, + }; + + auto conv_fn = std::make_shared(); + auto result = conv_fn->Apply({input, weight, bias}); + ASSERT_EQ(result.size(), 1); + EXPECT_EQ(result[0]->Dims(), (std::vector{2, 3, 2, 2})); + test::ExpectTensorNear(result[0], expected, kGoldenAbsError); +} + +TEST_P(AutogradConvForwardTest, ConvForwardRejectsInvalidShapes) { + EXPECT_DEATH( + { + auto input = std::make_shared(std::vector{1, 2, 3, 3}, DataType::kFLOAT32, GetDevice()); + auto weight = std::make_shared(std::vector{1, 3, 2, 2}, DataType::kFLOAT32, GetDevice()); + auto conv_fn = std::make_shared(); + auto result = conv_fn->Apply({input, weight}); + (void)result; + }, + ""); + + EXPECT_DEATH( + { + auto input = std::make_shared(std::vector{1, 1, 2, 2}, DataType::kFLOAT32, GetDevice()); + auto weight = std::make_shared(std::vector{1, 1, 3, 3}, DataType::kFLOAT32, GetDevice()); + auto conv_fn = std::make_shared(); + auto result = conv_fn->Apply({input, weight}); + (void)result; + }, + ""); +} + +// An empty batch is a degenerate but valid call: the kernel must not launch a zero-block +// configuration; the output tensor keeps the (N=0, ...) shape and is returned untouched. +TEST_P(AutogradConvForwardTest, ConvForwardEmptyBatch) { + auto input = std::make_shared(std::vector{0, 2, 4, 4}, DataType::kFLOAT32, GetDevice()); + auto weight = std::make_shared(std::vector{3, 2, 3, 3}, DataType::kFLOAT32, GetDevice(), true); + weight->Fill(1.0f); + + auto conv_fn = std::make_shared(); + auto result = conv_fn->Apply({input, weight}); + ASSERT_EQ(result.size(), 1); + EXPECT_EQ(result[0]->Dims(), (std::vector{0, 3, 2, 2})); + EXPECT_EQ(result[0]->NumElements(), 0); +} + +INFINI_TRAIN_REGISTER_TEST(AutogradConvForwardTest); diff --git a/tests/autograd/test_autograd_conv_train.cc b/tests/autograd/test_autograd_conv_train.cc new file mode 100644 index 000000000..d2bb6999e --- /dev/null +++ b/tests/autograd/test_autograd_conv_train.cc @@ -0,0 +1,95 @@ +#include +#include +#include +#include + +#include "gtest/gtest.h" + +#include "infini_train/include/nn/modules/conv.h" +#include "infini_train/include/nn/modules/linear.h" +#include "infini_train/include/nn/modules/loss.h" +#include "infini_train/include/optimizer.h" +#include "infini_train/include/tensor.h" + +#include "tests/common/test_utils.h" + +using namespace infini_train; + +// Minimal end-to-end training step: Conv2d -> Flatten -> Linear -> CrossEntropy, one +// loss->Backward() through the autograd graph and one SGD step. Loss, gradients and +// parameters are inspected through host copies, so one body serves both devices. +class AutogradConvTrainTest : public infini_train::test::InfiniTrainTest {}; + +TEST_P(AutogradConvTrainTest, ConvTrainStepUpdatesParams) { + const Device device = GetDevice(); + const Device host = Device(); + constexpr float kLearningRate = 0.1f; + + std::vector input_values; + for (int idx = 0; idx < 2 * 1 * 8 * 8; ++idx) { input_values.push_back((idx * 37) % 101 / 50.0f - 1.0f); } + auto input + = std::make_shared(input_values.data(), std::vector{2, 1, 8, 8}, DataType::kFLOAT32, device); + auto target = std::make_shared(std::vector{2}, DataType::kINT64, device); + target->Fill(0); + + auto conv = std::make_shared(1, 2, 3, true, device); + auto fc = std::make_shared(2 * 6 * 6, 3, true, device); + auto loss_fn = std::make_shared(); + + auto conv_out = (*conv)({input})[0]; + ASSERT_EQ(conv_out->Dims(), (std::vector{2, 2, 6, 6})); + auto flat = conv_out->Flatten(1); + ASSERT_EQ(flat->Dims(), (std::vector{2, 72})); + auto logits = (*fc)({flat})[0]; + ASSERT_EQ(logits->Dims(), (std::vector{2, 3})); + auto loss = (*loss_fn)({logits, target})[0]; + ASSERT_TRUE(loss->Dims().empty()); + const auto loss_cpu = loss->To(host); + EXPECT_TRUE(std::isfinite(*static_cast(loss_cpu.DataPtr()))); + + loss->Backward(); + + std::vector> params = conv->Parameters(); + for (auto ¶m : fc->Parameters()) { params.push_back(param); } + ASSERT_EQ(params.size(), 4); + + // Every parameter must have received a finite, non-trivial gradient through the graph. + for (const auto ¶m : params) { + const auto &grad = param->grad(); + ASSERT_NE(grad, nullptr); + ASSERT_EQ(grad->Dims(), param->Dims()); + const auto grad_cpu = grad->To(host); + const auto *grad_data = static_cast(grad_cpu.DataPtr()); + float max_abs = 0.0f; + for (size_t idx = 0; idx < grad_cpu.NumElements(); ++idx) { + ASSERT_TRUE(std::isfinite(grad_data[idx])); + max_abs = std::max(max_abs, std::fabs(grad_data[idx])); + } + EXPECT_GT(max_abs, 0.0f); + } + + // Snapshot into freshly allocated host buffers: a same-device To(host) view would alias + // the live storage and silently track the optimizer's in-place update. + std::vector> old_values; + for (const auto ¶m : params) { + auto old_value = std::make_shared(param->Dims(), param->Dtype(), host); + old_value->CopyFrom(*param); + old_values.push_back(old_value); + } + + auto optimizer = std::make_shared(params, kLearningRate); + optimizer->Step(); + + for (size_t p = 0; p < params.size(); ++p) { + const auto new_cpu = params[p]->To(host); + const auto grad_cpu = params[p]->grad()->To(host); + const auto *old_data = static_cast(old_values[p]->DataPtr()); + const auto *grad_data = static_cast(grad_cpu.DataPtr()); + const auto *new_data = static_cast(new_cpu.DataPtr()); + for (size_t idx = 0; idx < params[p]->NumElements(); ++idx) { + EXPECT_FLOAT_EQ(new_data[idx], old_data[idx] - kLearningRate * grad_data[idx]) << "param " << p; + } + } +} + +INFINI_TRAIN_REGISTER_TEST(AutogradConvTrainTest); diff --git a/tests/autograd/test_autograd_flatten_grad.cc b/tests/autograd/test_autograd_flatten_grad.cc new file mode 100644 index 000000000..d9917dcb0 --- /dev/null +++ b/tests/autograd/test_autograd_flatten_grad.cc @@ -0,0 +1,65 @@ +#include +#include +#include + +#include "gtest/gtest.h" + +#include "infini_train/include/nn/parallel/global.h" +#include "infini_train/include/tensor.h" + +#include "tests/common/test_utils.h" + +using namespace infini_train; + +// Gradient contract of Tensor::Flatten: it is a pure reshape (via the no-op view), +// so d(loss)/d(x) must be the reshape-back of the downstream gradient. +// The composition x -> flatten -> x*x -> sum makes the expected grad = 2*x. +class AutogradFlattenGradTest : public infini_train::test::InfiniTrainTest {}; + +namespace { +void CheckFlattenGrad(const Device &device, const std::vector &input_dims, int64_t start, int64_t end, + const std::vector &expected_flat_dims) { + const Device host = Device(); + + std::vector values; + for (size_t idx = 0; idx < static_cast( + std::accumulate(input_dims.begin(), input_dims.end(), 1, std::multiplies())); + ++idx) { + values.push_back(((idx % 17) - 8) / 4.0f); + } + auto x = std::make_shared(input_dims, DataType::kFLOAT32, device, true); + auto storage = std::make_shared(values.data(), input_dims, DataType::kFLOAT32, device); + x->CopyFrom(*storage); + + auto flat = x->Flatten(start, end); + ASSERT_EQ(flat->Dims(), expected_flat_dims); + + auto sq = flat->Mul(flat); + auto reduced = sq; + while (reduced->Dims().size() > 0) { reduced = reduced->Sum(0); } + reduced->Backward(); + + const auto &grad = x->grad(); + ASSERT_NE(grad, nullptr); + ASSERT_EQ(grad->Dims(), input_dims); + const auto grad_cpu = grad->To(host); + const float *grad_data = static_cast(grad_cpu.DataPtr()); + for (size_t idx = 0; idx < values.size(); ++idx) { + EXPECT_FLOAT_EQ(grad_data[idx], 2.0f * values[idx]) << "gradient mismatch at flat index " << idx; + } +} +} // namespace + +// 4-D conv-output layout: (N, C, H, W) flattening the channel dim with start=1, +// the same call pattern as the MNIST demo. +TEST_P(AutogradFlattenGradTest, FlattenGradFourDConvLayout) { + CheckFlattenGrad(GetDevice(), {2, 1, 4, 4}, 1, -1, {2, 16}); +} + +TEST_P(AutogradFlattenGradTest, FlattenGradThreeDStartEnd) { CheckFlattenGrad(GetDevice(), {2, 3, 4}, 0, 1, {6, 4}); } + +TEST_P(AutogradFlattenGradTest, FlattenGradTwoDStartEnd) { CheckFlattenGrad(GetDevice(), {3, 4}, 0, 1, {12}); } + +TEST_P(AutogradFlattenGradTest, FlattenGradOneD) { CheckFlattenGrad(GetDevice(), {5}, 0, 0, {5}); } + +INFINI_TRAIN_REGISTER_TEST(AutogradFlattenGradTest); diff --git a/tests/autograd/test_autograd_linear_backward.cc b/tests/autograd/test_autograd_linear_backward.cc index 9ce88eee7..cd5af4ad3 100644 --- a/tests/autograd/test_autograd_linear_backward.cc +++ b/tests/autograd/test_autograd_linear_backward.cc @@ -48,4 +48,108 @@ TEST_P(AutogradLinearBackwardTest, LinearBackwardNoBias) { EXPECT_EQ(grad_inputs.size(), 2); } +// Non-square (bs=4, out_features=3) with distinct values per element: the bias +// gradient must be the column sum of grad_output ([18, 22, 26] below), which a +// row-sum misread would report as [6, 22, 38]. Constant-valued grad_output +// cannot distinguish the two reductions, so values must vary. +TEST_P(AutogradLinearBackwardTest, LinearBackwardValues) { + const std::vector input_values{1.0f, 2.0f, 3.0f, 4.0f, 5.0f, 6.0f, 7.0f, 8.0f}; + const std::vector weight_values{1.0f, 0.0f, 0.0f, 1.0f, 1.0f, 1.0f}; + const std::vector grad_values{0.0f, 1.0f, 2.0f, 3.0f, 4.0f, 5.0f, 6.0f, 7.0f, 8.0f, 9.0f, 10.0f, 11.0f}; + + auto input + = std::make_shared(input_values.data(), std::vector{4, 2}, DataType::kFLOAT32, GetDevice()); + input->set_requires_grad(true); + auto weight + = std::make_shared(weight_values.data(), std::vector{3, 2}, DataType::kFLOAT32, GetDevice()); + weight->set_requires_grad(true); + auto bias = std::make_shared(std::vector{3}, DataType::kFLOAT32, GetDevice(), true); + bias->Fill(0.0f); + auto grad + = std::make_shared(grad_values.data(), std::vector{4, 3}, DataType::kFLOAT32, GetDevice()); + + auto linear_fn = std::make_shared(); + linear_fn->Apply({input, weight, bias}); + auto grad_inputs = linear_fn->Backward({grad}); + + ASSERT_EQ(grad_inputs.size(), 3); + ASSERT_NE(grad_inputs[0], nullptr); + ASSERT_NE(grad_inputs[1], nullptr); + ASSERT_NE(grad_inputs[2], nullptr); + + // grad_bias = grad_output.sum(dim=0) = [0+3+6+9, 1+4+7+10, 2+5+8+11] + EXPECT_EQ(grad_inputs[2]->Dims(), (std::vector{3})); + test::ExpectTensorFloatEqual(grad_inputs[2], {18.0f, 22.0f, 26.0f}); + // grad_input = grad_output * weight + EXPECT_EQ(grad_inputs[0]->Dims(), (std::vector{4, 2})); + test::ExpectTensorFloatEqual(grad_inputs[0], {2.0f, 3.0f, 8.0f, 9.0f, 14.0f, 15.0f, 20.0f, 21.0f}); + // grad_weight = grad_output^T * input + EXPECT_EQ(grad_inputs[1]->Dims(), (std::vector{3, 2})); + test::ExpectTensorFloatEqual(grad_inputs[1], {102.0f, 120.0f, 118.0f, 140.0f, 134.0f, 160.0f}); +} + +// Larger non-square shape with a deterministic pseudo-random pattern, checked +// against a double-precision host reference of the sample-dimension sum. +TEST_P(AutogradLinearBackwardTest, LinearBackwardBiasValues) { + constexpr int64_t bs = 16; + constexpr int64_t out_features = 7; + constexpr int64_t in_features = 3; + + std::vector grad_values(bs * out_features); + for (size_t i = 0; i < grad_values.size(); ++i) { grad_values[i] = static_cast((i * 13 + 7) % 23) - 11.0f; } + + auto input = std::make_shared(std::vector{bs, in_features}, DataType::kFLOAT32, GetDevice(), true); + input->Fill(1.0f); + auto weight = std::make_shared(std::vector{out_features, in_features}, DataType::kFLOAT32, + GetDevice(), true); + weight->Fill(1.0f); + auto bias = std::make_shared(std::vector{out_features}, DataType::kFLOAT32, GetDevice(), true); + bias->Fill(0.0f); + auto grad = std::make_shared(grad_values.data(), std::vector{bs, out_features}, DataType::kFLOAT32, + GetDevice()); + + auto linear_fn = std::make_shared(); + linear_fn->Apply({input, weight, bias}); + auto grad_inputs = linear_fn->Backward({grad}); + + ASSERT_EQ(grad_inputs.size(), 3); + std::vector expected_bias(out_features, 0.0); + for (int64_t i = 0; i < bs; ++i) { + for (int64_t j = 0; j < out_features; ++j) { expected_bias[j] += grad_values[i * out_features + j]; } + } + const std::vector expected_bias_f32(expected_bias.begin(), expected_bias.end()); + test::ExpectTensorNear(grad_inputs[2], expected_bias_f32, 1e-4f); +} + +// The bf16 branch accumulates in fp32 and returns a promoted fp32 grad_bias; +// bf16 kernels are CUDA-only (the CPU Linear path is fp32-only). +TEST_P(AutogradLinearBackwardTest, LinearBackwardBiasBFloat16) { + SKIP_CPU(); + const std::vector grad_values{0.0f, 1.0f, 2.0f, 3.0f, 4.0f, 5.0f, 6.0f, 7.0f, 8.0f, 9.0f, 10.0f, 11.0f}; + + auto grad_f32 + = std::make_shared(grad_values.data(), std::vector{4, 3}, DataType::kFLOAT32, GetDevice()); + auto grad = std::make_shared(grad_f32->To(DataType::kBFLOAT16)); + + auto input_f32 = std::make_shared(std::vector{4, 2}, DataType::kFLOAT32, GetDevice(), true); + input_f32->Fill(1.0f); + auto input = std::make_shared(input_f32->To(DataType::kBFLOAT16)); + input->set_requires_grad(true); + auto weight_f32 = std::make_shared(std::vector{3, 2}, DataType::kFLOAT32, GetDevice(), true); + weight_f32->Fill(1.0f); + auto weight = std::make_shared(weight_f32->To(DataType::kBFLOAT16)); + weight->set_requires_grad(true); + auto bias = std::make_shared(std::vector{3}, DataType::kBFLOAT16, GetDevice(), true); + bias->Fill(0.0f); + + auto linear_fn = std::make_shared(); + linear_fn->Apply({input, weight, bias}); + auto grad_inputs = linear_fn->Backward({grad}); + + ASSERT_EQ(grad_inputs.size(), 3); + // bf16 inputs 0..11 are exact; the fp32 accumulation must reproduce [18, 22, 26]. + EXPECT_EQ(grad_inputs[2]->Dtype(), DataType::kFLOAT32); + test::ExpectTensorFloatEqual(grad_inputs[2], {18.0f, 22.0f, 26.0f}); +} + INFINI_TRAIN_REGISTER_TEST(AutogradLinearBackwardTest); diff --git a/tests/autograd/test_autograd_relu_backward.cc b/tests/autograd/test_autograd_relu_backward.cc new file mode 100644 index 000000000..e76c7b781 --- /dev/null +++ b/tests/autograd/test_autograd_relu_backward.cc @@ -0,0 +1,111 @@ +#include +#include +#include +#include + +#include "gtest/gtest.h" + +#include "infini_train/include/autograd/activations.h" +#include "infini_train/include/nn/parallel/global.h" +#include "infini_train/include/tensor.h" + +#include "tests/common/test_utils.h" + +using namespace infini_train; + +namespace { +uint32_t FloatBits(float value) { + uint32_t bits; + std::memcpy(&bits, &value, sizeof(bits)); + return bits; +} + +// grad respect to input: NaN passes grad through, 0/-0/negative mask to 0, positive keeps grad. +float ExpectedReluGrad(float x, float grad) { return x <= 0.0f ? 0.0f : grad; } +} // namespace + +class AutogradReLUBackwardTest : public infini_train::test::InfiniTrainTest {}; + +TEST_P(AutogradReLUBackwardTest, ReLUBackwardNaNTransparentAndZeroMask) { + const std::vector input_values{std::nanf(""), + -0.0f, + 0.0f, + 1.5f, + -2.5f, + std::numeric_limits::infinity(), + -std::numeric_limits::infinity(), + 1e-30f, + -1e-30f, + std::numeric_limits::max(), + -std::numeric_limits::max()}; + const std::vector grad_values{1.0f, 2.0f, 3.0f, 4.0f, 5.0f, 6.0f, 7.0f, 8.0f, 9.0f, 10.0f, 11.0f}; + + auto input + = std::make_shared(input_values.data(), std::vector{static_cast(input_values.size())}, + DataType::kFLOAT32, GetDevice()); + auto grad + = std::make_shared(grad_values.data(), std::vector{static_cast(grad_values.size())}, + DataType::kFLOAT32, GetDevice()); + + auto relu_fn = std::make_shared(); + auto result = relu_fn->Apply({input}); + ASSERT_EQ(result.size(), 1); + auto grad_inputs = relu_fn->Backward({grad}); + ASSERT_EQ(grad_inputs.size(), 1); + ASSERT_EQ(grad_inputs[0]->Dims(), input->Dims()); + + const auto host_grad_input = grad_inputs[0]->To(Device()); + const float *out = static_cast(host_grad_input.DataPtr()); + + for (size_t idx = 0; idx < input_values.size(); ++idx) { + const float expected = ExpectedReluGrad(input_values[idx], grad_values[idx]); + if (std::isnan(input_values[idx])) { + // The gradient must be passed through bit-for-bit, not zeroed. + EXPECT_EQ(FloatBits(out[idx]), FloatBits(grad_values[idx])) << "NaN gradient must pass through at " << idx; + } else { + EXPECT_EQ(FloatBits(out[idx]), FloatBits(expected)) << "gradient bit mismatch at position " << idx; + } + } +} + +// Two-dimensional case with a mixed-sign input, checking mask and passthrough jointly. +TEST_P(AutogradReLUBackwardTest, ReLUBackwardTwoD) { + const std::vector input_values{-1.0f, 2.0f, -0.0f, 4.0f, -5.0f, 6.0f}; + const std::vector grad_values{1.0f, 2.0f, 3.0f, 4.0f, 5.0f, 6.0f}; + auto input + = std::make_shared(input_values.data(), std::vector{2, 3}, DataType::kFLOAT32, GetDevice()); + auto grad + = std::make_shared(grad_values.data(), std::vector{2, 3}, DataType::kFLOAT32, GetDevice()); + + auto relu_fn = std::make_shared(); + auto result = relu_fn->Apply({input}); + ASSERT_EQ(result.size(), 1); + auto grad_inputs = relu_fn->Backward({grad}); + ASSERT_EQ(grad_inputs.size(), 1); + ASSERT_EQ(grad_inputs[0]->Dims(), (std::vector{2, 3})); + test::ExpectTensorFloatEqual(grad_inputs[0], {0.0f, 2.0f, 0.0f, 4.0f, 0.0f, 6.0f}); +} + +// Gradients entering at the masked (zero) positions must stay 0 regardless of the +// incoming grad value, and vice versa for the active positions. +TEST_P(AutogradReLUBackwardTest, ReLUBackwardGradValuesIndependentOfMask) { + const std::vector input_values{-3.0f, -0.5f, -0.0f, 0.0f, 1.0f, 7.0f}; + const std::vector grad_values{-2.0f, 100.0f, -5.0f, 3.0f, -1.0f, 0.25f}; + auto input + = std::make_shared(input_values.data(), std::vector{6}, DataType::kFLOAT32, GetDevice()); + auto grad = std::make_shared(grad_values.data(), std::vector{6}, DataType::kFLOAT32, GetDevice()); + + auto relu_fn = std::make_shared(); + auto result = relu_fn->Apply({input}); + ASSERT_EQ(result.size(), 1); + auto grad_inputs = relu_fn->Backward({grad}); + ASSERT_EQ(grad_inputs.size(), 1); + + const auto host_grad_input = grad_inputs[0]->To(Device()); + const float *out = static_cast(host_grad_input.DataPtr()); + for (size_t idx = 0; idx < input_values.size(); ++idx) { + EXPECT_EQ(out[idx], ExpectedReluGrad(input_values[idx], grad_values[idx])) << "position " << idx; + } +} + +INFINI_TRAIN_REGISTER_TEST(AutogradReLUBackwardTest); diff --git a/tests/autograd/test_autograd_relu_forward.cc b/tests/autograd/test_autograd_relu_forward.cc new file mode 100644 index 000000000..bc7d54580 --- /dev/null +++ b/tests/autograd/test_autograd_relu_forward.cc @@ -0,0 +1,124 @@ +#include +#include +#include +#include + +#include "gtest/gtest.h" + +#include "infini_train/include/autograd/activations.h" +#include "infini_train/include/nn/modules/activations.h" +#include "infini_train/include/nn/parallel/global.h" +#include "infini_train/include/tensor.h" + +#include "tests/common/test_utils.h" + +using namespace infini_train; + +namespace { +uint32_t FloatBits(float value) { + uint32_t bits; + std::memcpy(&bits, &value, sizeof(bits)); + return bits; +} + +// Reference for the finite-valued part: NaN and -0 must survive untouched, +// negatives collapse to +0. +float ExpectedRelu(float x) { return x < 0.0f ? 0.0f : x; } +} // namespace + +class AutogradReLUForwardTest : public infini_train::test::InfiniTrainTest {}; + +// A single-case battery of NaN, -0, 0, +/-values, infinities and the fp32 extremes. +// NaN and -0 are asserted bit- or positionally because they have no ordering. +TEST_P(AutogradReLUForwardTest, ReLUForwardNaNZeroExtremes) { + const std::vector input_values{std::nanf(""), + -0.0f, + 0.0f, + 1.5f, + -2.5f, + std::numeric_limits::infinity(), + -std::numeric_limits::infinity(), + 1e-30f, + -1e-30f, + std::numeric_limits::max(), + -std::numeric_limits::max()}; + auto input + = std::make_shared(input_values.data(), std::vector{static_cast(input_values.size())}, + DataType::kFLOAT32, GetDevice()); + + auto relu_fn = std::make_shared(); + auto result = relu_fn->Apply({input}); + ASSERT_EQ(result.size(), 1); + ASSERT_EQ(result[0]->Dims(), input->Dims()); + + const auto host_result = result[0]->To(Device()); + const float *out = static_cast(host_result.DataPtr()); + + for (size_t idx = 0; idx < input_values.size(); ++idx) { + if (std::isnan(input_values[idx])) { + EXPECT_TRUE(std::isnan(out[idx])) << "NaN position " << idx << " must stay NaN"; + } else if (std::signbit(input_values[idx]) && input_values[idx] == 0.0f) { + EXPECT_EQ(out[idx], 0.0f) << "-0 position " << idx; + EXPECT_TRUE(std::signbit(out[idx])) << "-0 signbit lost at position " << idx; + } else { + EXPECT_EQ(FloatBits(out[idx]), FloatBits(ExpectedRelu(input_values[idx]))) + << "finite bit mismatch at position " << idx; + } + } +} + +TEST_P(AutogradReLUForwardTest, ReLUForwardModuleTwoD) { + const std::vector input_values{-1.0f, 0.0f, 2.0f, 3.5f, -4.25f, 0.5f}; + auto input + = std::make_shared(input_values.data(), std::vector{2, 3}, DataType::kFLOAT32, GetDevice()); + + auto relu = std::make_shared(); + auto result = (*relu)({input}); + ASSERT_EQ(result.size(), 1); + ASSERT_EQ(result[0]->Dims(), (std::vector{2, 3})); + test::ExpectTensorFloatEqual(result[0], {0.0f, 0.0f, 2.0f, 3.5f, 0.0f, 0.5f}); +} + +// ReLU is elementwise; the 4-D conv-output layout must keep its shape. +TEST_P(AutogradReLUForwardTest, ReLUForwardFourD) { + std::vector input_values; + for (int idx = 0; idx < 2 * 2 * 4 * 4; ++idx) { input_values.push_back((idx * 7) % 13 / 3.0f - 2.0f); } + auto input = std::make_shared(input_values.data(), std::vector{2, 2, 4, 4}, DataType::kFLOAT32, + GetDevice()); + + auto relu_fn = std::make_shared(); + auto result = relu_fn->Apply({input}); + ASSERT_EQ(result.size(), 1); + ASSERT_EQ(result[0]->Dims(), (std::vector{2, 2, 4, 4})); + + const auto host_result = result[0]->To(Device()); + const float *out = static_cast(host_result.DataPtr()); + for (size_t idx = 0; idx < input_values.size(); ++idx) { + EXPECT_EQ(out[idx], ExpectedRelu(input_values[idx])) << "position " << idx; + } +} + +// An empty batch is a valid 0-element tensor: the kernel must not touch memory. +TEST_P(AutogradReLUForwardTest, ReLUForwardEmptyBatch) { + auto input = std::make_shared(std::vector{0, 3}, DataType::kFLOAT32, GetDevice()); + auto relu_fn = std::make_shared(); + auto result = relu_fn->Apply({input}); + ASSERT_EQ(result.size(), 1); + EXPECT_EQ(result[0]->Dims(), (std::vector{0, 3})); + EXPECT_EQ(result[0]->NumElements(), 0); +} + +// The op is declared FP32-only; other dtypes must be rejected loudly. +TEST_P(AutogradReLUForwardTest, ReLUForwardRejectsNonFloat) { + ONLY_CPU(); + auto input = std::make_shared(std::vector{2, 3}, DataType::kFLOAT16, GetDevice()); + EXPECT_DEATH( + { + auto relu_fn = std::make_shared(); + auto result = relu_fn->Apply({input}); + (void)result; + }, + ""); +} + +INFINI_TRAIN_REGISTER_TEST(AutogradReLUForwardTest); diff --git a/tests/autograd/test_autograd_relu_train.cc b/tests/autograd/test_autograd_relu_train.cc new file mode 100644 index 000000000..1d0d130f7 --- /dev/null +++ b/tests/autograd/test_autograd_relu_train.cc @@ -0,0 +1,120 @@ +#include +#include +#include +#include + +#include "gtest/gtest.h" + +#include "infini_train/include/nn/modules/activations.h" +#include "infini_train/include/nn/modules/conv.h" +#include "infini_train/include/nn/modules/linear.h" +#include "infini_train/include/nn/modules/loss.h" +#include "infini_train/include/nn/parallel/global.h" +#include "infini_train/include/optimizer.h" +#include "infini_train/include/tensor.h" + +#include "tests/common/test_utils.h" + +using namespace infini_train; + +// Minimal end-to-end training step with ReLU in the middle: +// Conv2d -> ReLU -> Flatten -> Linear -> CrossEntropy, one loss->Backward() through the +// autograd graph and one SGD step, on CPU and CUDA. +class AutogradReLUTrainTest : public infini_train::test::InfiniTrainTest {}; + +TEST_P(AutogradReLUTrainTest, ReluConvFlattenLinearCrossEntropySgdStep) { + const Device device = GetDevice(); + const Device host = Device(); + constexpr float kLearningRate = 0.1f; + + std::vector input_values; + for (int idx = 0; idx < 2 * 1 * 8 * 8; ++idx) { input_values.push_back((idx * 37) % 101 / 50.0f - 1.0f); } + auto input + = std::make_shared(input_values.data(), std::vector{2, 1, 8, 8}, DataType::kFLOAT32, device); + auto target = std::make_shared(std::vector{2}, DataType::kINT64, device); + target->Fill(0); + + auto conv = std::make_shared(1, 2, 3, true, device); + auto relu = std::make_shared(); + auto fc = std::make_shared(2 * 6 * 6, 3, true, device); + auto loss_fn = std::make_shared(); + + auto conv_out = (*conv)({input})[0]; + ASSERT_EQ(conv_out->Dims(), (std::vector{2, 2, 6, 6})); + auto relu_out = (*relu)({conv_out})[0]; + ASSERT_EQ(relu_out->Dims(), (std::vector{2, 2, 6, 6})); + + // The pre-activations must contain negatives so the ReLU mask actually fires. + bool saw_negative = false; + bool saw_masked_zero = false; + const auto conv_cpu = conv_out->To(host); + const auto relu_cpu = relu_out->To(host); + const float *conv_data = static_cast(conv_cpu.DataPtr()); + const float *relu_data = static_cast(relu_cpu.DataPtr()); + for (size_t idx = 0; idx < conv_cpu.NumElements(); ++idx) { + ASSERT_TRUE(std::isfinite(relu_data[idx])); + if (conv_data[idx] < 0.0f) { + saw_negative = true; + } + if (relu_data[idx] == 0.0f) { + saw_masked_zero = true; + } + } + ASSERT_TRUE(saw_negative); + ASSERT_TRUE(saw_masked_zero); + + auto flat = relu_out->Flatten(1); + ASSERT_EQ(flat->Dims(), (std::vector{2, 72})); + auto logits = (*fc)({flat})[0]; + ASSERT_EQ(logits->Dims(), (std::vector{2, 3})); + auto loss = (*loss_fn)({logits, target})[0]; + ASSERT_TRUE(loss->Dims().empty()); + const auto loss_cpu = loss->To(host); + EXPECT_TRUE(std::isfinite(*static_cast(loss_cpu.DataPtr()))); + + loss->Backward(); + + std::vector> params = conv->Parameters(); + for (auto ¶m : fc->Parameters()) { params.push_back(param); } + ASSERT_EQ(params.size(), 4); + + // Every parameter must have received a finite, non-trivial gradient through the graph. + for (const auto ¶m : params) { + const auto &grad = param->grad(); + ASSERT_NE(grad, nullptr); + ASSERT_EQ(grad->Dims(), param->Dims()); + const auto grad_cpu = grad->To(host); + const auto *grad_data = static_cast(grad_cpu.DataPtr()); + float max_abs = 0.0f; + for (size_t idx = 0; idx < grad_cpu.NumElements(); ++idx) { + ASSERT_TRUE(std::isfinite(grad_data[idx])); + max_abs = std::max(max_abs, std::fabs(grad_data[idx])); + } + EXPECT_GT(max_abs, 0.0f); + } + + std::vector> old_values_cpu; + for (const auto ¶m : params) { + // CopyFrom (not To) so the snapshot owns its buffer: on CPU To(home device) + // returns a view, which would alias the in-place update below. + auto old_value = std::make_shared(param->Dims(), param->Dtype(), host); + old_value->CopyFrom(*param); + old_values_cpu.push_back(old_value); + } + + auto optimizer = std::make_shared(params, kLearningRate); + optimizer->Step(); + + for (size_t p = 0; p < params.size(); ++p) { + const auto new_cpu = params[p]->To(host); + const auto grad_cpu = params[p]->grad()->To(host); + const auto *old_data = static_cast(old_values_cpu[p]->DataPtr()); + const auto *grad_data = static_cast(grad_cpu.DataPtr()); + const auto *new_data = static_cast(new_cpu.DataPtr()); + for (size_t idx = 0; idx < params[p]->NumElements(); ++idx) { + EXPECT_FLOAT_EQ(new_data[idx], old_data[idx] - kLearningRate * grad_data[idx]) << "param " << p; + } + } +} + +INFINI_TRAIN_REGISTER_TEST(AutogradReLUTrainTest);