[PIR] Fix 1 ut in PIR mode (test_imperative_optimizer_v2) (#65760)

SigureMo · web-flow · commit 62d727e4628c · 2024-07-09T10:27:10.000+08:00
diff --git a/paddle/fluid/pir/dialect/op_generator/api_gen.py b/paddle/fluid/pir/dialect/op_generator/api_gen.py
@@ -93,8 +93,8 @@
     {handle_optional_inputs}
     {in_combine}
     {compute_op}
-    {handle_optional_outputs}
     {set_null_type}
+    {handle_optional_outputs}
     {out_split}
     {return_result}"""
 
diff --git a/paddle/fluid/pir/dialect/op_generator/ops_api_gen.py b/paddle/fluid/pir/dialect/op_generator/ops_api_gen.py
@@ -133,6 +133,7 @@
     'assign_value',
     'share_data_',
     'onednn_to_paddle_layout',
+    'lars_momentum_',
     'lrn',
     'multi_gru',
     'matmul_with_flatten',
@@ -199,8 +200,6 @@
     'fusion_lstm',
     'fusion_seqpool_cvm_concat',
     'nce',
-    'lars_momentum',
-    'lars_momentum_',
     'lrn',
     'max_pool2d_v2',
     'partial_sum',
diff --git a/paddle/phi/ops/yaml/inconsistent/static_ops.yaml b/paddle/phi/ops/yaml/inconsistent/static_ops.yaml
@@ -428,6 +428,20 @@
     param: [x, num_hash, mod_by]
     data_type: x
 
+- op : lars_momentum_
+  args: (Tensor[] param, Tensor[] grad, Tensor[] velocity, Tensor[] learning_rate, Tensor[] master_param, float mu, float lars_coeff=0.001f, float[] lars_weight_decay={0.0005f}, float epsilon=0.0f, bool multi_precision=false, float rescale_grad=1.0f)
+  output: Tensor[](param_out){param.size()}, Tensor[](velocity_out){param.size()}, Tensor[](master_param_out){param.size()}
+  infer_meta:
+    func: LarsMomentumInferMeta
+    param: [param, velocity, learning_rate, grad, master_param, lars_weight_decay, mu, lars_coeff, epsilon, multi_precision, rescale_grad]
+  kernel:
+    func: lars_momentum
+    param: [param, velocity, learning_rate, grad, master_param, lars_weight_decay, mu, lars_coeff, epsilon, multi_precision, rescale_grad]
+    data_type: param
+  optional: master_param, master_param_out
+  inplace : (param -> param_out), (velocity -> velocity_out), (master_param -> master_param_out)
+  traits : pir::SideEffectTrait
+
 - op : less_equal
   args : (Tensor x, Tensor y)
   output : Tensor(out)
@@ -938,19 +952,6 @@
   optional: dropout1_seed, dropout2_seed, linear1_bias, linear2_bias, ln1_scale, ln1_bias, ln2_scale, ln2_bias, ln2_mean, ln2_variance, ln1_mean, ln1_variance, ln1_out
   backward: fused_feedforward_grad
 
-- op: lars_momentum
-  args: (Tensor[] param, Tensor[] grad, Tensor[] velocity, Tensor[] learning_rate, Tensor[] master_param, float mu, float lars_coeff=0.001f, float[] lars_weight_decay={0.0005f}, float epsilon=0.0f, bool multi_precision=false, float rescale_grad=1.0f)
-  output: Tensor[](param_out){param.size()}, Tensor[](velocity_out){param.size()}, Tensor[](master_param_out){param.size()}
-  infer_meta:
-    func: LarsMomentumInferMeta
-    param: [param, velocity, learning_rate, grad, master_param, lars_weight_decay, mu, lars_coeff, epsilon, multi_precision, rescale_grad]
-  kernel:
-    func: lars_momentum
-    param: [param, velocity, learning_rate, grad, master_param, lars_weight_decay, mu, lars_coeff, epsilon, multi_precision, rescale_grad]
-    data_type: param
-  optional: master_param, master_param_out
-  inplace : master_param -> master_param_out
-
 - op: moving_average_abs_max_scale
   args: (Tensor x, Tensor in_accum, Tensor in_state, float moving_rate=0.9f, bool is_test=false)
   output: Tensor(out), Tensor(out_scale), Tensor(out_state), Tensor(out_accum)
diff --git a/paddle/phi/ops/yaml/op_compat.yaml b/paddle/phi/ops/yaml/op_compat.yaml
@@ -4283,7 +4283,7 @@
   outputs :
     out : Out
 
-- op: lars_momentum
+- op: lars_momentum_ (lars_momentum)
   inputs:
     {param : Param, grad : Grad, velocity : Velocity, learning_rate : LearningRate, master_param : MasterParam}
   outputs :
diff --git a/python/paddle/incubate/optimizer/lars_momentum.py b/python/paddle/incubate/optimizer/lars_momentum.py
@@ -13,9 +13,12 @@
 # limitations under the License.
 import warnings
 
-from paddle import _legacy_C_ops
+from paddle import _C_ops, _legacy_C_ops
 from paddle.base import framework
-from paddle.base.framework import in_dygraph_mode
+from paddle.framework import (
+    in_dynamic_mode,
+    in_pir_mode,
+)
 from paddle.optimizer import Optimizer
 
 
@@ -81,6 +84,7 @@ class LarsMomentumOptimizer(Optimizer):
             ...     feed={"inp": np_inp},
             ...     fetch_list=[out.name])
     """
+
     _velocity_acc_str = "velocity"
 
     def __init__(
@@ -184,7 +188,7 @@ def _append_optimize_op(self, block, param_and_grad):
             inputs["MasterParam"] = master_weight
             outputs["MasterParamOut"] = master_weight
 
-        if in_dygraph_mode():
+        if in_dynamic_mode():
             tmp, tmp2 = _legacy_C_ops.lars_momentum(
                 [param_and_grad[0]],
                 [param_and_grad[1]],
@@ -205,6 +209,21 @@ def _append_optimize_op(self, block, param_and_grad):
                 "rescale_grad",
                 self._rescale_grad,
             )
+        elif in_pir_mode():
+            _, _, _ = _C_ops.lars_momentum_(
+                [param_and_grad[0]],
+                [param_and_grad[1]],
+                [velocity_acc],
+                [lr],
+                master_weight,
+                self._momentum,
+                self._lars_coeff,
+                [_lars_weight_decay],
+                self._epsilon,
+                find_master,
+                self._rescale_grad,
+            )
+            return None
         else:
             # create the momentum optimize op
             momentum_op = block.append_op(
diff --git a/python/paddle/optimizer/adagrad.py b/python/paddle/optimizer/adagrad.py
@@ -17,6 +17,11 @@
 import warnings
 from typing import TYPE_CHECKING, Sequence
 
+from paddle import _C_ops
+from paddle.framework import (
+    in_dynamic_or_pir_mode,
+)
+
 from ..base import framework
 from .optimizer import Optimizer
 
@@ -122,6 +127,7 @@ class Adagrad(Optimizer):
             >>> adagrad.clear_grad()
 
     """
+
     type: str
     initial_accumulator_value: float
     _moment_acc_str = "moment"
@@ -210,29 +216,44 @@ def _append_optimize_op(self, block, param_and_grad):
             else None
         )
 
-        # Create the adagrad optimizer op
-        inputs = {
-            "Param": param_and_grad[0],
-            "Grad": param_and_grad[1],
-            "Moment": moment_acc,
-            "LearningRate": self._create_param_lr(param_and_grad),
-        }
-
-        outputs = {"ParamOut": param_and_grad[0], "MomentOut": moment_acc}
-
-        if find_master:
-            inputs["MasterParam"] = master_weight
-            outputs["MasterParamOut"] = master_weight
-
-        adagrad_op = block.append_op(
-            type=self.type,
-            inputs=inputs,
-            outputs=outputs,
-            attrs={"epsilon": self._epsilon, "multi_precision": find_master},
-            stop_gradient=True,
-        )
+        if in_dynamic_or_pir_mode():
+            _, _, _ = _C_ops.adagrad_(
+                param_and_grad[0],
+                param_and_grad[1],
+                moment_acc,
+                self._create_param_lr(param_and_grad),
+                master_weight if find_master else None,
+                self._epsilon,
+                find_master,
+            )
+            return None
+        else:
+            # Create the adagrad optimizer op
+            inputs = {
+                "Param": param_and_grad[0],
+                "Grad": param_and_grad[1],
+                "Moment": moment_acc,
+                "LearningRate": self._create_param_lr(param_and_grad),
+            }
+
+            outputs = {"ParamOut": param_and_grad[0], "MomentOut": moment_acc}
+
+            if find_master:
+                inputs["MasterParam"] = master_weight
+                outputs["MasterParamOut"] = master_weight
+
+            adagrad_op = block.append_op(
+                type=self.type,
+                inputs=inputs,
+                outputs=outputs,
+                attrs={
+                    "epsilon": self._epsilon,
+                    "multi_precision": find_master,
+                },
+                stop_gradient=True,
+            )
 
-        return adagrad_op
+            return adagrad_op
 
     def _update_param_group(self, parameters):
         self._epsilon = parameters.get('epsilon', self._default_dict['epsilon'])
diff --git a/python/paddle/optimizer/adamax.py b/python/paddle/optimizer/adamax.py
@@ -264,7 +264,7 @@ def _append_optimize_op(self, block, param_and_grad):
         beta1_pow_acc = self._get_accumulator_master(
             self._beta1_pow_acc_str, param_and_grad[0]
         )
-        if framework.in_dygraph_mode():
+        if framework.in_dynamic_or_pir_mode():
             _C_ops.adamax_(
                 param_and_grad[0],
                 param_and_grad[1],
@@ -329,6 +329,12 @@ def _finish_update(self, block, parameters_and_grads):
                             beta1_pow_acc, self._beta1, 0.0, True
                         )
                         beta1_pow_acc.copy_(tmp, False)
+                elif framework.in_pir_mode():
+                    with param.block.program._optimized_guard([param, grad]):
+                        beta1_pow_acc = self._get_accumulator_master(
+                            self._beta1_pow_acc_str, param
+                        )
+                        _C_ops.scale_(beta1_pow_acc, self._beta1, 0.0, True)
                 else:
                     with param.block.program._optimized_guard(
                         [param, grad]
diff --git a/test/deprecated/legacy_test/CMakeLists.txt b/test/deprecated/legacy_test/CMakeLists.txt
@@ -659,14 +659,12 @@ set_tests_properties(test_imperative_ptb_rnn PROPERTIES TIMEOUT 120)
 set_tests_properties(test_masked_select_op PROPERTIES TIMEOUT 120)
 set_tests_properties(test_sigmoid_cross_entropy_with_logits_op
                      PROPERTIES TIMEOUT 120)
-set_tests_properties(test_imperative_optimizer_v2 PROPERTIES TIMEOUT 150)
 set_tests_properties(test_sgd_op_deprecated PROPERTIES TIMEOUT 250)
 set_tests_properties(test_generator_dataloader_deprecated PROPERTIES TIMEOUT
                                                                      120)
 set_tests_properties(test_reduce_op PROPERTIES TIMEOUT 500)
 set_tests_properties(test_program_prune_backward_deprecated PROPERTIES TIMEOUT
                                                                        120)
-set_tests_properties(test_imperative_optimizer_v2 PROPERTIES TIMEOUT 250)
 set_tests_properties(test_decoupled_py_reader_deprecated PROPERTIES TIMEOUT 120)
 set_tests_properties(test_fuse_bn_act_pass_deprecated PROPERTIES TIMEOUT 120)
 set_tests_properties(test_conv2d_api_deprecated PROPERTIES TIMEOUT 120)
diff --git a/test/legacy_test/CMakeLists.txt b/test/legacy_test/CMakeLists.txt
@@ -872,6 +872,7 @@ set_tests_properties(test_group_norm_op PROPERTIES TIMEOUT 1000)
 set_tests_properties(test_imperative_transformer_sorted_gradient
                      PROPERTIES TIMEOUT 120)
 set_tests_properties(test_imperative_optimizer PROPERTIES TIMEOUT 250)
+set_tests_properties(test_imperative_optimizer_v2 PROPERTIES TIMEOUT 250)
 set_tests_properties(test_activation_op PROPERTIES TIMEOUT 270)
 set_tests_properties(test_normal PROPERTIES TIMEOUT 120)
 set_tests_properties(test_conv2d_op PROPERTIES TIMEOUT 120)
diff --git a/test/legacy_test/test_imperative_optimizer_v2.py b/test/legacy_test/test_imperative_optimizer_v2.py
@@ -20,6 +20,7 @@
 
 import paddle
 from paddle import base
+from paddle.autograd.backward_utils import ValueDict
 from paddle.base import core
 from paddle.distributed.fleet.meta_optimizers import DGCMomentumOptimizer
 from paddle.pir_utils import test_with_pir_api
@@ -28,6 +29,30 @@
 # In dygraph, don't support ModelAverage, DGCMomentumOptimizer, ExponentialMovingAverage, PipelineOptimizer, LookaheadOptimizer, RecomputeOptimizer.
 
 
+def create_parameter_mapping(startup_program, main_program):
+    startup_params = {}
+    main_params = {}
+    parameter_mapping = ValueDict()
+    for op in startup_program.global_block().ops:
+        if op.name() == "builtin.set_parameter":
+            name = op.attrs()["parameter_name"]
+            param = op.operand(0).source()
+            startup_params[name] = param
+
+    for op in main_program.global_block().ops:
+        if op.name() == "builtin.parameter":
+            name = op.attrs()["parameter_name"]
+            param = op.result(0)
+            main_params[name] = param
+
+    assert len(startup_params) == len(main_params)
+    for name, startup_param in startup_params.items():
+        assert name in main_params
+        main_param = main_params[name]
+        parameter_mapping[main_param] = startup_param
+    return parameter_mapping
+
+
 class MLP(paddle.nn.Layer):
     def __init__(self, param_attr=None, bias_attr=None):
         super().__init__()
@@ -73,7 +98,13 @@ def _check_exception(self, exception_message, place=None):
         try:
             paddle.disable_static()
             paddle.seed(seed)
-            paddle.framework.random._manual_program_seed(seed)
+            if paddle.framework.use_pir_api():
+                with paddle.pir_utils.OldIrGuard():
+                    # Note: dygraph use self.main_program.global_block().create_parameter(), it's need manual seed to old Program
+                    paddle.framework.random._manual_program_seed(seed)
+                paddle.framework.random._manual_program_seed(seed)
+            else:
+                paddle.framework.random._manual_program_seed(seed)
             mlp = MLP()
             optimizer = self.get_optimizer_dygraph(
                 parameter_list=mlp.parameters()
@@ -96,7 +127,13 @@ def _check_mlp(self, place=None):
 
         paddle.disable_static(place)
         paddle.seed(seed)
-        paddle.framework.random._manual_program_seed(seed)
+        if paddle.framework.use_pir_api():
+            with paddle.pir_utils.OldIrGuard():
+                # Note: dygraph use self.main_program.global_block().create_parameter(), it's need manual seed to old Program
+                paddle.framework.random._manual_program_seed(seed)
+            paddle.framework.random._manual_program_seed(seed)
+        else:
+            paddle.framework.random._manual_program_seed(seed)
 
         mlp = MLP()
         optimizer = self.get_optimizer_dygraph(parameter_list=mlp.parameters())
@@ -150,7 +187,13 @@ def _check_mlp(self, place=None):
         paddle.enable_static()
         with new_program_scope():
             paddle.seed(seed)
-            paddle.framework.random._manual_program_seed(seed)
+            if paddle.framework.use_pir_api():
+                with paddle.pir_utils.OldIrGuard():
+                    # Note: dygraph use self.main_program.global_block().create_parameter(), it's need manual seed to old Program
+                    paddle.framework.random._manual_program_seed(seed)
+                paddle.framework.random._manual_program_seed(seed)
+            else:
+                paddle.framework.random._manual_program_seed(seed)
 
             if place is None:
                 place = (
@@ -181,16 +224,30 @@ def _check_mlp(self, place=None):
             # initialize params and fetch them
             static_param_init_value = {}
             static_param_name_list = []
+            static_params = []
             for param in mlp.parameters():
                 static_param_name_list.append(param.name)
+                static_params.append(param)
+
+            if paddle.framework.use_pir_api():
+                parameter_mapping = create_parameter_mapping(
+                    paddle.static.default_startup_program(),
+                    paddle.static.default_main_program(),
+                )
+                startup_params = [
+                    parameter_mapping[param] for param in static_params
+                ]
+            else:
+                startup_params = static_params
 
             out = exe.run(
-                base.default_startup_program(),
-                fetch_list=static_param_name_list,
+                paddle.static.default_startup_program(),
+                fetch_list=startup_params,
             )
 
-            for i in range(len(static_param_name_list)):
-                static_param_init_value[static_param_name_list[i]] = out[i]
+            for i in range(len(static_params)):
+                param_name = static_param_name_list[i]
+                static_param_init_value[param_name] = out[i]
 
             for batch_id, data in enumerate(train_reader()):
                 if batch_id >= self.batch_num:
@@ -205,8 +262,8 @@ def _check_mlp(self, place=None):
                     .reshape([128, 1])
                 )
 
-                fetch_list = [avg_loss.name]
-                fetch_list.extend(static_param_name_list)
+                fetch_list = [avg_loss]
+                fetch_list.extend(static_params)
                 out = exe.run(
                     base.default_main_program(),
                     feed={"pixel": static_x_data, "label": y_data},
