Support kwargs in forward pre hook (#71283)

* support kwargs in forward pre hook

* fix args passing

* move test case to test_imperative_hook_for_layer.py

* add test after remove hook

* let auto parallel pp llama test case use kwargs after global layer

* rename var and delete useless code

* format code

* fix bug

* fix bug

* remove kwargs flag when removing hook

* fix hook id

Author: liufengwei0103

python/paddle/distributed/auto_parallel/intermediate/pipeline_parallel.py

@@ -217,29 +217,35 @@ def forward_post_hook(layer, input, output):
                     "layer output can only be tensor or list/tuple of tensor"
                 )
 
-        def forward_pre_hook(layer, input):
+        def forward_pre_hook(layer, args, kwargs):
             pp_idx = getattr(layer, "pipeline_stage_index", 0)
-            new_input = []
-            for t in input:
+            new_args = []
+            new_kwargs = {}
+
+            def reshard_tensor_args(t):
                 if is_tensor(t) and t.is_dist() and t.process_mesh == g_mesh:
-                    new_input.append(
-                        dist.reshard(
-                            t,
-                            self.get_mesh(pp_idx),
-                            [dist.Replicate(), dist.Replicate()],
-                        )
+                    return dist.reshard(
+                        t,
+                        self.get_mesh(pp_idx),
+                        [dist.Replicate(), dist.Replicate()],
                     )
-                else:
-                    new_input.append(t)
-            return tuple(new_input)
+                return t
+
+            for arg in args:
+                new_args.append(reshard_tensor_args(arg))
+
+            for key, arg in kwargs.items():
+                new_kwargs[key] = reshard_tensor_args(arg)
+
+            return (new_args, new_kwargs)
 
         for layer_name in self.global_spec:
             layer = self.get_layer_by_name(layer_name)
             layer.register_forward_post_hook(forward_post_hook)
 
         for layer_name in self.pipeline_layers:
             layer = self.get_layer_by_name(layer_name)
-            layer.register_forward_pre_hook(forward_pre_hook)
+            layer.register_forward_pre_hook(forward_pre_hook, with_kwargs=True)
 
 
 def pipeline_parallel(model, optimizer=None, config=None):

python/paddle/nn/layer/layers.py

@@ -338,17 +338,29 @@ class HookRemoveHelper:
     next_hook_id: int = 0
 
     def __init__(
-        self, hooks: typing.OrderedDict[int, Callable[..., Any]]
+        self,
+        hooks: typing.OrderedDict[int, Callable[..., Any]],
+        *,
+        extra_hook_dict: Any = None,
     ) -> None:
         self._hooks_ref = weakref.ref(hooks)
         self._hook_id = HookRemoveHelper.next_hook_id
         HookRemoveHelper.next_hook_id += 1
 
+        self._extra_hooks_ref = None
+        if extra_hook_dict is not None:
+            self._extra_hooks_ref = weakref.ref(extra_hook_dict)
+
     def remove(self) -> None:
         hooks = self._hooks_ref()
         if hooks is not None and self._hook_id in hooks:
             del hooks[self._hook_id]
 
+        if self._extra_hooks_ref is not None:
+            extra_hooks = self._extra_hooks_ref()
+            if extra_hooks is not None and self._hook_id in extra_hooks:
+                del extra_hooks[self._hook_id]
+
 
 class Layer:
     """
@@ -437,6 +449,9 @@ def __init__(
         self._forward_post_hooks: typing.OrderedDict[int, _ForwardPostHook] = (
             OrderedDict()
         )
+        self._forward_pre_hooks_with_kwargs_flag: typing.OrderedDict[
+            int, bool
+        ] = OrderedDict()
 
         # only used in AMP Training
         self._cast_to_low_precision = True
@@ -696,7 +711,7 @@ def register_forward_post_hook(
         return hook_remove_helper
 
     def register_forward_pre_hook(
-        self, hook: _ForwardPreHook
+        self, hook: _ForwardPreHook, *, with_kwargs: bool = False
     ) -> HookRemoveHelper:
         """
 
@@ -748,8 +763,15 @@ def register_forward_pre_hook(
                 >>> # hook change the linear's input to input * 2, so out0 is equal to out1.
                 >>> assert (out0.numpy() == out1.numpy()).any()
         """
-        hook_remove_helper = HookRemoveHelper(self._forward_pre_hooks)
+        hook_remove_helper = HookRemoveHelper(
+            self._forward_pre_hooks,
+            extra_hook_dict=self._forward_pre_hooks_with_kwargs_flag,
+        )
         self._forward_pre_hooks[hook_remove_helper._hook_id] = hook
+        if with_kwargs:
+            self._forward_pre_hooks_with_kwargs_flag[
+                hook_remove_helper._hook_id
+            ] = True
         return hook_remove_helper
 
     def create_parameter(
@@ -1490,12 +1512,27 @@ def _build_once(self, *args: Any, **kwargs: Any) -> None:
         pass
 
     def _dygraph_call_func(self, *inputs: Any, **kwargs: Any) -> Any:
-        for forward_pre_hook in self._forward_pre_hooks.values():
-            hook_result = forward_pre_hook(self, inputs)
-            if hook_result is not None:
-                if not isinstance(hook_result, tuple):
-                    hook_result = (hook_result,)
-                inputs = hook_result
+
+        for hook_id, forward_pre_hook in self._forward_pre_hooks.items():
+            if hook_id in self._forward_pre_hooks_with_kwargs_flag:
+                args_kwargs_result = forward_pre_hook(self, inputs, kwargs)
+                if args_kwargs_result is not None:
+                    if (
+                        isinstance(args_kwargs_result, tuple)
+                        and len(args_kwargs_result) == 2
+                    ):
+                        inputs, kwargs = args_kwargs_result
+                    else:
+                        raise RuntimeError(
+                            "forward pre-hook must return None or a tuple "
+                            f"of (new_args, new_kwargs), but got {args_kwargs_result}."
+                        )
+            else:
+                hook_result = forward_pre_hook(self, inputs)
+                if hook_result is not None:
+                    if not isinstance(hook_result, tuple):
+                        hook_result = (hook_result,)
+                    inputs = hook_result
 
         if not self._built:
             self._build_once(*inputs, **kwargs)

test/auto_parallel/hybrid_strategy/single_llama_model.py

@@ -205,7 +205,9 @@ def forward(self, input_ids):
         global_tensor = self.global_layer(None)
 
         for idx, (decoder_layer) in enumerate(self.layers):
-            hidden_states = decoder_layer(hidden_states, global_tensor)
+            hidden_states = decoder_layer(
+                hidden_states=hidden_states, global_tensor=global_tensor
+            )
 
         hidden_states = self.norm(hidden_states)
 

test/legacy_test/test_imperative_hook_for_layer.py

@@ -224,5 +224,40 @@ def test_forward_hook(self):
                 self.assertFalse(call_forward_pre_hook)
 
 
+def forward_pre_hook_with_kwargs(layer, args, kwargs):
+    kwargs['x'] = kwargs['x'] * 2
+    return (args, kwargs)
+
+
+class SimpleNetWithKWArgs(paddle.nn.Layer):
+    def __init__(
+        self,
+    ):
+        super().__init__()
+
+    def forward(self, x, y):
+        z = x + y
+
+        return z
+
+
+class TestHookWithKWArgs(unittest.TestCase):
+    def test_kwargs_hook(self):
+        net = SimpleNetWithKWArgs()
+        remove_handler = net.register_forward_pre_hook(
+            forward_pre_hook_with_kwargs, with_kwargs=True
+        )
+
+        x = paddle.randn((2, 3))
+        y = paddle.randn((2, 3))
+
+        out = net(x=x, y=y)
+        np.testing.assert_allclose(out.numpy(), (x * 2 + y).numpy())
+
+        remove_handler.remove()
+        out = net(x=x, y=y)
+        np.testing.assert_allclose(out.numpy(), (x + y).numpy())
+
+
 if __name__ == '__main__':
     unittest.main()