Node register post hook (#71033)

wanghuancoder · web-flow · commit e74360ebe40b · 2025-02-11T10:28:27.000+08:00
* support grad_fn._register_post_hook
diff --git a/paddle/fluid/eager/api/manual/eager_manual/nodes/add_n_node.cc b/paddle/fluid/eager/api/manual/eager_manual/nodes/add_n_node.cc
@@ -109,6 +109,10 @@ AddNGradNodeFinal::operator()(
         INPUT_PRINT_TEMPLATE, input_str, output_str);
   }
 
+  if (HasNodePostHook()) {
+    returns = ApplyNodePostHooks(returns, hooked_grads);
+  }
+
   if (NeedComplexToRealConversion()) HandleComplexGradToRealGrad(&returns);
   return returns;
 }
diff --git a/paddle/fluid/eager/api/manual/eager_manual/nodes/conv2d_nodes.cc b/paddle/fluid/eager/api/manual/eager_manual/nodes/conv2d_nodes.cc
@@ -224,6 +224,10 @@ Conv2dGradNodeFinal::operator()(
         INPUT_PRINT_TEMPLATE, input_str, output_str);
   }
 
+  if (HasNodePostHook()) {
+    returns = ApplyNodePostHooks(returns, hooked_grads);
+  }
+
   // Return
   if (NeedComplexToRealConversion()) HandleComplexGradToRealGrad(&returns);
   return returns;
@@ -412,6 +416,10 @@ Conv2dDoubleGradNodeFinal::operator()(
         INPUT_PRINT_TEMPLATE, input_str, output_str);
   }
 
+  if (HasNodePostHook()) {
+    returns = ApplyNodePostHooks(returns, hooked_grads);
+  }
+
   // Return
   if (NeedComplexToRealConversion()) HandleComplexGradToRealGrad(&returns);
   return returns;
diff --git a/paddle/fluid/eager/api/manual/eager_manual/nodes/multiply_node.cc b/paddle/fluid/eager/api/manual/eager_manual/nodes/multiply_node.cc
@@ -254,6 +254,10 @@ MultiplyGradNode::operator()(
         INPUT_PRINT_TEMPLATE, input_str, output_str);
   }
 
+  if (HasNodePostHook()) {
+    returns = ApplyNodePostHooks(returns, hooked_grads);
+  }
+
   // Return
   if (NeedComplexToRealConversion()) HandleComplexGradToRealGrad(&returns);
   return returns;
@@ -524,6 +528,10 @@ MultiplyDoubleGradNode::operator()(
         INPUT_PRINT_TEMPLATE, input_str, output_str);
   }
 
+  if (HasNodePostHook()) {
+    returns = ApplyNodePostHooks(returns, hooked_grads);
+  }
+
   // Return
   if (NeedComplexToRealConversion()) HandleComplexGradToRealGrad(&returns);
   return returns;
@@ -681,6 +689,10 @@ MultiplyGradNode::operator()(
         INPUT_PRINT_TEMPLATE, input_str, output_str);
   }
 
+  if (HasNodePostHook()) {
+    returns = ApplyNodePostHooks(returns, hooked_grads);
+  }
+
   // Return
   if (NeedComplexToRealConversion()) HandleComplexGradToRealGrad(&returns);
   return returns;
diff --git a/paddle/fluid/eager/api/manual/eager_manual/nodes/sync_batch_norm_node.cc b/paddle/fluid/eager/api/manual/eager_manual/nodes/sync_batch_norm_node.cc
@@ -254,6 +254,10 @@ SyncBatchNormGradNode::operator()(
         INPUT_PRINT_TEMPLATE, input_str, output_str);
   }
 
+  if (HasNodePostHook()) {
+    returns = ApplyNodePostHooks(returns, hooked_grads);
+  }
+
   // Return
   if (NeedComplexToRealConversion()) HandleComplexGradToRealGrad(&returns);
   return returns;
@@ -485,6 +489,10 @@ SyncBatchNormGradNode::operator()(
         INPUT_PRINT_TEMPLATE, input_str, output_str);
   }
 
+  if (HasNodePostHook()) {
+    returns = ApplyNodePostHooks(returns, hooked_grads);
+  }
+
   // Return
   if (NeedComplexToRealConversion()) HandleComplexGradToRealGrad(&returns);
   return returns;
diff --git a/paddle/fluid/eager/auto_code_generator/generator/eager_gen.py b/paddle/fluid/eager/auto_code_generator/generator/eager_gen.py
@@ -362,8 +362,12 @@ class {} : public egr::GradNodeBase {{
   VLOG(4) << \"Finish AD API GRAD: {}";
   VLOG(6) << "gradnode_ptr = " << this;
   // LOG IF DEBUG
-
 {}
+
+  if (HasNodePostHook()) {{
+    returns = ApplyNodePostHooks(returns, hooked_grads);
+  }}
+
   // Return
 {}
 }}
