【AutoParallel】Add auto_cast support in auto_parallel (#60158)

heavyrain-lzy · web-flow · commit 256dc4d1c0fb · 2023-12-22T17:32:00.000+08:00
* add auto_cast support in auto_parallel

* change code according to review

* fix bug in recompute function
diff --git a/python/paddle/distributed/auto_parallel/interface.py b/python/paddle/distributed/auto_parallel/interface.py
@@ -221,7 +221,14 @@ def __call__(self, *args, **kwargs):
             with paddle.static.name_scope(
                 f'/auto_parallel/rc_{_g_recompute_idx}'
             ):
-                output = self._op(*args, **kwargs)
+                if paddle.base.dygraph.base.in_to_static_mode():
+                    output = (
+                        paddle.jit.dy2static.convert_call_func.convert_call(
+                            self._op
+                        )(*args, **kwargs)
+                    )
+                else:
+                    output = self._op(*args, **kwargs)
 
             return output
 
@@ -245,7 +252,14 @@ def __init__(self, run_function):
 
         def __call__(self, *args, **kwargs):
             with paddle.static.name_scope('/exclude_rc'):
-                output = self._run_function(*args, **kwargs)
+                if paddle.base.dygraph.base.in_to_static_mode():
+                    output = (
+                        paddle.jit.dy2static.convert_call_func.convert_call(
+                            self._run_function
+                        )(*args, **kwargs)
+                    )
+                else:
+                    output = self._run_function(*args, **kwargs)
 
             return output
 
diff --git a/python/paddle/distributed/auto_parallel/static/dist_op.py b/python/paddle/distributed/auto_parallel/static/dist_op.py
@@ -258,7 +258,12 @@ def __call__(self, *args, **kwargs):
         default_prog = paddle.static.default_main_program()
         cur_block = default_prog.current_block()
         op_size = len(cur_block.ops)
-        output = self._serial_op(*args, **kwargs)
+        if paddle.base.dygraph.base.in_to_static_mode():
+            output = paddle.jit.dy2static.convert_call_func.convert_call(
+                self._serial_op
+            )(*args, **kwargs)
+        else:
+            output = self._serial_op(*args, **kwargs)
         new_op_size = len(cur_block.ops)
 
         if isinstance(output, (tuple, list)):
diff --git a/python/paddle/distributed/auto_parallel/static/helper.py b/python/paddle/distributed/auto_parallel/static/helper.py
@@ -17,10 +17,17 @@
 from collections import defaultdict
 
 from paddle.jit import not_to_static, to_static
-from paddle.jit.dy2static.program_translator import StaticFunction
+from paddle.jit.dy2static.program_translator import (
+    ProgramTranslator,
+    StaticFunction,
+)
 from paddle.jit.dy2static.utils import as_not_paddle_func
 from paddle.nn import Layer
 from paddle.static import Parameter, global_scope, program_guard
+from paddle.static.amp.fp16_utils import (
+    DEFAULT_AMP_OPTIONS,
+    prepare_op_amp_options,
+)
 
 from .converter import Converter
 from .utils import get_logger, to_list
@@ -244,8 +251,14 @@ def build_program(self, mode):
 
         # NOTE(dev): Because @to_static is a Lazy mechanism, so we explicitly call this to trigger
         # generating Program IR immediately.
-        getattr(self.proxy_layer, func_name).concrete_program  # noqa: B018
-
+        concrete_program = getattr(
+            self.proxy_layer, func_name
+        ).concrete_program  # noqa: B018
+        prepare_op_amp_options(
+            concrete_program.main_program,
+            ProgramTranslator.get_instance()._amp_records,
+            DEFAULT_AMP_OPTIONS,
+        )
         self._build_startup_program()
 
     def _build_startup_program(self):
diff --git a/python/paddle/distributed/passes/auto_parallel_amp.py b/python/paddle/distributed/passes/auto_parallel_amp.py
@@ -146,6 +146,8 @@ def _is_in_black_fp32_varnames(self, op):
             return _is_in_fp32_varnames(op, self._amp_list)
 
     def _op_keep_fp32_input(self, op, in_name):
+        if not op.amp_options.enable:
+            return True
         if self._is_float16:
             return _keep_fp32_input(op, in_name)
         else:
@@ -156,6 +158,8 @@ def _op_keep_fp32_input(self, op, in_name):
             return False
 
     def _op_keep_fp32_output(self, op, out_name):
+        if not op.amp_options.enable:
+            return True
         if self._is_float16:
             return _keep_fp32_output(op, out_name)
         else:
diff --git a/python/paddle/distributed/passes/auto_parallel_fp16.py b/python/paddle/distributed/passes/auto_parallel_fp16.py
@@ -69,9 +69,31 @@ def set_op_dtype_to_fp16(op):
         op._set_attr('dtype', __target_dtype__)
 
 
+def set_auto_cast_attr(cast_op, block):
+    in_name = cast_op.input('X')[0]
+    out_name = cast_op.output('Out')[0]
+    in_var = block._find_var_recursive(in_name)
+    out_var = block._find_var_recursive(out_name)
+    assert (
+        in_var is not None and out_var is not None
+    ), f"in_var {in_name} or out_var {out_name} is None of cast op"
+    if is_forward_op(cast_op):
+        cast_op._set_attr('in_dtype', in_var.dtype)
+        cast_op._set_attr('out_dtype', out_var.dtype)
+    elif is_backward_op(cast_op):
+        in_var_fw = block._find_var_recursive(in_name[: in_name.find("@")])
+        out_var_fw = block._find_var_recursive(out_name[: out_name.find("@")])
+        cast_op._set_attr('in_dtype', in_var_fw.dtype)
+        cast_op._set_attr('out_dtype', out_var_fw.dtype)
+        in_var.desc.set_dtype(in_var_fw.dtype)
+        out_var.desc.set_dtype(out_var_fw.dtype)
+
+
 # adapot for backward op
 # TODO check if bf16 and fp16 still share the same logic
 def _keep_fp32_input(op, in_name):
+    if not op.amp_options.enable:
+        return True
     op_type = op.type
     if op_type == 'batch_norm':
         # Scale, Bias, Mean, Variance should be float32.
@@ -101,6 +123,9 @@ def _keep_fp32_input(op, in_name):
 
 # TODO check if bf16 and fp16 still share the same logic
 def _keep_fp32_output(op, out_name):
+    # TODO(lizhiyu02): Support 'custom_white_list' adn 'custom_black_list' in amp_options
+    if not op.amp_options.enable:
+        return True
     op_type = op.type
     if op_type in ['batch_norm', 'fused_bn_add_activation']:
         return out_name != 'Y'
@@ -141,6 +166,7 @@ def __init__(
         self.grad_op_to_op_map = (
             self.dist_context.dist_op_context.grad_op_id_to_op_id
         )
+        self.forward_op_to_amp_options = {}
         if input_data_var_names:
             self.input_data_var_names = input_data_var_names
         else:
@@ -176,7 +202,7 @@ def _build_state(self):
                         self.out_var_op_deps[name].extend(
                             [op.desc.original_id()]
                         )
-
+                self._mark_amp_options_info(op)
                 self._mark_op(op)
 
         # set forward tensor dtype
@@ -189,6 +215,25 @@ def _build_state(self):
 
         return self.is_train
 
+    def _mark_amp_options_info(self, op):
+        """
+        Mark amp options info for backward ops according to forward ops
+        """
+        if is_forward_op(op):
+            self.forward_op_to_amp_options[
+                op.desc.original_id()
+            ] = op.amp_options
+        elif is_backward_op(op):
+            if op.desc.original_id() in self.grad_op_to_op_map:
+                if (
+                    self.grad_op_to_op_map[op.desc.original_id()]
+                    in self.forward_op_to_amp_options.keys()
+                ):
+                    amp_option = self.forward_op_to_amp_options[
+                        self.grad_op_to_op_map[op.desc.original_id()]
+                    ]
+                    op.set_amp_options(amp_option)
+
     def _mark_op(self, op):
         if op.type in __amp_skip_ops__:
             return
@@ -210,7 +255,7 @@ def _mark_op(self, op):
                         self._op_fp16_dict[op.desc.original_id()] = True
                     return
 
-            if __amp_utils__._need_keep_fp32(
+            if not op.amp_options.enable or __amp_utils__._need_keep_fp32(
                 op, self.amp_list.unsupported_list, self.use_fp16_guard
             ):
                 self._op_fp16_dict[op.desc.original_id()] = False
@@ -253,6 +298,11 @@ def set_var_to_fp16(self, var_name, block):
 
     def resolute_tensor_dtype(self, block):
         for op in block.ops:
+            # 'amp_options' flag has highest priority
+            if not op.amp_options.enable:
+                if op.type == "cast":
+                    set_auto_cast_attr(op, block)
+                continue
             if is_forward_op(op):
                 # NOTE (JZ-LIANG) un-expected cast op when user call "+, -, *, /" in python
                 if (
@@ -795,7 +845,6 @@ def _apply_single_impl(self, main_program, startup_program, context):
         # NOTE don't not change input data dtype, since it is controled by dataloader
         # and which is out of control of FP16 Pass
         input_data_var_names = [var.name for var in self.get_attr("input_data")]
-
         with paddle.static.program_guard(main_program, startup_program):
             fp16_state = FP16State(
                 main_program,
