[Dy2St][PIR] Support nested structure in PyLayer forward inputs (#69642)

SigureMo · web-flow · commit 4da5a832dc00 · 2024-12-11T15:21:09.000+08:00
diff --git a/paddle/fluid/pir/dialect/operator/ir/manual_pylayer_op.cc b/paddle/fluid/pir/dialect/operator/ir/manual_pylayer_op.cc
@@ -128,12 +128,14 @@ void PyLayerOp::Print(pir::IrPrinter &printer) {
   printer.PrintAttributeMap(*op);
   os << " -> ";
   printer.PrintOpReturnType(*op);
-  os << "{";
+  os << " {\n";
+  printer.AddIndentation();
   for (auto &item : forward_block()) {
-    os << "\n  ";
     printer.PrintOperation(item);
+    os << "\n";
   }
-  os << "\n }";
+  printer.DecreaseIndentation();
+  os << printer.indentation() << "}";
 }
 
 void PyLayerOp::VerifySig() {
diff --git a/python/paddle/jit/dy2static/py_layer.py b/python/paddle/jit/dy2static/py_layer.py
@@ -21,20 +21,22 @@
 from paddle.base.libpaddle.pir import build_pipe_for_pylayer
 from paddle.common_ops_import import LayerHelper
 from paddle.static.nn import static_pylayer
+from paddle.utils import flatten, pack_sequence_as
 
 from .program_translator import convert_to_static, unwrap_decorators
 
 
 class StaticPyLayerContext:
     def __init__(self):
         self.saved_vars = []
+        self.saved_vars_structure = None
 
         if in_pir_mode():
             self.tuple_push_op_name = "cf.tuple_push"
             self.tuple_pop_op_name = "cf.tuple_pop"
 
     def __setattr__(self, attr: str, value: object):
-        attr_allow_list = ["saved_vars"]
+        attr_allow_list = ["saved_vars", "saved_vars_structure"]
         if (
             in_pir_mode()
             and attr not in attr_allow_list
@@ -68,9 +70,14 @@ def __setattr__(self, attr: str, value: object):
 
     def save_for_backward(self, *tensors):
         if in_pir_mode():
+            self.saved_vars_structure = tensors
+            flatten_tensors = flatten(tensors)
+            tensor_elements = list(
+                filter(lambda x: isinstance(x, pir.Value), flatten_tensors)
+            )
             current_insert_point = pir.get_current_insertion_point()
             current_block = current_insert_point.block()
-            build_pipe_for_pylayer(current_block, tensors)
+            build_pipe_for_pylayer(current_block, tensor_elements)
         else:
             for tensor in tensors:
                 assert isinstance(tensor, Variable)
@@ -84,6 +91,16 @@ def saved_tensor(self):
             for op in current_block.ops:
                 if op.name() == self.tuple_pop_op_name:
                     out_list = op.as_tuple_pop_op().pop_all_values()
+            if self.saved_vars_structure is not None:
+                flattened_structure = flatten(self.saved_vars_structure)
+                value_cursor = 0
+                for i, tensor in enumerate(flattened_structure):
+                    if isinstance(tensor, pir.Value):
+                        flattened_structure[i] = out_list[value_cursor]
+                        value_cursor += 1
+                out_list = pack_sequence_as(
+                    self.saved_vars_structure, flattened_structure
+                )
         else:
             helper = LayerHelper("StaticPyLayerContext")
             out_list = []
diff --git a/python/paddle/static/nn/static_pylayer.py b/python/paddle/static/nn/static_pylayer.py
@@ -269,6 +269,11 @@ def __call__(self, *output_grads):
             input_grads = (input_grads,)
 
         self._hook_check_func(output_grads, input_grads)
+        input_grads = [
+            input_grad
+            for input_grad in flatten(input_grads)
+            if isinstance(input_grad, (paddle.pir.Value, type(None)))
+        ]
 
         return input_grads
 
@@ -369,7 +374,7 @@ def static_pylayer(forward_fn, inputs, backward_fn=None, name=None):
 
     if in_pir_mode():
         fwd_inputs = [
-            inp for inp in inputs if isinstance(inp, paddle.pir.Value)
+            inp for inp in flatten(inputs) if isinstance(inp, paddle.pir.Value)
         ]
         pylayer_op = build_pylayer_op(fwd_inputs)
         outputs = None
@@ -404,9 +409,14 @@ def hook_inputs_outputs_check_function(output_grads, input_grads):
                     for x in flatten(inputs)
                     if isinstance(x, paddle.pir.Value)
                 ]
+                input_grads = [
+                    x
+                    for x in flatten(input_grads)
+                    if isinstance(x, (paddle.pir.Value, type(None)))
+                ]
                 if len(input_grads) != len(forward_inputs):
                     raise ValueError(
-                        f"The number of input grads should be equal to the number of inputs, but got {len(input_grads)} and {len(inputs)}."
+                        f"The number of input grads should be equal to the number of inputs, but got {len(input_grads)} and {len(forward_inputs)}."
                     )
                 for inp_grad, fwd_input in zip(input_grads, forward_inputs):
                     # NOTE: inp_grad will be None if fwd_input.stop_gradients=True
diff --git a/test/dygraph_to_static/test_pylayer.py b/test/dygraph_to_static/test_pylayer.py
@@ -812,5 +812,163 @@ def test_wrong_usage(self):
             static_layer(x)
 
 
+class NestedStructurePyLayer(PyLayer):
+    @staticmethod
+    def forward(ctx, x, y):
+        ctx.save_for_backward(x, y)
+        x1 = paddle.tanh(x[0])
+        y1 = paddle.tanh(x[1])
+        z1 = paddle.tanh(y)
+        return [x1, y1, z1]
+
+    @staticmethod
+    def backward(ctx, *grad1):
+        x0, x1 = ctx.saved_tensor()
+        x_grad = grad1[0] * (1 - paddle.square(x0[0]))
+        y_grad = grad1[1] * (1 - paddle.square(x0[1]))
+        z_grad = grad1[2] * (1 - paddle.square(x1))
+
+        return [x_grad, y_grad], z_grad
+
+
+class NestedStructurePyLayerModel(paddle.nn.Layer):
+    def __init__(self):
+        super().__init__()
+        self.w0 = self.create_parameter(shape=[42, 42])
+        self.w1 = self.create_parameter(shape=[42, 42])
+        self.w2 = self.create_parameter(shape=[42, 42])
+
+    def forward(self, x):
+        y1 = paddle.matmul(x, self.w0)
+        y2 = paddle.matmul(x, self.w1)
+        y3 = paddle.matmul(x, self.w2)
+
+        z = NestedStructurePyLayer.apply([y1, y2], y3)
+        return z[0] + z[1] + z[2]
+
+
+class TestNestedStructurePyLayer(unittest.TestCase):
+    def test_nested_structure(self):
+        input = paddle.randn([2, 42]).astype("float32")
+        input.stop_gradient = False
+
+        model = NestedStructurePyLayerModel()
+        dygraph_res = model(input)
+        dygraph_res.backward()
+        dygraph_input_grads = [
+            paddle.assign(input.grad),
+            paddle.assign(model.w0.grad),
+            paddle.assign(model.w1.grad),
+            paddle.assign(model.w2.grad),
+        ]
+        input.clear_grad()
+        model.w0.clear_grad()
+        model.w1.clear_grad()
+        model.w2.clear_grad()
+
+        static_model = paddle.jit.to_static(model, full_graph=True)
+        static_res = static_model(input)
+        static_res.backward()
+        static_input_grads = [
+            paddle.assign(input.grad),
+            paddle.assign(model.w0.grad),
+            paddle.assign(model.w1.grad),
+            paddle.assign(model.w2.grad),
+        ]
+        input.clear_grad()
+        model.w0.clear_grad()
+        model.w1.clear_grad()
+        model.w2.clear_grad()
+        for i, (dygraph_grad, static_grad) in enumerate(
+            zip(dygraph_input_grads, static_input_grads)
+        ):
+            np.testing.assert_allclose(
+                dygraph_grad.numpy(),
+                static_grad.numpy(),
+                rtol=1e-5,
+                atol=0,
+                err_msg=f"dygraph_grad[{i}]: {dygraph_grad} \n static_grad[{i}]: {static_grad}",
+            )
+
+
+class NestedStructureWithNonePyLayer(PyLayer):
+    @staticmethod
+    def forward(ctx, x, y):
+        ctx.save_for_backward(x, y)
+        x1 = paddle.tanh(x[0])
+        y1 = paddle.tanh(x[1])
+        z1 = paddle.tanh(y)
+        return [x1, y1, z1]
+
+    @staticmethod
+    def backward(ctx, *grad1):
+        x0, x1 = ctx.saved_tensor()
+        x_grad = grad1[0] * (1 - paddle.square(x0[0]))
+        z_grad = grad1[2] * (1 - paddle.square(x1))
+
+        return [x_grad, None], z_grad
+
+
+class NestedStructureWithNonePyLayerModel(paddle.nn.Layer):
+    def __init__(self):
+        super().__init__()
+        self.w0 = self.create_parameter(shape=[42, 42])
+        self.w1 = self.create_parameter(shape=[42, 42])
+        self.w2 = self.create_parameter(shape=[42, 42])
+
+    def forward(self, x):
+        y1 = paddle.matmul(x, self.w0)
+        y2 = paddle.matmul(x, self.w1)
+        y2.stop_gradient = True
+        y3 = paddle.matmul(x, self.w2)
+
+        z = NestedStructurePyLayer.apply([y1, y2], y3)
+        return z[0] + z[1] + z[2]
+
+
+class TestNestedStructureWithNonePyLayer(unittest.TestCase):
+    def test_nested_structure(self):
+        input = paddle.randn([2, 42]).astype("float32")
+        input.stop_gradient = False
+
+        model = NestedStructurePyLayerModel()
+        dygraph_res = model(input)
+        dygraph_res.backward()
+        dygraph_input_grads = [
+            paddle.assign(input.grad),
+            paddle.assign(model.w0.grad),
+            paddle.assign(model.w1.grad),
+            paddle.assign(model.w2.grad),
+        ]
+        input.clear_grad()
+        model.w0.clear_grad()
+        model.w1.clear_grad()
+        model.w2.clear_grad()
+
+        static_model = paddle.jit.to_static(model, full_graph=True)
+        static_res = static_model(input)
+        static_res.backward()
+        static_input_grads = [
+            paddle.assign(input.grad),
+            paddle.assign(model.w0.grad),
+            paddle.assign(model.w1.grad),
+            paddle.assign(model.w2.grad),
+        ]
+        input.clear_grad()
+        model.w0.clear_grad()
+        model.w1.clear_grad()
+        model.w2.clear_grad()
+        for i, (dygraph_grad, static_grad) in enumerate(
+            zip(dygraph_input_grads, static_input_grads)
+        ):
+            np.testing.assert_allclose(
+                dygraph_grad.numpy(),
+                static_grad.numpy(),
+                rtol=1e-5,
+                atol=0,
+                err_msg=f"dygraph_grad[{i}]: {dygraph_grad} \n static_grad[{i}]: {static_grad}",
+            )
+
+
 if __name__ == "__main__":
     unittest.main()