diff --git a/paddle/fluid/eager/grad_node_info.cc b/paddle/fluid/eager/grad_node_info.cc
@@ -867,4 +867,49 @@ uintptr_t GradNodeBase::GetPtr() const {
   return reinterpret_cast<uintptr_t>(this);
 }
 
+int64_t GradNodeBase::RegisterNodePostHook(
+    std::shared_ptr<NodePostHookBase>&& hook) {
+  post_hooks_.emplace(next_post_hook_id_, std::move(hook));
+  return next_post_hook_id_++;
+}
+
+bool GradNodeBase::RemoveNodePostHook(int64_t hook_id) {
+  auto remove_cnt = post_hooks_.erase(hook_id);
+  if (remove_cnt == 0) {
+    return false;
+  }
+  return true;
+}
+
+bool GradNodeBase::HasNodePostHook() { return !post_hooks_.empty(); }
+
+paddle::small_vector<std::vector<paddle::Tensor>, egr::kSlotSmallVectorSize>
+GradNodeBase::ApplyNodePostHooks(
+    const paddle::small_vector<std::vector<paddle::Tensor>,
+                               egr::kSlotSmallVectorSize>& grad_outputs,
+    const paddle::small_vector<std::vector<paddle::Tensor>,
+                               egr::kSlotSmallVectorSize>& grad_inputs) {
+  paddle::small_vector<std::vector<paddle::Tensor>, kSlotSmallVectorSize> outs =
+      grad_outputs;
+  for (auto& iter : post_hooks_) {
+    auto hook = iter.second;
+    outs = (*hook)(outs, grad_inputs);
+  }
+
+  for (size_t i = 0; i < outs.size(); i++) {
+    if (outs[i].empty() && (!grad_outputs[i].empty())) {
+      outs[i].resize(grad_outputs[i].size());
+    }
+
+    for (size_t j = 0; j < outs[i].size(); j++) {
+      if (!outs[i][j].defined() || !outs[i][j].initialized()) {
+        outs[i][j] = grad_outputs[i][j];
+      } else {
+        CheckTensor(grad_outputs[i][j], outs[i][j]);
+      }
+    }
+  }
+
+  return outs;
+}
 }  // namespace egr
diff --git a/paddle/fluid/eager/grad_node_info.h b/paddle/fluid/eager/grad_node_info.h
@@ -340,6 +340,16 @@ class GradNodeBase {
     is_run_auto_parallel_ = is_run_auto_parallel;
   }
 
+  int64_t RegisterNodePostHook(std::shared_ptr<NodePostHookBase>&& hook);
+  bool RemoveNodePostHook(int64_t hook_id);
+  bool HasNodePostHook();
+  paddle::small_vector<std::vector<paddle::Tensor>, egr::kSlotSmallVectorSize>
+  ApplyNodePostHooks(
+      const paddle::small_vector<std::vector<paddle::Tensor>,
+                                 egr::kSlotSmallVectorSize>& grad_outputs,
+      const paddle::small_vector<std::vector<paddle::Tensor>,
+                                 egr::kSlotSmallVectorSize>& grad_inputs);
+
  private:
   // bwd_out_meta_ is used to record Grad output info for backward
   paddle::small_vector<std::vector<GradSlotMeta>, kSlotSmallVectorSize>
