【Hackathon 9th No.3】Add bias check for fused_layer_norm (#74851)

* Add bias check for fused_layer_norm

* refine code

Author: ooooo-create

paddle/phi/infermeta/multiary.cc

@@ -2647,6 +2647,33 @@ void FusedLayerNormInferMeta(const MetaTensor& x,
                             x_dims_vec[i],
                             residual_dims_vec[i]));
     }
+    if (bias) {
+      std::vector<int64_t> bias_dims_vec = common::vectorize(bias.dims());
+      PADDLE_ENFORCE_EQ(
+          x_dims_size - begin_norm_axis,
+          bias_dims_vec.size(),
+          common::errors::InvalidArgument(
+              "The normalized size of Input(X) must be equal to the size "
+              "of Bias, but received normalized size of Input(X) is [%d], "
+              "received size of Bias is [%d]",
+              x_dims_size - begin_norm_axis,
+              bias_dims_vec.size()));
+      for (size_t i = begin_norm_axis; i < x_dims_size; ++i) {
+        if (x_dims_vec[i] == -1 || bias_dims_vec[i - begin_norm_axis] == -1 ||
+            x_dims_vec[i] == 0)
+          continue;
+
+        PADDLE_ENFORCE_EQ(x_dims_vec[i],
+                          bias_dims_vec[i - begin_norm_axis],
+                          common::errors::InvalidArgument(
+                              "The normalized dimension of Input(X) and Bias "
+                              "must match at axis %d, but received Input(X) "
+                              "dimension is [%d], Bias dimension is [%d]",
+                              i,
+                              x_dims_vec[i],
+                              bias_dims_vec[i - begin_norm_axis]));
+      }
+    }
   }
 
   int64_t rows = 1;
@@ -2666,6 +2693,18 @@ void FusedLayerNormInferMeta(const MetaTensor& x,
               normalized_dims,
               norm_weight.dims()[0]));
     }
+    if (norm_bias) {
+      PADDLE_ENFORCE_EQ(
+          normalized_dims,
+          norm_bias.dims()[0],
+          common::errors::InvalidArgument(
+              "The normalized size of Input(X) must equal to be "
+              "the size of Bias, but received "
+              "normalized size of Input(X) is [%d], received size "
+              "of Bias is [%d]",
+              normalized_dims,
+              norm_bias.dims()[0]));
+    }
   }
 
   auto out_dims = common::make_ddim(x_dims_vec);

test/legacy_test/test_fused_layernorm_op.py

@@ -1229,5 +1229,41 @@ def setUp(self):
         self.quant_min_bound = -127
 
 
+@unittest.skipIf(
+    not core.is_compiled_with_cuda() and not paddle.is_compiled_with_rocm(),
+    "core is not compiled with CUDA or ROCM",
+)
+class TestFusedLayerNorm_ZeroSize_Error(unittest.TestCase):
+    def test_bias_error(self):
+        with paddle.base.dygraph.guard():
+            x = paddle.randn([16, 256], dtype="float32")
+            bias = paddle.randn([0], dtype="float32")
+            residual = paddle.rand([16, 256], "float32")
+            self.assertRaises(
+                ValueError,
+                paddle.incubate.nn.functional.fused_layer_norm,
+                x=x,
+                norm_weight=paddle.randn([256], dtype="float32"),
+                norm_bias=paddle.randn([256], dtype="float32"),
+                epsilon=1e-06,
+                begin_norm_axis=1,
+                bias=bias,
+                residual=residual,
+            )
+
+            bias = paddle.randn([256], dtype="float32")
+            self.assertRaises(
+                ValueError,
+                paddle.incubate.nn.functional.fused_layer_norm,
+                x=x,
+                norm_weight=paddle.randn([256], dtype="float32"),
+                norm_bias=paddle.randn([0], dtype="float32"),
+                epsilon=1e-06,
+                begin_norm_axis=1,
+                bias=bias,
+                residual=residual,
+            )
+
+
 if __name__ == "__main__":
     unittest.main()