Dynamic amp support sync_batch_norm op (#32770)

PaddlePaddle · May 10, 2021 · 23ab01e · 23ab01e
1 parent beab956
commit 23ab01e
Showing 1 changed file with 4 additions and 2 deletions.
diff --git a/paddle/fluid/imperative/amp_auto_cast.cc b/paddle/fluid/imperative/amp_auto_cast.cc
@@ -160,7 +160,8 @@ NameVarBaseMap AutoCastInputs(const std::string& op_type,
   if (AmpOperators::Instance().GetMutableAllowOps()->count(op_type)) {
     for (auto& pair : new_ins) {
       // NOTE(zhiqiu): batch_norm and layer_norm support only input x is fp16.
-      if ((op_type == "batch_norm" || op_type == "layer_norm") &&
+      if ((op_type == "batch_norm" || op_type == "layer_norm" ||
+           op_type == "sync_batch_norm") &&
           pair.first != "X") {
         continue;
       }
@@ -191,7 +192,8 @@ NameVarBaseMap AutoCastInputs(const std::string& op_type,
     }
     for (auto& pair : new_ins) {
       // NOTE(zhiqiu): batch_norm and layer_norm support only input x is fp16.
-      if ((op_type == "batch_norm" || op_type == "layer_norm") &&
+      if ((op_type == "batch_norm" || op_type == "layer_norm" ||
+           op_type == "sync_batch_norm") &&
           pair.first == "X" && dst_type == framework::proto::VarType::FP32) {
         continue;
       }