chore: import upstream snapshot with attribution

paddle/phi/kernels/selected_rows/impl/add_n_kernel_impl.h

@@ -0,0 +1,90 @@
+// Copyright (c) 2022 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.
+
+#pragma once
+
+#include "paddle/phi/kernels/selected_rows/add_n_kernel.h"
+
+#include "paddle/phi/core/kernel_registry.h"
+#include "paddle/phi/core/tensor_utils.h"
+#include "paddle/phi/kernels/funcs/math_function.h"
+#include "paddle/phi/kernels/funcs/selected_rows_functor.h"
+
+namespace phi {
+namespace sr {
+template <typename T, typename Context>
+void AddNKernel(const Context &dev_ctx,
+                const std::vector<const SelectedRows *> &x,
+                SelectedRows *out) {
+  dev_ctx.template Alloc<T>(out->mutable_value());
+
+  bool in_place = false;
+  if (x.size() > 0 && x[0]->value().Holder() == out->value().Holder()) {
+    in_place = true;
+  }
+
+  if (in_place && x.size() < 2) {
+    return;
+  }
+
+  std::vector<const SelectedRows *> inputs;
+  SelectedRows temp_in0;
+
+  if (in_place) {
+    auto &in0 = *x[0];
+    temp_in0.set_height(in0.height());
+    temp_in0.set_rows(in0.rows());
+    Copy<Context>(
+        dev_ctx, in0.value(), in0.place(), false, temp_in0.mutable_value());
+    inputs.push_back(&temp_in0);
+    for (size_t i = 1; i < x.size(); ++i) {
+      auto &in = *x[i];
+      if (in.rows().size() > 0) {
+        inputs.push_back(&in);
+      }
+    }
+  } else {
+    for (auto in_var : x) {
+      auto &in = *in_var;
+      if (in.rows().size() > 0) {
+        inputs.push_back(in_var);
+      }
+    }
+  }
+
+  out->mutable_rows()->clear();
+
+  bool has_data = false;
+  for (auto &in : inputs) {
+    if (in->rows().size() > 0) {
+      has_data = true;
+      break;
+    }
+  }
+  if (has_data) {
+    funcs::scatter::MergeAdd<Context, T> merge_add;
+    merge_add(dev_ctx, inputs, out);
+
+    out->SyncIndex();
+
+  } else {
+    // no data, just set a empty out tensor.
+    auto *out_dense = out->mutable_value();
+    out_dense->clear();
+    out_dense->Resize({0});
+    dev_ctx.template Alloc<T>(out_dense);
+  }
+}
+}  // namespace sr
+}  // namespace phi

paddle/phi/kernels/selected_rows/impl/clip_by_norm_kernel_impl.h

@@ -0,0 +1,45 @@
+// Copyright (c) 2022 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.
+
+#pragma once
+
+#include "paddle/phi/core/dense_tensor.h"
+#include "paddle/phi/core/device_context.h"
+#include "paddle/phi/core/selected_rows.h"
+#include "paddle/phi/kernels/clip_by_norm_kernel.h"
+#include "paddle/phi/kernels/funcs/selected_rows_functor.h"
+#include "paddle/phi/kernels/selected_rows/clip_by_norm_kernel.h"
+
+namespace phi {
+namespace sr {
+
+template <typename T, typename Context>
+void ClipByNormKernel(const Context& dev_ctx,
+                      const SelectedRows& x,
+                      float max_norm,
+                      SelectedRows* out) {
+  SelectedRows merged_input;
+  funcs::scatter::MergeAdd<Context, T> merge_func;
+  merge_func(dev_ctx, x, &merged_input);
+  auto input = &(merged_input.value());
+  out->set_rows(merged_input.rows());
+  out->set_height(merged_input.height());
+  auto out_tensor = out->mutable_value();
+  out_tensor->Resize(merged_input.value().dims());
+  return phi::ClipByNormKernel<T, Context>(
+      dev_ctx, *input, max_norm, out_tensor);
+}
+
+}  // namespace sr
+}  // namespace phi

paddle/phi/kernels/selected_rows/impl/clip_kernel_impl.h

