[XPU] Add rms_norm and rms_norm_grad op (#63989)

cqulilujia · web-flow · commit 456366042f0b · 2024-08-29T15:13:51.000+08:00
* [XPU] Add rms_norm and rms_norm_grad op

* use python backword for test
diff --git a/paddle/phi/backends/xpu/xpu2_op_list.cc b/paddle/phi/backends/xpu/xpu2_op_list.cc
@@ -1222,7 +1222,13 @@ XPUOpMap& get_kl2_ops() {
                      phi::DataType::INT64})},
       {"randint", XPUKernelSet({phi::DataType::INT32, phi::DataType::INT64})},
       {"rms_norm",
-       XPUKernelSet({phi::DataType::FLOAT32, phi::DataType::FLOAT16})},
+       XPUKernelSet({phi::DataType::FLOAT32,
+                     phi::DataType::FLOAT16,
+                     phi::DataType::BFLOAT16})},
+      {"rms_norm_grad",
+       XPUKernelSet({phi::DataType::FLOAT32,
+                     phi::DataType::FLOAT16,
+                     phi::DataType::BFLOAT16})},
       {"group_norm", XPUKernelSet({phi::DataType::FLOAT32})},
       {"group_norm_grad", XPUKernelSet({phi::DataType::FLOAT32})},
       {"meshgrid",
diff --git a/paddle/phi/backends/xpu/xpu3_op_list.cc b/paddle/phi/backends/xpu/xpu3_op_list.cc
@@ -876,13 +876,27 @@ XPUOpMap& get_kl3_ops() {
                      phi::DataType::INT32,
                      phi::DataType::INT64,
                      phi::DataType::BFLOAT16})},
+      {"rms_norm",
+       XPUKernelSet({phi::DataType::FLOAT32,
+                     phi::DataType::FLOAT16,
+                     phi::DataType::BFLOAT16})},
+      {"rms_norm_grad",
+       XPUKernelSet({phi::DataType::FLOAT32,
+                     phi::DataType::FLOAT16,
+                     phi::DataType::BFLOAT16})},
       {"rnn_grad", XPUKernelSet({phi::DataType::FLOAT32})},
       {"roi_align", XPUKernelSet({phi::DataType::FLOAT32})},
       {"roi_align_grad", XPUKernelSet({phi::DataType::FLOAT32})},
       {"roll", XPUKernelSet({phi::DataType::FLOAT32})},
       {"roll_grad", XPUKernelSet({phi::DataType::FLOAT32})},
