fix(pir):fix bug layer_norm kernel error when run in pir with amp_o2 use_promote (#68543)

blacksheep-Aristotle · web-flow · commit dcca128ebf17 · 2024-10-10T16:25:23.000+08:00
when use_promote , layer_norm bias and scale should cast to fp32 when input is fp32.
diff --git a/paddle/fluid/eager/amp_auto_cast.h b/paddle/fluid/eager/amp_auto_cast.h
@@ -66,11 +66,6 @@ inline paddle::Tensor AmpAutoCast(const std::string& input_name,
   VLOG(6) << "AMP AmpAutoCasts: op_name(" << op_name << ") input(" << input_name
           << ") dst_dtype(" << phi::DataTypeToString(dst_dtype) << ").";
 
-  if ((op_name == "batch_norm" || op_name == "layer_norm" ||
-       op_name == "sync_batch_norm") &&
-      input_name != "X") {
-    return input;
-  }
   if (op_name == "fused_softmax_mask" && input_name == "Mask" &&
       input.dtype() == phi::DataType::FLOAT32) {
     return input;
@@ -86,6 +81,17 @@ inline paddle::Tensor AmpAutoCast(const std::string& input_name,
         return input;
       }
     }
+    if ((op_name == "batch_norm" || op_name == "layer_norm" ||
+         op_name == "sync_batch_norm" || op_name == "weight_only_linear") &&
+        input_name != "x") {
+      return input;
+    }
+  } else if (dst_dtype == phi::DataType::BFLOAT16) {
+    if ((op_name == "batch_norm" || op_name == "layer_norm" ||
+         op_name == "sync_batch_norm" || op_name == "weight_only_linear") &&
+        input_name != "x") {
+      return input;
+    }
   }
 
   if (NeedCast(input, dst_dtype)) {
diff --git a/paddle/fluid/imperative/amp_utils.h b/paddle/fluid/imperative/amp_utils.h
@@ -279,14 +279,8 @@ inline T AmpAutoCast(const std::string& input_name,
                      const phi::DataType& dst_dtype,
                      const std::string& op_name,
                      bool trace_backward = true) {
-  VLOG(6) << "AMP AmpAutoCasts:"
-          << " input(" << input_name << ") dst_dtype("
-          << phi::DataTypeToString(dst_dtype) << ").";
-  if ((op_name == "batch_norm" || op_name == "layer_norm" ||
-       op_name == "sync_batch_norm" || op_name == "weight_only_linear") &&
-      input_name != "x") {
-    return input;
-  }
+  VLOG(6) << "AMP AmpAutoCasts: op_name(" << op_name << ")input(" << input_name
+          << ") dst_dtype(" << phi::DataTypeToString(dst_dtype) << ").";
 
   if (dst_dtype == phi::DataType::FLOAT16) {
     if (op_name == "run_program") {
@@ -305,6 +299,17 @@ inline T AmpAutoCast(const std::string& input_name,
         return input;
       }
     }
+    if ((op_name == "batch_norm" || op_name == "layer_norm" ||
+         op_name == "sync_batch_norm" || op_name == "weight_only_linear") &&
+        input_name != "x") {
+      return input;
+    }
+  } else if (dst_dtype == phi::DataType::BFLOAT16) {
+    if ((op_name == "batch_norm" || op_name == "layer_norm" ||
+         op_name == "sync_batch_norm" || op_name == "weight_only_linear") &&
+        input_name != "x") {
+      return input;
+    }
   }
   if (NeedCast(input, dst_dtype)) {
     VLOG(6) << "Input : " << input.impl() << "NeedCast";