@@ -0,0 +1,61 @@
+// Copyright (c) 2022 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.
+
+#pragma once
+
+#include "paddle/phi/common/scalar.h"
+#include "paddle/phi/core/dense_tensor.h"
+#include "paddle/phi/core/device_context.h"
+#include "paddle/phi/core/selected_rows.h"
+#include "paddle/phi/kernels/funcs/selected_rows_functor.h"
+#include "paddle/phi/kernels/selected_rows/clip_kernel.h"
+
+namespace phi {
+namespace sr {
+
+template <typename T, typename Context>
+void ClipSparseKernel(const Context& dev_ctx,
+                      const SelectedRows& x,
+                      const Scalar& min,
+                      const Scalar& max,
+                      SelectedRows* out) {
+  auto max_ = max.to<T>();
+  auto min_ = min.to<T>();
+
+  PADDLE_ENFORCE_LE(
+      min_,
+      max_,
+      errors::InvalidArgument("max should be greater than or equal to min. "
+                              "But received min = %f, max = %f",
+                              static_cast<float>(min_),
+                              static_cast<float>(max_)));
+
+  PADDLE_ENFORCE_NE(&x,
+                    out,
+                    errors::InvalidArgument("Inplace clip is not allowed "
+                                            "when x is SelectedRows"));
+  funcs::scatter::MergeAdd<Context, T> merge_func;
+  merge_func(dev_ctx, x, out);
+  auto* out_tensor = out->mutable_value();
+  auto* out_data = out_tensor->data<T>();
+  int64_t numel = out_tensor->numel();
+  phi::Transform<Context> trans;
+  trans(dev_ctx,
+        out_data,
+        out_data + numel,
+        out_data,
+        ClipFunctor<T>(min_, max_));
+}
+}  // namespace sr
+}  // namespace phi

paddle/phi/kernels/selected_rows/impl/dgc_clip_by_norm_kernel_impl.h

@@ -0,0 +1,51 @@
+// 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
+//
+//     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.
+
+#pragma once
+#include "glog/logging.h"
+#include "paddle/phi/kernels/selected_rows/clip_by_norm_kernel.h"
+
+namespace phi {
+namespace sr {
+template <typename T, typename Context>
+void DGCClipByNormKernel(const Context& dev_ctx,
+                         const SelectedRows& x_in,
+                         const DenseTensor& current_step_in,
+                         float max_norm,
+                         float rampup_begin_step,
+                         SelectedRows* out) {
+  if (static_cast<int>(rampup_begin_step) < 0) {
+    return;
+  }
+
+  auto current_step_tensor = &current_step_in;
+  auto* current_step = current_step_tensor->data<T>();
+
+  VLOG(10) << "current_step:" << *current_step
+           << ", rampup_begin_step:" << rampup_begin_step;
+
+  if (static_cast<int>(*current_step) < static_cast<int>(rampup_begin_step)) {
+    VLOG(10) << "current_step:" << *current_step
+             << " < rampup_begin_step:" << rampup_begin_step
+             << " so doesn't use dgc_clip_by_norm";
+    return;
+  }
+
+  auto* x = &x_in;
+  SelectedRows* output_selected_rows = out;
+  return phi::sr::ClipByNormKernel<T>(
+      dev_ctx, *x, max_norm, output_selected_rows);
+}
+}  // namespace sr
+}  // namespace phi

paddle/phi/kernels/selected_rows/impl/ftrl_kernel_impl.h

