[Dy2St] Clean unused inputs and outputs for backward (#66278)

SigureMo · web-flow · commit f624a938b14a · 2024-07-20T00:13:17.000+08:00
diff --git a/paddle/fluid/eager/to_static/run_program_op_func.h b/paddle/fluid/eager/to_static/run_program_op_func.h
@@ -80,6 +80,29 @@ static void clear_unused_out_var_in_backward(
   delete garbages;
 }
 
+static void pir_clear_unused_out_var_in_backward(
+    const std::vector<pir::Value>& fo,
+    const pir::Block* forward_block,
+    const pir::Block* backward_block,
+    paddle::framework::Scope* scope) {
+  auto out_names = details::GetNameFromValue(forward_block, fo, false, true);
+  std::deque<std::shared_ptr<paddle::memory::Allocation>>* garbages =
+      new std::deque<std::shared_ptr<paddle::memory::Allocation>>();
+  for (auto out_name : out_names) {
+    if (!backward_block->kwargs().count(out_name)) {
+      auto var = scope->FindVar(out_name);
+      if (var == nullptr) {
+        continue;
+      }
+      if (var->IsType<phi::DenseTensor>()) {
+        garbages->emplace_back(
+            var->GetMutable<phi::DenseTensor>()->MoveMemoryHolder());
+      }
+    }
+  }
+  delete garbages;
+}
+
 static std::vector<paddle::Tensor> filter_unused_input_var_in_backward(
     const std::vector<paddle::Tensor>& x,
     const std::vector<std::string>& x_names,
@@ -255,24 +278,17 @@ inline void pir_run_program_ad_func(
   if (attrs.count("is_test")) {
     is_test = PADDLE_GET_CONST(bool, attrs.at("is_test"));
   }
-  std::shared_ptr<PirGradNodeRunProgram> grad_node;
   VLOG(2) << "start run run_program with require_any_grad = "
           << require_any_grad << ", is_test = " << is_test;
-
-  if (!is_test && require_any_grad) {
-    // Create GradOpNode (1 means [out_grad], 2 means [x_grad, paramx_grad])
-    grad_node = std::make_shared<PirGradNodeRunProgram>(1, 2);
-  }
-
+  auto x_tmp = Trans2ContiguousTensors(x);
+  auto params_tmp = Trans2ContiguousTensors(params);
   // Call forward function
   // if require_any_grad is False, don't save any middle vars.
   int64_t place_hash_key = 0x9e3779b9;
   for (const paddle::Tensor& tensor : x) {
     int64_t device_type = static_cast<int64_t>(tensor.place().GetType());
     place_hash_key = hash_with_seed(place_hash_key, device_type);
   }
-  auto x_tmp = Trans2ContiguousTensors(x);
-  auto params_tmp = Trans2ContiguousTensors(params);
   PirRunProgramAPI(x_tmp,
                    params_tmp,
                    out,
@@ -281,6 +297,9 @@ inline void pir_run_program_ad_func(
                    attrs,
                    place_hash_key);
   if (!is_test && require_any_grad) {
+    // Create GradOpNode (1 means [out_grad], 2 means [x_grad, paramx_grad])
+    auto grad_node = std::make_shared<PirGradNodeRunProgram>(1, 2);
+
     // Set place hash keys for backward
     grad_node->SetPlaceHashKey(place_hash_key);
 
@@ -292,6 +311,17 @@ inline void pir_run_program_ad_func(
     // Set TensorWrappers
     grad_node->SetFwdX(filter_x);
 
+    std::shared_ptr<::pir::Program> backward_program = PADDLE_GET_CONST(
+        std::shared_ptr<::pir::Program>, attrs.at("backward_program"));
+    std::shared_ptr<::pir::Program> forward_program = PADDLE_GET_CONST(
+        std::shared_ptr<::pir::Program>, attrs.at("forward_program"));
+    auto forward_outputs =
+        PADDLE_GET_CONST(std::vector<::pir::Value>, attrs.at("fo"));
+    pir_clear_unused_out_var_in_backward(forward_outputs,
+                                         forward_program->block(),
+                                         backward_program->block(),
+                                         step_scope[0]);
+
     grad_node->SetFwdParams(params_tmp);
 
     grad_node->SetStepScope(step_scope);  // just for set useable.
diff --git a/paddle/fluid/eager/to_static/run_program_op_node.h b/paddle/fluid/eager/to_static/run_program_op_node.h
@@ -524,14 +524,14 @@ inline void PirRunProgramAPI(
         program_id,
         global_inner_scope,
         place_hash_key);
-    // Step 3. get all eager gc vars
-    // std::set<std::string> skip_eager_delete_vars =
-    // paddle::framework::details::ParseSafeEagerDeletionSkipVarsSet(
-    // *backward_program);
-
+    // Step 3. get all eager gc vars (skip_names = backward_inputs -
+    // no_need_buffers)
+    std::vector<std::string> skip_names;
     // update interpretercore skip_gc_var
-    auto skip_names = details::GetNameFromValue(
-        forward_program->block(), middle_values, false, true);
+    std::vector<pir::Value> kwargs_values;
+    for (auto &kwarg : backward_program->block()->kwargs()) {
+      skip_names.push_back(kwarg.first);
+    }
     auto skip_names_set =
         std::set<std::string>(skip_names.begin(), skip_names.end());
     auto no_need_buffer_values = PADDLE_GET_CONST(std::vector<::pir::Value>,
@@ -545,9 +545,7 @@ inline void PirRunProgramAPI(
     skip_names = details::GetNameFromValue(
         forward_program->block(), output_values, false, true);
     skip_names_set.insert(skip_names.begin(), skip_names.end());
-    skip_names = details::GetNameFromValue(
-        forward_program->block(), input_values, true, false);
-    skip_names_set.insert(skip_names.begin(), skip_names.end());
+
     details::print_collection(skip_names_set);
     interpreter_core->SetSkipGcVars(skip_names_set);
 
@@ -757,6 +755,7 @@ inline void RunProgramAPI(
     // all out_vars are skip_eager_var
     skip_eager_delete_vars.insert(output_names.begin(), output_names.end());
     // update interpretercore skip_gc_var
+    details::print_collection(skip_eager_delete_vars);
     interpreter_core->SetSkipGcVars(skip_eager_delete_vars);
 
     std::set<std::string> input_vars;
