[XPU] use fp16 in adamw optimizer (#57077)

houj04 · web-flow · commit 774162bd0450 · 2023-09-12T16:39:05.000+08:00
diff --git a/paddle/phi/core/dense_tensor.cc b/paddle/phi/core/dense_tensor.cc
@@ -232,6 +232,9 @@ void DenseTensor::set_meta(const DenseTensorMeta& meta) {
   } else {
     meta_.strides = meta.strides;
   }
+#ifdef PADDLE_WITH_XPU
+  meta_.scale_value = meta.scale_value;
+#endif
 }
 
 /* @jim19930609: This interface will be further modified until we finalized the
diff --git a/paddle/phi/core/dense_tensor_impl.cc b/paddle/phi/core/dense_tensor_impl.cc
@@ -392,6 +392,9 @@ DenseTensor& DenseTensor::ShareDataWith(const DenseTensor& src) {
   meta_.offset = src.meta_.offset;
   meta_.use_gpudnn = src.meta_.use_gpudnn;
   meta_.strides = src.meta_.strides;
+#ifdef PADDLE_WITH_XPU
+  meta_.scale_value = src.meta_.scale_value;
+#endif
   storage_properties_ =
       std::move(CopyStorageProperties(src.storage_properties_));
 #ifdef PADDLE_WITH_DNNL
diff --git a/paddle/phi/core/tensor_meta.cc b/paddle/phi/core/tensor_meta.cc
@@ -118,19 +118,30 @@ DDim DenseTensorMeta::calc_strides(const DDim& dims) {
   }
 }
 
-DenseTensorMeta::DenseTensorMeta() { use_gpudnn = true; }
+DenseTensorMeta::DenseTensorMeta() {
+  use_gpudnn = true;
+#ifdef PADDLE_WITH_XPU
+  scale_value = -1.0f;
+#endif
+}
 
 DenseTensorMeta::DenseTensorMeta(DataType dtype, const DDim& dims)
     : dims(dims), dtype(dtype) {
   strides = calc_strides(dims);
   use_gpudnn = true;
+#ifdef PADDLE_WITH_XPU
+  scale_value = -1.0f;
+#endif
 }
 
 DenseTensorMeta::DenseTensorMeta(DataType dtype,
                                  const DDim& dims,
                                  const DDim& strides)
     : dims(dims), dtype(dtype), strides(strides) {
   use_gpudnn = true;
+#ifdef PADDLE_WITH_XPU
+  scale_value = -1.0f;
+#endif
 }
 
 DenseTensorMeta::DenseTensorMeta(DataType dtype,
@@ -140,6 +151,9 @@ DenseTensorMeta::DenseTensorMeta(DataType dtype,
     : dims(dims), dtype(dtype), layout(layout), offset(offset) {
   strides = calc_strides(dims);
   use_gpudnn = true;
+#ifdef PADDLE_WITH_XPU
+  scale_value = -1.0f;
+#endif
 }
 
 DenseTensorMeta::DenseTensorMeta(DataType dtype,
@@ -150,6 +164,9 @@ DenseTensorMeta::DenseTensorMeta(DataType dtype,
     : dims(dims), dtype(dtype), layout(layout), lod(lod), offset(offset) {
   strides = calc_strides(dims);
   use_gpudnn = true;
+#ifdef PADDLE_WITH_XPU
+  scale_value = -1.0f;
+#endif
 }
 
 DenseTensorMeta::DenseTensorMeta(const DenseTensorMeta& other) {
@@ -165,6 +182,9 @@ DenseTensorMeta::DenseTensorMeta(const DenseTensorMeta& other) {
   } else {
     strides = other.strides;
   }
+#ifdef PADDLE_WITH_XPU
+  scale_value = other.scale_value;
+#endif
 }
 
 DenseTensorMeta& DenseTensorMeta::operator=(const DenseTensorMeta& other) {
@@ -180,6 +200,9 @@ DenseTensorMeta& DenseTensorMeta::operator=(const DenseTensorMeta& other) {
   } else {
     strides = other.strides;
   }
+#ifdef PADDLE_WITH_XPU
+  scale_value = other.scale_value;
+#endif
   return *this;
 }
 
@@ -197,7 +220,9 @@ DenseTensorMeta& DenseTensorMeta::operator=(  // NOLINT
   } else {
     strides = std::move(other.strides);
   }
-
+#ifdef PADDLE_WITH_XPU
+  scale_value = other.scale_value;
+#endif
   return *this;
 }
 
diff --git a/paddle/phi/core/tensor_meta.h b/paddle/phi/core/tensor_meta.h
@@ -82,13 +82,23 @@ struct DenseTensorMeta {
   LoD lod;
   size_t offset{0};
   DDim strides;
+
+#ifdef PADDLE_WITH_XPU
+  // for per tensor scale
+  float scale_value{-1.0f};
+#endif
 };
 
 inline bool operator==(const DenseTensorMeta& lhs, const DenseTensorMeta& rhs) {
   return (lhs.is_scalar == rhs.is_scalar) && lhs.use_gpudnn == rhs.use_gpudnn &&
          (lhs.dims == rhs.dims) && (lhs.dtype == rhs.dtype) &&
          (lhs.layout == rhs.layout) && (lhs.lod == rhs.lod) &&
+#ifdef PADDLE_WITH_XPU
+         (lhs.offset == rhs.offset) && (lhs.strides == rhs.strides) &&
+         (lhs.scale_value == rhs.scale_value);
+#else
          (lhs.offset == rhs.offset) && (lhs.strides == rhs.strides);
+#endif
 }
 
 struct StringTensorMeta {
diff --git a/paddle/phi/kernels/xpu/adamw_kernel.cc b/paddle/phi/kernels/xpu/adamw_kernel.cc
@@ -20,11 +20,30 @@
 
 #include "paddle/phi/backends/xpu/enforce_xpu.h"
 #include "paddle/phi/backends/xpu/xpu_context.h"
+#include "paddle/phi/common/memory_utils.h"
 #include "paddle/phi/core/kernel_registry.h"
 #include "paddle/phi/core/tensor_utils.h"
 
 namespace phi {
 
+template <typename Context>
+float GetAbsMax(const Context& dev_ctx,
+                const float* input,
+                float* buffer_xpu,
+                int64_t numel) {
+  float buffer_cpu[6];
+  // int findmax(Context* ctx, const T* x, float* maxptr, int64_t len);
+  int r = xpu::findmax<float>(dev_ctx.x_context(), input, buffer_xpu, numel);
+  PADDLE_ENFORCE_XDNN_SUCCESS(r, "findmax");
+  memory_utils::Copy(CPUPlace(),
+                     static_cast<void*>(buffer_cpu),
+                     dev_ctx.GetPlace(),
+                     static_cast<void*>(buffer_xpu),
+                     sizeof(float) * 6);
+  float* max_value = std::max_element(buffer_cpu, buffer_cpu + 6);
+  return *max_value;
+}
+
 template <typename T, typename Context>
 void AdamwDenseKernel(const Context& dev_ctx,
                       const DenseTensor& param,
@@ -52,6 +71,98 @@ void AdamwDenseKernel(const Context& dev_ctx,
                       DenseTensor* beta1_pow_out,
                       DenseTensor* beta2_pow_out,
                       DenseTensor* master_param_outs) {
+  // check moment_dtype
+  auto moment1_dtype = moment1.dtype();
+  auto moment2_dtype = moment2.dtype();
+  PADDLE_ENFORCE_EQ(moment1_dtype,
+                    moment1_out->dtype(),
+                    errors::InvalidArgument(
+                        "moment1.dtype does not match moment1_out->dtype"));
+  PADDLE_ENFORCE_EQ(moment2_dtype,
+                    moment2_out->dtype(),
+                    errors::InvalidArgument(
+                        "moment2.dtype does not match moment2_out->dtype"));
+  PADDLE_ENFORCE_EQ(
+      moment1_dtype,
+      moment2_dtype,
+      errors::InvalidArgument("moment1.dtype does not match moment2.dtype"));
+
+  bool moment_in_fp16 = false;
+  if (moment1_dtype == phi::DataType::FLOAT16) {
+    moment_in_fp16 = true;
+  } else {
+    PADDLE_ENFORCE_EQ(
+        moment1_dtype,
+        phi::DataType::FLOAT32,
+        errors::InvalidArgument("moment1.dtype is neither fp32 nor fp16"));
+  }
+
+  float* moment1_input_for_xdnn = nullptr;
+  float* moment2_input_for_xdnn = nullptr;
+  float* moment1_output_for_xdnn = nullptr;
+  float* moment2_output_for_xdnn = nullptr;
+
+  xpu::ctx_guard RAII_GUARD(dev_ctx.x_context());
+  if (moment_in_fp16) {
+    // allocate temp buffer on XPU
+    moment1_input_for_xdnn = RAII_GUARD.alloc_l3_or_gm<float>(moment1.numel());
+    PADDLE_ENFORCE_XDNN_NOT_NULL(moment1_input_for_xdnn);
+    moment2_input_for_xdnn = RAII_GUARD.alloc_l3_or_gm<float>(moment2.numel());
+    PADDLE_ENFORCE_XDNN_NOT_NULL(moment2_input_for_xdnn);
+    moment1_output_for_xdnn =
+        RAII_GUARD.alloc_l3_or_gm<float>(moment1_out->numel());
+    PADDLE_ENFORCE_XDNN_NOT_NULL(moment1_output_for_xdnn);
+    moment2_output_for_xdnn =
+        RAII_GUARD.alloc_l3_or_gm<float>(moment2_out->numel());
+    PADDLE_ENFORCE_XDNN_NOT_NULL(moment2_output_for_xdnn);
+
+    int r = 0;
+    using XPUType16 = typename XPUTypeTrait<phi::dtype::float16>::Type;
+
+    // cast moment1 and moment2, from fp16 to fp32
+    // int cast(Context* ctx, const TX* x, TY* y, int64_t len);
+    r = xpu::cast<XPUType16, float>(
+        dev_ctx.x_context(),
+        reinterpret_cast<const XPUType16*>(
+            moment1.template data<phi::dtype::float16>()),
+        moment1_input_for_xdnn,
+        moment1.numel());
+    PADDLE_ENFORCE_XDNN_SUCCESS(r, "cast moment1 from fp16 to float");
+    r = xpu::cast<XPUType16, float>(
+        dev_ctx.x_context(),
+        reinterpret_cast<const XPUType16*>(
+            moment2.template data<phi::dtype::float16>()),
+        moment2_input_for_xdnn,
+        moment2.numel());
+    PADDLE_ENFORCE_XDNN_SUCCESS(r, "cast moment2 from fp16 to float");
+
+    // de-scale using meta's scale_value
+    // int scale(Context* ctx, const T* x, T* y, int64_t len, bool
+    // bias_after_scale, float _scale, float _bias);
+    phi::DenseTensorMeta moment1_meta = moment1.meta();
+    if (moment1_meta.scale_value > 0) {
+      r = xpu::scale<float>(dev_ctx.x_context(),
+                            moment1_input_for_xdnn,
+                            moment1_input_for_xdnn,
+                            moment1.numel(),
+                            false,
+                            1.0f / moment1_meta.scale_value,
+                            0.0f);
+      PADDLE_ENFORCE_XDNN_SUCCESS(r, "de-scale for moment1");
+    }
+    phi::DenseTensorMeta moment2_meta = moment2.meta();
+    if (moment2_meta.scale_value > 0) {
+      r = xpu::scale<float>(dev_ctx.x_context(),
+                            moment2_input_for_xdnn,
+                            moment2_input_for_xdnn,
+                            moment2.numel(),
+                            false,
+                            1.0f / moment2_meta.scale_value,
+                            0.0f);
+      PADDLE_ENFORCE_XDNN_SUCCESS(r, "de-scale for moment2");
+    }
+  }
+
   using XPUType = typename XPUTypeTrait<T>::Type;
   bool skip_update_ = false;
   if (skip_update.is_initialized()) {
@@ -94,7 +205,7 @@ void AdamwDenseKernel(const Context& dev_ctx,
   if (!with_decay) {
     coeff = static_cast<float>(0.0);
   }
-  xpu::ctx_guard RAII_GUARD(dev_ctx.x_context());
+
   float* new_lr = RAII_GUARD.alloc_l3_or_gm<float>(learning_rate.numel());
   PADDLE_ENFORCE_XDNN_NOT_NULL(new_lr);
   int r = 0;
@@ -107,17 +218,23 @@ void AdamwDenseKernel(const Context& dev_ctx,
                  0.0f);
   PADDLE_ENFORCE_XDNN_SUCCESS(r, "scale");
 
+  // int adamw(Context* ctx, const T* g, const float* mom1, const float* mom2,
+  // const T* param, const float* beta1_pow, const float* beta2_pow, const
+  // float* lr, float* moment1_out, float* moment2_out, T* param_out, float
+  // beta1, float beta2, float epsilon, float coeff, int64_t n);
   r = xpu::adamw(
       dev_ctx.x_context(),
       reinterpret_cast<const XPUType*>(grad.template data<T>()),
-      moment1.template data<float>(),
-      moment2.template data<float>(),
+      moment_in_fp16 ? moment1_input_for_xdnn : moment1.template data<float>(),
+      moment_in_fp16 ? moment2_input_for_xdnn : moment2.template data<float>(),
       reinterpret_cast<const XPUType*>(param.template data<T>()),
       beta1_pow_ptr,
       beta2_pow_ptr,
       new_lr,
-      dev_ctx.template Alloc<float>(moment1_out),
-      dev_ctx.template Alloc<float>(moment2_out),
+      moment_in_fp16 ? moment1_output_for_xdnn
+                     : dev_ctx.template Alloc<float>(moment1_out),
+      moment_in_fp16 ? moment2_output_for_xdnn
+                     : dev_ctx.template Alloc<float>(moment2_out),
       reinterpret_cast<XPUType*>(dev_ctx.template Alloc<T>(param_out)),
       beta1_,
       beta2_,
@@ -126,6 +243,75 @@ void AdamwDenseKernel(const Context& dev_ctx,
       param.numel());
   PADDLE_ENFORCE_XDNN_SUCCESS(r, "adamw");
 
+  if (moment_in_fp16) {
+    int r = 0;
+    using XPUType16 = typename XPUTypeTrait<phi::dtype::float16>::Type;
+
+    // findmax and calculate scale_value for moment1 and moment2
+    float* buffer_for_findmax = RAII_GUARD.alloc_l3_or_gm<float>(6);
+
+    // for moment1
+    float moment1_max = GetAbsMax<Context>(dev_ctx,
+                                           moment1_output_for_xdnn,
+                                           buffer_for_findmax,
+                                           moment1_out->numel());
+    float moment1_scale_value = 65504.0f / moment1_max / 2.0f;
+    // int scale(Context* ctx, const T* x, T* y, int64_t len, bool
+    // bias_after_scale, float _scale, float _bias);
+    r = xpu::scale<float>(dev_ctx.x_context(),
+                          moment1_output_for_xdnn,
+                          moment1_output_for_xdnn,
+                          moment1_out->numel(),
+                          false,
+                          moment1_scale_value,
+                          0.0f);
+    PADDLE_ENFORCE_XDNN_SUCCESS(
+        r, "scale before convert to fp16, for moment1_output_for_xdnn");
+    // write to meta info
+    phi::DenseTensorMeta moment1_out_meta = moment1_out->meta();
+    moment1_out_meta.scale_value = moment1_scale_value;
+    moment1_out->set_meta(moment1_out_meta);
+
+    // for moment2
+    float moment2_max = GetAbsMax<Context>(dev_ctx,
+                                           moment2_output_for_xdnn,
+                                           buffer_for_findmax,
+                                           moment2_out->numel());
+    float moment2_scale_value = 65504.0f / moment2_max / 2.0f;
+    // int scale(Context* ctx, const T* x, T* y, int64_t len, bool
+    // bias_after_scale, float _scale, float _bias);
+    r = xpu::scale<float>(dev_ctx.x_context(),
+                          moment2_output_for_xdnn,
+                          moment2_output_for_xdnn,
+                          moment2_out->numel(),
+                          false,
+                          moment2_scale_value,
+                          0.0f);
+    PADDLE_ENFORCE_XDNN_SUCCESS(
+        r, "scale before convert to fp16, for moment2_output_for_xdnn");
+    // write to meta info
+    phi::DenseTensorMeta moment2_out_meta = moment2_out->meta();
+    moment2_out_meta.scale_value = moment2_scale_value;
+    moment2_out->set_meta(moment2_out_meta);
+
+    // cast moment1 and moment2 output, from fp32 to fp16
+    // int cast(Context* ctx, const TX* x, TY* y, int64_t len);
+    r = xpu::cast<float, XPUType16>(
+        dev_ctx.x_context(),
+        moment1_output_for_xdnn,
+        reinterpret_cast<XPUType16*>(
+            dev_ctx.template Alloc<phi::dtype::float16>(moment1_out)),
+        moment1.numel());
+    PADDLE_ENFORCE_XDNN_SUCCESS(r, "cast moment1_out from float to fp16");
+    r = xpu::cast<float, XPUType16>(
+        dev_ctx.x_context(),
+        moment2_output_for_xdnn,
+        reinterpret_cast<XPUType16*>(
+            dev_ctx.template Alloc<phi::dtype::float16>(moment2_out)),
+        moment2.numel());
+    PADDLE_ENFORCE_XDNN_SUCCESS(r, "cast moment2_out from float to fp16");
+  }
+
   if (!use_global_beta_pow) {
     // update in cpu
     if (beta1_pow.place() == CPUPlace() && beta2_pow.place() == CPUPlace()) {
diff --git a/python/paddle/optimizer/adamw.py b/python/paddle/optimizer/adamw.py
@@ -341,8 +341,25 @@ def _add_moments_pows(self, p):
         acc_dtype = p.dtype
         if self._is_dtype_fp16_or_bf16(acc_dtype):
             acc_dtype = core.VarDesc.VarType.FP32
-        self._add_accumulator(self._moment1_acc_str, p, dtype=acc_dtype)
-        self._add_accumulator(self._moment2_acc_str, p, dtype=acc_dtype)
+        if core.is_compiled_with_xpu():
+            import os
+
+            xpu_adamw_moment_dtype = os.getenv(
+                "xpu_adamw_moment_dtype", default="fp32"
+            )
+            if xpu_adamw_moment_dtype == "fp16":
+                self._add_accumulator(
+                    self._moment1_acc_str, p, dtype=core.VarDesc.VarType.FP16
+                )
+                self._add_accumulator(
+                    self._moment2_acc_str, p, dtype=core.VarDesc.VarType.FP16
+                )
+            else:
+                self._add_accumulator(self._moment1_acc_str, p, dtype=acc_dtype)
+                self._add_accumulator(self._moment2_acc_str, p, dtype=acc_dtype)
+        else:
+            self._add_accumulator(self._moment1_acc_str, p, dtype=acc_dtype)
+            self._add_accumulator(self._moment2_acc_str, p, dtype=acc_dtype)
         self._add_accumulator(
             name=self._beta1_pow_acc_str,
             param=p,

Original file line number	Diff line number	Diff line change
`@@ -232,6 +232,9 @@ void DenseTensor::set_meta(const DenseTensorMeta& meta) {`
`232`	`232`	`} else {`
`233`	`233`	`meta_.strides = meta.strides;`
`234`	`234`	`}`
	`235`	`+#ifdef PADDLE_WITH_XPU`
	`236`	`+ meta_.scale_value = meta.scale_value;`
	`237`	`+#endif`
`235`	`238`	`}`
`236`	`239`
`237`	`240`	`/* @jim19930609: This interface will be further modified until we finalized the`