@@ -0,0 +1,176 @@
+// 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
+//
+//     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.
+
+#pragma once
+#include "paddle/phi/core/dense_tensor.h"
+#include "paddle/phi/core/kernel_registry.h"
+#include "paddle/phi/kernels/funcs/eigen/common.h"
+#include "paddle/phi/kernels/funcs/for_range.h"
+#include "paddle/phi/kernels/funcs/selected_rows_functor.h"
+
+namespace phi {
+namespace sr {
+
+template <typename T, int MajorType = Eigen::RowMajor>
+using EigenVector = EigenVector<T, MajorType>;
+
+template <typename T>
+class SparseFTRLFunctor {
+ private:
+  const T* g_;
+  const T* p_;
+  const T* s_acc_;
+  const T* l_acc_;
+  const T* lr_;
+  const T l1_;
+  const T l2_;
+  const T lr_power_;
+  const int64_t* rows_;
+  const int64_t row_numel_;
+  T* p_out_;
+  T* s_acc_out_;
+  T* l_acc_out_;
+
+ public:
+  SparseFTRLFunctor(const T* g,
+                    const T* p,
+                    const T* s_acc,
+                    const T* lr,
+                    const T l1,
+                    const T l2,
+                    const T lr_power,
+                    const int64_t* rows,
+                    int64_t row_numel,
+                    T* p_out,
+                    T* s_acc_out,
+                    T* l_acc_out)
+      : g_(g),
+        p_(p),
+        s_acc_(s_acc),
+        lr_(lr),
+        l1_(l1),
+        l2_(l2),
+        lr_power_(lr_power),
+        rows_(rows),
+        row_numel_(row_numel),
+        p_out_(p_out),
+        s_acc_out_(s_acc_out),
+        l_acc_out_(l_acc_out) {}
+
+  inline HOSTDEVICE void operator()(size_t i) {
+    auto j = rows_[i / row_numel_] * row_numel_ + i % row_numel_;
+    const T g = g_[i];
+    const T p = p_[j];
+    const T s_acc = s_acc_[j];
+    const T lr = lr_[0];
+
+    auto new_acc = s_acc + g * g;
+
+    if (lr_power_ == static_cast<T>(-0.5)) {
+      l_acc_out_[j] += g - (std::sqrt(new_acc) - std::sqrt(s_acc)) / lr * p;
+    } else {
+      l_acc_out_[j] +=
+          g - (std::pow(new_acc, -lr_power_) - std::pow(s_acc, -lr_power_)) /
+                  lr * p;
+    }
+
+    auto l_acc = l_acc_out_[j];
+
+    if (std::fabs(l_acc) > l1_) {
+      auto x = -l_acc;
+      if (l_acc >= static_cast<T>(0)) {
+        x += l1_;
+      } else {
+        x -= l1_;
+      }
+
+      auto y = static_cast<T>(2) * l2_;
+      if (lr_power_ == static_cast<T>(-0.5)) {
+        y += std::sqrt(new_acc) / lr;
+      } else {
+        y += std::pow(new_acc, -lr_power_) / lr;
+      }
+
+      auto pre_shrink = x / y;
+      p_out_[j] = pre_shrink;
+    } else {
+      p_out_[j] = static_cast<T>(0);
+    }
+
+    s_acc_out_[j] += g * g;
+  }
+};
+
+template <typename T, typename Context>
+void FTRLOpKernel(const Context& dev_ctx,
+                  const DenseTensor& param,
+                  const DenseTensor& squared_accumulator,
+                  const DenseTensor& linear_accumulator,
+                  const SelectedRows& grad_in,
+                  const DenseTensor& learning_rate,
+                  float l1_in,
+                  float l2_in,
+                  float lr_power_in,
+                  DenseTensor* param_out,
+                  DenseTensor* squared_accum_out,
+                  DenseTensor* linear_accum_out) {
+  auto* lr_in = &learning_rate;
+
+  auto* param_in = &param;
+  auto* sq_accum_in = &squared_accumulator;
+
+  auto* sq_accum_out = squared_accum_out;
+  auto* lin_accum_out = linear_accum_out;
+
+  dev_ctx.template Alloc<T>(param_out);
+  dev_ctx.template Alloc<T>(sq_accum_out);
+  dev_ctx.template Alloc<T>(lin_accum_out);
+
+  auto l1 = static_cast<T>(l1_in) + static_cast<T>(1e-10);
+  auto l2 = static_cast<T>(l2_in) + static_cast<T>(1e-10);
+  auto lr_power = static_cast<T>(lr_power_in);
+
+  auto grad = &grad_in;
+
+  SelectedRows tmp_merged_grad;
+  SelectedRows* merged_grad = &tmp_merged_grad;
+  funcs::scatter::MergeAdd<Context, T> merge_func;
+  merge_func(dev_ctx, *grad, merged_grad);
+
+  auto* merged_rows = merged_grad->mutable_rows();
+  phi::MixVector<int64_t> mixv_merged_rows(merged_rows);
+  const int64_t* rows = mixv_merged_rows.Data(dev_ctx.GetPlace());
+  auto row_numel = static_cast<int64_t>(merged_grad->value().dims()[1]);
+  auto row_height = static_cast<int64_t>(merged_grad->rows().size());
+
+  funcs::ForRange<Context> for_range(static_cast<const Context&>(dev_ctx),
+                                     row_numel * row_height);
+
+  SparseFTRLFunctor<T> functor(merged_grad->value().data<T>(),
+                               param_in->data<T>(),
+                               sq_accum_in->data<T>(),
+                               lr_in->data<T>(),
+                               l1,
+                               l2,
+                               lr_power,
+                               rows,
+                               row_numel,
+                               dev_ctx.template Alloc<T>(param_out),
+                               dev_ctx.template Alloc<T>(sq_accum_out),
+                               dev_ctx.template Alloc<T>(lin_accum_out));
+  for_range(functor);
+}
+
+}  // namespace sr
+}  // namespace phi