-      {"rsqrt", XPUKernelSet({phi::DataType::FLOAT32})},
-      {"rsqrt_grad", XPUKernelSet({phi::DataType::FLOAT32})},
+      {"rsqrt",
+       XPUKernelSet({phi::DataType::FLOAT32,
+                     phi::DataType::FLOAT16,
+                     phi::DataType::BFLOAT16})},
+      {"rsqrt_grad",
+       XPUKernelSet({phi::DataType::FLOAT32,
+                     phi::DataType::FLOAT16,
+                     phi::DataType::BFLOAT16})},
       {"save",
        XPUKernelSet({phi::DataType::FLOAT32,
                      phi::DataType::FLOAT64,
diff --git a/paddle/phi/infermeta/backward.cc b/paddle/phi/infermeta/backward.cc
@@ -13,6 +13,7 @@ See the License for the specific language governing permissions and
 limitations under the License. */
 
 #include "paddle/phi/infermeta/backward.h"
+#include "glog/logging.h"
 #include "paddle/phi/common/type_traits.h"
 #include "paddle/phi/core/utils/data_type.h"
 #include "paddle/phi/kernels/funcs/axis_utils.h"
@@ -1421,6 +1422,23 @@ void ReshapeDoubleGradInferMeta(const MetaTensor& out_grad,
   }
 }
 
+void RmsNormGradInferMeta(const MetaTensor& x,
+                          const MetaTensor& norm_weight,
+                          const MetaTensor& norm_bias,
+                          MetaTensor* x_grad,
+                          MetaTensor* norm_weight_grad,
+                          MetaTensor* norm_bias_grad) {
+  if (x_grad) {
+    x_grad->share_meta(x);
+  }
+  if (norm_weight && norm_weight_grad) {
+    norm_weight_grad->share_meta(norm_weight);
+  }
+  if (norm_bias && norm_bias_grad) {
+    norm_bias_grad->share_meta(norm_bias);
+  }
+}
+
 void RnnGradInferMeta(const MetaTensor& x,
                       const std::vector<const MetaTensor*>& pre_state,
                       const std::vector<const MetaTensor*>& weight_list,
diff --git a/paddle/phi/infermeta/backward.h b/paddle/phi/infermeta/backward.h
@@ -518,6 +518,13 @@ void ReshapeDoubleGradInferMeta(const MetaTensor& out_grad,
                                 const MetaTensor& x_grad_grad,
                                 MetaTensor* out_grad_grad);
 
+void RmsNormGradInferMeta(const MetaTensor& x,
+                          const MetaTensor& norm_weight,
+                          const MetaTensor& norm_bias,
+                          MetaTensor* x_grad,
+                          MetaTensor* norm_weight_grad,
+                          MetaTensor* norm_bias_grad);
+
 void RnnGradInferMeta(const MetaTensor& x,
                       const std::vector<const MetaTensor*>& pre_state,
                       const std::vector<const MetaTensor*>& weight_list,
diff --git a/paddle/phi/infermeta/multiary.cc b/paddle/phi/infermeta/multiary.cc
@@ -4714,10 +4714,10 @@ void RmsNormInferMeta(const MetaTensor& x,
   }
   PADDLE_ENFORCE_EQ(normalized_dims,
                     norm_weight.dims()[0],
-                    common::errors::InvalidArgument(
-                        "The normalized size of Input(X) must equal to be"
-                        "the size of Weight, but received"
-                        "normalized size of Input(X) is [%d], received size"
+                    phi::errors::InvalidArgument(
+                        "The normalized size of Input(X) must equal to be "
+                        "the size of Weight, but received "
+                        "normalized size of Input(X) is [%d], received size "
                         "of Weight is [%d]",
                         normalized_dims,
                         norm_weight.dims()[0]));
@@ -4744,25 +4744,12 @@ void RmsNormInferMeta(const MetaTensor& x,
     inv_var->set_layout(x.layout());
   }
 
-  residual_out->set_dims(out_dims);
-  residual_out->set_dtype(x.dtype());
-  residual_out->set_layout(x.layout());
-  residual_out->share_lod(x);
-}
-
-void RmsNormGradInferMeta(const MetaTensor& x,
-                          const MetaTensor& norm_weight,
-                          MetaTensor* x_grad,
-                          MetaTensor* norm_weight_grad) {
-  x_grad->set_dtype(x.dtype());
-  x_grad->set_layout(x.layout());
-  x_grad->share_lod(x);
-  x_grad->set_dims(x.dims());
-
-  norm_weight_grad->set_dtype(norm_weight.dtype());
-  norm_weight_grad->set_layout(norm_weight.layout());
-  norm_weight_grad->share_lod(norm_weight);
-  norm_weight_grad->set_dims(norm_weight.dims());
+  if (residual != nullptr) {
+    residual_out->set_dims(out_dims);
+    residual_out->set_dtype(x.dtype());
+    residual_out->set_layout(x.layout());
+    residual_out->share_lod(x);
+  }
 }
 
 void RmspropInferMeta(const MetaTensor& param,
diff --git a/paddle/phi/infermeta/multiary.h b/paddle/phi/infermeta/multiary.h
@@ -908,11 +908,6 @@ void RmsNormInferMeta(const MetaTensor& x,
                       MetaTensor* residual_out,
                       MetaTensor* inv_var);
 
-void RmsNormGradInferMeta(const MetaTensor& x,
-                          const MetaTensor& norm_weight,
-                          MetaTensor* x_grad,
-                          MetaTensor* norm_weight_grad);
-
 void RmspropInferMeta(const MetaTensor& param,
                       const MetaTensor& mean_square,
                       const MetaTensor& grad,
diff --git a/paddle/phi/kernels/gpu/rms_norm_grad_kernel.cu b/paddle/phi/kernels/gpu/rms_norm_grad_kernel.cu
@@ -162,12 +162,13 @@ void RmsNormGradKernel(const Context& dev_ctx,
                        const DenseTensor& norm_weight,
                        const paddle::optional<DenseTensor>& norm_bias,
                        const DenseTensor& inv_var,
-                       const DenseTensor& dy,
+                       const DenseTensor& out_grad,
                        const float epsilon,
                        const int begin_norm_axis,
                        const float quant_scale,
-                       DenseTensor* grad_x,
-                       DenseTensor* grad_norm_weight) {
+                       DenseTensor* x_grad,
+                       DenseTensor* norm_weight_grad,
+                       DenseTensor* norm_bias_grad) {
   if (bias || residual || norm_bias) {
     PADDLE_THROW(common::errors::Unimplemented(
         "bias or residual or norm_bias is not supported yet"));
@@ -179,10 +180,10 @@ void RmsNormGradKernel(const Context& dev_ctx,
                                      x,
                                      norm_weight,
                                      inv_var,
-                                     dy,
+                                     out_grad,
                                      epsilon,
-                                     grad_x,
-                                     grad_norm_weight,
+                                     x_grad,
+                                     norm_weight_grad,
                                      begin_norm_axis);
 }
 }  // namespace phi
diff --git a/paddle/phi/kernels/rms_norm_grad_kernel.h b/paddle/phi/kernels/rms_norm_grad_kernel.h
@@ -0,0 +1,38 @@
+// Copyright (c) 2024 PaddlePaddle Authors. All Rights Reserved.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     https://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+#pragma once
+
+#include "paddle/phi/core/dense_tensor.h"
+#include "paddle/phi/core/selected_rows.h"
+
+namespace phi {
+
+template <typename T, typename Context>
+void RmsNormGradKernel(const Context& dev_ctx,
+                       const DenseTensor& x,
+                       const paddle::optional<DenseTensor>& bias,
+                       const paddle::optional<DenseTensor>& residual,
+                       const DenseTensor& norm_weight,
+                       const paddle::optional<DenseTensor>& norm_bias,
+                       const DenseTensor& inv_var,
+                       const DenseTensor& out_grad,
+                       const float epsilon,
+                       const int begin_norm_axis,
+                       const float quant_scale,
+                       DenseTensor* x_grad,
+                       DenseTensor* norm_weight_grad,
+                       DenseTensor* norm_bias_grad);
+
+}  // namespace phi
diff --git a/paddle/phi/kernels/xpu/rms_norm_grad_kernel.cc b/paddle/phi/kernels/xpu/rms_norm_grad_kernel.cc
@@ -0,0 +1,133 @@
+// Copyright (c) 2024 PaddlePaddle Authors. All Rights Reserved.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     https://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+#include "paddle/phi/kernels/rms_norm_grad_kernel.h"
+#include "paddle/phi/backends/xpu/enforce_xpu.h"
+#include "paddle/phi/core/kernel_registry.h"
+
+namespace phi {
+
+template <typename T, typename Context>
+void RmsNormGradKernel(const Context& dev_ctx,
+                       const DenseTensor& x,
+                       const paddle::optional<DenseTensor>& bias,
+                       const paddle::optional<DenseTensor>& residual,
+                       const DenseTensor& norm_weight,
+                       const paddle::optional<DenseTensor>& norm_bias,
+                       const DenseTensor& inv_var,
+                       const DenseTensor& out_grad,
+                       const float epsilon,
+                       const int begin_norm_axis,
+                       const float quant_scale,
+                       DenseTensor* x_grad,
+                       DenseTensor* norm_weight_grad,
+                       DenseTensor* norm_bias_grad) {
+  if (bias || residual) {
+    PADDLE_THROW(phi::errors::Unimplemented(
+        "bias or residual is not supported in XPU rms_norm_grad yet"));
+  }
+  if (quant_scale > 0.0f) {
+    PADDLE_THROW(phi::errors::Unimplemented(
+        "quantization is not supported in XPU rms_norm_grad yet"));
+  }
+  using XPUType = typename XPUTypeTrait<T>::Type;
+  xpu::ctx_guard RAII_GUARD(dev_ctx.x_context());
+
+  const T* x_data = x.data<T>();
+  const T* norm_weight_data = norm_weight.data<T>();
+  const float* inv_var_data = inv_var.data<float>();
+  const T* out_grad_data = out_grad.data<T>();
+
+  dev_ctx.template Alloc<T>(x_grad);
+  T* x_grad_data = x_grad->data<T>();
+
+  T* norm_weight_grad_data = nullptr;
+  if (norm_weight_grad) {
+    dev_ctx.template Alloc<T>(norm_weight_grad);
+    norm_weight_grad_data = norm_weight_grad->data<T>();
+  }
+
+  T* norm_bias_grad_data = nullptr;
+  if (norm_bias && norm_bias_grad) {
+    dev_ctx.template Alloc<T>(norm_bias_grad);
+    norm_bias_grad_data = norm_bias_grad->data<T>();
+  }
+
+  int32_t rows = 1;
+  int32_t cols = 1;
+  for (int i = 0; i < begin_norm_axis; i++) {
+    rows *= x.dims()[i];
+  }
+
+  for (int i = begin_norm_axis; i < x.dims().size(); i++) {
+    cols *= x.dims()[i];
+  }
+
+  PADDLE_ENFORCE_EQ(
+      cols,
+      norm_weight.dims()[0],
+      phi::errors::InvalidArgument(
+          "The product from begin_norm_axis to last_axis of input tensor x, "
+          "i.e., cols(%d)"
+          "must be equal to the norm_weight tensor's dimension(%d).",
+          cols,
+          norm_weight.dims()[0]));
+
+  if (norm_bias) {
+    PADDLE_ENFORCE_EQ(
+        cols,
+        norm_bias.get().dims()[0],
+        phi::errors::InvalidArgument(
+            "The product from begin_norm_axis to the last axis of input tensor "
+            "x, "
+            "i.e., cols(%d)"
+            "must be equal to the norm_bias tensor's dimension(%d).",
+            cols,
+            norm_bias.get().dims()[0]));
+  }
+
+  PADDLE_ENFORCE_EQ(rows,
+                    inv_var.numel(),
+                    phi::errors::InvalidArgument(
+                        "The product from begin_norm_axis to the last axis of "
+                        "input tensor x, "
+                        "i.e., rows(%d) "
+                        "must be equal to the inv_var tensor's numel(%d). ",
+                        rows,
+                        inv_var.numel()));
+
+  int r = baidu::xpu::api::rms_layer_norm_grad<XPUType, XPUType>(
+      dev_ctx.x_context(),
+      reinterpret_cast<const XPUType*>(x_data),
+      reinterpret_cast<const XPUType*>(out_grad_data),
+      reinterpret_cast<XPUType*>(x_grad_data),
+      rows,
+      cols,
+      epsilon,
+      reinterpret_cast<const XPUType*>(norm_weight_data),
+      inv_var_data,
+      reinterpret_cast<XPUType*>(norm_weight_grad_data),
+      reinterpret_cast<XPUType*>(norm_bias_grad_data),
+      false);
+  PADDLE_ENFORCE_XDNN_SUCCESS(r, "rms_layer_norm_grad");
+}
+}  // namespace phi
+
+PD_REGISTER_KERNEL(rms_norm_grad,
+                   XPU,
+                   ALL_LAYOUT,
+                   phi::RmsNormGradKernel,
+                   float,
+                   phi::dtype::float16,
+                   phi::dtype::bfloat16) {}
diff --git a/paddle/phi/kernels/xpu/rms_norm_kernel.cc b/paddle/phi/kernels/xpu/rms_norm_kernel.cc
diff --git a/paddle/phi/ops/yaml/backward.yaml b/paddle/phi/ops/yaml/backward.yaml
diff --git a/test/xpu/test_rms_norm_op_xpu.py b/test/xpu/test_rms_norm_op_xpu.py
