[AMP] fix bf16 amp training error (#54571)

paddle/fluid/eager/amp_auto_cast.h

@@ -69,15 +69,16 @@ inline paddle::Tensor AmpAutoCast(const std::string& input_name,
   VLOG(6) << "AMP AmpAutoCasts:"
           << " input(" << input_name << ") dst_dtype("
           << phi::DataTypeToString(dst_dtype) << ").";
-  if (dst_dtype == phi::DataType::FLOAT16) {
-    if (op_name == "run_program") {
-      return input;
-    }
+
   if ((op_name == "batch_norm" || op_name == "layer_norm" ||
        op_name == "sync_batch_norm") &&
       input_name != "X") {
     return input;
   }
+  if (dst_dtype == phi::DataType::FLOAT16) {
+    if (op_name == "run_program") {
+      return input;
+    }
     if ((op_name == "fused_attention" || op_name == "fused_feedforward")) {
       if (input_name == "LnScale" || input_name == "LnBias" ||
           input_name == "Ln2Scale" || input_name == "Ln2Bias" ||
@@ -86,6 +87,7 @@ inline paddle::Tensor AmpAutoCast(const std::string& input_name,
       }
     }
   }
+
   if (NeedCast(input, dst_dtype)) {
     paddle::framework::AttributeMap cast_attrs = {
         {"in_dtype", paddle::framework::TransToProtoVarType(input.dtype())},

paddle/fluid/eager/amp_utils.h

@@ -26,19 +26,24 @@ static inline phi::DataType GetPromoteType(
                                kSlotSmallVectorSize>& amp_tensors_vector,
     const phi::DataType& amp_dtype) {
   auto dst_type = amp_dtype;
-  if (egr::Controller::Instance().GetCurrentTracer()->GetAmpDtype() ==
-      "float16") {
+  // only consider the dtype of input(X).
   if (op_name == "batch_norm" || op_name == "layer_norm" ||
-        op_name == "sync_batch_norm") {
+      op_name == "sync_batch_norm" ||
+      op_name == "moving_average_abs_max_scale") {
     if (amp_tensors_vector[0][0].dtype() == phi::DataType::FLOAT32) {
       dst_type = phi::DataType::FLOAT32;
     }
-    } else if (op_name == "fused_attention") {
+    return dst_type;
+  }
+
+  if (egr::Controller::Instance().GetCurrentTracer()->GetAmpDtype() ==
+      "float16") {
+    if (op_name == "fused_attention") {
       for (size_t i = 0; i < amp_tensors_vector.size(); i++) {
         if (i != 3 || i != 4 || i != 9 || i != 10) {
           if (amp_tensors_vector[i][0].dtype() == phi::DataType::FLOAT32) {
             dst_type = phi::DataType::FLOAT32;
-            break;
+            return dst_type;
           }
         }
       }
@@ -47,21 +52,13 @@ static inline phi::DataType GetPromoteType(
         if (i != 7 || i != 8 || i != 9 || i != 10) {
           if (amp_tensors_vector[i][0].dtype() == phi::DataType::FLOAT32) {
             dst_type = phi::DataType::FLOAT32;
-            break;
+            return dst_type;
           }
         }
       }
-    } else {
-      for (const auto& tensors : amp_tensors_vector) {
-        for (const auto& tensor : tensors) {
-          if (tensor.dtype() == phi::DataType::FLOAT32) {
-            dst_type = tensor.dtype();
-            break;
-          }
-        }
     }
   }
-  } else {
+
   for (const auto& tensors : amp_tensors_vector) {
     for (const auto& tensor : tensors) {
       if (tensor.dtype() == phi::DataType::FLOAT32) {
@@ -70,14 +67,7 @@ static inline phi::DataType GetPromoteType(
       }
     }
   }
-  }
-  // NOTE(juncai): moving_average_abs_max_scale only consider the dtype of
-  // input(X)
-  if (op_name == "moving_average_abs_max_scale") {
-    if (amp_tensors_vector[0][0].dtype() == phi::DataType::FLOAT16) {
-      dst_type = phi::DataType::FLOAT16;
-    }
-  }
+
   return dst_type;
 }
 

paddle/fluid/eager/eager_amp_auto_cast.h

@@ -89,15 +89,16 @@ inline paddle::Tensor EagerAmpAutoCast(const std::string& input_name,
   VLOG(6) << "AMP AmpAutoCasts:"
           << " input(" << egr::EagerUtils::TensorStr(input) << " to dst_dtype("
           << phi::DataTypeToString(dst_dtype) << ").";
-  if (dst_dtype == phi::DataType::FLOAT16) {
-    if (op_name == "run_program") {
-      return input;
-    }
   if ((op_name == "batch_norm" || op_name == "layer_norm" ||
        op_name == "sync_batch_norm") &&
       input_name != "x") {
     return input;
   }
+
+  if (dst_dtype == phi::DataType::FLOAT16) {
+    if (op_name == "run_program") {
+      return input;
+    }
     if ((op_name == "fused_attention" || op_name == "fused_feedforward")) {
       if (input_name == "LnScale" || input_name == "LnBias" ||
           input_name == "Ln2Scale" || input_name == "Ln2Bias" ||