[PIR AMP] Fix optimizer apply_gradients (#66778)

0x45f · web-flow · commit 2dfbc8479960 · 2024-08-07T10:49:14.000+08:00
diff --git a/python/paddle/incubate/optimizer/lars_momentum.py b/python/paddle/incubate/optimizer/lars_momentum.py
@@ -125,8 +125,8 @@ def __init__(
         self._master_weights = {}
 
     def _create_accumulators(self, block, parameters):
-        assert isinstance(block, framework.Block)
-
+        if not isinstance(block, (framework.Block, pir.Block)):
+            raise TypeError("block is not instance of Block.")
         for p in parameters:
             if self._multi_precision and self._is_dtype_fp16_or_bf16(p.dtype):
                 master_p = self._create_master_weight(p)
@@ -143,7 +143,8 @@ def _create_accumulators(self, block, parameters):
             self._add_accumulator(self._velocity_acc_str, p)
 
     def _append_optimize_op(self, block, param_and_grad):
-        assert isinstance(block, framework.Block)
+        if not isinstance(block, (framework.Block, pir.Block)):
+            raise TypeError("block is not instance of Block.")
         _lars_weight_decay = self._lars_weight_decay
         param_name = param_and_grad[0].name
         if len(self._exclude_from_weight_decay) > 0:
diff --git a/python/paddle/optimizer/adagrad.py b/python/paddle/optimizer/adagrad.py
@@ -17,7 +17,7 @@
 import warnings
 from typing import TYPE_CHECKING, Sequence
 
-from paddle import _C_ops
+from paddle import _C_ops, pir
 from paddle.framework import (
     in_dynamic_or_pir_mode,
 )
@@ -164,8 +164,8 @@ def __init__(
         }
 
     def _create_accumulators(self, block, parameters):
-        assert isinstance(block, framework.Block)
-
+        if not isinstance(block, (framework.Block, pir.Block)):
+            raise TypeError("block is not instance of Block.")
         if isinstance(parameters, dict):
             parameters = self._update_param_group(parameters)
 
@@ -197,8 +197,8 @@ def _create_accumulators(self, block, parameters):
             self._already_create_accumulator.add(p.name)
 
     def _append_optimize_op(self, block, param_and_grad):
-        assert isinstance(block, framework.Block)
-
+        if not isinstance(block, (framework.Block, pir.Block)):
+            raise TypeError("block is not instance of Block.")
         if isinstance(param_and_grad, dict):
             param_and_grad = self._update_param_group(param_and_grad)
 
diff --git a/python/paddle/optimizer/adamax.py b/python/paddle/optimizer/adamax.py
@@ -16,7 +16,7 @@
 import warnings
 from typing import TYPE_CHECKING, Sequence
 
-from paddle import _C_ops
+from paddle import _C_ops, pir
 
 from ..base import core, framework
 from ..base.dygraph import no_grad
@@ -241,7 +241,7 @@ def _create_accumulators(self, block, parameters):
             self._already_create_accumulator.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)
 
@@ -315,7 +315,7 @@ def _append_optimize_op(self, block, param_and_grad):
 
     def _finish_update(self, block, parameters_and_grads):
         """Update Beta1 Power accumulator"""
-        assert isinstance(block, framework.Block)
+        assert isinstance(block, (framework.Block, pir.Block))
         if isinstance(parameters_and_grads, list):
             for param, grad in parameters_and_grads:
                 if grad is None or param.stop_gradient is True:
diff --git a/python/paddle/optimizer/asgd.py b/python/paddle/optimizer/asgd.py
@@ -17,7 +17,7 @@
 from typing import TYPE_CHECKING, Sequence
 
 import paddle
-from paddle import _C_ops
+from paddle import _C_ops, pir
 from paddle.tensor.creation import to_tensor
 
 from ..base import framework
@@ -140,7 +140,8 @@ def __init__(
         self._n_tensor = None
 
     def _create_accumulators(self, block, parameters):
-        assert isinstance(block, framework.Block)
+        if not isinstance(block, (framework.Block, pir.Block)):
+            raise TypeError("block is not instance of Block.")
         if isinstance(parameters, dict):
             parameters = self._update_param_group(parameters)
 
diff --git a/python/paddle/optimizer/lamb.py b/python/paddle/optimizer/lamb.py
@@ -16,7 +16,7 @@
 
 from typing import TYPE_CHECKING, Callable, Sequence
 
-from paddle import _C_ops
+from paddle import _C_ops, pir
 from paddle.base.executor import global_scope
 
 from ..base import core, framework
@@ -190,7 +190,8 @@ def _get_parameter(self, name, scope=None):
         return p_t, master_p_t
 
     def _create_accumulators(self, block, parameters):
-        assert isinstance(block, framework.Block)
+        if not isinstance(block, (framework.Block, pir.Block)):
+            raise TypeError("block is not instance of Block.")
         if isinstance(parameters, dict):
             parameters = self._update_param_group(parameters)
 
@@ -237,7 +238,8 @@ def _add_moments_pows(self, p):
         )
 
     def _append_optimize_op(self, block, param_and_grad):
-        assert isinstance(block, framework.Block)
+        if not isinstance(block, (framework.Block, pir.Block)):
+            raise TypeError("block is not instance of Block.")
         if isinstance(param_and_grad, dict):
             param_and_grad = self._update_param_group(param_and_grad)
 
diff --git a/python/paddle/optimizer/momentum.py b/python/paddle/optimizer/momentum.py
@@ -18,7 +18,7 @@
 from typing import TYPE_CHECKING, Sequence
 
 import paddle
-from paddle import _C_ops
+from paddle import _C_ops, pir
 from paddle.framework import in_dynamic_or_pir_mode
 from paddle.regularizer import L2Decay
 
@@ -271,7 +271,8 @@ def _create_regularization_of_grad(self, param, grad, regularization=None):
         )
 
     def _append_optimize_op(self, block, param_and_grad):
