Add multiprecision for adadelta op (#50131)

AnnaTrainingG · web-flow · commit a8a2b7f47d61 · 2023-03-06T11:25:08.000+08:00
diff --git a/paddle/fluid/operators/optimizers/adadelta_op.cc b/paddle/fluid/operators/optimizers/adadelta_op.cc
@@ -39,12 +39,17 @@ class AdadeltaOpMaker : public framework::OpProtoAndCheckerMaker {
     AddInput("AvgSquaredGrad", "(Tensor) Input average of squared gradient");
     AddInput("AvgSquaredUpdate",
              "(Tensor) Input average of squared parameter updates");
+    AddInput("MasterParam", "FP32 master weight for AMP.").AsDispensable();
 
     AddOutput("ParamOut", "(Tensor) Output parameter");
     AddOutput("AvgSquaredGradOut",
               "(Tensor) Output average of squared gradient");
     AddOutput("AvgSquaredUpdateOut",
               "(Tensor) Output average of squared parameter updates");
+    AddOutput("MasterParamOut",
+              "The updated FP32 master weight for AMP. "
+              "It shared memory with Input(MasterParam).")
+        .AsDispensable();
 
     AddAttr<float>("rho",
                    "(float, default 0.95) Exponential decay rate "
@@ -54,6 +59,10 @@ class AdadeltaOpMaker : public framework::OpProtoAndCheckerMaker {
                    "(float, default 1.0e-6) Constant for "
                    "numerical stability")
         .SetDefault(1.0e-6f);
+    AddAttr<bool>("multi_precision",
+                  "(bool, default false) "
+                  "Whether to use multi-precision during weight updating.")
+        .SetDefault(false);
     AddComment(R"DOC(
 Adadelta Optimizer.
 
diff --git a/paddle/fluid/pybind/eager_generator.h b/paddle/fluid/pybind/eager_generator.h
@@ -206,6 +206,8 @@ std::map<std::string, std::set<std::string>> op_ins_map = {
      {"Q", "K", "V", "Offset", "Columns", "KeyPaddingMask", "AttnMask"}},
     {"sgd", {"Param", "LearningRate", "Grad", "MasterParam"}},
     {"adagrad", {"Param", "Grad", "Moment", "LearningRate", "MasterParam"}},
+    {"adadelta",
+     {"Param", "Grad", "AvgSquaredGrad", "AvgSquaredUpdate", "MasterParam"}},
     {"graph_khop_sampler", {"Row", "Eids", "Col_Ptr", "X"}},
     {"nce",
      {"Input",
@@ -311,6 +313,11 @@ std::map<std::string, std::set<std::string>> op_outs_map = {
       "SavedMean",
       "SavedVariance",
       "ReserveSpace"}},
+    {"adadelta",
+     {"ParamOut",
+      "AvgSquaredGradOut",
+      "AvgSquaredUpdateOut",
+      "MasterParamOut"}},
     {"unique", {"Out", "Index", "Indices", "Counts"}},
     {"unique_consecutive", {"Out", "Index", "Counts"}},
     {"generate_proposals", {"RpnRois", "RpnRoiProbs", "RpnRoisNum"}},
@@ -400,7 +407,11 @@ std::map<std::string, std::set<std::string>> op_passing_outs_map = {
       "MeanGradOut",
       "MasterParamOut"}},
     {"ftrl", {"ParamOut", "SquaredAccumOut", "LinearAccumOut"}},
-    {"adadelta", {"ParamOut", "AvgSquaredGradOut", "AvgSquaredUpdateOut"}},
+    {"adadelta",
+     {"ParamOut",
+      "AvgSquaredGradOut",
+      "AvgSquaredUpdateOut",
+      "MasterParamOut"}},
     {"adagrad", {"ParamOut", "MomentOut", "MasterParamOut"}},
     {"adamax", {"ParamOut", "MomentOut", "InfNormOut"}},
     {"dpsgd", {"ParamOut"}},
diff --git a/paddle/phi/api/yaml/legacy_ops.yaml b/paddle/phi/api/yaml/legacy_ops.yaml
@@ -20,13 +20,15 @@
     data_type : x
 
 - op : adadelta_
-  args : (Tensor param, Tensor grad, Tensor avg_squared_grad, Tensor avg_squared_update, float rho, float epsilon)
-  output : Tensor(param_out), Tensor(moment_out), Tensor(inf_norm_out)
+  args : (Tensor param, Tensor grad, Tensor avg_squared_grad, Tensor avg_squared_update, Tensor master_param, float rho, float epsilon, bool multi_precision)
+  output : Tensor(param_out), Tensor(moment_out), Tensor(inf_norm_out), Tensor(master_param_out)
   infer_meta :
     func : AdadeltaInferMeta
   kernel :
     func : adadelta
-  inplace : (param -> param_out), (avg_squared_grad -> moment_out), (avg_squared_update -> inf_norm_out)
+    data_type : param
+  optional : master_param
+  inplace : (param -> param_out), (avg_squared_grad -> moment_out), (avg_squared_update -> inf_norm_out), (master_param -> master_param_out)
 
 - op : adagrad_
   args : (Tensor param, Tensor grad, Tensor moment, Tensor learning_rate, Tensor master_param, float epsilon, bool multi_precision)
diff --git a/paddle/phi/infermeta/multiary.cc b/paddle/phi/infermeta/multiary.cc
@@ -38,11 +38,14 @@ void AdadeltaInferMeta(const MetaTensor& param,
                        const MetaTensor& grad,
                        const MetaTensor& avg_squared_grad,
                        const MetaTensor& avg_squared_update,
+                       const MetaTensor& master_param,
                        float rho,
                        float epsilon,
+                       bool multi_precision,
                        MetaTensor* param_out,
                        MetaTensor* avg_squared_grad_out,
-                       MetaTensor* avg_squared_update_out) {
+                       MetaTensor* avg_squared_update_out,
+                       MetaTensor* master_param_out) {
   auto param_dims = param.dims();
   PADDLE_ENFORCE_EQ(
       param_dims,
diff --git a/paddle/phi/infermeta/multiary.h b/paddle/phi/infermeta/multiary.h
@@ -43,11 +43,14 @@ void AdadeltaInferMeta(const MetaTensor& param,
                        const MetaTensor& grad,
                        const MetaTensor& avg_squared_grad,
                        const MetaTensor& avg_squared_update,
+                       const MetaTensor& master_param,
                        float rho,
                        float epsilon,
+                       bool multi_precision,
                        MetaTensor* param_out,
                        MetaTensor* avg_squared_grad_out,
-                       MetaTensor* avg_squared_update_out);
+                       MetaTensor* avg_squared_update_out,
+                       MetaTensor* master_param_outs);
 
 void AdagradInferMeta(const MetaTensor& param,
                       const MetaTensor& grad,
diff --git a/paddle/phi/kernels/adadelta_kernel.h b/paddle/phi/kernels/adadelta_kernel.h
@@ -24,10 +24,13 @@ void AdadeltaKernel(const Context& dev_ctx,
                     const DenseTensor& grad,
                     const DenseTensor& avg_squared_grad,
                     const DenseTensor& avg_squared_update,
+                    const paddle::optional<DenseTensor>& master_param,
                     float rho,
                     float epsilon,
+                    bool multi_precision,
                     DenseTensor* param_out,
                     DenseTensor* avg_squared_grad_out,
-                    DenseTensor* avg_squared_update_out);
+                    DenseTensor* avg_squared_update_out,
+                    DenseTensor* master_param_outs);
 
 }  // namespace phi
diff --git a/paddle/phi/kernels/gpu/adadelta_kernel.cu b/paddle/phi/kernels/gpu/adadelta_kernel.cu
@@ -18,5 +18,10 @@
 #include "paddle/phi/core/kernel_registry.h"
 #include "paddle/phi/kernels/impl/adadelta_kernel_impl.h"
 
-PD_REGISTER_KERNEL(
-    adadelta, GPU, ALL_LAYOUT, phi::AdadeltaKernel, float, double) {}
+PD_REGISTER_KERNEL(adadelta,
+                   GPU,
+                   ALL_LAYOUT,
+                   phi::AdadeltaKernel,
+                   float,
+                   double,
+                   phi::dtype::float16) {}
diff --git a/paddle/phi/kernels/impl/adadelta_kernel_impl.h b/paddle/phi/kernels/impl/adadelta_kernel_impl.h
@@ -14,6 +14,7 @@
 
 #pragma once
 
+#include "paddle/phi/common/amp_type_traits.h"
 #include "paddle/phi/kernels/adadelta_kernel.h"
 #include "paddle/phi/kernels/funcs/eigen/common.h"
 #include "paddle/phi/kernels/funcs/eigen/eigen_function.h"
@@ -26,40 +27,58 @@ void AdadeltaKernel(const Context& dev_ctx,
                     const DenseTensor& grad,
                     const DenseTensor& avg_squared_grad,
                     const DenseTensor& avg_squared_update,
+                    const paddle::optional<DenseTensor>& master_param,
                     float rho,
                     float epsilon,
+                    bool multi_precision,
                     DenseTensor* param_out,
                     DenseTensor* avg_squared_grad_out,
-                    DenseTensor* avg_squared_update_out) {
+                    DenseTensor* avg_squared_update_out,
+                    DenseTensor* master_param_outs) {
+  using MPDType = typename phi::dtype::template MPTypeTrait<T>::Type;
   dev_ctx.template Alloc<T>(param_out);
-  dev_ctx.template Alloc<T>(avg_squared_grad_out);
-  dev_ctx.template Alloc<T>(avg_squared_update_out);
+  dev_ctx.template Alloc<MPDType>(avg_squared_grad_out);
+  dev_ctx.template Alloc<MPDType>(avg_squared_update_out);
 
-  T rho_ = static_cast<T>(rho);
-  T epsilon_ = static_cast<T>(epsilon);
+  MPDType rho_ = static_cast<MPDType>(rho);
+  MPDType epsilon_ = static_cast<MPDType>(epsilon);
 
   auto eigen_param = EigenVector<T>::Flatten(param);
   auto eigen_grad = EigenVector<T>::Flatten(grad);
   // Squared gradient accumulator
-  auto eigen_avg_squared_grad = EigenVector<T>::Flatten(avg_squared_grad);
+  auto eigen_avg_squared_grad = EigenVector<MPDType>::Flatten(avg_squared_grad);
   // Squared updates accumulator
-  auto eigen_avg_squared_update = EigenVector<T>::Flatten(avg_squared_update);
+  auto eigen_avg_squared_update =
+      EigenVector<MPDType>::Flatten(avg_squared_update);
   auto eigen_param_out = EigenVector<T>::Flatten(*param_out);
   auto eigen_avg_squared_grad_out =
-      EigenVector<T>::Flatten(*avg_squared_grad_out);
+      EigenVector<MPDType>::Flatten(*avg_squared_grad_out);
   auto eigen_avg_squared_update_out =
-      EigenVector<T>::Flatten(*avg_squared_update_out);
+      EigenVector<MPDType>::Flatten(*avg_squared_update_out);
   auto& place = *dev_ctx.eigen_device();
 
+  auto eigen_grad_cast = eigen_grad.template cast<MPDType>();
+
   eigen_avg_squared_grad_out.device(place) =
-      rho_ * eigen_avg_squared_grad + (1 - rho_) * eigen_grad.square();
+      rho_ * eigen_avg_squared_grad + (1 - rho_) * eigen_grad_cast.square();
   auto update = -((eigen_avg_squared_update + epsilon_) /
                   (eigen_avg_squared_grad_out + epsilon_))
                      .sqrt() *
-                eigen_grad;
+                eigen_grad_cast;
   eigen_avg_squared_update_out.device(place) =
       rho_ * eigen_avg_squared_update + (1 - rho_) * update.square();
-  eigen_param_out.device(place) = eigen_param + update;
+
+  if (multi_precision) {
+    auto eigen_master_param_out =
+        EigenVector<MPDType>::Flatten(*master_param_outs);
+    auto eigen_master_param = EigenVector<MPDType>::Flatten(*master_param);
+
+    eigen_master_param_out.device(place) = eigen_master_param + update;
+    eigen_param_out.device(place) =
+        (eigen_param.template cast<MPDType>() + update).template cast<T>();
+  } else {
+    eigen_param_out.device(place) = eigen_param + update.template cast<T>();
+  }
 }
 
 }  // namespace phi
diff --git a/paddle/phi/kernels/xpu/adadelta_kernel.cc b/paddle/phi/kernels/xpu/adadelta_kernel.cc
@@ -25,11 +25,14 @@ void AdadeltaKernel(const Context& dev_ctx,
                     const DenseTensor& grad,
                     const DenseTensor& avg_squared_grad,
                     const DenseTensor& avg_squared_update,
+                    const paddle::optional<DenseTensor>& master_param,
                     float rho,
                     float epsilon,
+                    bool multi_precision,
                     DenseTensor* param_out,
                     DenseTensor* avg_squared_grad_out,
-                    DenseTensor* avg_squared_update_out) {
+                    DenseTensor* avg_squared_update_out,
+                    DenseTensor* master_param_outs) {
   dev_ctx.template Alloc<T>(param_out);
   dev_ctx.template Alloc<T>(avg_squared_grad_out);
   dev_ctx.template Alloc<T>(avg_squared_update_out);
diff --git a/paddle/phi/ops/compat/adadelta_sig.cc b/paddle/phi/ops/compat/adadelta_sig.cc
@@ -0,0 +1,36 @@
+// Copyright (c) 2023 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
+//
+//     http://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/core/compat/op_utils.h"
+
+namespace phi {
+
+KernelSignature AdadeltaOpArgumentMapping(const ArgumentMappingContext& ctx) {
+  if (ctx.IsDenseTensorInput("Grad")) {
+    return KernelSignature(
+        "adadelta",
+        {"Param", "Grad", "AvgSquaredGrad", "AvgSquaredUpdate", "MasterParam"},
+        {"rho", "epsilon", "multi_precision"},
+        {"ParamOut",
+         "AvgSquaredGradOut",
+         "AvgSquaredUpdateOut",
+         "MasterParamOut"});
+  }
+
+  return KernelSignature("unregistered", {}, {}, {});
+}
+
+}  // namespace phi
+
+PD_REGISTER_ARG_MAPPING_FN(adadelta, phi::AdadeltaOpArgumentMapping);
diff --git a/python/paddle/fluid/optimizer.py b/python/paddle/fluid/optimizer.py
diff --git a/python/paddle/fluid/tests/unittests/test_adadelta_op.py b/python/paddle/fluid/tests/unittests/test_adadelta_op.py
diff --git a/python/paddle/optimizer/adadelta.py b/python/paddle/optimizer/adadelta.py