paddle/phi/kernels/selected_rows/impl/get_tensor_from_selected_rows_kernel_impl.h

@@ -0,0 +1,32 @@
+// 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
+//
+//     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.
+
+#pragma once
+#include "paddle/phi/backends/gpu/gpu_context.h"
+#include "paddle/phi/core/dense_tensor.h"
+#include "paddle/phi/core/kernel_registry.h"
+#include "paddle/phi/kernels/funcs/selected_rows_functor.h"
+
+namespace phi {
+namespace sr {
+template <typename T, typename Context>
+void GetTensorFromSelectedRowsKernel(const Context& dev_ctx,
+                                     const SelectedRows& x,
+                                     DenseTensor* out) {
+  out->Resize(x.value().dims());
+  dev_ctx.template Alloc<T>(out);
+  phi::Copy(dev_ctx, x.value(), dev_ctx.GetPlace(), false, out);
+}
+}  // namespace sr
+}  // namespace phi

paddle/phi/kernels/selected_rows/impl/isfinite_kernel_impl.h

@@ -0,0 +1,39 @@
+// Copyright (c) 2022 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.
+
+#pragma once
+
+#include "paddle/phi/kernels/funcs/isfinite_functor.h"
+#include "paddle/phi/kernels/selected_rows/isfinite_kernel.h"
+
+namespace phi {
+
+#define DEFINE_ISFINITE_SR(isfinite)                                      \
+  template <typename T, typename Context>                                 \
+  void isfinite##SR(                                                      \
+      const Context& dev_ctx, const SelectedRows& x, SelectedRows* out) { \
+    if (out && out->numel() == 0) {                                       \
+      dev_ctx.template Alloc<bool>(out);                                  \
+      return;                                                             \
+    }                                                                     \
+    dev_ctx.template Alloc<bool>(out);                                    \
+    Isinf##Kernel<T, Context>(dev_ctx, x.value(), out->mutable_value());  \
+  }
+
+DEFINE_ISFINITE_SR(Isinf)
+DEFINE_ISFINITE_SR(Isnan)
+DEFINE_ISFINITE_SR(Isfinite)
+#undef DEFINE_ISFINITE_SR
+
+}  // namespace phi

paddle/phi/kernels/selected_rows/impl/lamb_kernel_impl.h

