[PIR]Migrate ClipGradByGlobalNorm into pir (#58500)

0x45f · web-flow · commit 84c2d025c841 · 2023-11-15T17:10:36.000+08:00
diff --git a/paddle/fluid/pir/dialect/op_generator/ops_api_gen.py b/paddle/fluid/pir/dialect/op_generator/ops_api_gen.py
@@ -90,6 +90,7 @@
     'fused_attention',
     'fused_feedforward',
     'self_dp_attention',
+    'get_tensor_from_selected_rows',
 ]
 
 NO_NEED_GEN_STATIC_ONLY_APIS = [
diff --git a/paddle/fluid/pir/dialect/operator/ir/ops.yaml b/paddle/fluid/pir/dialect/operator/ir/ops.yaml
@@ -67,6 +67,8 @@
 - op : get_tensor_from_selected_rows
   args : (Tensor x)
   output : Tensor(out)
+  infer_meta :
+    func : UnchangedInferMeta
   kernel:
     func: get_tensor_from_selected_rows {selected_rows -> dense}
 
diff --git a/python/paddle/nn/clip.py b/python/paddle/nn/clip.py
@@ -21,8 +21,14 @@
 from paddle import _C_ops
 from paddle.base import core, framework, unique_name
 from paddle.base.data_feeder import check_variable_and_dtype
+from paddle.base.libpaddle import DataType
 from paddle.common_ops_import import Variable, check_type, default_main_program
-from paddle.framework import LayerHelper, in_dynamic_mode
+from paddle.framework import (
+    LayerHelper,
+    in_dynamic_mode,
+    in_dynamic_or_pir_mode,
+    in_pir_mode,
+)
 from paddle.tensor.layer_function_generator import templatedoc
 
 __all__ = []
@@ -114,7 +120,7 @@ def merge_selected_rows(x, name=None):
             ...     type=base.core.VarDesc.VarType.SELECTED_ROWS)
             >>> y = paddle.nn.clip.merge_selected_rows(var)
     """
-    if in_dynamic_mode():
+    if in_dynamic_or_pir_mode():
         return _C_ops.merge_selected_rows(x)
 
     helper = LayerHelper("merge_selected_rows", **locals())
@@ -173,6 +179,8 @@ def get_tensor_from_selected_rows(x, name=None):
             >>> x = block.create_var(name="X", dtype="float32", persistable=True, type=base.core.VarDesc.VarType.SELECTED_ROWS)
             >>> z = paddle.nn.clip.get_tensor_from_selected_rows(x)
     """
+    if in_pir_mode():
+        return _C_ops.get_tensor_from_selected_rows(x)
 
     check_type(x, 'x', Variable, 'get_tensor_from_selected_rows')
     if x.type != core.VarDesc.VarType.SELECTED_ROWS:
@@ -211,6 +219,10 @@ def _cast_to_mp_type_if_enabled(x):
         or x.dtype == core.VarDesc.VarType.BF16
     ) and _clip_by_global_norm_using_mp_type():
         return x.astype(core.VarDesc.VarType.FP32)
+    elif (
+        x.dtype == DataType.FLOAT16 or x.dtype == DataType.BFLOAT16
+    ) and _clip_by_global_norm_using_mp_type():
+        return x.astype(DataType.FP32)
     else:
         return x
 
@@ -227,7 +239,7 @@ def _squared_l2_norm(x):
         sum_square = paddle.sum(square)
         return sum_square
 
-    if in_dynamic_mode():
+    if in_dynamic_or_pir_mode():
         return _C_ops.squared_l2_norm(x)
 
     op_type = 'squared_l2_norm'
@@ -339,12 +351,17 @@ def __str__(self):
     def _dygraph_clip(self, params_grads):
         raise NotImplementedError
 
+    def _pir_clip(self, params_grads):
+        raise NotImplementedError
+
     def _static_clip(self, params_grads):
         raise NotImplementedError
 
     def __call__(self, params_grads):
         if in_dynamic_mode():
             return self._dygraph_clip(params_grads)
+        elif in_pir_mode():
+            return self._pir_clip(params_grads)
         else:
             for p, g in params_grads:
                 if getattr(p, 'gradient_clip_attr', None) is not None:
@@ -749,6 +766,97 @@ def async_add_n(var_list):
 
         return params_and_grads
 
