perf: reduce redundant CUDA initialization and synchronization

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<const float *>(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<const float *>(loss_cpu.DataPtr())[0];
                 LOG(INFO) << "Rank " << rank.GlobalRank() << ": finish backward";
             }
 

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<const float *>(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<const float *>(loss_cpu.DataPtr())[0];
                 LOG(INFO) << "Rank " << rank.GlobalRank() << ": finish backward";
             }
 

example/mixtral/main.cc

@@ -133,11 +133,13 @@ int main(int argc, char *argv[]) {
             auto y = std::make_shared<Tensor>(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<const float *>(loss_cpu.DataPtr())[0] / grad_accum_steps;
             loss = loss / static_cast<float>(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<const float *>(loss_cpu.DataPtr())[0];
         }
         optimizer->Step();
 

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<float *>(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;
         }

infini_train/src/kernels/cuda/concat.cu

@@ -186,7 +186,7 @@ std::vector<std::shared_ptr<Tensor>> ConcatBackward(const std::shared_ptr<Tensor
     grads.reserve(input_dims_list.size());
     for (const auto &dvec : input_dims_list) {
         auto t = std::make_shared<Tensor>(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);
     }
 

infini_train/src/kernels/cuda/cross_entropy.cu

@@ -204,7 +204,7 @@ std::shared_ptr<Tensor> CrossEntropyBackward(const std::shared_ptr<Tensor> &inpu
                                  DataTypeList<INFINI_ALL_FLOATING_TYPES>>(
         {target->Dtype(), input_casted->Dtype()},
         [=]<typename Ttarget, typename Tinput>() {
-            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<const Tinput *>(grad_output->DataPtr());
             const Ttarget *target_ptr = static_cast<const Ttarget *>(target->DataPtr());
             const Tinput *input_ptr = static_cast<const Tinput *>(input_casted->DataPtr());

infini_train/src/kernels/cuda/layernorm.cu

@@ -89,8 +89,7 @@ LayerNormForward(const std::shared_ptr<Tensor> &input, const std::shared_ptr<Ten
     core::cuda::DispatchCudaFunc<INFINI_ALL_FLOATING_TYPES>(
         dtype,
         [=]<typename T>() {
-            mean->Fill(0.0);
-            rstd->Fill(0.0);
+            // Each token block writes its mean and rstd exactly once; no Fill is needed.
             LayerNormForwardKernel<BLOCK_SIZE><<<num_blocks, threads_per_block, 0, cuda_stream>>>(
                 static_cast<const T *>(input->DataPtr()), static_cast<const T *>(weight->DataPtr()),
                 static_cast<const T *>(bias->DataPtr()), static_cast<float *>(mean->DataPtr()),
@@ -183,7 +182,7 @@ LayerNormBackward(const std::shared_ptr<Tensor> &input, const std::shared_ptr<Te
     core::cuda::DispatchCudaFunc<INFINI_ALL_FLOATING_TYPES>(
         dtype,
         [=]<typename T>() {
-            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<BLOCK_SIZE><<<num_blocks, threads_per_block, 0, cuda_stream>>>(

infini_train/src/kernels/cuda/linear.cu

@@ -65,6 +65,7 @@ std::shared_ptr<Tensor> LinearForward(const std::shared_ptr<Tensor> &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<Tensor> LinearForward(const std::shared_ptr<Tensor> &input, cons
                     static_cast<T *>(output->DataPtr()), static_cast<const T *>(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<Tensor> LinearForward(const std::shared_ptr<Tensor> &input, cons
                               .x = static_cast<const float *>(input->DataPtr()),
                               .y = static_cast<float *>(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<Tensor> LinearForward(const std::shared_ptr<Tensor> &input, cons
                 .C = output->DataPtr(),
                 .ldc = static_cast<int>(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,

infini_train/src/kernels/cuda/reduction.cu

@@ -181,7 +181,7 @@ std::shared_ptr<Tensor> ReduceOpBackward(const std::shared_ptr<Tensor> &grad_out
     core::cuda::DispatchCudaFunc<INFINI_ALL_FLOATING_TYPES>(
         dtype,
         [=]<typename T>() {
-            grad_input->Fill(0.0);
+            // The backward kernel assigns every grad_input element on all reduction branches; no Fill is needed.
             GenericReduceBackwardKernel<<<num_blocks, threads_per_block, 0, cuda_stream>>>(
                 static_cast<T *>(grad_input->DataPtr()), static_cast<const T *>(grad_output->DataPtr()),
                 input ? static_cast<const T *>(input->DataPtr()) : nullptr,

infini_train/src/kernels/cuda/slice.cu

@@ -48,8 +48,7 @@ std::shared_ptr<Tensor> SliceForward(const std::shared_ptr<Tensor> &input, const
 
     auto dtype = input->Dtype();
     auto new_tensor = std::make_shared<Tensor>(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<int64_t> src_strides(dims.size(), 0), dst_strides(new_dims.size(), 0);
     int64_t stride = 1;

infini_train/src/kernels/cuda/softmax.cu

@@ -198,7 +198,7 @@ std::shared_ptr<Tensor> SoftmaxBackward(const std::shared_ptr<Tensor> &grad_outp
     CHECK(dim >= 0 && dim < output->Dims().size());
 
     auto grad_input = std::make_shared<Tensor>(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);),

infini_train/src/kernels/cuda/split.cu

@@ -114,6 +114,7 @@ std::shared_ptr<Tensor> LaunchSplitBackward(const std::vector<int64_t> &input_di
     const auto &grad = grad_outputs[0];
     auto dtype = grad->Dtype();
     auto grad_input = std::make_shared<Tensor>(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<int64_t>());

infini_train/src/kernels/cuda/stack.cu

@@ -113,7 +113,7 @@ std::vector<std::shared_ptr<Tensor>> StackBackward(const std::vector<int64_t> &i
     std::vector<std::shared_ptr<Tensor>> grads;
     for (int i = 0; i < num_inputs; ++i) {
         auto t = std::make_shared<Tensor>(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);
     }
 

infini_train/src/nn/parallel/pp/pipeline_schedule.cc

@@ -255,9 +255,10 @@ float PipelineSchedule::StepMicroBatches(const std::vector<std::shared_ptr<Tenso
                         {activations[task.local_chunk_idx][mb][0], std::make_shared<Tensor>(target_on_device)})[0];
                     loss = loss / n;
                 }
-                total_loss += static_cast<const float *>(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<const float *>(loss->To(Device()).DataPtr())[0];
             } else {
                 auto out_tensor = activations[task.local_chunk_idx][mb][0];
 

tests/autograd/test_autograd_linear_forward.cc

@@ -18,11 +18,12 @@ TEST_P(AutogradLinearForwardTest, LinearForward) {
     auto weight = std::make_shared<Tensor>(std::vector<int64_t>{4, 3}, DataType::kFLOAT32, GetDevice(), true);
     weight->Fill(1.0f);
     auto bias = std::make_shared<Tensor>(std::vector<int64_t>{4}, DataType::kFLOAT32, GetDevice(), true);
-    bias->Fill(0.0f);
+    bias->Fill(2.0f);
     auto linear_fn = std::make_shared<autograd::Linear>();
     auto result = linear_fn->Apply({input, weight, bias});
     EXPECT_EQ(result.size(), 1);
     EXPECT_EQ(result[0]->Dims(), (std::vector<int64_t>{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<int64_t>{2, 4}));
+    test::ExpectTensorEqual(result[0], 3.0f);
 }
 
 TEST_P(AutogradLinearForwardTest, LinearBatch) {

tests/autograd/test_autograd_loss.cc

@@ -0,0 +1,51 @@
+#include <cmath>
+#include <vector>
+
+#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<float> logits_values{1.0f, 2.0f, 3.0f, 2.0f, 1.0f, 0.0f};
+    auto logits
+        = std::make_shared<Tensor>(logits_values.data(), std::vector<int64_t>{2, 3}, DataType::kFLOAT32, GetDevice());
+    auto target = std::make_shared<Tensor>(std::vector<int64_t>{2}, DataType::kINT64, GetDevice());
+    target->Fill(0);
+
+    auto cross_entropy = std::make_shared<autograd::CrossEntropy>();
+    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<Tensor>(std::vector<int64_t>{}, 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<float> 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);

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) {

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<Tensor>(std::vector<int64_t>{2, 3}, DataType::kFLOAT32, GetDevice(), true);
-    a->Fill(1.0f);
+    std::vector<float> values{1.0f, 3.0f, 2.0f, 4.0f, 0.0f, -1.0f};
+    auto a = std::make_shared<Tensor>(values.data(), std::vector<int64_t>{2, 3}, DataType::kFLOAT32, GetDevice());
     auto max_fn = std::make_shared<autograd::Max>(1, false);
     auto result = max_fn->Apply({a});
     auto grad = std::make_shared<Tensor>(std::vector<int64_t>{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<Tensor>(std::vector<int64_t>{2, 3}, DataType::kFLOAT32, GetDevice(), true);
-    a->Fill(1.0f);
+    std::vector<float> values{1.0f, 3.0f, 2.0f, 4.0f, 0.0f, -1.0f};
+    auto a = std::make_shared<Tensor>(values.data(), std::vector<int64_t>{2, 3}, DataType::kFLOAT32, GetDevice());
     auto min_fn = std::make_shared<autograd::Min>(1, false);
     auto result = min_fn->Apply({a});
     auto grad = std::make_shared<Tensor>(std::vector<int64_t>{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);

tests/autograd/test_autograd_softmax_backward.cc

@@ -17,21 +17,28 @@ TEST_P(AutogradSoftmaxBackwardTest, SoftmaxBackward) {
     a->Fill(1.0f);
     auto softmax_fn = std::make_shared<autograd::Softmax>(1);
     auto result = softmax_fn->Apply({a});
-    auto grad = std::make_shared<Tensor>(std::vector<int64_t>{2, 3}, DataType::kFLOAT32, GetDevice(), true);
-    grad->Fill(1.0f);
+    std::vector<float> grad_values{1.0f, 2.0f, 4.0f, 4.0f, 2.0f, 1.0f};
+    auto grad
+        = std::make_shared<Tensor>(grad_values.data(), std::vector<int64_t>{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) {
     auto a = std::make_shared<Tensor>(std::vector<int64_t>{4, 3}, DataType::kFLOAT32, GetDevice(), true);
     a->Fill(1.0f);
     auto softmax_fn = std::make_shared<autograd::Softmax>(0);
     auto result = softmax_fn->Apply({a});
-    auto grad = std::make_shared<Tensor>(std::vector<int64_t>{4, 3}, DataType::kFLOAT32, GetDevice(), true);
-    grad->Fill(1.0f);
+    std::vector<float> 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<Tensor>(grad_values.data(), std::vector<int64_t>{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);