-        assert isinstance(block, framework.Block)
+        if not isinstance(block, (framework.Block, pir.Block)):
+            raise TypeError("block is not instance of Block.")
         if isinstance(param_and_grad, dict):
             param_and_grad = self._update_param_group(param_and_grad)
 
diff --git a/python/paddle/optimizer/nadam.py b/python/paddle/optimizer/nadam.py
@@ -17,7 +17,7 @@
 import warnings
 from typing import TYPE_CHECKING, Sequence
 
-from paddle import _C_ops
+from paddle import _C_ops, pir
 from paddle.base.libpaddle import DataType
 
 from ..base import core, framework
@@ -236,8 +236,8 @@ def _add_moments_pows(self, p):
         self._add_accumulator(self._moment2_acc_str, p, dtype=acc_dtype)
 
     def _create_accumulators(self, block, parameters):
-        if not isinstance(block, framework.Block):
-            raise TypeError("block is not instance of framework.Block.")
+        if not isinstance(block, (framework.Block, pir.Block)):
+            raise TypeError("block is not instance of Block.")
 
         if isinstance(parameters, dict):
             parameters = parameters.get('params')
@@ -263,8 +263,8 @@ def _create_accumulators(self, block, parameters):
             self._already_create_accumulator.add(p.name)
 
     def _append_optimize_op(self, block, param_and_grad):
-        if not isinstance(block, framework.Block):
-            raise TypeError("block is not instance of framework.Block.")
+        if not isinstance(block, (framework.Block, pir.Block)):
+            raise TypeError("block is not instance of Block.")
 
         if isinstance(param_and_grad, dict):
             param_and_grad = self._update_param_group(param_and_grad)