@@ -0,0 +1,371 @@
+// Copyright (c) 2022 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.
+#pragma once
+#include "glog/logging.h"
+
+#include "paddle/phi/common/data_type.h"
+#include "paddle/phi/core/dense_tensor.h"
+#include "paddle/phi/core/selected_rows.h"
+#include "paddle/phi/kernels/full_kernel.h"
+#include "paddle/phi/kernels/funcs/lamb_functors.h"
+#include "paddle/phi/kernels/funcs/selected_rows_functor.h"
+
+namespace phi {
+namespace sr {
+
+template <typename T, typename MT, typename Context, bool IsMultiPrecision>
+void ComputeRowImpl(const Context& dev_ctx,
+                    const DenseTensor& param,
+                    const SelectedRows& grad,
+                    const DenseTensor& lr,
+                    const DenseTensor& mom1,
+                    const DenseTensor& mom2,
+                    const DenseTensor& beta1_pow,
+                    const DenseTensor& beta2_pow,
+                    const optional<DenseTensor>& master_param_opt,
+                    const optional<DenseTensor>& skip_update_opt,
+                    float weight_decay_f,
+                    float beta1_f,
+                    float beta2_f,
+                    float epsilon_f,
+                    bool always_adapt,
+                    bool multi_precision,
+                    DenseTensor* param_out,
+                    DenseTensor* mom1_out,
+                    DenseTensor* mom2_out,
+                    DenseTensor* beta1_pow_out,
+                    DenseTensor* beta2_pow_out,
+                    DenseTensor* master_param_out);
+
+template <typename T, typename Context>
+void LambKernel(const Context& dev_ctx,
+                const DenseTensor& param,
+                const SelectedRows& grad,
+                const DenseTensor& learning_rate,
+                const DenseTensor& moment1,
+                const DenseTensor& moment2,
+                const DenseTensor& beta1_pow,
+                const DenseTensor& beta2_pow,
+                const optional<DenseTensor>& master_param,
+                const optional<DenseTensor>& skip_update,
+                float weight_decay,
+                float beta1,
+                float beta2,
+                float epsilon,
+                bool always_adapt,
+                bool multi_precision,
+                DenseTensor* param_out,
+                DenseTensor* moment1_out,
+                DenseTensor* moment2_out,
+                DenseTensor* beta1_pow_out,
+                DenseTensor* beta2_pow_out,
+                DenseTensor* master_param_outs) {
+  using MT = typename phi::dtype::MPTypeTrait<T>::Type;
+  if (multi_precision) {
+    ComputeRowImpl<T, MT, Context, true>(dev_ctx,
+                                         param,
+                                         grad,
+                                         learning_rate,
+                                         moment1,
+                                         moment2,
+                                         beta1_pow,
+                                         beta2_pow,
+                                         master_param,
+                                         skip_update,
+                                         weight_decay,
+                                         beta1,
+                                         beta2,
+                                         epsilon,
+                                         always_adapt,
+                                         multi_precision,
+                                         param_out,
+                                         moment1_out,
+                                         moment2_out,
+                                         beta1_pow_out,
+                                         beta2_pow_out,
+                                         master_param_outs);
+  } else {
+    ComputeRowImpl<T, T, Context, false>(dev_ctx,
+                                         param,
+                                         grad,
+                                         learning_rate,
+                                         moment1,
+                                         moment2,
+                                         beta1_pow,
+                                         beta2_pow,
+                                         master_param,
+                                         skip_update,
+                                         weight_decay,
+                                         beta1,
+                                         beta2,
+                                         epsilon,
+                                         always_adapt,
+                                         multi_precision,
+                                         param_out,
+                                         moment1_out,
+                                         moment2_out,
+                                         beta1_pow_out,
+                                         beta2_pow_out,
+                                         master_param_outs);
+  }
+}
+
+template <typename T, typename MT, typename Context, bool IsMultiPrecision>
+void ComputeRowImpl(const Context& dev_ctx,
+                    const DenseTensor& param,
+                    const SelectedRows& grad,
+                    const DenseTensor& lr,
+                    const DenseTensor& mom1,
+                    const DenseTensor& mom2,
+                    const DenseTensor& beta1_pow,
+                    const DenseTensor& beta2_pow,
+                    const optional<DenseTensor>& master_param_opt,
+                    const optional<DenseTensor>& skip_update_opt,
+                    float weight_decay_f,
+                    float beta1_f,
+                    float beta2_f,
+                    float epsilon_f,
+                    bool always_adapt,
+                    bool multi_precision UNUSED,
+                    DenseTensor* param_out,
+                    DenseTensor* mom1_out,
+                    DenseTensor* mom2_out,
+                    DenseTensor* beta1_pow_out,
+                    DenseTensor* beta2_pow_out,
+                    DenseTensor* master_param_out) {
+  if (!IsMultiPrecision) {
+    constexpr auto kIsSameType = std::is_same<T, MT>::value;
+    PADDLE_ENFORCE_EQ(
+        kIsSameType,
+        true,
+        common::errors::InvalidArgument(
+            "When multi_precision=False, T and MT must be the same type."));
+  }
+
+  const auto* master_param =
+      IsMultiPrecision ? master_param_opt.get_ptr() : nullptr;
+  const auto* skip_update = skip_update_opt.get_ptr();
+  const bool* skip_update_flag = skip_update && skip_update->IsInitialized()
+                                     ? skip_update->data<bool>()
+                                     : nullptr;
+  if (skip_update_flag &&
+      skip_update->place().GetType() == AllocationType::CPU &&
+      (*skip_update_flag)) {
+    return;
+  }
+
+  auto weight_decay = static_cast<MT>(weight_decay_f);
+  auto beta1 = static_cast<MT>(beta1_f);
+  auto beta2 = static_cast<MT>(beta2_f);
+  auto epsilon = static_cast<MT>(epsilon_f);
+  auto numel = param.numel();
+  funcs::ForRange<Context> for_range(dev_ctx, numel);
+  DenseTensor trust_ratio_div;
+  trust_ratio_div.Resize(param.dims());
+  /*auto trust_ratio_div =
+      dev_ctx.AllocateTmpTensor<MT, DeviceContext>(param.dims(), dev_ctx);*/
+  auto* trust_ratio_div_ptr = dev_ctx.template Alloc<MT>(&trust_ratio_div);
+
+  const void* param_ptr = param.data();
+  const void* master_param_ptr = master_param ? master_param->data() : nullptr;
+  void* param_out_ptr = dev_ctx.template Alloc<T>(param_out);
+  void* master_param_out_ptr =
+      master_param_out ? dev_ctx.template Alloc<MT>(master_param_out) : nullptr;
+  // Update moments
+  bool should_update_beta_pow_later = false;
+  const MT *beta1_pow_ptr = nullptr, *beta2_pow_ptr = nullptr;
+  MT *beta1_pow_out_ptr = nullptr, *beta2_pow_out_ptr = nullptr;
+  VLOG(10) << "Beta1Pow place: " << beta1_pow.place()
+           << " , Beta2Pow place: " << beta2_pow.place();
+  // Diff from here
+  PADDLE_ENFORCE_EQ(IsMultiPrecision,
+                    false,
+                    common::errors::Unimplemented(
+                        "SelectedRows gradient is not supported when "
+                        "multi_precision=True."));
+  constexpr bool kIsSameType = std::is_same<T, MT>::value;
+  PADDLE_ENFORCE_EQ(kIsSameType,
+                    true,
+                    common::errors::Unimplemented(
+                        "SelectedRows gradient is not supported when "
+                        "multi_precision=True."));
+  if (grad.rows().size() == 0) {
+    VLOG(3) << "grad row size is 0!!";
+    return;
+  }
+
+  std::vector<int64_t> cpu_rows(grad.rows().begin(), grad.rows().end());
+  bool is_strict_sorted = true;
+  for (size_t i = 1; i < cpu_rows.size(); ++i) {
+    if (cpu_rows[i - 1] >= cpu_rows[i]) {
+      is_strict_sorted = false;
+      break;
+    }
+  }
+
+  SelectedRows tmp_grad_merge;
+  const SelectedRows* grad_merge_ptr;
+  if (is_strict_sorted) {
+    grad_merge_ptr = &grad;
+  } else {
+    // merge duplicated rows if any.
+    // The rows of grad_merge have been sorted inside MergeAdd functor
+    funcs::scatter::MergeAdd<Context, T> merge_func;
+    merge_func(dev_ctx, grad, &tmp_grad_merge, true);
+    grad_merge_ptr = &tmp_grad_merge;
+  }
+
+  auto& grad_merge = *grad_merge_ptr;
+  auto& grad_tensor = grad_merge.value();
+  const T* grad_data = grad_tensor.template data<T>();
+  auto* grad_merge_rows = &grad_merge.rows();
+  phi::MixVector<int64_t> mixv_grad_merge_rows(grad_merge_rows);
+  const int64_t* rows = mixv_grad_merge_rows.Data(dev_ctx.GetPlace());
+  auto row_numel = grad_tensor.numel() / grad_merge.rows().size();
+  if (dev_ctx.GetPlace().GetType() == AllocationType::GPU &&
+      beta1_pow.place() == phi::CPUPlace() &&
+      beta2_pow.place() == phi::CPUPlace()) {
+    SparseLambMomentREGUpdateFunctor<T> moment_update_functor(
+        static_cast<T>(weight_decay),
+        static_cast<T>(beta1),
+        static_cast<T>(beta2),
+        static_cast<T>(epsilon),
+        *beta1_pow.template data<T>(),
+        *beta2_pow.template data<T>(),
+        mom1.template data<T>(),
+        dev_ctx.template Alloc<T>(mom1_out),
+        mom2.template data<T>(),
+        dev_ctx.template Alloc<T>(mom2_out),
+        grad_data,
+        param.template data<T>(),
+        trust_ratio_div.template data<T>(),
+        rows,
+        row_numel,
+        grad_merge.rows().size(),
+        skip_update_flag);
+    for_range(moment_update_functor);
+    T* beta1_pow_out_data = dev_ctx.template HostAlloc<T>(beta1_pow_out);
+    beta1_pow_out_data[0] =
+        static_cast<T>(beta1) * beta1_pow.template data<T>()[0];
+    T* beta2_pow_out_data = dev_ctx.template HostAlloc<T>(beta2_pow_out);
+    beta2_pow_out_data[0] =
+        static_cast<T>(beta2) * beta2_pow.template data<T>()[0];
+  } else {
+    beta1_pow_ptr = beta1_pow.template data<MT>();
+    beta2_pow_ptr = beta2_pow.template data<MT>();
+    beta1_pow_out_ptr = dev_ctx.template Alloc<MT>(beta1_pow_out);
+    beta2_pow_out_ptr = dev_ctx.template Alloc<MT>(beta2_pow_out);
+    should_update_beta_pow_later = true;
+    SparseLambMomentMENUpdateFunctor<T> moment_update_functor(
+        static_cast<T>(weight_decay),
+        static_cast<T>(beta1),
+        static_cast<T>(beta2),
+        static_cast<T>(epsilon),
+        reinterpret_cast<const T*>(beta1_pow_ptr),
+        reinterpret_cast<const T*>(beta2_pow_ptr),
+        mom1.template data<T>(),
+        dev_ctx.template Alloc<T>(mom1_out),
+        mom2.template data<T>(),
+        dev_ctx.template Alloc<T>(mom2_out),
+        grad_data,
+        param.template data<T>(),
+        trust_ratio_div.template data<T>(),
+        rows,
+        row_numel,
+        grad_merge.rows().size(),
+        skip_update_flag);
+    for_range(moment_update_functor);
+  }
+  // Same from here
+  // Update parameter
+  // The code in the following part is exactly the same as that in
+  // paddle/phi/kernels/impl/lamb_kernel_impl.h Please modify it together
+
+  // DenseTensor p_norm_t;
+  // p_norm_t.Resize({1});
+  // auto* p_norm_ptr = dev_ctx.template Alloc<MT>(&p_norm_t);
+
+  // DenseTensor trust_ratio_div_norm_t;
+  // trust_ratio_div_norm_t.Resize({1});
+  // auto* trust_ratio_div_norm_ptr =
+  //     dev_ctx.template Alloc<MT>(&trust_ratio_div_norm_t);
+
+  DenseTensor p_norm_t;
+  DataType dtype = phi::CppTypeToDataType<MT>::Type();
+  FullKernel<MT, Context>(
+      dev_ctx, std::vector<int64_t>({1}), 0, dtype, &p_norm_t);
+  auto* p_norm_ptr = p_norm_t.data<MT>();
+
+  DenseTensor trust_ratio_div_norm_t;
+  FullKernel<MT, Context>(
+      dev_ctx, std::vector<int64_t>({1}), 0, dtype, &trust_ratio_div_norm_t);
+  auto* trust_ratio_div_norm_ptr = trust_ratio_div_norm_t.data<MT>();
+
+  // TODO(zengjinle): remove the following Eigen operations when
+  // *skip_update == true.
+  if (weight_decay > static_cast<MT>(0) || always_adapt) {
+    memory_utils::Buffer buffer(dev_ctx.GetPlace());
+    funcs::SquaredL2Norm(dev_ctx,
+                         reinterpret_cast<const MT*>(
+                             IsMultiPrecision ? master_param_ptr : param_ptr),
+                         p_norm_ptr,
+                         numel,
+                         &buffer);
+    funcs::SquaredL2Norm(
+        dev_ctx, trust_ratio_div_ptr, trust_ratio_div_norm_ptr, numel, &buffer);
+  }
+
+  if (VLOG_IS_ON(1)) {
+    const auto& name = "Param";
+    auto pn = funcs::ToVector(p_norm_ptr, 1, dev_ctx.GetPlace());
+    auto tn = funcs::ToVector(trust_ratio_div_norm_ptr, 1, dev_ctx.GetPlace());
+    auto dtype = DataTypeToString(CppTypeToDataType<T>::Type());
+    VLOG(1) << "Param " << dtype << " " << name << " pn = " << pn[0]
+            << " , tn = " << tn[0];
+  }
+
+#define CALL_PADDLE_UPDATE_LAMB_PARAM_FUNC(__should_update_beta_pow)          \
+  do {                                                                        \
+    LambParamUpdateFunctor<T, MT, IsMultiPrecision, __should_update_beta_pow> \
+        param_update_functor(lr.template data<MT>(),                          \
+                             static_cast<const T*>(param_ptr),                \
+                             static_cast<const MT*>(master_param_ptr),        \
+                             p_norm_ptr,                                      \
+                             trust_ratio_div_ptr,                             \
+                             trust_ratio_div_norm_ptr,                        \
+                             static_cast<T*>(param_out_ptr),                  \
+                             static_cast<MT*>(master_param_out_ptr),          \
+                             skip_update_flag);                               \
+    if (__should_update_beta_pow) {                                           \
+      param_update_functor.SetBetaPows(beta1_pow_ptr,                         \
+                                       beta2_pow_ptr,                         \
+                                       beta1_pow_out_ptr,                     \
+                                       beta2_pow_out_ptr,                     \
+                                       beta1,                                 \
+                                       beta2);                                \
+    }                                                                         \
+    for_range(param_update_functor);                                          \
+  } while (0)
+
+  if (should_update_beta_pow_later) {
+    CALL_PADDLE_UPDATE_LAMB_PARAM_FUNC(true);
+  } else {
+    CALL_PADDLE_UPDATE_LAMB_PARAM_FUNC(false);
+  }
+
+#undef CALL_PADDLE_UPDATE_LAMB_PARAM_FUNC
+}
+
+}  // namespace sr
+}  // namespace phi

