wehub-resource-sync
1443 lines
56 KiB
Python
1443 lines
56 KiB
Python
# 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 __future__ import annotations
|
|
|
|
import copy
|
|
import os
|
|
import warnings
|
|
from typing import (
|
|
TYPE_CHECKING,
|
|
Any,
|
|
Literal,
|
|
Protocol,
|
|
TypeVar,
|
|
overload,
|
|
)
|
|
|
|
import paddle
|
|
from paddle.base import core
|
|
from paddle.base.framework import (
|
|
_current_expected_place,
|
|
_dygraph_tracer,
|
|
dygraph_only,
|
|
in_dynamic_or_pir_mode,
|
|
in_pir_mode,
|
|
)
|
|
from paddle.base.wrapped_decorator import signature_safe_contextmanager
|
|
from paddle.static.amp.decorator import OptimizerWithMixedPrecision
|
|
|
|
from .amp_lists import black_list, white_list
|
|
|
|
if TYPE_CHECKING:
|
|
from collections.abc import Callable, Generator
|
|
from contextlib import AbstractContextManager
|
|
from typing import TypeAlias, TypeGuard
|
|
|
|
from paddle import Tensor
|
|
from paddle._typing import PlaceLike
|
|
from paddle._typing.dtype_like import _DTypeLiteral
|
|
from paddle.nn import Layer
|
|
from paddle.nn.layer.layers import _StateDict
|
|
from paddle.static import Operator, Program
|
|
|
|
_AmpLevelLiteral = Literal["O0", "OD", "O1", "O2"]
|
|
_CustomList: TypeAlias = list[str] | tuple[str, ...] | set[str]
|
|
|
|
class _OptimizerLike(Protocol):
|
|
def minimize(
|
|
self,
|
|
loss: Tensor,
|
|
startup_program: Program,
|
|
parameters: list[Tensor],
|
|
no_grad_set: set[Tensor],
|
|
) -> tuple[list[Operator], list[tuple[Tensor, Tensor]]]: ...
|
|
|
|
def step(
|
|
self, closure: Callable[[], Tensor] | None
|
|
) -> Tensor | None: ...
|
|
|
|
def set_state_dict(self, state_dict: dict[str, Tensor]) -> None: ...
|
|
|
|
def clear_grad(self, set_to_zero: bool) -> None: ...
|
|
|
|
|
|
_ModelsT = TypeVar("_ModelsT", "Layer", list["Layer"])
|
|
_OptimizersT = TypeVar("_OptimizersT", "_OptimizerLike", list["_OptimizerLike"])
|
|
|
|
|
|
AMP_RELATED_FLAGS = [
|
|
'FLAGS_cudnn_exhaustive_search',
|
|
'FLAGS_conv_workspace_size_limit',
|
|
'FLAGS_cudnn_batchnorm_spatial_persistent',
|
|
]
|
|
|
|
AMP_RELATED_FLAGS_SETTING = {
|
|
'FLAGS_cudnn_exhaustive_search': 1,
|
|
'FLAGS_conv_workspace_size_limit': 1000,
|
|
'FLAGS_cudnn_batchnorm_spatial_persistent': 1,
|
|
}
|
|
|
|
AMP_LEVEL = core.AmpLevel
|
|
_g_amp_state_ = None
|
|
|
|
|
|
def amp_state():
|
|
global _g_amp_state_
|
|
return _g_amp_state_
|
|
|
|
|
|
class AMPGlobalState:
|
|
model_parameters: list[Tensor]
|
|
use_master_grad: bool
|
|
already_register_final_backward_hook: bool
|
|
already_classify_params_meshes: bool
|
|
mesh2params: dict[paddle.distributed.ProcessMesh | None, list[Tensor]]
|
|
amp_dtype: _DTypeLiteral
|
|
|
|
def __init__(self) -> None:
|
|
self.model_parameters = []
|
|
self.use_master_grad = False
|
|
self.already_register_final_backward_hook = False
|
|
self.already_classify_params_meshes = False # For dist
|
|
self.mesh2params = {} # For dist
|
|
self.amp_dtype = 'float32'
|
|
|
|
def __setattr__(self, name: str, val: Any) -> None:
|
|
self.__dict__[name] = val
|
|
|
|
|
|
_amp_global_state = AMPGlobalState()
|
|
|
|
|
|
def amp_global_state() -> AMPGlobalState:
|
|
return _amp_global_state
|
|
|
|
|
|
# NOTE(zhiqiu): similar as paddle.static.amp.fp16_lists.AutoMixedPrecisionLists._update_list
|
|
# The reason why not use AutoMixedPrecisionLists is that custom_black_varnames is not suitable for imperative mode.
|
|
def _update_list(
|
|
custom_white_list: _CustomList,
|
|
custom_black_list: _CustomList,
|
|
level: _AmpLevelLiteral = 'O1',
|
|
dtype: _DTypeLiteral = 'float16',
|
|
) -> tuple[set[str], set[str]]:
|
|
"""
|
|
Update black and white list according to users' custom list.
|
|
"""
|
|
if level == 'O0':
|
|
_white_list = set()
|
|
_black_list = set()
|
|
return _white_list, _black_list
|
|
_white_list = copy.copy(white_list()[dtype][level])
|
|
_black_list = copy.copy(black_list()[dtype][level])
|
|
if custom_white_list and custom_black_list:
|
|
for op_name in custom_white_list:
|
|
if op_name in custom_black_list:
|
|
raise ValueError("Custom white list overlap custom black list")
|
|
if custom_white_list:
|
|
for op_name in custom_white_list:
|
|
if op_name in _black_list:
|
|
_black_list.remove(op_name)
|
|
_white_list.add(op_name)
|
|
if custom_black_list:
|
|
for op_name in custom_black_list:
|
|
if op_name in _white_list:
|
|
_white_list.remove(op_name)
|
|
_black_list.add(op_name)
|
|
return _white_list, _black_list
|
|
|
|
|
|
def _in_amp_guard() -> bool:
|
|
"""
|
|
Judge whether current code block is in `amp_guard` context.
|
|
"""
|
|
tracer = _dygraph_tracer()
|
|
if tracer:
|
|
if tracer._amp_level == core.AmpLevel.O1:
|
|
return True
|
|
else:
|
|
return False
|
|
else:
|
|
return False
|
|
|
|
|
|
def _in_pure_fp16_guard() -> bool:
|
|
tracer = _dygraph_tracer()
|
|
return tracer and tracer._amp_level == core.AmpLevel.O2
|
|
|
|
|
|
def _is_gpu_float16_supported() -> bool:
|
|
"""
|
|
Judge whether current gpu support float16 amp.
|
|
"""
|
|
prop = paddle.device.cuda.get_device_capability()
|
|
return prop[0] >= 7 or paddle.is_compiled_with_rocm()
|
|
|
|
|
|
def _is_gpu_bfloat16_supported() -> bool:
|
|
"""
|
|
Judge whether current gpu support bfloat16 amp.
|
|
"""
|
|
prop = paddle.device.cuda.get_device_capability()
|
|
cuda_version = paddle.version.cuda()
|
|
if cuda_version is not None and cuda_version != 'False':
|
|
cuda_version_check = int(cuda_version.split('.')[0]) >= 11
|
|
else:
|
|
cuda_version_check = False
|
|
return prop[0] >= 8 and cuda_version_check or paddle.is_compiled_with_rocm()
|
|
|
|
|
|
def _is_xpu_float16_supported() -> bool:
|
|
"""
|
|
Judge whether current xpu device support float16 amp.
|
|
Only XPU2 and XPU3 support float16 amp.
|
|
"""
|
|
place = _current_expected_place()
|
|
return (
|
|
core.get_xpu_device_version(place.get_device_id())
|
|
>= core.XPUVersion.XPU2
|
|
)
|
|
|
|
|
|
def _is_xpu_bfloat16_supported() -> bool:
|
|
"""
|
|
Judge whether current xpu device support bfloat16 amp.
|
|
Only XPU3 support bfloat16 amp.
|
|
Although XPU2 supports bfloat16 computing, but XPU2's bfloat16 operators haven't been widely covered.
|
|
"""
|
|
place = _current_expected_place()
|
|
return (
|
|
core.get_xpu_device_version(place.get_device_id())
|
|
>= core.XPUVersion.XPU3
|
|
)
|
|
|
|
|
|
def _is_custom_device_bfloat16_supported() -> bool:
|
|
"""
|
|
Judge whether current custom device support bfloat16 amp.
|
|
"""
|
|
place = _current_expected_place()
|
|
return (
|
|
place.get_device_type() == 'npu'
|
|
or place.get_device_type() == 'intel_hpu'
|
|
or place.get_device_type() == 'iluvatar_gpu'
|
|
or place.get_device_type() == 'metax_gpu'
|
|
)
|
|
|
|
|
|
def need_keep_fp32(layer: Layer, dtype: str) -> bool:
|
|
need_keep_fp32 = False
|
|
# Highest priority. Because all the layers except BN will use bfloat16 params in bfloat16 training,
|
|
# here we provide a option to keep fp32 param.
|
|
if not layer._cast_to_low_precision:
|
|
need_keep_fp32 = True
|
|
# The BN layers will keep fp32
|
|
elif isinstance(
|
|
layer,
|
|
(
|
|
paddle.nn.BatchNorm,
|
|
paddle.nn.BatchNorm1D,
|
|
paddle.nn.BatchNorm2D,
|
|
paddle.nn.BatchNorm3D,
|
|
paddle.nn.SyncBatchNorm,
|
|
),
|
|
):
|
|
need_keep_fp32 = True
|
|
# layer._dtype is used to set params dtype. BF16 will use bf16 params.
|
|
elif (layer._dtype == 'float16') or (
|
|
(dtype == 'float16')
|
|
and isinstance(
|
|
layer,
|
|
(
|
|
paddle.nn.LayerNorm,
|
|
paddle.nn.InstanceNorm1D,
|
|
paddle.nn.InstanceNorm2D,
|
|
paddle.nn.InstanceNorm3D,
|
|
),
|
|
)
|
|
):
|
|
need_keep_fp32 = True
|
|
|
|
return need_keep_fp32
|
|
|
|
|
|
def set_excluded_layers(
|
|
models: list[Layer],
|
|
excluded_layers: Layer | list[Layer | type[Layer]] | type[Layer],
|
|
) -> None:
|
|
excluded_layers_instances = []
|
|
excluded_layers_types = []
|
|
error_message = "excluded_layers must be either a nn.Layer instance/type or a list of nn.Layer instances/types."
|
|
if excluded_layers is None:
|
|
excluded_layers = []
|
|
elif isinstance(excluded_layers, paddle.nn.Layer):
|
|
excluded_layers_instances = [excluded_layers]
|
|
elif isinstance(excluded_layers, type) and issubclass(
|
|
excluded_layers, paddle.nn.Layer
|
|
):
|
|
excluded_layers_types = [excluded_layers]
|
|
elif isinstance(excluded_layers, list):
|
|
for item in excluded_layers:
|
|
if isinstance(item, paddle.nn.Layer):
|
|
excluded_layers_instances.append(item)
|
|
elif issubclass(item, paddle.nn.Layer):
|
|
excluded_layers_types.append(item)
|
|
else:
|
|
raise TypeError(error_message)
|
|
else:
|
|
raise TypeError(error_message)
|
|
|
|
for idx in range(len(excluded_layers_instances)):
|
|
for layer in excluded_layers_instances[idx].sublayers(
|
|
include_self=True
|
|
):
|
|
layer._cast_to_low_precision = False
|
|
excluded_layers_types = tuple(excluded_layers_types)
|
|
for idx in range(len(models)):
|
|
for layer in models[idx].sublayers(include_self=True):
|
|
if isinstance(layer, excluded_layers_types):
|
|
layer._cast_to_low_precision = False
|
|
|
|
|
|
def _pir_apply(
|
|
self: Layer,
|
|
func: Callable[[Tensor, _DTypeLiteral], Tensor | None],
|
|
dtype: _DTypeLiteral,
|
|
include_sublayers: bool = True,
|
|
) -> None:
|
|
if include_sublayers:
|
|
for layer in self.children():
|
|
_pir_apply(layer, func, dtype, include_sublayers)
|
|
|
|
for key, param in self._parameters.items():
|
|
if param is not None:
|
|
param_applied = func(param, dtype)
|
|
|
|
for key, buf in self._buffers.items():
|
|
if buf is not None:
|
|
self._buffers[key] = func(buf, dtype)
|
|
|
|
self._dtype = dtype
|
|
|
|
|
|
def _pir_transform(t: Tensor, dtype: str) -> None:
|
|
main = paddle.static.default_main_program()
|
|
startup = paddle.static.default_startup_program()
|
|
with paddle.static.program_guard(startup):
|
|
block = startup.global_block()
|
|
for op in block.ops:
|
|
if (
|
|
op.name() == 'builtin.set_parameter'
|
|
and op.attrs()['parameter_name'] == t.name
|
|
):
|
|
param = op.operand(0).source()
|
|
cast_param = paddle.cast(param, dtype)
|
|
cast_param.persistable = True
|
|
paddle._pir_ops.update_parameter(cast_param, t.name)
|
|
block.remove_op(op)
|
|
break
|
|
main.set_parameters_from(startup)
|
|
with paddle.static.program_guard(main):
|
|
paddle.pir.reset_insertion_point_to_start()
|
|
block = main.global_block()
|
|
cast_param = paddle._pir_ops.parameter(t.name)
|
|
cast_param.trainable = t.trainable
|
|
cast_param.stop_gradient = t.stop_gradient
|
|
cast_param.persistable = t.persistable
|
|
cast_param.optimize_attr = t.optimize_attr
|
|
cast_param.regularizer = t.regularizer
|
|
cast_param.do_model_average = t.do_model_average
|
|
cast_param.need_clip = t.need_clip
|
|
cast_param.is_distributed = t.is_distributed
|
|
cast_param.is_parameter = t.is_parameter
|
|
op = t.get_defining_op()
|
|
t.replace_all_uses_with(cast_param)
|
|
block.remove_op(op)
|
|
t.value_assign(cast_param)
|
|
|
|
|
|
def _pir_to_impl(
|
|
self: Layer,
|
|
dtype: _DTypeLiteral,
|
|
include_sublayers: bool,
|
|
floating_only: bool,
|
|
) -> Layer:
|
|
def transform(t: Tensor, dtype: _DTypeLiteral) -> Tensor | None:
|
|
if floating_only and (not paddle.is_floating_point(t)):
|
|
return t
|
|
return _pir_transform(t, dtype)
|
|
|
|
with warnings.catch_warnings():
|
|
warnings.filterwarnings("ignore", category=UserWarning)
|
|
_pir_apply(self, transform, dtype, include_sublayers)
|
|
|
|
self._dtype = dtype
|
|
return self
|
|
|
|
|
|
def amp_initialize(
|
|
models: list[Layer],
|
|
dtype: _DTypeLiteral,
|
|
excluded_layers: Layer | list[Layer | type[Layer]] | type[Layer],
|
|
) -> list[Layer]:
|
|
set_excluded_layers(models, excluded_layers)
|
|
for idx in range(len(models)):
|
|
for layer in models[idx].sublayers(include_self=True):
|
|
if need_keep_fp32(layer, dtype):
|
|
continue
|
|
if dtype == "float16" and isinstance(
|
|
layer,
|
|
(
|
|
paddle.incubate.nn.FusedFeedForward,
|
|
paddle.incubate.nn.FusedMultiHeadAttention,
|
|
),
|
|
):
|
|
layer._amp_decorate(dtype=dtype)
|
|
continue
|
|
|
|
if in_pir_mode():
|
|
_pir_to_impl(
|
|
layer,
|
|
dtype=dtype,
|
|
include_sublayers=False,
|
|
floating_only=True,
|
|
)
|
|
else:
|
|
layer._to_impl(
|
|
dtype=dtype, include_sublayers=False, floating_only=True
|
|
)
|
|
return models
|
|
|
|
|
|
def check_models(models: list[Layer]) -> None:
|
|
for model in models:
|
|
if not isinstance(model, paddle.nn.Layer):
|
|
raise RuntimeError(
|
|
f"Current train mode is pure fp16, models should be paddle.nn.Layer, but receive {type(model)}."
|
|
)
|
|
if isinstance(model, paddle.DataParallel):
|
|
raise RuntimeError(
|
|
"For distributed AMP training, you should first use paddle.amp.decorate() to decorate origin model, and then call paddle.DataParallel get distributed model."
|
|
)
|
|
|
|
|
|
def _is_valid_optimizer(optimizer: Any) -> TypeGuard[_OptimizerLike]:
|
|
from paddle.distributed.fleet.meta_optimizers.dygraph_optimizer.dygraph_sharding_optimizer import (
|
|
DygraphShardingOptimizer,
|
|
DygraphShardingOptimizerV2,
|
|
)
|
|
|
|
return isinstance(
|
|
optimizer,
|
|
(
|
|
paddle.optimizer.Optimizer,
|
|
DygraphShardingOptimizer,
|
|
DygraphShardingOptimizerV2,
|
|
),
|
|
)
|
|
|
|
|
|
def check_optimizers(optimizers: list[Any]) -> None:
|
|
for optimizer in optimizers:
|
|
if not _is_valid_optimizer(optimizer):
|
|
raise RuntimeError(
|
|
f"Current train mode is pure fp16, optimizers should be paddle.optimizer.Optimizer or DygraphShardingOptimizer, but receive {type(optimizer)}."
|
|
)
|
|
|
|
|
|
@signature_safe_contextmanager
|
|
def amp_guard(
|
|
enable: bool = True,
|
|
custom_white_list: _CustomList | None = None,
|
|
custom_black_list: _CustomList | None = None,
|
|
level: _AmpLevelLiteral = 'O1',
|
|
dtype: _DTypeLiteral = 'float16',
|
|
use_promote: bool = True,
|
|
) -> Generator[None, None, None]:
|
|
"""
|
|
Create a context which enables auto-mixed-precision(AMP) of operators executed in dynamic graph mode.
|
|
If enabled, the input data type (float32 or float16) of each operator is decided
|
|
by autocast algorithm for better performance.
|
|
|
|
Commonly, it is used together with `GradScaler` to achieve Auto-Mixed-Precision in
|
|
imperative mode. It is used together with `decorator` to achieve Pure fp16 in imperative mode.
|
|
|
|
Args:
|
|
enable(bool, optional): Enable auto-mixed-precision or not. Default is True.
|
|
custom_white_list(set|list|tuple|None, optional): The custom white_list. It's the set of ops that support
|
|
fp16 calculation and are considered numerically-safe and performance-critical. These ops
|
|
will be converted to fp16.
|
|
custom_black_list(set|list|tuple|None, optional): The custom black_list. The set of ops that support fp16
|
|
calculation and are considered numerically-dangerous and whose effects may also be
|
|
observed in downstream ops. These ops will not be converted to fp16.
|
|
level(str, optional): Auto mixed precision level. Accepted values are "O1" and "O2": O1 represent mixed precision, the input data type of each operator will be casted by white_list and black_list;
|
|
O2 represent Pure fp16, all operators parameters and input data will be casted to fp16, except operators in black_list, don't support fp16 kernel and batchnorm. Default is O1(amp).
|
|
dtype(str|core.DataType, optional): Whether to use 'float16' or 'bfloat16'. Default is 'float16'.
|
|
use_promote(bool, optional): Whether op's dtype is 'float32', accord 'Promote to the Widest' principle, use 'float32' to calculate.
|
|
Only active on 'AMP-02'. Default is True.
|
|
|
|
Examples:
|
|
|
|
.. code-block:: pycon
|
|
|
|
>>> # doctest: +REQUIRES(env:GPU)
|
|
>>> import paddle
|
|
|
|
>>> data = paddle.uniform([10, 3, 32, 32], paddle.float32, -1, 1)
|
|
>>> conv2d = paddle.nn.Conv2D(3, 2, 3, bias_attr=False)
|
|
>>> conv2d = paddle.amp.amp_decorate(models=conv2d, level='O2')
|
|
>>> with paddle.amp.amp_guard():
|
|
... conv = conv2d(data)
|
|
... print(conv.dtype)
|
|
>>> # doctest: +SKIP("This has diff in xdoctest env")
|
|
paddle.float16
|
|
>>> # doctest: -SKIP
|
|
>>> with paddle.amp.amp_guard(enable=False):
|
|
... conv = conv2d(data)
|
|
... print(conv.dtype)
|
|
>>> # doctest: +SKIP("This has diff in xdoctest env")
|
|
paddle.float32
|
|
>>> # doctest: -SKIP
|
|
"""
|
|
assert in_dynamic_or_pir_mode(), (
|
|
"We only support 'amp_guard' in dynamic or pir mode."
|
|
)
|
|
|
|
amp_state = locals()
|
|
global _g_amp_state_
|
|
original_state = _g_amp_state_
|
|
_g_amp_state_ = amp_state
|
|
|
|
# check amp_level: O0-O2
|
|
level = level.upper()
|
|
if level not in ['O0', 'OD', 'O1', 'O2']:
|
|
raise ValueError("level should be O0, OD, O1 or O2.")
|
|
|
|
# check amp_dtype: float16 or bfloat16
|
|
if isinstance(dtype, paddle.base.core.DataType):
|
|
dtype = dtype.name
|
|
dtype = dtype.lower()
|
|
if enable:
|
|
if dtype not in ['float16', 'bfloat16']:
|
|
raise ValueError(
|
|
"If enable amp, dtype should be 'float16' or 'bfloat16'."
|
|
)
|
|
|
|
amp_dtype = dtype
|
|
amp_global_state().amp_dtype = amp_dtype
|
|
|
|
if level == 'OD':
|
|
amp_level = AMP_LEVEL.OD
|
|
elif level == 'O1':
|
|
amp_level = AMP_LEVEL.O1
|
|
elif level == 'O2':
|
|
amp_level = AMP_LEVEL.O2
|
|
elif level == 'O0':
|
|
amp_level = AMP_LEVEL.O0
|
|
|
|
_white_list, _black_list = _update_list(
|
|
custom_white_list, custom_black_list, level, dtype
|
|
)
|
|
|
|
if in_pir_mode():
|
|
if not enable:
|
|
amp_level = AMP_LEVEL.O0
|
|
amp_dtype = "float32"
|
|
amp_attrs = core._get_amp_attrs()
|
|
# set amp level
|
|
original_amp_level = amp_attrs._amp_level
|
|
amp_attrs._amp_level = amp_level
|
|
# set amp op list
|
|
original_white_list, original_black_list = core._get_amp_op_list()
|
|
core._set_amp_op_list(_white_list, _black_list)
|
|
# set amp dtype
|
|
original_amp_dtype = amp_attrs._amp_dtype
|
|
amp_attrs._amp_dtype = amp_dtype
|
|
# switch promote
|
|
if amp_level == AMP_LEVEL.O2:
|
|
original_use_promote = amp_attrs._use_promote
|
|
amp_attrs._use_promote = use_promote
|
|
|
|
try:
|
|
yield
|
|
finally:
|
|
_g_amp_state_ = original_state
|
|
amp_attrs._amp_level = original_amp_level
|
|
core._set_amp_op_list(original_white_list, original_black_list)
|
|
amp_attrs._amp_dtype = original_amp_dtype
|
|
if amp_level == AMP_LEVEL.O2:
|
|
amp_attrs._use_promote = original_use_promote
|
|
|
|
else:
|
|
# check tracer
|
|
tracer = _dygraph_tracer()
|
|
if not tracer:
|
|
raise ValueError(
|
|
"current_tracer is None, maybe it is not in imperative mode."
|
|
)
|
|
# check device_type:
|
|
# NOTE: Now, amp only support gpu for float16 and bfloat16, xpu for float16, npu for float16 and bfloat16.
|
|
# Maybe we will support cpu for bfloat16.
|
|
if enable and not (
|
|
tracer._expected_place.is_gpu_place()
|
|
or tracer._expected_place.is_xpu_place()
|
|
or tracer._expected_place.is_custom_place()
|
|
):
|
|
warnings.warn(
|
|
f'amp_guard can only be enabled on CUDAPlace, XPUPlace, and CustomPlace, current place is {tracer._expected_place}, so it makes no effect.'
|
|
)
|
|
enable = False
|
|
if enable:
|
|
# For xpu:
|
|
if tracer._expected_place.is_xpu_place():
|
|
if (dtype == 'float16') and not _is_xpu_float16_supported():
|
|
xpu_version = core.get_xpu_device_version(
|
|
_current_expected_place().get_device_id()
|
|
)
|
|
warnings.warn(
|
|
f'{core.XPUVersion(xpu_version)} does not support float16 amp.'
|
|
)
|
|
enable = False
|
|
elif (dtype == 'bfloat16') and not _is_xpu_bfloat16_supported():
|
|
xpu_version = core.get_xpu_device_version(
|
|
_current_expected_place().get_device_id()
|
|
)
|
|
warnings.warn(
|
|
f'{core.XPUVersion(xpu_version)} does not support bfloat16 amp.'
|
|
)
|
|
enable = False
|
|
# For custom device:
|
|
if (
|
|
tracer._expected_place.is_custom_place()
|
|
and not _is_custom_device_bfloat16_supported()
|
|
and (dtype == 'bfloat16')
|
|
):
|
|
warnings.warn('CustomPlace only support float16 amp.')
|
|
enable = False
|
|
# For gpu float16: Compute Capability should >= 7.
|
|
# For gpu bfloat16: Compute Capability should >= 8 & CUDA Version should >= 11.
|
|
if tracer._expected_place.is_gpu_place():
|
|
if (dtype == 'float16') and not _is_gpu_float16_supported():
|
|
prop = paddle.device.cuda.get_device_capability()
|
|
warnings.warn(
|
|
f"For float16, amp only support NVIDIA GPU with Compute Capability 7.0 or higher, current GPU is: {paddle.device.cuda.get_device_name()}, with Compute Capability: {prop[0]}.{prop[1]}."
|
|
)
|
|
enable = False
|
|
elif (dtype == 'bfloat16') and not _is_gpu_bfloat16_supported():
|
|
prop = paddle.device.cuda.get_device_capability()
|
|
cuda_version = paddle.version.cuda()
|
|
warnings.warn(
|
|
f"For bfloat16, amp only support NVIDIA GPU with Compute Capability 8.0 or higher and CUDA Version 11.0 or higher, current GPU is: {paddle.device.cuda.get_device_name()}, with Compute Capability: {prop[0]}.{prop[1]}, current CUDA Version is: {cuda_version}."
|
|
)
|
|
enable = False
|
|
|
|
if not enable:
|
|
amp_level = AMP_LEVEL.O0
|
|
amp_dtype = "float32"
|
|
|
|
# master_grad_hook will run at the end of backward.
|
|
# Since backward_final_hook will be cleared once they have been
|
|
# done, we should register the hook every step.
|
|
if (
|
|
amp_global_state().use_master_grad
|
|
and not amp_global_state().already_register_final_backward_hook
|
|
):
|
|
|
|
def _dtensor_from_local(local_tensor, mesh, placements):
|
|
global_dims = list(local_tensor.shape)
|
|
for idx, placement in enumerate(placements):
|
|
if placement.is_shard():
|
|
global_dims[placement.get_dim()] = (
|
|
global_dims[placement.get_dim()] * mesh.shape[idx]
|
|
)
|
|
place = paddle.framework._current_expected_place()
|
|
place = paddle.framework._get_paddle_place(place)
|
|
|
|
return paddle.Tensor(
|
|
local_tensor,
|
|
dims=global_dims,
|
|
process_mesh=mesh,
|
|
placements=placements,
|
|
place=place,
|
|
)
|
|
|
|
def master_grad_hook():
|
|
# NOTE(lizhiyu): To support semi-auto of dygraph mode, we must
|
|
# classify the params of model into different classes according to their process_mesh.
|
|
# Otherwise, fault will occur.
|
|
if not amp_global_state().already_classify_params_meshes:
|
|
for param in amp_global_state().model_parameters:
|
|
if param is not None and param.process_mesh is not None:
|
|
if (
|
|
param.process_mesh
|
|
not in amp_global_state().mesh2params
|
|
):
|
|
amp_global_state().mesh2params[
|
|
param.process_mesh
|
|
] = [param]
|
|
else:
|
|
amp_global_state().mesh2params[
|
|
param.process_mesh
|
|
].append(param)
|
|
amp_global_state().already_classify_params_meshes = True
|
|
|
|
if len(amp_global_state().mesh2params):
|
|
for _, params in amp_global_state().mesh2params.items():
|
|
core.eager.set_master_grads(params)
|
|
else:
|
|
core.eager.set_master_grads(
|
|
amp_global_state().model_parameters
|
|
)
|
|
|
|
amp_global_state().already_register_final_backward_hook = False
|
|
|
|
def _update_main_grad_hook(param):
|
|
@paddle.autograd.no_grad()
|
|
def param_hook(tmp_grad):
|
|
if tmp_grad is not None and tmp_grad._is_initialized():
|
|
if param.main_grad is None:
|
|
tmp = core.eager.Tensor(
|
|
value=tmp_grad._local_value()
|
|
.cast(paddle.float32)
|
|
.value(),
|
|
place=tmp_grad.place,
|
|
name="main_grad@" + param.name,
|
|
)
|
|
param.main_grad = _dtensor_from_local(
|
|
tmp,
|
|
tmp_grad.process_mesh,
|
|
tmp_grad.placements,
|
|
)
|
|
else:
|
|
param.main_grad._local_value().add_(
|
|
tmp_grad._local_value()
|
|
)
|
|
tmp_grad._clear_data()
|
|
|
|
return param_hook
|
|
|
|
if os.getenv("FLAGS_enable_tensor_fusion") in [
|
|
"True",
|
|
"true",
|
|
"1",
|
|
]:
|
|
for param in amp_global_state().model_parameters:
|
|
if not hasattr(param, "main_grad"):
|
|
param.main_grad = None
|
|
param._register_grad_hook(_update_main_grad_hook(param))
|
|
os.environ["FLAGS_enable_tensor_fusion"] = "0"
|
|
else:
|
|
core.eager._add_backward_final_hook(master_grad_hook)
|
|
amp_global_state().already_register_final_backward_hook = True
|
|
|
|
if tracer:
|
|
# enable auto_cast
|
|
original_amp_level = tracer._amp_level
|
|
tracer._amp_level = amp_level
|
|
|
|
# set amp op list
|
|
original_white_list, original_black_list = tracer._get_amp_op_list()
|
|
tracer._set_amp_op_list(_white_list, _black_list)
|
|
|
|
# TODO(zhiqiu) set amp related flags automatically in this guard
|
|
# Currently, if FLAGS_cudnn_batchnorm_spatial_persistent is set True in amp_guard,
|
|
# batch_norm can run in fast mode, but batch_norm_grad can not if backward if not executed inside amp_guard.
|
|
# So, users need to set related flags manually.
|
|
|
|
# original_flags = get_flags(AMP_RELATED_FLAGS)
|
|
# set_flags(AMP_RELATED_FLAGS_SETTING)
|
|
|
|
# set amp dtype
|
|
original_amp_dtype = tracer._amp_dtype
|
|
tracer._amp_dtype = amp_dtype
|
|
|
|
# switch promote
|
|
if amp_level == AMP_LEVEL.O2:
|
|
original_use_promote = tracer._use_promote
|
|
tracer._use_promote = use_promote
|
|
|
|
# restore status
|
|
try:
|
|
yield
|
|
finally:
|
|
if tracer:
|
|
_g_amp_state_ = original_state
|
|
tracer._amp_level = original_amp_level
|
|
tracer._set_amp_op_list(
|
|
original_white_list, original_black_list
|
|
)
|
|
# set_flags(original_flags)
|
|
tracer._amp_dtype = original_amp_dtype
|
|
if amp_level == AMP_LEVEL.O2:
|
|
tracer._use_promote = original_use_promote
|
|
|
|
|
|
class StateDictHook:
|
|
def __init__(self, save_dtype: str) -> None:
|
|
self._save_dtype = save_dtype
|
|
|
|
def __call__(self, state_dict: _StateDict) -> None:
|
|
with paddle.base.framework._dygraph_guard(paddle.base.dygraph.Tracer()):
|
|
for key in state_dict:
|
|
param = state_dict[key]
|
|
if not isinstance(param, paddle.Tensor):
|
|
continue
|
|
if paddle.is_floating_point(param):
|
|
param_applied = paddle.cast(param, self._save_dtype)
|
|
param_applied.name = param.name
|
|
state_dict[key] = param_applied
|
|
|
|
|
|
def _set_multi_precision(
|
|
optimizer: _OptimizerLike, multi_precision: bool
|
|
) -> None:
|
|
from paddle.distributed.fleet.meta_optimizers.dygraph_optimizer.dygraph_sharding_optimizer import (
|
|
DygraphShardingOptimizer,
|
|
DygraphShardingOptimizerV2,
|
|
)
|
|
|
|
optimizer = (
|
|
optimizer._inner_opt
|
|
if isinstance(
|
|
optimizer, (DygraphShardingOptimizer, DygraphShardingOptimizerV2)
|
|
)
|
|
else optimizer
|
|
)
|
|
if hasattr(optimizer, "_multi_precision"):
|
|
optimizer._multi_precision = multi_precision
|
|
|
|
|
|
@overload
|
|
def amp_decorate(
|
|
models: _ModelsT,
|
|
optimizers: _OptimizersT = ...,
|
|
level: _AmpLevelLiteral = ...,
|
|
dtype: _DTypeLiteral = ...,
|
|
master_weight: bool | None = ...,
|
|
save_dtype: _DTypeLiteral | None = ...,
|
|
master_grad: bool = ...,
|
|
excluded_layers: (
|
|
Layer | list[Layer | type[Layer]] | type[Layer] | None
|
|
) = ...,
|
|
) -> tuple[_ModelsT, _OptimizersT]: ...
|
|
|
|
|
|
@overload
|
|
def amp_decorate(
|
|
models: _ModelsT,
|
|
optimizers: None = ...,
|
|
level: _AmpLevelLiteral = ...,
|
|
dtype: _DTypeLiteral = ...,
|
|
master_weight: bool | None = ...,
|
|
save_dtype: _DTypeLiteral | None = ...,
|
|
master_grad: bool = ...,
|
|
excluded_layers: (
|
|
Layer | list[Layer | type[Layer]] | type[Layer] | None
|
|
) = ...,
|
|
) -> _ModelsT: ...
|
|
|
|
|
|
@dygraph_only
|
|
def amp_decorate(
|
|
models: _ModelsT,
|
|
optimizers: _OptimizersT | None = None,
|
|
level: _AmpLevelLiteral = 'O1',
|
|
dtype: _DTypeLiteral = 'float16',
|
|
master_weight: bool | None = None,
|
|
save_dtype: _DTypeLiteral | None = None,
|
|
master_grad: bool = False,
|
|
excluded_layers: (
|
|
Layer | list[Layer | type[Layer]] | type[Layer] | None
|
|
) = None,
|
|
) -> tuple[_ModelsT, _OptimizersT] | _ModelsT:
|
|
"""
|
|
Decorate models and optimizers for auto-mixed-precision. When level is O1(amp), the decorate will do nothing.
|
|
When level is O2(pure fp16), the decorate will cast all parameters of models to FP16, except BatchNorm, InstanceNorm and LayerNorm.
|
|
|
|
Commonly, it is used together with `amp_guard` to achieve Pure fp16 in imperative mode.
|
|
|
|
Args:
|
|
models(Layer|list of Layer, optional): The defined models by user, models must be either a single model or a list of models. Default is None.
|
|
optimizers(Optimizer|list of Optimizer|None, optional): The defined optimizers by user, optimizers must be either a single optimizer or a list of optimizers. Default is None.
|
|
level(str, optional): Auto mixed precision level. Accepted values are "O1" and "O2": O1 represent mixed precision, the decorator will do nothing;
|
|
O2 represent Pure fp16/bf16, the decorator will cast all parameters of models to FP16/BF16, except BatchNorm, InstanceNorm and LayerNorm. Default is O1(amp)
|
|
dtype(str, optional): Whether to use 'float16' or 'bfloat16'. Default is 'float16'.
|
|
master_weight(bool|None, optional): For level='O2', whether to use multi-precision during weight updating. If master_weight is None, in O2 level optimizer will use multi-precision. Default is None.
|
|
save_dtype(str|None, optional): The save model parameter dtype when use `paddle.save` or `paddle.jit.save`,it should be float16, bfloat16, float32, float64 or None.
|
|
The save_dtype will not change model parameters dtype, it just change the state_dict dtype. When save_dtype is None, the save dtype is same as model dtype. Default is None.
|
|
master_grad(bool, optional): For level='O2', whether to use float32 weight gradients for calculations such as gradient clipping, weight decay, and weight updates. If master_grad is enabled, the weight
|
|
gradients will be float32 dtype after the back propagation. Default is False, there is only float16 weight gradients.
|
|
excluded_layers(Layer|list of Layer, optional): Specify the layers not to be decorated. The weights of these layers will always keep float32 when level is O2. `excluded_layers` can be specified as
|
|
an Layer instance/type or a list of Layer instances/types. Default is None, the weights of the whole model will be casted to float16 or bfloat16.
|
|
|
|
Examples:
|
|
|
|
.. code-block:: pycon
|
|
|
|
>>> # doctest: +REQUIRES(env:GPU)
|
|
>>> # Demo1: single model and optimizer:
|
|
>>> import paddle
|
|
>>> paddle.device.set_device('gpu')
|
|
|
|
>>> model = paddle.nn.Conv2D(3, 2, 3, bias_attr=False)
|
|
>>> optimizer = paddle.optimizer.SGD(parameters=model.parameters())
|
|
|
|
>>> model, optimizer = paddle.amp.amp_decorate(models=model, optimizers=optimizer, level='O2')
|
|
|
|
>>> data = paddle.rand([10, 3, 32, 32])
|
|
|
|
>>> with paddle.amp.amp_guard(enable=True, custom_white_list=None, custom_black_list=None, level='O2'):
|
|
... output = model(data)
|
|
... print(output.dtype)
|
|
paddle.float16
|
|
|
|
>>> # Demo2: multi models and optimizers:
|
|
>>> model2 = paddle.nn.Conv2D(3, 2, 3, bias_attr=False)
|
|
>>> optimizer2 = paddle.optimizer.Adam(parameters=model2.parameters())
|
|
|
|
>>> models, optimizers = paddle.amp.amp_decorate(models=[model, model2], optimizers=[optimizer, optimizer2], level='O2')
|
|
|
|
>>> data = paddle.rand([10, 3, 32, 32])
|
|
|
|
>>> with paddle.amp.amp_guard(enable=True, custom_white_list=None, custom_black_list=None, level='O2'):
|
|
... output = models[0](data)
|
|
... output2 = models[1](data)
|
|
... print(output.dtype)
|
|
... print(output2.dtype)
|
|
paddle.float16
|
|
paddle.float16
|
|
|
|
>>> # Demo3: optimizers is None:
|
|
>>> model3 = paddle.nn.Conv2D(3, 2, 3, bias_attr=False)
|
|
>>> optimizer3 = paddle.optimizer.Adam(parameters=model2.parameters())
|
|
|
|
>>> model = paddle.amp.amp_decorate(models=model3, level='O2')
|
|
|
|
>>> data = paddle.rand([10, 3, 32, 32])
|
|
|
|
>>> with paddle.amp.amp_guard(enable=True, custom_white_list=None, custom_black_list=None, level='O2'):
|
|
... output = model(data)
|
|
... print(output.dtype)
|
|
paddle.float16
|
|
"""
|
|
if level not in ['O1', 'O2']:
|
|
raise ValueError(
|
|
"level should be O1 or O2, O1 represent AMP train mode, O2 represent Pure fp16 train mode."
|
|
)
|
|
if dtype not in ['float16', 'bfloat16']:
|
|
raise ValueError("dtype only support float16 or bfloat16.")
|
|
|
|
if level == 'O1':
|
|
if optimizers is None:
|
|
return models
|
|
else:
|
|
return models, optimizers
|
|
|
|
# check tracer
|
|
tracer = _dygraph_tracer()
|
|
if not tracer:
|
|
raise ValueError(
|
|
"current_tracer is None, maybe it is not in imperative mode."
|
|
)
|
|
|
|
# check device_type:
|
|
if not (
|
|
tracer._expected_place.is_gpu_place()
|
|
or tracer._expected_place.is_xpu_place()
|
|
or tracer._expected_place.is_custom_place()
|
|
):
|
|
if optimizers is None:
|
|
return models
|
|
else:
|
|
return models, optimizers
|
|
# For xpu:
|
|
if tracer._expected_place.is_xpu_place():
|
|
if (dtype == 'float16' and not _is_xpu_float16_supported()) or (
|
|
dtype == 'bfloat16' and not _is_xpu_bfloat16_supported()
|
|
):
|
|
if optimizers is None:
|
|
return models
|
|
else:
|
|
return models, optimizers
|
|
# For custom device:
|
|
if (
|
|
tracer._expected_place.is_custom_place()
|
|
and not _is_custom_device_bfloat16_supported()
|
|
and (dtype == 'bfloat16')
|
|
):
|
|
if optimizers is None:
|
|
return models
|
|
else:
|
|
return models, optimizers
|
|
# For gpu float16: Compute Capability should >= 7.
|
|
# For gpu bfloat16: Compute Capability should >= 8 & CUDA Version should >= 11.
|
|
if tracer._expected_place.is_gpu_place():
|
|
if (dtype == 'float16' and not _is_gpu_float16_supported()) or (
|
|
dtype == 'bfloat16' and not _is_gpu_bfloat16_supported()
|
|
):
|
|
if optimizers is None:
|
|
return models
|
|
else:
|
|
return models, optimizers
|
|
|
|
models_is_list = False
|
|
if isinstance(models, paddle.nn.Layer):
|
|
models_is_list = False
|
|
models = [models]
|
|
check_models(models)
|
|
elif isinstance(models, list):
|
|
check_models(models)
|
|
models_is_list = True
|
|
else:
|
|
raise TypeError(
|
|
"models must be either a single model or a list of models."
|
|
)
|
|
|
|
# initialize parameters of the model.
|
|
amp_initialize(models=models, dtype=dtype, excluded_layers=excluded_layers)
|
|
|
|
if optimizers is not None:
|
|
# check optimizers
|
|
optimizers_is_list = False
|
|
if _is_valid_optimizer(optimizers):
|
|
optimizers_is_list = False
|
|
optimizers = [optimizers]
|
|
check_optimizers(optimizers)
|
|
elif isinstance(optimizers, list):
|
|
check_optimizers(optimizers)
|
|
optimizers_is_list = True
|
|
else:
|
|
raise TypeError(
|
|
"optimizers must be either a single optimizer or a list of optimizers."
|
|
)
|
|
# support master_weight
|
|
use_multi_precision = master_weight is not False
|
|
for opt in optimizers:
|
|
_set_multi_precision(opt, use_multi_precision)
|
|
|
|
# support master_grad
|
|
if master_grad:
|
|
amp_global_state().use_master_grad = True
|
|
for idx in range(len(models)):
|
|
amp_global_state().model_parameters.extend(models[idx].parameters())
|
|
|
|
if save_dtype is not None:
|
|
if save_dtype not in ['float16', 'bfloat16', 'float32', 'float64']:
|
|
raise ValueError(
|
|
f"save_dtype can only be float16 float32 or float64, but your input save_dtype is {save_dtype}."
|
|
)
|
|
for idx in range(len(models)):
|
|
for layer in models[idx].sublayers(include_self=True):
|
|
layer.register_state_dict_hook(StateDictHook(save_dtype))
|
|
|
|
if models_is_list:
|
|
if optimizers is not None:
|
|
if optimizers_is_list:
|
|
return models, optimizers
|
|
else:
|
|
return models, optimizers[0]
|
|
else:
|
|
return models
|
|
else:
|
|
if optimizers is not None:
|
|
if optimizers_is_list:
|
|
return models[0], optimizers
|
|
else:
|
|
return models[0], optimizers[0]
|
|
else:
|
|
return models[0]
|
|
|
|
|
|
def autocast(
|
|
device_type: str | None,
|
|
dtype: _DTypeLiteral = 'float16',
|
|
enabled: bool = True,
|
|
cache_enabled: bool = True,
|
|
) -> AbstractContextManager:
|
|
"""
|
|
Create a context which enables auto-mixed-precision(AMP) of operators executed in dynamic graph mode.
|
|
If enabled, the input data type (float32, float16 or bfloat16) of each operator is decided
|
|
by autocast algorithm for better performance.
|
|
|
|
Commonly, it is used together with `GradScaler` and `decorator` to achieve Auto-Mixed-Precision in
|
|
imperative mode.
|
|
|
|
Args:
|
|
device_type(str, optional): Device type.But because the paddle does not distinguish between devices, this parameter does not work
|
|
enable(bool, optional): Enable auto-mixed-precision or not. Default is True.
|
|
dtype(str, optional): Whether to use 'float16' or 'bfloat16'. Default is 'float16'.
|
|
cache_enabled(bool, optional): whether to enable cache or not. Default is True. But this parameter is not used
|
|
|
|
Note:
|
|
paddle.cuda.amp.
|
|
|
|
Examples:
|
|
|
|
.. code-block:: pycon
|
|
|
|
>>> # doctest: +REQUIRES(env:GPU)
|
|
>>> import paddle
|
|
|
|
>>> conv2d = paddle.nn.Conv2D(3, 2, 3, bias_attr=False)
|
|
>>> data = paddle.rand([10, 3, 32, 32])
|
|
|
|
>>> with paddle.amp.auto_cast():
|
|
... conv = conv2d(data)
|
|
... print(conv.dtype)
|
|
>>> # doctest: +SKIP("This has diff in xdoctest env")
|
|
paddle.float16
|
|
>>> # doctest: -SKIP
|
|
|
|
>>> with paddle.amp.auto_cast(enable=False):
|
|
... conv = conv2d(data)
|
|
... print(conv.dtype)
|
|
>>> # doctest: +SKIP("This has diff in xdoctest env")
|
|
paddle.float32
|
|
>>> # doctest: -SKIP
|
|
|
|
"""
|
|
return auto_cast(enable=enabled, dtype=dtype)
|
|
|
|
|
|
def auto_cast(
|
|
enable: bool = True,
|
|
custom_white_list: _CustomList | None = None,
|
|
custom_black_list: _CustomList | None = None,
|
|
level: _AmpLevelLiteral = 'O1',
|
|
dtype: _DTypeLiteral = 'float16',
|
|
use_promote: bool = True,
|
|
) -> AbstractContextManager:
|
|
"""
|
|
Create a context which enables auto-mixed-precision(AMP) of operators executed in dynamic graph mode.
|
|
If enabled, the input data type (float32, float16 or bfloat16) of each operator is decided
|
|
by autocast algorithm for better performance.
|
|
|
|
Commonly, it is used together with `GradScaler` and `decorator` to achieve Auto-Mixed-Precision in
|
|
imperative mode.
|
|
|
|
Args:
|
|
enable(bool, optional): Enable auto-mixed-precision or not. Default is True.
|
|
custom_white_list(set|list|tuple|None, optional): A default white list is already set. Usually there is no need to set custom white list.
|
|
The set of ops should be considered numerically-safe and performance-critical. These ops will be converted to float16/bfloat16.
|
|
custom_black_list(set|list|tuple|None, optional): A default black list is already set. You can set a custom black list according to the model.
|
|
The set of ops are considered numerically-dangerous and whose effects may also be observed in downstream ops. These ops will not be
|
|
converted to float16/bfloat16.
|
|
level(str, optional): Auto mixed precision level. Accepted values are "O1", "O2" and "OD": At the O1 level, operators in the white list
|
|
will use float16/bfloat16 inputs for calculations, and operators in the black list will use float32 inputs for calculations. At the O2
|
|
level, model's parameters will be casted to float16/bfloat16 by using `decorator`, and operators that have all float16/bfloat16 inputs
|
|
will be converted to float16/bfloat16, and that have any float32 input will be converted to float32. For the OD level, operators in
|
|
default white list will compute in float16/bfloat16, and the others will compute in float32. Default is O1.
|
|
dtype(str, optional): Whether to use 'float16' or 'bfloat16'. Default is 'float16'.
|
|
use_promote(bool, optional): Whether to promotes to fp32 when op has any float32 inputs. It is only supported when amp level is O2. Default is True.
|
|
|
|
Examples:
|
|
|
|
.. code-block:: pycon
|
|
|
|
>>> # doctest: +REQUIRES(env:GPU)
|
|
>>> import paddle
|
|
|
|
>>> conv2d = paddle.nn.Conv2D(3, 2, 3, bias_attr=False)
|
|
>>> data = paddle.rand([10, 3, 32, 32])
|
|
|
|
>>> with paddle.amp.auto_cast():
|
|
... conv = conv2d(data)
|
|
... print(conv.dtype)
|
|
>>> # doctest: +SKIP("This has diff in xdoctest env")
|
|
paddle.float16
|
|
>>> # doctest: -SKIP
|
|
|
|
>>> with paddle.amp.auto_cast(enable=False):
|
|
... conv = conv2d(data)
|
|
... print(conv.dtype)
|
|
>>> # doctest: +SKIP("This has diff in xdoctest env")
|
|
paddle.float32
|
|
>>> # doctest: -SKIP
|
|
|
|
>>> with paddle.amp.auto_cast(custom_black_list={'conv2d'}):
|
|
... conv = conv2d(data)
|
|
... print(conv.dtype)
|
|
>>> # doctest: +SKIP("This has diff in xdoctest env")
|
|
paddle.float32
|
|
>>> # doctest: -SKIP
|
|
|
|
>>> a = paddle.rand([2, 3])
|
|
>>> b = paddle.rand([2, 3])
|
|
>>> with paddle.amp.auto_cast(custom_white_list={'elementwise_add'}):
|
|
... c = a + b
|
|
... print(c.dtype)
|
|
>>> # doctest: +SKIP("This has diff in xdoctest env")
|
|
paddle.float16
|
|
>>> # doctest: -SKIP
|
|
|
|
>>> with paddle.amp.auto_cast(custom_white_list={'elementwise_add'}, level='O2'):
|
|
... d = a + b
|
|
... print(d.dtype)
|
|
>>> # doctest: +SKIP("This has diff in xdoctest env")
|
|
paddle.float16
|
|
>>> # doctest: -SKIP
|
|
|
|
"""
|
|
return amp_guard(
|
|
enable, custom_white_list, custom_black_list, level, dtype, use_promote
|
|
)
|
|
|
|
|
|
@overload
|
|
def decorate(
|
|
models: _ModelsT,
|
|
optimizers: _OptimizersT,
|
|
level: _AmpLevelLiteral = ...,
|
|
dtype: _DTypeLiteral = ...,
|
|
master_weight: bool | None = ...,
|
|
save_dtype: _DTypeLiteral | None = ...,
|
|
master_grad: bool = ...,
|
|
excluded_layers: (
|
|
Layer | list[Layer | type[Layer]] | type[Layer] | None
|
|
) = ...,
|
|
) -> tuple[_ModelsT, _OptimizersT]: ...
|
|
|
|
|
|
@overload
|
|
def decorate(
|
|
models: _ModelsT,
|
|
optimizers: None = ...,
|
|
level: _AmpLevelLiteral = ...,
|
|
dtype: _DTypeLiteral = ...,
|
|
master_weight: bool | None = ...,
|
|
save_dtype: _DTypeLiteral | None = ...,
|
|
master_grad: bool = ...,
|
|
excluded_layers: (
|
|
Layer | list[Layer | type[Layer]] | type[Layer] | None
|
|
) = ...,
|
|
) -> _ModelsT: ...
|
|
|
|
|
|
def decorate(
|
|
models: _ModelsT,
|
|
optimizers: _OptimizersT | None = None,
|
|
level: _AmpLevelLiteral = 'O1',
|
|
dtype: _DTypeLiteral = 'float16',
|
|
master_weight: bool | None = None,
|
|
save_dtype: _DTypeLiteral | None = None,
|
|
master_grad: bool = False,
|
|
excluded_layers: (
|
|
Layer | list[Layer | type[Layer]] | type[Layer] | None
|
|
) = None,
|
|
) -> tuple[_ModelsT, _OptimizersT] | _ModelsT:
|
|
"""
|
|
Decorate models and optimizers for auto-mixed-precision. When level is O1(amp), the decorate will do nothing.
|
|
When level is O2(pure float16/bfloat16), the decorate will cast all parameters of models to float16/bfloat16, except BatchNorm, InstanceNorm and LayerNorm.
|
|
|
|
Commonly, it is used together with `auto_cast` to achieve Pure float16/bfloat16 in imperative mode.
|
|
|
|
Args:
|
|
models(Layer|list of Layer): The defined models by user, models must be either a single model or a list of models. Default is None.
|
|
optimizers(Optimizer|list of Optimizer|None, optional): The defined optimizers by user, optimizers must be either a single optimizer or a list of optimizers. Default is None.
|
|
level(str, optional): Auto mixed precision level. Accepted values are 'O1' and 'O2': O1 represent mixed precision, the decorator will do nothing;
|
|
O2 represent Pure float16/bfloat16, the decorator will cast all parameters of models to float16/bfloat16, except BatchNorm, InstanceNorm and LayerNorm. Default is O1(amp)
|
|
dtype(str, optional): Whether to use 'float16' or 'bfloat16'. Default is 'float16'.
|
|
master_weight(bool, optional): For level='O2', whether to use multi-precision during weight updating. If master_weight is None, in O2 level optimizer will use multi-precision. Default is None.
|
|
save_dtype(str|None, optional): The save model parameter dtype when use `paddle.save` or `paddle.jit.save`,it should be float16, bfloat16, float32, float64 or None.
|
|
The save_dtype will not change model parameters dtype, it just change the state_dict dtype. When save_dtype is None, the save dtype is same as model dtype. Default is None.
|
|
master_grad(bool, optional): For level='O2', whether to use float32 weight gradients for calculations such as gradient clipping, weight decay, and weight updates. If master_grad is enabled, the weight
|
|
gradients will be float32 dtype after the back propagation. Default is False, there is only float16 weight gradients.
|
|
excluded_layers(Layer|list of Layer, optional): Specify the layers not to be decorated. The weights of these layers will always keep float32 when level is O2. `excluded_layers` can be specified as
|
|
an Layer instance/type or a list of Layer instances/types. Default is None, the weights of the whole model will be casted to float16 or bfloat16.
|
|
|
|
Examples:
|
|
|
|
.. code-block:: pycon
|
|
|
|
>>> # doctest: +REQUIRES(env:GPU)
|
|
>>> # Demo1: single model and optimizer:
|
|
>>> import paddle
|
|
>>> paddle.device.set_device('gpu')
|
|
|
|
>>> model = paddle.nn.Conv2D(3, 2, 3, bias_attr=False)
|
|
>>> optimizer = paddle.optimizer.SGD(parameters=model.parameters())
|
|
|
|
>>> model, optimizer = paddle.amp.decorate(models=model, optimizers=optimizer, level='O2')
|
|
|
|
>>> data = paddle.rand([10, 3, 32, 32])
|
|
|
|
>>> with paddle.amp.auto_cast(enable=True, custom_white_list=None, custom_black_list=None, level='O2'):
|
|
... output = model(data)
|
|
... print(output.dtype)
|
|
paddle.float16
|
|
|
|
>>> # Demo2: multi models and optimizers:
|
|
>>> model2 = paddle.nn.Conv2D(3, 2, 3, bias_attr=False)
|
|
>>> optimizer2 = paddle.optimizer.Adam(parameters=model2.parameters())
|
|
|
|
>>> models, optimizers = paddle.amp.decorate(models=[model, model2], optimizers=[optimizer, optimizer2], level='O2')
|
|
|
|
>>> data = paddle.rand([10, 3, 32, 32])
|
|
|
|
>>> with paddle.amp.auto_cast(enable=True, custom_white_list=None, custom_black_list=None, level='O2'):
|
|
... output = models[0](data)
|
|
... output2 = models[1](data)
|
|
... print(output.dtype)
|
|
... print(output2.dtype)
|
|
paddle.float16
|
|
paddle.float16
|
|
|
|
>>> # Demo3: optimizers is None:
|
|
>>> model3 = paddle.nn.Conv2D(3, 2, 3, bias_attr=False)
|
|
>>> optimizer3 = paddle.optimizer.Adam(parameters=model3.parameters())
|
|
|
|
>>> model = paddle.amp.decorate(models=model3, level='O2')
|
|
|
|
>>> data = paddle.rand([10, 3, 32, 32])
|
|
|
|
>>> with paddle.amp.auto_cast(enable=True, custom_white_list=None, custom_black_list=None, level='O2'):
|
|
... output = model(data)
|
|
... print(output.dtype)
|
|
paddle.float16
|
|
|
|
"""
|
|
|
|
if paddle.framework.in_pir_mode():
|
|
assert not isinstance(models, (list, tuple))
|
|
assert not isinstance(optimizers, (list, tuple))
|
|
amp_global_state().use_master_grad = master_grad
|
|
if level in ['O0', 'OD', 'O1']:
|
|
if optimizers is None:
|
|
return models
|
|
else:
|
|
optimizers = OptimizerWithMixedPrecision(
|
|
optimizer=optimizers,
|
|
amp_lists=None,
|
|
level=level,
|
|
dtype=dtype,
|
|
init_loss_scaling=1.0,
|
|
incr_every_n_steps=None,
|
|
decr_every_n_nan_or_inf=None,
|
|
incr_ratio=None,
|
|
decr_ratio=None,
|
|
use_dynamic_loss_scaling=False,
|
|
use_amp_guard=None,
|
|
use_master_grad=master_grad,
|
|
use_promote=None,
|
|
)
|
|
return models, optimizers
|
|
elif level == 'O2':
|
|
amp_initialize(
|
|
models=[models], dtype=dtype, excluded_layers=excluded_layers
|
|
)
|
|
use_multi_precision = master_weight is not False
|
|
_set_multi_precision(optimizers, use_multi_precision)
|
|
if optimizers is None:
|
|
return models
|
|
else:
|
|
optimizers = OptimizerWithMixedPrecision(
|
|
optimizer=optimizers,
|
|
amp_lists=None,
|
|
level=level,
|
|
dtype=dtype,
|
|
init_loss_scaling=1.0,
|
|
incr_every_n_steps=None,
|
|
decr_every_n_nan_or_inf=None,
|
|
incr_ratio=None,
|
|
decr_ratio=None,
|
|
use_dynamic_loss_scaling=False,
|
|
use_amp_guard=None,
|
|
use_master_grad=master_grad,
|
|
use_promote=None,
|
|
)
|
|
return models, optimizers
|
|
else:
|
|
raise ValueError("level should be O0, OD, O1 or O2.")
|
|
else:
|
|
return amp_decorate(
|
|
models,
|
|
optimizers,
|
|
level,
|
|
dtype,
|
|
master_weight,
|
|
save_dtype,
|
|
master_grad,
|
|
excluded_layers,
|
|
)
|
|
|
|
|
|
def is_autocast_enabled(device_type: PlaceLike | None = None) -> bool:
|
|
"""
|
|
Check whether auto-mixed-precision is enabled in the current context.
|
|
|
|
Args:
|
|
device_type (PlaceLike, optional): The device type to check. This argument is ignored for all devices sharing the same AMP state in paddlepaddle.
|
|
|
|
Returns:
|
|
bool: True if auto-mixed-precision is enabled, False otherwise.
|
|
|
|
Examples:
|
|
.. code-block:: pycon
|
|
|
|
>>> # doctest: +REQUIRES(env:GPU)
|
|
>>> # Demo1: Check if auto-mixed-precision is enabled by default
|
|
>>> import paddle
|
|
>>> paddle.device.set_device('gpu')
|
|
>>> print(paddle.is_autocast_enabled())
|
|
False
|
|
|
|
>>> # Demo2: Enable auto-mixed-precision and check again
|
|
>>> with paddle.amp.auto_cast():
|
|
... print(paddle.is_autocast_enabled())
|
|
True
|
|
"""
|
|
if in_pir_mode():
|
|
amp_attrs = core._get_amp_attrs()
|
|
return amp_attrs._amp_level != AMP_LEVEL.O0
|
|
else:
|
|
tracer = _dygraph_tracer()
|
|
if tracer:
|
|
return tracer._amp_level != core.AmpLevel.O0
|
|
return False
|
|
|
|
|
|
def get_autocast_dtype(device_type: PlaceLike | None = None) -> _DTypeLiteral:
|
|
"""
|
|
Get the auto-mixed-precision dtype in the current context if autocast is enabled else default AMP dtype(float16).
|
|
|
|
Args:
|
|
device_type (PlaceLike, optional): The device type to check. This argument is ignored for all devices sharing the same AMP state in paddlepaddle.
|
|
|
|
Returns:
|
|
_DTypeLiteral: The current AMP dtype.
|
|
|
|
Examples:
|
|
.. code-block:: pycon
|
|
|
|
>>> # doctest: +REQUIRES(env:GPU)
|
|
>>> # Demo1: Get default auto-mixed-precision dtype
|
|
>>> import paddle
|
|
>>> paddle.device.set_device('gpu')
|
|
>>> print(paddle.get_autocast_dtype())
|
|
float16
|
|
|
|
>>> # Demo2: Enable auto-mixed-precision and get the dtype
|
|
>>> with paddle.amp.auto_cast():
|
|
... print(paddle.get_autocast_dtype())
|
|
float16
|
|
"""
|
|
if not is_autocast_enabled():
|
|
return "float16"
|
|
if in_pir_mode():
|
|
amp_attrs = core._get_amp_attrs()
|
|
return amp_attrs._amp_dtype
|
|
else:
|
|
tracer = _dygraph_tracer()
|
|
return tracer._amp_dtype
|