chore: import upstream snapshot with attribution
This commit is contained in:
wehub-resource-sync
@@ -0,0 +1,28 @@
|
||||
# Copyright (c) 2020 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.
|
||||
|
||||
from . import ( # noqa: F401
|
||||
bf16,
|
||||
debugging,
|
||||
decorator,
|
||||
fp16_lists,
|
||||
fp16_utils,
|
||||
)
|
||||
from .decorator import decorate # noqa: F401
|
||||
from .fp16_lists import AutoMixedPrecisionLists, CustomOpLists # noqa: F401
|
||||
from .fp16_utils import ( # noqa: F401
|
||||
cast_model_to_fp16,
|
||||
cast_parameters_to_fp16,
|
||||
fp16_guard,
|
||||
)
|
||||
@@ -0,0 +1,176 @@
|
||||
# Copyright (c) 2020 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 paddle
|
||||
from paddle import _C_ops
|
||||
from paddle.base.data_feeder import check_type, check_variable_and_dtype
|
||||
from paddle.base.framework import (
|
||||
Variable,
|
||||
in_dynamic_or_pir_mode,
|
||||
)
|
||||
from paddle.base.layer_helper import LayerHelper
|
||||
|
||||
|
||||
def check_finite_and_unscale(x, scale, name=None, float_status=None):
|
||||
"""
|
||||
Check if input X contains all finite data, if yes, scale it by input Scale.
|
||||
|
||||
$$Out = X / scale$$
|
||||
|
||||
If any tensor in X contains Inf or Nan, the Out will generate a indicator.
|
||||
FoundInfinite will be 1 (True), and Out will not be scaled. In this case, the data of
|
||||
Out should not be used, and its data may not be deterministic.
|
||||
Otherwise, FoundInfinite will be 0 (False).
|
||||
|
||||
Args:
|
||||
x(list|tuple): The input tensors of check_finite_and_unscale operator.
|
||||
scale: The scale of check_finite_and_unscale operator.
|
||||
float_status(Tensor): (Only used on NPU) The float status to check overflow.
|
||||
"""
|
||||
|
||||
if in_dynamic_or_pir_mode():
|
||||
x, found_inf = _C_ops.check_finite_and_unscale_(x, scale)
|
||||
return x, found_inf
|
||||
|
||||
helper = LayerHelper("check_finite_and_unscale", **locals())
|
||||
found_inf = helper.create_variable_for_type_inference(dtype='bool')
|
||||
check_type(x, 'x', (tuple, list), 'check_finite_and_unscale')
|
||||
for e in x:
|
||||
check_variable_and_dtype(
|
||||
e,
|
||||
"x",
|
||||
['float16', 'float32', 'float64', 'uint16'],
|
||||
'check_finite_and_unscale',
|
||||
)
|
||||
|
||||
inputs = {'X': x, 'Scale': scale}
|
||||
outputs = {'Out': x, 'FoundInfinite': found_inf}
|
||||
helper.append_op(
|
||||
type='check_finite_and_unscale', inputs=inputs, outputs=outputs
|
||||
)
|
||||
|
||||
return x, found_inf
|
||||
|
||||
|
||||
def update_loss_scaling(
|
||||
x,
|
||||
found_inf,
|
||||
prev_loss_scaling,
|
||||
num_good_steps,
|
||||
num_bad_steps,
|
||||
incr_every_n_steps,
|
||||
decr_every_n_nan_or_inf,
|
||||
incr_ratio,
|
||||
decr_ratio,
|
||||
stop_update=False,
|
||||
name=None,
|
||||
):
|
||||
"""
|
||||
Update loss scaling according to overall gradients. If all gradients is
|
||||
finite after incr_every_n_steps, loss scaling will increase by incr_ratio.
|
||||
Otherwise, loss scaling will decrease by decr_ratio after
|
||||
decr_every_n_nan_or_inf steps and each step some gradients are infinite.
|
||||
|
||||
Args:
|
||||
x(list|tuple): The input tensors of update_loss_scaling operator.
|
||||
found_inf (Variable): A boolean variable indicates whether
|
||||
there is any infinite gradient.
|
||||
prev_loss_scaling (Variable): Previous loss scaling.
|
||||
num_good_steps (Variable): A variable accumulates good steps in which
|
||||
all gradients are finite.
|
||||
num_bad_steps (Variable): A variable accumulates bad steps in which
|
||||
some gradients are infinite.
|
||||
incr_every_n_steps (int): A variable represents increasing loss
|
||||
scaling every n consecutive steps with
|
||||
finite gradients.
|
||||
decr_every_n_nan_or_inf (int): A variable represents decreasing
|
||||
loss scaling every n accumulated
|
||||
steps with nan or inf gradients.
|
||||
incr_ratio(float): The multiplier to use when increasing the loss
|
||||
scaling.
|
||||
decr_ratio(float): The less-than-one-multiplier to use when decreasing
|
||||
loss scaling.
|
||||
"""
|
||||
if in_dynamic_or_pir_mode():
|
||||
_C_ops.update_loss_scaling_(
|
||||
x,
|
||||
found_inf,
|
||||
prev_loss_scaling,
|
||||
num_good_steps,
|
||||
num_bad_steps,
|
||||
incr_every_n_steps,
|
||||
decr_every_n_nan_or_inf,
|
||||
incr_ratio,
|
||||
decr_ratio,
|
||||
stop_update,
|
||||
)
|
||||
return x
|
||||
|
||||
check_variable_and_dtype(
|
||||
prev_loss_scaling,
|
||||
"prev_loss_scaling",
|
||||
['float32', 'float64'],
|
||||
"update_loss_scaling",
|
||||
)
|
||||
check_type(x, 'x', (tuple, list), 'update_loss_scaling')
|
||||
for e in x:
|
||||
check_variable_and_dtype(
|
||||
e,
|
||||
"x",
|
||||
['float16', 'float32', 'float64', 'uint16'],
|
||||
'update_loss_scaling',
|
||||
)
|
||||
if e.dtype in [paddle.float16, paddle.bfloat16]:
|
||||
assert prev_loss_scaling.dtype == paddle.float32, (
|
||||
"The dtype of prev_loss_scaling should be float32 when the dtype of x is float16 or bfloat16."
|
||||
)
|
||||
else:
|
||||
assert prev_loss_scaling.dtype == e.dtype, (
|
||||
"The dtype of prev_loss_scaling should be equal to the dtype of x."
|
||||
)
|
||||
|
||||
helper = LayerHelper("update_loss_scaling", **locals())
|
||||
|
||||
inputs = {
|
||||
'X': x,
|
||||
'FoundInfinite': found_inf,
|
||||
'PrevLossScaling': prev_loss_scaling,
|
||||
'InGoodSteps': num_good_steps,
|
||||
'InBadSteps': num_bad_steps,
|
||||
}
|
||||
|
||||
outputs = {
|
||||
'Out': x,
|
||||
'LossScaling': prev_loss_scaling,
|
||||
'OutGoodSteps': num_good_steps,
|
||||
'OutBadSteps': num_bad_steps,
|
||||
}
|
||||
|
||||
attrs = {
|
||||
'incr_every_n_steps': incr_every_n_steps,
|
||||
'decr_every_n_nan_or_inf': decr_every_n_nan_or_inf,
|
||||
'incr_ratio': incr_ratio,
|
||||
'decr_ratio': decr_ratio,
|
||||
}
|
||||
|
||||
if isinstance(stop_update, Variable):
|
||||
inputs['StopUpdate'] = stop_update
|
||||
else:
|
||||
attrs['stop_update'] = stop_update
|
||||
|
||||
helper.append_op(
|
||||
type='update_loss_scaling', inputs=inputs, outputs=outputs, attrs=attrs
|
||||
)
|
||||
|
||||
return x
|
||||
@@ -0,0 +1,28 @@
|
||||
# Copyright (c) 2021 PaddlePaddle Authors. All Rights Reserve.
|
||||
#
|
||||
# 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.
|
||||
|
||||
from . import ( # noqa: F401
|
||||
amp_lists,
|
||||
amp_utils,
|
||||
decorator,
|
||||
)
|
||||
from .amp_lists import AutoMixedPrecisionListsBF16 # noqa: F401
|
||||
from .amp_utils import ( # noqa: F401
|
||||
bf16_guard,
|
||||
cast_model_to_bf16,
|
||||
cast_parameters_to_bf16,
|
||||
convert_float_to_uint16,
|
||||
rewrite_program_bf16,
|
||||
)
|
||||
from .decorator import decorate_bf16 # noqa: F401
|
||||
@@ -0,0 +1,110 @@
|
||||
# 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 copy
|
||||
|
||||
from paddle.amp.amp_lists import BF16_WHITE_LIST
|
||||
from paddle.base import core
|
||||
|
||||
from ..fp16_lists import (
|
||||
black_list as black_list_fp16,
|
||||
gray_list as gray_list_fp16,
|
||||
white_list as white_list_fp16,
|
||||
)
|
||||
|
||||
|
||||
class AutoMixedPrecisionListsBF16:
|
||||
"""
|
||||
AutoMixedPrecisionListsBF16 is a class for fp32/bf16 op types list. The lists are used for an
|
||||
algorithm which determines op's execution mode (fp32 or bf16).It can update pre-defined
|
||||
fp32 list and bf16 list according to users' custom fp32 bf16 lists.
|
||||
|
||||
Args:
|
||||
custom_bf16_list (set): Users' custom bf16 list.
|
||||
custom_fp32_list (set): Users' custom fp32 list.
|
||||
custom_fp32_varnames (set): Users' custom fp32 variables' names.
|
||||
|
||||
Examples:
|
||||
.. code-block:: pycon
|
||||
|
||||
>>> import paddle
|
||||
>>> paddle.enable_static()
|
||||
>>> with paddle.static.amp.bf16.bf16_guard():
|
||||
... paddle.static.amp.bf16.AutoMixedPrecisionListsBF16(custom_fp32_list={'lstm'})
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
custom_bf16_list=None,
|
||||
custom_fp32_list=None,
|
||||
custom_fp32_varnames=None,
|
||||
):
|
||||
self._custom_bf16_list = custom_bf16_list
|
||||
self._custom_fp32_list = custom_fp32_list
|
||||
self.bf16_list = copy.copy(bf16_list)
|
||||
self.fp32_list = copy.copy(fp32_list)
|
||||
self.gray_list = copy.copy(gray_list)
|
||||
self.bf16_initializer_list = copy.copy(bf16_initializer_list)
|
||||
self.unsupported_list = copy.copy(unsupported_list)
|
||||
self.fp32_varnames = copy.copy(custom_fp32_varnames)
|
||||
self._update_list()
|
||||
|
||||
def _update_list(self):
|
||||
"""
|
||||
Update fp32 and bf16 list according to users' custom list.
|
||||
"""
|
||||
if self._custom_bf16_list and self._custom_fp32_list:
|
||||
for op_name in self._custom_bf16_list:
|
||||
if op_name in self._custom_fp32_list:
|
||||
raise ValueError(
|
||||
"Custom bf16 list overlap custom fp32 list"
|
||||
)
|
||||
if self._custom_bf16_list:
|
||||
for op_name in self._custom_bf16_list:
|
||||
if op_name in self.fp32_list:
|
||||
self.fp32_list.remove(op_name)
|
||||
elif op_name in self.gray_list:
|
||||
self.gray_list.remove(op_name)
|
||||
self.bf16_list.add(op_name)
|
||||
if self._custom_fp32_list:
|
||||
for op_name in self._custom_fp32_list:
|
||||
if op_name in self.bf16_list:
|
||||
self.bf16_list.remove(op_name)
|
||||
elif op_name in self.gray_list:
|
||||
self.gray_list.remove(op_name)
|
||||
self.fp32_list.add(op_name)
|
||||
self.unsupported_list.add(op_name)
|
||||
|
||||
|
||||
bf16_initializer_list = {'fill_constant', 'uniform_random'}
|
||||
|
||||
# always bf16
|
||||
bf16_list = BF16_WHITE_LIST
|
||||
|
||||
# depends on the prev_op type
|
||||
gray_list = gray_list_fp16
|
||||
|
||||
_, _, _sys_unsupported_bf16_list = core.op_supported_infos(
|
||||
'CPU', core.VarDesc.VarType.BF16
|
||||
)
|
||||
unsupported_list = _sys_unsupported_bf16_list
|
||||
|
||||
fp32_list = black_list_fp16.copy().copy()
|
||||
fp32_list |= white_list_fp16
|
||||
fp32_list |= gray_list_fp16
|
||||
|
||||
fp32_list -= bf16_list
|
||||
fp32_list -= gray_list
|
||||
unsupported_list -= bf16_list
|
||||
unsupported_list -= gray_list
|
||||
@@ -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
|
||||
@@ -0,0 +1,343 @@
|
||||
# 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 types
|
||||
import warnings
|
||||
|
||||
import paddle
|
||||
from paddle.base import core, default_main_program, program_guard, unique_name
|
||||
|
||||
from .amp_lists import AutoMixedPrecisionListsBF16
|
||||
from .amp_utils import (
|
||||
cast_model_to_bf16,
|
||||
cast_parameters_to_bf16,
|
||||
rewrite_program_bf16,
|
||||
)
|
||||
|
||||
|
||||
class OptimizerWithMixedPrecision:
|
||||
"""
|
||||
Optimizer with mixed-precision (MP) training. This is a wrapper of a common
|
||||
optimizer, plus the support of mixed-precision pre-training. The object
|
||||
of this class almost has the same behavior as the common optimizer, with the
|
||||
methods `minimize()`, `backward()`, `apply_gradients()` implemented.
|
||||
Additionally, it enables the MP training automatically, i.e, the creation
|
||||
and maintenance of master parameters, scaling of loss, etc.
|
||||
|
||||
Args:
|
||||
optimizer (Optimizer): A common Optimizer object.
|
||||
amp_lists (CustomOpLists): An CustomOpLists object.
|
||||
use_pure_bf16(bool): Whether to use the pure bf16 training.
|
||||
use_bf16_guard(bool): Whether to use `bf16_guard` when constructing the program.
|
||||
|
||||
"""
|
||||
|
||||
def __init__(self, optimizer, amp_lists, use_pure_bf16, use_bf16_guard):
|
||||
self._optimizer = optimizer
|
||||
self._amp_lists = amp_lists
|
||||
self._param_grads = None
|
||||
self._train_program = None
|
||||
|
||||
self._learning_rate = optimizer._learning_rate
|
||||
self._learning_rate_map = optimizer._learning_rate_map
|
||||
self._use_pure_bf16 = use_pure_bf16
|
||||
self._use_bf16_guard = use_bf16_guard
|
||||
self._to_bf16_var_names = None
|
||||
|
||||
def _init_amp_var(self):
|
||||
# Ensure the data type of learning rate vars is float32 (same as the
|
||||
# master parameter dtype)
|
||||
if isinstance(self._optimizer._learning_rate, float):
|
||||
self._optimizer._learning_rate_map[default_main_program()] = (
|
||||
paddle.static.create_global_var(
|
||||
name=unique_name.generate("learning_rate"),
|
||||
shape=[1],
|
||||
value=float(self._optimizer._learning_rate),
|
||||
dtype='float32',
|
||||
persistable=True,
|
||||
)
|
||||
)
|
||||
|
||||
def backward(
|
||||
self,
|
||||
loss,
|
||||
startup_program=None,
|
||||
parameter_list=None,
|
||||
no_grad_set=None,
|
||||
callbacks=None,
|
||||
):
|
||||
"""
|
||||
Backward propagation or auto differentiation for gradients' computation.
|
||||
|
||||
Args:
|
||||
loss (Variable): The loss Variable to minimize.
|
||||
startup_program (Program|None): The startup Program for initializing
|
||||
parameters in `parameter_list`.
|
||||
parameter_list (list|None): A list of Variables to update.
|
||||
no_grad_set (set|None): A set of Variables should be ignored.
|
||||
callbacks (list|None): A list of callable objects to run when appending
|
||||
backward operator for one parameter.
|
||||
|
||||
Returns:
|
||||
A list of (param, grad), which is a tuple of a parameter and its
|
||||
gradient respectively, and the scaled loss.
|
||||
"""
|
||||
train_program = loss.block.program
|
||||
self._train_program = train_program
|
||||
|
||||
with program_guard(self._train_program, startup_program):
|
||||
self._init_amp_var()
|
||||
|
||||
if self._use_pure_bf16:
|
||||
self._to_bf16_var_names = cast_model_to_bf16(
|
||||
self._train_program,
|
||||
startup_program,
|
||||
self._amp_lists,
|
||||
self._use_bf16_guard,
|
||||
)
|
||||
else:
|
||||
rewrite_program_bf16(self._train_program, self._amp_lists)
|
||||
|
||||
if loss.dtype != core.VarDesc.VarType.FP32:
|
||||
loss = loss.astype('float32')
|
||||
|
||||
params_grads = self._optimizer.backward(
|
||||
loss, startup_program, parameter_list, no_grad_set, callbacks
|
||||
)
|
||||
return params_grads
|
||||
|
||||
def amp_init(
|
||||
self, place, scope=None, test_program=None, use_bf16_test=False
|
||||
):
|
||||
"""
|
||||
Init the amp training, such as cast fp32 parameters to bf16 type.
|
||||
|
||||
Args:
|
||||
place(CPUPlace): place is used to initialize
|
||||
bf16 parameters with fp32 values.
|
||||
scope(Scope): The scope is used to find fp32 parameters.
|
||||
test_program(Program): The program is used for testing.
|
||||
use_bf16_test(bool): Whether to use bf16 testing.
|
||||
|
||||
Examples:
|
||||
.. code-block:: pycon
|
||||
|
||||
>>> # doctest: +SKIP("paddle.static.amp module doesn't support PIR mode")
|
||||
>>> import numpy as np
|
||||
>>> import paddle
|
||||
>>> import paddle.nn.functional as F
|
||||
>>> paddle.enable_static()
|
||||
|
||||
>>> def run_example_code():
|
||||
... place = paddle.CPUPlace()
|
||||
... exe = paddle.static.Executor(place)
|
||||
... 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)
|
||||
... # 1) Use bf16_guard to control the range of bf16 kernels used.
|
||||
... 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)
|
||||
... # 2) Create the optimizer and set `multi_precision` to True.
|
||||
... # Setting `multi_precision` to True can avoid the poor accuracy
|
||||
... # or the slow convergence in a way.
|
||||
... optimizer = paddle.optimizer.Momentum(learning_rate=0.01, multi_precision=True)
|
||||
... # 3) These ops in `custom_black_list` will keep in the float32 computation type.
|
||||
... amp_list = paddle.static.amp.CustomOpLists(custom_black_list=['pool2d'])
|
||||
... # 4) The entry of Paddle AMP.
|
||||
... # Enable pure bf16 training by setting `use_pure_bf16` to True.
|
||||
... optimizer = paddle.static.amp.bf16.decorate_bf16(
|
||||
... optimizer,
|
||||
... amp_list,
|
||||
... use_pure_bf16=True,
|
||||
... )
|
||||
... # If you don't use the default_startup_program(), you should pass
|
||||
... # your defined `startup_program` into `minimize`.
|
||||
... optimizer.minimize(loss)
|
||||
... exe.run(paddle.static.default_startup_program())
|
||||
... # 5) Use `amp_init` after FP32 parameters initialization(such as `exe.run(startup_program)`).
|
||||
... # If you want to perform the testing process, you should pass `test_program` into `amp_init`.
|
||||
... optimizer.amp_init(place, scope=paddle.static.global_scope())
|
||||
|
||||
>>> run_example_code()
|
||||
|
||||
"""
|
||||
assert self._train_program is not None, (
|
||||
"Please call the minimize method first."
|
||||
)
|
||||
if self._use_pure_bf16:
|
||||
cast_parameters_to_bf16(
|
||||
place, self._train_program, scope, self._to_bf16_var_names
|
||||
)
|
||||
if test_program is not None:
|
||||
if self._use_pure_bf16:
|
||||
cast_model_to_bf16(
|
||||
test_program,
|
||||
amp_lists=self._amp_lists,
|
||||
use_bf16_guard=self._use_bf16_guard,
|
||||
)
|
||||
elif use_bf16_test:
|
||||
rewrite_program_bf16(test_program, amp_lists=self._amp_lists)
|
||||
|
||||
def apply_gradients(self, params_grads):
|
||||
"""
|
||||
Apply gradients.
|
||||
|
||||
Args:
|
||||
params_grads (list): A list of params.
|
||||
|
||||
Returns:
|
||||
A list of optimize operators.
|
||||
"""
|
||||
|
||||
return self._optimizer.apply_gradients(params_grads)
|
||||
|
||||
def apply_optimize(self, loss, startup_program, params_grads):
|
||||
program = loss.block.program
|
||||
with program_guard(program, startup_program):
|
||||
optimize_ops = self.apply_gradients(params_grads)
|
||||
return optimize_ops
|
||||
|
||||
def minimize(
|
||||
self, loss, startup_program=None, parameter_list=None, no_grad_set=None
|
||||
):
|
||||
"""
|
||||
Perform optimization by minimizing the given loss.
|
||||
|
||||
Args:
|
||||
loss (Variable): The loss Variable.
|
||||
startup_program (Program): startup_program for initializing parameters
|
||||
in `parameter_list`.
|
||||
parameter_list (list): list of Variables to update.
|
||||
no_grad_set (set|None): set of Variables should be ignored.
|
||||
|
||||
Returns:
|
||||
The scaled loss by scaling factor, the list of optimize ops, and a
|
||||
list of scaled parameters and gradients.
|
||||
"""
|
||||
opt_dict = self._optimizer.__class__.__dict__
|
||||
if 'minimize' in opt_dict and isinstance(
|
||||
opt_dict['minimize'], types.FunctionType
|
||||
):
|
||||
warnings.warn(
|
||||
"The decorated optimizer has its own `minimize` method, but it will not be executed."
|
||||
)
|
||||
|
||||
params_grads = self.backward(
|
||||
loss,
|
||||
startup_program=startup_program,
|
||||
parameter_list=parameter_list,
|
||||
no_grad_set=no_grad_set,
|
||||
)
|
||||
|
||||
optimize_ops = self.apply_optimize(loss, startup_program, params_grads)
|
||||
|
||||
return optimize_ops, params_grads
|
||||
|
||||
|
||||
def decorate_bf16(
|
||||
optimizer, amp_lists=None, use_pure_bf16=False, use_bf16_guard=None
|
||||
):
|
||||
"""
|
||||
Decorate the given optimizer to adapt to the mixed-precision training.
|
||||
|
||||
Args:
|
||||
optimizer(Optimizer): A common Optimizer.
|
||||
amp_lists (CustomOpLists): An CustomOpLists object.
|
||||
use_pure_bf16(bool): Whether to use the pure bf16 training. Default False.
|
||||
use_bf16_guard(bool): Whether to use `bf16_guard` when constructing the program.
|
||||
Default None, which means that its value equals to `use_pure_bf16`.
|
||||
|
||||
Returns:
|
||||
An optimizer acting like a normal one but with mixed-precision training
|
||||
enabled.
|
||||
|
||||
Examples:
|
||||
.. code-block:: pycon
|
||||
:name: example-1
|
||||
|
||||
# fp32&bf16 list based strategy example
|
||||
>>> # doctest: +SKIP("paddle.static.amp module doesn't support PIR mode")
|
||||
>>> import paddle
|
||||
>>> import paddle.static as static
|
||||
|
||||
>>> paddle.enable_static()
|
||||
|
||||
>>> data = static.data(name='X', shape=[None, 1], dtype='float32')
|
||||
>>> hidden = static.nn.fc(x=data, size=10)
|
||||
>>> loss = paddle.mean(hidden)
|
||||
>>> optimizer = paddle.optimizer.Adam(learning_rate=0.001)
|
||||
|
||||
>>> mp_optimizer = static.amp.bf16.decorate_bf16(optimizer=optimizer)
|
||||
|
||||
>>> ops, param_grads = mp_optimizer.minimize(loss)
|
||||
|
||||
|
||||
|
||||
|
||||
.. code-block:: pycon
|
||||
:name: example-2
|
||||
|
||||
# pure bf16 training example
|
||||
>>> # doctest: +SKIP("paddle.static.amp module doesn't support PIR mode")
|
||||
>>> import numpy as np
|
||||
>>> import paddle
|
||||
>>> import paddle.nn.functional as F
|
||||
|
||||
>>> def run_example_code():
|
||||
... place = paddle.CPUPlace()
|
||||
... exe = paddle.static.Executor(place)
|
||||
... 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)
|
||||
... # 1) Use bf16_guard to control the range of bf16 kernels used.
|
||||
... 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)
|
||||
... # 2) Create the optimizer and set `multi_precision` to True.
|
||||
... # Setting `multi_precision` to True can avoid the poor accuracy
|
||||
... # or the slow convergence in a way.
|
||||
... optimizer = paddle.optimizer.Momentum(learning_rate=0.01, multi_precision=True)
|
||||
... # 3) These ops in `custom_black_list` will keep in the float32 computation type.
|
||||
... amp_list = paddle.static.amp.CustomOpLists(custom_black_list=['pool2d'])
|
||||
... # 4) The entry of Paddle AMP.
|
||||
... # Enable pure bf16 training by setting `use_pure_bf16` to True.
|
||||
... optimizer = paddle.static.amp.bf16.decorate_bf16(
|
||||
... optimizer,
|
||||
... amp_list,
|
||||
... use_pure_bf16=True,
|
||||
... )
|
||||
... # If you don't use the default_startup_program(), you should pass
|
||||
... # your defined `startup_program` into `minimize`.
|
||||
... optimizer.minimize(loss)
|
||||
... exe.run(paddle.static.default_startup_program())
|
||||
... # 5) Use `amp_init` after FP32 parameters initialization(such as `exe.run(startup_program)`).
|
||||
... # If you want to perform the testing process, you should pass `test_program` into `amp_init`.
|
||||
... optimizer.amp_init(place, scope=paddle.static.global_scope())
|
||||
>>> run_example_code()
|
||||
|
||||
"""
|
||||
if amp_lists is None:
|
||||
amp_lists = AutoMixedPrecisionListsBF16()
|
||||
|
||||
if use_bf16_guard is None:
|
||||
use_bf16_guard = use_pure_bf16
|
||||
|
||||
mp_optimizer = OptimizerWithMixedPrecision(
|
||||
optimizer, amp_lists, use_pure_bf16, use_bf16_guard
|
||||
)
|
||||
|
||||
return mp_optimizer
|
||||
@@ -0,0 +1,279 @@
|
||||
# Copyright (c) 2023 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 copy
|
||||
import logging
|
||||
|
||||
import paddle
|
||||
from paddle.base.log_helper import get_logger
|
||||
|
||||
_logger = get_logger(
|
||||
__name__, logging.INFO, fmt='%(asctime)s-%(levelname)s: %(message)s'
|
||||
)
|
||||
|
||||
|
||||
class OperatorStatsUnit:
|
||||
def __init__(self):
|
||||
self.op_type = None
|
||||
self.fp32_calls = 0
|
||||
self.fp16_calls = 0
|
||||
self.bf16_calls = 0
|
||||
self.other_calls = 0
|
||||
|
||||
def update(self, dtype):
|
||||
if dtype is None:
|
||||
self.other_calls = self.other_calls + 1
|
||||
else:
|
||||
if dtype == paddle.float32:
|
||||
self.fp32_calls = self.fp32_calls + 1
|
||||
elif dtype == paddle.float16:
|
||||
self.fp16_calls = self.fp16_calls + 1
|
||||
elif dtype == paddle.bfloat16:
|
||||
self.bf16_calls = self.bf16_calls + 1
|
||||
else:
|
||||
self.other_calls = self.other_calls + 1
|
||||
|
||||
def addto(self, another):
|
||||
self.fp32_calls += another.fp32_calls
|
||||
self.fp16_calls += another.fp16_calls
|
||||
self.bf16_calls += another.bf16_calls
|
||||
self.other_calls += another.other_calls
|
||||
|
||||
def convert_to_list(self):
|
||||
return [
|
||||
self.fp16_calls,
|
||||
self.bf16_calls,
|
||||
self.fp32_calls,
|
||||
self.other_calls,
|
||||
]
|
||||
|
||||
|
||||
def _is_floating_point(dtype):
|
||||
if dtype in [
|
||||
paddle.base.core.VarDesc.VarType.FP64,
|
||||
paddle.base.core.VarDesc.VarType.FP32,
|
||||
paddle.base.core.VarDesc.VarType.FP16,
|
||||
paddle.base.core.VarDesc.VarType.BF16,
|
||||
]:
|
||||
return True
|
||||
else:
|
||||
return False
|
||||
|
||||
|
||||
def _get_var_dtype_from_block(block, op, arg_name, is_input):
|
||||
var_names = op.input(arg_name) if is_input else op.output(arg_name)
|
||||
assert isinstance(var_names, list)
|
||||
if len(var_names) == 0:
|
||||
return None
|
||||
|
||||
var_name = var_names[0]
|
||||
try:
|
||||
var = block._var_recursive(var_name)
|
||||
return var.dtype
|
||||
except:
|
||||
_logger.warning(
|
||||
"Operator < {} > gets {} < {} : {} > error!".format(
|
||||
op.type, "input" if is_input else "output", arg_name, var_name
|
||||
)
|
||||
)
|
||||
return None
|
||||
|
||||
|
||||
def _extract_compute_dtype(op, block):
|
||||
var_name = None
|
||||
compute_dtype = None
|
||||
for in_name in op.input_names:
|
||||
var_dtype = _get_var_dtype_from_block(block, op, in_name, True)
|
||||
if var_dtype is None:
|
||||
continue
|
||||
|
||||
if compute_dtype is None:
|
||||
compute_dtype = var_dtype
|
||||
else:
|
||||
if compute_dtype != var_dtype:
|
||||
if _is_floating_point(compute_dtype) and _is_floating_point(
|
||||
var_dtype
|
||||
):
|
||||
_logger.warning(
|
||||
f"Operator < {op.type} > has different input data types, input_names = {op.input_names}, output_names = {op.output_names}."
|
||||
)
|
||||
elif _is_floating_point(var_dtype):
|
||||
# When there are multiple inputs, such as embedding
|
||||
# (ids is integer, w is floating-point), the kernel
|
||||
# dtype is normally decided by the input of floating-point.
|
||||
compute_dtype = var_dtype
|
||||
|
||||
for out_name in op.output_names:
|
||||
var_dtype = _get_var_dtype_from_block(block, op, out_name, False)
|
||||
if var_dtype is None:
|
||||
continue
|
||||
|
||||
if compute_dtype is None:
|
||||
# Kernel dtype is mostly decided by the input's dtype.
|
||||
# When the operator has no input, it might have an attr
|
||||
# such as dtype to specify the output's dtype.
|
||||
compute_dtype = var_dtype
|
||||
else:
|
||||
if compute_dtype != var_dtype:
|
||||
if _is_floating_point(compute_dtype) and _is_floating_point(
|
||||
var_dtype
|
||||
):
|
||||
_logger.warning(
|
||||
f"Operator < {op.type} > has different input / output data types, input_names = {op.input_names}, output_names = {op.output_names}."
|
||||
)
|
||||
return compute_dtype
|
||||
|
||||
|
||||
def _merge_op_stats(op_stats_list):
|
||||
merged_op_stats_dict = {}
|
||||
for each_op_stats_dict in op_stats_list:
|
||||
for op_type, unit in each_op_stats_dict.items():
|
||||
if merged_op_stats_dict.get(op_type, None) is None:
|
||||
merged_op_stats_dict[op_type] = copy.copy(unit)
|
||||
else:
|
||||
merged_op_stats_dict[op_type].addto(unit)
|
||||
return merged_op_stats_dict
|
||||
|
||||
|
||||
def _get_op_stats_list(program):
|
||||
def _is_special_ops_with_input_x(op_type):
|
||||
# operators have input X and have inputs different dtypes.
|
||||
special_op_list = ['cast', 'batch_norm', 'instance_norm', 'layer_norm']
|
||||
if op_type in special_op_list:
|
||||
return True
|
||||
if op_type.replace("_grad", "") in special_op_list:
|
||||
return True
|
||||
return False
|
||||
|
||||
op_stats_list = []
|
||||
for block in program.blocks:
|
||||
block_op_stats_dict = {}
|
||||
for op in block.ops:
|
||||
if block_op_stats_dict.get(op.type, None) is None:
|
||||
unit = OperatorStatsUnit()
|
||||
block_op_stats_dict[op.type] = unit
|
||||
else:
|
||||
unit = block_op_stats_dict[op.type]
|
||||
|
||||
if op.type in [
|
||||
'create_py_reader',
|
||||
'read',
|
||||
'create_double_buffer_reader',
|
||||
]:
|
||||
compute_dtype = None
|
||||
elif _is_special_ops_with_input_x(op.type):
|
||||
# Not check the input and output dtype difference for this operators.
|
||||
compute_dtype = _get_var_dtype_from_block(block, op, 'X', True)
|
||||
elif "Param" in op.input_names:
|
||||
# Specify compute_dtype for optimizers.
|
||||
compute_dtype = _get_var_dtype_from_block(
|
||||
block, op, 'Param', True
|
||||
)
|
||||
else:
|
||||
compute_dtype = _extract_compute_dtype(op, block)
|
||||
unit.update(dtype=compute_dtype)
|
||||
op_stats_list.append(block_op_stats_dict)
|
||||
return op_stats_list
|
||||
|
||||
|
||||
def collect_operator_stats(program=None, print_subblocks=False):
|
||||
"""
|
||||
Collect the number of operators for different data types through parsing
|
||||
the program. The statistical data are categorized according to four data
|
||||
types, namely float32, float16, bfloat16 and others.
|
||||
|
||||
Args:
|
||||
program(Program, optional): The program to parse. Default None, and the default main_program will be parsed.
|
||||
print_subblocks(bool, optional): Whether to print the operator stats for each subblock. Default False.
|
||||
|
||||
Examples:
|
||||
|
||||
.. code-block:: pycon
|
||||
|
||||
>>> import paddle
|
||||
>>> paddle.enable_static()
|
||||
|
||||
>>> class SimpleConvNet(paddle.nn.Layer):
|
||||
... def __init__(self):
|
||||
... super().__init__()
|
||||
... self.conv = paddle.nn.Conv2D(in_channels=1, out_channels=6, kernel_size=3)
|
||||
... self.linear = paddle.nn.Linear(in_features=26, out_features=10)
|
||||
...
|
||||
... def forward(self, x):
|
||||
... out = self.conv(x)
|
||||
... out = paddle.nn.functional.relu(out)
|
||||
... out = self.linear(out)
|
||||
... out = paddle.nn.functional.softmax(out)
|
||||
... return out
|
||||
|
||||
>>> main_program = paddle.static.Program()
|
||||
>>> startup_program = paddle.static.Program()
|
||||
>>> with paddle.utils.unique_name.guard():
|
||||
... with paddle.static.program_guard(main_program, startup_program):
|
||||
... model = SimpleConvNet()
|
||||
... x = paddle.static.data(name='input', shape=[None, 1, 28, 28], dtype='float32')
|
||||
... out = model(x)
|
||||
... loss = paddle.mean(out)
|
||||
... optimizer = paddle.optimizer.AdamW()
|
||||
... optimizer = paddle.static.amp.decorate(optimizer)
|
||||
... optimizer.minimize(loss)
|
||||
>>> paddle.static.amp.debugging.collect_operator_stats(main_program)
|
||||
<------------------------------------------------ op list of all blocks ------------------------------------------------->
|
||||
<------------------------------------------------------- op list -------------------------------------------------------->
|
||||
<--------------- Op Name ---------------- | -- FP16 Calls --- | -- BF16 Calls --- | --- FP32 Calls--- | -- Other Calls -->
|
||||
adamw | 0 | 0 | 4 | 0
|
||||
cast | 5 | 0 | 6 | 0
|
||||
check_finite_and_unscale | 0 | 0 | 1 | 0
|
||||
conv2d | 1 | 0 | 0 | 0
|
||||
conv2d_grad | 1 | 0 | 0 | 0
|
||||
elementwise_add | 2 | 0 | 0 | 0
|
||||
elementwise_add_grad | 2 | 0 | 0 | 0
|
||||
elementwise_mul | 0 | 0 | 1 | 0
|
||||
elementwise_mul_grad | 0 | 0 | 1 | 0
|
||||
fill_constant | 0 | 0 | 1 | 0
|
||||
matmul_v2 | 1 | 0 | 0 | 0
|
||||
matmul_v2_grad | 1 | 0 | 0 | 0
|
||||
memcpy | 0 | 0 | 0 | 1
|
||||
reduce_mean | 0 | 0 | 1 | 0
|
||||
reduce_mean_grad | 0 | 0 | 1 | 0
|
||||
relu | 1 | 0 | 0 | 0
|
||||
relu_grad | 1 | 0 | 0 | 0
|
||||
reshape2 | 0 | 0 | 1 | 0
|
||||
reshape2_grad | 0 | 0 | 1 | 0
|
||||
softmax | 0 | 0 | 1 | 0
|
||||
softmax_grad | 0 | 0 | 1 | 0
|
||||
update_loss_scaling | 0 | 0 | 1 | 0
|
||||
<----------------------------------------------------- op count: 22 ----------------------------------------------------->
|
||||
"""
|
||||
|
||||
def _convert_to_list(op_stats_unit_dict):
|
||||
for key, value in op_stats_unit_dict.items():
|
||||
op_stats_unit_dict[key] = value.convert_to_list()
|
||||
return op_stats_unit_dict
|
||||
|
||||
if program is None:
|
||||
program = paddle.static.default_main_program()
|
||||
|
||||
op_stats_list = _get_op_stats_list(program)
|
||||
merged_op_stats = _merge_op_stats(op_stats_list)
|
||||
if print_subblocks and len(op_stats_list) > 1:
|
||||
for i in range(len(op_stats_list)):
|
||||
print("<{:-^120}>".format(" op list of block " + str(i) + " "))
|
||||
paddle.amp.debugging._print_operator_stats(
|
||||
_convert_to_list(op_stats_list[i])
|
||||
)
|
||||
print("<{:-^120}>".format(" op list of all blocks "))
|
||||
paddle.amp.debugging._print_operator_stats(
|
||||
_convert_to_list(merged_op_stats)
|
||||
)
|
||||
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,273 @@
|
||||
# Copyright (c) 2019 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 copy
|
||||
import logging
|
||||
|
||||
import paddle
|
||||
from paddle.amp.amp_lists import (
|
||||
EXTRA_BLACK_LIST,
|
||||
FP16_BLACK_LIST,
|
||||
FP16_WHITE_LIST,
|
||||
)
|
||||
from paddle.base import core
|
||||
from paddle.base.log_helper import get_logger
|
||||
|
||||
_logger = get_logger(
|
||||
__name__, logging.INFO, fmt='%(asctime)s-%(levelname)s: %(message)s'
|
||||
)
|
||||
|
||||
black_list = FP16_BLACK_LIST
|
||||
_extra_black_list = EXTRA_BLACK_LIST
|
||||
white_list = FP16_WHITE_LIST
|
||||
|
||||
|
||||
def check_amp_dtype(dtype):
|
||||
"""
|
||||
Check amp_dtype: float16 or bfloat16
|
||||
"""
|
||||
if isinstance(dtype, str):
|
||||
dtype = dtype.lower()
|
||||
if dtype not in ['float16', 'bfloat16']:
|
||||
raise ValueError(
|
||||
"If enable AMP, dtype should be 'float16' or 'bfloat16'."
|
||||
)
|
||||
return dtype
|
||||
|
||||
|
||||
def get_low_precision_vartype(dtype):
|
||||
if isinstance(dtype, (core.VarDesc.VarType, core.DataType)):
|
||||
return dtype
|
||||
elif isinstance(dtype, str):
|
||||
dtype = dtype.lower()
|
||||
if dtype == "float16":
|
||||
var_type = core.VarDesc.VarType.FP16
|
||||
elif dtype == "bfloat16":
|
||||
var_type = core.VarDesc.VarType.BF16
|
||||
else:
|
||||
raise ValueError(
|
||||
"If enable AMP, dtype should be 'float16' or 'bfloat16'."
|
||||
)
|
||||
return var_type
|
||||
else:
|
||||
raise TypeError(
|
||||
f"The type of dtype is expected to be string or core.VarDesc.VarType, but received {type(dtype)}."
|
||||
)
|
||||
|
||||
|
||||
def get_low_precision_dtypestr(dtype):
|
||||
if isinstance(dtype, str):
|
||||
return check_amp_dtype(dtype)
|
||||
elif isinstance(dtype, (core.VarDesc.VarType, core.DataType)):
|
||||
if dtype == paddle.float16:
|
||||
return "float16"
|
||||
elif dtype == paddle.bfloat16:
|
||||
return "bfloat16"
|
||||
else:
|
||||
raise ValueError(
|
||||
"If enable AMP, dtype should be core.VarDesc.VarType.FP16 or core.VarDesc.VarType.BF16."
|
||||
)
|
||||
else:
|
||||
raise TypeError(
|
||||
f"The type of dtype is expected to be string or core.VarDesc.VarType, but received {type(dtype)}."
|
||||
)
|
||||
|
||||
|
||||
def _get_sys_unsupported_list(dtype):
|
||||
var_type = get_low_precision_vartype(dtype)
|
||||
|
||||
# The set of ops that don't support fp16 calculation
|
||||
device = None
|
||||
if core.is_compiled_with_xpu():
|
||||
device = 'XPU'
|
||||
elif isinstance(
|
||||
paddle.framework._current_expected_place(), paddle.CustomPlace
|
||||
):
|
||||
device = paddle.framework._current_expected_place().get_device_type()
|
||||
else:
|
||||
device = 'GPU'
|
||||
all_ops, _, sys_unsupported_list = core.op_supported_infos(device, var_type)
|
||||
|
||||
# sys_unsupported_list will include the following ops.
|
||||
supported_fp16_list = {
|
||||
"conditional_block",
|
||||
"conditional_block_infer",
|
||||
"select_input",
|
||||
"while",
|
||||
"cast",
|
||||
"tensor_array_to_tensor",
|
||||
"lod_array_length",
|
||||
"write_to_array",
|
||||
}
|
||||
sys_unsupported_list -= supported_fp16_list
|
||||
|
||||
return device, sys_unsupported_list, all_ops
|
||||
|
||||
|
||||
def _get_unsupported_list(dtype):
|
||||
# The set of ops that don't support fp16 calculation
|
||||
_, _sys_unsupported_list, _sys_all_list = _get_sys_unsupported_list(dtype)
|
||||
return _sys_unsupported_list, _sys_all_list
|
||||
|
||||
|
||||
# The three sets listed below are changed dynamically. They don't contain all
|
||||
# paddle ops currently.
|
||||
|
||||
# The set of ops that support fp16 calculation and are considered numerically-
|
||||
# safe and performance-critical. These ops are always converted to fp16.
|
||||
|
||||
_only_supported_fp16_list = {'resnet_unit', 'fused_bn_add_activation'}
|
||||
|
||||
|
||||
def _get_white_list(dtype):
|
||||
white_list_for_dtype = copy.copy(FP16_WHITE_LIST)
|
||||
if dtype == 'float16':
|
||||
white_list_for_dtype = white_list_for_dtype | _only_supported_fp16_list
|
||||
return white_list_for_dtype
|
||||
|
||||
|
||||
def _get_black_list():
|
||||
_black_list = copy.copy(FP16_BLACK_LIST)
|
||||
_black_list = _black_list | EXTRA_BLACK_LIST
|
||||
return _black_list
|
||||
|
||||
|
||||
class AutoMixedPrecisionLists:
|
||||
"""
|
||||
AutoMixedPrecisionLists is a class for black/white list. It can update
|
||||
pre-defined black list and white list according to users' custom black
|
||||
white lists. The lists are used for an algorithm which determines op's
|
||||
execution mode (fp32, fp16 or bf16).
|
||||
|
||||
Args:
|
||||
custom_white_list (set): Users' custom white list.
|
||||
custom_black_list (set): Users' custom black list.
|
||||
custom_black_varnames (set): Users' custom black variables' names.
|
||||
dtype (str): the low precision dtype, which can be set to 'float16' or 'bfloat16'.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
custom_white_list=None,
|
||||
custom_black_list=None,
|
||||
custom_black_varnames=None,
|
||||
dtype="float16",
|
||||
):
|
||||
self.amp_dtype = check_amp_dtype(dtype)
|
||||
self._custom_white_list = custom_white_list
|
||||
self._custom_black_list = custom_black_list
|
||||
self.white_list = copy.copy(_get_white_list(self.amp_dtype))
|
||||
self.black_list = copy.copy(_get_black_list())
|
||||
self.gray_list = copy.copy(gray_list)
|
||||
unsupported_list, sys_all_list = _get_unsupported_list(self.amp_dtype)
|
||||
self.unsupported_list = copy.copy(unsupported_list)
|
||||
self.all_list = copy.copy(sys_all_list)
|
||||
self.black_varnames = copy.copy(custom_black_varnames)
|
||||
self._update_list()
|
||||
|
||||
def _update_list(self):
|
||||
"""
|
||||
Update black and white list according to users' custom list.
|
||||
"""
|
||||
_logger.debug(f"---- custom_white_list {self._custom_white_list} ---- ")
|
||||
_logger.debug(f"---- custom_black_list {self._custom_black_list} ---- ")
|
||||
_logger.debug(f"---- custom_black_varnames {self.black_varnames} ---- ")
|
||||
if self._custom_white_list and self._custom_black_list:
|
||||
for op_name in self._custom_white_list:
|
||||
if op_name in self._custom_black_list:
|
||||
raise ValueError(
|
||||
f"The given custom_white_list overlaps custom_black_list with < {op_name} >!"
|
||||
)
|
||||
if self._custom_white_list:
|
||||
for op_name in self._custom_white_list:
|
||||
if op_name in self.black_list:
|
||||
self.black_list.remove(op_name)
|
||||
elif op_name in self.gray_list:
|
||||
self.gray_list.remove(op_name)
|
||||
self.white_list.add(op_name)
|
||||
if self._custom_black_list:
|
||||
for op_name in self._custom_black_list:
|
||||
if op_name in self.white_list:
|
||||
self.white_list.remove(op_name)
|
||||
elif op_name in self.gray_list:
|
||||
self.gray_list.remove(op_name)
|
||||
self.black_list.add(op_name)
|
||||
self.unsupported_list.add(op_name)
|
||||
device, sys_unsupported_list, _ = _get_sys_unsupported_list(
|
||||
self.amp_dtype
|
||||
)
|
||||
actual_unsupported_list = []
|
||||
for op_name in sys_unsupported_list:
|
||||
if op_name in self.white_list:
|
||||
actual_unsupported_list.append(op_name)
|
||||
if len(actual_unsupported_list) > 0:
|
||||
_logger.warning(
|
||||
f"On current {device}, {self.amp_dtype} is not supported for operators < {actual_unsupported_list} > in white_list!"
|
||||
)
|
||||
|
||||
|
||||
# This set contains two types of ops. All ops supported fp16 calculation. One
|
||||
# of two types is considered numerically-safe, but may be made unsafe by an
|
||||
# upstream blacklist op. Another type do not have numerically-significant
|
||||
# effects, like stack, flatten2.
|
||||
gray_list = {
|
||||
'elementwise_add',
|
||||
'elementwise_sub',
|
||||
'elementwise_mul',
|
||||
'elementwise_div',
|
||||
'elementwise_max',
|
||||
'elementwise_min',
|
||||
'elementwise_pow',
|
||||
'elementwise_mod',
|
||||
'elementwise_floordiv',
|
||||
'batch_norm',
|
||||
'layer_norm',
|
||||
'tanh',
|
||||
'sigmoid',
|
||||
'top_k',
|
||||
'pool2d',
|
||||
'pool3d',
|
||||
'dropout',
|
||||
'relu',
|
||||
'relu6',
|
||||
'leaky_relu',
|
||||
'soft_relu',
|
||||
'flatten2',
|
||||
'stack',
|
||||
'unstack',
|
||||
'uniform_random',
|
||||
'uniform_random_batch_size_like',
|
||||
'gaussian_random',
|
||||
'slice',
|
||||
'rank',
|
||||
'scale',
|
||||
'transpose2',
|
||||
'reshape2',
|
||||
'gather',
|
||||
'fill_constant',
|
||||
'get_tensor_from_selected_rows',
|
||||
'sign',
|
||||
'cast',
|
||||
'fused_bn_add_activation',
|
||||
'c_identity',
|
||||
'c_concat',
|
||||
'all_reduce',
|
||||
'concat',
|
||||
'split',
|
||||
'fused_feedforward',
|
||||
'fused_attention',
|
||||
'fused_multi_transformer',
|
||||
}
|
||||
|
||||
CustomOpLists = AutoMixedPrecisionLists
|
||||
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,142 @@
|
||||
# Copyright (c) 2023 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.
|
||||
|
||||
# The implementation refers to https://arpitbhayani.me/blogs/function-overloading.
|
||||
# Note: it is customed for paddle.static.amp.decorate function.
|
||||
|
||||
import inspect
|
||||
import logging
|
||||
from enum import Enum
|
||||
|
||||
from paddle.base.log_helper import get_logger
|
||||
|
||||
_logger = get_logger(
|
||||
__name__, logging.INFO, fmt='%(asctime)s-%(levelname)s: %(message)s'
|
||||
)
|
||||
|
||||
|
||||
class FunctionType(Enum):
|
||||
FP16_ONLY = 0
|
||||
COMMON = 1
|
||||
|
||||
|
||||
class Function:
|
||||
"""
|
||||
Function is a wrap over standard python function
|
||||
An instance of this Function class is also callable
|
||||
just like the python function that it wrapped.
|
||||
When the instance is "called" like a function it fetches
|
||||
the function to be invoked from the virtual namespace and then
|
||||
invokes the same.
|
||||
"""
|
||||
|
||||
def __init__(self, fn):
|
||||
self.fn = fn
|
||||
|
||||
def __call__(self, *args, **kwargs):
|
||||
"""
|
||||
Overriding the __call__ function which makes the
|
||||
instance callable.
|
||||
"""
|
||||
# fetching the function to be invoked from the virtual namespace
|
||||
# through the arguments.
|
||||
fn = Namespace.get_instance().get(*args, **kwargs)
|
||||
# invoking the wrapped function and returning the value.
|
||||
return fn(*args, **kwargs)
|
||||
|
||||
|
||||
class Namespace:
|
||||
"""
|
||||
Namespace is the singleton class that is responsible
|
||||
for holding all the functions.
|
||||
"""
|
||||
|
||||
__instance = None
|
||||
|
||||
def __init__(self):
|
||||
if self.__instance is None:
|
||||
self.function_map = {}
|
||||
Namespace.__instance = self
|
||||
else:
|
||||
raise Exception("cannot instantiate Namespace again.")
|
||||
|
||||
@staticmethod
|
||||
def get_instance():
|
||||
if Namespace.__instance is None:
|
||||
Namespace()
|
||||
return Namespace.__instance
|
||||
|
||||
def register(self, fn, key):
|
||||
"""
|
||||
Register the function in the virtual namespace and return
|
||||
an instance of callable Function that wraps the function fn.
|
||||
|
||||
Args:
|
||||
fn (function): the native python function handle.
|
||||
key (FunctionType): the specified type.
|
||||
"""
|
||||
assert isinstance(key, FunctionType), (
|
||||
f"The type of key is expected to be FunctionType, but received {type(key)}."
|
||||
)
|
||||
func = Function(fn)
|
||||
self.function_map[key] = fn
|
||||
return func
|
||||
|
||||
def get(self, *args, **kwargs):
|
||||
"""
|
||||
Get the matching function from the virtual namespace according to the actual arguments.
|
||||
Return None if it did not find any matching function.
|
||||
"""
|
||||
_logger.debug(f"get function: args={args}, kwargs={kwargs}")
|
||||
satisfied_function_keys = set(self.function_map.keys())
|
||||
num_actual_args = len(args) + len(kwargs)
|
||||
for func_key in self.function_map.keys():
|
||||
if func_key not in satisfied_function_keys:
|
||||
continue
|
||||
fn = self.function_map[func_key]
|
||||
specs = inspect.getfullargspec(fn)
|
||||
if len(specs) < len(args) + len(kwargs):
|
||||
# Remove the not satisfied function according to the number of actual arguments.
|
||||
_logger.debug(
|
||||
f"fn={fn} (key={func_key}) is not satisfied and removed."
|
||||
)
|
||||
satisfied_function_keys.remove(func_key)
|
||||
continue
|
||||
if len(kwargs) > 0:
|
||||
# Remove the not satisfied function according to argument keys in kwargs.
|
||||
for arg_name, value in kwargs.items():
|
||||
if arg_name not in specs.args:
|
||||
_logger.debug(
|
||||
f"fn={fn} (key={func_key}) is not satisfied and removed."
|
||||
)
|
||||
satisfied_function_keys.remove(func_key)
|
||||
break
|
||||
if len(satisfied_function_keys) == 1:
|
||||
key = next(iter(satisfied_function_keys))
|
||||
elif len(args) >= 3 and isinstance(args[2], float):
|
||||
key = FunctionType.FP16_ONLY
|
||||
else:
|
||||
key = FunctionType.COMMON
|
||||
return self.function_map.get(key)
|
||||
|
||||
|
||||
def overload(key):
|
||||
"""overload is the decorator that wraps the function
|
||||
and returns a callable object of type Function.
|
||||
"""
|
||||
|
||||
def decorator(fn):
|
||||
return Namespace.get_instance().register(fn, key)
|
||||
|
||||
return decorator
|
||||
Reference in New Issue
Block a user