paddle/phi/kernels/selected_rows/impl/load_kernel_impl.h

@@ -0,0 +1,47 @@
+// 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
+//
+//     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.
+
+#pragma once
+
+#include <string>
+
+#include "paddle/phi/core/framework/selected_rows_serialize.h"
+#include "paddle/phi/core/kernel_registry.h"
+#include "paddle/phi/core/tensor_utils.h"
+
+namespace phi::sr {
+
+template <typename T, typename Context>
+void LoadSelectedRowsKernel(const Context& dev_ctx,
+                            const std::string& file_path,
+                            int64_t seek,
+                            const std::vector<int64_t>& shape,
+                            bool load_as_fp16,
+                            SelectedRows* out) {
+  // FIXME(yuyang18): We save variable to local file now, but we should change
+  // it to save an output stream.
+  std::ifstream fin(file_path, std::ios::binary);
+  PADDLE_ENFORCE_EQ(
+      static_cast<bool>(fin),
+      true,
+      errors::Unavailable("Load operator fail to open file %s, please check "
+                          "whether the model file is complete or damaged.",
+                          file_path));
+  PADDLE_ENFORCE_NOT_NULL(
+      out,
+      errors::InvalidArgument("The variable to be loaded cannot be found."));
+
+  phi::DeserializeFromStream(fin, out, dev_ctx);
+}
+}  // namespace phi::sr