+    def _pir_clip(self, params_grads):
+        params_and_grads = []
+        sum_square_list = []
+        sum_square_list_fp16 = []
+        sum_square_list_fp32 = []
+        for p, g in params_grads:
+            if g is None:
+                continue
+            if getattr(p, 'need_clip', True) is False:
+                continue
+            merge_grad = g
+
+            if in_pir_mode() and g.is_selected_row_type():
+                merge_grad = merge_selected_rows(g)
+                merge_grad = get_tensor_from_selected_rows(merge_grad)
+
+            sum_square = _squared_l2_norm(merge_grad)
+            if (
+                sum_square.dtype == DataType.FLOAT16
+                or sum_square.dtype == DataType.BFLOAT16
+            ):
+                sum_square_list_fp16.append(sum_square)
+            elif sum_square.dtype == DataType.FLOAT32:
+                sum_square_list_fp32.append(sum_square)
+            else:
+                sum_square_list.append(sum_square)
+
+        # all parameters have been filterd out
+        if (
+            len(sum_square_list)
+            + len(sum_square_list_fp16)
+            + len(sum_square_list_fp32)
+            == 0
+        ):
+            return params_grads
+
+        def async_add_n(var_list):
+            return paddle.stack(var_list).sum()
+
+        sum_dtype = 'float64' if len(sum_square_list) > 0 else "float32"
+        global_norm_var = []
+        if len(sum_square_list_fp16) > 0:
+            global_norm_var_fp16 = async_add_n(sum_square_list_fp16)
+            global_norm_var.append(global_norm_var_fp16.astype(sum_dtype))
+        if len(sum_square_list_fp32) > 0:
+            global_norm_var_fp32 = async_add_n(sum_square_list_fp32)
+            if sum_dtype == 'float32':
+                global_norm_var.append(global_norm_var_fp32)
+            else:
+                global_norm_var.append(global_norm_var_fp32.astype(sum_dtype))
+        if len(sum_square_list) > 0:
+            global_norm_var_fp64 = async_add_n(sum_square_list)
+            global_norm_var.append(global_norm_var_fp64)
+        global_norm_var = async_add_n(global_norm_var)
+        global_norm_var = paddle.sqrt(global_norm_var)
+        max_global_norm = paddle.full(
+            shape=[], dtype=global_norm_var.dtype, fill_value=self.clip_norm
+        )
+
+        need_clip = False
+        if not self.auto_skip_clip:  # always apply clip
+            need_clip = True
+            clip_var = paddle.divide(
+                x=max_global_norm,
+                y=paddle.maximum(x=global_norm_var, y=max_global_norm),
+            )
+        elif global_norm_var > max_global_norm:
+            # only when global_norm_var > max_global_norm, grad need clip
+            need_clip = True
+            clip_var = paddle.divide(x=max_global_norm, y=global_norm_var)
+
+        for p, g in params_grads:
+            if g is None:
+                continue
+            if getattr(p, 'need_clip', True) is False:
+                params_and_grads.append((p, g))
+                continue
+            # TODO(wangxi): use inplace elementwise_mul
+            if need_clip:
+                clip_input = (
+                    clip_var.astype(g.dtype)
+                    if clip_var.dtype != g.dtype
+                    else clip_var
+                )
+                new_grad = paddle.multiply(g, clip_input)
+                params_and_grads.append((p, new_grad))
+            else:
+                params_and_grads.append((p, g))
+
+        return params_and_grads
+
     def _static_clip(self, params_grads):
         params_and_grads = []
         sum_square_list = []
@@ -925,6 +1033,9 @@ def _process_context(self, context, param, grad):
         if grad.type == core.VarDesc.VarType.SELECTED_ROWS:
             merge_grad = merge_selected_rows(grad)
             merge_grad = get_tensor_from_selected_rows(merge_grad)
+        elif in_pir_mode() and grad.is_selected_row_type():
+            merge_grad = merge_selected_rows(grad)
+            merge_grad = get_tensor_from_selected_rows(merge_grad)
 
         local_norm_var = _squared_l2_norm(merge_grad)
         context[self.group_name].append(local_norm_var)
@@ -947,6 +1058,10 @@ def async_add_n(var_list):
             assert group_scale_var.shape == (1,)
             self.context[group_scale_name] = group_scale_var
 