diff --git a/python/paddle/optimizer/optimizer.py b/python/paddle/optimizer/optimizer.py
@@ -1318,19 +1318,19 @@ def _pir_create_optimization_pass(
             their internal state.
         """
 
-        global_block = framework.default_main_program().global_block()
+        global_block = paddle.static.default_main_program().global_block()
         target_block = global_block
 
-        start = len(target_block.ops)
+        last_op = target_block.ops[-1]
 
         self._create_global_learning_rate()
 
-        params_grads_device_map = (
-            parameters_and_grads['params']
-            if isinstance(parameters_and_grads, dict)
-            else parameters_and_grads
-        )
-        self._update_param_device_map(params_grads_device_map, target_block)
+        # params_grads_device_map = (
+        #     parameters_and_grads['params']
+        #     if isinstance(parameters_and_grads, dict)
+        #     else parameters_and_grads
+        # )
+        # self._update_param_device_map(params_grads_device_map, target_block)
 
         if isinstance(parameters_and_grads, list):
             self._create_accumulators(
@@ -1373,8 +1373,8 @@ def _pir_create_optimization_pass(
         self._finish_update(target_block, parameters_and_grads)
         paddle.base.core._set_warmup(False)
 
-        end = len(target_block.ops)
-        return target_block._slice_ops(start, end)
+        start_index = target_block.ops.index(last_op) + 1
+        return target_block.ops[start_index:]
 
     def backward(
         self,
@@ -1527,7 +1527,10 @@ def apply_gradients(
             params_grads, self.regularization
         )
 
-        optimize_ops = self._create_optimization_pass(params_grads)
+        if in_pir_mode():
+            optimize_ops = self._pir_create_optimization_pass(params_grads)
+        else:
+            optimize_ops = self._create_optimization_pass(params_grads)
         return optimize_ops
 
     def _apply_optimize(
diff --git a/python/paddle/optimizer/radam.py b/python/paddle/optimizer/radam.py
@@ -17,7 +17,7 @@
 import warnings
 from typing import TYPE_CHECKING, Sequence
 
-from paddle import _C_ops
+from paddle import _C_ops, pir
 from paddle.base.libpaddle import DataType
 
 from ..base import core, framework
@@ -232,8 +232,8 @@ def _add_moments_pows(self, p):
         self._add_accumulator(self._moment2_acc_str, p, dtype=acc_dtype)
 
     def _create_accumulators(self, block, parameters):
-        if not isinstance(block, framework.Block):
-            raise TypeError("block is not instance of framework.Block.")
+        if not isinstance(block, (framework.Block, pir.Block)):
+            raise TypeError("block is not instance of Block.")
 
         if isinstance(parameters, dict):
             parameters = parameters.get('params')
@@ -259,8 +259,8 @@ def _create_accumulators(self, block, parameters):
             self._already_create_accumulator.add(p.name)
 
     def _append_optimize_op(self, block, param_and_grad):
-        if not isinstance(block, framework.Block):
-            raise TypeError("block is not instance of framework.Block.")
+        if not isinstance(block, (framework.Block, pir.Block)):
+            raise TypeError("block is not instance of Block.")
 
         if isinstance(param_and_grad, dict):
             param_and_grad = self._update_param_group(param_and_grad)
diff --git a/python/paddle/optimizer/rmsprop.py b/python/paddle/optimizer/rmsprop.py
@@ -19,7 +19,7 @@
 
 from typing_extensions import NotRequired
 
-from paddle import _C_ops
+from paddle import _C_ops, pir
 
 from ..base import framework
 from ..base.framework import in_dynamic_or_pir_mode
@@ -224,8 +224,8 @@ def __init__(
         }
 
     def _create_accumulators(self, block, parameters):
-        if not isinstance(block, framework.Block):
-            raise TypeError("block is not instance of framework.Block.")
+        if not isinstance(block, (framework.Block, pir.Block)):
+            raise TypeError("block is not instance of Block.")
 
         if isinstance(parameters, dict):
             parameters = parameters.get('params')
@@ -255,8 +255,8 @@ def _create_accumulators(self, block, parameters):
             self._already_create_accumulator.add(p.name)
 
     def _append_optimize_op(self, block, param_and_grad):
-        if not isinstance(block, framework.Block):
-            raise TypeError("block is not instance of framework.Block.")
+        if not isinstance(block, (framework.Block, pir.Block)):
+            raise TypeError("block is not instance of Block.")
 
         if isinstance(param_and_grad, dict):
             param_and_grad = self._update_param_group(param_and_grad)
diff --git a/python/paddle/optimizer/rprop.py b/python/paddle/optimizer/rprop.py
@@ -17,7 +17,7 @@
 import warnings
 from typing import TYPE_CHECKING, Sequence
 
-from paddle import _C_ops
+from paddle import _C_ops, pir
 from paddle.tensor.creation import to_tensor
 
 from ..base import framework
@@ -152,14 +152,15 @@ def __init__(
         self._sign = True
 
     def _to_tensor(self, block, dtype):
-        assert isinstance(block, framework.Block)
+        assert isinstance(block, (framework.Block, pir.Block))
         self._learning_rate_range = to_tensor(
             self._learning_rate_range, dtype=dtype
         )
         self._etas = to_tensor(self._etas, dtype=dtype)
 
     def _create_accumulators(self, block, parameters):
-        assert isinstance(block, framework.Block)
+        if not isinstance(block, (framework.Block, pir.Block)):
+            raise TypeError("block is not instance of Block.")
         if isinstance(parameters, dict):
             parameters = self._update_param_group(parameters)
 
diff --git a/test/amp/test_pir_amp.py b/test/amp/test_pir_amp.py
@@ -164,7 +164,10 @@ def test_linear_amp_o2(self):
                     out = linear(x)
                     loss = paddle.mean(out)
                 scaled = scaler.scale(loss)
-                scaler.minimize(optimizer, scaled, startup_program=startup)
+                opt_ops, _ = scaler.minimize(
+                    optimizer, scaled, startup_program=startup
+                )
+            np.testing.assert_equal(len(opt_ops), 8)
             cast_op_count = 0
             for op in main.global_block().ops:
                 if op.name() == 'pd_op.cast':
diff --git a/test/legacy_test/test_lambv2_op.py b/test/legacy_test/test_lambv2_op.py
@@ -25,7 +25,7 @@
 
 class LAMBOptimizer(paddle.optimizer.Lamb):
     def _append_optimize_op(self, block, param_and_grad):
-        assert isinstance(block, base.framework.Block)
+        assert isinstance(block, (base.framework.Block, paddle.pir.Block))
         block.program._use_lamb = True
 
         m = moment1 = self._get_accumulator(
