[PIR AMP]Fix opt.minimize error when no using grad scaler (#64284)

0x45f · web-flow · commit 06ac6ab35509 · 2024-05-14T20:11:03.000+08:00
diff --git a/python/paddle/amp/auto_cast.py b/python/paddle/amp/auto_cast.py
@@ -26,7 +26,6 @@
 )
 from paddle.base.wrapped_decorator import signature_safe_contextmanager
 from paddle.static.amp.decorator import OptimizerWithMixedPrecision
-from paddle.static.amp.fp16_lists import AutoMixedPrecisionLists
 
 from .amp_lists import black_list, white_list
 
@@ -1035,15 +1034,15 @@ def decorate(
                     amp_lists=None,
                     level=level,
                     dtype=dtype,
-                    init_loss_scaling=2.0**16,
-                    incr_every_n_steps=2000,
-                    decr_every_n_nan_or_inf=1,
-                    incr_ratio=2.0,
-                    decr_ratio=0.5,
+                    init_loss_scaling=1.0,
+                    incr_every_n_steps=None,
+                    decr_every_n_nan_or_inf=None,
+                    incr_ratio=None,
+                    decr_ratio=None,
                     use_dynamic_loss_scaling=False,
                     use_amp_guard=None,
                     use_master_grad=master_grad,
-                    use_promote=True,
+                    use_promote=None,
                 )
                 return models, optimizers
         elif level == 'O2':
@@ -1057,16 +1056,18 @@ def decorate(
             else:
                 optimizers = OptimizerWithMixedPrecision(
                     optimizer=optimizers,
-                    amp_lists=AutoMixedPrecisionLists(dtype=dtype),
+                    amp_lists=None,
                     level=level,
                     dtype=dtype,
-                    init_loss_scaling=2**15,
+                    init_loss_scaling=1.0,
+                    incr_every_n_steps=None,
+                    decr_every_n_nan_or_inf=None,
+                    incr_ratio=None,
+                    decr_ratio=None,
                     use_dynamic_loss_scaling=False,
-                    incr_every_n_steps=1000,
-                    decr_every_n_nan_or_inf=2,
-                    incr_ratio=2.0,
-                    decr_ratio=0.8,
+                    use_amp_guard=None,
                     use_master_grad=master_grad,
+                    use_promote=None,
                 )
                 return models, optimizers
         else:
diff --git a/python/paddle/static/amp/decorator.py b/python/paddle/static/amp/decorator.py
@@ -271,6 +271,8 @@ def backward(
                 self._train_program, startup_program
             ):
                 self._init_amp_var()
+                if self._scaled_loss is None:
+                    self._scaled_loss = loss
                 params_grads = self._optimizer.backward(
                     self._scaled_loss,
                     startup_program,
diff --git a/test/amp/test_pir_amp.py b/test/amp/test_pir_amp.py
@@ -64,6 +64,46 @@ def test_linear_amp_o1(self):
             np.testing.assert_equal(len(_white_list), 0)
             np.testing.assert_equal(len(_black_list), 0)
 
+    def test_linear_amp_o2_without_scaler(self):
+        if not core.is_compiled_with_cuda():
+            return
+        with paddle.pir_utils.IrGuard():
+            startup = paddle.static.Program()
+            main = paddle.static.Program()
+            with paddle.static.program_guard(main, startup):
+                x = paddle.static.data('x', [3, 4], 'float32')
+                linear = paddle.nn.Linear(4, 5)
+                optimizer = paddle.optimizer.Adam(
+                    learning_rate=0.001, parameters=linear.parameters()
+                )
+                linear, optimizer = paddle.amp.decorate(
+                    models=linear,
+                    optimizers=optimizer,
+                    level='O2',
+                    master_weight=True,
+                    master_grad=True,
+                )
+
+                with paddle.amp.auto_cast(
+                    level='O2', dtype='float16', use_promote=True
+                ):
+                    out = linear(x)
+                    loss = paddle.mean(out)
+                optimizer.minimize(loss)
+            cast_op_count = 0
+            for op in main.global_block().ops:
+                if op.name() == 'pd_op.cast':
+                    cast_op_count += 1
+            np.testing.assert_equal(cast_op_count, 3)
+            place = paddle.CUDAPlace(0)
+            exe = paddle.static.Executor(place)
+            exe.run(startup)
+            result = exe.run(
+                main,
+                feed={'x': np.random.rand(3, 4).astype('float32')},
+                fetch_list=[loss],
+            )
+
     def test_linear_amp_o2(self):
         if not core.is_compiled_with_cuda():
             return
