support inplace in dygraph eager_fluid state

[pangyoki]

PR types

New features

PR changes

Others

Describe

动态图中间态添加inplace策略。

python-c

实现要点：

• 返回inplace输出tensor时，应该直接将对应inplace输入tensor的PyObject返回，不需新建一个PyObject对象，否则会导致python端的inplace输入输出的id不一致。

static PyObject * eager_api_exp_(PyObject *self, PyObject *args, PyObject *kwargs)
{
  PyThreadState *tstate = nullptr;
  try
  {
    
    auto& X = GetTensorFromArgs("exp", "X", args, 0, false);
    framework::AttributeMap attrs;
    ConstructAttrMapFromPyArgs("exp", args, 1, PyTuple_GET_SIZE(args) , attrs);
    tstate = PyEval_SaveThread();
    auto out = exp__dygraph_function(X, attrs);
    PyEval_RestoreThread(tstate);
    tstate = nullptr;
    ssize_t arg_id = GetIdxFromCoreOpsInfoMap(core_ops_args_info, "exp", "X");
    ssize_t return_id = GetIdxFromCoreOpsInfoMap(core_ops_returns_info, "exp", "Out");
    return ToPyObject(out, return_id, args, arg_id);
  }
  catch(...) {
    if (tstate) {
      PyEval_RestoreThread(tstate);
    }
    ThrowExceptionToPython(std::current_exception());
    return nullptr;
  }
}

动态图层

实现要点：

• 自动代码生成的动态图执行流程有变（不影响功能）：
  • 原流程：先跑TraceOp，然后生成输入和输出的auto_grad meta，然后构反向。
  • 新流程：因为需要为inplace op做check_inplace的检查，check_inplace检查在执行TraceOp之前比较好。另一方面，check_inplace需要使用输入auto_grad meta信息。所以流程变为：先创建输入的auto_grad meta信息，然后check_inplace，然后执行TraceOp生成输出，然后创建输出的auto_grad meta信息，然后构反向。
• 实现inplace功能：TraceOp执行时，直接使用输入的EagerVariable替换输出。后续不重新创建输出Tensor，直接使用输入Tensor代替输出。
• 实现inplace问题：因为EagerVariable内不会改变输入Tensor的meta信息（导致inplace reshape无法改变ddim信息），因此新增了ModifyInplaceInput修改inplace tensor的meta信息。

paddle::experimental::Tensor exp__dygraph_function(paddle::experimental::Tensor& X, const paddle::framework::AttributeMap& attr_map) {

  paddle::platform::RecordEvent dygraph_entrance_record_event("exp dygraph", paddle::platform::TracerEventType::Operator, 1);
  VLOG(3) << "Running Eager Forward Op: exp";
  // Dygraph Forward Pass

  std::map<std::string, std::vector<std::shared_ptr<egr::EagerVariable>>> ins = { { "X", egr::EagerUtils::TrySyncToVars(X) } };

  std::map<std::string, std::vector<std::shared_ptr<egr::EagerVariable>>> outs = { { "Out", ins["X"] } };


  // Prepare Autograd Meta 
  egr::AutogradMeta* p_autograd_X = egr::EagerUtils::nullable_autograd_meta(X);

  bool trace_backward = egr::Controller::Instance().HasGrad();

  bool require_any_grad = egr::EagerUtils::ComputeRequireGrad(trace_backward, p_autograd_X);
  // Check Inplace
  egr::EagerUtils::CheckInplace(X, p_autograd_X, require_any_grad);

  paddle::framework::AttributeMap attrs = attr_map;
  paddle::framework::AttributeMap default_attrs;
  egr::Controller::Instance().GetCurrentTracer()->TraceOp("exp", ins, outs, attrs, 
     egr::Controller::Instance().GetExpectedPlace(),
     &default_attrs, true, {{"X", "Out"}});

  egr::EagerUtils::ModifyInplaceInput(outs["Out"][0], &X);
  X.bump_inplace_version();
  VLOG(3) << "Tensor(" << X.name() << ") uses Inplace Strategy.";

  {
    paddle::platform::RecordEvent node_creation_record_event("exp node_creation", paddle::platform::TracerEventType::Operator, 1);
    p_autograd_X = egr::EagerUtils::autograd_meta(&X);
    if(require_any_grad) {
      VLOG(6) << " Construct Grad for exp "; 
      egr::EagerUtils::PassStopGradient(false, p_autograd_X);
      // Create GradOpNode
      auto grad_node = std::make_shared<GradNodeexp>(1, 1);

      // Set Attributes
      grad_node->SetAttrMap(std::move(attrs));
      grad_node->SetDefaultAttrMap(std::move(default_attrs));

      // Set Tensor Wrappers
      grad_node->SetTensorWrapperOut(X, false);

      grad_node->SetGradOutMeta(p_autograd_X, 0);
      if(p_autograd_X) grad_node->AddEdges(p_autograd_X, 0);
      egr::EagerUtils::SetOutRankWithSlot(p_autograd_X, 0);
      egr::EagerUtils::SetHistory(p_autograd_X, grad_node);
      grad_node->SetGradInMeta(p_autograd_X, 0);
      egr::EagerUtils::CheckAndRetainGrad(X);

    }
  }

  return X;

}

反向检测

• TensorWrapper中加入snapshot_inplace_version_快照信息。
• 在执行反向GradNode，从TensorWrapper recover出Tensor时，进行inplace_version反向检测。比较snapshot_inplace_version_与Tensor的current_inplace_version_是否一致。

示例

• inplace操作

with _test_eager_guard():
    a = paddle.rand([2,3])
    a.stop_gradient=False
    b = a * 2
    c = b ** 2
    d = c.exp_()
    print(c.inplace_version)
    # 1
    print(id(c) == id(d))
    # True
    d.sum().backward()

• stop_gradient=False的叶子节点做inplace操作报错

with _test_eager_guard():
    a = paddle.rand([2,3])
    a.stop_gradient=False
    a.reshape_([-1])

# ValueError: (InvalidArgument) Leaf Var () that doesn't stop gradient can't use inplace strategy.

• 反向检测报错

with _test_eager_guard():
    a = paddle.rand([2,3])
    a.stop_gradient=False
    b = a * 2
    c = b ** 2
    d = b.exp_()
    c.sum().backward()

# RuntimeError: (PermissionDenied) Tensor '' used in gradient computation has been modified by an inplace operation. Its version is 1 but the expected version is 0. Please fix your code to void calling an inplace operator after using the Tensor which will used in gradient computation.

[paddle-bot-old]

✅ This PR's description meets the template requirements!
Please wait for other CI results.

[paddle-bot-old]

Thanks for your contribution!
Please wait for the result of CI firstly. See Paddle CI Manual for details.

[jim19930609]

Looks like we're gonna check inplace version anyway, let's move this function "check_inplace_version" out.

[pangyoki]

done in PR #41118

[jim19930609]

single underscore "_" in function name?

[JiabinYang]

its ok if this method indicate _inplace_version api in python

[jim19930609]

Better add "" at the very end of the function name, like "scale_dygraph_function" for inplaced scale

[jim19930609]

Better store hard-coded op name in a static set

[pangyoki]

done in PR #41118

[JiabinYang]

LGTM

[XieYunshen]

LGTM for set_tests_properties(test_inplace_eager_fluid PROPERTIES TIMEOUT 120)