[Dy2St]Modify jit.load into Lazy Initialization Mode for backward program (#57240)

Aurelius84 · web-flow · commit 4920462600b2 · 2023-09-20T19:48:41.000+08:00
* [Dy2St]Modify jit.load into Lazy Initialization Mode for backward program

* fix is_test

* fix typo

* fix logic

* fix build scope logic
diff --git a/paddle/fluid/eager/to_static/run_program_op_func.h b/paddle/fluid/eager/to_static/run_program_op_func.h
@@ -140,8 +140,11 @@ inline void run_program_ad_func(
   RunProgramAPI(
       x_tmp, params_tmp, out, step_scope, dout, require_any_grad, attrs);
   VLOG(2) << "start run run_program grad";
-
-  if (require_any_grad) {
+  auto is_test = false;
+  if (attrs.count("is_test")) {
+    is_test = PADDLE_GET_CONST(bool, attrs.at("is_test"));
+  }
+  if (!is_test && require_any_grad) {
     auto x_names =
         PADDLE_GET_CONST(std::vector<std::string>, attrs.at("x_names"));
 
diff --git a/paddle/fluid/eager/to_static/run_program_op_node.h b/paddle/fluid/eager/to_static/run_program_op_node.h
@@ -315,14 +315,16 @@ static void ShareTensorsFromScopeByValue(
 static void ShareTensorsFromScopeWithPartialBlock(
     const std::vector<Tensor *> &tensors,
     const paddle::framework::BlockDesc &forward_global_block,
-    const paddle::framework::BlockDesc &backward_global_block,
+    const paddle::framework::BlockDesc *backward_global_block,
     paddle::framework::Scope *scope) {
   for (size_t i = 0; i < tensors.size(); ++i) {
     auto &name = tensors[i]->name();
+    bool in_forward_block = forward_global_block.HasVar(name);
+    bool in_backward_block =
+        backward_global_block && backward_global_block->HasVar(name);
     if (name == paddle::framework::kEmptyVarName ||
         name == paddle::framework::kFakeVarName ||
-        (!forward_global_block.HasVar(name) &&
-         !backward_global_block.HasVar(name))) {
+        (!in_forward_block && !in_backward_block)) {
       VLOG(2) << "find tensor name is " << name << ", skip it!";
       continue;
     }
@@ -660,10 +662,16 @@ inline void RunProgramAPI(
 
   auto *forward_global_block = PADDLE_GET_CONST(
       paddle::framework::BlockDesc *, attrs.at("forward_global_block"));
-  auto *backward_global_block = PADDLE_GET_CONST(
-      paddle::framework::BlockDesc *, attrs.at("backward_global_block"));
   auto *forward_program = forward_global_block->Program();
-  auto *backward_program = backward_global_block->Program();
+
+  paddle::framework::BlockDesc *backward_global_block = nullptr;
+  paddle::framework::ProgramDesc *backward_program = nullptr;
+
+  if (!is_test) {
+    backward_global_block = PADDLE_GET_CONST(paddle::framework::BlockDesc *,
+                                             attrs.at("backward_global_block"));
+    backward_program = backward_global_block->Program();
+  }
 
   auto &interpretercore_info_cache =
       paddle::framework::InterpreterCoreInfoCache::Instance();
@@ -710,9 +718,12 @@ inline void RunProgramAPI(
               global_inner_scope);
     }
     // Step 3. get all eager gc vars
-    std::set<std::string> skip_eager_delete_vars =
-        paddle::framework::details::ParseSafeEagerDeletionSkipVarsSet(
-            *backward_program);
+    std::set<std::string> skip_eager_delete_vars;
+    if (!is_test) {
+      skip_eager_delete_vars =
+          paddle::framework::details::ParseSafeEagerDeletionSkipVarsSet(
+              *backward_program);
+    }
 
     // all out_vars are skip_eager_var
     skip_eager_delete_vars.insert(output_names.begin(), output_names.end());
@@ -765,19 +776,15 @@ inline void RunProgramAPI(
         1);
     interpreter_core->Run({});
   }
-
+  VLOG(3) << paddle::framework::GenScopeTreeDebugInfo(out_scope_vec->front());
   {
     paddle::platform::RecordEvent record_event(
         "fetch_and_gc", paddle::platform::TracerEventType::UserDefined, 1);
     // Get Output
     details::ShareTensorsFromScopeWithPartialBlock(
-        out, *forward_global_block, *backward_global_block, global_inner_scope);
-    details::ShareTensorsFromScopeWithPartialBlock(dout,
-                                                   *forward_global_block,
-                                                   *backward_global_block,
-                                                   global_inner_scope);
-
-    VLOG(3) << paddle::framework::GenScopeTreeDebugInfo(out_scope_vec->front());
+        out, *forward_global_block, backward_global_block, global_inner_scope);
+    details::ShareTensorsFromScopeWithPartialBlock(
+        dout, *forward_global_block, backward_global_block, global_inner_scope);
 
     if (is_test || !require_any_grad) {
       VLOG(4) << "don't require any grad, set this scope can reused";
@@ -939,11 +946,11 @@ inline void RunProgramGradAPI(
     // Step 4. get outputs
     details::ShareTensorsFromScopeWithPartialBlock(x_grad,
                                                    *forward_global_block,
-                                                   *backward_global_block,
+                                                   backward_global_block,
                                                    global_inner_scope);
     details::ShareTensorsFromScopeWithPartialBlock(params_grad,
                                                    *forward_global_block,
-                                                   *backward_global_block,
+                                                   backward_global_block,
                                                    global_inner_scope);
     VLOG(4) << "after backward gc all vars";
     global_inner_scope->SetCanReused(true);
diff --git a/paddle/fluid/framework/executor_cache.cc b/paddle/fluid/framework/executor_cache.cc
@@ -356,7 +356,7 @@ std::shared_ptr<InterpreterCore> CreateNewIRInterpreterCoreInfoToCache(
 std::unique_ptr<::pir::Program> ConstructFowardIrProgram(
     const paddle::framework::BlockDesc *forward_global_block,
     const paddle::framework::BlockDesc *backward_global_block,
-    const std::vector<std::string> output_names,
+    const std::vector<std::string> &output_names,
     const std::vector<paddle::Tensor> &x,
     const std::vector<std::string> &x_names,
     const std::vector<paddle::Tensor> &params,
@@ -415,19 +415,21 @@ std::unique_ptr<::pir::Program> ConstructFowardIrProgram(
   }
 
   std::set<std::string> set_parameter_names;
-  for (auto op_desc : backward_global_block->Program()->Block(0).AllOps()) {
-    for (const auto &n : op_desc->Inputs()) {
-      const auto &input_var_names = n.second;
-      for (const auto &var_name : input_var_names) {
-        set_parameter_names.insert(var_name);
-      }
-    }
-  }
-
   for (auto &t : output_names) {
     set_parameter_names.insert(t);
   }
 
+  if (backward_global_block != nullptr) {
+    for (auto op_desc : backward_global_block->Program()->Block(0).AllOps()) {
+      for (const auto &n : op_desc->Inputs()) {
+        const auto &input_var_names = n.second;
+        for (const auto &var_name : input_var_names) {
+          set_parameter_names.insert(var_name);
+        }
+      }
+    }
+  }
+
   for (auto &name : set_parameter_names) {
     if (!set_output_names.count(name)) {
       continue;
@@ -443,7 +445,6 @@ std::unique_ptr<::pir::Program> ConstructFowardIrProgram(
     op_desc->SetInput("x", {name});
     op_desc->SetOutput("out", {"@EMPTY@"});
   }
-
   paddle::translator::ProgramTranslator program_translator(&local_program,
                                                            program.get());
 
diff --git a/paddle/fluid/framework/executor_cache.h b/paddle/fluid/framework/executor_cache.h
@@ -253,7 +253,7 @@ std::shared_ptr<InterpreterCore> CreateNewIRInterpreterCoreInfoToCache(
 std::unique_ptr<::pir::Program> ConstructFowardIrProgram(
     const paddle::framework::BlockDesc* forward_global_block,
     const paddle::framework::BlockDesc* backward_global_block,
-    const std::vector<std::string> output_names,
+    const std::vector<std::string>& output_names,
     const std::vector<paddle::Tensor>& x,
     const std::vector<std::string>& x_names,
     const std::vector<paddle::Tensor>& params,
diff --git a/python/paddle/jit/translated_layer.py b/python/paddle/jit/translated_layer.py
@@ -347,15 +347,11 @@ def __init__(self, program_desc):
         self._suffix_varname_dict = None
         # forward program
         self._infer_program_desc = self._preprocess(program_desc)
-        # forward + backward program
-        self._train_program_desc = self._append_backward_desc(
-            self._infer_program_desc
-        )
 
     # forward:
     @switch_to_static_graph
     def _create_forward_train_program(self):
-        whole_program = _build_program_by_desc(self._train_program_desc)
+        whole_program = _build_program_by_desc(self.train_program)
         end_op_index = self._infer_program_desc.block(0).op_size()
         if end_op_index > 0:
             return add_build_strategy_for(whole_program, 0, end_op_index)
@@ -369,7 +365,7 @@ def _forward_program_desc(self):
     # backward
     @switch_to_static_graph
     def _create_backward_train_program(self):
-        whole_program = _build_program_by_desc(self._train_program_desc)
+        whole_program = _build_program_by_desc(self.train_program)
         start_op_index = self._infer_program_desc.block(0).op_size() + len(
             self._output_descs
         )
@@ -389,9 +385,9 @@ def _backward_program_desc(self):
     def infer_program(self):
         return self._infer_program_desc
 
-    @property
+    @LazyInitialized
     def train_program(self):
-        return self._train_program_desc
+        return self._append_backward_desc(self._infer_program_desc)
 
     @property
     def forward_program(self):
@@ -1010,10 +1006,15 @@ def _run_dygraph(instance, input, program_holder):
             (
                 'forward_global_block',
                 forward_program.block(0),
-                'backward_global_block',
-                program_holder.backward_program.block(0),
             )
         )
+        if not instance._is_test:
+            attrs.extend(
+                (
+                    'backward_global_block',
+                    program_holder.backward_program.block(0),
+                )
+            )
 
     _legacy_C_ops.run_program(
         _valid_vars(input_vars),
@@ -1055,7 +1056,6 @@ def _run_static_graph(input, program_holder, trace_program):
         trace_program, exclude=param_var_names
     )
     trace_program.flush()
-    output_names = [var.name() for var in program_holder.output_descs]
     # append blocks from 'trace_program'
     _append_block(
         main_program,