+        if in_pir_mode():
+            grad = paddle.multiply(grad, self.context[group_scale_name])
+            return param, grad
+
         # inplace
         param.block.append_op(
             type='elementwise_mul',
diff --git a/test/legacy_test/test_gradient_clip.py b/test/legacy_test/test_gradient_clip.py
@@ -159,6 +159,133 @@ def backward_and_optimize(self, cost):
         pass
 
 
+class TestPirGradientClipByGlobalNorm(TestGradientClip):
+    def init(self):
+        self.clip_norm = 0.2
+
+    def check_clip_result(self, out, out_clip):
+        global_norm = 0
+        for v in out:
+            global_norm += np.sum(np.square(v))
+        global_norm = np.sqrt(global_norm)
+        scale = self.clip_norm / np.maximum(self.clip_norm, global_norm)
+        res = []
+        for i in range(len(out)):
+            out[i] = scale * out[i]
+
+        for u, v in zip(out, out_clip):
+            np.testing.assert_allclose(
+                u,
+                v,
+                rtol=1e-05,
+                atol=1e-08,
+                err_msg=f'gradient clip by global norm has wrong results!, \nu={u}\nv={v}\ndiff={u - v}',
+            )
+
+    def _run(self, place, dtype='float32'):
+        paddle.seed(2023)
+        prog = paddle.static.Program()
+        startup_program = paddle.static.Program()
+        with paddle.static.program_guard(
+            main_program=prog, startup_program=startup_program
+        ):
+            image = paddle.static.data(
+                name="a", shape=[-1, 784], dtype='float32'
+            )
+            label = paddle.static.data(name="b", shape=[-1, 1], dtype='int64')
+            hidden_linear = paddle.nn.Linear(784, 32)
+            if dtype != 'float32':
+                image_cast = paddle.cast(image, dtype)
+                hidden = paddle.nn.functional.relu(hidden_linear(image_cast))
+            else:
+                hidden = paddle.nn.functional.relu(hidden_linear(image))
+
+            predict_linear = paddle.nn.Linear(32, 10)
+            predict = paddle.nn.functional.softmax(predict_linear(hidden))
+
+            cost = paddle.nn.functional.cross_entropy(
+                input=predict, label=label, reduction='none', use_softmax=False
+            )
+            avg_cost = paddle.mean(cost)
+
+            grad_list = paddle.autograd.ir_backward.grad(
+                avg_cost, prog.global_block().all_parameters()
+            )
+
+            train_reader = paddle.batch(
+                paddle.dataset.mnist.train(), batch_size=3
+            )
+            exe = base.Executor(place)
+            exe.run(startup_program)
+            data = next(train_reader())
+            a = np.array([i[0] for i in data]).astype('float32')
+            b = np.array([i[1] for i in data]).reshape(3, 1).astype('int64')
+            out = exe.run(prog, feed={'a': a, 'b': b}, fetch_list=grad_list)
+            return out
+
+    def _run_clip(self, place, dtype='float32'):
+        paddle.seed(2023)
+        prog = paddle.static.Program()
+        startup_program = paddle.static.Program()
+        with paddle.static.program_guard(
+            main_program=prog, startup_program=startup_program
+        ):
+            image = paddle.static.data(
+                name="a", shape=[-1, 784], dtype='float32'
+            )
+            label = paddle.static.data(name="b", shape=[-1, 1], dtype='int64')
+            hidden_linear = paddle.nn.Linear(784, 32)
+            if dtype != 'float32':
+                image_cast = paddle.cast(image, dtype)
+                hidden = paddle.nn.functional.relu(hidden_linear(image_cast))
+            else:
+                hidden = paddle.nn.functional.relu(hidden_linear(image))
+
+            predict_linear = paddle.nn.Linear(32, 10)
+            predict = paddle.nn.functional.softmax(predict_linear(hidden))
+
+            cost = paddle.nn.functional.cross_entropy(
+                input=predict, label=label, reduction='none', use_softmax=False
+            )
+            avg_cost = paddle.mean(cost)
+
+            params = prog.global_block().all_parameters()
+            grad_list = paddle.autograd.ir_backward.grad(avg_cost, params)
+
+            p_g_clip = self.clip_gradient(list(zip(params, grad_list)))
+
+            grad_clip_list = [elem[1] for elem in p_g_clip]
+            train_reader = paddle.batch(
+                paddle.dataset.mnist.train(), batch_size=3
+            )
+            exe = base.Executor(place)
+            exe.run(startup_program)
+            data = next(train_reader())
+            a = np.array([i[0] for i in data]).astype('float32')
+            b = np.array([i[1] for i in data]).reshape(3, 1).astype('int64')
+            out_clip = exe.run(
+                prog, feed={'a': a, 'b': b}, fetch_list=grad_clip_list
+            )
+            return out_clip
+
+    def check_gradient_clip(self, place, dtype='float32'):
+        out = self._run(place, dtype)
+        out_clip = self._run_clip(place, dtype)
+        self.check_clip_result(out, out_clip)
+
+    def test_new_gradient_clip(self):
+        def func(params_grads):
+            clip = paddle.nn.ClipGradByGlobalNorm(clip_norm=self.clip_norm)
+            return clip(params_grads)
+
+        self.clip_gradient = func
+        with paddle.pir_utils.IrGuard():
+            self.check_gradient_clip(base.CPUPlace())
+
+    def check_sparse_gradient_clip(self, place):
+        pass
+
+
 class TestGradientClipByGlobalNorm(TestGradientClip):
     def init(self):
         self.clip_norm = 0.2

Original file line number	Diff line number	Diff line change
`@@ -90,6 +90,7 @@`
`90`	`90`	`'fused_attention',`
`91`	`91`	`'fused_feedforward',`
`92`	`92`	`'self_dp_attention',`
	`93`	`+ 'get_tensor_from_selected_rows',`
`93`	`94`	`]`
`94`	`95`
`95`	`96`	`NO_NEED_GEN_STATIC_ONLY_APIS = [`