[Fluid] Move distributed_fused_lamb_init to phi (#55993)

paddle/fluid/operators/optimizers/distributed_fused_lamb_init_op.cc

@@ -12,7 +12,8 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
-#include "paddle/fluid/operators/optimizers/distributed_fused_lamb_init_op.h"
+#include "paddle/fluid/framework/op_registry.h"
+#include "paddle/fluid/framework/operator.h"
 
 namespace paddle {
 namespace operators {
@@ -116,9 +117,3 @@ namespace ops = paddle::operators;
 REGISTER_OP_WITHOUT_GRADIENT(distributed_fused_lamb_init,
                              ops::DistributedFusedLambInitOp,
                              ops::DistributedFusedLambInitOpMaker);
-
-PD_REGISTER_STRUCT_KERNEL(distributed_fused_lamb_init,
-                          CPU,
-                          ALL_LAYOUT,
-                          ops::DistributedFusedLambInitOpKernel,
-                          float) {}

paddle/phi/kernels/distributed_fused_lamb_init_kernel.h

+// 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.
+
+#pragma once
+
+#include "paddle/phi/common/scalar.h"
+#include "paddle/phi/core/dense_tensor.h"
+
+namespace phi {
+
+template <typename T, typename Context>
+void DistributedFusedLambInitOpKernel(
+    const Context& dev_ctx,
+    const std::vector<const DenseTensor*>& param,
+    const std::vector<const DenseTensor*>& grad,
+    float beta1,
+    float beta2,
+    const std::vector<int>& apply_weight_decay,
+    int alignment,
+    int rank,
+    int nranks,
+    DenseTensor* fp32_fused_param,
+    DenseTensor* fp32_fused_grad,
+    DenseTensor* fp16_fused_param,
+    DenseTensor* fp16_fused_grad,
+    DenseTensor* moment1,
+    DenseTensor* moment2,
+    DenseTensor* beta1_pow,
+    DenseTensor* beta2_pow,
+    DenseTensor* fused_param_offsets,
+    DenseTensor* fp32_shard_fused_param_offsets,
+    DenseTensor* fp16_shard_fused_param_offsets,
+    DenseTensor* param_info,
+    DenseTensor* param_order,
+    std::vector<DenseTensor*> param_out,
+    std::vector<DenseTensor*> master_param_out,
+    std::vector<DenseTensor*> grad_out,
+    DenseTensor* global_scale,
+    DenseTensor* step);
+
+}  // namespace phi

paddle/phi/kernels/fusion/cpu/distributed_fused_lamb_init_kernel.cc

+// 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/kernels/distributed_fused_lamb_init_kernel.h"
+#include "paddle/phi/core/errors.h"
+#include "paddle/phi/core/kernel_registry.h"
+
+namespace phi {
+namespace fusion {
+
+template <typename T, typename Context>
+void DistributedFusedLambInitOpKernel(
+    const Context& dev_ctx,
+    const std::vector<const DenseTensor*>& param,
+    const std::vector<const DenseTensor*>& grad,
+    float beta1,
+    float beta2,
+    const std::vector<int>& apply_weight_decay,
+    int alignment,
+    int rank,
+    int nranks,
+    DenseTensor* fp32_fused_param,
+    DenseTensor* fp32_fused_grad,
+    DenseTensor* fp16_fused_param,
+    DenseTensor* fp16_fused_grad,
+    DenseTensor* moment1,
+    DenseTensor* moment2,
+    DenseTensor* beta1_pow,
+    DenseTensor* beta2_pow,
+    DenseTensor* fused_param_offsets,
+    DenseTensor* fp32_shard_fused_param_offsets,
+    DenseTensor* fp16_shard_fused_param_offsets,
+    DenseTensor* param_info,
+    DenseTensor* param_order,
+    std::vector<DenseTensor*> param_out,
+    std::vector<DenseTensor*> master_param_out,
+    std::vector<DenseTensor*> grad_out,
+    DenseTensor* global_scale,
+    DenseTensor* step) {
+  PADDLE_THROW(phi::errors::Unavailable(
+      "Do not support expert count op for cpu kernel now."));
+}
+}  // namespace fusion
+}  // namespace phi
+
+PD_REGISTER_KERNEL(distributed_fused_lamb_init,
+                   CPU,
+                   ALL_LAYOUT,
+                   phi::fusion::DistributedFusedLambInitOpKernel,
+                   float) {
+  kernel->OutputAt(0).SetDataType(phi::DataType::FLOAT32);
+  kernel->OutputAt(1).SetDataType(phi::DataType::FLOAT32);
+  kernel->OutputAt(2).SetDataType(phi::DataType::FLOAT16);
+  kernel->OutputAt(3).SetDataType(phi::DataType::FLOAT16);
+  kernel->OutputAt(4).SetDataType(phi::DataType::FLOAT32);
+  kernel->OutputAt(5).SetDataType(phi::DataType::FLOAT32);
+  kernel->OutputAt(6).SetDataType(phi::DataType::FLOAT32);
+  kernel->OutputAt(7).SetDataType(phi::DataType::FLOAT32);
+  kernel->OutputAt(8).SetDataType(phi::DataType::INT32);
+  kernel->OutputAt(9).SetDataType(phi::DataType::INT32);
+  kernel->OutputAt(10).SetDataType(phi::DataType::INT32);
+  kernel->OutputAt(11).SetDataType(phi::DataType::INT32);
+  kernel->OutputAt(12).SetDataType(phi::DataType::INT32);
+  kernel->OutputAt(13).SetDataType(kernel_key.dtype());
+  kernel->OutputAt(14).SetDataType(phi::DataType::FLOAT32);
+  kernel->OutputAt(15).SetDataType(kernel_key.dtype());
+  kernel->OutputAt(16).SetDataType(phi::DataType::FLOAT32);
+  kernel->OutputAt(17).SetDataType(phi::DataType::INT64);
+}

paddle/fluid/operators/optimizers/cast_with_ptr.h → paddle/phi/kernels/fusion/gpu/cast_with_ptr.h

@@ -14,28 +14,24 @@
 
 #pragma once
 
-#include "paddle/fluid/platform/device/gpu/gpu_launch_config.h"
-#include "paddle/fluid/platform/device_context.h"
-#include "paddle/fluid/platform/enforce.h"
 #include "paddle/phi/api/include/tensor.h"
+#include "paddle/phi/backends/gpu/gpu_launch_config.h"
 #include "paddle/phi/core/ddim.h"
+#include "paddle/phi/core/enforce.h"
 #include "paddle/phi/kernels/funcs/elementwise_base.h"
 
-namespace paddle {
-namespace operators {
-namespace details {
+namespace phi {
 
 template <typename InT, typename OutT>
 struct CastFunctor {
   HOSTDEVICE OutT operator()(InT x) const { return static_cast<OutT>(x); }
 };
-
 template <typename InT, typename OutT, int VecSize>
 static void VecCastKernel(const phi::GPUContext &ctx,
                           const InT *x,
                           OutT *y,
                           size_t n) {
-  auto config = platform::GetGpuLaunchConfig1D(ctx, n, VecSize);
+  auto config = phi::backends::gpu::GetGpuLaunchConfig1D(ctx, n, VecSize);
   auto block = config.GetGridSize();
   auto thread = config.GetBlockSize();
   auto main_offset = n / (VecSize * thread) * VecSize * thread;
@@ -50,8 +46,6 @@ static void VecCastKernel(const phi::GPUContext &ctx,
           in_arr, out_arr, n, main_offset, VecSize, FunctorT());
 }
 
-}  // namespace details
-
 template <typename InT, typename OutT>
 static void LaunchCastKernel(const phi::GPUContext &ctx,
                              const InT *x,
@@ -61,20 +55,19 @@ static void LaunchCastKernel(const phi::GPUContext &ctx,
   PADDLE_ENFORCE_NE(
       static_cast<const void *>(x),
       static_cast<void *>(y),
-      platform::errors::InvalidArgument("Inplace cast is not supported yet."));
+      errors::InvalidArgument("Inplace cast is not supported yet."));
   int vec_size = std::min(phi::GetVectorizedSize(x), phi::GetVectorizedSize(y));
   switch (vec_size) {
     case 4:
-      return details::VecCastKernel<InT, OutT, 4>(ctx, x, y, n);
+      return VecCastKernel<InT, OutT, 4>(ctx, x, y, n);
     case 2:
-      return details::VecCastKernel<InT, OutT, 2>(ctx, x, y, n);
+      return VecCastKernel<InT, OutT, 2>(ctx, x, y, n);
     case 1:
-      return details::VecCastKernel<InT, OutT, 1>(ctx, x, y, n);
+      return VecCastKernel<InT, OutT, 1>(ctx, x, y, n);
     default:
-      PADDLE_THROW(platform::errors::InvalidArgument(
-          "The vectorized size must be 1, 2 or 4."));
+      PADDLE_THROW(
+          errors::InvalidArgument("The vectorized size must be 1, 2 or 4."));
   }
 }
 
-}  // namespace operators
-}  // namespace paddle
+}  // namespace phi

paddle/fluid/operators/optimizers/distributed_fused_lamb_init_op.cu → paddle/phi/kernels/fusion/gpu/distributed_fused_lamb_init_kernel.cu

-// Copyright (c) 2021 PaddlePaddle Authors. All Rights Reserved.
+// 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.
@@ -12,24 +12,24 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
-#include "paddle/fluid/operators/optimizers/distributed_fused_lamb_init_op.h"
-#include "paddle/fluid/memory/memcpy.h"
-#include "paddle/fluid/operators/optimizers/cast_with_ptr.h"
-#include "paddle/fluid/platform/device/gpu/gpu_launch_config.h"
+#include "paddle/phi/kernels/distributed_fused_lamb_init_kernel.h"
 #include "paddle/phi/common/data_type.h"
+#include "paddle/phi/core/enforce.h"
+#include "paddle/phi/core/kernel_registry.h"
 #include "paddle/phi/kernels/funcs/algorithm.h"
 #include "paddle/phi/kernels/funcs/math_function.h"
 #include "paddle/phi/kernels/funcs/tensor_to_string.h"
+#include "paddle/phi/kernels/fusion/gpu/cast_with_ptr.h"
 
-namespace paddle {
-namespace operators {
+namespace phi {
+namespace fusion {
 
 using phi::funcs::FlattenToString;
 using phi::funcs::ToVector;
 
 struct ParamGradInfo {
-  phi::DenseTensor *param_t{nullptr};
-  phi::DenseTensor *grad_t{nullptr};
+  DenseTensor *param_t{nullptr};
+  DenseTensor *grad_t{nullptr};
   size_t idx{0};
   size_t numel{0};
   size_t numel_with_padding{0};
@@ -82,20 +82,17 @@ static void GetParamGradShardInfo(const std::vector<ParamGradInfo> &infos,
                                   size_t *start_numel_offset,
                                   size_t *end_numel_offset) {
   VLOG(10) << "NumelOffset: "
-           << string::join_strings(infos, ",", [](const ParamGradInfo &info) {
+           << paddle::string::join_strings(
+                  infos, ",", [](const ParamGradInfo &info) {
                     return info.numel_offset;
                   });
   VLOG(10) << "start_size = " << start_size << " , end_size = " << end_size;
 
   if (infos.empty()) {
     PADDLE_ENFORCE_EQ(
-        start_size,
-        0,
-        platform::errors::InvalidArgument("start_size should be 0."));
+        start_size, 0, errors::InvalidArgument("start_size should be 0."));
     PADDLE_ENFORCE_EQ(
-        end_size,
-        0,
-        platform::errors::InvalidArgument("end_size should be 0."));
+        end_size, 0, errors::InvalidArgument("end_size should be 0."));
     *start_idx = 0;
     *end_idx = 0;
     *start_numel_offset = 0;
@@ -103,10 +100,10 @@ static void GetParamGradShardInfo(const std::vector<ParamGradInfo> &infos,
     return;
   }
 
-  PADDLE_ENFORCE_LT(start_size,
+  PADDLE_ENFORCE_LT(
+      start_size,
       end_size,
-                    platform::errors::InvalidArgument(
-                        "start_size should be less than end_size."));
+      errors::InvalidArgument("start_size should be less than end_size."));
   size_t n = infos.size();
   ParamGradInfoNumelOffsetCompFunctor comp;
   auto i = static_cast<size_t>(
@@ -116,7 +113,7 @@ static void GetParamGradShardInfo(const std::vector<ParamGradInfo> &infos,
     PADDLE_ENFORCE_GT(
         i,
         0,
-        platform::errors::InvalidArgument(
+        errors::InvalidArgument(
             "Cannot find suitable sharding which is between [%d, %d)",
             start_size,
             end_size));
@@ -125,7 +122,7 @@ static void GetParamGradShardInfo(const std::vector<ParamGradInfo> &infos,
   PADDLE_ENFORCE_LT(
       i,
       n,
-      platform::errors::InvalidArgument(
+      errors::InvalidArgument(
           "Cannot find suitable sharding which is between [%d, %d)",
           start_size,
           end_size));
@@ -136,10 +133,10 @@ static void GetParamGradShardInfo(const std::vector<ParamGradInfo> &infos,
       infos.begin());
   *end_idx = j - 1;
   *end_numel_offset = end_size - infos[j - 1].numel_offset;
-  PADDLE_ENFORCE_GT(*end_numel_offset,
+  PADDLE_ENFORCE_GT(
+      *end_numel_offset,
       0,
-                    platform::errors::InvalidArgument(
-                        "Internal error when sharding, this may be a bug "
+      errors::InvalidArgument("Internal error when sharding, this may be a bug "
                               "caused by empty parameter."));
   VLOG(10) << "Sharding [start_size=" << start_size << ", end_size=" << end_size
            << "): " << (*start_idx) << ":" << (*start_numel_offset) << " -> "
@@ -154,7 +151,7 @@ static size_t FillAlignmentPaddingInfo(std::vector<ParamGradInfo> *infos,
   PADDLE_ENFORCE_EQ(
       alignment % sizeof_dtype,
       0,
-      platform::errors::InvalidArgument(
+      errors::InvalidArgument(
           "The attr(alignment) should be exactly divided by sizeof(T) %d.",
           sizeof_dtype));
   alignment /= sizeof_dtype;
@@ -182,41 +179,41 @@ static size_t FillAlignmentPaddingInfo(std::vector<ParamGradInfo> *infos,
 
 template <typename T>
 static T *TensorFillConstant(const phi::GPUContext &dev_ctx,
-                             phi::DenseTensor *tensor,
-                             const framework::DDim &dims,
+                             DenseTensor *tensor,
+                             const DDim &dims,
                              T value) {
   tensor->Resize(dims);
-  auto *ptr = tensor->mutable_data<T>(dev_ctx.GetPlace());
+  auto *ptr = dev_ctx.template Alloc<T>(tensor);
   phi::funcs::SetConstant<phi::GPUContext, T> set_constant;
   set_constant(dev_ctx, tensor, value);
   return ptr;
 }
 
-static phi::DenseTensor CastDataForInitedTensor(const phi::GPUContext &dev_ctx,
-                                                phi::DenseTensor *origin,
-                                                phi::DenseTensor *fused_out,
+static DenseTensor CastDataForInitedTensor(const phi::GPUContext &dev_ctx,
+                                           DenseTensor *origin,
+                                           DenseTensor *fused_out,
                                            size_t numel_offset) {
-  PADDLE_ENFORCE_EQ(origin->IsInitialized(),
+  PADDLE_ENFORCE_EQ(
+      origin->IsInitialized(),
       true,
-                    platform::errors::InvalidArgument(
-                        "The tensor to be cast should be initialized."));
+      errors::InvalidArgument("The tensor to be cast should be initialized."));
 
   PADDLE_ENFORCE_EQ(fused_out->dtype(),
                     phi::DataType::FLOAT32,
-                    platform::errors::InvalidArgument(
+                    errors::InvalidArgument(
                         "The dst tensor to be cast should be FP32 tensor."));
   PADDLE_ENFORCE_EQ(origin->dtype(),
                     phi::DataType::FLOAT16,
-                    platform::errors::InvalidArgument(
+                    errors::InvalidArgument(
                         "The src tensor to be cast should be FP16 tensor."));
   auto *dst = fused_out->data<float>() + numel_offset;
-  auto *src = origin->data<platform::float16>();
+  auto *src = origin->data<dtype::float16>();
   auto numel = origin->numel();
   LaunchCastKernel(dev_ctx, src, dst, numel);
   VLOG(10) << "Cast from FP32 -> FP16, range: [" << numel_offset << ", "
            << numel_offset + numel << ")"
            << " , total: [0, " << fused_out->numel() << ")";
-  framework::DDim fused_out_dim = fused_out->dims();
+  DDim fused_out_dim = fused_out->dims();
   auto fused_out_numel = fused_out->numel();
   fused_out->Resize({fused_out_numel});
   auto sliced_tensor = fused_out->Slice(numel_offset, numel + numel_offset);
@@ -224,45 +221,40 @@ static phi::DenseTensor CastDataForInitedTensor(const phi::GPUContext &dev_ctx,
   return sliced_tensor;
 }
 
-static phi::DenseTensor CopyAndShareBufferForInitedTensor(
-    phi::DenseTensor *origin,
-    phi::DenseTensor *fused_out,
-    size_t numel_offset,
-    gpuStream_t stream) {
+static DenseTensor CopyAndShareBufferForInitedTensor(
+    const phi::GPUContext &dev_ctx,
+    DenseTensor *origin,
+    DenseTensor *fused_out,
+    size_t numel_offset) {
   PADDLE_ENFORCE_EQ(
       origin->IsInitialized(),
       true,
-      platform::errors::InvalidArgument(
+      errors::InvalidArgument(
           "The tensor to be copied and shared data should be initialized."));
   auto dtype = fused_out->type();
   PADDLE_ENFORCE_EQ(origin->type(),
                     dtype,
-                    platform::errors::InvalidArgument(
+                    errors::InvalidArgument(
                         "The tensor to be copied and shared data should be "
                         "have the same data type."));
   auto place = fused_out->place();
   PADDLE_ENFORCE_EQ(
       origin->place(),
       place,
-      platform::errors::InvalidArgument("The tensor to be copied and shared "
+      errors::InvalidArgument("The tensor to be copied and shared "
                               "data should be have the same place."));
   PADDLE_ENFORCE_EQ(
-      platform::is_gpu_place(place),
+      dev_ctx.GetPlace().GetType() == phi::AllocationType::GPU,
       true,
-      platform::errors::InvalidArgument(
+      errors::InvalidArgument(
           "The tensor to be copied and shared data should be on GPU place."));
 
   auto numel = origin->numel();
-  framework::DDim fused_out_dim = fused_out->dims();
+  DDim fused_out_dim = fused_out->dims();
   auto fused_out_numel = fused_out->numel();
   auto sliced_tensor = fused_out->Resize({fused_out_numel})
                            .Slice(numel_offset, numel + numel_offset);
-  memory::Copy(place,
-               sliced_tensor.data(),
-               place,
-               origin->data(),
-               numel * phi::SizeOf(dtype),
-               stream);
+  phi::Copy(dev_ctx, *origin, dev_ctx.GetPlace(), false, &sliced_tensor);
   origin->ShareBufferWith(sliced_tensor);
   fused_out->Resize(fused_out_dim);
   VLOG(10) << "Copy and share buffer, range: [" << numel_offset << ", "
@@ -271,17 +263,17 @@ static phi::DenseTensor CopyAndShareBufferForInitedTensor(
   return sliced_tensor;
 }
 
-static void ShareBufferForNonInitedTensor(phi::DenseTensor *origin,
-                                          phi::DenseTensor *fused_out,
+static void ShareBufferForNonInitedTensor(DenseTensor *origin,
+                                          DenseTensor *fused_out,
                                           size_t numel_offset,
-                                          const framework::DDim &dims) {
+                                          const DDim &dims) {
   PADDLE_ENFORCE_EQ(
       origin->IsInitialized(),
       false,
-      platform::errors::InvalidArgument(
+      errors::InvalidArgument(
           "The tensor to be shared data should not be initialized."));
 
-  framework::DDim fused_out_dim = fused_out->dims();
+  DDim fused_out_dim = fused_out->dims();
   auto fused_out_numel = fused_out->numel();
   auto numel = phi::product(dims);
   *origin = fused_out->Resize({fused_out_numel})
@@ -294,10 +286,11 @@ static void ShareBufferForNonInitedTensor(phi::DenseTensor *origin,
 }
 
 template <typename T>
-static void CopyVectorToCPUTensor(const std::vector<T> &src,
-                                  phi::DenseTensor *dst) {
+static void CopyVectorToCPUTensor(const phi::GPUContext &dev_ctx,
+                                  const std::vector<T> &src,
+                                  DenseTensor *dst) {
   dst->Resize({static_cast<int64_t>(src.size())});
-  T *dst_ptr = dst->mutable_data<T>(platform::CPUPlace());
+  T *dst_ptr = dev_ctx.template HostAlloc<T>(dst);
   const T *src_ptr = src.data();
   auto nbytes = src.size() * sizeof(T);
   std::memcpy(dst_ptr, src_ptr, nbytes);
@@ -339,98 +332,115 @@ static T ClipByBound(T x, T low_value, T high_value) {
   return x;
 }
 
-template <typename T>
-class DistributedFusedLambInitOpKernel<T, phi::GPUContext>
-    : public framework::OpKernel<T> {
- public:
-  void Compute(const framework::ExecutionContext &ctx) const override {
+template <typename T, typename Context>
+void DistributedFusedLambInitOpKernel(
+    const Context &dev_ctx,
+    const std::vector<const DenseTensor *> &param,
+    const std::vector<const DenseTensor *> &grad,
+    float beta1,
+    float beta2,
+    const std::vector<int> &apply_weight_decay,
+    int alignment,
+    int rank,
+    int nranks,
+    DenseTensor *fp32_fused_param,
+    DenseTensor *fp32_fused_grad,
+    DenseTensor *fp16_fused_param,
+    DenseTensor *fp16_fused_grad,
+    DenseTensor *moment1,
+    DenseTensor *moment2,
+    DenseTensor *beta1_pow,
+    DenseTensor *beta2_pow,
+    DenseTensor *fused_param_offsets,
+    DenseTensor *fp32_shard_fused_param_offsets,
+    DenseTensor *fp16_shard_fused_param_offsets,
+    DenseTensor *param_info,
+    DenseTensor *param_order,
+    std::vector<DenseTensor *> param_out,
+    std::vector<DenseTensor *> master_param_out,
+    std::vector<DenseTensor *> grad_out,
+    DenseTensor *global_scale,
+    DenseTensor *step) {
   VLOG(10) << "starts to run DistributedFusedLambInitOp";
-    auto &dev_ctx = ctx.template device_context<phi::GPUContext>();
-    auto place = ctx.GetPlace();
+  auto place = dev_ctx.GetPlace();
   auto stream = dev_ctx.stream();
 
   // Step 1: Check Input(Param) and Output(ParamOut), Input(Grad) and
   // Output(GradOut)
-    auto params = ctx.MultiInput<phi::DenseTensor>("Param");
-    auto grads = ctx.MultiInput<phi::DenseTensor>("Grad");
-    auto master_params = ctx.MultiOutput<phi::DenseTensor>("MasterParamOut");
   std::vector<ParamGradInfo> fp32_infos, fp16_infos;
   {
-      PADDLE_ENFORCE_EQ(params.size(),
-                        grads.size(),
-                        platform::errors::InvalidArgument(
-                            "The parameter number and parameter gradient "
+    PADDLE_ENFORCE_EQ(
+        param.size(),
+        grad.size(),
+        errors::InvalidArgument("The parameter number and parameter gradient "
                                 "number should be the same."));
 
-      auto params_out = ctx.MultiOutput<phi::DenseTensor>("ParamOut");
-      auto grads_out = ctx.MultiOutput<phi::DenseTensor>("GradOut");
     PADDLE_ENFORCE_EQ(
-          params.size(),
-          params_out.size(),
-          platform::errors::InvalidArgument("Input(Param) and Output(ParamOut) "
+        param.size(),
+        param_out.size(),
+        errors::InvalidArgument("Input(Param) and Output(ParamOut) "
                                 "should have the same number."));
     PADDLE_ENFORCE_EQ(
-          grads.size(),
-          grads_out.size(),
-          platform::errors::InvalidArgument(
+        grad.size(),
+        grad_out.size(),
+        errors::InvalidArgument(
             "Input(Grad) and Output(GradOut) should have the same number."));
-      size_t n = params.size();
+    size_t n = param.size();
     VLOG(10) << "parameter number: " << n;
     for (size_t i = 0; i < n; ++i) {
-        auto *p = params[i];
-        auto *g = grads[i];
-        auto *p_out = params_out[i];
-        auto *g_out = grads_out[i];
+      auto *p = param[i];
+      auto *g = grad[i];
+      auto *p_out = param_out[i];
+      auto *g_out = grad_out[i];
 
       PADDLE_ENFORCE_NOT_NULL(
           p,
-            platform::errors::InvalidArgument(
-                "The %d-th parameter should not be nullptr.", i));
+          errors::InvalidArgument("The %d-th parameter should not be nullptr.",
+                                  i));
       PADDLE_ENFORCE_EQ(p->IsInitialized(),
                         true,
-                          platform::errors::InvalidArgument(
+                        errors::InvalidArgument(
                             "The %d-th parameter should be initialized.", i));
       PADDLE_ENFORCE_EQ(
           p->place(),
           place,
-            platform::errors::InvalidArgument(
-                "The %d-th parameter is not initialized on the right place.",
-                i));
-        PADDLE_ENFORCE_EQ(p,
+          errors::InvalidArgument(
+              "The %d-th parameter is not initialized on the right place.", i));
+      PADDLE_ENFORCE_EQ(
+          p,
           p_out,
-                          platform::errors::InvalidArgument(
-                              "The %d-th Input(Param) and Output(ParamOut) "
+          errors::InvalidArgument("The %d-th Input(Param) and Output(ParamOut) "
                                   "should be the same tensor.",
                                   i));
 
       auto dtype = p->dtype();
       PADDLE_ENFORCE_NOT_NULL(
           g,
-            platform::errors::InvalidArgument(
-                "The %d-th gradient should not be nullptr.", i));
+          errors::InvalidArgument("The %d-th gradient should not be nullptr.",
+                                  i));
       PADDLE_ENFORCE_EQ(g,
                         g_out,
-                          platform::errors::InvalidArgument(
+                        errors::InvalidArgument(
                             "The %d-th Input(Grad) and Output(Grad) should "
                             "be the same tensor."));
       auto numel = p->numel();
-        PADDLE_ENFORCE_GT(numel,
+      PADDLE_ENFORCE_GT(
+          numel,
           0,
-                          platform::errors::InvalidArgument(
-                              "The %d-th Input(Param) have no elements."));
+          errors::InvalidArgument("The %d-th Input(Param) have no elements."));
 
       void *g_data = nullptr;
       if (g->IsInitialized()) {
         PADDLE_ENFORCE_EQ(g->dtype(),
                           dtype,
-                            platform::errors::InvalidArgument(
+                          errors::InvalidArgument(
                               "The %d-th Input(Param) and Input(Grad) should "
                               "have the same data type %s.",
                               i,
                               dtype));
         PADDLE_ENFORCE_EQ(g->dims(),
                           p->dims(),
-                            platform::errors::InvalidArgument(
+                          errors::InvalidArgument(
                               "The %d-th Input(Param) and Input(Grad) should "
                               "have the same shape.",
                               i));
@@ -445,8 +455,8 @@ class DistributedFusedLambInitOpKernel<T, phi::GPUContext>
         fp16_infos.emplace_back();
         info = &fp16_infos.back();
       } else {
-          PADDLE_THROW(platform::errors::InvalidArgument(
-              "Unsupported data type %s.", dtype));
+        PADDLE_THROW(
+            errors::InvalidArgument("Unsupported data type %s.", dtype));
       }
 
       VLOG(10) << "Found " << dtype << " parameter " << i << " shape=["
@@ -462,108 +472,102 @@ class DistributedFusedLambInitOpKernel<T, phi::GPUContext>
       info->numel_offset = 0;        // not determined yet
     }
   }
-    const auto &apply_weight_decay =
-        ctx.Attr<std::vector<int>>("apply_weight_decay");
+
   size_t fp32_wd_end_idx =
       ReorderParamGradInfoList(apply_weight_decay, &fp32_infos);
   size_t fp16_wd_end_idx =
       ReorderParamGradInfoList(apply_weight_decay, &fp16_infos);
 
-    auto *param_order_t = ctx.Output<phi::DenseTensor>("ParamOrder");
   auto param_num = fp32_infos.size() + fp16_infos.size();
-    param_order_t->Resize({static_cast<int16_t>(param_num)});
-    auto *param_order = param_order_t->mutable_data<int>(platform::CPUPlace());
+  param_order->Resize({static_cast<int16_t>(param_num)});
+  auto *param_order_t = dev_ctx.template HostAlloc<int>(param_order);
   for (size_t i = 0; i < fp32_infos.size(); ++i) {
-      param_order[i] = static_cast<int>(fp32_infos[i].idx);
+    param_order_t[i] = static_cast<int>(fp32_infos[i].idx);
   }
   for (size_t i = 0; i < fp16_infos.size(); ++i) {
-      param_order[i + fp32_infos.size()] = static_cast<int>(fp16_infos[i].idx);
+    param_order_t[i + fp32_infos.size()] = static_cast<int>(fp16_infos[i].idx);
   }
 
   VLOG(10) << "Fill ParamGradInfo ends";
 
   // Step 2: determine the numel_with_padding and numel_offset
-    auto rank = ctx.Attr<int>("rank");
-    auto nranks = ctx.Attr<int>("nranks");
-    auto alignment = ctx.Attr<int>("alignment");
   VLOG(10) << "rank = " << rank << ", nranks = " << nranks
            << " , alignment = " << alignment;
   if (alignment <= 0) {
-      alignment = platform::GpuMinChunkSize();
+    alignment = phi::backends::gpu::GpuMinChunkSize();
   }
-    PADDLE_ENFORCE_GE(alignment,
+  PADDLE_ENFORCE_GE(
+      alignment,
       1,
-                      platform::errors::InvalidArgument(
-                          "The attr(alignment) should be larger than 0."));
-    PADDLE_ENFORCE_EQ(alignment & (alignment - 1),
+      errors::InvalidArgument("The attr(alignment) should be larger than 0."));
+  PADDLE_ENFORCE_EQ(
+      alignment & (alignment - 1),
       0,
-                      platform::errors::InvalidArgument(
-                          "The attr(alignment) should be the power of 2."));
-    PADDLE_ENFORCE_GE(
-        rank,
+      errors::InvalidArgument("The attr(alignment) should be the power of 2."));
+  PADDLE_ENFORCE_GE(rank,
                     0,
-        platform::errors::InvalidArgument(
+                    errors::InvalidArgument(
                         "The attr(rank) should be equal to or larger than 0."));
   PADDLE_ENFORCE_LT(
       rank,
       nranks,
-        platform::errors::InvalidArgument(
+      errors::InvalidArgument(
           "The attr(rank) should be less than the attr(nranks)."));
   // NOTE: We guarantee that both fp32_numel and fp16_numel can be exactly
   // divided by alignment and nranks.
   auto fp32_numel = FillAlignmentPaddingInfo(
       &fp32_infos, alignment, nranks, phi::DataType::FLOAT32);
-    VLOG(10) << "FP32 ParamGradInfo: " << string::join_strings(fp32_infos, " ");
+  VLOG(10) << "FP32 ParamGradInfo: "
+           << paddle::string::join_strings(fp32_infos, " ");
   auto fp16_numel = FillAlignmentPaddingInfo(
       &fp16_infos, alignment, nranks, phi::DataType::FLOAT16);
-    VLOG(10) << "FP16 ParamGradInfo: " << string::join_strings(fp16_infos, " ");
+  VLOG(10) << "FP16 ParamGradInfo: "
+           << paddle::string::join_strings(fp16_infos, " ");
   auto total_numel = fp32_numel + fp16_numel;
-    PADDLE_ENFORCE_LT(
-        total_numel,
+  PADDLE_ENFORCE_LT(total_numel,
                     std::numeric_limits<int>::max(),
-        platform::errors::InvalidArgument("Too many parameter number."));
+                    errors::InvalidArgument("Too many parameter number."));
 
   auto fp32_numel_each_device = fp32_numel / nranks;
   auto fp16_numel_each_device = fp16_numel / nranks;
   auto numel_each_device = fp32_numel_each_device + fp16_numel_each_device;
   VLOG(10) << "Fill padding ends. total_numel = " << total_numel
-             << ", fp32_numel = " << fp32_numel
-             << ", fp16_numel = " << fp16_numel
+           << ", fp32_numel = " << fp32_numel << ", fp16_numel = " << fp16_numel
            << ", fp32_numel_each_device = " << fp32_numel_each_device
            << ", fp16_numel_each_device = " << fp16_numel_each_device;
 
   // Step 3: allocate output tensor and do initialization
   float *fused_fp32_param = nullptr, *fused_fp32_grad = nullptr;
-    platform::float16 *fused_fp16_param = nullptr, *fused_fp16_grad = nullptr;
-    phi::DenseTensor *fp32_p_t = nullptr, *fp16_p_t = nullptr,
-                     *fp32_g_t = nullptr, *fp16_g_t = nullptr;
-    std::vector<phi::DenseTensor *> fp16_master_params;
+  dtype::float16 *fused_fp16_param = nullptr, *fused_fp16_grad = nullptr;
+  DenseTensor *fp32_p_t = nullptr, *fp16_p_t = nullptr, *fp32_g_t = nullptr,
+              *fp16_g_t = nullptr;
+  std::vector<DenseTensor *> fp16_master_params;
   if (total_numel > 0) {
-      fp32_p_t = ctx.Output<phi::DenseTensor>("FP32FusedParam");
+    fp32_p_t = fp32_fused_param;
     fused_fp32_param = TensorFillConstant<float>(
         dev_ctx, fp32_p_t, {static_cast<int64_t>(total_numel)}, 0.0f);
   }
 
   if (fp32_numel > 0) {
-      fp32_g_t = ctx.Output<phi::DenseTensor>("FP32FusedGrad");
+    fp32_g_t = fp32_fused_grad;
     fused_fp32_grad = TensorFillConstant<float>(
         dev_ctx, fp32_g_t, {static_cast<int64_t>(fp32_numel)}, 0.0f);
   }
 
   if (fp16_numel > 0) {
-      fp16_p_t = ctx.Output<phi::DenseTensor>("FP16FusedParam");
-      fused_fp16_param = TensorFillConstant<platform::float16>(
-          dev_ctx,
+    fp16_p_t = fp16_fused_param;
+    fused_fp16_param =
+        TensorFillConstant<dtype::float16>(dev_ctx,
                                            fp16_p_t,
                                            {static_cast<int64_t>(fp16_numel)},
-          static_cast<platform::float16>(0));
+                                           static_cast<dtype::float16>(0));
 
-      fp16_g_t = ctx.Output<phi::DenseTensor>("FP16FusedGrad");
-      fused_fp16_grad = TensorFillConstant<platform::float16>(
-          dev_ctx,
+    fp16_g_t = fp16_fused_grad;
+    fused_fp16_grad =
+        TensorFillConstant<dtype::float16>(dev_ctx,
                                            fp16_g_t,
                                            {static_cast<int64_t>(fp16_numel)},
-          static_cast<platform::float16>(0));
+                                           static_cast<dtype::float16>(0));
   }
   VLOG(10) << "Allocate FP32FusedParam/Grad, FP16FusedParam/Grad ends";
 
@@ -573,16 +577,20 @@ class DistributedFusedLambInitOpKernel<T, phi::GPUContext>
   // (3) For FP32FusedGrad/FP16FusedGrad, memcpy if gradient has been inited
   for (const auto &info : fp32_infos) {
     auto sliced_tensor = CopyAndShareBufferForInitedTensor(
-          info.param_t, fp32_p_t, info.numel_offset, stream);
-      master_params[info.idx]->Resize(info.param_t->dims());
-      master_params[info.idx]->ShareBufferWith(sliced_tensor);
-      PADDLE_ENFORCE_EQ(master_params[info.idx]->mutable_data<float>(place),
-                        sliced_tensor.data<float>(),
-                        platform::errors::InvalidArgument(
-                            "Invalid master weight tensor pointer."));
+        dev_ctx, info.param_t, fp32_p_t, info.numel_offset);
+    master_param_out[info.idx]->Resize(info.param_t->dims());
+    master_param_out[info.idx]->ShareBufferWith(sliced_tensor);
+    float *master_param_tmp =
+        dev_ctx.template Alloc<float>(master_param_out[info.idx]);
+    float *sliced_tensor_tmp = reinterpret_cast<float *>(sliced_tensor.data());
+    PADDLE_ENFORCE_EQ(
+        master_param_tmp,
+        sliced_tensor_tmp,
+        errors::InvalidArgument("Invalid master weight tensor pointer."));
+
     if (info.grad_t->IsInitialized()) {
       CopyAndShareBufferForInitedTensor(
-            info.grad_t, fp32_g_t, info.numel_offset, stream);
+          dev_ctx, info.grad_t, fp32_g_t, info.numel_offset);
     } else {
       ShareBufferForNonInitedTensor(
           info.grad_t, fp32_g_t, info.numel_offset, info.param_t->dims());
@@ -600,19 +608,22 @@ class DistributedFusedLambInitOpKernel<T, phi::GPUContext>
     auto master_weight_offset = info.numel_offset + fp16_numel_offset;
     auto sliced_tensor = CastDataForInitedTensor(
         dev_ctx, info.param_t, fp32_p_t, master_weight_offset);
-      master_params[info.idx]->Resize(info.param_t->dims());
-      master_params[info.idx]->ShareBufferWith(sliced_tensor);
+    master_param_out[info.idx]->Resize(info.param_t->dims());
+    master_param_out[info.idx]->ShareBufferWith(sliced_tensor);
 
     CopyAndShareBufferForInitedTensor(
-          info.param_t, fp16_p_t, info.numel_offset, stream);
-      PADDLE_ENFORCE_EQ(master_params[info.idx]->mutable_data<float>(place),
-                        sliced_tensor.data<float>(),
-                        platform::errors::InvalidArgument(
-                            "Invalid master weight tensor pointer."));
+        dev_ctx, info.param_t, fp16_p_t, info.numel_offset);
+    float *master_param_tmp =
+        dev_ctx.template Alloc<float>(master_param_out[info.idx]);
+    float *sliced_tensor_tmp = reinterpret_cast<float *>(sliced_tensor.data());
+    PADDLE_ENFORCE_EQ(
+        master_param_tmp,
+        sliced_tensor_tmp,
+        errors::InvalidArgument("Invalid master weight tensor pointer."));
 
     if (info.grad_t->IsInitialized()) {
       CopyAndShareBufferForInitedTensor(
-            info.grad_t, fp16_g_t, info.numel_offset, stream);
+          dev_ctx, info.grad_t, fp16_g_t, info.numel_offset);
     } else {
       ShareBufferForNonInitedTensor(
           info.grad_t, fp16_g_t, info.numel_offset, info.param_t->dims());
@@ -621,22 +632,12 @@ class DistributedFusedLambInitOpKernel<T, phi::GPUContext>
   VLOG(10) << "Copy/share data for Param/Grad ends";
 
   // Step 4: For Moment1, Moment2, Beta1Pow, Beta2Pow, just fill constant
-    TensorFillConstant<float>(dev_ctx,
-                              ctx.Output<phi::DenseTensor>("Moment1"),
-                              {static_cast<int64_t>(numel_each_device)},
-                              0.0f);
-    TensorFillConstant<float>(dev_ctx,
-                              ctx.Output<phi::DenseTensor>("Moment2"),
-                              {static_cast<int64_t>(numel_each_device)},
-                              0.0f);
-    TensorFillConstant<float>(dev_ctx,
-                              ctx.Output<phi::DenseTensor>("Beta1Pow"),
-                              {1},
-                              ctx.Attr<float>("beta1"));
-    TensorFillConstant<float>(dev_ctx,
-                              ctx.Output<phi::DenseTensor>("Beta2Pow"),
-                              {1},
-                              ctx.Attr<float>("beta2"));
+  TensorFillConstant<float>(
+      dev_ctx, moment1, {static_cast<int64_t>(numel_each_device)}, 0.0f);
+  TensorFillConstant<float>(
+      dev_ctx, moment2, {static_cast<int64_t>(numel_each_device)}, 0.0f);
+  TensorFillConstant<float>(dev_ctx, beta1_pow, {1}, beta1);
+  TensorFillConstant<float>(dev_ctx, beta2_pow, {1}, beta2);
   VLOG(10) << "Init Moment and BetaPow ends";
 
   // Step 5: Do sharding
@@ -665,34 +666,33 @@ class DistributedFusedLambInitOpKernel<T, phi::GPUContext>
   size_t total_local_param_num = fp32_local_param_num + fp16_local_param_num;
   VLOG(10) << "Found the sharding arguments";
 
-    auto *param_info_t = ctx.Output<phi::DenseTensor>("ParamInfo");
-    param_info_t->Resize({8});
-    auto *param_info = param_info_t->mutable_data<int>(platform::CPUPlace());
-    param_info[0] = static_cast<int>(fp32_start_idx);
-    param_info[1] = static_cast<int>(fp32_local_param_num);
-    param_info[2] = static_cast<int>(fp32_infos.size());
-    param_info[3] = ClipByBound<int>(fp32_wd_end_idx,
+  param_info->Resize({8});
+  auto *param_info_t = dev_ctx.template HostAlloc<int>(param_info);
+  param_info_t[0] = static_cast<int>(fp32_start_idx);
+  param_info_t[1] = static_cast<int>(fp32_local_param_num);
+  param_info_t[2] = static_cast<int>(fp32_infos.size());
+  param_info_t[3] = ClipByBound<int>(fp32_wd_end_idx,
                                      fp32_start_idx,
                                      fp32_start_idx + fp32_local_param_num) -
                     static_cast<int>(fp32_start_idx);
-    param_info[4] = static_cast<int>(fp16_start_idx + fp32_infos.size());
-    param_info[5] = static_cast<int>(fp16_local_param_num);
-    param_info[6] = static_cast<int>(fp16_infos.size());
-    param_info[7] = ClipByBound<int>(fp16_wd_end_idx,
+  param_info_t[4] = static_cast<int>(fp16_start_idx + fp32_infos.size());
+  param_info_t[5] = static_cast<int>(fp16_local_param_num);
+  param_info_t[6] = static_cast<int>(fp16_infos.size());
+  param_info_t[7] = ClipByBound<int>(fp16_wd_end_idx,
                                      fp16_start_idx,
                                      fp16_start_idx + fp16_local_param_num) -
                     static_cast<int>(fp16_start_idx);
 
-    VLOG(10) << "Start FP32 idx: " << param_info[0];
-    VLOG(10) << "Local FP32 param num: " << param_info[1];
-    VLOG(10) << "Global FP32 param num: " << param_info[2];
+  VLOG(10) << "Start FP32 idx: " << param_info_t[0];
+  VLOG(10) << "Local FP32 param num: " << param_info_t[1];
+  VLOG(10) << "Global FP32 param num: " << param_info_t[2];
 
-    VLOG(10) << "Start FP16 idx: " << param_info[4];
-    VLOG(10) << "Local FP16 param num: " << param_info[5];
-    VLOG(10) << "Global FP16 param num: " << param_info[6];
+  VLOG(10) << "Start FP16 idx: " << param_info_t[4];
+  VLOG(10) << "Local FP16 param num: " << param_info_t[5];
+  VLOG(10) << "Global FP16 param num: " << param_info_t[6];
 
   std::vector<int> numel_offsets;
-    numel_offsets.reserve(params.size() + 1);
+  numel_offsets.reserve(param.size() + 1);
   for (const auto &info : fp32_infos) {
     numel_offsets.push_back(info.numel_offset);
   }
@@ -701,8 +701,8 @@ class DistributedFusedLambInitOpKernel<T, phi::GPUContext>
   }
   numel_offsets.push_back(fp32_numel + fp16_numel);
   PADDLE_ENFORCE_EQ(numel_offsets.size(),
-                      params.size() + 1,
-                      platform::errors::InvalidArgument(
+                    param.size() + 1,
+                    errors::InvalidArgument(
                         "The numel_offsets number must be one larger than "
                         "the parameter number."));
   VLOG(10) << "Total numel offset: " << FlattenToString(numel_offsets);
@@ -723,13 +723,12 @@ class DistributedFusedLambInitOpKernel<T, phi::GPUContext>
       end_n = std::min(end_n, fp32_end_numel_offset);
     }
 
-      PADDLE_ENFORCE_NE(valid_start_n,
+    PADDLE_ENFORCE_NE(
+        valid_start_n,
         end_n,
-                        platform::errors::InvalidArgument(
-                            "Indices sharding error. This may be a bug."));
-      VLOG(10) << "FP32 Partial numel = ["
-               << valid_start_n + fp32_infos[i].numel << ","
-               << end_n + fp32_infos[i].numel;
+        errors::InvalidArgument("Indices sharding error. This may be a bug."));
+    VLOG(10) << "FP32 Partial numel = [" << valid_start_n + fp32_infos[i].numel
+             << "," << end_n + fp32_infos[i].numel;
     auto len = end_n - valid_start_n;
     fp32_partial_numel_offsets.push_back(fp32_partial_numel_offsets.back() +
                                          len);
@@ -750,48 +749,56 @@ class DistributedFusedLambInitOpKernel<T, phi::GPUContext>
       end_n = std::min(end_n, fp16_end_numel_offset);
     }
 
-      PADDLE_ENFORCE_NE(valid_start_n,
+    PADDLE_ENFORCE_NE(
+        valid_start_n,
         end_n,
-                        platform::errors::InvalidArgument(
-                            "Indices sharding error. This may be a bug."));
+        errors::InvalidArgument("Indices sharding error. This may be a bug."));
     auto len = end_n - valid_start_n;
     fp16_partial_numel_offsets.push_back(fp16_partial_numel_offsets.back() +
                                          len);
   }
 
-    CopyVectorToCPUTensor(numel_offsets,
-                          ctx.Output<phi::DenseTensor>("FusedParamOffsets"));
+  CopyVectorToCPUTensor(dev_ctx, numel_offsets, fused_param_offsets);
   CopyVectorToCPUTensor(
-        fp32_partial_numel_offsets,
-        ctx.Output<phi::DenseTensor>("FP32ShardFusedParamOffsets"));
+      dev_ctx, fp32_partial_numel_offsets, fp32_shard_fused_param_offsets);
   CopyVectorToCPUTensor(
-        fp16_partial_numel_offsets,
-        ctx.Output<phi::DenseTensor>("FP16ShardFusedParamOffsets"));
+      dev_ctx, fp16_partial_numel_offsets, fp16_shard_fused_param_offsets);
 
-    auto *global_scale = ctx.Output<phi::DenseTensor>("GlobalScale");
   if (!global_scale->IsInitialized()) {
     TensorFillConstant<float>(dev_ctx, global_scale, {1}, 1.0f);
   }
   VLOG(10) << "Init global scale ends";
 
-    TensorFillConstant<int64_t>(dev_ctx,
-                                ctx.Output<phi::DenseTensor>("Step"),
-                                {1},
-                                static_cast<int64_t>(0));
+  TensorFillConstant<int64_t>(dev_ctx, step, {1}, static_cast<int64_t>(0));
 
   dev_ctx.Wait();
   VLOG(10) << "Wait for H2D copy";
-  }
-};
-
-}  // namespace operators
-}  // namespace paddle
+}
 
-namespace ops = paddle::operators;
-namespace plat = paddle::platform;
+}  // namespace fusion
+}  // namespace phi
 
-PD_REGISTER_STRUCT_KERNEL(distributed_fused_lamb_init,
+PD_REGISTER_KERNEL(distributed_fused_lamb_init,
                    GPU,
                    ALL_LAYOUT,
-                          ops::DistributedFusedLambInitOpKernel,
-                          float) {}
+                   phi::fusion::DistributedFusedLambInitOpKernel,
+                   float) {
+  kernel->OutputAt(0).SetDataType(phi::DataType::FLOAT32);
+  kernel->OutputAt(1).SetDataType(phi::DataType::FLOAT32);
+  kernel->OutputAt(2).SetDataType(phi::DataType::FLOAT16);
+  kernel->OutputAt(3).SetDataType(phi::DataType::FLOAT16);
+  kernel->OutputAt(4).SetDataType(phi::DataType::FLOAT32);
+  kernel->OutputAt(5).SetDataType(phi::DataType::FLOAT32);
+  kernel->OutputAt(6).SetDataType(phi::DataType::FLOAT32);
+  kernel->OutputAt(7).SetDataType(phi::DataType::FLOAT32);
+  kernel->OutputAt(8).SetDataType(phi::DataType::INT32);
+  kernel->OutputAt(9).SetDataType(phi::DataType::INT32);
+  kernel->OutputAt(10).SetDataType(phi::DataType::INT32);
+  kernel->OutputAt(11).SetDataType(phi::DataType::INT32);
+  kernel->OutputAt(12).SetDataType(phi::DataType::INT32);
+  kernel->OutputAt(13).SetDataType(kernel_key.dtype());
+  kernel->OutputAt(14).SetDataType(phi::DataType::FLOAT32);
+  kernel->OutputAt(15).SetDataType(kernel_key.dtype());
+  kernel->OutputAt(16).SetDataType(phi::DataType::FLOAT32);
+  kernel->OutputAt(17).SetDataType(phi::DataType::INT64);
+}

paddle/fluid/operators/optimizers/distributed_fused_lamb_init_op.h → paddle/phi/ops/compat/distributed_fused_lamb_init_sig.cc

-// Copyright (c) 2021 PaddlePaddle Authors. All Rights Reserved.
+// 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.
@@ -12,22 +12,37 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
-#pragma once
+#include "paddle/phi/core/compat/op_utils.h"
 
-#include "paddle/fluid/framework/op_registry.h"
-#include "paddle/fluid/framework/operator.h"
+namespace phi {
 
-namespace paddle {
-namespace operators {
+KernelSignature DistributedFusedLambInitOpArgumentMapping(
+    const ArgumentMappingContext& ctx UNUSED) {
+  return KernelSignature(
+      "distributed_fused_lamb_init",
+      {"Param", "Grad"},
+      {"beta1", "beta2", "apply_weight_decay", "alignment", "rank", "nranks"},
+      {"FP32FusedParam",
+       "FP32FusedGrad",
+       "FP16FusedParam",
+       "FP16FusedGrad",
+       "Moment1",
+       "Moment2",
+       "Beta1Pow",
+       "Beta2Pow",
+       "FusedParamOffsets",
+       "FP32ShardFusedParamOffsets",
+       "FP16ShardFusedParamOffsets",
+       "ParamInfo",
+       "ParamOrder",
+       "ParamOut",
+       "MasterParamOut",
+       "GradOut",
+       "GlobalScale",
+       "Step"});
+}
 
-template <typename T, typename DevCtx>
-class DistributedFusedLambInitOpKernel : public framework::OpKernel<T> {
- public:
-  void Compute(const framework::ExecutionContext &ctx) const override {
-    PADDLE_THROW(platform::errors::Unimplemented(
-        "The distributed_fused_lamb_init operator does not support CPU yet."));
-  }
-};
+}  // namespace phi
 
-}  // namespace operators
-}  // namespace paddle
+PD_REGISTER_ARG_MAPPING_FN(distributed_fused_lamb_init,
+                           phi::DistributedFusedLambInitOpArgumentMapping);