[XPU] only save adamw scale_value when flag is on (#67077)

houj04 · web-flow · commit a6f16a158569 · 2024-08-07T11:12:29.000+08:00
* [XPU] only save adamw scale_value when flag is on

* fix ut.
diff --git a/python/paddle/optimizer/optimizer.py b/python/paddle/optimizer/optimizer.py
@@ -365,9 +365,13 @@ def state_dict(self) -> dict[str, Tensor]:
                     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()
+                        xpu_adamw_moment_dtype = os.getenv(
+                            "xpu_adamw_moment_dtype", default="fp32"
+                        )
+                        if xpu_adamw_moment_dtype == "fp16":
+                            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:
@@ -433,9 +437,13 @@ def set_state_dict(self, state_dict: dict[str, Tensor]) -> None:
                 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)
+                    xpu_adamw_moment_dtype = os.getenv(
+                        "xpu_adamw_moment_dtype", default="fp32"
                     )
+                    if xpu_adamw_moment_dtype == "fp16":
+                        tensor.set_xpu_scale_value(
+                            state_dict.get(var_tmp.name + ".SCALE_VALUE", -1.0)
+                        )
                 var.set_value(state_dict[var_tmp.name])
 
     def get_opti_var_name_list(self) -> list[str]:
@@ -1023,11 +1031,15 @@ def _add_accumulator(
 
                     # 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
-                            )
+                        xpu_adamw_moment_dtype = os.getenv(
+                            "xpu_adamw_moment_dtype", default="fp32"
                         )
+                        if xpu_adamw_moment_dtype == "fp16":
+                            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
@@ -28,9 +28,14 @@ def test_tensor_scale_value(self):
         # read modified scale_value
         self.assertEqual(x.get_tensor().get_xpu_scale_value(), -1.25)
 
-    def test_state_dict(self):
-        os.environ["xpu_adamw_moment_dtype"] = "fp16"
-
+    def _test_state_dict(self):
+        xpu_adamw_moment_dtype = os.getenv(
+            "xpu_adamw_moment_dtype", default="fp32"
+        )
+        if xpu_adamw_moment_dtype == "fp16":
+            use_fp16 = True
+        else:
+            use_fp16 = False
         linear = paddle.nn.Linear(10, 10)
         inp = paddle.rand([10, 10], dtype="float32")
         out = linear(inp)
@@ -51,8 +56,24 @@ def test_state_dict(self):
 
         # read scale_value in state dict
         state_dict_1 = adam.state_dict()
-        self.assertTrue("linear_0.w_0_moment1_0.SCALE_VALUE" in state_dict_1)
-        self.assertTrue("linear_0.b_0_moment1_0.SCALE_VALUE" in state_dict_1)
+        if use_fp16:
+            self.assertTrue(
+                "linear_0.w_0_moment1_0.SCALE_VALUE" in state_dict_1
+            )
+            self.assertTrue(
+                "linear_0.b_0_moment1_0.SCALE_VALUE" in state_dict_1
+            )
+        else:
+            self.assertTrue(
+                "linear_0.w_0_moment1_0.SCALE_VALUE" not in state_dict_1
+            )
+            self.assertTrue(
+                "linear_0.b_0_moment1_0.SCALE_VALUE" not in state_dict_1
+            )
+
+        if not use_fp16:
+            # do not need "overwrite" and "check overwritten value" below
+            return
 
         # overwrite scale_value
         state_dict_1["linear_0.w_0_moment1_0.SCALE_VALUE"] = 0.75
@@ -70,6 +91,12 @@ def test_state_dict(self):
             state_dict_2["linear_0.b_0_moment1_0.SCALE_VALUE"], 12.3125
         )
 
+    def test_state_dict(self):
+        os.environ["xpu_adamw_moment_dtype"] = "fp16"
+        self._test_state_dict()
+        os.environ["xpu_adamw_moment_dtype"] = "fp32"
+        self._test_state_dict()
+
 
 if __name__ == '__main__':
     paddle.disable_static()
