chore: import upstream snapshot with attribution
This commit is contained in:
wehub-resource-sync
@@ -0,0 +1,600 @@
|
||||
# Copyright (c) 2021 PaddlePaddle Authors. All Rights Reserved.
|
||||
#
|
||||
# Licensed under the Apache License, Version 2.0 (the "License");
|
||||
# you may not use this file except in compliance with the License.
|
||||
# You may obtain a copy of the License at
|
||||
|
||||
#
|
||||
# http://www.apache.org/licenses/LICENSE-2.0
|
||||
#
|
||||
# Unless required by applicable law or agreed to in writing, software
|
||||
# distributed under the License is distributed on an "AS IS" BASIS,
|
||||
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
# See the License for the specific language governing permissions and
|
||||
# limitations under the License.
|
||||
|
||||
import collections
|
||||
import logging
|
||||
import struct
|
||||
|
||||
import numpy as np
|
||||
|
||||
import paddle
|
||||
from paddle.base import core, framework, global_scope
|
||||
from paddle.base.log_helper import get_logger
|
||||
from paddle.base.wrapped_decorator import signature_safe_contextmanager
|
||||
|
||||
from ..fp16_utils import (
|
||||
_rename_arg,
|
||||
_rename_op_input,
|
||||
find_true_post_op,
|
||||
find_true_prev_op,
|
||||
)
|
||||
from .amp_lists import AutoMixedPrecisionListsBF16
|
||||
|
||||
_logger = get_logger(
|
||||
__name__, logging.INFO, fmt='%(asctime)s-%(levelname)s: %(message)s'
|
||||
)
|
||||
|
||||
_valid_types = [
|
||||
core.VarDesc.VarType.DENSE_TENSOR,
|
||||
core.VarDesc.VarType.SELECTED_ROWS,
|
||||
core.VarDesc.VarType.DENSE_TENSOR_ARRAY,
|
||||
]
|
||||
|
||||
_bf16_guard_pattern = "__use_bf16__"
|
||||
|
||||
|
||||
def convert_float_to_uint16(in_list):
|
||||
in_list = np.asarray(in_list)
|
||||
out = np.vectorize(
|
||||
lambda x: struct.unpack('<I', struct.pack('<f', x))[0] >> 16,
|
||||
otypes=[np.uint16],
|
||||
)(in_list.flat)
|
||||
return np.reshape(out, in_list.shape)
|
||||
|
||||
|
||||
def _dtype_to_str(dtype):
|
||||
"""
|
||||
Convert specific variable type to its corresponding string.
|
||||
|
||||
Args:
|
||||
dtype (VarType): Variable type.
|
||||
"""
|
||||
if dtype == paddle.bfloat16:
|
||||
return 'bf16'
|
||||
else:
|
||||
return 'fp32'
|
||||
|
||||
|
||||
def _insert_cast_op(block, op, idx, src_dtype, dest_dtype):
|
||||
"""
|
||||
Insert cast op and rename args of input and output.
|
||||
|
||||
Args:
|
||||
block (Program): The block in which the operator is.
|
||||
op (Operator): The operator to insert cast op.
|
||||
idx (int): The index of current operator.
|
||||
src_dtype (VarType): The input variable dtype of cast op.
|
||||
dest_dtype (VarType): The output variable dtype of cast op.
|
||||
|
||||
Returns:
|
||||
num_cast_op (int): The number of cast ops that have been inserted.
|
||||
"""
|
||||
num_cast_ops = 0
|
||||
|
||||
for in_name in op.input_names:
|
||||
if src_dtype == paddle.float32 and op.type in [
|
||||
'batch_norm',
|
||||
'fused_bn_add_activation',
|
||||
'layer_norm',
|
||||
]:
|
||||
if in_name not in {'X', 'Z'}:
|
||||
continue
|
||||
for in_var_name in op.input(in_name):
|
||||
in_var = block.var(in_var_name)
|
||||
if in_var.type not in _valid_types or in_var.dtype == dest_dtype:
|
||||
continue
|
||||
if in_var.dtype == src_dtype:
|
||||
cast_name = in_var.name + '.cast_' + _dtype_to_str(dest_dtype)
|
||||
out_var = block.vars.get(cast_name)
|
||||
if out_var is None or out_var.dtype != dest_dtype:
|
||||
out_var = block.create_var(
|
||||
name=cast_name,
|
||||
dtype=dest_dtype,
|
||||
persistable=False,
|
||||
stop_gradient=in_var.stop_gradient,
|
||||
)
|
||||
|
||||
block._insert_op(
|
||||
idx,
|
||||
type="cast",
|
||||
inputs={"X": in_var},
|
||||
outputs={"Out": out_var},
|
||||
attrs={
|
||||
"in_dtype": in_var.dtype,
|
||||
"out_dtype": out_var.dtype,
|
||||
},
|
||||
)
|
||||
num_cast_ops += 1
|
||||
_rename_arg(op, in_var.name, out_var.name)
|
||||
else:
|
||||
if op.has_attr('in_dtype'):
|
||||
op._set_attr('in_dtype', dest_dtype)
|
||||
if src_dtype == paddle.float32 and dest_dtype == paddle.bfloat16:
|
||||
for out_name in op.output_names:
|
||||
if (
|
||||
op.type
|
||||
in ['batch_norm', 'fused_bn_add_activation', 'layer_norm']
|
||||
and out_name != 'Y'
|
||||
):
|
||||
continue
|
||||
for out_var_name in op.output(out_name):
|
||||
out_var = block.var(out_var_name)
|
||||
if out_var.type not in _valid_types:
|
||||
continue
|
||||
if out_var.dtype == paddle.float32:
|
||||
out_var.desc.set_dtype(core.VarDesc.VarType.BF16)
|
||||
if op.has_attr('out_dtype'):
|
||||
op._set_attr('out_dtype', core.VarDesc.VarType.BF16)
|
||||
return num_cast_ops
|
||||
|
||||
|
||||
def _insert_cast_post_op(
|
||||
block, op, idx, src_dtype, dest_dtype, target_name, op_var_rename_map
|
||||
):
|
||||
num_cast_ops = 0
|
||||
target_var = block.var(target_name)
|
||||
if target_var.type not in _valid_types or target_var.dtype == dest_dtype:
|
||||
return num_cast_ops
|
||||
|
||||
assert target_var.dtype == src_dtype, (
|
||||
f"The real dtype({_dtype_to_str(target_var.dtype)}) is not equal to the src dtype({_dtype_to_str(src_dtype)})"
|
||||
)
|
||||
|
||||
cast_name = target_var.name + '.cast_' + _dtype_to_str(dest_dtype)
|
||||
cast_var = block.vars.get(cast_name)
|
||||
if cast_var is None or cast_var.dtype != dest_dtype:
|
||||
cast_var = block.create_var(
|
||||
name=cast_name,
|
||||
dtype=dest_dtype,
|
||||
persistable=False,
|
||||
stop_gradient=target_var.stop_gradient,
|
||||
)
|
||||
block._insert_op(
|
||||
idx,
|
||||
type="cast",
|
||||
inputs={"X": target_var},
|
||||
outputs={"Out": cast_var},
|
||||
attrs={"in_dtype": target_var.dtype, "out_dtype": cast_var.dtype},
|
||||
)
|
||||
num_cast_ops += 1
|
||||
op_var_rename_map[block.idx][target_var.name] = cast_var.name
|
||||
|
||||
return num_cast_ops
|
||||
|
||||
|
||||
def _is_in_fp32_varnames(op, amp_lists):
|
||||
if not amp_lists.fp32_varnames:
|
||||
return False
|
||||
|
||||
for in_name in op.input_arg_names:
|
||||
if in_name in amp_lists.fp32_varnames:
|
||||
return True
|
||||
|
||||
for out_name in op.output_arg_names:
|
||||
if out_name in amp_lists.fp32_varnames:
|
||||
return True
|
||||
|
||||
return False
|
||||
|
||||
|
||||
def _need_keep_fp32(op, unsupported_op_list, use_bf16_guard):
|
||||
if op.type in unsupported_op_list:
|
||||
# the highest priority condition: If ops don't have bf16 computing kernels,
|
||||
# they must be executed in fp32 calculation pattern.
|
||||
return True
|
||||
|
||||
# process ops about learning rate
|
||||
in_out_arg_names = []
|
||||
in_out_arg_names.extend(list(op.input_arg_names))
|
||||
in_out_arg_names.extend(list(op.output_arg_names))
|
||||
for name in in_out_arg_names:
|
||||
if "learning_rate" in name:
|
||||
return True
|
||||
|
||||
if use_bf16_guard:
|
||||
if op.has_attr("op_namescope") and (
|
||||
_bf16_guard_pattern in op.attr("op_namescope")
|
||||
):
|
||||
# op in bf16 guard
|
||||
return False
|
||||
else:
|
||||
# op not in bf16 guard
|
||||
return True
|
||||
else:
|
||||
return False
|
||||
|
||||
|
||||
@signature_safe_contextmanager
|
||||
def bf16_guard():
|
||||
"""
|
||||
As for the pure bf16 training, if users set `use_bf16_guard` to True,
|
||||
only those ops created in the context manager `bf16_guard` will be
|
||||
transformed as float16 type.
|
||||
|
||||
Examples:
|
||||
.. code-block:: pycon
|
||||
|
||||
>>> import numpy as np
|
||||
>>> import paddle
|
||||
>>> import paddle.nn.functional as F
|
||||
>>> paddle.enable_static()
|
||||
>>> data = paddle.static.data(name='X', shape=[None, 1, 28, 28], dtype='float32')
|
||||
>>> conv2d = paddle.static.nn.conv2d(input=data, num_filters=6, filter_size=3)
|
||||
|
||||
>>> with paddle.static.amp.bf16.bf16_guard():
|
||||
... bn = paddle.static.nn.batch_norm(input=conv2d, act="relu")
|
||||
... pool = F.max_pool2d(bn, kernel_size=2, stride=2)
|
||||
... hidden = paddle.static.nn.fc(pool, size=10)
|
||||
... loss = paddle.mean(hidden)
|
||||
"""
|
||||
with framework.name_scope(prefix=_bf16_guard_pattern):
|
||||
yield
|
||||
|
||||
|
||||
def are_post_ops_bf16(post_ops, keep_fp32_ops):
|
||||
for post_op in post_ops:
|
||||
for op in post_op:
|
||||
if op in keep_fp32_ops:
|
||||
return False
|
||||
return True
|
||||
|
||||
|
||||
def cast_initializers_to_bf16(
|
||||
startup_prog,
|
||||
amp_lists,
|
||||
block,
|
||||
all_ops,
|
||||
keep_fp32_ops,
|
||||
to_bf16_var_names=None,
|
||||
):
|
||||
prepend_ops = startup_prog.global_block().ops
|
||||
for op in prepend_ops:
|
||||
if str(op.type) in amp_lists.bf16_initializer_list:
|
||||
change_op = True
|
||||
op_post_ops = []
|
||||
op_out_vars = []
|
||||
for out_name in op.output_names:
|
||||
for out_var_name in op.output(out_name):
|
||||
out_var = block.var(out_var_name)
|
||||
post_op = find_true_post_op(all_ops, op, out_var_name, True)
|
||||
|
||||
if out_var is None or out_var.type not in _valid_types:
|
||||
change_op = False
|
||||
break
|
||||
op_post_ops.append(post_op)
|
||||
op_out_vars.append(out_var)
|
||||
|
||||
if change_op and are_post_ops_bf16(op_post_ops, keep_fp32_ops):
|
||||
for out_var in op_out_vars:
|
||||
if out_var.dtype == paddle.float32:
|
||||
out_var.desc.set_dtype(core.VarDesc.VarType.BF16)
|
||||
if (
|
||||
to_bf16_var_names is not None
|
||||
and out_var.name in to_bf16_var_names
|
||||
):
|
||||
to_bf16_var_names.remove(out_var.name)
|
||||
if (
|
||||
op.has_attr('dtype')
|
||||
and op.attr('dtype') == core.VarDesc.VarType.FP32
|
||||
):
|
||||
op._set_attr('dtype', core.VarDesc.VarType.BF16)
|
||||
|
||||
|
||||
def cast_model_to_bf16(
|
||||
program, startup_prog=None, amp_lists=None, use_bf16_guard=True
|
||||
):
|
||||
"""
|
||||
Traverse all ops in the whole model and set their inputs and outputs
|
||||
to the bf16 data type. This function will do some special processing for
|
||||
the batch normalization, which will keep the batchnorm's computations in FP32.
|
||||
Args:
|
||||
program (Program): The used program.
|
||||
amp_lists (AutoMixedPrecisionListsBF16): An AutoMixedPrecisionListsBF16 object.
|
||||
use_bf16_guard(bool): Determine whether to use `bf16_guard` when
|
||||
constructing the program. Default True.
|
||||
"""
|
||||
|
||||
if amp_lists is None:
|
||||
amp_lists = AutoMixedPrecisionListsBF16()
|
||||
global_block = program.global_block()
|
||||
keep_fp32_ops = set()
|
||||
to_bf16_var_names = set()
|
||||
to_bf16_pre_cast_ops = set()
|
||||
origin_ops = []
|
||||
for block in program.blocks:
|
||||
origin_ops.extend(block.ops)
|
||||
|
||||
for block in program.blocks:
|
||||
ops = block.ops
|
||||
for op in ops:
|
||||
if op.type == 'create_py_reader' or op.type == 'read':
|
||||
continue
|
||||
if _need_keep_fp32(op, amp_lists.unsupported_list, use_bf16_guard):
|
||||
keep_fp32_ops.add(op)
|
||||
continue # processed below
|
||||
for in_name in op.input_names:
|
||||
if op.type in {
|
||||
'batch_norm',
|
||||
'fused_bn_add_activation',
|
||||
'layer_norm',
|
||||
} and in_name not in {'X', 'Z'}:
|
||||
continue
|
||||
for in_var_name in op.input(in_name):
|
||||
in_var = None
|
||||
try:
|
||||
in_var = block.var(in_var_name)
|
||||
except ValueError as e:
|
||||
_logger.debug(
|
||||
f"-- {e}, try to get it in the global block --"
|
||||
)
|
||||
in_var = global_block.var(in_var_name)
|
||||
if in_var is not None:
|
||||
_logger.debug(
|
||||
f"-- var {in_var_name} is got in the global block --"
|
||||
)
|
||||
|
||||
if in_var is None or in_var.type not in _valid_types:
|
||||
continue
|
||||
|
||||
if in_var.dtype == paddle.float32:
|
||||
in_var.desc.set_dtype(core.VarDesc.VarType.BF16)
|
||||
to_bf16_var_names.add(in_var_name)
|
||||
|
||||
_logger.debug(
|
||||
f"-- op type: {op.type}, in var name: {in_var_name}, in var dtype: {in_var.dtype} --"
|
||||
)
|
||||
|
||||
for out_name in op.output_names:
|
||||
if (
|
||||
op.type
|
||||
in {'batch_norm', 'fused_bn_add_activation', 'layer_norm'}
|
||||
and out_name != 'Y'
|
||||
):
|
||||
continue
|
||||
for out_var_name in op.output(out_name):
|
||||
out_var = None
|
||||
try:
|
||||
out_var = block.var(out_var_name)
|
||||
except ValueError as e:
|
||||
_logger.debug(
|
||||
f"-- {e}, try to get it in the global block --"
|
||||
)
|
||||
out_var = global_block.var(out_var_name)
|
||||
if out_var is not None:
|
||||
_logger.debug(
|
||||
f"-- var {out_var_name} is got in the global block --"
|
||||
)
|
||||
|
||||
if out_var is None or out_var.type not in _valid_types:
|
||||
continue
|
||||
|
||||
if out_var.dtype == paddle.float32:
|
||||
out_var.desc.set_dtype(core.VarDesc.VarType.BF16)
|
||||
|
||||
_logger.debug(
|
||||
f"-- op type: {op.type}, out var name: {out_var_name}, out var dtype: {out_var.dtype} --"
|
||||
)
|
||||
for attr_name in ['in_dtype', 'out_dtype', 'dtype']:
|
||||
if (
|
||||
op.has_attr(attr_name)
|
||||
and op.attr(attr_name) == paddle.float32
|
||||
):
|
||||
op._set_attr(attr_name, core.VarDesc.VarType.BF16)
|
||||
|
||||
if startup_prog is not None:
|
||||
cast_initializers_to_bf16(
|
||||
startup_prog,
|
||||
amp_lists,
|
||||
global_block,
|
||||
ops,
|
||||
keep_fp32_ops,
|
||||
to_bf16_var_names,
|
||||
)
|
||||
|
||||
# process ops in keep_fp32_ops
|
||||
op_var_rename_map = [
|
||||
collections.OrderedDict() for _ in range(len(program.blocks))
|
||||
]
|
||||
for block in program.blocks:
|
||||
ops = block.ops
|
||||
idx = 0
|
||||
while idx < len(ops):
|
||||
op = ops[idx]
|
||||
num_cast_ops = 0
|
||||
if op not in keep_fp32_ops:
|
||||
if op in to_bf16_pre_cast_ops:
|
||||
in_var_cast_num = _insert_cast_op(
|
||||
block,
|
||||
op,
|
||||
idx,
|
||||
core.VarDesc.VarType.FP32,
|
||||
core.VarDesc.VarType.BF16,
|
||||
)
|
||||
num_cast_ops += in_var_cast_num
|
||||
else:
|
||||
pre_cast_num = _insert_cast_op(
|
||||
block,
|
||||
op,
|
||||
idx,
|
||||
core.VarDesc.VarType.BF16,
|
||||
core.VarDesc.VarType.FP32,
|
||||
)
|
||||
num_cast_ops += pre_cast_num
|
||||
for out_var_name in op.output_arg_names:
|
||||
out_var = block.vars.get(out_var_name)
|
||||
if out_var is None or out_var.type not in _valid_types:
|
||||
continue
|
||||
if out_var.dtype == paddle.bfloat16:
|
||||
out_var.desc.set_dtype(core.VarDesc.VarType.FP32)
|
||||
post_ops = find_true_post_op(ops, op, out_var_name)
|
||||
for post_op in post_ops:
|
||||
if post_op in keep_fp32_ops:
|
||||
continue
|
||||
post_cast_num = _insert_cast_post_op(
|
||||
block,
|
||||
op,
|
||||
idx + pre_cast_num + 1,
|
||||
core.VarDesc.VarType.FP32,
|
||||
core.VarDesc.VarType.BF16,
|
||||
out_var_name,
|
||||
op_var_rename_map,
|
||||
)
|
||||
num_cast_ops += post_cast_num
|
||||
idx += num_cast_ops + 1
|
||||
|
||||
_rename_op_input(program, op_var_rename_map, origin_ops, keep_fp32_ops)
|
||||
return to_bf16_var_names
|
||||
|
||||
|
||||
def cast_parameters_to_bf16(place, program, scope=None, to_bf16_var_names=None):
|
||||
"""
|
||||
Traverse all parameters in the whole model and set them to the BF16 data type.
|
||||
Whereas, this function will keep parameters of batchnorms in FP32.
|
||||
Args:
|
||||
place(base.CPUPlace|base.CUDAPlace): `place` is used to restore the BF16 weight tensors.
|
||||
program (Program): The used program.
|
||||
scope(base.Scope, optional): `scope` is used to get the FP32 weight tensor values.
|
||||
Default is None.
|
||||
to_bf16_var_names(set|list, optional): The data types of vars in `to_bf16_var_names`
|
||||
will be set to BF16. Usually, it is the returned
|
||||
value of `cast_model_to_bf16` API.
|
||||
"""
|
||||
all_parameters = []
|
||||
for block in program.blocks:
|
||||
all_parameters.extend(block.all_parameters())
|
||||
|
||||
bf16_var_names = to_bf16_var_names if to_bf16_var_names else set()
|
||||
var_scope = scope if scope else global_scope()
|
||||
for param in all_parameters:
|
||||
if param.name in bf16_var_names:
|
||||
_logger.debug(f"---- cast {param.name} to bf16 dtype ----")
|
||||
param_t = var_scope.find_var(param.name).get_tensor()
|
||||
data = np.array(param_t)
|
||||
param_t.set(convert_float_to_uint16(data), place)
|
||||
|
||||
|
||||
def rewrite_program_bf16(main_prog, amp_lists=None):
|
||||
"""
|
||||
Traverse all ops in current block and insert cast op according to
|
||||
which set current op belongs to.
|
||||
|
||||
1. When an op belongs to the fp32 list, add it to fp32 set
|
||||
2. When an op belongs to the bf16 list, add it to bf16 set
|
||||
3. When an op belongs to the gray list. If one
|
||||
of its inputs is the output of fp32 set op or fp32 list op,
|
||||
add it to fp32 set. If all of its previous ops are not fp32
|
||||
op and one of its inputs is the output of bf16 set op or
|
||||
bf16 list op, add it to bf16 set.
|
||||
4. When an op isn't in the lists, add it to fp32 op set.
|
||||
5. Add necessary cast ops to make sure that fp32 set op will be
|
||||
computed in fp32 mode, while bf16 set op will be computed in
|
||||
bf16 mode.
|
||||
|
||||
Args:
|
||||
main_prog (Program): The main program for training.
|
||||
"""
|
||||
if amp_lists is None:
|
||||
amp_lists = AutoMixedPrecisionListsBF16()
|
||||
block = main_prog.global_block()
|
||||
ops = block.ops
|
||||
bf16_op_set = set()
|
||||
fp32_op_set = set()
|
||||
for op in ops:
|
||||
# NOTE(zhiqiu): 'create_py_reader' and 'read' is used in non-iterable DataLoader,
|
||||
# we don't need to handle reader op and the input of 'create_py_reader' is not
|
||||
# in block, which may result in errors.
|
||||
# See GeneratorLoader._init_non_iterable() for details.
|
||||
if op.type == 'create_py_reader' or op.type == 'read':
|
||||
continue
|
||||
|
||||
if amp_lists.fp32_varnames is not None and _is_in_fp32_varnames(
|
||||
op, amp_lists
|
||||
):
|
||||
fp32_op_set.add(op)
|
||||
continue
|
||||
|
||||
if op.type in amp_lists.fp32_list:
|
||||
fp32_op_set.add(op)
|
||||
elif op.type in amp_lists.bf16_list:
|
||||
bf16_op_set.add(op)
|
||||
elif op.type in amp_lists.gray_list:
|
||||
is_fp32_op = False
|
||||
is_bf16_op = False
|
||||
for in_name in op.input_names:
|
||||
# if this op has inputs
|
||||
if in_name:
|
||||
for in_var_name in op.input(in_name):
|
||||
in_var = block.var(in_var_name)
|
||||
# this in_var isn't the output of other op
|
||||
if in_var.op is None:
|
||||
continue
|
||||
elif in_var.op is op:
|
||||
prev_op = find_true_prev_op(ops, op, in_var_name)
|
||||
if prev_op is None:
|
||||
continue
|
||||
else:
|
||||
prev_op = in_var.op
|
||||
# if it's one of inputs
|
||||
if (
|
||||
prev_op in fp32_op_set
|
||||
or prev_op.type in amp_lists.fp32_list
|
||||
):
|
||||
is_fp32_op = True
|
||||
elif (
|
||||
prev_op in bf16_op_set
|
||||
or prev_op.type in amp_lists.bf16_list
|
||||
):
|
||||
is_bf16_op = True
|
||||
if is_fp32_op:
|
||||
fp32_op_set.add(op)
|
||||
elif is_bf16_op:
|
||||
bf16_op_set.add(op)
|
||||
else:
|
||||
pass
|
||||
else:
|
||||
# For numerical safe, we apply fp32 computation on ops that
|
||||
# are not determined which list they should stay.
|
||||
fp32_op_set.add(op)
|
||||
|
||||
idx = 0
|
||||
while idx < len(ops):
|
||||
op = ops[idx]
|
||||
num_cast_ops = 0
|
||||
if op in fp32_op_set:
|
||||
num_cast_ops = _insert_cast_op(
|
||||
block,
|
||||
op,
|
||||
idx,
|
||||
core.VarDesc.VarType.BF16,
|
||||
core.VarDesc.VarType.FP32,
|
||||
)
|
||||
elif op in bf16_op_set:
|
||||
if (
|
||||
op.has_attr('dtype')
|
||||
and op.attr('dtype') == core.VarDesc.VarType.FP32
|
||||
):
|
||||
op._set_attr('dtype', core.VarDesc.VarType.BF16)
|
||||
|
||||
num_cast_ops = _insert_cast_op(
|
||||
block,
|
||||
op,
|
||||
idx,
|
||||
core.VarDesc.VarType.FP32,
|
||||
core.VarDesc.VarType.BF16,
|
||||
)
|
||||
else:
|
||||
pass
|
||||
|
||||
idx += num_cast_ops + 1
|
||||
Reference in New Issue
Block a user