@@ -361,6 +371,9 @@ class GradNodeBase {
       gradient_hooks_;
   int64_t next_hook_id_{0};
 
+  std::map<int64_t, std::shared_ptr<NodePostHookBase>> post_hooks_;
+  int64_t next_post_hook_id_{0};
+
   // We handle complex to real conversion only if any complex GradIn is involved
   bool need_complex_to_real_ = false;
 
diff --git a/paddle/fluid/eager/hooks.h b/paddle/fluid/eager/hooks.h
@@ -19,6 +19,7 @@
 #include <utility>
 #include <vector>
 
+#include "paddle/fluid/eager/type_defs.h"
 #include "paddle/fluid/platform/enforce.h"
 #include "paddle/phi/api/include/tensor.h"
 namespace egr {
@@ -124,4 +125,16 @@ class SavedTensorsHooks {
   bool is_enable_{false};
 };
 
+class NodePostHookBase {
+ public:
+  virtual ~NodePostHookBase() = default;
+  virtual paddle::small_vector<std::vector<paddle::Tensor>,
+                               egr::kSlotSmallVectorSize>
+  operator()(
+      const paddle::small_vector<std::vector<paddle::Tensor>,
+                                 egr::kSlotSmallVectorSize>& grad_outputs,
+      const paddle::small_vector<std::vector<paddle::Tensor>,
+                                 egr::kSlotSmallVectorSize>& grad_inputs) = 0;
+};
+
 }  // namespace egr
diff --git a/paddle/fluid/pybind/eager_utils.cc b/paddle/fluid/pybind/eager_utils.cc
@@ -2720,6 +2720,95 @@ void* UnPackHook::operator()(void* packed_value, void* other) {
   return reinterpret_cast<void*>(ret);
 }
 
+PyObject* ToPyObject(
+    const paddle::small_vector<std::vector<paddle::Tensor>,
+                               egr::kSlotSmallVectorSize>& grads) {
+  PyObject* args = nullptr;
+  args = PyTuple_New(grads.size());
+
+  for (size_t i = 0; i < grads.size(); i++) {
+    if (grads[i].size() == 0) {
+      Py_INCREF(Py_None);
+      PyTuple_SET_ITEM(args, i, Py_None);
+    } else if (grads[i].size() == 1) {
+      PyTuple_SET_ITEM(args, i, ToPyObject(grads[i][0]));
+    } else {
+      PyTuple_SET_ITEM(args, i, ToPyObject(grads[i]));
+    }
+  }
+
+  return args;
+}
+
+paddle::small_vector<std::vector<paddle::Tensor>, egr::kSlotSmallVectorSize>
+CastPyArg2SmallVectorOfVectorOfTensor(PyObject* obj) {
+  paddle::small_vector<std::vector<paddle::Tensor>, egr::kSlotSmallVectorSize>
+      result;
+  if (PyList_Check(obj)) {
+    Py_ssize_t len = PyList_Size(obj);
+    PyObject* item = nullptr;
+    for (Py_ssize_t i = 0; i < len; i++) {
+      item = PyList_GetItem(obj, i);
+      if (PyObject_TypeCheck(item, p_tensor_type)) {
+        std::vector<paddle::Tensor> tensors;
+        tensors.push_back(reinterpret_cast<TensorObject*>(item)->tensor);
+        result.emplace_back(tensors);
+      } else if (item == Py_None) {
+        // emplace empty Tensor for None
+        std::vector<paddle::Tensor> tensors;
+        result.emplace_back(tensors);
+      } else {
+        result.emplace_back(CastPyArg2VectorOfTensor(obj, 0));
+      }
+    }
+  } else if (PyTuple_Check(obj)) {
+    Py_ssize_t len = PyTuple_Size(obj);
+    PyObject* item = nullptr;
+    for (Py_ssize_t i = 0; i < len; i++) {
+      item = PyTuple_GetItem(obj, i);
+      if (PyObject_TypeCheck(item, p_tensor_type)) {
+        std::vector<paddle::Tensor> tensors;
+        tensors.push_back(reinterpret_cast<TensorObject*>(item)->tensor);
+        result.emplace_back(tensors);
+      } else if (item == Py_None) {
+        // emplace empty Tensor for None
+        std::vector<paddle::Tensor> tensors;
+        result.emplace_back(tensors);
+      } else {
+        result.emplace_back(CastPyArg2VectorOfTensor(obj, 0));
+      }
+    }
+  } else {
+    PADDLE_THROW(common::errors::InvalidType(
+        "argument must be "
+        "list or tuple, but got %s",
+        reinterpret_cast<PyTypeObject*>(obj->ob_type)->tp_name));
+  }
+  return result;
+}
+
+paddle::small_vector<std::vector<paddle::Tensor>, egr::kSlotSmallVectorSize>
+NodePostHook::operator()(
+    const paddle::small_vector<std::vector<paddle::Tensor>,
+                               egr::kSlotSmallVectorSize>& grad_outputs,
+    const paddle::small_vector<std::vector<paddle::Tensor>,
+                               egr::kSlotSmallVectorSize>& grad_inputs) {
+  bool grad_tmp = egr::Controller::Instance().HasGrad();
+  egr::Controller::Instance().SetHasGrad(false);
+  ::pybind11::gil_scoped_acquire gil;
+  PyObject* args = PyTuple_New(2);
+  PADDLE_ENFORCE_NOT_NULL(
+      args, common::errors::External(pybind11::detail::error_string().c_str()));
+  PyTuple_SET_ITEM(args, 0, ToPyObject(grad_outputs));
+  PyTuple_SET_ITEM(args, 1, ToPyObject(grad_inputs));
+  PyObject* ret = PyObject_Call(hook_.ptr(), args, nullptr);
+  PADDLE_ENFORCE_NOT_NULL(
+      ret, common::errors::External(pybind11::detail::error_string().c_str()));
+  Py_XDECREF(args);
+  egr::Controller::Instance().SetHasGrad(grad_tmp);
+  return CastPyArg2SmallVectorOfVectorOfTensor(ret);
+}
+
 /* ------------------ for SetStaticOpArgPreCastHook ----------------------- */
 
 static Py_tss_t static_op_arg_pre_cast_hook_key = {0, 0};
diff --git a/paddle/fluid/pybind/eager_utils.h b/paddle/fluid/pybind/eager_utils.h
@@ -241,6 +241,27 @@ class UnPackHook : public egr::UnPackHookBase {
  private:
   PyObject* hook_;
 };
+
+#pragma GCC visibility push(hidden)
+class NodePostHook : public egr::NodePostHookBase {
+ public:
+  explicit NodePostHook(py::object hook) : hook_(hook) {}
+
+  ~NodePostHook() {}
+
+  paddle::small_vector<std::vector<paddle::Tensor>, egr::kSlotSmallVectorSize>
+  operator()(
+      const paddle::small_vector<std::vector<paddle::Tensor>,
+                                 egr::kSlotSmallVectorSize>& grad_outputs,
+      const paddle::small_vector<std::vector<paddle::Tensor>,
+                                 egr::kSlotSmallVectorSize>& grad_inputs)
+      override;
+
+ private:
+  py::object hook_;
+};
+#pragma GCC visibility pop
+
 template <typename Tuple, size_t N>
 struct TupleTensorResult {
   static void Run(const Tuple& out, PyObject* result) {
diff --git a/paddle/fluid/pybind/pybind.cc b/paddle/fluid/pybind/pybind.cc
@@ -236,6 +236,7 @@ limitations under the License. */
 #include "paddle/fluid/inference/tensorrt/pir/declare_plugin.h"
 #include "paddle/fluid/platform/tensorrt/trt_plugin.h"
 #endif
+#include "paddle/fluid/eager/accumulation/accumulation_node.h"
 
 COMMON_DECLARE_bool(use_mkldnn);
 COMMON_DECLARE_string(prim_backward_blacklist);
@@ -1192,6 +1193,24 @@ PYBIND11_MODULE(libpaddle, m) {
           }
         });
 
+  class NodePostHookRemoveHelper {
+   public:
+    NodePostHookRemoveHelper(std::shared_ptr<egr::GradNodeBase> node,
+                             int64_t hook_id)
+        : node_(node), hook_id_(hook_id) {}
+    ~NodePostHookRemoveHelper() = default;
+    bool remove() { return node_->RemoveNodePostHook(hook_id_); }
+
+   private:
+    std::shared_ptr<egr::GradNodeBase> node_;
+    int64_t hook_id_;
+  };
+
+  py::class_<NodePostHookRemoveHelper,
+             std::shared_ptr<NodePostHookRemoveHelper>>(
+      m, "NodePostHookRemoveHelper")
+      .def("remove", &NodePostHookRemoveHelper::remove);
+
   py::class_<egr::GradNodeBase, std::shared_ptr<egr::GradNodeBase>>(
       m, "GradNodeBase")
       .def("name",
@@ -1208,9 +1227,20 @@ PYBIND11_MODULE(libpaddle, m) {
            [](const std::shared_ptr<egr::GradNodeBase> &self) {
              return self->InputMeta();
            })
-      .def("output_meta", [](const std::shared_ptr<egr::GradNodeBase> &self) {
-        return self->OutputMeta();
-      });
+      .def("output_meta",
+           [](const std::shared_ptr<egr::GradNodeBase> &self) {
+             return self->OutputMeta();
+           })
+      .def("_register_post_hook",
+           [](const std::shared_ptr<egr::GradNodeBase> &self, py::object hook) {
+             if (std::dynamic_pointer_cast<egr::GradNodeAccumulation>(self)) {
+               PADDLE_THROW(common::errors::InvalidArgument(
+                   "Could not register hook for GradNodeAccumulation."));
+             }
+             int64_t hook_id = self->RegisterNodePostHook(
+                 std::make_shared<NodePostHook>(hook));
+             return std::make_shared<NodePostHookRemoveHelper>(self, hook_id);
+           });
 
 #if defined(PADDLE_WITH_CUDA) || defined(PADDLE_WITH_HIP)
   m.def("cudnn_version", &platform::DnnVersion);
diff --git a/test/legacy_test/test_node_register_hook.py b/test/legacy_test/test_node_register_hook.py

Original file line number	Diff line number	Diff line change
`@@ -109,6 +109,10 @@ AddNGradNodeFinal::operator()(`
`109`	`109`	`INPUT_PRINT_TEMPLATE, input_str, output_str);`
`110`	`110`	`}`
`111`	`111`
	`112`	`+ if (HasNodePostHook()) {`
	`113`	`+ returns = ApplyNodePostHooks(returns, hooked_grads);`
	`114`	`+ }`
	`115`	`+`
`112`	`116`	`if (NeedComplexToRealConversion()) HandleComplexGradToRealGrad(&returns);`
`113`	`117`	`return returns;`
`114`	`118`	`}`