[XPU] refine scale_value in AdamW (#58805)

houj04 · web-flow · commit 93b4f21b456b · 2023-11-13T16:31:51.000+08:00
* [XPU] fix adamw optimizer state_dict

* fix typo

* rename scale to scale_value

* revert unnecessary codes.

* remove scale_value from DenseTensorMeta

* fix typo
diff --git a/paddle/fluid/pybind/eager_utils.cc b/paddle/fluid/pybind/eager_utils.cc
@@ -750,9 +750,9 @@ PyObject* ToPyObject(int64_t value) { return PyLong_FromLongLong(value); }
 
 PyObject* ToPyObject(size_t value) { return PyLong_FromSize_t(value); }
 
-PyObject* ToPyObject(float value) { return PyLong_FromDouble(value); }
+PyObject* ToPyObject(float value) { return PyFloat_FromDouble(value); }
 
-PyObject* ToPyObject(double value) { return PyLong_FromDouble(value); }
+PyObject* ToPyObject(double value) { return PyFloat_FromDouble(value); }
 
 PyObject* ToPyObject(const char* value) { return PyUnicode_FromString(value); }
 
diff --git a/paddle/fluid/pybind/tensor.cc b/paddle/fluid/pybind/tensor.cc
@@ -444,6 +444,24 @@ void BindTensor(pybind11::module &m) {  // NOLINT
       .def("_set_complex128_element", TensorSetElement<paddle::complex128>)
       .def("_get_complex128_element", TensorGetElement<paddle::complex128>)
       .def("_place", [](phi::DenseTensor &self) { return self.place(); })
+#ifdef PADDLE_WITH_XPU
+      .def("get_xpu_scale_value",
+           [](phi::DenseTensor &self) {
+             if (self.storage_properties_initialized()) {
+               const phi::XPUStorageProperties &sp =
+                   self.storage_properties<phi::XPUStorageProperties>();
+               return sp.xpu_scale_value;
+             } else {
+               return phi::XPUStorageProperties::default_xpu_scale_value;
+             }
+           })
+      .def("set_xpu_scale_value",
+           [](phi::DenseTensor &self, float new_value) {
+             std::unique_ptr<phi::StorageProperties> sp =
+                 std::make_unique<phi::XPUStorageProperties>(new_value);
+             self.set_storage_properties(std::move(sp));
+           })
+#endif
       .def("_dtype",
            [](phi::DenseTensor &self) {
              return framework::TransToProtoVarType(self.type());
diff --git a/paddle/phi/core/dense_tensor.cc b/paddle/phi/core/dense_tensor.cc
@@ -233,9 +233,6 @@ 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
@@ -312,6 +309,9 @@ template const NPUStorageProperties& DenseTensor::storage_properties() const;
 #ifdef PADDLE_WITH_DNNL
 template const OneDNNStorageProperties& DenseTensor::storage_properties() const;
 #endif
+#ifdef PADDLE_WITH_XPU
+template const XPUStorageProperties& DenseTensor::storage_properties() const;
+#endif
 
 bool DenseTensor::storage_properties_initialized() const {
   return storage_properties_ != nullptr;
diff --git a/paddle/phi/core/dense_tensor_impl.cc b/paddle/phi/core/dense_tensor_impl.cc
@@ -392,9 +392,6 @@ 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/storage_properties.cc b/paddle/phi/core/storage_properties.cc
@@ -33,6 +33,13 @@ std::unique_ptr<StorageProperties> CopyStorageProperties(
       result->mem_desc =
           static_cast<OneDNNStorageProperties*>(sp.get())->mem_desc;
       return result;
+#endif
+#ifdef PADDLE_WITH_XPU
+    } else if (XPUStorageProperties::classof(sp.get())) {
+      auto result = std::make_unique<XPUStorageProperties>();
+      result->xpu_scale_value =
+          static_cast<XPUStorageProperties*>(sp.get())->xpu_scale_value;
+      return result;
 #endif
     } else {
       return nullptr;
diff --git a/paddle/phi/core/storage_properties.h b/paddle/phi/core/storage_properties.h
@@ -49,6 +49,20 @@ struct NPUStorageProperties
   DDim storage_dims;
 };
 
+#ifdef PADDLE_WITH_XPU
+struct XPUStorageProperties
+    : public StorageProperties,
+      public TypeInfoTraits<StorageProperties, XPUStorageProperties> {
+  XPUStorageProperties() = default;
+  explicit XPUStorageProperties(float value) : xpu_scale_value(value) {}
+  virtual ~XPUStorageProperties() = default;
+  static const char* name() { return "XPUStorageProperties"; }
+  static constexpr float default_xpu_scale_value = -1.0f;
+
+  float xpu_scale_value{default_xpu_scale_value};
+};
+#endif
+
 // Add OneDNNStorageProperties firstly for unittest covergae
 #ifdef PADDLE_WITH_DNNL
 struct OneDNNStorageProperties
diff --git a/paddle/phi/core/tensor_meta.cc b/paddle/phi/core/tensor_meta.cc
@@ -118,30 +118,19 @@ DDim DenseTensorMeta::calc_strides(const DDim& dims) {
   }
 }
 
-DenseTensorMeta::DenseTensorMeta() {
-  use_gpudnn = true;
-#ifdef PADDLE_WITH_XPU
-  scale_value = -1.0f;
-#endif
-}
+DenseTensorMeta::DenseTensorMeta() { use_gpudnn = true; }
 
 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,
@@ -151,9 +140,6 @@ 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,
@@ -164,9 +150,6 @@ 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) {
@@ -182,9 +165,6 @@ 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) {
@@ -200,9 +180,6 @@ DenseTensorMeta& DenseTensorMeta::operator=(const DenseTensorMeta& other) {
   } else {
     strides = other.strides;
   }
-#ifdef PADDLE_WITH_XPU
-  scale_value = other.scale_value;
-#endif
   return *this;
 }
 
@@ -220,9 +197,6 @@ 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
@@ -83,23 +83,13 @@ struct TEST_API 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/core/utils/type_info.cc b/paddle/phi/core/utils/type_info.cc
@@ -73,6 +73,10 @@ template class TypeInfoTraits<phi::DeviceContext, XPUContext>;
 template class TypeInfoTraits<phi::StorageProperties, OneDNNStorageProperties>;
 #endif
 
+#ifdef PADDLE_WITH_XPU
+template class TypeInfoTraits<phi::StorageProperties, XPUStorageProperties>;
+#endif
+
 template class TypeInfoTraits<phi::StorageProperties, NPUStorageProperties>;
 
 }  // namespace phi
diff --git a/paddle/phi/kernels/xpu/adamw_kernel.cc b/paddle/phi/kernels/xpu/adamw_kernel.cc
@@ -136,28 +136,38 @@ void AdamwDenseKernel(const Context& dev_ctx,
         moment2.numel());
     PADDLE_ENFORCE_XDNN_SUCCESS(r, "cast moment2 from fp16 to float");
 
-    // de-scale using meta's scale_value
+    // acquire xpu_scale_value
+    float moment1_scale_value = XPUStorageProperties::default_xpu_scale_value;
+    if (moment1.storage_properties_initialized()) {
+      moment1_scale_value =
+          moment1.storage_properties<XPUStorageProperties>().xpu_scale_value;
+    }
+    float moment2_scale_value = XPUStorageProperties::default_xpu_scale_value;
+    if (moment2.storage_properties_initialized()) {
+      moment2_scale_value =
+          moment2.storage_properties<XPUStorageProperties>().xpu_scale_value;
+    }
+
+    // de-scale using 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) {
+    if (moment1_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,
+                            1.0f / moment1_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) {
+    if (moment2_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,
+                            1.0f / moment2_scale_value,
                             0.0f);
       PADDLE_ENFORCE_XDNN_SUCCESS(r, "de-scale for moment2");
     }
@@ -267,10 +277,10 @@ void AdamwDenseKernel(const Context& dev_ctx,
                           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);
+    // write to moment1_out
+    std::unique_ptr<phi::StorageProperties> moment1_out_sp =
+        std::make_unique<phi::XPUStorageProperties>(moment1_scale_value);
+    moment1_out->set_storage_properties(std::move(moment1_out_sp));
 
     // for moment2
     float moment2_max = GetAbsMax<Context>(dev_ctx,
@@ -289,10 +299,10 @@ void AdamwDenseKernel(const Context& dev_ctx,
                           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);
+    // write to moment2_out
+    std::unique_ptr<phi::StorageProperties> moment2_out_sp =
+        std::make_unique<phi::XPUStorageProperties>(moment2_scale_value);
+    moment2_out->set_storage_properties(std::move(moment2_out_sp));
 
     // cast moment1 and moment2 output, from fp32 to fp16
     // int cast(Context* ctx, const TX* x, TY* y, int64_t len);
diff --git a/python/paddle/optimizer/optimizer.py b/python/paddle/optimizer/optimizer.py
@@ -323,6 +323,11 @@ def state_dict(self):
         for k, v in self._accumulators.items():
             for para_name, var_tmp in v.items():
                 state_dict[var_tmp.name] = var_tmp
+                # save scale value for xpu
+                if core.is_compiled_with_xpu():
+                    state_dict[
+                        var_tmp.name + ".SCALE_VALUE"
+                    ] = var_tmp.get_tensor().get_xpu_scale_value()
         # if has master weight and then save master weight
         if hasattr(self, "_master_weights"):
             if len(self._master_weights) != 0:
@@ -382,8 +387,15 @@ def set_state_dict(self, state_dict):
                 assert (
                     var_tmp.name in state_dict
                 ), f"optimizer Tensor {var_tmp.name} not found"
+
                 var = var_tmp.value()
                 tensor = var.get_tensor()
+                # load scale value for xpu
+                if core.is_compiled_with_xpu():
+                    tensor.set_xpu_scale_value(
+                        state_dict.get(var_tmp.name + ".SCALE_VALUE", -1.0)
+                    )
+
                 model_np = np.array(tensor)
 
                 load_para = state_dict[var_tmp.name]
@@ -909,6 +921,14 @@ def _add_accumulator(
                     ), f"Optimizer set error, {var_name} should in state dict"
                     var.set_value(self._accumulators_holder.pop(var_name))
 
+                    # load scale value for xpu
+                    if core.is_compiled_with_xpu():
+                        var.get_tensor().set_xpu_scale_value(
+                            self._accumulators_holder.get(
+                                var_name + ".SCALE_VALUE", -1.0
+                            )
+                        )
+
         self._accumulators[name][param.name] = var
         return var
 
diff --git a/test/xpu/test_adamw_fp16_xpu.py b/test/xpu/test_adamw_fp16_xpu.py
