Files
paddlepaddle--paddle/python/paddle/amp/grad_scaler.py
T
2026-07-13 12:40:42 +08:00

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)