[PIR]Migrate AdamW into pir (#58946)

0x45f · web-flow · commit b0c684c8ba1b · 2023-11-15T11:01:53.000+08:00
diff --git a/python/paddle/optimizer/adamw.py b/python/paddle/optimizer/adamw.py
@@ -17,11 +17,19 @@
 from collections.abc import Callable
 
 import paddle
+from paddle import pir
+from paddle.base.libpaddle import DataType
+from paddle.pir import OpResult
 
 from .. import _C_ops
 from ..base import core, framework
 from ..base.dygraph import base as imperative_base
-from ..base.framework import Parameter, Variable
+from ..base.framework import (
+    Parameter,
+    Variable,
+    in_dynamic_or_pir_mode,
+    in_pir_mode,
+)
 from ..nn.clip import GradientClipBase
 from .lr import LRScheduler
 from .optimizer import Optimizer
@@ -174,7 +182,7 @@ def __init__(
         if not 0 <= epsilon:
             raise ValueError("Invaild value of epsilon, expect epsilon >= 0.")
         if not isinstance(weight_decay, float) and not isinstance(
-            weight_decay, framework.Variable
+            weight_decay, (framework.Variable, OpResult)
         ):
             raise TypeError("weight_decay should be float or Tensor.")
         if lr_ratio is not None:
@@ -307,7 +315,7 @@ def _add_param_group(self, param_group):
             different optimization options.
         """
         params = param_group['params']
-        if isinstance(params, Parameter):
+        if isinstance(params, (Parameter, pir.core.ParameterMeta)):
             param_group['params'] = [params]
         elif isinstance(params, set):
             raise TypeError(
@@ -340,7 +348,9 @@ def _add_param_group(self, param_group):
     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
+            acc_dtype = (
+                DataType.FLOAT32 if in_pir_mode() else core.VarDesc.VarType.FP32
+            )
         if core.is_compiled_with_xpu():
             import os
 
@@ -365,7 +375,7 @@ def _add_moments_pows(self, p):
             param=p,
             dtype=acc_dtype,
             fill_value=0.9
-            if isinstance(self._beta1, Variable)
+            if isinstance(self._beta1, (Variable, OpResult))
             else self._beta1,
             shape=[1],
             type=core.VarDesc.VarType.LOD_TENSOR,
@@ -376,15 +386,15 @@ def _add_moments_pows(self, p):
             param=p,
             dtype=acc_dtype,
             fill_value=0.999
-            if isinstance(self._beta2, Variable)
+            if isinstance(self._beta2, (Variable, OpResult))
             else self._beta2,
             shape=[1],
             type=core.VarDesc.VarType.LOD_TENSOR,
             device='cpu',
         )
 
     def _create_accumulators(self, block, parameters):
-        assert isinstance(block, framework.Block)
+        assert isinstance(block, (framework.Block, pir.Block))
         if isinstance(parameters, dict):
             parameters = self._update_param_group(parameters)
 
@@ -409,7 +419,7 @@ def _create_accumulators(self, block, parameters):
             self._already_create_accumulater.add(p.name)
 
     def _append_optimize_op(self, block, param_and_grad):
-        assert isinstance(block, framework.Block)
+        assert isinstance(block, (framework.Block, pir.Block))
         if isinstance(param_and_grad, dict):
             param_and_grad = self._update_param_group(param_and_grad)
         param, grad = param_and_grad
@@ -445,7 +455,7 @@ def _append_optimize_op(self, block, param_and_grad):
         lr = self._create_param_lr(param_and_grad)
 
         # create the adamw optimize op
-        if framework.in_dygraph_mode():
+        if in_dynamic_or_pir_mode():
             lr_ratio_ = (
                 1.0
                 if self._lr_ratio is None
diff --git a/python/paddle/optimizer/optimizer.py b/python/paddle/optimizer/optimizer.py
@@ -760,14 +760,14 @@ def _create_master_weight(self, param):
         if param.name in self._master_weights:
             var = self._master_weights[param.name]
         else:
-            assert isinstance(self.helper, LayerHelper)
             var_name = self._gen_master_weight_var_name(param)
             if in_pir_mode():
                 var = paddle.cast(param, 'float32')
             elif framework.in_dygraph_mode():
                 var = paddle.cast(param, 'float32')
                 var.name = var_name
             else:
+                assert isinstance(self.helper, LayerHelper)
                 var = paddle.static.create_global_var(
                     name=var_name,
                     shape=param.shape,
diff --git a/test/legacy_test/test_adamw_op.py b/test/legacy_test/test_adamw_op.py
@@ -151,7 +151,7 @@ def setUp(self):
         }
 
     def test_check_output(self):
-        self.check_output()
+        self.check_output(check_pir=True)
 
 
 @unittest.skipIf(
@@ -208,7 +208,7 @@ def setUp(self):
         }
 
     def test_check_output(self):
-        self.check_output_with_place(core.CUDAPlace(0))
+        self.check_output_with_place(core.CUDAPlace(0), check_pir=True)
 
 
 class TestAdamWOp(unittest.TestCase):
@@ -278,6 +278,47 @@ def test_adamw_op(self):
         assert rets[0] is not None
         paddle.disable_static()
 
+    def test_pir_adam_op(self):
+        with paddle.pir_utils.IrGuard():
+            place = base.CPUPlace()
+            shape = [2, 3, 8, 8]
+            exe = base.Executor(place)
+            train_prog = paddle.static.Program()
+            startup = paddle.static.Program()
+            with paddle.static.program_guard(train_prog, startup):
+                with base.unique_name.guard():
+                    data = paddle.static.data(name="data", shape=shape)
+                    conv_layer = paddle.nn.Conv2D(3, 8, 3)
+                    conv = conv_layer(data)
+                    loss = paddle.mean(conv)
+
+                    beta1 = paddle.pir.core.create_parameter(
+                        'float32',
+                        [1],
+                        initializer=paddle.nn.initializer.Constant(0.85),
+                    )
+                    beta2 = paddle.pir.core.create_parameter(
+                        'float32',
+                        [1],
+                        initializer=paddle.nn.initializer.Constant(0.95),
+                    )
+                    betas = [beta1, beta2]
+                    opt = paddle.optimizer.AdamW(
+                        learning_rate=1e-5,
+                        beta1=beta1,
+                        beta2=beta2,
+                        weight_decay=0.01,
+                        epsilon=1e-8,
+                    )
+                    opt.minimize(loss)
+
+            exe.run(startup)
+            data_np = np.random.random(shape).astype('float32')
+            rets = exe.run(
+                train_prog, feed={"data": data_np}, fetch_list=[loss]
+            )
+            assert rets[0] is not None
+
     def test_adamw_op_invalid_input(self):
         paddle.disable_static()
         linear = paddle.nn.Linear(10, 10)
@@ -753,6 +794,7 @@ def get_numpy_output(param, grad, moment1, moment2, lr_ratio, t):
             np.testing.assert_allclose(linear2.weight.numpy(), fc2_w, rtol=1e-6)
             np.testing.assert_allclose(linear2.bias.numpy(), fc2_b, rtol=1e-6)
 
+    # @test_with_pir_api
     def test_adamw_op(self):
         paddle.enable_static()
         place = base.CUDAPlace(0)
@@ -763,9 +805,9 @@ def test_adamw_op(self):
         weight_decay = 0.01
         epsilon = 1e-8
 
-        train_prog = base.Program()
-        startup = base.Program()
-        with base.program_guard(train_prog, startup):
+        train_prog = paddle.static.Program()
+        startup = paddle.static.Program()
+        with paddle.static.program_guard(train_prog, startup):
             with base.unique_name.guard():
                 x = paddle.static.data(
                     name='x', shape=[None, 10], dtype='float32'
