根据反馈修复代码问题；新增作业提示

Author: JYMiracle305

docs/训练营作业介绍.md

@@ -1,3 +1,9 @@
+# 特别提示
+
+1. 如果大家在作业项目开发过程中发现项目本身的问题，不要修改用例代码，请及时反馈给助教。
+
+2. 建议大家尽量独立完成作业，完成用例的同时做到"最小化修改"，原则上在作业要求注释的位置开发即可。另外，作业报告内可以写下对代码实现的理解、算法或者系统的示意图以及开发过程中的问题记录，形式不限。
+
 # 本地自测
 
 1. 拉取代码仓库 ``git clone git@github.com:InfiniTensor/TinyInfiniTrain.git --recursive``

infini_train/src/autograd/function.cc

@@ -58,7 +58,7 @@ std::vector<std::shared_ptr<Tensor>> Function::Apply(const std::vector<std::shar
         auto &output_tensor = output_tensors[output_idx];
         output_tensor->set_requires_grad(output_requires_grad);
         output_tensor->set_is_leaf(false);
-        output_tensor->set_grad_fn(shared_from_this());
+        output_tensor->set_grad_fn(output_requires_grad ? shared_from_this() : nullptr);
         output_tensor->set_output_idx(output_idx);
     }
 
@@ -77,8 +77,8 @@ void Function::BackwardPartial(const std::shared_ptr<Tensor> &grad_output, int g
     if (grad_outputs_reached_ == grad_outputs_.size()
         && (dependencies_reached_ == dependencies_number_ || dependencies_number_ == 0)) {
         auto grad_inputs = Backward(grad_outputs_);
-        // saved_tensors_.clear();
-        // grad_outputs_.clear();
+        saved_tensors_.clear();
+        grad_outputs_.clear();
         CHECK_EQ(grad_inputs.size(), next_functions_.size());
         for (int idx = 0; idx < grad_inputs.size(); ++idx) {
             auto &grad_input = grad_inputs[idx];

test/optimizer/test_adam.cc

@@ -58,59 +58,3 @@ TEST(AdamOptimizerTest, MomentumAccumulation) {
         EXPECT_NEAR(param_history[t-1] - 1.0f, expected_update, 1e-5);
     }
 }
-
-#ifdef USE_CUDA
-
-TEST(AdamOptimizerTest, BasicParameterUpdateCuda) {
-    auto param = std::make_shared<Tensor>(std::vector<int64_t>{3}, DataType::kFLOAT32,
-        Device(DeviceType::kCUDA, 0));
-    param->Fill(1.0f); // 初始参数值 [1.0, 1.0, 1.0]
-    param->RequiresGrad();
-    
-    auto grad = std::make_shared<Tensor>(param->Dims(), param->Dtype());
-    grad->Fill(1.0f);
-    float* grad_data = static_cast<float*>(param->grad()->DataPtr());
-    std::memcpy(grad_data, grad->DataPtr(), grad->SizeInBytes());
-
-    optimizers::Adam optimizer({param}, 0.001f, 0.9f, 0.999f, 1e-8);
-
-    optimizer.Step();
-
-    float* param_data = static_cast<float*>(param->DataPtr());
-    for (int i = 0; i < 3; ++i) {
-        EXPECT_LT(param_data[i], 1.0f); // 参数值应该减小
-    }
-}
-
-TEST(AdamOptimizerTest, MomentumAccumulationCuda) {
-    auto param = std::make_shared<Tensor>(std::vector<int64_t>{1}, DataType::kFLOAT32,
-        Device(DeviceType::kCUDA, 0));
-    param->Fill(1.0f);
-    param->RequiresGrad();
-    param->grad()->Fill(0.5f);
-
-    float learning_rate = 1e-3, beta1 = 0.9, beta2 = 0.999, eps = 1e-8;
-
-    optimizers::Adam optimizer({param}, learning_rate, beta1, beta2, eps);
-    
-    std::vector<float> param_history;
-    for (int i = 0; i < 3; ++i) {
-        optimizer.Step();
-        param_history.push_back(static_cast<float*>(param->DataPtr())[0]);
-    }
-
-    EXPECT_LT(param_history[1], param_history[0]);
-    EXPECT_LT(param_history[2], param_history[1]);
-
-    float m = 0, v = 0, expected_update = 0;
-    for (int t = 1; t <= 3; ++t) {
-        m = beta1 * m + (1 - beta1) * 0.5f;       // 一阶动量
-        v = beta2 * v + (1 - beta2) * 0.25f;      // 二阶动量
-        float m_hat = m / (1.0f - std::pow(beta1, t));  // 动态校正因子
-        float v_hat = v / (1.0f - std::pow(beta2, t));
-
-        expected_update -= learning_rate * m_hat / (std::sqrt(v_hat) + 1e-8f);
-        EXPECT_NEAR(param_history[t-1] - 1.0f, expected_update, 1e-5);
-    }
-}
-#endif

test/optimizer/test_adam_cuda.cc

@@ -2,6 +2,8 @@
 #include<iostream>
 #include<vector>
 
+#include "cuda_runtime_api.h"
+
 #include "gtest/gtest.h"
 #include "infini_train/include/tensor.h"
 #include "infini_train/include/device.h"
@@ -18,13 +20,14 @@ TEST(AdamOptimizerTest, BasicParameterUpdateCuda) {
     auto grad = std::make_shared<Tensor>(param->Dims(), param->Dtype());
     grad->Fill(1.0f);
     float* grad_data = static_cast<float*>(param->grad()->DataPtr());
-    std::memcpy(grad_data, grad->DataPtr(), grad->SizeInBytes());
+    cudaMemcpy(grad_data, grad->DataPtr(), grad->SizeInBytes(), cudaMemcpyDefault);
 
     optimizers::Adam optimizer({param}, 0.001f, 0.9f, 0.999f, 1e-8);
 
     optimizer.Step();
 
-    float* param_data = static_cast<float*>(param->DataPtr());
+    auto param_cpu = param->To(Device(DeviceType::kCPU, 0));
+    float* param_data = static_cast<float*>(param_cpu.DataPtr());
     for (int i = 0; i < 3; ++i) {
         EXPECT_LT(param_data[i], 1.0f); // 参数值应该减小
     }
@@ -44,7 +47,8 @@ TEST(AdamOptimizerTest, MomentumAccumulationCuda) {
     std::vector<float> param_history;
     for (int i = 0; i < 3; ++i) {
         optimizer.Step();
-        param_history.push_back(static_cast<float*>(param->DataPtr())[0]);
+        auto param_cpu = param->To(Device(DeviceType::kCPU, 0));
+        param_history.push_back(static_cast<float*>(param_cpu.DataPtr())[0]);
     }
 
     EXPECT_LT(param_history[1], param_history[0]);