[PIR AMP]Adapt test/amp uts in PIR (#62745)

0x45f · web-flow · commit 294b3cf8f63d · 2024-03-20T15:28:02.000+08:00
diff --git a/paddle/fluid/framework/new_executor/pir_interpreter.cc b/paddle/fluid/framework/new_executor/pir_interpreter.cc
@@ -81,6 +81,7 @@ COMMON_DECLARE_bool(dynamic_static_unified_comm);
 
 COMMON_DECLARE_bool(enable_pir_in_executor);
 COMMON_DECLARE_bool(enable_pir_in_executor_trace_run);
+COMMON_DECLARE_int32(low_precision_op_list);
 
 #define CREATE_INSTR(instr_name)                                   \
   vec_instruction_base_.emplace_back(std::make_unique<instr_name>( \
@@ -89,6 +90,21 @@ COMMON_DECLARE_bool(enable_pir_in_executor_trace_run);
 namespace paddle {
 namespace framework {
 
+void RecordLowPrecisionOp(const InstructionBase* instr_node) {
+  if (FLAGS_low_precision_op_list) {
+    std::string op_name = instr_node->Name();
+    ::pir::Operation* op = instr_node->Operation();
+    if (op->HasAttribute("kernel_key")) {
+      phi::KernelKey kernel_key =
+          op->attribute("kernel_key")
+              .dyn_cast<paddle::dialect::KernelAttribute>()
+              .data();
+      phi::KernelFactory::Instance().AddToLowPrecisionKernelList(
+          op_name, kernel_key.dtype());
+    }
+  }
+}
+
 PirInterpreter::PirInterpreter(const platform::Place& place,
                                const std::vector<std::string>& fetch_var_names,
                                const ::pir::Block* ir_block,
@@ -1735,6 +1751,9 @@ void PirInterpreter::RunInstructionBase(InstructionBase* instr_node) {
       }
     }
 #endif
+
+    RecordLowPrecisionOp(instr_node);
+
     VLOG(2) << "\nbegin: " << __func__ << " OP id:" << instr_node->Id()
             << " name:" << instr_node->Name() << " type:"
             << (instr_node->KernelType() == OpFuncType::kCpuSync
diff --git a/paddle/fluid/pir/dialect/op_generator/api_gen.py b/paddle/fluid/pir/dialect/op_generator/api_gen.py
@@ -105,7 +105,7 @@
         auto op_name = phi::TransToFluidOpName("{op_name}");
         paddle::small_vector<std::vector<pir::Value>, egr::kSlotSmallVectorSize> amp_values_vector = {{ {no_optional_inputs} }};
         {optional_inputs}
-        auto amp_dst_dtype = paddle::imperative::GetAmpDestDtype("{op_name}", amp_values_vector);
+        auto amp_dst_dtype = paddle::imperative::GetAmpDestDtype(op_name, amp_values_vector);
         {new_inputs}
         {{
             paddle::imperative::AutoCastGuard guard(egr::Controller::Instance().GetCurrentAmpAttrs(), paddle::imperative::AmpLevel::O0);
diff --git a/python/paddle/amp/auto_cast.py b/python/paddle/amp/auto_cast.py
@@ -257,6 +257,7 @@ def _pir_transform(t, dtype):
                 break
     main.set_parameters_from(startup)
     with paddle.static.program_guard(main):
+        paddle.pir.reset_insertion_point_to_start()
         block = main.global_block()
         cast_param = paddle._pir_ops.parameter(t.name)
         cast_param.stop_gradient = t.stop_gradient
diff --git a/python/paddle/amp/debugging.py b/python/paddle/amp/debugging.py
@@ -21,7 +21,6 @@
 import paddle
 from paddle import _C_ops
 from paddle.base import core
-from paddle.base.framework import dygraph_only
 
 from ..framework import LayerHelper, in_dynamic_or_pir_mode
 
@@ -455,7 +454,6 @@ def _print_operator_stats(op_count_dict):
     print("<{:-^120}>\n".format(" op count: " + str(total_ops) + " "))
 
 
-@dygraph_only
 def enable_operator_stats_collection():
     """
     Enable to collect the number of operators for different data types.
@@ -494,7 +492,6 @@ def enable_operator_stats_collection():
     paddle.set_flags({'FLAGS_low_precision_op_list': 1})
 
 
-@dygraph_only
 def disable_operator_stats_collection():
     """
     Disable the collection the number of operators for different data types.
@@ -535,7 +532,6 @@ def disable_operator_stats_collection():
     paddle.set_flags({'FLAGS_low_precision_op_list': 0})
 
 
-@dygraph_only
 @contextlib.contextmanager
 def collect_operator_stats():
     """
diff --git a/python/paddle/amp/grad_scaler.py b/python/paddle/amp/grad_scaler.py
@@ -117,6 +117,8 @@ def __init__(
 
         self._enable = enable
         self._use_dynamic_loss_scaling = False
+        self._init_loss_scaling = 1.0
+        self._scale = None
 
         if self._enable:
             assert incr_ratio > 1.0, "The incr_ratio must be > 1.0."
@@ -206,6 +208,7 @@ def scale(self, var):
         ):
             self._enable = False
             self._use_dynamic_loss_scaling = False
+            self._init_loss_scaling = 1.0
             warnings.warn(
                 'It is not recommended to use dynamic loss scaling for %s, so GradScaler is disable by default.'
                 % (amp_global_state().amp_dtype)
diff --git a/python/paddle/optimizer/adadelta.py b/python/paddle/optimizer/adadelta.py
@@ -15,10 +15,10 @@
 import warnings
 
 from paddle import _C_ops
+from paddle.base.framework import in_dynamic_or_pir_mode
 
 from ..base import framework
 from ..base.dygraph import no_grad
-from ..framework import in_dynamic_mode
 from .optimizer import Optimizer
 
 __all__ = []
@@ -190,7 +190,7 @@ def _append_optimize_op(self, block, param_and_grad):
             else None
         )
 
-        if in_dynamic_mode():
+        if in_dynamic_or_pir_mode():
             with no_grad():
                 _C_ops.adadelta_(
                     param_and_grad[0],
diff --git a/python/paddle/optimizer/adam.py b/python/paddle/optimizer/adam.py
@@ -327,6 +327,9 @@ def _append_optimize_op(self, block, param_and_grad):
                 if not isinstance(self._beta2, Variable)
                 else self._beta2.item(0)
             )
+            found_inf = (
+                self._get_auxiliary_var('found_inf') if in_pir_mode() else None
+            )
 
             _, _, _, _, _, _ = _C_ops.adam_(
                 param_and_grad[0],
@@ -337,7 +340,7 @@ def _append_optimize_op(self, block, param_and_grad):
                 beta1_pow_acc,
                 beta2_pow_acc,
                 master_weight,
-                None,
+                found_inf,
                 _beta1,
                 _beta2,
                 self._epsilon,
diff --git a/python/paddle/optimizer/adamw.py b/python/paddle/optimizer/adamw.py
@@ -473,6 +473,10 @@ def _append_optimize_op(self, block, param_and_grad):
                 else self._beta2.item(0)
             )
 
+            found_inf = (
+                self._get_auxiliary_var('found_inf') if in_pir_mode() else None
+            )
+
             _, _, _, _, _, _ = _C_ops.adamw_(
                 param_and_grad[0],
                 param_and_grad[1],
@@ -482,7 +486,7 @@ def _append_optimize_op(self, block, param_and_grad):
                 beta1_pow_acc,
                 beta2_pow_acc,
                 master_weight,
-                None,
+                found_inf,
                 _beta1,
                 _beta2,
                 self._epsilon,
diff --git a/python/paddle/optimizer/optimizer.py b/python/paddle/optimizer/optimizer.py
@@ -813,10 +813,13 @@ def get_param_from_startup(startup, name):
                     )
                     var = paddle.cast(startup_param, 'float32')
                     var.persistable = True
-                    paddle._pir_ops.set_parameter(var, var_name)
-                main_program.set_parameters_from(startup_program)
+                    paddle._pir_ops.set_persistable_value(var, var_name)
                 with paddle.static.program_guard(main_program):
-                    var = paddle._pir_ops.parameter(var_name)
+                    paddle.pir.reset_insertion_point_to_start()
+                    var = paddle.static.data(
+                        var_name, var.shape, var.dtype, core.Place()
+                    )
+                    var.persistable = True
             elif framework.in_dygraph_mode():
                 var = paddle.cast(param, 'float32')
                 var.name = var_name
@@ -848,21 +851,28 @@ def _gen_master_weight_var_name(self, param):
 
     def _create_master_grad(self, grad):
         assert self._is_dtype_fp16_or_bf16(grad.dtype)
-        if grad.name in self._master_grads:
-            var = self._master_grads[grad.name]
+        if in_pir_mode():
+            if grad in self._master_grads:
+                var = self._master_grads[grad]
+            else:
+                var = paddle.cast(grad, 'float32')
+                self._master_grads[grad] = var
         else:
-            var_name = grad.name + "_fp32_master"
-            var_name = unique_name.generate(var_name)
-            var = grad.block.create_var(
-                name=var_name,
-                shape=grad.shape,
-                value=0,
-                dtype='float32',
-                lod_level=grad.lod_level,
-                persistable=grad.persistable,
-                is_data=grad.is_data,
-            )
-            self._master_grads[grad.name] = var
+            if grad.name in self._master_grads:
+                var = self._master_grads[grad.name]
+            else:
+                var_name = grad.name + "_fp32_master"
+                var_name = unique_name.generate(var_name)
+                var = grad.block.create_var(
+                    name=var_name,
+                    shape=grad.shape,
+                    value=0,
+                    dtype='float32',
+                    lod_level=grad.lod_level,
+                    persistable=grad.persistable,
+                    is_data=grad.is_data,
+                )
+                self._master_grads[grad.name] = var
         return var
 
     def _create_accumulators(self, block, parameters):
diff --git a/python/paddle/static/amp/decorator.py b/python/paddle/static/amp/decorator.py
@@ -485,8 +485,7 @@ def _append_cast_to_master_grad_op(self, param_grads):
             for p, g in param_grads:
                 if g not in self._optimizer._master_grads:
                     if self._optimizer._is_dtype_fp16_or_bf16(g.dtype):
-                        master_g = paddle.cast(g, 'float32')
-                        self._optimizer._master_grads[g] = master_g
+                        master_g = self._optimizer._create_master_grad(g)
                         params_master_grads.append((p, master_g))
                     else:
                         params_master_grads.append((p, g))
diff --git a/test/amp/test_amp_api.py b/test/amp/test_amp_api.py
@@ -30,13 +30,14 @@
     "run test when gpu's compute capability is at least 7.0.",
 )
 class TestAutoCast(AmpTestBase):
-    def setUp(self):
+    def init_net(self):
         self._conv = paddle.nn.Conv2D(
             in_channels=1, out_channels=6, kernel_size=3, bias_attr=False
         )
         self._linear = paddle.nn.Linear(in_features=4, out_features=4)
 
     def test_amp_OD_level(self):
+        self.init_net()
         with paddle.amp.auto_cast(level='OD'):
             out1 = self._conv(paddle.rand(shape=[1, 1, 6, 6], dtype='float32'))
             out2 = out1 + paddle.rand(shape=out1.shape, dtype='float16')
@@ -46,6 +47,23 @@ def test_amp_OD_level(self):
         self.assertEqual(out2.dtype, paddle.float32)
         self.assertEqual(out3.dtype, paddle.float32)
 
+    def test_pir_amp_OD_level(self):
+        with paddle.pir_utils.IrGuard():
+            with paddle.static.program_guard(
+                paddle.static.Program(), paddle.static.Program()
+            ):
+                self.init_net()
+                with paddle.amp.auto_cast(level='OD'):
+                    out1 = self._conv(
+                        paddle.rand(shape=[1, 1, 6, 6], dtype='float32')
+                    )
+                    out2 = out1 + paddle.rand(shape=out1.shape, dtype='float16')
+                    out3 = self._linear(out2)
+
+                self.assertEqual(out1.dtype, core.DataType.FLOAT16)
+                self.assertEqual(out2.dtype, core.DataType.FLOAT32)
+                self.assertEqual(out3.dtype, core.DataType.FLOAT32)
+
 
 class SimpleConvNet(nn.Layer):
     def __init__(self):
@@ -169,6 +187,49 @@ def test_amp_grad_scaler(self):
         self.assertTrue('scale' not in op_list)
         self.assertTrue('check_finite_and_unscale' not in op_list)
 
+    def test_pir_amp_grad_scaler(self):
+        with paddle.pir_utils.IrGuard():
+            startup = paddle.static.Program()
+            main = paddle.static.Program()
+            with paddle.static.program_guard(main, startup):
+                model = paddle.nn.Conv2D(3, 2, 3)
+                optimizer = paddle.optimizer.SGD(
+                    learning_rate=0.01, parameters=model.parameters()
+                )
+                model, optimizer = paddle.amp.decorate(
+                    models=model,
+                    optimizers=optimizer,
+                )
+                scaler = paddle.amp.GradScaler()
+                data = paddle.static.data('data', [1, 3, 8, 8], dtype='float32')
+
+                with paddle.amp.auto_cast(
+                    custom_black_list=['conv2d'], dtype='bfloat16'
+                ):
+                    out = model(data)
+                    loss = out.mean()
+                scaled = scaler.scale(loss)
+                scaler.minimize(optimizer, scaled)
+
+                place = paddle.CUDAPlace(0)
+                exe = paddle.static.Executor(place)
+                exe.run(startup)
+                paddle.amp.debugging.enable_operator_stats_collection()
+                exe.run(
+                    main,
+                    feed={'data': np.random.rand(1, 3, 8, 8).astype('float32')},
+                    fetch_list=[loss],
+                )
+                paddle.amp.debugging.disable_operator_stats_collection()
+                op_list = paddle.base.core.get_low_precision_op_list()
+
+                self.assertEqual(scaler._enable, False)
+                self.assertEqual(scaler._use_dynamic_loss_scaling, False)
+                self.assertTrue('pd_op.scale' not in op_list)
+                self.assertTrue(
+                    'pd_op.check_finite_and_unscale_' not in op_list
+                )
+
 
 @unittest.skipIf(
     not core.is_compiled_with_cuda()
diff --git a/test/amp/test_amp_decorate.py b/test/amp/test_amp_decorate.py
diff --git a/test/amp/test_amp_list.py b/test/amp/test_amp_list.py
diff --git a/test/amp/test_amp_master_grad.py b/test/amp/test_amp_master_grad.py
diff --git a/test/amp/test_amp_master_weight.py b/test/amp/test_amp_master_weight.py