paddle/phi/kernels/selected_rows/impl/save_kernel_impl.h

@@ -0,0 +1,63 @@
+// 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
+//
+//     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.
+
+#pragma once
+
+#include <stdint.h>
+
+#include <fstream>
+#include <numeric>
+#include <string>
+#include <vector>
+
+#include "paddle/phi/core/framework/dense_tensor_serialize.h"
+#include "paddle/phi/core/framework/selected_rows_serialize.h"
+#include "paddle/phi/core/framework/var_type_helper.h"
+#include "paddle/phi/core/kernel_registry.h"
+#include "paddle/phi/core/tensor_utils.h"
+#include "paddle/phi/kernels/cast_kernel.h"
+
+namespace phi {
+
+template <typename T, typename Context>
+void SaveSelectedRowsKernel(const Context& dev_ctx,
+                            const SelectedRows& x,
+                            const std::string& file_path,
+                            bool overwrite,
+                            bool save_as_fp16) {
+  PADDLE_ENFORCE_EQ(
+      FileExists(file_path) && !overwrite,
+      false,
+      common::errors::PreconditionNotMet(
+          "%s exists!, cannot save to it when overwrite is set to false.",
+          file_path));
+  PADDLE_ENFORCE_EQ(save_as_fp16,
+                    false,
+                    common::errors::Unimplemented(
+                        "SelectedRows is not supported to save as float16."));
+
+  MkDirRecursively(DirName(file_path).c_str());
+
+  // FIXME(yuyang18): We save variable to local file now, but we should change
+  // it to save an output stream.
+  std::ofstream fout(file_path, std::ios::binary);
+  PADDLE_ENFORCE_EQ(static_cast<bool>(fout),
+                    true,
+                    common::errors::Unavailable(
+                        "Cannot open %s to save variables.", file_path));
+  phi::SerializeToStream(fout, x, dev_ctx);
+  fout.close();
+}
+
+}  // namespace phi

paddle/phi/kernels/selected_rows/impl/share_data_kernel_impl.h

@@ -0,0 +1,29 @@
+// 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
+//
+//     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.
+
+#pragma once
+#include "paddle/phi/core/kernel_registry.h"
+
+namespace phi {
+namespace sr {
+template <typename T, typename Context>
+void ShareDataKernel(const Context &dev_ctx,
+                     const SelectedRows &x,
+                     SelectedRows *out) {
+  out->set_rows(x.rows());
+  out->set_height(x.height());
+  out->mutable_value()->ShareDataWith(x.value());
+}
+}  // namespace sr
+}  // namespace phi