     auto grad = ctx.Input<framework::LoDTensor>("Grad");
     auto learning_rate = ctx.Input<framework::LoDTensor>("LearningRate");
 
+    int64_t numel = param->numel();
+    int grid = (numel + LARS_BLOCK_SIZE - 1) / LARS_BLOCK_SIZE;
     const framework::Tensor* master_param = nullptr;
     framework::Tensor* master_param_out = nullptr;
+    const MT* master_param_data = nullptr;
+    MT* master_param_out_data = nullptr;
+
     if (multi_precision) {
       bool has_master =
           ctx.HasInput("MasterParam") && ctx.HasOutput("MasterParamOut");

template <typename T>
using MultiPrecisionType = typename details::MPTypeTrait<T>::Type;

template <typename T, typename MT>
__global__ void MomentumLarsKernel(
    const T* p, const T* g, const MT* v,
    const MultiPrecisionType<T>* learning_rate, const MT mu, const int64_t num,
    const MT lars_coeff, const MT lars_weight_decay,
    const MultiPrecisionType<T>* p_norm, const MultiPrecisionType<T>* g_norm,
    T* p_out, MT* v_out, const MT epsilon, const MT* master_p, MT* master_p_out,
    const MultiPrecisionType<T> rescale_grad) {
  const MT lr = static_cast<MT>(learning_rate[0]);
  MT local_lr = lr;
  const MT p_n = static_cast<MT>(p_norm[0]);
  const MT g_n = static_cast<MT>(g_norm[0]);

  if (lars_weight_decay > static_cast<MT>(0) && p_n > static_cast<MT>(0) &&
      g_n > static_cast<MT>(0)) {
    local_lr =
        lr * lars_coeff * p_n / (g_n + lars_weight_decay * p_n + epsilon);
  }
  CUDA_KERNEL_LOOP(i, num) {
    MT grad = static_cast<MT>(g[i]) * static_cast<MT>(rescale_grad);
    MT param = master_p ? master_p[i] : static_cast<MT>(p[i]);

    MT v_new = v[i] * mu + local_lr * (grad + lars_weight_decay * param);
    MT p_new = param - v_new;

    v_out[i] = v_new;
    p_out[i] = static_cast<T>(p_new);
    if (master_p_out) master_p_out[i] = p_new;
  }
}

template <typename DeviceContext, typename T>
class LarsMomentumOpCUDAKernel : public framework::OpKernel<T> {
  using MPDType = MultiPrecisionType<T>;

 public:
  void Compute(const framework::ExecutionContext& ctx) const override {
    const bool multi_precision = ctx.Attr<bool>("multi_precision");
    if (multi_precision) {
      InnerCompute<MPDType>(ctx, multi_precision);
    } else {
      InnerCompute<T>(ctx, multi_precision);
    }
  }

 private:
  template <typename MT>
  void InnerCompute(const framework::ExecutionContext& ctx,
                    const bool multi_precision) const {
    auto param_out = ctx.Output<framework::LoDTensor>("ParamOut");
    auto velocity_out = ctx.Output<framework::LoDTensor>("VelocityOut");
    auto param = ctx.Input<framework::LoDTensor>("Param");


#include "paddle/fluid/framework/op_registry.h"
#include "paddle/fluid/operators/amp/fp16_type_traits.h"
#include "paddle/fluid/operators/math/math_cuda_utils.h"
#include "paddle/fluid/operators/optimizers/lars_momentum_op.h"
#include "paddle/fluid/platform/fast_divmod.h"

#if defined(__NVCC__) && CUDA_VERSION >= 11000
/* Once CUDA_VERSION is beyond 11.0, cooperative_groups can be involved in
   without adding --rdc=true compile flag, then L2_norm cuda kernel can be
   set as a __device__ kernel rather than global kernel. On the contrary,
   the compile flag shall be set in old version, which may affect the cuda
   kernel performance in paddle, consequently, L2_norm kernel shall be set
   as a __global__ kernel.
*/
#include <cooperative_groups.h>
#define LARS_FUNCTION_FLAG __device__
#else
#define LARS_FUNCTION_FLAG __global__
#endif

#ifdef __HIPCC__
#define LARS_BLOCK_SIZE 256
#else
#define LARS_BLOCK_SIZE 512
#endif

namespace paddle {
namespace operators {
