diff --git a/example/gpt2/main.cc b/example/gpt2/main.cc index 60c0c908..1fce1414 100644 --- a/example/gpt2/main.cc +++ b/example/gpt2/main.cc @@ -479,10 +479,12 @@ void Train(const nn::parallel::Rank &rank) { LOG(INFO) << "Rank " << rank.GlobalRank() << ": finish loss forward"; - auto loss_cpu = loss->To(Device()); - lossf += static_cast(loss_cpu.DataPtr())[0]; LOG(INFO) << "Rank " << rank.GlobalRank() << ": start backward"; loss->Backward(); + // Defer the loss D2H copy until after backward; reading it earlier would synchronize CUDA + // between forward and backward. + auto loss_cpu = loss->To(Device()); + lossf += static_cast(loss_cpu.DataPtr())[0]; LOG(INFO) << "Rank " << rank.GlobalRank() << ": finish backward"; } diff --git a/example/llama3/main.cc b/example/llama3/main.cc index 302e0808..fc543ec8 100644 --- a/example/llama3/main.cc +++ b/example/llama3/main.cc @@ -456,10 +456,12 @@ void Train(const nn::parallel::Rank &rank) { LOG(INFO) << "Rank " << rank.GlobalRank() << ": finish loss forward"; - auto loss_cpu = loss->To(Device()); - lossf += static_cast(loss_cpu.DataPtr())[0]; LOG(INFO) << "Rank " << rank.GlobalRank() << ": start backward"; loss->Backward(); + // Defer the loss D2H copy until after backward; reading it earlier would synchronize CUDA + // between forward and backward. + auto loss_cpu = loss->To(Device()); + lossf += static_cast(loss_cpu.DataPtr())[0]; LOG(INFO) << "Rank " << rank.GlobalRank() << ": finish backward"; } diff --git a/example/mixtral/main.cc b/example/mixtral/main.cc index 84853865..ceabfad9 100644 --- a/example/mixtral/main.cc +++ b/example/mixtral/main.cc @@ -133,11 +133,13 @@ int main(int argc, char *argv[]) { auto y = std::make_shared(y_cpu->To(train_device)); auto logits = (*model)({x})[0]; auto loss = (*loss_fn)({logits, y})[0]; - auto loss_cpu = loss->To(Device()); - lossf += static_cast(loss_cpu.DataPtr())[0] / grad_accum_steps; loss = loss / static_cast(grad_accum_steps); autocast_guard.Disable(); loss->Backward(); + // Defer the loss D2H copy until after backward; reading it earlier would synchronize CUDA + // between forward and backward. + auto loss_cpu = loss->To(Device()); + lossf += static_cast(loss_cpu.DataPtr())[0]; } optimizer->Step(); diff --git a/example/mnist/main.cc b/example/mnist/main.cc index e62257d7..7744e094 100644 --- a/example/mnist/main.cc +++ b/example/mnist/main.cc @@ -70,6 +70,10 @@ int main(int argc, char *argv[]) { optimizer.ZeroGrad(); auto loss = loss_fn.Forward({outputs[0], new_label}); + loss[0]->Backward(); + + // Defer the loss D2H copy until after backward; reading it earlier would synchronize CUDA + // between forward and backward. auto loss_cpu = loss[0]->To(cpu_device); float current_loss = static_cast(loss_cpu.DataPtr())[0]; total_loss += current_loss; @@ -79,7 +83,6 @@ int main(int argc, char *argv[]) { << " loss: " << current_loss; } - loss[0]->Backward(); optimizer.Step(); train_idx += 1; } diff --git a/infini_train/src/kernels/cuda/concat.cu b/infini_train/src/kernels/cuda/concat.cu index a7fa7490..16b5da65 100644 --- a/infini_train/src/kernels/cuda/concat.cu +++ b/infini_train/src/kernels/cuda/concat.cu @@ -186,7 +186,7 @@ std::vector> ConcatBackward(const std::shared_ptr(dvec, dtype, device); - t->Fill(0.0); + // ConcatBackwardKernel maps every grad_output element to exactly one grad tensor element; no Fill is needed. grads.push_back(t); } diff --git a/infini_train/src/kernels/cuda/cross_entropy.cu b/infini_train/src/kernels/cuda/cross_entropy.cu index 8ba2785b..96f987fe 100644 --- a/infini_train/src/kernels/cuda/cross_entropy.cu +++ b/infini_train/src/kernels/cuda/cross_entropy.cu @@ -204,7 +204,7 @@ std::shared_ptr CrossEntropyBackward(const std::shared_ptr &inpu DataTypeList>( {target->Dtype(), input_casted->Dtype()}, [=]() { - grad_input->Fill(0.0); + // One sample block writes all of its num_classes gradient elements; no Fill is needed. const Tinput *output_grad_ptr = static_cast(grad_output->DataPtr()); const Ttarget *target_ptr = static_cast(target->DataPtr()); const Tinput *input_ptr = static_cast(input_casted->DataPtr()); diff --git a/infini_train/src/kernels/cuda/layernorm.cu b/infini_train/src/kernels/cuda/layernorm.cu index 12f3c08d..05d43c37 100644 --- a/infini_train/src/kernels/cuda/layernorm.cu +++ b/infini_train/src/kernels/cuda/layernorm.cu @@ -89,8 +89,7 @@ LayerNormForward(const std::shared_ptr &input, const std::shared_ptr( dtype, [=]() { - mean->Fill(0.0); - rstd->Fill(0.0); + // Each token block writes its mean and rstd exactly once; no Fill is needed. LayerNormForwardKernel<<>>( static_cast(input->DataPtr()), static_cast(weight->DataPtr()), static_cast(bias->DataPtr()), static_cast(mean->DataPtr()), @@ -183,7 +182,7 @@ LayerNormBackward(const std::shared_ptr &input, const std::shared_ptr( dtype, [=]() { - grad_input->Fill(0.0); + // Each token block writes its complete grad_input slice; no Fill is needed. grad_weight->Fill(0.0); grad_bias->Fill(0.0); LayerNormBackwardKernel<<>>( diff --git a/infini_train/src/kernels/cuda/linear.cu b/infini_train/src/kernels/cuda/linear.cu index 2725053e..1b4c1819 100644 --- a/infini_train/src/kernels/cuda/linear.cu +++ b/infini_train/src/kernels/cuda/linear.cu @@ -65,6 +65,7 @@ std::shared_ptr LinearForward(const std::shared_ptr &input, cons infini_train::core::GetDeviceGuardImpl(device.type())->GetStream(device)) ->cuda_stream(); + const float beta = bias ? 1.0f : 0.0f; if (bias) { CHECK_EQ(bias->Dims().size(), 1); CHECK_EQ(bias->Dims()[0], out_features); @@ -78,9 +79,8 @@ std::shared_ptr LinearForward(const std::shared_ptr &input, cons static_cast(output->DataPtr()), static_cast(bias->DataPtr()), bs, out_features); }, "CUDA LinearForward"); - } else { - output->Fill(0.0); } + // In the no-bias path, beta=0 makes cuBLAS fully overwrite output; no Fill is needed. // When bs==1 and fp32, use cublasSgemv (more efficient than GEMM for matrix-vector). // cublasSgemv does not support bf16, so bf16 falls through to Gemm. @@ -94,7 +94,7 @@ std::shared_ptr LinearForward(const std::shared_ptr &input, cons .x = static_cast(input->DataPtr()), .y = static_cast(output->DataPtr()), .alpha = 1.0f, - .beta = 1.0f, // output already initialized with bias or zero above + .beta = beta, }); } else { // cuBLAS is colmun-major @@ -125,7 +125,7 @@ std::shared_ptr LinearForward(const std::shared_ptr &input, cons .C = output->DataPtr(), .ldc = static_cast(out_features), .alpha = 1.0f, - .beta = 1.0f, // bias already written into output; beta=1 accumulates + .beta = beta, .batch_count = 1, .input_dtype = dtype, .output_dtype = dtype, diff --git a/infini_train/src/kernels/cuda/reduction.cu b/infini_train/src/kernels/cuda/reduction.cu index c56470e3..b9fb57e0 100644 --- a/infini_train/src/kernels/cuda/reduction.cu +++ b/infini_train/src/kernels/cuda/reduction.cu @@ -181,7 +181,7 @@ std::shared_ptr ReduceOpBackward(const std::shared_ptr &grad_out core::cuda::DispatchCudaFunc( dtype, [=]() { - grad_input->Fill(0.0); + // The backward kernel assigns every grad_input element on all reduction branches; no Fill is needed. GenericReduceBackwardKernel<<>>( static_cast(grad_input->DataPtr()), static_cast(grad_output->DataPtr()), input ? static_cast(input->DataPtr()) : nullptr, diff --git a/infini_train/src/kernels/cuda/slice.cu b/infini_train/src/kernels/cuda/slice.cu index d030d73a..f71b4be6 100644 --- a/infini_train/src/kernels/cuda/slice.cu +++ b/infini_train/src/kernels/cuda/slice.cu @@ -48,8 +48,7 @@ std::shared_ptr SliceForward(const std::shared_ptr &input, const auto dtype = input->Dtype(); auto new_tensor = std::make_shared(new_dims, dtype, input->GetDevice()); - // NOTE(zbl): must initialize with 0 - new_tensor->Fill(0.0); + // SliceForwardKernel writes every output index in [0, total_elements); no Fill is needed. std::vector src_strides(dims.size(), 0), dst_strides(new_dims.size(), 0); int64_t stride = 1; diff --git a/infini_train/src/kernels/cuda/softmax.cu b/infini_train/src/kernels/cuda/softmax.cu index a7f3b612..ce709d38 100644 --- a/infini_train/src/kernels/cuda/softmax.cu +++ b/infini_train/src/kernels/cuda/softmax.cu @@ -198,7 +198,7 @@ std::shared_ptr SoftmaxBackward(const std::shared_ptr &grad_outp CHECK(dim >= 0 && dim < output->Dims().size()); auto grad_input = std::make_shared(output_dims, promoted_type, output->GetDevice()); - grad_input->Fill(0.0); + // For non-empty tensors, the grid covers every outer/axis/inner index exactly once; no Fill is needed. switch (promoted_type) { DISPATCH_CASE(WRAP(LaunchBackward<256, float>(grad_input, grad_output_promoted, output_promoted, dim);), diff --git a/infini_train/src/kernels/cuda/split.cu b/infini_train/src/kernels/cuda/split.cu index bda0dd70..b52513a2 100644 --- a/infini_train/src/kernels/cuda/split.cu +++ b/infini_train/src/kernels/cuda/split.cu @@ -114,6 +114,7 @@ std::shared_ptr LaunchSplitBackward(const std::vector &input_di const auto &grad = grad_outputs[0]; auto dtype = grad->Dtype(); auto grad_input = std::make_shared(input_dims, dtype, grad->GetDevice()); + // Keep initialization: defensive early returns in SplitBackwardKernel can leave elements unwritten. grad_input->Fill(0.0); int64_t N = std::accumulate(input_dims.begin(), input_dims.begin() + dim, 1, std::multiplies()); diff --git a/infini_train/src/kernels/cuda/stack.cu b/infini_train/src/kernels/cuda/stack.cu index 841940ea..2ffe7795 100644 --- a/infini_train/src/kernels/cuda/stack.cu +++ b/infini_train/src/kernels/cuda/stack.cu @@ -113,7 +113,7 @@ std::vector> StackBackward(const std::vector &i std::vector> grads; for (int i = 0; i < num_inputs; ++i) { auto t = std::make_shared(base_dims, dtype, grad_output->GetDevice()); - t->Fill(0.0); + // StackBackwardKernel writes every element of every grad tensor exactly once; no Fill is needed. grads.push_back(t); } diff --git a/infini_train/src/nn/parallel/pp/pipeline_schedule.cc b/infini_train/src/nn/parallel/pp/pipeline_schedule.cc index 090b7b15..c7c1d16f 100644 --- a/infini_train/src/nn/parallel/pp/pipeline_schedule.cc +++ b/infini_train/src/nn/parallel/pp/pipeline_schedule.cc @@ -255,9 +255,10 @@ float PipelineSchedule::StepMicroBatches(const std::vector(target_on_device)})[0]; loss = loss / n; } - total_loss += static_cast(loss->To(Device()).DataPtr())[0]; - loss->Backward(); + // Defer the loss D2H copy until after backward; reading it earlier would synchronize CUDA + // between forward and backward. + total_loss += static_cast(loss->To(Device()).DataPtr())[0]; } else { auto out_tensor = activations[task.local_chunk_idx][mb][0]; diff --git a/tests/autograd/test_autograd_linear_forward.cc b/tests/autograd/test_autograd_linear_forward.cc index 89d7c537..18c55142 100644 --- a/tests/autograd/test_autograd_linear_forward.cc +++ b/tests/autograd/test_autograd_linear_forward.cc @@ -18,11 +18,12 @@ TEST_P(AutogradLinearForwardTest, LinearForward) { auto weight = std::make_shared(std::vector{4, 3}, DataType::kFLOAT32, GetDevice(), true); weight->Fill(1.0f); auto bias = std::make_shared(std::vector{4}, DataType::kFLOAT32, GetDevice(), true); - bias->Fill(0.0f); + bias->Fill(2.0f); auto linear_fn = std::make_shared(); auto result = linear_fn->Apply({input, weight, bias}); EXPECT_EQ(result.size(), 1); EXPECT_EQ(result[0]->Dims(), (std::vector{2, 4})); + test::ExpectTensorEqual(result[0], 5.0f); } TEST_P(AutogradLinearForwardTest, LinearNoBias) { @@ -34,6 +35,7 @@ TEST_P(AutogradLinearForwardTest, LinearNoBias) { auto result = linear_fn->Apply({input, weight}); EXPECT_EQ(result.size(), 1); EXPECT_EQ(result[0]->Dims(), (std::vector{2, 4})); + test::ExpectTensorEqual(result[0], 3.0f); } TEST_P(AutogradLinearForwardTest, LinearBatch) { diff --git a/tests/autograd/test_autograd_loss.cc b/tests/autograd/test_autograd_loss.cc new file mode 100644 index 00000000..a66c9e99 --- /dev/null +++ b/tests/autograd/test_autograd_loss.cc @@ -0,0 +1,51 @@ +#include +#include + +#include "gtest/gtest.h" + +#include "infini_train/include/autograd/loss.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; + +class AutogradLossTest : public infini_train::test::InfiniTrainTest {}; + +TEST_P(AutogradLossTest, CrossEntropyForwardAndBackwardValues) { + std::vector logits_values{1.0f, 2.0f, 3.0f, 2.0f, 1.0f, 0.0f}; + auto logits + = std::make_shared(logits_values.data(), std::vector{2, 3}, DataType::kFLOAT32, GetDevice()); + auto target = std::make_shared(std::vector{2}, DataType::kINT64, GetDevice()); + target->Fill(0); + + auto cross_entropy = std::make_shared(); + auto result = cross_entropy->Apply({logits, target}); + ASSERT_EQ(result.size(), 1); + ASSERT_TRUE(result[0]->Dims().empty()); + + const float row0_sum = std::exp(1.0f) + std::exp(2.0f) + std::exp(3.0f); + const float row1_sum = std::exp(2.0f) + std::exp(1.0f) + std::exp(0.0f); + const float expected_loss = 0.5f * ((std::log(row0_sum) - 1.0f) + (std::log(row1_sum) - 2.0f)); + test::ExpectTensorNear(result[0], expected_loss, 1e-5f); + + auto grad_output = std::make_shared(std::vector{}, DataType::kFLOAT32, GetDevice()); + grad_output->Fill(1.0f); + auto grad_inputs = cross_entropy->Backward({grad_output}); + ASSERT_EQ(grad_inputs.size(), 2); + ASSERT_NE(grad_inputs[0], nullptr); + EXPECT_EQ(grad_inputs[1], nullptr); + + std::vector expected_grad(logits_values.size()); + for (int row = 0; row < 2; ++row) { + const float sum = row == 0 ? row0_sum : row1_sum; + for (int col = 0; col < 3; ++col) { + const float probability = std::exp(logits_values[row * 3 + col]) / sum; + expected_grad[row * 3 + col] = 0.5f * (probability - (col == 0 ? 1.0f : 0.0f)); + } + } + test::ExpectTensorNear(grad_inputs[0], expected_grad, 1e-5f); +} + +INFINI_TRAIN_REGISTER_TEST(AutogradLossTest); diff --git a/tests/autograd/test_autograd_normalization_backward.cc b/tests/autograd/test_autograd_normalization_backward.cc index d3926ce8..7d486985 100644 --- a/tests/autograd/test_autograd_normalization_backward.cc +++ b/tests/autograd/test_autograd_normalization_backward.cc @@ -25,6 +25,9 @@ TEST_P(AutogradNormalizationBackwardTest, LayerNormBackward) { grad->Fill(1.0f); auto grad_inputs = layernorm_fn->Backward({grad}); EXPECT_EQ(grad_inputs.size(), 3); + test::ExpectTensorNear(grad_inputs[0], 0.0f, 1e-5f); + test::ExpectTensorNear(grad_inputs[1], 0.0f, 1e-5f); + test::ExpectTensorNear(grad_inputs[2], 6.0f, 1e-5f); } TEST_P(AutogradNormalizationBackwardTest, LayerNormBackwardZeroBias) { diff --git a/tests/autograd/test_autograd_reduction_backward.cc b/tests/autograd/test_autograd_reduction_backward.cc index 04709dd2..43e10a77 100644 --- a/tests/autograd/test_autograd_reduction_backward.cc +++ b/tests/autograd/test_autograd_reduction_backward.cc @@ -21,6 +21,7 @@ TEST_P(AutogradReductionBackwardTest, SumBackward) { grad->Fill(1.0f); auto grad_inputs = sum_fn->Backward({grad}); EXPECT_EQ(grad_inputs.size(), 1); + test::ExpectTensorEqual(grad_inputs[0], {1.0f, 1.0f, 1.0f, 1.0f, 1.0f, 1.0f}); } TEST_P(AutogradReductionBackwardTest, MeanBackward) { @@ -32,28 +33,32 @@ TEST_P(AutogradReductionBackwardTest, MeanBackward) { grad->Fill(1.0f); auto grad_inputs = mean_fn->Backward({grad}); EXPECT_EQ(grad_inputs.size(), 1); + test::ExpectTensorEqual(grad_inputs[0], + {1.0f / 3.0f, 1.0f / 3.0f, 1.0f / 3.0f, 1.0f / 3.0f, 1.0f / 3.0f, 1.0f / 3.0f}); } TEST_P(AutogradReductionBackwardTest, MaxBackward) { - auto a = std::make_shared(std::vector{2, 3}, DataType::kFLOAT32, GetDevice(), true); - a->Fill(1.0f); + std::vector values{1.0f, 3.0f, 2.0f, 4.0f, 0.0f, -1.0f}; + auto a = std::make_shared(values.data(), std::vector{2, 3}, DataType::kFLOAT32, GetDevice()); auto max_fn = std::make_shared(1, false); auto result = max_fn->Apply({a}); auto grad = std::make_shared(std::vector{2}, DataType::kFLOAT32, GetDevice(), true); grad->Fill(1.0f); auto grad_inputs = max_fn->Backward({grad}); EXPECT_EQ(grad_inputs.size(), 1); + test::ExpectTensorEqual(grad_inputs[0], {0.0f, 1.0f, 0.0f, 1.0f, 0.0f, 0.0f}); } TEST_P(AutogradReductionBackwardTest, MinBackward) { - auto a = std::make_shared(std::vector{2, 3}, DataType::kFLOAT32, GetDevice(), true); - a->Fill(1.0f); + std::vector values{1.0f, 3.0f, 2.0f, 4.0f, 0.0f, -1.0f}; + auto a = std::make_shared(values.data(), std::vector{2, 3}, DataType::kFLOAT32, GetDevice()); auto min_fn = std::make_shared(1, false); auto result = min_fn->Apply({a}); auto grad = std::make_shared(std::vector{2}, DataType::kFLOAT32, GetDevice(), true); grad->Fill(1.0f); auto grad_inputs = min_fn->Backward({grad}); EXPECT_EQ(grad_inputs.size(), 1); + test::ExpectTensorEqual(grad_inputs[0], {1.0f, 0.0f, 0.0f, 0.0f, 0.0f, 1.0f}); } TEST_P(AutogradReductionBackwardTest, SumBackwardKeepDim) { @@ -65,6 +70,7 @@ TEST_P(AutogradReductionBackwardTest, SumBackwardKeepDim) { grad->Fill(1.0f); auto grad_inputs = sum_fn->Backward({grad}); EXPECT_EQ(grad_inputs.size(), 1); + test::ExpectTensorEqual(grad_inputs[0], {1.0f, 1.0f, 1.0f, 1.0f, 1.0f, 1.0f}); } TEST_P(AutogradReductionBackwardTest, MeanBackwardKeepDim) { @@ -76,6 +82,8 @@ TEST_P(AutogradReductionBackwardTest, MeanBackwardKeepDim) { grad->Fill(1.0f); auto grad_inputs = mean_fn->Backward({grad}); EXPECT_EQ(grad_inputs.size(), 1); + test::ExpectTensorEqual(grad_inputs[0], + {1.0f / 3.0f, 1.0f / 3.0f, 1.0f / 3.0f, 1.0f / 3.0f, 1.0f / 3.0f, 1.0f / 3.0f}); } INFINI_TRAIN_REGISTER_TEST(AutogradReductionBackwardTest); diff --git a/tests/autograd/test_autograd_softmax_backward.cc b/tests/autograd/test_autograd_softmax_backward.cc index fb18f00e..012fa5d8 100644 --- a/tests/autograd/test_autograd_softmax_backward.cc +++ b/tests/autograd/test_autograd_softmax_backward.cc @@ -17,10 +17,13 @@ TEST_P(AutogradSoftmaxBackwardTest, SoftmaxBackward) { a->Fill(1.0f); auto softmax_fn = std::make_shared(1); auto result = softmax_fn->Apply({a}); - auto grad = std::make_shared(std::vector{2, 3}, DataType::kFLOAT32, GetDevice(), true); - grad->Fill(1.0f); + std::vector grad_values{1.0f, 2.0f, 4.0f, 4.0f, 2.0f, 1.0f}; + auto grad + = std::make_shared(grad_values.data(), std::vector{2, 3}, DataType::kFLOAT32, GetDevice()); auto grad_inputs = softmax_fn->Backward({grad}); EXPECT_EQ(grad_inputs.size(), 1); + test::ExpectTensorNear(grad_inputs[0], + {-4.0f / 9.0f, -1.0f / 9.0f, 5.0f / 9.0f, 5.0f / 9.0f, -1.0f / 9.0f, -4.0f / 9.0f}, 1e-6f); } TEST_P(AutogradSoftmaxBackwardTest, SoftmaxBackwardDim0) { @@ -28,10 +31,14 @@ TEST_P(AutogradSoftmaxBackwardTest, SoftmaxBackwardDim0) { a->Fill(1.0f); auto softmax_fn = std::make_shared(0); auto result = softmax_fn->Apply({a}); - auto grad = std::make_shared(std::vector{4, 3}, DataType::kFLOAT32, GetDevice(), true); - grad->Fill(1.0f); + 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, 12.0f}; + auto grad + = std::make_shared(grad_values.data(), std::vector{4, 3}, DataType::kFLOAT32, GetDevice()); auto grad_inputs = softmax_fn->Backward({grad}); EXPECT_EQ(grad_inputs.size(), 1); + test::ExpectTensorNear( + grad_inputs[0], + {-1.125f, -1.125f, -1.125f, -0.375f, -0.375f, -0.375f, 0.375f, 0.375f, 0.375f, 1.125f, 1.125f, 1.125f}, 1e-6f); } INFINI_TRAIN_REGISTER_TEST(AutogradSoftmaxBackwardTest); diff --git a/tests/autograd/test_autograd_transform_backward.cc b/tests/autograd/test_autograd_transform_backward.cc index cdb3b367..16c8b73d 100644 --- a/tests/autograd/test_autograd_transform_backward.cc +++ b/tests/autograd/test_autograd_transform_backward.cc @@ -1,3 +1,4 @@ +#include #include #include "gtest/gtest.h" @@ -23,4 +24,56 @@ TEST_P(AutogradTransformBackwardTest, TransposeBackward) { EXPECT_EQ(grad_inputs.size(), 1); } +TEST_P(AutogradTransformBackwardTest, SplitBackwardValues) { + auto input = std::make_shared(std::vector{2, 5}, DataType::kFLOAT32, GetDevice(), true); + auto split = std::make_shared(2, 1); + auto outputs = split->Apply({input}); + ASSERT_EQ(outputs.size(), 3); + + std::vector> grad_outputs; + for (size_t i = 0; i < outputs.size(); ++i) { + auto grad = std::make_shared(outputs[i]->Dims(), DataType::kFLOAT32, GetDevice()); + grad->Fill(static_cast(i + 1)); + grad_outputs.push_back(grad); + } + + auto grad_inputs = split->Backward(grad_outputs); + ASSERT_EQ(grad_inputs.size(), 1); + test::ExpectTensorEqual(grad_inputs[0], {1.0f, 1.0f, 2.0f, 2.0f, 3.0f, 1.0f, 1.0f, 2.0f, 2.0f, 3.0f}); +} + +TEST_P(AutogradTransformBackwardTest, StackBackwardValues) { + auto a = std::make_shared(std::vector{2, 3}, DataType::kFLOAT32, GetDevice(), true); + auto b = std::make_shared(std::vector{2, 3}, DataType::kFLOAT32, GetDevice(), true); + auto stack = std::make_shared(1); + auto outputs = stack->Apply({a, b}); + ASSERT_EQ(outputs.size(), 1); + + std::vector grad_values(12); + std::iota(grad_values.begin(), grad_values.end(), 0.0f); + auto grad_output + = std::make_shared(grad_values.data(), outputs[0]->Dims(), DataType::kFLOAT32, GetDevice()); + auto grad_inputs = stack->Backward({grad_output}); + ASSERT_EQ(grad_inputs.size(), 2); + test::ExpectTensorEqual(grad_inputs[0], {0.0f, 1.0f, 2.0f, 6.0f, 7.0f, 8.0f}); + test::ExpectTensorEqual(grad_inputs[1], {3.0f, 4.0f, 5.0f, 9.0f, 10.0f, 11.0f}); +} + +TEST_P(AutogradTransformBackwardTest, ConcatBackwardValues) { + auto a = std::make_shared(std::vector{2, 2}, DataType::kFLOAT32, GetDevice(), true); + auto b = std::make_shared(std::vector{2, 1}, DataType::kFLOAT32, GetDevice(), true); + auto concat = std::make_shared(1); + auto outputs = concat->Apply({a, b}); + ASSERT_EQ(outputs.size(), 1); + + std::vector grad_values(6); + std::iota(grad_values.begin(), grad_values.end(), 0.0f); + auto grad_output + = std::make_shared(grad_values.data(), outputs[0]->Dims(), DataType::kFLOAT32, GetDevice()); + auto grad_inputs = concat->Backward({grad_output}); + ASSERT_EQ(grad_inputs.size(), 2); + test::ExpectTensorEqual(grad_inputs[0], {0.0f, 1.0f, 3.0f, 4.0f}); + test::ExpectTensorEqual(grad_inputs[1], {2.0f, 5.0f}); +} + INFINI_TRAIN_REGISTER_TEST(AutogradTransformBackwardTest); diff --git a/tests/autograd/test_autograd_transform_forward.cc b/tests/autograd/test_autograd_transform_forward.cc index 680d8b0d..797cf6e6 100644 --- a/tests/autograd/test_autograd_transform_forward.cc +++ b/tests/autograd/test_autograd_transform_forward.cc @@ -1,3 +1,4 @@ +#include #include #include "gtest/gtest.h" @@ -22,12 +23,14 @@ TEST_P(AutogradTransformForwardTest, TransposeForward) { } TEST_P(AutogradTransformForwardTest, SliceForward) { - auto a = std::make_shared(std::vector{4, 4}, DataType::kFLOAT32, GetDevice(), true); - a->Fill(1.0f); + std::vector values(16); + std::iota(values.begin(), values.end(), 0.0f); + auto a = std::make_shared(values.data(), std::vector{4, 4}, DataType::kFLOAT32, GetDevice()); auto slice_fn = std::make_shared(std::vector{1, 1}, std::vector{3, 3}, std::vector{1, 1}); auto result = slice_fn->Apply({a}); EXPECT_EQ(result.size(), 1); + test::ExpectTensorEqual(result[0], {5.0f, 6.0f, 9.0f, 10.0f}); } TEST_P(AutogradTransformForwardTest, SplitForward) { diff --git a/tests/checkpoint/test_checkpoint_serialization.cc b/tests/checkpoint/test_checkpoint_serialization.cc index e21952f4..f5939c47 100644 --- a/tests/checkpoint/test_checkpoint_serialization.cc +++ b/tests/checkpoint/test_checkpoint_serialization.cc @@ -47,9 +47,7 @@ TEST_P(CheckpointSerializationTest, SaveAndLoadModelFP32) { EXPECT_EQ(loaded.global_step, 42); EXPECT_EQ(loaded.consumed_batches, 100); - auto w1_cpu = model2->parameter("weight")->To(Device()); - const float *data = static_cast(w1_cpu.DataPtr()); - for (int i = 0; i < 6; ++i) { EXPECT_NEAR(data[i], 0.42f, 1e-6); } + test::ExpectTensorNear(model2->parameter("weight"), 0.42f, 1e-6f); std::filesystem::remove_all(dir); } diff --git a/tests/common/test_utils.h b/tests/common/test_utils.h index 06a8bf04..f01bb5cb 100644 --- a/tests/common/test_utils.h +++ b/tests/common/test_utils.h @@ -1,15 +1,68 @@ #pragma once +#include +#include + #if defined(USE_CUDA) #include #endif #include "infini_train/include/device.h" +#include "infini_train/include/tensor.h" #include "gtest/gtest.h" namespace infini_train { namespace test { +inline void ExpectTensorEqual(const std::shared_ptr &tensor, const std::vector &expected) { + ASSERT_NE(tensor, nullptr); + ASSERT_EQ(tensor->Dtype(), DataType::kFLOAT32); + + auto cpu = tensor->To(Device()); + ASSERT_EQ(cpu.NumElements(), expected.size()); + const float *actual = static_cast(cpu.DataPtr()); + for (size_t i = 0; i < expected.size(); ++i) { + EXPECT_FLOAT_EQ(actual[i], expected[i]) << "Mismatch at flat index " << i; + } +} + +inline void ExpectTensorEqual(const std::shared_ptr &tensor, float expected) { + ASSERT_NE(tensor, nullptr); + ASSERT_EQ(tensor->Dtype(), DataType::kFLOAT32); + + auto cpu = tensor->To(Device()); + const float *actual = static_cast(cpu.DataPtr()); + for (size_t i = 0; i < cpu.NumElements(); ++i) { + EXPECT_FLOAT_EQ(actual[i], expected) << "Mismatch at flat index " << i; + } +} + +inline void ExpectTensorNear(const std::shared_ptr &tensor, const std::vector &expected, + float tolerance) { + ASSERT_NE(tensor, nullptr); + ASSERT_EQ(tensor->Dtype(), DataType::kFLOAT32); + ASSERT_GE(tolerance, 0.0f); + + auto cpu = tensor->To(Device()); + ASSERT_EQ(cpu.NumElements(), expected.size()); + const float *actual = static_cast(cpu.DataPtr()); + for (size_t i = 0; i < expected.size(); ++i) { + EXPECT_NEAR(actual[i], expected[i], tolerance) << "Mismatch at flat index " << i; + } +} + +inline void ExpectTensorNear(const std::shared_ptr &tensor, float expected, float tolerance) { + ASSERT_NE(tensor, nullptr); + ASSERT_EQ(tensor->Dtype(), DataType::kFLOAT32); + ASSERT_GE(tolerance, 0.0f); + + auto cpu = tensor->To(Device()); + const float *actual = static_cast(cpu.DataPtr()); + for (size_t i = 0; i < cpu.NumElements(); ++i) { + EXPECT_NEAR(actual[i], expected, tolerance) << "Mismatch at flat index " << i; + } +} + #if defined(USE_CUDA) #define REQUIRE_MIN_DEVICES(n) \ do { \ diff --git a/tests/tensor/cpu_only/test_tensor_copy.cc b/tests/tensor/cpu_only/test_tensor_copy.cc index e8d31e08..fbb811f5 100644 --- a/tests/tensor/cpu_only/test_tensor_copy.cc +++ b/tests/tensor/cpu_only/test_tensor_copy.cc @@ -5,6 +5,8 @@ #include "infini_train/include/tensor.h" +#include "tests/common/test_utils.h" + using namespace infini_train; class TensorCopyCpuTest : public ::testing::Test {}; @@ -16,6 +18,5 @@ TEST_F(TensorCopyCpuTest, CopiesCPUToCPU) { = std::make_shared(std::vector{2, 3}, DataType::kFLOAT32, Device(Device::DeviceType::kCPU, 0)); source->Fill(1.0f); target->CopyFrom(source); - auto *target_data = static_cast(target->DataPtr()); - for (int i = 0; i < 6; ++i) { EXPECT_FLOAT_EQ(target_data[i], 1.0f); } + test::ExpectTensorEqual(target, 1.0f); } diff --git a/tests/tensor/test_tensor_delete.cc b/tests/tensor/test_tensor_delete.cc index ef2758b9..fa2fd27a 100644 --- a/tests/tensor/test_tensor_delete.cc +++ b/tests/tensor/test_tensor_delete.cc @@ -25,17 +25,10 @@ TEST_P(TensorDeleteTest, MoveTransferKeepsData) { auto tensor = std::make_shared(std::vector{2, 3}, DataType::kFLOAT32, GetDevice()); tensor->Fill(5.0f); - // Verify data via CPU copy for any device. - auto cpu_copy - = std::make_shared(std::vector{2, 3}, DataType::kFLOAT32, Device(Device::DeviceType::kCPU, 0)); - cpu_copy->CopyFrom(tensor); - auto moved = std::move(tensor); EXPECT_EQ(tensor, nullptr); ASSERT_NE(moved, nullptr); - - auto *data = static_cast(cpu_copy->DataPtr()); - for (int i = 0; i < 6; ++i) { EXPECT_FLOAT_EQ(data[i], 5.0f); } + test::ExpectTensorEqual(moved, 5.0f); } TEST_P(TensorDeleteTest, NullifiesPointerOnMove) {