Fix backward program kwargs error when process inplace value (#63939)

0x45f · web-flow · commit 66ea6656f4d3 · 2024-05-07T20:41:55.000+08:00
diff --git a/paddle/fluid/imperative/amp_utils.h b/paddle/fluid/imperative/amp_utils.h
@@ -301,6 +301,12 @@ inline T AmpAutoCast(const std::string& input_name,
           input_name == "Ln1Scale" || input_name == "Ln1Bias") {
         return input;
       }
+      if (input_name == "ln_scale" || input_name == "ln_bias" ||
+          input_name == "ln_scale_2" || input_name == "ln_bias_2" ||
+          input_name == "ln1_scale" || input_name == "ln1_bias" ||
+          input_name == "ln2_scale" || input_name == "ln2_bias") {
+        return input;
+      }
     }
   }
   if (NeedCast(input, dst_dtype)) {
diff --git a/paddle/fluid/pybind/eager_functions.cc b/paddle/fluid/pybind/eager_functions.cc
@@ -288,13 +288,13 @@ PyObject* eager_api_get_grads_types(PyObject* self,
   EAGER_TRY
   auto tensor_list = CastPyArg2VectorOfTensor(PyTuple_GET_ITEM(args, 0), 0);
 
-  std::vector<int> ret;
+  std::vector<phi::DataType> ret;
 
   for (auto& tensor : tensor_list) {
     VLOG(6) << "Get grad for tensor: " << tensor.name();
     auto meta = egr::EagerUtils::nullable_autograd_meta(tensor);
     if (!meta || meta->StopGradient()) {
-      ret.emplace_back(-1);
+      ret.emplace_back(phi::DataType::UNDEFINED);
       continue;
     }
 
@@ -304,11 +304,10 @@ PyObject* eager_api_get_grads_types(PyObject* self,
           (tensor.dtype() == phi::DataType::FLOAT32 ||
            tensor.dtype() == phi::DataType::FLOAT16 ||
            tensor.dtype() == phi::DataType::BFLOAT16)) {
-        ret.emplace_back(
-            paddle::framework::TransToProtoVarType(tensor.dtype()));
+        ret.emplace_back(tensor.dtype());
       }
     } else {
-      ret.emplace_back(-1);
+      ret.emplace_back(phi::DataType::UNDEFINED);
     }
   }
 
diff --git a/paddle/fluid/pybind/eager_utils.cc b/paddle/fluid/pybind/eager_utils.cc
@@ -1083,6 +1083,14 @@ PyObject* ToPyObject(const phi::DataType& dtype) {
   return obj.ptr();
 }
 
+PyObject* ToPyObject(const std::vector<phi::DataType>& dtypes) {
+  PyObject* result = PyList_New((Py_ssize_t)dtypes.size());
+  for (size_t i = 0; i < dtypes.size(); i++) {
+    PyList_SET_ITEM(result, static_cast<Py_ssize_t>(i), ToPyObject(dtypes[i]));
+  }
+  return result;
+}
+
 PyObject* ToPyObject(const pir::Value& value) {
   auto obj = ::pybind11::cast(value);
   obj.inc_ref();
diff --git a/paddle/fluid/pybind/eager_utils.h b/paddle/fluid/pybind/eager_utils.h
@@ -148,7 +148,8 @@ PyObject* ToPyObject(const phi::distributed::Placements& value);
 PyObject* ToPyObject(const phi::SelectedRows* value);
 PyObject* ToPyObject(const paddle::framework::proto::VarType::Type& dtype);
 PyObject* ToPyObject(const paddle::framework::proto::VarType& type);
-PyObject* ToPyObject(const phi::DataType& type);
+PyObject* ToPyObject(const phi::DataType& dtype);
+PyObject* ToPyObject(const std::vector<phi::DataType>& dtypes);
 PyObject* ToPyObject(const void* value);
 PyObject* ToPyObject(const std::unordered_map<int, int>& value);
 PyObject* ToPyObject(
diff --git a/paddle/phi/infermeta/fusion.cc b/paddle/phi/infermeta/fusion.cc
@@ -1307,6 +1307,7 @@ void FusedFeedForwardInferMeta(const MetaTensor& x,
     ln2_variance->set_dims(mean_dim);
   }
   out->share_lod(x);
+  out->set_dtype(x.dtype());
 }
 
 static bool IsUnaryCompound(const std::vector<std::string>& functor_list) {
diff --git a/python/paddle/incubate/nn/layer/fused_transformer.py b/python/paddle/incubate/nn/layer/fused_transformer.py
@@ -15,7 +15,6 @@
 
 import paddle
 from paddle.base import core
-from paddle.base.core import VarDesc
 from paddle.base.dygraph import no_grad
 from paddle.base.framework import convert_np_dtype_to_dtype_
 from paddle.framework import in_dynamic_mode
@@ -48,11 +47,12 @@ def _to_dtype(t, dtype):
     if not paddle.is_floating_point(t):
         return t
 
-    if type(dtype) is not VarDesc.VarType:
+    if not isinstance(dtype, (core.VarDesc.VarType, core.DataType)):
         dtype = convert_np_dtype_to_dtype_(dtype)
 
     if t.place.is_gpu_place():
-        size_dtype = core.size_of_dtype(dtype)
+        proto_dtype = paddle.base.framework.convert_to_proto_type(dtype)
+        size_dtype = core.size_of_dtype(proto_dtype)
         waiting_alloc_memory = (
             ((np.prod(t.shape) * size_dtype) / 256 + 1) * 256 * 1.2
         )
diff --git a/python/paddle/jit/dy2static/pir_partial_program.py b/python/paddle/jit/dy2static/pir_partial_program.py
@@ -378,17 +378,22 @@ def has_name(value):
                 ):
                     old_name = name_defining_op.attrs()['output_name']
                     new_name = value2name[ufset.find_root(value)]
+                    if old_name == new_name:
+                        continue
                     paddle.core.pir.reset_shadow_output_name(
                         name_defining_op, new_name
                     )
                     if program2 is None:
                         continue
                     block = program2.global_block()
                     kwargs = block.kwargs()
-                    if old_name in kwargs and new_name not in kwargs:
-                        new_value = block.add_kwarg(
-                            new_name, kwargs[old_name].type()
-                        )
+                    if old_name in kwargs:
+                        if new_name not in kwargs:
+                            new_value = block.add_kwarg(
+                                new_name, kwargs[old_name].type()
+                            )
+                        else:
+                            new_value = kwargs[new_name]
                         kwargs[old_name].replace_all_uses_with(new_value)
                         block.erase_kwarg(old_name)
 
diff --git a/python/paddle/optimizer/adam.py b/python/paddle/optimizer/adam.py
@@ -32,8 +32,6 @@
 
 __all__ = []
 
-GRAD_TYPES = [int(paddle.float32), int(paddle.float16), int(paddle.bfloat16)]
-
 
 class Adam(Optimizer):
     r"""
@@ -570,13 +568,16 @@ def _append_optimize_multi_tensor_op(
                 params = [pair[0] for pair in parameters_and_grads]
                 grads_types = core.eager.get_grads_types(params)
                 for index, tp in enumerate(grads_types):
-                    if tp == GRAD_TYPES[0]:
+                    if tp == core.DataType.FLOAT32:
                         grad_dict['FP32_LODTensor'].append(
                             parameters_and_grads[index][1]
                         )
                         lr = self._create_param_lr(parameters_and_grads[index])
                         lr_dict['FP32_LODTensor'].append(lr)
-                    elif tp == GRAD_TYPES[1] or tp == GRAD_TYPES[2]:
+                    elif (
+                        tp == core.DataType.FLOAT16
+                        or tp == core.DataType.BFLOAT16
+                    ):
                         grad_dict['FP16_LODTensor'].append(
                             parameters_and_grads[index][1]
                         )

Original file line number	Diff line number	Diff line change
`@@ -301,6 +301,12 @@ inline T AmpAutoCast(const std::string& input_name,`
`301`	`301`	`input_name == "Ln1Scale" \|\| input_name == "Ln1Bias") {`
`302`	`302`	`return input;`
`303`	`303`	`}`
	`304`	`+ if (input_name == "ln_scale" \|\| input_name == "ln_bias" \|\|`
	`305`	`+ input_name == "ln_scale_2" \|\| input_name == "ln_bias_2" \|\|`
	`306`	`+ input_name == "ln1_scale" \|\| input_name == "ln1_bias" \|\|`
	`307`	`+ input_name == "ln2_scale" \|\| input_name == "ln2_bias") {`
	`308`	`+ return input;`
	`309`	`+ }`
`304`	`310`	`}`
`305`	`311`	`}`
`306`	`312`	`if (NeedCast(input, dst_dtype)) {`
Original file line number	Diff line number	Diff line change
`@@ -1307,6 +1307,7 @@ void FusedFeedForwardInferMeta(const MetaTensor& x,`
`1307`	`1307`	`ln2_variance->set_dims(mean_dim);`
`1308`	`1308`	`}`
`1309`	`1309`	`out->share_lod(x);`
	`1310`	`+ out->set_dtype(x.dtype());`
`1310`	`1311`	`}`
`1311`	`1312`
`1312`	`1313`	`static bool IsUnaryCompound(const std::vector<std::string>& functor_list) {`