[Eager] Remove retain_grad_flag in accumulation_nade, add is_new_grad args in...

[Eager] Remove retain_grad_flag in accumulation_nade, add is_new_grad args in operator (#42240) (#42290)

paddle/fluid/eager/accumulation/accumulation_node.cc

@@ -24,7 +24,7 @@
 #include "paddle/fluid/platform/errors.h"
 
 #include "glog/logging.h"
-DECLARE_bool(retain_grad_for_all_tensor);
+
 namespace egr {
 
 static void CopyOrAddTensor(paddle::experimental::Tensor* tensor,
@@ -41,7 +41,7 @@ static void CopyOrAddTensor(paddle::experimental::Tensor* tensor,
 std::vector<std::vector<paddle::experimental::Tensor>> GradNodeAccumulation::
 operator()(
     std::vector<std::vector<paddle::experimental::Tensor>>& grads,  // NOLINT
-    bool create_graph) {
+    bool create_graph, bool is_new_grad) {
   VLOG(3) << "Running Eager Backward Node: GradNodeAccumulation";
   PADDLE_ENFORCE(grads.size() == 1,
                  paddle::platform::errors::Fatal(
@@ -63,7 +63,7 @@ operator()(
     grad_out = grads[0][0];
   }
 
-  if (!weak_grad_.expired() && FLAGS_retain_grad_for_all_tensor) {
+  if (!weak_grad_.expired() && !is_new_grad) {
     auto grad = weak_grad_.lock();
     CopyOrAddTensor(grad.get(), grad_out);
   }

paddle/fluid/eager/accumulation/accumulation_node.h

@@ -39,7 +39,7 @@ class GradNodeAccumulation : public GradNodeBase {
   // Functor: perform backward computations
   virtual std::vector<std::vector<paddle::experimental::Tensor>> operator()(
       std::vector<std::vector<paddle::experimental::Tensor>>& grads,  // NOLINT
-      bool create_graph = false) override;
+      bool create_graph = false, bool is_new_grad = false) override;
 
   void ClearTensorWrappers() override { VLOG(6) << "Do nothing here now"; }
 

paddle/fluid/eager/api/generated/eager_generated/backwards/scale_node.cc

@@ -147,7 +147,7 @@ void GradNodeScale::SetAttributes_scale(float scale) { scale_ = scale; }
 std::vector<std::vector<paddle::experimental::Tensor>> GradNodeScale::
 operator()(
     std::vector<std::vector<paddle::experimental::Tensor>>& grads,  // NOLINT
-    bool create_graph) {
+    bool create_graph, bool is_new_grad) {
   // 1. Check Output Size
   PADDLE_ENFORCE(
       ((grads.size() == 1) && (grads[0].size() == 1)),

paddle/fluid/eager/api/generated/eager_generated/backwards/scale_node.h

@@ -40,7 +40,7 @@ class GradNodeScale : public GradNodeBase {
   // Functor: perform backward computations
   virtual std::vector<std::vector<paddle::experimental::Tensor>> operator()(
       std::vector<std::vector<paddle::experimental::Tensor>>& grads,  // NOLINT
-      bool create_graph = false) override;
+      bool create_graph = false, bool is_new_grad = false) override;
 
   void ClearTensorWrappers() override { VLOG(6) << "Do nothing here now"; }
 

paddle/fluid/eager/auto_code_generator/eager_generator.cc

@@ -2444,7 +2444,7 @@ static std::string GenerateGradNodeCCContents(
       "std::vector<std::vector<paddle::experimental::Tensor>> "
       "GradNode%s::operator()("
       "std::vector<std::vector<paddle::experimental::Tensor>>& grads, bool "
-      "create_graph) {\n"
+      "create_graph, bool is_new_grad) {\n"
       "%s"
       "%s"
       "\n}";
@@ -2490,7 +2490,7 @@ static std::string GenerateGradNodeHeaderContents(
       "  virtual std::vector<std::vector<paddle::experimental::Tensor>> "
       "operator()("
       "std::vector<std::vector<paddle::experimental::Tensor>>& grads, bool "
-      "create_graph = false) "
+      "create_graph = false, bool is_new_grad = false) "
       "override;\n"
       "\n"
       "  void ClearTensorWrappers() override { \n"

paddle/fluid/eager/auto_code_generator/final_state_generator/eager_gen.py

@@ -119,7 +119,7 @@ class {} : public egr::GradNodeBase {{
   ~{}() override = default;
 
   virtual std::vector<std::vector<paddle::experimental::Tensor>> operator()(
-      std::vector<std::vector<paddle::experimental::Tensor>>& grads, bool create_graph = false) override;
+      std::vector<std::vector<paddle::experimental::Tensor>>& grads, bool create_graph = false, bool is_new_grad = false) override;
   std::string name() override {{ return \"{}\"; }}
   
   void ClearTensorWrappers() override {{
@@ -149,7 +149,7 @@ class {} : public egr::GradNodeBase {{
 
 GRAD_FUNCTION_TEMPLATE = \
 """
-std::vector<std::vector<paddle::experimental::Tensor>> {}::operator()(std::vector<std::vector<paddle::experimental::Tensor>>& grads, bool create_graph) {{
+std::vector<std::vector<paddle::experimental::Tensor>> {}::operator()(std::vector<std::vector<paddle::experimental::Tensor>>& grads, bool create_graph, bool is_new_grad) {{
     // Fill Zero For GradIn Tensors
 {}
 

paddle/fluid/eager/backward.cc

@@ -690,7 +690,7 @@ std::vector<paddle::experimental::Tensor> RunBackward(
     VLOG(6) << "Run Backward Kernel with GradTensorHolder.";
     // Run Pre Backward Node and get outputs
     std::vector<std::vector<paddle::experimental::Tensor>> grad_output_tensors =
-        (*node)(node_input_buffer->Buffers(), create_graph);
+        (*node)(node_input_buffer->Buffers(), create_graph, is_general_grad);
 
     // retain_grad or not
     if (!retain_graph) {

paddle/fluid/eager/custom_operator/custom_operator_node.cc

@@ -20,8 +20,9 @@
 
 namespace egr {
 std::vector<std::vector<paddle::experimental::Tensor>> RunCustomOpNode::
-operator()(std::vector<std::vector<paddle::experimental::Tensor>>& grads,
-           bool create_graph) {  // NOLINT
+operator()(
+    std::vector<std::vector<paddle::experimental::Tensor>>& grads,  // NOLINT
+    bool create_graph, bool is_new_grad) {
   paddle::CustomOpKernelContext ctx;
   auto grad_inputs_name = paddle::framework::OpMetaInfoHelper::GetInputs(
       egr::Controller::Instance().GetOpMetaInfoMap().at(op_type_)[1]);

paddle/fluid/eager/custom_operator/custom_operator_node.h

@@ -39,7 +39,7 @@ class RunCustomOpNode : public GradNodeBase {
   virtual std::vector<std::vector<paddle::experimental::Tensor>>
   operator()(                                                         // NOLINT
       std::vector<std::vector<paddle::experimental::Tensor>>& grads,  // NOLINT
-      bool create_graph = false)                                      // NOLINT
+      bool create_graph = false, bool is_new_grad = false)            // NOLINT
       override;
 
   std::string name() {

paddle/fluid/eager/grad_node_info.h

@@ -109,7 +109,7 @@ class GradNodeBase {
    * **/
   virtual std::vector<std::vector<paddle::experimental::Tensor>> operator()(
       std::vector<std::vector<paddle::experimental::Tensor>>& grads,  // NOLINT
-      bool create_graph = false) = 0;
+      bool create_graph = false, bool is_new_grad = false) = 0;
 
   virtual void ClearTensorWrappers() = 0;
 

paddle/fluid/eager/pylayer/py_layer_node.cc

@@ -32,7 +32,7 @@ namespace egr {
 std::vector<std::vector<paddle::experimental::Tensor>> GradNodePyLayer::
 operator()(
     std::vector<std::vector<paddle::experimental::Tensor>>& grads,  // NOLINT
-    bool create_graph) {
+    bool create_graph, bool is_new_grad) {
   VLOG(3) << "Running Eager Backward Node: " << name();
 
   std::vector<std::vector<paddle::experimental::Tensor>> hooked_grads =

paddle/fluid/eager/pylayer/py_layer_node.h

@@ -36,7 +36,7 @@ class GradNodePyLayer : public GradNodeBase {
 
   virtual std::vector<std::vector<paddle::experimental::Tensor>> operator()(
       std::vector<std::vector<paddle::experimental::Tensor>>& grads,  // NOLINT
-      bool create_graph = false) override;
+      bool create_graph = false, bool is_new_grad = false) override;
 
   void ClearTensorWrappers() override { VLOG(6) << "Do nothing here now"; }
 

paddle/fluid/eager/tests/data_structure_tests/grad_node_test.h

@@ -33,7 +33,7 @@ class GradTestNode : public egr::GradNodeBase {
   std::string name() override { return "GradTestNode"; }
   std::vector<std::vector<paddle::experimental::Tensor>> operator()(
       std::vector<std::vector<paddle::experimental::Tensor>>& grads,  // NOLINT
-      bool create_graph = false) override {
+      bool create_graph = false, bool is_new_grad = false) override {
     val_ = std::dynamic_pointer_cast<phi::DenseTensor>(grads[0][0].impl())
                ->data<float>()[0];
     phi::DenseTensorMeta meta =

paddle/fluid/eager/to_static/run_program_op_node.h

@@ -366,7 +366,7 @@ class GradNodeRunProgram : public egr::GradNodeBase {
   // Functor: perform backward computations
   virtual std::vector<std::vector<paddle::experimental::Tensor>> operator()(
       std::vector<std::vector<paddle::experimental::Tensor>> &grads,  // NOLINT
-      bool create_graph) override {
+      bool create_graph, bool is_new_grad) override {
     VLOG(3) << "Running Eager Backward Node: GradNodeRunProgram";
     std::vector<std::vector<paddle::experimental::Tensor>> hooked_grads =
         GradNodeRunProgram::ApplyGradientHooks(grads);

python/paddle/fluid/tests/unittests/test_tensor_register_hook.py

@@ -462,11 +462,9 @@ class TestTensorRegisterHook(unittest.TestCase):
         x.register_hook(double_print_hook)
 
         y = x * x
-        fluid.set_flags({'FLAGS_retain_grad_for_all_tensor': False})
         # Since y = x * x, dx = 2 * x
         dx = paddle.grad(
             outputs=[y], inputs=[x], create_graph=True, retain_graph=True)[0]
-        fluid.set_flags({'FLAGS_retain_grad_for_all_tensor': True})
 
         z = y + dx
         self.assertTrue(x.grad is None)