Support initializing specific grad tensors to zero for selected operators

paddle/fluid/eager/accumulation/accumulation_node.cc

@@ -39,8 +39,9 @@ static void CopyOrAddTensor(paddle::experimental::Tensor* tensor,
 }
 
 std::vector<std::vector<paddle::experimental::Tensor>> GradNodeAccumulation::
-operator()(const std::vector<std::vector<paddle::experimental::Tensor>>& grads,
- bool create_graph) {
+operator()(
+ std::vector<std::vector<paddle::experimental::Tensor>>& grads, // NOLINT
+ bool create_graph) {
  VLOG(3) << "Running Eager Backward Node: GradNodeAccumulation";
  PADDLE_ENFORCE(grads.size() == 1,
  paddle::platform::errors::Fatal(

paddle/fluid/eager/accumulation/accumulation_node.h

@@ -35,7 +35,7 @@ class GradNodeAccumulation : public GradNodeBase {
 
  // Functor: perform backward computations
  virtual std::vector<std::vector<paddle::experimental::Tensor>> operator()(
- const std::vector<std::vector<paddle::experimental::Tensor>>& grads,
+ std::vector<std::vector<paddle::experimental::Tensor>>& grads, // NOLINT
  bool create_graph = false) override;
 
  void ClearTensorWrappers() override { VLOG(6) << "Do nothing here now"; }

paddle/fluid/eager/api/generated/eager_generated/backwards/scale_node.cc

@@ -145,8 +145,9 @@ void GradNodeScale::SetTensorWrappers_X(
 void GradNodeScale::SetAttributes_scale(float scale) { scale_ = scale; }
 
 std::vector<std::vector<paddle::experimental::Tensor>> GradNodeScale::
-operator()(const std::vector<std::vector<paddle::experimental::Tensor>>& grads,
- bool create_graph) {
+operator()(
+ std::vector<std::vector<paddle::experimental::Tensor>>& grads, // NOLINT
+ bool create_graph) {
  // 1. Check Output Size
  PADDLE_ENFORCE(
  ((grads.size() == 1) && (grads[0].size() == 1)),

paddle/fluid/eager/api/generated/eager_generated/backwards/scale_node.h

@@ -39,7 +39,7 @@ class GradNodeScale : public GradNodeBase {
 
  // Functor: perform backward computations
  virtual std::vector<std::vector<paddle::experimental::Tensor>> operator()(
- const std::vector<std::vector<paddle::experimental::Tensor>>& grads,
+ std::vector<std::vector<paddle::experimental::Tensor>>& grads, // NOLINT
  bool create_graph = false) override;
 
  void ClearTensorWrappers() override { VLOG(6) << "Do nothing here now"; }

paddle/fluid/eager/auto_code_generator/eager_generator.cc

@@ -47,6 +47,9 @@ std::unordered_map<std::string, std::vector<std::string>>
 static std::unordered_map<std::string, paddle::framework::AttributeMap>
  operators_with_attrs = {};
 
+static std::unordered_set<std::string> ops_to_fill_zero_for_empty_grads = {
+ "split"};
+
 /* --- Black Ops list that's NO NEED to apply code generation --- */
 static std::unordered_set<std::string> black_ops_list = {"run_program"};
 
@@ -2243,11 +2246,21 @@ static std::string GenerateGradNodeCCContents(
  // [Generation] Get Full Grad Function
  const char* GRAD_FUNCTION_TEMPLATE =
  "std::vector<std::vector<paddle::experimental::Tensor>> "
- "GradNode%s::operator()(const "
- "std::vector<std::vector<paddle::experimental::Tensor>>& grads, "
- "bool create_graph) {\n%s\n}";
- std::string grad_function_str = paddle::string::Sprintf(
- GRAD_FUNCTION_TEMPLATE, fwd_op_type, generated_grad_function_body);
+ "GradNode%s::operator()("
+ "std::vector<std::vector<paddle::experimental::Tensor>>& grads, bool "
+ "create_graph) {\n"
+ "%s"
+ "%s"
+ "\n}";
+ std::string fill_zero_str = "";
+ if (ops_to_fill_zero_for_empty_grads.count(fwd_op_type)) {
+ fill_zero_str =
+ "egr::EagerUtils::FillZeroForEmptyGradInputs(&grads, "
+ "this->InputMeta());\n";
+ }
+ std::string grad_function_str =
+ paddle::string::Sprintf(GRAD_FUNCTION_TEMPLATE, fwd_op_type,
+ fill_zero_str, generated_grad_function_body);
 
  VLOG(6) << "Generated returns";
 
@@ -2279,9 +2292,9 @@ static std::string GenerateGradNodeHeaderContents(
  " ~GradNode%s() override { VLOG(6) << \" Destruct GradNode%s \"; }\n"
  "\n"
  " virtual std::vector<std::vector<paddle::experimental::Tensor>> "
- "operator()(const "
- "std::vector<std::vector<paddle::experimental::Tensor>>& grads, const "
- "bool create_graph = false) "
+ "operator()("
+ "std::vector<std::vector<paddle::experimental::Tensor>>& grads, bool "
+ "create_graph = false) "
  "override;\n"
  "\n"
  " void ClearTensorWrappers() override { \n"

paddle/fluid/eager/auto_code_generator/final_state_generator/eager_gen.py

@@ -17,6 +17,8 @@
 import argparse
 import os
 
+ops_to_fill_zero_for_empty_grads = set(list("split"))
+
 # For API dispatch used at python-level
 # { op_name : [arg_name, ...] }
 core_ops_returns_info = {}
@@ -598,7 +600,8 @@ class {} : public egr::GradNodeBase {{
  ~{}() override = default;
 
  virtual std::vector<std::vector<paddle::experimental::Tensor>> operator()(
- const std::vector<std::vector<paddle::experimental::Tensor>>& grads, bool create_graph = false) override;
+ std::vector<std::vector<paddle::experimental::Tensor>>& grads, bool create_graph = false) override;
+
  std::string name() override {{ return \" {} \"; }}
 
  void ClearTensorWrappers() override {{
@@ -656,10 +659,11 @@ def GenerateNodeDefinition(fwd_api_name, bwd_api_name, backward_fwd_input_map,
  for _, (ttype, fwd_position,
  grad_api_position) in backward_grad_input_map.items():
  if IsPlainTensorType(ttype):
- grad_api_args[grad_api_position] = f"grads[{fwd_position}][0]"
+ grad_api_args[
+ grad_api_position] = f"hooked_grads[{fwd_position}][0]"
  else:
  assert IsVectorTensorType(ttype)
- grad_api_args[grad_api_position] = f"grads[{fwd_position}]"
+ grad_api_args[grad_api_position] = f"hooked_grads[{fwd_position}]"
 
  for name, _, _, grad_api_position in backward_attrs_list:
  saved_attribute_name = GetSavedName(name)
@@ -687,23 +691,30 @@ def GenerateNodeDefinition(fwd_api_name, bwd_api_name, backward_fwd_input_map,
 
  grad_node_name = GetGradNodeName(fwd_api_name)
 
+ fill_zero_str = ""
+ if fwd_api_name in ops_to_fill_zero_for_empty_grads:
+ fill_zero_str = "egr::EagerUtils::FillZeroForEmptyGradInputs(&grads, this->InputMeta());\n"
+
  if len(namespace) > 0:
  grad_api_namespace = f"paddle::experimental::{namespace}"
  else:
  grad_api_namespace = f"paddle::experimental"
 
  FUNCTION_TEMPLATE = """
-std::vector<std::vector<paddle::experimental::Tensor>> {}::operator()(const std::vector<std::vector<paddle::experimental::Tensor>>& grads, bool create_graph) {{
+std::vector<std::vector<paddle::experimental::Tensor>> {}::operator()(std::vector<std::vector<paddle::experimental::Tensor>>& grads, bool create_graph) {{
+ {}
+ auto hooked_grads = ApplyGradientHooks(grads);
+
  // Call grad_api function
- VLOG(3) << \"Finally State Running: \" << \"{}\"; 
+ VLOG(3) << \"Final State Running: \" << \"{}\"; 
  auto grad_api_returns = {}::{}({});
  {}
 }}
  """
 
  node_definition_str = FUNCTION_TEMPLATE.format(
- grad_node_name, grad_node_name, grad_api_namespace, bwd_api_name,
- grad_api_args_str, returns_str)
+ grad_node_name, fill_zero_str, grad_node_name, grad_api_namespace,
+ bwd_api_name, grad_api_args_str, returns_str)
 
  return node_definition_str
 

paddle/fluid/eager/custom_operator/custom_operator_node.cc

@@ -20,8 +20,8 @@
 
 namespace egr {
 std::vector<std::vector<paddle::experimental::Tensor>> RunCustomOpNode::
-operator()(const std::vector<std::vector<paddle::experimental::Tensor>>& grads,
- bool create_graph) {
+operator()(std::vector<std::vector<paddle::experimental::Tensor>>& grads,
+ bool create_graph) { // NOLINT
  paddle::CustomOpKernelContext ctx;
  auto grad_inputs_name = paddle::framework::OpMetaInfoHelper::GetInputs(
  egr::Controller::Instance().GetOpMetaInfoMap().at(op_type_)[1]);

paddle/fluid/eager/custom_operator/custom_operator_node.h

@@ -37,8 +37,9 @@ class RunCustomOpNode : public GradNodeBase {
 
  // Functor: perform backward computations
  virtual std::vector<std::vector<paddle::experimental::Tensor>> operator()(
- const std::vector<std::vector<paddle::experimental::Tensor>>& grads,
- bool create_graph) override;
+ std::vector<std::vector<paddle::experimental::Tensor>>& grads,
+ bool create_graph = false) // NOLINT
+ override;
 
  std::string name() {
  return paddle::string::Sprintf("RunCustomOpNode: %s_grad", op_type_);

paddle/fluid/eager/grad_node_info.cc

@@ -102,6 +102,7 @@ const std::vector<std::vector<GradSlotMeta>>& GradNodeBase::OutputMeta() const {
 
 void GradNodeBase::SetGradInMeta(const paddle::experimental::Tensor& fwd_out,
  size_t slot_rank) {
+ VLOG(6) << "Set GradSlotMeta for Grad Inputs";
  auto* fwd_out_meta = egr::EagerUtils::nullable_autograd_meta(fwd_out);
  PADDLE_ENFORCE_LE(
  slot_rank, (bwd_in_meta_.size() - 1),
@@ -117,6 +118,12 @@ void GradNodeBase::SetGradInMeta(const paddle::experimental::Tensor& fwd_out,
  auto& meta = metas[0];
  meta.SetStopGradient(fwd_out_meta->StopGradient());
 
+ if (!fwd_out.is_initialized()) {
+ VLOG(6)
+ << "Skip Configuring GradSlotMeta for uninitialized GradInput Tensor";
+ return;
+ }
+
  // Record TensorMeta
  if (phi::DenseTensor::classof(fwd_out.impl().get())) {
  // Only Copy Meta
@@ -128,7 +135,9 @@ void GradNodeBase::SetGradInMeta(const paddle::experimental::Tensor& fwd_out,
  paddle::platform::errors::Fatal(
  "Attempting to copy DenseTensorMeta with phi::DataType::UNDEFINED,"
  "which is illegal."));
+
  meta.SetTensorMeta(dense_tensor->meta());
+ meta.SetPlace(fwd_out.inner_place());
 
  if (paddle::framework::IsComplexType(
  paddle::framework::TransToProtoVarType(dense_tensor->type()))) {
@@ -143,6 +152,7 @@ void GradNodeBase::SetGradInMeta(const paddle::experimental::Tensor& fwd_out,
 void GradNodeBase::SetGradInMeta(
  const std::vector<paddle::experimental::Tensor>& fwd_out,
  size_t slot_rank) {
+ VLOG(6) << "Set GradSlotMeta for Grad Inputs";
  size_t slot_size = fwd_out.size();
  PADDLE_ENFORCE_LE(
  slot_rank, (bwd_in_meta_.size() - 1),
@@ -172,6 +182,12 @@ void GradNodeBase::SetGradInMeta(
  meta.SetStopGradient(fwd_out_meta->StopGradient());
  }
 
+ if (!fwd_out_tensor.is_initialized()) {
+ VLOG(6)
+ << "Skip Configuring GradSlotMeta for uninitialized GradInput Tensor";
+ return;
+ }
+
  // Record TensorMeta
  if (phi::DenseTensor::classof(fwd_out_tensor.impl().get())) {
  // Only Copy Meta
@@ -184,6 +200,8 @@ void GradNodeBase::SetGradInMeta(
  "with phi::DataType::UNDEFINED,"
  "which is illegal."));
  meta.SetTensorMeta(dense_tensor->meta());
+ meta.SetPlace(fwd_out_tensor.inner_place());
+
  if (paddle::framework::IsComplexType(
  paddle::framework::TransToProtoVarType(dense_tensor->type()))) {
  need_complex_to_real_ = true;
@@ -228,6 +246,7 @@ void GradNodeBase::SetGradOutMeta(const paddle::experimental::Tensor& fwd_in,
  "with phi::DataType::UNDEFINED,"
  "which is illegal."));
  meta.SetTensorMeta(dense_tensor->meta());
+ meta.SetPlace(fwd_in.inner_place());
  }
  } else {
  VLOG(6) << "Unable to initialize the DenseTensorMeta of GradSlotMeta with "
@@ -272,6 +291,7 @@ void GradNodeBase::SetGradOutMeta(
  "phi::DataType::UNDEFINED,"
  "which is illegal."));
  meta.SetTensorMeta(dense_tensor->meta());
+ meta.SetPlace(fwd_in_tensor.inner_place());
  }
  } else {
  VLOG(6) << "Unable to initialize the DenseTensorMeta of GradSlotMeta "

paddle/fluid/eager/grad_node_info.h

@@ -76,8 +76,12 @@ class GradSlotMeta {
  return *meta_.get();
  }
 
+ void SetPlace(const phi::Place& place) { place_ = place; }
+ const phi::Place& GetPlace() const { return place_; }
+
  private:
  bool stop_gradient_{false};
+ phi::Place place_;
  std::shared_ptr<phi::DenseTensorMeta> meta_ = nullptr;
 };
 
@@ -102,7 +106,7 @@ class GradNodeBase {
  * is better choice to fit this format.
  * **/
  virtual std::vector<std::vector<paddle::experimental::Tensor>> operator()(
- const std::vector<std::vector<paddle::experimental::Tensor>>& grads,
+ std::vector<std::vector<paddle::experimental::Tensor>>& grads, // NOLINT
  bool create_graph = false) = 0;
 
  virtual void ClearTensorWrappers() = 0;

paddle/fluid/eager/grad_tensor_holder.h

@@ -53,7 +53,7 @@ class GradTensorHolder {
  return buffer_[pos];
  }
 
- const std::vector<std::vector<paddle::experimental::Tensor>>& Buffers() {
+ std::vector<std::vector<paddle::experimental::Tensor>>& Buffers() {
  return buffer_;
  }
 

paddle/fluid/eager/tests/data_structure_tests/accumulation_node_test.cc

@@ -80,13 +80,15 @@ TEST(AccumulationNode, Tensor) {
  grad_meta->SetStopGradient(false);
 
  // operator()
- paddle::experimental::Tensor ret_et0 = node->operator()({{et0}})[0][0];
+ std::vector<std::vector<paddle::experimental::Tensor>> et0_vec = {{et0}};
+ paddle::experimental::Tensor ret_et0 = node->operator()(et0_vec)[0][0];
  auto* ret_et0_ptr =
  std::dynamic_pointer_cast<phi::DenseTensor>(ret_et0.impl())
  ->data<paddle::platform::float16>();
  CHECK_EQ(ret_et0_ptr[0], paddle::platform::float16(10.0f));
 
- paddle::experimental::Tensor ret_et1 = node->operator()({{et1}})[0][0];
+ std::vector<std::vector<paddle::experimental::Tensor>> et1_vec = {{et1}};
+ paddle::experimental::Tensor ret_et1 = node->operator()(et1_vec)[0][0];
 
  auto* ret_et1_ptr =
  std::dynamic_pointer_cast<phi::DenseTensor>(ret_et1.impl())
@@ -121,7 +123,7 @@ TEST(AccumulationNode, Tensor) {
  std::make_shared<egr::CppTensorVoidHook>(reduce_hook_1));
 
  // operator()
- paddle::experimental::Tensor _ret = node->operator()({{et0}})[0][0];
+ paddle::experimental::Tensor _ret = node->operator()(et0_vec)[0][0];
 
  // Check operator() result, should be 36.0
  auto* _ret_ptr = std::dynamic_pointer_cast<phi::DenseTensor>(_ret.impl())

paddle/fluid/eager/tests/data_structure_tests/grad_node_test.h

@@ -32,7 +32,7 @@ class GradTestNode : public egr::GradNodeBase {
  GradTestNode() : GradNodeBase() { val_ = 1.0; }
  std::string name() override { return "GradTestNode"; }
  std::vector<std::vector<paddle::experimental::Tensor>> operator()(
- const std::vector<std::vector<paddle::experimental::Tensor>>& grads,
+ std::vector<std::vector<paddle::experimental::Tensor>>& grads,
  bool create_graph = false) override {
  val_ = std::dynamic_pointer_cast<phi::DenseTensor>(grads[0][0].impl())
  ->data<float>()[0];

paddle/fluid/eager/tests/task_tests/eager_utils_test.cc

@@ -247,4 +247,20 @@ TEST(EagerUtils, GetGradAccumulationNode) {
  ASSERT_ANY_THROW(egr::EagerUtils::GetGradAccumulationNode(t0));
 }
 
+TEST(EagerUtils, FillZeroForEmptyGradInputs) {
+ std::vector<std::vector<paddle::experimental::Tensor>> grads = {
+ std::vector<paddle::experimental::Tensor>(1)};
+ std::vector<std::vector<GradSlotMeta>> slot_metas = {
+ std::vector<GradSlotMeta>(1)};
+
+ phi::DenseTensorMeta tensor_meta;
+ tensor_meta.dtype = paddle::experimental::DataType::FLOAT32;
+ tensor_meta.dims = {2, 4};
+ slot_metas[0][0].SetTensorMeta(tensor_meta);
+ slot_metas[0][0].SetPlace(phi::CPUPlace());
+
+ EagerUtils::FillZeroForEmptyGradInputs(&grads, slot_metas);
+ eager_test::CompareTensorWithValue<float>(grads[0][0], 0.0);
+}
+
 } // namespace egr

paddle/fluid/eager/to_static/run_program_op_node.h

@@ -370,7 +370,7 @@ class GradNodeRunProgram : public egr::GradNodeBase {
  ~GradNodeRunProgram() override = default;
  // Functor: perform backward computations
  virtual std::vector<std::vector<paddle::experimental::Tensor>> operator()(
- const std::vector<std::vector<paddle::experimental::Tensor>> &grads,
+ std::vector<std::vector<paddle::experimental::Tensor>> &grads, // NOLINT
  bool create_graph) override {
  VLOG(3) << "Running Eager Backward Node: GradNodeRunProgram";
  PADDLE_ENFORCE_EQ(

paddle/fluid/eager/utils.cc

@@ -20,6 +20,7 @@
 
 #include "paddle/phi/api/all.h"
 #include "paddle/phi/common/layout.h"
+#include "paddle/phi/core/compat/convert_utils.h"
 #include "paddle/phi/core/tensor_meta.h"
 
 #include "paddle/fluid/framework/data_layout.h"
@@ -392,4 +393,28 @@ std::shared_ptr<egr::GradNodeBase> EagerUtils::GetGradAccumulationNode(
  }
 }
 
+void EagerUtils::FillZeroForEmptyGradInputs(
+ std::vector<std::vector<paddle::experimental::Tensor>>* in_grads,
+ const std::vector<std::vector<GradSlotMeta>>& grad_in_metas) {
+ for (size_t i = 0; i < in_grads->size(); i++) {
+ for (size_t j = 0; j < (*in_grads)[0].size(); j++) {
+ paddle::experimental::Tensor& grad = (*in_grads)[i][j];
+ if (!grad.is_initialized()) {
+ const GradSlotMeta& grad_in_meta = grad_in_metas[i][j];
+ PADDLE_ENFORCE(
+ grad_in_meta.HasTensorMeta(),
+ paddle::platform::errors::Fatal(
+ "Unable to fill empty grad inputs due to empty GradSlotMeta"));
+
+ const auto& tensor_meta = grad_in_meta.GetTensorMeta();
+ phi::Place place = grad_in_meta.GetPlace();
+
+ auto tensor_with_zero = paddle::experimental::full(
+ phi::vectorize(tensor_meta.dims), 0.0, tensor_meta.dtype, place);
+ grad.set_impl(tensor_with_zero.impl());
+ }
+ }
+ }
+}
+
 } // namespace egr