[Dy2St] Fix unused grad value error when split forward backward (#66557)

0x45f · SigureMo · web-flow · commit b9d956356ca5 · 2024-08-08T10:28:06.000+08:00
---------

Co-authored-by: SigureMo &lt;sigure.qaq@gmail.com&gt;
diff --git a/paddle/fluid/framework/new_executor/instruction/control_flow/if_instruction.cc b/paddle/fluid/framework/new_executor/instruction/control_flow/if_instruction.cc
@@ -118,7 +118,7 @@ IfInstruction::IfInstruction(size_t id,
       outputs.emplace(value, GetValueIds(value, *value_exec_info));
     }
     if (value.use_count() > 0) {
-      VLOG(0) << "value " << i << " use conutn != 0";
+      VLOG(6) << "value " << i << " use conutn != 0";
       is_last_op = false;
     }
   }
diff --git a/paddle/fluid/pybind/pir.cc b/paddle/fluid/pybind/pir.cc
@@ -1651,36 +1651,40 @@ std::map<int, int> GetOpInplaceInfo(const pir::Operation *op) {
 std::pair<std::vector<pir::Value>, std::unordered_set<pir::Value>>
 AnalysisMiddleVariable(const Program &program,
                        const std::vector<pir::Value> &forward_inputs,
+                       const std::vector<pir::Value> &backward_outputs,
                        const std::vector<int> &forward_range,
                        const std::vector<int> &backward_range) {
   std::vector<pir::Value> middle_values;
 
-  std::unordered_set<pir::Value> backward_inputs;
+  std::unordered_set<pir::Value> backward_used_values;
   std::unordered_set<pir::Value> x_or_param(forward_inputs.begin(),
                                             forward_inputs.end());
+  for (const auto &value : backward_outputs) {
+    backward_used_values.insert(value);
+  }
   range_block_do(
-      program.block(), backward_range, [&backward_inputs](Operation *op) {
+      program.block(), backward_range, [&backward_used_values](Operation *op) {
         pir::Walk(op, [&](Operation *inner_op) {
           for (auto &t : inner_op->operands()) {
-            backward_inputs.insert(t.source());
+            backward_used_values.insert(t.source());
           }
         });
       });
 
   range_block_do(
       program.block(),
       forward_range,
-      [&middle_values, &backward_inputs, &x_or_param](Operation *op) {
+      [&middle_values, &backward_used_values, &x_or_param](Operation *op) {
         pir::Walk(op, [&](Operation *inner_op) {
           for (auto &t : inner_op->results()) {
             auto v = Value(t.Value::impl());
-            if (backward_inputs.count(v) && !x_or_param.count(v)) {
+            if (backward_used_values.count(v) && !x_or_param.count(v)) {
               middle_values.push_back(v);
             }
           }
         });
       });
-  return std::make_pair(middle_values, backward_inputs);
+  return std::make_pair(middle_values, backward_used_values);
 }
 
 void mapping_value(const std::vector<pir::Value> &origin,
@@ -1828,16 +1832,24 @@ SplitedResult SplitForwardBackward(
     forward_in_out_values.insert(
         forward_in_out_values.end(), v->begin(), v->end());
   }
+  std::vector<pir::Value> backward_out_values;
+  for (auto &v : std::vector({&forward_inputs_grads, &forward_params_grads})) {
+    backward_out_values.insert(backward_out_values.end(), v->begin(), v->end());
+  }
 
   std::vector<pir::Value> fx, fp, fm, fo, bx, bp, bm, bo_g, bx_g, bp_g, bo;
   std::vector<pir::Value> no_need_buffer_values;
   pir::IrContext *ctx = pir::IrContext::Instance();
   auto forward_program = std::make_shared<Program>(ctx);
   auto backward_program = std::make_shared<Program>(ctx);
   std::vector<pir::Value> middle_values;
-  std::unordered_set<pir::Value> backward_inputs;
-  std::tie(middle_values, backward_inputs) = AnalysisMiddleVariable(
-      program, forward_in_out_values, forward_range, backward_range);
+  std::unordered_set<pir::Value> backward_used_values;
+  std::tie(middle_values, backward_used_values) =
+      AnalysisMiddleVariable(program,
+                             forward_in_out_values,
+                             backward_out_values,
+                             forward_range,
+                             backward_range);
 
   pir::Block &backward_block = *backward_program->block();
   bool has_backward = (backward_range[1] > backward_range[0]);
@@ -1897,19 +1909,20 @@ SplitedResult SplitForwardBackward(
       create_output_fn(forward_value_map, forward_program, "output_"));
 
   auto create_kwarg_fn = [&backward_block,
-                          &backward_inputs,
+                          &backward_used_values,
                           &backward_value_map,
                           &forward_value_map](const std::string &prefix) {
     auto counter = std::make_shared<size_t>(0);
     return [&backward_block,
-            &backward_inputs,
+            &backward_used_values,
             &backward_value_map,
             &forward_value_map,
             &prefix,
             counter](const pir::Value &v) {
       // NOTE(SigureMo): Ensure counter++ executed in each iteration.
       auto default_name = prefix + std::to_string((*counter)++);
-      if (v && !backward_value_map.count(v) && (backward_inputs.count(v))) {
+      if (v && !backward_value_map.count(v) &&
+          (backward_used_values.count(v))) {
         backward_value_map[v] = backward_block.AddKwarg(
             name_analysis::TryGetValueFirstName(forward_value_map[v])
                 .value_or(default_name),
diff --git a/test/dygraph_to_static/test_grad.py b/test/dygraph_to_static/test_grad.py
@@ -20,6 +20,7 @@
 from dygraph_to_static_utils import (
     Dy2StTestBase,
     test_legacy_and_pt_and_pir,
+    test_pir_only,
 )
 
 import paddle
@@ -154,5 +155,32 @@ def tearDown(self):
         self.temp_dir.cleanup()
 
 
+class UnuseGradVarLayer(paddle.nn.Layer):
+    def __init__(self):
+        super().__init__()
+
+    def forward(self, var_0, var_1):
+        var_1 = var_1 + 1
+        return var_0, var_1
+
+
+class TestUnuseGradVar(Dy2StTestBase):
+    @test_pir_only
+    def test_run(self):
+        layer = UnuseGradVarLayer()
+        layer = paddle.jit.to_static(layer)
+
+        x = paddle.to_tensor([1.0])
+        y = paddle.to_tensor([2.0])
+        x.stop_gradient = False
+        y.stop_gradient = False
+
+        out1, out2 = layer(x, y)
+        out = out1 + out2
+        out.backward()
+        np.testing.assert_array_equal(out.numpy(), [4])
+        np.testing.assert_array_equal(x.grad.numpy(), [1])
+
+
 if __name__ == '__main__':
     unittest.main()

Original file line number	Diff line number	Diff line change
`@@ -118,7 +118,7 @@ IfInstruction::IfInstruction(size_t id,`
`118`	`118`	`outputs.emplace(value, GetValueIds(value, *value_exec_info));`
`119`	`119`	`}`
`120`	`120`	`if (value.use_count() > 0) {`
`121`		`- VLOG(0) << "value " << i << " use conutn != 0";`
	`121`	`+ VLOG(6) << "value " << i << " use conutn != 0";`
`122`	`122`	`is_last_op = false;`
`123`	`123`	`}`
`124`	`124`	`}`