[Dy2St] Record patched name to avoid rollback failures (#69940)

SigureMo · web-flow · commit 173d018f8a2e · 2024-12-04T21:55:13.000+08:00
diff --git a/python/paddle/jit/api.py b/python/paddle/jit/api.py
@@ -335,6 +335,7 @@ def decorated(python_func):
                 logging_utils.warn(
                     f"`{class_name}.forward` has already been decorated somewhere. It will be redecorated to replace previous one."
                 )
+            function._original_funcs["forward"] = function.forward
             function.forward = decorated(function.forward)
             return function
         else:
diff --git a/python/paddle/jit/dy2static/program_translator.py b/python/paddle/jit/dy2static/program_translator.py
@@ -433,6 +433,8 @@ def __init__(self, function, input_spec=None, **kwargs):
         self._cuda_graph_capture_mode = ""
         self._cuda_graph_pool_id = 0
         self._property = kwargs.get("property", False)
+        # Note: Record the patched method name for rollback.
+        self._patched_name = None
         self._get_debug_name()
 
     def _get_debug_name(self) -> str:
@@ -679,19 +681,21 @@ def rollback_impl(class_instance):
             return self._dygraph_function
 
         # only rollback sub-functions on path of top _dygraph_function
-        func_name = self._dygraph_function.__name__
-        assert (
-            func_name in self.class_instance._original_funcs
-        ), f"Not Found function '{func_name}' in class '{self.class_instance.__class__}'."
-        func = self.class_instance._original_funcs[func_name]
-        setattr(
-            self.class_instance, func_name, func.__get__(self.class_instance)
+        fn_name = (
+            self._patched_name
+            if self._patched_name is not None
+            else self._dygraph_function.__name__
         )
+        assert (
+            fn_name in self.class_instance._original_funcs
+        ), f"Not Found function '{fn_name}' in class '{self.class_instance.__class__}'."
+        func = self.class_instance._original_funcs[fn_name]
+        setattr(self.class_instance, fn_name, func.__get__(self.class_instance))
 
         for sublayer in self.class_instance.sublayers(include_self=False):
             rollback_impl(sublayer)
 
-        return getattr(self.class_instance, func_name)
+        return getattr(self.class_instance, fn_name)
 
     def __deepcopy__(self, memo):
         """
diff --git a/python/paddle/nn/layer/layers.py b/python/paddle/nn/layer/layers.py
@@ -1758,6 +1758,12 @@ def _remove_if_exist(*dicts):
                 if name in d:
                     del d[name]
 
+        if isinstance(
+            value, paddle.jit.dy2static.program_translator.StaticFunction
+        ):
+            object.__setattr__(self, name, value)
+            value._patched_name = name
+            return
         if isinstance(getattr(type(self), name, None), property):
             object.__setattr__(self, name, value)
         params = self.__dict__.get('_parameters', None)
diff --git a/test/dygraph_to_static/test_rollback.py b/test/dygraph_to_static/test_rollback.py
@@ -144,5 +144,37 @@ def test_rollback(self):
         self.assertTrue(not isinstance(net.func, StaticFunction))
 
 
+class FuncRollbackWithPatchedFunction(paddle.nn.Layer):
+    def __init__(self) -> None:
+        super().__init__()
+
+    def forward(self, x):
+        return x + 1
+
+
+def patched_fn(self, x):
+    return x + 2
+
+
+FuncRollbackWithPatchedFunction.forward = patched_fn
+
+
+class TestRollBackWithPatchedFunction(Dy2StTestBase):
+    @test_ast_only
+    def test_rollback(self):
+        x = paddle.zeros([2, 2])
+        net = FuncRollbackWithPatchedFunction()
+        dy_out = net(x)
+        static_net = paddle.jit.to_static(net, full_graph=True)
+        st_out = static_net(x)
+        static_net.forward.rollback()
+        dy_out_rollback = net(x)
+
+        self.assertTrue(not isinstance(net.forward, StaticFunction))
+
+        np.testing.assert_array_equal(dy_out.numpy(), st_out.numpy())
+        np.testing.assert_array_equal(dy_out.numpy(), dy_out_rollback.numpy())
+
+
 if __name__ == "__main__":
     unittest.main()

Original file line number	Diff line number	Diff line change
`@@ -335,6 +335,7 @@ def decorated(python_func):`
`335`	`335`	`logging_utils.warn(`
`336`	`336`	f"`{class_name}.forward` has already been decorated somewhere. It will be redecorated to replace previous one."
`337`	`337`	`)`
	`338`	`+ function._original_funcs["forward"] = function.forward`
`338`	`339`	`function.forward = decorated(function.forward)`
`339`	`340`	`return function`
`340`	`341`	`else:`