1436 lines
56 KiB
Python
1436 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 warnings
|
|
from collections import defaultdict
|
|
from enum import Enum
|
|
from typing import (
|
|
TYPE_CHECKING,
|
|
Any,
|
|
TypedDict,
|
|
overload,
|
|
)
|
|
|
|
import numpy as np
|
|
|
|
import paddle
|
|
from paddle import _C_ops, _legacy_C_ops
|
|
from paddle.base import core, unique_name
|
|
from paddle.base.data_feeder import check_type
|
|
from paddle.base.framework import Operator, _dygraph_tracer, in_pir_mode
|
|
from paddle.framework import in_dynamic_mode
|
|
from paddle.utils.decorator_utils import grad_scaler_decorator
|
|
|
|
from .auto_cast import amp_global_state
|
|
|
|
if TYPE_CHECKING:
|
|
from paddle import Tensor
|
|
from paddle.optimizer.optimizer import Optimizer
|
|
from paddle.static.amp.decorator import OptimizerWithMixedPrecision
|
|
|
|
class _ScaleStateDict(TypedDict):
|
|
scale: Tensor
|
|
incr_ratio: float
|
|
decr_ratio: float
|
|
incr_every_n_steps: int
|
|
decr_every_n_nan_or_inf: int
|
|
incr_count: int
|
|
decr_count: int
|
|
use_dynamic_loss_scaling: bool
|
|
|
|
|
|
class OptimizerState(Enum):
|
|
INIT = 0
|
|
UNSCALED = 1
|
|
STEPPED = 2
|
|
|
|
|
|
def _refresh_optimizer_state():
|
|
return {"state": OptimizerState.INIT}
|
|
|
|
|
|
class AmpScaler:
|
|
"""
|
|
AmpScaler is used for Auto-Mixed-Precision training/inferring in imperative
|
|
mode. It controls the scaling of loss, helps avoiding numerical overflow.
|
|
The object of this class has seventeen methods `scale()`, `unscale_()`, `minimize()` and `get`/`set` api of parameters.
|
|
|
|
`scale()` is used to multiply the loss by a scale ratio.
|
|
`unscale_()` is used to unscale the gradients of parameters, multiplies the gradients of parameters by 1/(scale ratio)
|
|
`minimize()` is similar as `optimizer.minimize()`, performs parameters updating, and it will update the loss_scaling.
|
|
|
|
Commonly, it is used together with `amp_guard` to achieve Auto-Mixed-Precision in
|
|
imperative mode.
|
|
|
|
Args:
|
|
enable(bool, optional): Enable loss scaling or not. Default is True.
|
|
init_loss_scaling (float, optional): The initial loss scaling factor. Default is 2**15.
|
|
incr_ratio(float, optional): The multiplier to use when increasing the loss
|
|
scaling. Default is 2.0.
|
|
decr_ratio(float, optional): The less-than-one-multiplier to use when decreasing
|
|
the loss scaling. Default is 0.5.
|
|
incr_every_n_steps(int, optional): Increases loss scaling every n consecutive
|
|
steps with finite gradients. Default is 1000.
|
|
decr_every_n_nan_or_inf(int, optional): Decreases loss scaling every n
|
|
accumulated steps with nan or inf gradients. Default is 2.
|
|
use_dynamic_loss_scaling(bool, optional): Whether to use dynamic loss scaling. If False, fixed loss_scaling is used. If True, the loss scaling is updated dynamically. Default is True.
|
|
Returns:
|
|
An AmpScaler object.
|
|
|
|
Examples:
|
|
|
|
.. code-block:: pycon
|
|
|
|
>>> import numpy as np
|
|
>>> import paddle
|
|
|
|
>>> data = np.random.uniform(-1, 1, [10, 3, 32, 32]).astype('float32')
|
|
>>> model = paddle.nn.Conv2D(3, 2, 3)
|
|
>>> optimizer = paddle.optimizer.SGD(
|
|
... learning_rate=0.01,
|
|
... parameters=model.parameters(),
|
|
... )
|
|
>>> scaler = paddle.amp.AmpScaler(init_loss_scaling=1024)
|
|
>>> data = paddle.to_tensor(data)
|
|
>>> with paddle.amp.amp_guard():
|
|
... conv = model(data)
|
|
... loss = paddle.mean(conv)
|
|
... scaled = scaler.scale(loss)
|
|
... scaled.backward()
|
|
... scaler.minimize(optimizer, scaled)
|
|
"""
|
|
|
|
def __init__(
|
|
self,
|
|
enable: bool = True,
|
|
init_loss_scaling: float = 2.0**15,
|
|
incr_ratio: float = 2.0,
|
|
decr_ratio: float = 0.5,
|
|
incr_every_n_steps: int = 1000,
|
|
decr_every_n_nan_or_inf: int = 1,
|
|
use_dynamic_loss_scaling: bool = True,
|
|
) -> None:
|
|
if in_dynamic_mode():
|
|
tracer = _dygraph_tracer()
|
|
if not tracer:
|
|
raise ValueError(
|
|
"current_tracer is None, maybe it is not in imperative mode."
|
|
)
|
|
|
|
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'AmpScaler can only be enabled on CUDAPlace, XPUPlace and CustomPlace, current place is {tracer._expected_place}, so it makes no effect.'
|
|
)
|
|
enable = False
|
|
|
|
self._enable = enable
|
|
self._use_dynamic_loss_scaling = False
|
|
self._init_loss_scaling = 1.0
|
|
self._scale = None
|
|
|
|
if self._enable:
|
|
assert incr_ratio > 1.0, "The incr_ratio must be > 1.0."
|
|
assert decr_ratio < 1.0, "The decr_ratio must be < 1.0."
|
|
|
|
self._init_loss_scaling = init_loss_scaling
|
|
self._incr_ratio = incr_ratio
|
|
self._decr_ratio = decr_ratio
|
|
self._incr_every_n_steps = incr_every_n_steps
|
|
self._decr_every_n_nan_or_inf = decr_every_n_nan_or_inf
|
|
self._incr_count = 0
|
|
self._decr_count = 0
|
|
self._use_dynamic_loss_scaling = use_dynamic_loss_scaling
|
|
|
|
if in_pir_mode():
|
|
self._scale = paddle.pir.core.create_persistable_value(
|
|
dtype='float32',
|
|
shape=[1],
|
|
name=unique_name.generate("loss_scaling"),
|
|
initializer=paddle.nn.initializer.ConstantInitializer(
|
|
value=self._init_loss_scaling
|
|
),
|
|
)
|
|
else:
|
|
self._found_inf = paddle.to_tensor(
|
|
np.array([0]).astype(np.bool_)
|
|
)
|
|
self._temp_found_inf_value_false = paddle.to_tensor(
|
|
np.array([0]).astype(np.bool_)
|
|
)
|
|
self._temp_found_inf_fp16 = paddle.to_tensor(
|
|
np.array([0]).astype(np.bool_)
|
|
)
|
|
self._temp_found_inf_bf16 = paddle.to_tensor(
|
|
np.array([0]).astype(np.bool_)
|
|
)
|
|
self._temp_found_inf_fp32 = paddle.to_tensor(
|
|
np.array([0]).astype(np.bool_)
|
|
)
|
|
self._scale = paddle.to_tensor(
|
|
np.array([self._init_loss_scaling]).astype(np.float32)
|
|
)
|
|
self._cache_found_inf = None
|
|
self._optimizer_states = defaultdict(_refresh_optimizer_state)
|
|
|
|
def scale(self, var: Tensor) -> Tensor:
|
|
"""
|
|
Multiplies a Tensor by the scale factor and returns scaled outputs.
|
|
If this instance of :class:`AmpScaler` is not enabled, output are returned unmodified.
|
|
|
|
Args:
|
|
var (Tensor): The Tensor to scale.
|
|
Returns:
|
|
The scaled Tensor or original Tensor.
|
|
|
|
Examples:
|
|
|
|
.. code-block:: pycon
|
|
|
|
>>> import numpy as np
|
|
>>> import paddle
|
|
>>> data = np.random.uniform(-1, 1, [10, 3, 32, 32]).astype('float32')
|
|
>>> model = paddle.nn.Conv2D(3, 2, 3)
|
|
>>> optimizer = paddle.optimizer.SGD(
|
|
... learning_rate=0.01,
|
|
... parameters=model.parameters(),
|
|
... )
|
|
>>> scaler = paddle.amp.AmpScaler(init_loss_scaling=1024)
|
|
>>> data = paddle.to_tensor(data)
|
|
>>> with paddle.amp.amp_guard():
|
|
... conv = model(data)
|
|
... loss = paddle.mean(conv)
|
|
... scaled = scaler.scale(loss)
|
|
... scaled.backward()
|
|
... scaler.minimize(optimizer, scaled)
|
|
"""
|
|
check_type(
|
|
var,
|
|
"var",
|
|
(paddle.Tensor, paddle.pir.Value),
|
|
'AmpScaler.scale()',
|
|
)
|
|
|
|
if (
|
|
self._enable
|
|
and amp_global_state().amp_dtype != 'float16'
|
|
and self._use_dynamic_loss_scaling
|
|
):
|
|
self._enable = False
|
|
self._use_dynamic_loss_scaling = False
|
|
self._init_loss_scaling = 1.0
|
|
warnings.warn(
|
|
f'It is not recommended to use dynamic loss scaling for {amp_global_state().amp_dtype}, so GradScaler is disable by default.'
|
|
)
|
|
|
|
if in_pir_mode():
|
|
if var.dtype != core.DataType.FLOAT32:
|
|
var = var.astype('float32')
|
|
if not self._use_dynamic_loss_scaling:
|
|
return var
|
|
scale_out = paddle._C_ops.multiply(var, self._scale)
|
|
multiply_op = scale_out.get_defining_op()
|
|
src_var_op = var.get_defining_op()
|
|
if multiply_op.dist_attr and src_var_op.dist_attr:
|
|
multiply_op.dist_attr = (
|
|
paddle.base.libpaddle.pir.create_op_dist_attribute(
|
|
multiply_op.dist_attr.process_mesh,
|
|
multiply_op.dist_attr.operands(),
|
|
multiply_op.dist_attr.results(),
|
|
src_var_op.dist_attr.chunk_id,
|
|
)
|
|
)
|
|
return scale_out
|
|
|
|
# NOTE(lizhiyu): We hack here to avoid changing the `dist_attr` of `self._scale` of 'no-calculation-rank'
|
|
if not self._enable or not var._is_initialized():
|
|
return var
|
|
|
|
return var * self._scale
|
|
|
|
def minimize(
|
|
self,
|
|
optimizer: Optimizer | OptimizerWithMixedPrecision,
|
|
*args: Any,
|
|
**kwargs: Any,
|
|
) -> tuple[list[Operator], list[tuple[Tensor, Tensor]]]:
|
|
"""
|
|
This function is similar as `Optimizer.minimize()`, which performs parameters updating.
|
|
|
|
If the scaled gradients of parameters contains NAN or INF, the parameters updating is skipped.
|
|
Otherwise, if `unscale_()` has not been called, it first unscales the scaled gradients of parameters, then updates the parameters.
|
|
|
|
Finally, the loss scaling ratio is updated.
|
|
|
|
Args:
|
|
optimizer(Optimizer): The optimizer used to update parameters.
|
|
args: Arguments, which will be forward to `Optimizer.minimize()`.
|
|
kwargs: Keyword arguments, which will be forward to `Optimizer.minimize()`.
|
|
|
|
Examples:
|
|
|
|
.. code-block:: pycon
|
|
|
|
>>> import numpy as np
|
|
>>> import paddle
|
|
>>> data = np.random.uniform(-1, 1, [10, 3, 32, 32]).astype('float32')
|
|
>>> model = paddle.nn.Conv2D(3, 2, 3)
|
|
>>> optimizer = paddle.optimizer.SGD(
|
|
... learning_rate=0.01,
|
|
... parameters=model.parameters(),
|
|
... )
|
|
>>> scaler = paddle.amp.AmpScaler(init_loss_scaling=1024)
|
|
>>> data = paddle.to_tensor(data)
|
|
>>> with paddle.amp.amp_guard():
|
|
... conv = model(data)
|
|
... loss = paddle.mean(conv)
|
|
... scaled = scaler.scale(loss)
|
|
... scaled.backward()
|
|
... scaler.minimize(optimizer, scaled)
|
|
"""
|
|
|
|
if in_pir_mode():
|
|
assert isinstance(
|
|
optimizer,
|
|
paddle.static.amp.decorator.OptimizerWithMixedPrecision,
|
|
)
|
|
optimizer._use_dynamic_loss_scaling = self._use_dynamic_loss_scaling
|
|
optimizer._init_loss_scaling = self._init_loss_scaling
|
|
optimizer._loss_scaling = self._scale
|
|
optimizer._scaled_loss = args[0]
|
|
if self._use_dynamic_loss_scaling:
|
|
optimizer._incr_every_n_steps = self._incr_every_n_steps
|
|
optimizer._decr_every_n_nan_or_inf = (
|
|
self._decr_every_n_nan_or_inf
|
|
)
|
|
optimizer._incr_ratio = self._incr_ratio
|
|
optimizer._decr_ratio = self._decr_ratio
|
|
optimizer._num_good_steps = None
|
|
optimizer._num_bad_steps = None
|
|
return optimizer.minimize(*args, **kwargs)
|
|
|
|
if not self._enable:
|
|
return optimizer.minimize(*args, **kwargs)
|
|
|
|
optimizer_state = self._optimizer_states[id(optimizer)]
|
|
|
|
# unscale the grad
|
|
if optimizer_state["state"] is OptimizerState.INIT:
|
|
self._unscale(optimizer)
|
|
|
|
optimize_ops, params_grads = (None, None)
|
|
|
|
if hasattr(optimizer, "_set_auxiliary_var"):
|
|
optimizer._set_auxiliary_var('found_inf', self._found_inf)
|
|
optimize_ops, params_grads = optimizer.minimize(*args, **kwargs)
|
|
# TODO: Fix to _cache_found_inf after PaddleNLP update
|
|
self._cache_found_inf = optimizer._get_auxiliary_var('found_inf')
|
|
else:
|
|
if self._found_inf:
|
|
self._cache_found_inf = True
|
|
else:
|
|
optimize_ops, params_grads = optimizer.minimize(*args, **kwargs)
|
|
self._cache_found_inf = False
|
|
|
|
if self._use_dynamic_loss_scaling:
|
|
# update the scale
|
|
self._update()
|
|
|
|
self._optimizer_states = defaultdict(_refresh_optimizer_state)
|
|
|
|
return optimize_ops, params_grads
|
|
|
|
def _unscale(self, optimizer):
|
|
"""
|
|
Unscale the gradients of parameters, multiplies the gradients of parameters by 1/(loss scaling ratio).
|
|
If this instance of :class:`GradScaler` is not enabled, output are returned unmodified.
|
|
Args:
|
|
optimizer(Optimizer): The optimizer used to update parameters.
|
|
Returns:
|
|
The unscaled parameters or original parameters.
|
|
"""
|
|
if not self._enable:
|
|
return
|
|
|
|
optimizer_state = self._optimizer_states[id(optimizer)]
|
|
|
|
if optimizer_state["state"] is OptimizerState.UNSCALED:
|
|
raise RuntimeError(
|
|
"unscale_() has already been called on this optimizer since the last update()."
|
|
)
|
|
elif optimizer_state["state"] is OptimizerState.STEPPED:
|
|
raise RuntimeError("unscale_() is being called after step().")
|
|
|
|
if getattr(optimizer, '_param_groups', None) and isinstance(
|
|
optimizer._param_groups[0], dict
|
|
):
|
|
param_grads = []
|
|
param_grads_fp16 = []
|
|
param_grads_bf16 = []
|
|
param_grads_fp32 = []
|
|
for group in optimizer._param_groups:
|
|
for param in group['params']:
|
|
if param._grad_ivar() is not None:
|
|
param_grads.append(param._grad_ivar())
|
|
if param._grad_ivar().dtype == paddle.float16:
|
|
param_grads_fp16.append(param._grad_ivar())
|
|
elif param._grad_ivar().dtype == paddle.bfloat16:
|
|
param_grads_bf16.append(param._grad_ivar())
|
|
else:
|
|
param_grads_fp32.append(param._grad_ivar())
|
|
else:
|
|
if in_dynamic_mode():
|
|
# It is very time-consuming to call c++ functions in a loop on the python side.
|
|
# We put this part of the code on the c++ side to improve the speed in eager mode.
|
|
(
|
|
param_grads_fp16,
|
|
param_grads_bf16,
|
|
param_grads_fp32,
|
|
) = core.eager.get_grads_lists(optimizer._parameter_list)
|
|
else:
|
|
# Keep the original code to support legacy mode.
|
|
# Delete the else branch when the legacy mode exits.
|
|
param_grads = [
|
|
param._grad_ivar()
|
|
for param in optimizer._parameter_list
|
|
if param._grad_ivar() is not None
|
|
]
|
|
param_grads_fp16 = [
|
|
param
|
|
for param in param_grads
|
|
if param.dtype == paddle.float16
|
|
]
|
|
param_grads_bf16 = [
|
|
param
|
|
for param in param_grads
|
|
if param.dtype == paddle.bfloat16
|
|
]
|
|
param_grads_fp32 = [
|
|
param
|
|
for param in param_grads
|
|
if param.dtype == paddle.float32
|
|
]
|
|
self._found_inf = self._temp_found_inf_value_false
|
|
if len(param_grads_fp16):
|
|
_legacy_C_ops.check_finite_and_unscale(
|
|
param_grads_fp16,
|
|
self._scale,
|
|
param_grads_fp16,
|
|
self._temp_found_inf_fp16,
|
|
)
|
|
self._found_inf = _C_ops.bitwise_or(
|
|
self._found_inf, self._temp_found_inf_fp16
|
|
)
|
|
if len(param_grads_bf16):
|
|
_legacy_C_ops.check_finite_and_unscale(
|
|
param_grads_bf16,
|
|
self._scale,
|
|
param_grads_bf16,
|
|
self._temp_found_inf_bf16,
|
|
)
|
|
self._found_inf = _C_ops.bitwise_or(
|
|
self._found_inf, self._temp_found_inf_bf16
|
|
)
|
|
if len(param_grads_fp32):
|
|
_legacy_C_ops.check_finite_and_unscale(
|
|
param_grads_fp32,
|
|
self._scale,
|
|
param_grads_fp32,
|
|
self._temp_found_inf_fp32,
|
|
)
|
|
self._found_inf = _C_ops.bitwise_or(
|
|
self._found_inf, self._temp_found_inf_fp32
|
|
)
|
|
|
|
optimizer_state["state"] = OptimizerState.UNSCALED
|
|
|
|
def _update(self):
|
|
"""
|
|
Updates the loss_scaling.
|
|
"""
|
|
if not self._enable:
|
|
return
|
|
|
|
if self._cache_found_inf:
|
|
self._incr_count = 0
|
|
self._decr_count = self._decr_count + 1
|
|
if self._decr_count == self._decr_every_n_nan_or_inf:
|
|
print(
|
|
f'Found inf or nan, current scale is: {float(self._scale)}, decrease to: {float(self._scale)}*{float(self._decr_ratio)}'
|
|
)
|
|
self._scale = self._scale * self._decr_ratio
|
|
self._decr_count = 0
|
|
else:
|
|
self._decr_count = 0
|
|
self._incr_count = self._incr_count + 1
|
|
if self._incr_count == self._incr_every_n_steps:
|
|
self._scale = self._scale * self._incr_ratio
|
|
self._incr_count = 0
|
|
|
|
return
|
|
|
|
def is_enable(self) -> bool:
|
|
"""
|
|
Enable loss scaling or not.
|
|
|
|
Returns:
|
|
bool: enable loss scaling return True else return False.
|
|
"""
|
|
return self._enable
|
|
|
|
def is_use_dynamic_loss_scaling(self) -> bool:
|
|
"""
|
|
Whether to use dynamic loss scaling.
|
|
|
|
Returns:
|
|
bool: if fixed loss_scaling is used return False, if the loss scaling is updated dynamically return true.
|
|
"""
|
|
return self._use_dynamic_loss_scaling
|
|
|
|
def get_init_loss_scaling(self) -> float:
|
|
"""
|
|
Return the initial loss scaling factor.
|
|
|
|
Returns:
|
|
float: the initial loss scaling factor.
|
|
"""
|
|
return self._init_loss_scaling
|
|
|
|
def set_init_loss_scaling(self, new_init_loss_scaling: int) -> None:
|
|
"""
|
|
Set the initial loss scaling factor by `new_init_loss_scaling`.
|
|
|
|
Args:
|
|
new_init_loss_scaling(int): The new_init_loss_scaling used to update initial loss scaling factor.s
|
|
"""
|
|
self._init_loss_scaling = new_init_loss_scaling
|
|
self._scale = paddle.to_tensor(
|
|
np.array([self._init_loss_scaling]).astype(np.float32)
|
|
)
|
|
|
|
def get_incr_ratio(self) -> float:
|
|
"""
|
|
Return the multiplier to use when increasing the loss scaling.
|
|
|
|
Returns:
|
|
float: the multiplier to use when increasing the loss scaling.
|
|
"""
|
|
return self._incr_ratio
|
|
|
|
def set_incr_ratio(self, new_incr_ratio: float) -> None:
|
|
"""
|
|
Set the multiplier to use when increasing the loss scaling by `new_incr_ratio`, `new_incr_ratio` should > 1.0.
|
|
|
|
Args:
|
|
new_incr_ratio(float): The new_incr_ratio used to update the multiplier to use when increasing the loss scaling.
|
|
"""
|
|
assert new_incr_ratio > 1.0, "The new_incr_ratio must be > 1.0."
|
|
self._incr_ratio = new_incr_ratio
|
|
|
|
def get_decr_ratio(self) -> float:
|
|
"""
|
|
Get the less-than-one-multiplier to use when decreasing the loss scaling.
|
|
|
|
Returns:
|
|
float: the less-than-one-multiplier to use when decreasing the loss scaling.
|
|
"""
|
|
return self._decr_ratio
|
|
|
|
def set_decr_ratio(self, new_decr_ratio: float) -> None:
|
|
"""
|
|
Set the less-than-one-multiplier to use when decreasing the loss scaling by `new_incr_ratio`, `new_decr_ratio` should < 1.0.
|
|
|
|
Args:
|
|
new_decr_ratio(float): The new_decr_ratio used to update the less-than-one-multiplier to use when decreasing the loss scaling.
|
|
"""
|
|
assert new_decr_ratio < 1.0, "The new_decr_ratio must be < 1.0."
|
|
self._decr_ratio = new_decr_ratio
|
|
|
|
def get_incr_every_n_steps(self) -> int:
|
|
"""
|
|
Return the num `n`, `n` represent increases loss scaling every `n` consecutive steps with finite gradients.
|
|
|
|
Returns:
|
|
int: the num `n`, `n` represent increases loss scaling every `n` consecutive steps with finite gradients.
|
|
"""
|
|
return self._incr_every_n_steps
|
|
|
|
def set_incr_every_n_steps(self, new_incr_every_n_steps: int) -> None:
|
|
"""
|
|
Set the num `n` by `new_incr_every_n_steps`, `n` represent increases loss scaling every `n` consecutive steps with finite gradients.
|
|
|
|
Args:
|
|
new_incr_every_n_steps(int): The new_incr_every_n_steps used to update the num `n`, `n` represent increases loss scaling every `n` consecutive steps with finite gradients.
|
|
"""
|
|
self._incr_every_n_steps = new_incr_every_n_steps
|
|
|
|
def get_decr_every_n_nan_or_inf(self) -> int:
|
|
"""
|
|
Return the num `n`, `n` represent decreases loss scaling every `n` accumulated steps with nan or inf gradients.
|
|
|
|
Returns:
|
|
int: the num `n`, `n` represent decreases loss scaling every `n` accumulated steps with nan or inf gradients.
|
|
"""
|
|
return self._decr_every_n_nan_or_inf
|
|
|
|
def set_decr_every_n_nan_or_inf(
|
|
self, new_decr_every_n_nan_or_inf: int
|
|
) -> None:
|
|
"""
|
|
Set the num `n` by `new_decr_every_n_nan_or_inf`, `n` represent decreases loss scaling every `n` accumulated steps with nan or inf gradients.
|
|
|
|
Args:
|
|
new_decr_every_n_nan_or_inf(int): The new_decr_every_n_nan_or_inf used to update the num `n`, `n` represent decreases loss scaling every `n` accumulated steps with nan or inf gradients.
|
|
"""
|
|
self._decr_every_n_nan_or_inf = new_decr_every_n_nan_or_inf
|
|
|
|
def state_dict(self) -> _ScaleStateDict:
|
|
"""
|
|
Returns the state of the scaler as a `dict`, If this instance is not enabled, returns an empty dict.
|
|
|
|
Returns:
|
|
A dict of scaler includes:
|
|
scale (tensor): The loss scaling factor.
|
|
incr_ratio(float): The multiplier to use when increasing the loss scaling.
|
|
decr_ratio(float): The less-than-one-multiplier to use when decreasing the loss scaling.
|
|
incr_every_n_steps(int): Increases loss scaling every n consecutive steps with finite gradients.
|
|
decr_every_n_nan_or_inf(int): Decreases loss scaling every n accumulated steps with nan or inf gradients.
|
|
incr_count(int): The number of recent consecutive unskipped steps.
|
|
decr_count(int): The number of recent consecutive skipped steps.
|
|
use_dynamic_loss_scaling(bool): Whether to use dynamic loss scaling. If False, fixed loss_scaling is used. If True, the loss scaling is updated dynamically. Default is True.
|
|
"""
|
|
return (
|
|
{
|
|
"scale": self._scale.numpy(),
|
|
"incr_ratio": self._incr_ratio,
|
|
"decr_ratio": self._decr_ratio,
|
|
"incr_every_n_steps": self._incr_every_n_steps,
|
|
"decr_every_n_nan_or_inf": self._decr_every_n_nan_or_inf,
|
|
"incr_count": self._incr_count,
|
|
"decr_count": self._decr_count,
|
|
"use_dynamic_loss_scaling": self._use_dynamic_loss_scaling,
|
|
}
|
|
if self._enable
|
|
else {}
|
|
)
|
|
|
|
def load_state_dict(self, state_dict: _ScaleStateDict) -> None:
|
|
"""
|
|
Loads the scaler state.
|
|
|
|
Args:
|
|
state_dict(dict): scaler state. Should be an object returned from a call to `AmpScaler.state_dict()`.
|
|
"""
|
|
if not self._enable:
|
|
return
|
|
|
|
if len(state_dict) == 0:
|
|
raise RuntimeError(
|
|
"The input state dict is empty, possibly because it was saved "
|
|
"from a disabled instance of GradScaler."
|
|
)
|
|
|
|
self._init_loss_scaling = state_dict["scale"][0]
|
|
self._scale = paddle.to_tensor(
|
|
np.array([self._init_loss_scaling]).astype(np.float32)
|
|
)
|
|
self._incr_ratio = state_dict["incr_ratio"]
|
|
self._decr_ratio = state_dict["decr_ratio"]
|
|
self._incr_every_n_steps = state_dict["incr_every_n_steps"]
|
|
self._decr_every_n_nan_or_inf = state_dict["decr_every_n_nan_or_inf"]
|
|
self._incr_count = state_dict["incr_count"]
|
|
self._decr_count = state_dict["decr_count"]
|
|
self._use_dynamic_loss_scaling = state_dict["use_dynamic_loss_scaling"]
|
|
|
|
|
|
class GradScaler(AmpScaler):
|
|
"""
|
|
GradScaler is used for Auto-Mixed-Precision training in dynamic graph mode.
|
|
It controls the scaling of loss, helps avoiding numerical overflow.
|
|
The object of this class has nineteen methods `scale()`, `unscale_()`, `minimize()`, `step()`, `update()` and `get`/`set` api of parameters.
|
|
|
|
`scale()` is used to multiply the loss by a scale ratio.
|
|
`unscale_()` is used to unscale the gradients of parameters, multiplies the gradients of parameters by 1/(scale ratio)
|
|
`minimize()` is similar as `optimizer.minimize()`, performs parameters updating, and it will update the loss_scaling, it equal to `step()` + `update()`.
|
|
`step()` is similar as `optimizer.step()`, which performs parameters updating.
|
|
`update` is used to update the loss_scaling.
|
|
|
|
|
|
Commonly, it is used together with `paddle.amp.auto_cast` to achieve Auto-Mixed-Precision in
|
|
dynamic graph mode.
|
|
|
|
This API supports three calling conventions:
|
|
|
|
``GradScaler(enable=True, init_loss_scaling=2.0**16, incr_ratio=2.0, decr_ratio=0.5, incr_every_n_steps=2000, decr_every_n_nan_or_inf=1, use_dynamic_loss_scaling=True)``
|
|
|
|
``GradScaler(device, init_scale=2.0**16, growth_factor=2.0, backoff_factor=0.5, growth_interval=2000, enabled=True)``
|
|
|
|
``GradScaler(init_scale=2.0**16, growth_factor=2.0, backoff_factor=0.5, growth_interval=2000, enabled=True)``
|
|
|
|
Args:
|
|
enable(bool, optional): Enable loss scaling or not. Default is True.
|
|
**Alias**: ``enabled``
|
|
init_loss_scaling (float, optional): The initial loss scaling factor. Default is 2.0**16.
|
|
**Alias**: ``init_scale``
|
|
incr_ratio(float, optional): The multiplier to use when increasing the loss
|
|
scaling. Default is 2.0.
|
|
**Alias**: ``growth_factor``
|
|
decr_ratio(float, optional): The less-than-one-multiplier to use when decreasing
|
|
the loss scaling. Default is 0.5.
|
|
**Alias**: ``backoff_factor``
|
|
incr_every_n_steps(int, optional): Increases loss scaling every n consecutive
|
|
steps with finite gradients. Default is 2000.
|
|
**Alias**: ``growth_interval``
|
|
decr_every_n_nan_or_inf(int, optional): Decreases loss scaling every n
|
|
accumulated steps with nan or inf gradients. Default is 1.
|
|
use_dynamic_loss_scaling(bool, optional): Whether to use dynamic loss scaling. If False, fixed loss_scaling is used. If True, the loss scaling is updated dynamically. Default is True.
|
|
|
|
Returns:
|
|
An GradScaler object.
|
|
|
|
Examples:
|
|
|
|
.. code-block:: pycon
|
|
|
|
>>> import paddle
|
|
|
|
>>> model = paddle.nn.Conv2D(3, 2, 3, bias_attr=True)
|
|
>>> optimizer = paddle.optimizer.SGD(
|
|
... learning_rate=0.01,
|
|
... parameters=model.parameters(),
|
|
... )
|
|
>>> scaler = paddle.amp.GradScaler(init_loss_scaling=1024)
|
|
>>> data = paddle.rand([10, 3, 32, 32])
|
|
|
|
>>> with paddle.amp.auto_cast():
|
|
... conv = model(data)
|
|
... loss = paddle.mean(conv)
|
|
|
|
>>> scaled = scaler.scale(loss) # scale the loss
|
|
>>> scaled.backward() # do backward
|
|
>>> scaler.minimize(optimizer, scaled) # update parameters
|
|
>>> optimizer.clear_grad()
|
|
"""
|
|
|
|
@overload
|
|
def __init__(
|
|
self,
|
|
enable: bool = True,
|
|
init_loss_scaling: float = 2.0**16,
|
|
incr_ratio: float = 2.0,
|
|
decr_ratio: float = 0.5,
|
|
incr_every_n_steps: int = 2000,
|
|
decr_every_n_nan_or_inf: int = 1,
|
|
use_dynamic_loss_scaling: bool = True,
|
|
) -> None: ...
|
|
|
|
@overload
|
|
def __init__(
|
|
self,
|
|
device: str,
|
|
init_scale: float = 2.0**16,
|
|
growth_factor: float = 2.0,
|
|
backoff_factor: float = 0.5,
|
|
growth_interval: int = 2000,
|
|
enabled: bool = True,
|
|
) -> None: ...
|
|
|
|
@overload
|
|
def __init__(
|
|
self,
|
|
init_scale: float = 2.0**16,
|
|
growth_factor: float = 2.0,
|
|
backoff_factor: float = 0.5,
|
|
growth_interval: int = 2000,
|
|
enabled: bool = True,
|
|
) -> None: ...
|
|
|
|
@grad_scaler_decorator
|
|
def __init__(
|
|
self,
|
|
enable: bool = True,
|
|
init_loss_scaling: float = 2.0**16,
|
|
incr_ratio: float = 2.0,
|
|
decr_ratio: float = 0.5,
|
|
incr_every_n_steps: int = 2000,
|
|
decr_every_n_nan_or_inf: int = 1,
|
|
use_dynamic_loss_scaling: bool = True,
|
|
) -> None:
|
|
super().__init__(
|
|
enable,
|
|
init_loss_scaling,
|
|
incr_ratio,
|
|
decr_ratio,
|
|
incr_every_n_steps,
|
|
decr_every_n_nan_or_inf,
|
|
use_dynamic_loss_scaling,
|
|
)
|
|
|
|
def scale(self, var: Tensor) -> Tensor:
|
|
"""
|
|
Multiplies a Tensor by the scale factor and returns scaled outputs.
|
|
If this instance of :class:`GradScaler` is not enabled, output are returned unmodified.
|
|
|
|
Args:
|
|
var (Tensor): The tensor to scale.
|
|
Returns:
|
|
The scaled tensor or original tensor.
|
|
|
|
Examples:
|
|
|
|
.. code-block:: pycon
|
|
|
|
>>> import paddle
|
|
>>> model = paddle.nn.Conv2D(3, 2, 3, bias_attr=True)
|
|
>>> optimizer = paddle.optimizer.SGD(
|
|
... learning_rate=0.01,
|
|
... parameters=model.parameters(),
|
|
... )
|
|
>>> scaler = paddle.amp.GradScaler(init_loss_scaling=1024)
|
|
>>> data = paddle.rand([10, 3, 32, 32])
|
|
>>> with paddle.amp.auto_cast():
|
|
... conv = model(data)
|
|
... loss = paddle.mean(conv)
|
|
>>> scaled = scaler.scale(loss) # scale the loss
|
|
>>> scaled.backward() # do backward
|
|
>>> scaler.minimize(optimizer, scaled) # update parameters
|
|
>>> optimizer.clear_grad()
|
|
"""
|
|
return super().scale(var)
|
|
|
|
def minimize(
|
|
self,
|
|
optimizer: Optimizer | OptimizerWithMixedPrecision,
|
|
*args: Any,
|
|
**kwargs: Any,
|
|
) -> tuple[list[Operator], list[tuple[Tensor, Tensor]]]:
|
|
"""
|
|
This function is similar as `optimizer.minimize()`, which performs parameters updating.
|
|
|
|
If the scaled gradients of parameters contains NAN or INF, the parameters updating is skipped.
|
|
Otherwise, if `unscale_()` has not been called, it first unscales the scaled gradients of parameters, then updates the parameters.
|
|
|
|
Finally, the loss scaling ratio is updated.
|
|
|
|
Args:
|
|
optimizer(Optimizer): The optimizer used to update parameters.
|
|
args: Arguments, which will be forward to `optimizer.minimize()`.
|
|
kwargs: Keyword arguments, which will be forward to `optimizer.minimize()`.
|
|
|
|
Examples:
|
|
|
|
.. code-block:: pycon
|
|
|
|
>>> import paddle
|
|
>>> model = paddle.nn.Conv2D(3, 2, 3, bias_attr=True)
|
|
>>> optimizer = paddle.optimizer.SGD(
|
|
... learning_rate=0.01,
|
|
... parameters=model.parameters(),
|
|
... )
|
|
>>> scaler = paddle.amp.GradScaler(init_loss_scaling=1024)
|
|
>>> data = paddle.rand([10, 3, 32, 32])
|
|
>>> with paddle.amp.auto_cast():
|
|
... conv = model(data)
|
|
... loss = paddle.mean(conv)
|
|
>>> scaled = scaler.scale(loss) # scale the loss
|
|
>>> scaled.backward() # do backward
|
|
>>> scaler.minimize(optimizer, scaled) # update parameters
|
|
>>> optimizer.clear_grad()
|
|
"""
|
|
return super().minimize(optimizer, *args, **kwargs)
|
|
|
|
def step(self, optimizer: Optimizer) -> None:
|
|
"""
|
|
This function is similar as `optimizer.step()`, which performs parameters updating.
|
|
|
|
If the scaled gradients of parameters contains NAN or INF, the parameters updating is skipped.
|
|
Otherwise, if `unscale_()` has not been called, it first unscales the scaled gradients of parameters, then updates the parameters.
|
|
|
|
Args:
|
|
optimizer(Optimizer): The optimizer used to update parameters.
|
|
|
|
Examples:
|
|
|
|
.. code-block:: pycon
|
|
|
|
>>> # doctest: +REQUIRES(env:GPU)
|
|
>>> import paddle
|
|
>>> paddle.device.set_device('gpu')
|
|
>>> model = paddle.nn.Conv2D(3, 2, 3, bias_attr=True)
|
|
>>> optimizer = paddle.optimizer.SGD(
|
|
... learning_rate=0.01,
|
|
... parameters=model.parameters(),
|
|
... )
|
|
>>> scaler = paddle.amp.GradScaler(init_loss_scaling=1024)
|
|
>>> data = paddle.rand([10, 3, 32, 32])
|
|
>>> with paddle.amp.auto_cast():
|
|
... conv = model(data)
|
|
... loss = paddle.mean(conv)
|
|
>>> scaled = scaler.scale(loss) # scale the loss
|
|
>>> scaled.backward() # do backward
|
|
>>> scaler.step(optimizer) # update parameters
|
|
>>> scaler.update() # update the loss scaling ratio
|
|
>>> optimizer.clear_grad()
|
|
"""
|
|
if not self._enable:
|
|
return optimizer.step()
|
|
|
|
optimizer_state = self._optimizer_states[id(optimizer)]
|
|
if optimizer_state["state"] is OptimizerState.STEPPED:
|
|
raise RuntimeError(
|
|
"step() has already been called since the last update()."
|
|
)
|
|
|
|
# unscale the grad
|
|
if optimizer_state["state"] is OptimizerState.INIT:
|
|
self._unscale(optimizer)
|
|
|
|
if hasattr(optimizer, "_set_auxiliary_var"):
|
|
optimizer._set_auxiliary_var('found_inf', self._found_inf)
|
|
optimizer.step()
|
|
self._cache_found_inf = optimizer._get_auxiliary_var('found_inf')
|
|
else:
|
|
if self._found_inf:
|
|
self._cache_found_inf = True
|
|
else:
|
|
optimizer.step()
|
|
self._cache_found_inf = False
|
|
|
|
optimizer_state["state"] = OptimizerState.STEPPED
|
|
|
|
if not self._use_dynamic_loss_scaling:
|
|
self._optimizer_states = defaultdict(_refresh_optimizer_state)
|
|
|
|
def update(self, new_scale: float | None = None) -> None:
|
|
"""
|
|
Updates the loss_scaling.
|
|
|
|
Args:
|
|
new_scale(float, optional): New loss scaling factor. If provided, the loss scaling factor
|
|
is directly set to ``new_scale`` and the internal step counts are reset.
|
|
Default is None.
|
|
|
|
Examples:
|
|
|
|
.. code-block:: pycon
|
|
|
|
>>> # doctest: +REQUIRES(env:GPU)
|
|
>>> import paddle
|
|
>>> paddle.device.set_device('gpu')
|
|
>>> model = paddle.nn.Conv2D(3, 2, 3, bias_attr=True)
|
|
>>> optimizer = paddle.optimizer.SGD(
|
|
... learning_rate=0.01,
|
|
... parameters=model.parameters(),
|
|
... )
|
|
>>> scaler = paddle.amp.GradScaler(init_loss_scaling=1024)
|
|
>>> data = paddle.rand([10, 3, 32, 32])
|
|
>>> with paddle.amp.auto_cast():
|
|
... conv = model(data)
|
|
... loss = paddle.mean(conv)
|
|
>>> scaled = scaler.scale(loss) # scale the loss
|
|
>>> scaled.backward() # do backward
|
|
>>> scaler.step(optimizer) # update parameters
|
|
>>> scaler.update() # update the loss scaling ratio
|
|
>>> optimizer.clear_grad()
|
|
"""
|
|
if not self._enable:
|
|
return
|
|
if new_scale is not None:
|
|
self._scale = paddle.to_tensor([new_scale], dtype='float32')
|
|
self._incr_count = 0
|
|
self._decr_count = 0
|
|
self._optimizer_states = defaultdict(_refresh_optimizer_state)
|
|
elif self._use_dynamic_loss_scaling:
|
|
self._update()
|
|
self._optimizer_states = defaultdict(_refresh_optimizer_state)
|
|
return
|
|
|
|
def unscale_(self, optimizer):
|
|
"""
|
|
Unscale the gradients of parameters, multiplies the gradients of parameters by 1/(loss scaling ratio).
|
|
If this instance of :class:`GradScaler` is not enabled, output are returned unmodified.
|
|
|
|
Args:
|
|
optimizer(Optimizer): The optimizer used to update parameters.
|
|
|
|
Returns:
|
|
The unscaled parameters or original parameters.
|
|
|
|
Examples:
|
|
|
|
.. code-block:: pycon
|
|
|
|
>>> # doctest: +REQUIRES(env:GPU)
|
|
>>> import paddle
|
|
>>> paddle.device.set_device('gpu')
|
|
>>> model = paddle.nn.Conv2D(3, 2, 3, bias_attr=True)
|
|
>>> optimizer = paddle.optimizer.SGD(
|
|
... learning_rate=0.01,
|
|
... parameters=model.parameters(),
|
|
... )
|
|
>>> scaler = paddle.amp.GradScaler(init_loss_scaling=1024)
|
|
>>> data = paddle.rand([10, 3, 32, 32])
|
|
>>> with paddle.amp.auto_cast():
|
|
... conv = model(data)
|
|
... loss = paddle.mean(conv)
|
|
>>> scaled = scaler.scale(loss) # scale the loss
|
|
>>> scaled.backward() # do backward
|
|
>>> scaler.unscale_(optimizer) # unscale the parameter
|
|
>>> scaler.step(optimizer)
|
|
>>> scaler.update()
|
|
>>> optimizer.clear_grad()
|
|
"""
|
|
return super()._unscale(optimizer)
|
|
|
|
def is_enable(self) -> bool:
|
|
"""
|
|
Enable loss scaling or not.
|
|
|
|
Returns:
|
|
bool: enable loss scaling return True else return False.
|
|
|
|
Examples:
|
|
|
|
.. code-block:: pycon
|
|
|
|
>>> # doctest: +REQUIRES(env:GPU, env:XPU)
|
|
>>> import paddle
|
|
>>> scaler = paddle.amp.GradScaler(
|
|
... enable=True,
|
|
... init_loss_scaling=1024,
|
|
... incr_ratio=2.0,
|
|
... decr_ratio=0.5,
|
|
... incr_every_n_steps=1000,
|
|
... decr_every_n_nan_or_inf=2,
|
|
... use_dynamic_loss_scaling=True,
|
|
... )
|
|
>>> enable = scaler.is_enable()
|
|
>>> print(enable)
|
|
True
|
|
"""
|
|
return super().is_enable()
|
|
|
|
def is_use_dynamic_loss_scaling(self) -> bool:
|
|
"""
|
|
Whether to use dynamic loss scaling.
|
|
|
|
Returns:
|
|
bool: if fixed loss_scaling is used return False, if the loss scaling is updated dynamically return True.
|
|
|
|
Examples:
|
|
.. code-block:: pycon
|
|
|
|
>>> # doctest: +REQUIRES(env:GPU, env:XPU)
|
|
>>> import paddle
|
|
>>> scaler = paddle.amp.GradScaler(
|
|
... enable=True,
|
|
... init_loss_scaling=1024,
|
|
... incr_ratio=2.0,
|
|
... decr_ratio=0.5,
|
|
... incr_every_n_steps=1000,
|
|
... decr_every_n_nan_or_inf=2,
|
|
... use_dynamic_loss_scaling=True,
|
|
... )
|
|
>>> use_dynamic_loss_scaling = scaler.is_use_dynamic_loss_scaling()
|
|
>>> print(use_dynamic_loss_scaling)
|
|
True
|
|
"""
|
|
return super().is_use_dynamic_loss_scaling()
|
|
|
|
def get_init_loss_scaling(self) -> float:
|
|
"""
|
|
Return the initial loss scaling factor.
|
|
|
|
Returns:
|
|
float: the initial loss scaling factor.
|
|
|
|
Examples:
|
|
.. code-block:: pycon
|
|
|
|
>>> # doctest: +REQUIRES(env:GPU, env:XPU)
|
|
>>> import paddle
|
|
>>> scaler = paddle.amp.GradScaler(
|
|
... enable=True,
|
|
... init_loss_scaling=1024,
|
|
... incr_ratio=2.0,
|
|
... decr_ratio=0.5,
|
|
... incr_every_n_steps=1000,
|
|
... decr_every_n_nan_or_inf=2,
|
|
... use_dynamic_loss_scaling=True,
|
|
... )
|
|
>>> init_loss_scaling = scaler.get_init_loss_scaling()
|
|
>>> print(init_loss_scaling)
|
|
1024
|
|
"""
|
|
return super().get_init_loss_scaling()
|
|
|
|
def set_init_loss_scaling(self, new_init_loss_scaling: int) -> None:
|
|
"""
|
|
Set the initial loss scaling factor by `new_init_loss_scaling`.
|
|
|
|
Args:
|
|
new_init_loss_scaling(float): The new_init_loss_scaling used to update initial loss scaling factor.
|
|
|
|
Examples:
|
|
.. code-block:: pycon
|
|
|
|
>>> # doctest: +REQUIRES(env:GPU, env:XPU)
|
|
>>> import paddle
|
|
>>> scaler = paddle.amp.GradScaler(
|
|
... enable=True,
|
|
... init_loss_scaling=1024,
|
|
... incr_ratio=2.0,
|
|
... decr_ratio=0.5,
|
|
... incr_every_n_steps=1000,
|
|
... decr_every_n_nan_or_inf=2,
|
|
... use_dynamic_loss_scaling=True,
|
|
... )
|
|
>>> print(scaler.get_init_loss_scaling())
|
|
1024
|
|
>>> new_init_loss_scaling = 1000
|
|
>>> scaler.set_init_loss_scaling(new_init_loss_scaling)
|
|
>>> print(scaler.get_init_loss_scaling())
|
|
1000
|
|
"""
|
|
super().set_init_loss_scaling(new_init_loss_scaling)
|
|
|
|
def get_incr_ratio(self) -> float:
|
|
"""
|
|
Return the multiplier to use when increasing the loss scaling.
|
|
|
|
Returns:
|
|
float: the multiplier to use when increasing the loss scaling.
|
|
|
|
Examples:
|
|
.. code-block:: pycon
|
|
|
|
>>> # doctest: +REQUIRES(env:GPU, env:XPU)
|
|
>>> import paddle
|
|
>>> scaler = paddle.amp.GradScaler(
|
|
... enable=True,
|
|
... init_loss_scaling=1024,
|
|
... incr_ratio=2.0,
|
|
... decr_ratio=0.5,
|
|
... incr_every_n_steps=1000,
|
|
... decr_every_n_nan_or_inf=2,
|
|
... use_dynamic_loss_scaling=True,
|
|
... )
|
|
>>> incr_ratio = scaler.get_incr_ratio()
|
|
>>> print(incr_ratio)
|
|
2.0
|
|
"""
|
|
return super().get_incr_ratio()
|
|
|
|
def set_incr_ratio(self, new_incr_ratio: float) -> None:
|
|
"""
|
|
Set the multiplier to use when increasing the loss scaling by `new_incr_ratio`, `new_incr_ratio` should > 1.0.
|
|
|
|
Args:
|
|
new_incr_ratio(float): The new_incr_ratio used to update the multiplier to use when increasing the loss scaling.
|
|
|
|
Examples:
|
|
.. code-block:: pycon
|
|
|
|
>>> # doctest: +REQUIRES(env:GPU, env:XPU)
|
|
>>> import paddle
|
|
>>> scaler = paddle.amp.GradScaler(
|
|
... enable=True,
|
|
... init_loss_scaling=1024,
|
|
... incr_ratio=2.0,
|
|
... decr_ratio=0.5,
|
|
... incr_every_n_steps=1000,
|
|
... decr_every_n_nan_or_inf=2,
|
|
... use_dynamic_loss_scaling=True,
|
|
... )
|
|
>>> print(scaler.get_incr_ratio())
|
|
2.0
|
|
>>> new_incr_ratio = 3.0
|
|
>>> scaler.set_incr_ratio(new_incr_ratio)
|
|
>>> print(scaler.get_incr_ratio())
|
|
3.0
|
|
"""
|
|
super().set_incr_ratio(new_incr_ratio)
|
|
|
|
def get_decr_ratio(self) -> float:
|
|
"""
|
|
Get the less-than-one-multiplier to use when decreasing the loss scaling.
|
|
|
|
Returns:
|
|
float: the less-than-one-multiplier to use when decreasing the loss scaling.
|
|
|
|
Examples:
|
|
.. code-block:: pycon
|
|
|
|
>>> # doctest: +REQUIRES(env:GPU, env:XPU)
|
|
>>> import paddle
|
|
>>> scaler = paddle.amp.GradScaler(
|
|
... enable=True,
|
|
... init_loss_scaling=1024,
|
|
... incr_ratio=2.0,
|
|
... decr_ratio=0.5,
|
|
... incr_every_n_steps=1000,
|
|
... decr_every_n_nan_or_inf=2,
|
|
... use_dynamic_loss_scaling=True,
|
|
... )
|
|
>>> decr_ratio = scaler.get_decr_ratio()
|
|
>>> print(decr_ratio)
|
|
0.5
|
|
"""
|
|
return super().get_decr_ratio()
|
|
|
|
def set_decr_ratio(self, new_decr_ratio: float) -> None:
|
|
"""
|
|
Set the less-than-one-multiplier to use when decreasing the loss scaling by `new_incr_ratio`, `new_decr_ratio` should < 1.0.
|
|
|
|
Args:
|
|
new_decr_ratio(float): The new_decr_ratio used to update the less-than-one-multiplier to use when decreasing the loss scaling.
|
|
|
|
Examples:
|
|
.. code-block:: pycon
|
|
|
|
>>> # doctest: +REQUIRES(env:GPU, env:XPU)
|
|
>>> import paddle
|
|
>>> scaler = paddle.amp.GradScaler(
|
|
... enable=True,
|
|
... init_loss_scaling=1024,
|
|
... incr_ratio=2.0,
|
|
... decr_ratio=0.5,
|
|
... incr_every_n_steps=1000,
|
|
... decr_every_n_nan_or_inf=2,
|
|
... use_dynamic_loss_scaling=True,
|
|
... )
|
|
>>> print(scaler.get_decr_ratio())
|
|
0.5
|
|
>>> new_decr_ratio = 0.1
|
|
>>> scaler.set_decr_ratio(new_decr_ratio)
|
|
>>> print(scaler.get_decr_ratio())
|
|
0.1
|
|
"""
|
|
super().set_decr_ratio(new_decr_ratio)
|
|
|
|
def get_incr_every_n_steps(self) -> int:
|
|
"""
|
|
Return the num `n`, `n` represent increases loss scaling every `n` consecutive steps with finite gradients.
|
|
|
|
Returns:
|
|
int: the num `n`, `n` represent increases loss scaling every `n` consecutive steps with finite gradients.
|
|
|
|
Examples:
|
|
.. code-block:: pycon
|
|
|
|
>>> # doctest: +REQUIRES(env:GPU, env:XPU)
|
|
>>> import paddle
|
|
>>> scaler = paddle.amp.GradScaler(
|
|
... enable=True,
|
|
... init_loss_scaling=1024,
|
|
... incr_ratio=2.0,
|
|
... decr_ratio=0.5,
|
|
... incr_every_n_steps=1000,
|
|
... decr_every_n_nan_or_inf=2,
|
|
... use_dynamic_loss_scaling=True,
|
|
... )
|
|
>>> incr_every_n_steps = scaler.get_incr_every_n_steps()
|
|
>>> print(incr_every_n_steps)
|
|
1000
|
|
"""
|
|
return super().get_incr_every_n_steps()
|
|
|
|
def set_incr_every_n_steps(self, new_incr_every_n_steps: int) -> None:
|
|
"""
|
|
Set the num `n` by `new_incr_every_n_steps`, `n` represent increases loss scaling every `n` consecutive steps with finite gradients.
|
|
|
|
Args:
|
|
new_incr_every_n_steps(int): The new_incr_every_n_steps used to update the num `n`, `n` represent increases loss scaling every `n` consecutive steps with finite gradients.
|
|
|
|
Examples:
|
|
.. code-block:: pycon
|
|
|
|
>>> # doctest: +REQUIRES(env:GPU, env:XPU)
|
|
>>> import paddle
|
|
>>> scaler = paddle.amp.GradScaler(
|
|
... enable=True,
|
|
... init_loss_scaling=1024,
|
|
... incr_ratio=2.0,
|
|
... decr_ratio=0.5,
|
|
... incr_every_n_steps=1000,
|
|
... decr_every_n_nan_or_inf=2,
|
|
... use_dynamic_loss_scaling=True,
|
|
... )
|
|
>>> print(scaler.get_incr_every_n_steps())
|
|
1000
|
|
>>> new_incr_every_n_steps = 2000
|
|
>>> scaler.set_incr_every_n_steps(new_incr_every_n_steps)
|
|
>>> print(scaler.get_incr_every_n_steps())
|
|
2000
|
|
"""
|
|
super().set_incr_every_n_steps(new_incr_every_n_steps)
|
|
|
|
def get_decr_every_n_nan_or_inf(self) -> int:
|
|
"""
|
|
Return the num `n`, `n` represent decreases loss scaling every `n` accumulated steps with nan or inf gradients.
|
|
|
|
Returns:
|
|
int: the num `n`, `n` represent decreases loss scaling every `n` accumulated steps with nan or inf gradients.
|
|
|
|
Examples:
|
|
.. code-block:: pycon
|
|
|
|
>>> # doctest: +REQUIRES(env:GPU, env:XPU)
|
|
>>> import paddle
|
|
>>> scaler = paddle.amp.GradScaler(
|
|
... enable=True,
|
|
... init_loss_scaling=1024,
|
|
... incr_ratio=2.0,
|
|
... decr_ratio=0.5,
|
|
... incr_every_n_steps=1000,
|
|
... decr_every_n_nan_or_inf=2,
|
|
... use_dynamic_loss_scaling=True,
|
|
... )
|
|
>>> decr_every_n_nan_or_inf = scaler.get_decr_every_n_nan_or_inf()
|
|
>>> print(decr_every_n_nan_or_inf)
|
|
2
|
|
"""
|
|
return super().get_decr_every_n_nan_or_inf()
|
|
|
|
def set_decr_every_n_nan_or_inf(
|
|
self, new_decr_every_n_nan_or_inf: int
|
|
) -> None:
|
|
"""
|
|
Set the num `n` by `new_decr_every_n_nan_or_inf`, `n` represent decreases loss scaling every `n` accumulated steps with nan or inf gradients.
|
|
|
|
Args:
|
|
new_decr_every_n_nan_or_inf(int): The new_decr_every_n_nan_or_inf used to update the num `n`, `n` represent decreases loss scaling every `n` accumulated steps with nan or inf gradients.
|
|
|
|
Examples:
|
|
.. code-block:: pycon
|
|
|
|
>>> # doctest: +REQUIRES(env:GPU, env:XPU)
|
|
>>> import paddle
|
|
>>> scaler = paddle.amp.GradScaler(
|
|
... enable=True,
|
|
... init_loss_scaling=1024,
|
|
... incr_ratio=2.0,
|
|
... decr_ratio=0.5,
|
|
... incr_every_n_steps=1000,
|
|
... decr_every_n_nan_or_inf=2,
|
|
... use_dynamic_loss_scaling=True,
|
|
... )
|
|
>>> print(scaler.get_decr_every_n_nan_or_inf())
|
|
2
|
|
>>> new_decr_every_n_nan_or_inf = 3
|
|
>>> scaler.set_decr_every_n_nan_or_inf(new_decr_every_n_nan_or_inf)
|
|
>>> print(scaler.get_decr_every_n_nan_or_inf())
|
|
3
|
|
"""
|
|
super().set_decr_every_n_nan_or_inf(new_decr_every_n_nan_or_inf)
|
|
|
|
is_enabled = is_enable
|
|
|
|
def get_scale(self) -> float:
|
|
"""
|
|
Return the current scale factor as a Python float.
|
|
|
|
If loss scaling is not enabled, returns 0.0.
|
|
|
|
Returns:
|
|
float: The current loss scaling factor, or 0.0 if disabled.
|
|
|
|
Examples:
|
|
|
|
.. code-block:: pycon
|
|
|
|
>>> # doctest: +REQUIRES(env:GPU, env:XPU)
|
|
>>> import paddle
|
|
>>> scaler = paddle.amp.GradScaler(init_loss_scaling=1024)
|
|
>>> print(scaler.get_scale())
|
|
1024.0
|
|
"""
|
|
if self._enable and self._scale is not None:
|
|
return float(self._scale)
|
|
return 0.0
|
|
|
|
get_growth_factor = get_incr_ratio
|
|
set_growth_factor = set_incr_ratio
|
|
get_backoff_factor = get_decr_ratio
|
|
set_backoff_factor = set_decr_ratio
|
|
get_growth_interval = get_incr_every_n_steps
|
|
set_growth_interval = set_incr_every_n_steps
|
|
|
|
def state_dict(self) -> _ScaleStateDict:
|
|
"""
|
|
Returns the state of the scaler as a `dict`, If this instance is not enabled, returns an empty dict.
|
|
|
|
Returns:
|
|
A dict of scaler includes:
|
|
scale (tensor): The loss scaling factor.
|
|
incr_ratio(float): The multiplier to use when increasing the loss scaling.
|
|
decr_ratio(float): The less-than-one-multiplier to use when decreasing the loss scaling.
|
|
incr_every_n_steps(int): Increases loss scaling every n consecutive steps with finite gradients.
|
|
decr_every_n_nan_or_inf(int): Decreases loss scaling every n accumulated steps with nan or inf gradients.
|
|
incr_count(int): The number of recent consecutive unskipped steps.
|
|
decr_count(int): The number of recent consecutive skipped steps.
|
|
use_dynamic_loss_scaling(bool): Whether to use dynamic loss scaling. If False, fixed loss_scaling is used. If True, the loss scaling is updated dynamically. Default is True.
|
|
|
|
Examples:
|
|
|
|
.. code-block:: pycon
|
|
|
|
>>> # doctest: +REQUIRES(env:GPU, env:XPU)
|
|
>>> import paddle
|
|
|
|
>>> scaler = paddle.amp.GradScaler(
|
|
... enable=True,
|
|
... init_loss_scaling=1024,
|
|
... incr_ratio=2.0,
|
|
... decr_ratio=0.5,
|
|
... incr_every_n_steps=1000,
|
|
... decr_every_n_nan_or_inf=2,
|
|
... use_dynamic_loss_scaling=True,
|
|
... )
|
|
>>> scaler_state = scaler.state_dict()
|
|
"""
|
|
return super().state_dict()
|
|
|
|
def load_state_dict(self, state_dict: _ScaleStateDict) -> None:
|
|
"""
|
|
Loads the scaler state.
|
|
|
|
Args:
|
|
state_dict(dict): scaler state. Should be an object returned from a call to `GradScaler.state_dict()`.
|
|
|
|
Examples:
|
|
|
|
.. code-block:: pycon
|
|
|
|
>>> # doctest: +REQUIRES(env:GPU, env:XPU)
|
|
>>> import paddle
|
|
|
|
>>> scaler = paddle.amp.GradScaler(
|
|
... enable=True,
|
|
... init_loss_scaling=1024,
|
|
... incr_ratio=2.0,
|
|
... decr_ratio=0.5,
|
|
... incr_every_n_steps=1000,
|
|
... decr_every_n_nan_or_inf=2,
|
|
... use_dynamic_loss_scaling=True,
|
|
... )
|
|
>>> scaler_state = scaler.state_dict()
|
|
>>> scaler.load_state_dict(scaler_state)
|
|
"""
|
|
super().load_state_dict(state_dict)
|