mindspore 2.1.0__cp38-cp38-manylinux1_x86_64.whl → 2.2.10__cp38-cp38-manylinux1_x86_64.whl

mindspore/train/amp.py

@@ -19,8 +19,8 @@ import mindspore as ms
 from mindspore import nn
 from mindspore import _checkparam as validator
 from mindspore.common import dtype as mstype
-from mindspore.nn.wrap.cell_wrapper import _TrainPipelineAccuStepCell
-from mindspore.nn.wrap.loss_scale import _TrainPipelineWithLossScaleCell
+from mindspore.nn.wrap.cell_wrapper import _TrainGradAccuStepCell
+from mindspore.nn.wrap.loss_scale import _TrainGradAccuWithLossScaleCell
 from mindspore.ops import functional as F
 from mindspore.parallel._utils import _get_pipeline_stages
 from mindspore.train.loss_scale_manager import DynamicLossScaleManager, LossScaleManager
@@ -30,9 +30,6 @@ from mindspore.ops import Primitive
 from mindspore import log as logger
 
 
-STREE = None
-
-
 AMP_WHITE_LIST = [
     nn.Conv1d,
     nn.Conv2d,
@@ -64,17 +61,19 @@ AMP_BLACK_LIST = [
     nn.LayerNorm
 ]
 
+MS_AMP_BY_REWRITE = False
+_amp_cast_op = P.Cast
 
 class _OutputTo16(nn.Cell):
     """Wrap cell for amp. Cast network output back to float16."""
-    def __init__(self, backbone):
+    def __init__(self, backbone, dtype=mstype.float16):
         super(_OutputTo16, self).__init__(auto_prefix=False)
         self._backbone = backbone
-        if isinstance(backbone, nn.Cell) and backbone.jit_config_dict:
-            self._jit_config_dict = backbone.jit_config_dict
+        self.dtype = dtype
+        self._get_attr_from_cell(backbone)
 
     def construct(self, *args, **kwargs):
-        return F.cast(self._backbone(*args, **kwargs), mstype.float16)
+        return F.cast(self._backbone(*args, **kwargs), self.dtype)
 
 
 class _OutputTo32(nn.Cell):
@@ -82,63 +81,73 @@ class _OutputTo32(nn.Cell):
     def __init__(self, backbone):
         super(_OutputTo32, self).__init__(auto_prefix=False)
         self._backbone = backbone
-        if isinstance(backbone, nn.Cell) and backbone.jit_config_dict:
-            self._jit_config_dict = backbone.jit_config_dict
+        self._get_attr_from_cell(backbone)
 
     def construct(self, *args, **kwargs):
         out = self._backbone(*args, **kwargs)
         return F.mixed_precision_cast(mstype.float32, out)
 
 
-def _allow_mix_precision(node, allowed_list) -> bool:
+
+def _allow_mix_precision(node, allowed_list, dtype) -> bool:
     """
     Check whether current node need do mix precision. Follow conditions need to be satisfied:
         1) Type of node is one of (Primitive, nn.Cell)
-        2) Node is not P.Cast()
+        2) Node is not Cast Op
         3) to_float(mindspore.float16) is not set in Cell
     """
-    if node.get_instance() in allowed_list:
+    node_inst = node.get_instance()
+    if node_inst in allowed_list:
         return True
+    if node.get_targets() is None:
+        return False
     if not issubclass(node.get_instance_type(), (Primitive, nn.Cell)):
         return False
-    if isinstance(node.get_instance(), P.Cast):
+    if isinstance(node_inst, _amp_cast_op):
         return False
     if issubclass(node.get_instance_type(), nn.Cell):
-        # if cell is already in allowed_list, it means to_float(mindspore.float16) is set by amp.
-        # if cell is not in allowed_list, but has to_float(mindspore.float16),
-        # it means to_float(mindspore.float16) is set by user.
-        if hasattr(node.get_instance(), "to_float_fp16") and node.get_instance().to_float_fp16:
+        # if cell is already in allowed_list, it means to_float() is set by amp.
+        # if cell is not in allowed_list, but has to_float(),
+        # it means to_float() is set by user.
+        to_float_flag = "bf16" if dtype == mstype.bfloat16 else "fp16"
+        if hasattr(node_inst, to_float_flag) and getattr(node_inst, to_float_flag):
             return False
     allowed_list.append(node.get_instance())
     return True
 
 
-def _insert_cast_operator_process(node, stree):
+def _insert_cast_operator_process(node, dtype):
     """insert cast for operators in white_list."""
+    dtype_str = "mindspore.bfloat16" if dtype == mstype.bfloat16 else "mindspore.float16"
     new_cast_node = None
-    # insert cast float16 before the primitive operators
+    stree = node.get_symbol_tree()
+    # insert cast fp16/bf16 before the primitive operators
     if issubclass(node.get_instance_type(), Primitive):
         for idx, arg in enumerate(node.get_args()):
             position = stree.before(node)
-            new_node = P.Cast()
-            cast_args = ms.rewrite.ScopedValue.create_name_values([arg.value, "mindspore.float16"], [arg.scope, ""])
-            cast_targets = ms.rewrite.ScopedValue.create_name_values([arg.value], [arg.scope])
+            new_node = _amp_cast_op()
+            cast_args = ms.rewrite.ScopedValue.create_name_values([arg.value, dtype_str], [arg.scope, ""])
+            arg_provider = node.get_handler().get_arg_providers()[idx]
+            if arg_provider and len(arg_provider[0].get_target_users(arg_provider[1])) > 1:
+                cast_targets = [stree.unique_name(str(arg))]
+            else:
+                cast_targets = ms.rewrite.ScopedValue.create_name_values([arg.value], [arg.scope])
             new_cast_node = ms.rewrite.Node.create_call_cell(new_node,
                                                              targets=cast_targets,
                                                              args=cast_args,
                                                              name='incast_{}{}'.format(node.get_name(), idx))
             stree.insert(position, new_cast_node)
             node.set_arg_by_node(idx, new_cast_node)
-    # insert cast float16 before the Cell operators
+    # insert cast fp16/bf16 before the Cell operators
     elif issubclass(node.get_instance_type(), nn.Cell):
-        node.get_instance().to_float(mstype.float16)
+        node.get_instance().to_float(dtype)
     # ignore if subclass is not one of (Primitive, nn.Cell)
     else:
         return
 
     # insert cast float32 after the operators
     position = stree.after(node)
-    new_node = P.Cast()
+    new_node = _amp_cast_op()
     cast_args = ms.rewrite.ScopedValue.create_name_values([node.get_targets()[0].value,
                                                            "mindspore.float32"])
     new_cast_node = ms.rewrite.Node.create_call_cell(new_node,
@@ -156,49 +165,102 @@ def _insert_cast_operator_process(node, stree):
                 user.set_arg_by_node(idx, new_cast_node)
 
 
-def _insert_cast_operator_white_list(stree, white_list):
+def _insert_cast_operator_white_list(stree, white_list, dtype):
     """insert cast for operators in white_list."""
     allowed_list = []
-    # Ignore if net called ".to_float(mindspore.float16)"
+    # Ignore if net called ".to_float(dtype)"
     net = stree.get_handler().get_origin_network()
-    if isinstance(net, nn.Cell) and hasattr(net, "to_float_fp16") and net.to_float_fp16:
+    to_float_flag = "bf16" if dtype == mstype.bfloat16 else "fp16"
+    if isinstance(net, nn.Cell) and hasattr(net, to_float_flag) and getattr(net, to_float_flag):
         return
-    for node in stree.nodes():
-        if node.get_targets() is None:
-            continue
+    node_list = []
+    node_list.extend(list(stree.nodes()))
+    while node_list:
+        node = node_list.pop()
         if node.get_node_type() == ms.rewrite.NodeType.CellContainer:
+            if MS_AMP_BY_REWRITE:
+                _insert_cast_for_cell_container(node, dtype, allowed_list, white_list=white_list)
             for n in node.get_handler().node_list:
                 if n.get_node_type() == ms.rewrite.NodeType.Tree:
                     _insert_cast_operator_white_list(ms.rewrite.TreeNodeHelper.get_sub_tree(ms.rewrite.Node(n)),
-                                                     white_list)
+                                                     white_list, dtype)
         elif node.get_node_type() == ms.rewrite.NodeType.Tree:
             substree = ms.rewrite.TreeNodeHelper.get_sub_tree(node)
-            _insert_cast_operator_white_list(substree, white_list)
-        elif node.get_instance_type() in white_list and _allow_mix_precision(node, allowed_list):
-            _insert_cast_operator_process(node, stree)
+            _insert_cast_operator_white_list(substree, white_list, dtype)
+        elif node.get_node_type() in [ms.rewrite.NodeType.CallFunction, ms.rewrite.NodeType.ControlFlow]:
+            if isinstance(node.get_handler(), ms.rewrite.node.NodeManager):
+                nodes = [ms.rewrite.Node(n) for n in node.get_handler().nodes()]
+                node_list.extend(nodes)
+        elif node.get_instance_type() in white_list and _allow_mix_precision(node, allowed_list, dtype):
+            _insert_cast_operator_process(node, dtype)
 
 
-def _need_removed_cast_pair(node):
+def _insert_cast_for_cell_container(cell_container, dtype, allowed_list, *, white_list=None, black_list=None):
+    """
+    Insert cast for cell containers.
+    Only one of white_list and black_list can be set.
+    """
+
+    class CastNet(nn.Cell):
+        """Cast net"""
+        def __init__(self, dtype):
+            super().__init__()
+            self.cast = _amp_cast_op()
+            self.dtype = dtype
+
+        def construct(self, x):
+            return self.cast(x, self.dtype)
+
+    cast_flag = False
+    current_node = None
+    stree = cell_container.get_symbol_tree()
+    for node in cell_container.get_handler().nodes():
+        current_node = ms.rewrite.Node(node)
+        if (white_list is not None and current_node.get_instance_type() in white_list) or \
+           (black_list is not None and current_node.get_instance_type() not in black_list) and \
+           (_allow_mix_precision(current_node, allowed_list, dtype)):
+            cast_flag = True
+            current_node.get_instance().to_float(dtype)
+        elif cast_flag:
+            # cast next node back to float32
+            current_node.get_instance().to_float(mstype.float32)
+            cast_flag = False
+    if cast_flag and current_node:
+        # if last node in cell_container is casted to fp16/bf16, insert a cast node to cast value back to fp32
+        cast_node = ms.rewrite.Node.create_call_cell(cell=CastNet(mstype.float32),
+                                                     args=[current_node.get_targets()[0]],
+                                                     targets=[current_node.get_targets()[0]],
+                                                     name=f"outcast_{cell_container.get_name()}")
+        stree.insert(stree.after(current_node), cast_node)
+
+
+def _need_removed_cast_pair(node, dtype):
     """check whether the cast pairs should be removed."""
-    cast_dtypes = ms.rewrite.ScopedValue.create_name_values(["mindspore.float16", "mindspore.float32"])
+    dtype_str = "mindspore.bfloat16" if dtype == mstype.bfloat16 else "mindspore.float16"
+    cast_dtypes = ms.rewrite.ScopedValue.create_name_values([dtype_str, "mindspore.float32"])
     cast_dtype_f16 = cast_dtypes[0]
     cast_dtype_f32 = cast_dtypes[1]
-    # current node should be P.Cast()(x, mindspore.float32)
-    if node.get_instance_type() != P.Cast:
+    # current node should be Cast Op to float32
+    if node.get_instance_type() != _amp_cast_op:
         return False
     node_cast_type = node.get_args()[1]
     if node_cast_type != cast_dtype_f32:
         return False
-    # all user nodes should be P.Cast()(x, mindspore.float16) or Cell with to_float(mindspore.float16)
+    # all user nodes should be Cast Op to dtype or Cell with to_float(dtype)
     if not node.get_users():
         return False
+    all_nodes = [ms.rewrite.Node(n) for n in node.get_handler().get_node_manager().nodes()]
     for user in node.get_users():
+        # If ControlFlow node(if statement) exists between current node and user node,
+        # cast pair should not be removed.
+        middle_nodes = all_nodes[all_nodes.index(node): all_nodes.index(user)]
+        if any([n.get_node_type() == ms.rewrite.NodeType.ControlFlow for n in middle_nodes]):
+            return False
         if isinstance(user.get_instance(), nn.Cell):
-            if not hasattr(user.get_instance(), "to_float_fp16"):
-                return False
-            if not user.get_instance().to_float_fp16:
+            to_float_flag = "bf16" if dtype == mstype.bfloat16 else "fp16"
+            if not (hasattr(user.get_instance(), to_float_flag) and getattr(user.get_instance(), to_float_flag)):
                 return False
-        elif user.get_instance_type() == P.Cast:
+        elif user.get_instance_type() == _amp_cast_op:
             user_cast_type = user.get_args()[1]
             if user_cast_type != cast_dtype_f16:
                 return False
@@ -207,11 +269,13 @@ def _need_removed_cast_pair(node):
     return True
 
 
-def _removed_cast_pair_process(stree, cast_f32_node):
+def _removed_cast_pair_process(cast_f32_node):
     """remove the duplicated cast operators."""
-    for user_node in cast_f32_node.get_users():
-        # remove cast f16 nodes
-        if user_node.get_instance_type() == P.Cast:
+    stree = cast_f32_node.get_symbol_tree()
+    cast_f32_users = cast_f32_node.get_users()
+    # remove cast f16 nodes
+    for user_node in cast_f32_users:
+        if user_node.get_instance_type() == _amp_cast_op:
             cast_f16_node = user_node
             # modify arguments using cast_f16's target[0] to cast_f32's args[0], which is f16 type
             for cast_f16_user in cast_f16_node.get_users():
@@ -229,34 +293,78 @@ def _removed_cast_pair_process(stree, cast_f32_node):
     stree.erase(cast_f32_node)
 
 
-def _remove_duplicated_cast(stree):
+def _remove_duplicated_cast(stree, dtype):
     """remove the duplicated cast operators."""
-    for node in stree.nodes():
-        if node.get_targets() is None:
-            continue
+    node_list = []
+    node_list.extend(list(stree.nodes()))
+    while node_list:
+        node = node_list.pop()
         if node.get_node_type() == ms.rewrite.NodeType.CellContainer:
             for n in node.get_handler().node_list:
                 if n.get_node_type() == ms.rewrite.NodeType.Tree:
-                    _remove_duplicated_cast(ms.rewrite.TreeNodeHelper.get_sub_tree(ms.rewrite.Node(n)))
+                    _remove_duplicated_cast(ms.rewrite.TreeNodeHelper.get_sub_tree(ms.rewrite.Node(n)), dtype)
         elif node.get_node_type() == ms.rewrite.NodeType.Tree:
             substree = ms.rewrite.TreeNodeHelper.get_sub_tree(node)
-            _remove_duplicated_cast(substree)
-        elif _need_removed_cast_pair(node):
-            _removed_cast_pair_process(stree, node)
+            _remove_duplicated_cast(substree, dtype)
+        elif node.get_node_type() in [ms.rewrite.NodeType.CallFunction, ms.rewrite.NodeType.ControlFlow]:
+            if isinstance(node.get_handler(), ms.rewrite.node.NodeManager):
+                nodes = [ms.rewrite.Node(n) for n in node.get_handler().nodes()]
+                node_list.extend(nodes)
+        elif _need_removed_cast_pair(node, dtype):
+            _removed_cast_pair_process(node)
 
 
-def _auto_white_list(network, white_list):
+def _auto_white_list(network, white_list, dtype):
     """process the white list of network."""
-    global STREE
-    STREE = ms.rewrite.SymbolTree.create(network)
-    _insert_cast_operator_white_list(STREE, white_list)
-    _remove_duplicated_cast(STREE)
-    return STREE.get_network()
+    stree = ms.rewrite.SymbolTree.create(network)
+    _insert_cast_operator_white_list(stree, white_list, dtype)
+    _remove_duplicated_cast(stree, dtype)
+    return stree.get_network()
 
 
-def _auto_black_list(network, black_list):
+def _insert_cast_operator_black_list(stree, black_list, dtype):
+    """insert cast for operators not in black_list."""
+    allowed_list = []
+    # Ignore if net called ".to_float(dtype)"
+    net = stree.get_handler().get_origin_network()
+    to_float_flag = "bf16" if dtype == mstype.bfloat16 else "fp16"
+    if isinstance(net, nn.Cell) and hasattr(net, to_float_flag) and getattr(net, to_float_flag):
+        return
+    for node in stree.nodes(all_nodes=True):
+        if node.get_targets() is None:
+            continue
+        if node.get_node_type() == ms.rewrite.NodeType.CellContainer:
+            _insert_cast_for_cell_container(node, dtype, allowed_list, black_list=black_list)
+        elif isinstance(node.get_handler().get_node_manager(), ms.rewrite.node.CellContainer):
+            # nodes in CellContainer are processed by _insert_cast_for_cell_container
+            continue
+        elif node.get_instance_type() not in black_list and _allow_mix_precision(node, allowed_list, dtype):
+            _insert_cast_operator_process(node, dtype)
+
+
+def _remove_duplicated_cast_rewrite(stree, dtype):
+    """remove the duplicated cast operators."""
+    for node in stree.nodes(all_nodes=True):
+        if _need_removed_cast_pair(node, dtype):
+            user_nodes = node.get_users()
+            # remove cast f16 nodes
+            for user_node in user_nodes:
+                if user_node.get_instance_type() == _amp_cast_op:
+                    stree.erase(user_node)
+            # remove the cast f32 node
+            stree.erase(node)
+
+
+def _auto_black_list_rewrite(network, black_list, dtype):
+    stree = ms.rewrite.SymbolTree.create(network)
+    _insert_cast_operator_black_list(stree, black_list, dtype)
+    _remove_duplicated_cast_rewrite(stree, dtype)
+    return stree.get_network()
+
+
+def _auto_black_list(network, black_list, dtype):
     """process the black list of network."""
-    network.to_float(mstype.float16)
+    network.to_float(dtype)
     cells = network.name_cells()
     change = False
     for name in cells:
@@ -264,30 +372,27 @@ def _auto_black_list(network, black_list):
         if subcell == network:
             continue
         if isinstance(subcell, tuple(black_list)):
-            network._cells[name] = _OutputTo16(subcell.to_float(mstype.float32))
+            network._cells[name] = _OutputTo16(subcell.to_float(mstype.float32), dtype)
             change = True
         else:
-            _auto_black_list(subcell, black_list)
+            _auto_black_list(subcell, black_list, dtype)
     if isinstance(network, nn.SequentialCell) and change:
         network.cell_list = list(network.cells())
+    return network
 
 
-def auto_mixed_precision(network, amp_level="O0"):
+def auto_mixed_precision(network, amp_level="O0", dtype=mstype.float16):
     """
     Returns a network processed with auto mixed precision.
 
     This interface will automatically perform mixed-precision processing on the input network, and the cells
-    and operators in the processed network will add precision conversion operations to calculate with float16 accuracy.
-    Inputs and parameters of cells and operators are converted to float16 type, and calculation results are converted
-    back to float32 type.
+    and operators in the processed network will add precision conversion operations to calculate with lower
+    precision: ``mstype.float16`` or ``mstype.bfloat16`` . Inputs and parameters of cells and operators are
+    converted to lower precision float, and calculation results are converted back to full precision float,
+    i.e. ``mstype.float32`` .
 
     The framework has a set of built-in blacklists and whitelists, and the `amp_level` determines which cells and
-    operators are specifically converted:
-
-    - When `amp_level="O0"` , no precision conversion is performed.
-    - When `amp_level="O1"` , only the cells and operators in the whitelist will be converted.
-    - When `amp_level="O2"` , all cells and operators except those in the blacklist will be converted.
-    - When `amp_level="O3"` , all cells and operators in the network are converted.
+    operators are specifically converted.
 
     The current built-in whitelist contents are:
 
@@ -305,26 +410,38 @@ def auto_mixed_precision(network, amp_level="O0"):
     :class:`mindspore.nn.LayerNorm`]
 
     For details on automatic mixed precision, refer to
-    `Automatic Mix Precision <https://www.mindspore.cn/tutorials/en/r2.1/advanced/mixed_precision.html>`_ .
+    `Automatic Mix Precision <https://www.mindspore.cn/tutorials/en/r2.2/advanced/mixed_precision.html>`_ .
+
+    Note:
+        - Repeatedly calling mixed-precision interfaces, such as `custom_mixed_precision` and `auto_mixed_precision`,
+          can result in a larger network hierarchy and slower performance.
+        - If interfaces like `Model` and `build_train_network` is used to train the network which is converted by
+          mixed-precision interfaces such as `custom_mixed_precision` and `auto_mixed_precision`, `amp_level`
+          need to be configured to ``O0`` to avoid the duplicated accuracy conversion.
 
     Args:
         network (Cell): Definition of the network.
         amp_level (str): Supports ["O0", "O1", "O2", "O3"]. Default: ``"O0"`` .
 
             - "O0": Do not change.
-            - "O1": Convert cells and operators in whitelist to float16 precision operations, and keep float32
+            - "O1": Convert cells and operators in whitelist to lower precision operations, and keep full
               precision operations for the rest.
-            - "O2": Keep float32 precision operations for cells and operators in blacklist, and convert the rest
-              to float16 precision operations.
-            - "O3": Cast network to float16.
+            - "O2": Keep full precision operations for cells and operators in blacklist, and convert the rest
+              to lower precision operations.
+            - "O3": Cast network to lower precision.
+
+        dtype (Type): The type used in lower precision calculations, can be ``mstype.float16`` or ``mstype.bfloat16`` ,
+            default: ``mstype.float16`` .
 
     Raises:
-        ValueError: If amp level is not supported.
+        TypeError: If `network` is not a Cell.
+        ValueError: If `dtype` is not one of ``mstype.float16`` , ``mstype.bfloat16`` .
+        ValueError: If `amp_level` is not within the supported range.
 
     Examples:
         >>> from mindspore import amp
         >>> # Define the network structure of LeNet5. Refer to
-        >>> # https://gitee.com/mindspore/docs/blob/r2.1/docs/mindspore/code/lenet.py
+        >>> # https://gitee.com/mindspore/docs/blob/r2.2/docs/mindspore/code/lenet.py
         >>> network = LeNet5()
         >>> amp_level = "O1"
         >>> net = amp.auto_mixed_precision(network, amp_level)
@@ -332,20 +449,37 @@ def auto_mixed_precision(network, amp_level="O0"):
     if not isinstance(network, nn.Cell):
         raise TypeError("The network type should be Cell.")
 
+    if dtype not in (mstype.float16, mstype.bfloat16):
+        raise ValueError(f"The dtype should be one of (mstype.float16, mstype.bfloat16), but got {dtype}.")
+
     if amp_level == "O0":
         return network
 
-    if amp_level == "O1":
-        return _auto_white_list(network, AMP_WHITE_LIST)
+    # Return network if the same amp level has already been configurated
+    if getattr(network, "_amp_level") in ("O1", "O2", "O3"):
+        logger.warning(f"The network's auto mixed-precision level is adjusted from {getattr(network, '_amp_level')} "
+                       f"to {amp_level}, and repeated calls to mixed-precision interfaces can cause performance "
+                       f"degradation.")
 
-    if amp_level == "O2":
-        _auto_black_list(network, AMP_BLACK_LIST)
+    if amp_level == "O1":
+        network = _auto_white_list(network, AMP_WHITE_LIST, dtype)
+    elif amp_level == "O2":
+        if MS_AMP_BY_REWRITE:
+            network = _auto_black_list_rewrite(network, AMP_BLACK_LIST, dtype)
+        else:
+            network = _auto_black_list(network, AMP_BLACK_LIST, dtype)
+            network = _OutputTo32(network)
     elif amp_level == "O3":
-        network.to_float(mstype.float16)
+        if MS_AMP_BY_REWRITE:
+            network = _auto_black_list_rewrite(network, [], dtype)
+        else:
+            network.to_float(dtype)
+            network = _OutputTo32(network)
     else:
         raise ValueError("The amp level {} is not supported".format(amp_level))
-    if amp_level in ("O2", "O3"):
-        network = _OutputTo32(network)
+
+    setattr(network, "_amp_level", amp_level)
+
     return network
 
 
@@ -436,8 +570,7 @@ def _add_loss_network(network, loss_fn, cast_model_type):
             super(WithLossCell, self).__init__(auto_prefix=False)
             self._backbone = backbone
             self._loss_fn = loss_fn
-            if isinstance(backbone, nn.Cell) and backbone.jit_config_dict:
-                self._jit_config_dict = backbone.jit_config_dict
+            self._get_attr_from_cell(backbone)
 
         def construct(self, data, label):
             out = self._backbone(data)
@@ -452,6 +585,39 @@ def _add_loss_network(network, loss_fn, cast_model_type):
     return network
 
 
+def _is_grad_accumulation(mcell):
+    if mcell.cls_name == "GradAccumulationCell":
+        return True
+    for cell in mcell.cells():
+        if _is_grad_accumulation(cell):
+            return True
+    return False
+
+
+def _auto_mixed_precision_process(network, config, level):
+    """Auto mixed precision process."""
+    if MS_AMP_BY_REWRITE:
+        if config["cast_model_type"] == mstype.float16 or level == "O2":
+            level = "O2" if config["keep_batchnorm_fp32"] else "O3"
+        elif config["cast_model_type"] == mstype.float32 and level in ("O2", "O3"):
+            # cast_model_type set by kwargs
+            level = "O0"
+        network = auto_mixed_precision(network, level)
+    else:
+        if config["cast_model_type"] == mstype.float16:
+            network.to_float(mstype.float16)
+
+            if config["keep_batchnorm_fp32"]:
+                _do_keep_batchnorm_fp32(network)
+        elif not config["keep_batchnorm_fp32"] and level == "O2":
+            network.to_float(mstype.float16)
+        elif config["cast_model_type"] == mstype.float32 and level in ("O2", "O3"):
+            pass
+        else:
+            network = auto_mixed_precision(network, level)
+    return network
+
+
 def build_train_network(network, optimizer, loss_fn=None, level='O0', boost_level='O0', **kwargs):
     """
     Build the mixed precision training cell automatically.
@@ -510,7 +676,7 @@ def build_train_network(network, optimizer, loss_fn=None, level='O0', boost_leve
     Examples:
         >>> from mindspore import amp, nn
         >>> # Define the network structure of LeNet5. Refer to
-        >>> # https://gitee.com/mindspore/docs/blob/r2.1/docs/mindspore/code/lenet.py
+        >>> # https://gitee.com/mindspore/docs/blob/r2.2/docs/mindspore/code/lenet.py
         >>> network = LeNet5()
         >>> net_loss = nn.SoftmaxCrossEntropyWithLogits(reduction="mean")
         >>> net_opt = nn.Momentum(network.trainable_params(), learning_rate=0.01, momentum=0.9)
@@ -525,17 +691,7 @@ def build_train_network(network, optimizer, loss_fn=None, level='O0', boost_leve
     _check_kwargs(kwargs)
     config = dict(_config_level.get(level), **kwargs)
 
-    if config["cast_model_type"] == mstype.float16:
-        network.to_float(mstype.float16)
-
-        if config["keep_batchnorm_fp32"]:
-            _do_keep_batchnorm_fp32(network)
-    elif not config["keep_batchnorm_fp32"] and level == "O2":
-        network.to_float(mstype.float16)
-    elif config["cast_model_type"] == mstype.float32 and level in ("O2", "O3"):
-        pass
-    else:
-        network = auto_mixed_precision(network, level)
+    network = _auto_mixed_precision_process(network, config, level)
 
     if loss_fn:
         network = _add_loss_network(network, loss_fn, config["cast_model_type"])
@@ -551,8 +707,8 @@ def build_train_network(network, optimizer, loss_fn=None, level='O0', boost_leve
                 raise ValueError("Only `loss_scale_manager=None` or "
                                  "`loss_scale_manager=FixedLossScaleManager(drop_overflow_update=False)`"
                                  "are supported on device `CPU`. ")
-            if _get_pipeline_stages() > 1:
-                network = _TrainPipelineWithLossScaleCell(network, optimizer,
+            if _get_pipeline_stages() > 1 or _is_grad_accumulation(network):
+                network = _TrainGradAccuWithLossScaleCell(network, optimizer,
                                                           scale_sense=update_cell).set_train()
             elif enable_boost:
                 network = boost.BoostTrainOneStepWithLossScaleCell(network, optimizer,
@@ -561,8 +717,8 @@ def build_train_network(network, optimizer, loss_fn=None, level='O0', boost_leve
                 network = nn.TrainOneStepWithLossScaleCell(network, optimizer,
                                                            scale_sense=update_cell).set_train()
             return network
-    if _get_pipeline_stages() > 1:
-        network = _TrainPipelineAccuStepCell(network, optimizer).set_train()
+    if _get_pipeline_stages() > 1 or _is_grad_accumulation(network):
+        network = _TrainGradAccuStepCell(network, optimizer).set_train()
     elif enable_boost:
         network = boost.BoostTrainOneStepCell(network, optimizer, loss_scale).set_train()
     else:
@@ -586,6 +742,23 @@ def get_white_list():
 
     Returns:
         list, A copy of internal white list.
+
+    Examples:
+        >>> from mindspore import amp
+        >>> white_list = amp.get_white_list()
+        >>> print(white_list)
+        [<class 'mindspore.nn.layer.conv.Conv1d'>, <class 'mindspore.nn.layer.conv.Conv2d'>,
+         <class 'mindspore.nn.layer.conv.Conv3d'>, <class 'mindspore.nn.layer.conv.Conv1dTranspose'>,
+         <class 'mindspore.nn.layer.conv.Conv2dTranspose'>, <class 'mindspore.nn.layer.conv.Conv3dTranspose'>,
+         <class 'mindspore.nn.layer.basic.Dense'>, <class 'mindspore.nn.layer.rnn_cells.LSTMCell'>,
+         <class 'mindspore.nn.layer.rnn_cells.RNNCell'>, <class 'mindspore.nn.layer.rnn_cells.GRUCell'>,
+         <class 'mindspore.ops.operations.nn_ops.Conv2D'>, <class 'mindspore.ops.operations.nn_ops.Conv3D'>,
+         <class 'mindspore.ops.operations.nn_ops.Conv2DTranspose'>,
+         <class 'mindspore.ops.operations.nn_ops.Conv3DTranspose'>,
+         <class 'mindspore.ops.operations.nn_ops.Conv2DBackpropInput'>,
+         <class 'mindspore.ops.operations.math_ops.MatMul'>, <class 'mindspore.ops.operations.math_ops.BatchMatMul'>,
+         <class 'mindspore.ops.operations.nn_ops.PReLU'>, <class 'mindspore.ops.operations.nn_ops.ReLU'>,
+         <class 'mindspore.ops.operations.math_ops.Ger'>]
     """
     white_list = AMP_WHITE_LIST.copy()
     return white_list
@@ -602,24 +775,31 @@ def get_black_list():
 
     Returns:
         list, A copy of internal black list.
+
+    Examples:
+        >>> from mindspore import amp
+        >>> black_list = amp.get_black_list()
+        >>> print(black_list)
+        [<class 'mindspore.nn.layer.normalization.BatchNorm1d'>, <class 'mindspore.nn.layer.normalization.BatchNorm2d'>,
+         <class 'mindspore.nn.layer.normalization.BatchNorm3d'>, <class 'mindspore.nn.layer.normalization.LayerNorm'>]
     """
     black_list = AMP_BLACK_LIST.copy()
     return black_list
 
 
-def custom_mixed_precision(network, *, white_list=None, black_list=None):
+def custom_mixed_precision(network, *, white_list=None, black_list=None, dtype=mstype.float16):
     """
     Custom mixed precision by setting whitelist or blacklist.
     When the `white_list` is provided, primitives and cells in `white_list` will perform the precision conversion.
-    When the `black_list` is provided, cells that are not in `black_list` will perform the pereision
-    conversion.
+    When the `black_list` is provided, cells that are not in `black_list` will perform the pereision conversion.
     Only one of `white_list` and `black_list` should be provided.
 
     Note:
-        - After using `custom_mixed_precision` for precision conversion, it is not supported to use other interfaces
-          for precision conversion again. If interfaces like `Model` and `build_train_network` is used to train
-          the converted network, `amp_level` need to be configured to ``O0`` to avoid the duplicated accuracy
-          conversion.
+        - Repeatedly calling mixed-precision interfaces, such as `custom_mixed_precision` and `auto_mixed_precision`,
+          can result in a larger network hierarchy and slower performance.
+        - If interfaces like `Model` and `build_train_network` is used to train the network which is converted by
+          mixed-precision interfaces such as `custom_mixed_precision` and `auto_mixed_precision`, `amp_level`
+          need to be configured to ``O0`` to avoid the duplicated accuracy conversion.
         - Primitives for blacklist is not support yet.
 
     Args:
@@ -628,6 +808,8 @@ def custom_mixed_precision(network, *, white_list=None, black_list=None):
             white list is not used.
         black_list (list[Cell], optional): Black list of custom mixed precision. Defaults: ``None`` , means
             black list is not used.
+        dtype (Type): The type used in lower precision calculations, can be ``mstype.float16`` or ``mstype.bfloat16`` ,
+            default: ``mstype.float16`` .
 
     Returns:
         network (Cell), A network supporting mixed precision.
@@ -635,12 +817,13 @@ def custom_mixed_precision(network, *, white_list=None, black_list=None):
     Raises:
         TypeError: The network type is not Cell.
         ValueError: Neither `white_list` nor `black_list` is provided.
+        ValueError: If `dtype` is not one of ``mstype.float16`` , ``mstype.bfloat16`` .
         ValueError: Both `white_list` and `black_list` are provided.
 
     Examples:
         >>> from mindspore import amp, nn
         >>> # Define the network structure of LeNet5. Refer to
-        >>> # https://gitee.com/mindspore/docs/blob/r2.1/docs/mindspore/code/lenet.py
+        >>> # https://gitee.com/mindspore/docs/blob/r2.2/docs/mindspore/code/lenet.py
         >>> net = LeNet5()
         >>> custom_white_list = amp.get_white_list()
         >>> custom_white_list.append(nn.Flatten)
@@ -656,13 +839,19 @@ def custom_mixed_precision(network, *, white_list=None, black_list=None):
         raise ValueError("For custom_mixed_precision, the white_list or black_list cannot be provided "
                          "at the same time, please provide one or the other.")
 
+    if dtype not in (mstype.float16, mstype.bfloat16):
+        raise ValueError(f"The dtype should be one of (mstype.float16, mstype.bfloat16), but got {dtype}.")
+
     if white_list is not None:
         _list_check(white_list, "white_list")
-        return _auto_white_list(network, white_list)
-
-    _list_check(black_list, "black_list")
-    _auto_black_list(network, black_list)
-    network = _OutputTo32(network)
+        network = _auto_white_list(network, white_list, dtype)
+    else:
+        _list_check(black_list, "black_list")
+        if MS_AMP_BY_REWRITE:
+            network = _auto_black_list_rewrite(network, black_list, dtype)
+        else:
+            network = _auto_black_list(network, black_list, dtype)
+            network = _OutputTo32(network)
     return network
 
 
@@ -693,3 +882,14 @@ def _list_check(custom_list: list, list_name: str):
         for elem in AMP_BLACK_LIST:
             if elem not in custom_list:
                 logger.warning(f"{elem} is removed from internal black list.")
+
+def _config_amp(*, enable_rewrite: bool = None, cast_op: type = None): # pylint: disable=unused-variable
+    """Configure auto mixed precision."""
+    global MS_AMP_BY_REWRITE
+    global _amp_cast_op
+
+    if enable_rewrite is not None:
+        MS_AMP_BY_REWRITE = enable_rewrite
+
+    if cast_op is not None:
+        _amp_cast_op = cast_op