wehub/paddlepaddle--paddle
1
0
Fork 0
Code Issues Pull Requests Actions 2 Deployments Models Agent Notes Packages Projects Releases Wiki Activity

chore: import upstream snapshot with attribution

Browse Source
This commit is contained in:
wehub-resource-sync
2026-07-13 12:40:42 +08:00
commit e25996e7db
15472 changed files with 3536181 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/paddle/optimizer/__init__.py
+50
View File
@@ -0,0 +1,50 @@
# Copyright (c) 2020 PaddlePaddle Authors. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
from . import lr # noqa: F401
from .adadelta import Adadelta
from .adagrad import Adagrad
from .adam import Adam
from .adamax import Adamax
from .adamw import AdamW
from .asgd import ASGD
from .lamb import Lamb
from .lbfgs import LBFGS
from .momentum import Momentum
from .muon import Muon
from .nadam import NAdam
from .optimizer import Optimizer
from .radam import RAdam
from .rmsprop import RMSProp
from .rprop import Rprop
from .sgd import SGD
__all__ = [
'Optimizer',
'Adagrad',
'Adam',
'AdamW',
'Adamax',
'ASGD',
'RAdam',
'RMSProp',
'Adadelta',
'SGD',
'Rprop',
'Momentum',
'NAdam',
'Lamb',
'LBFGS',
'Muon',
]
python/paddle/optimizer/adadelta.py
+268
View File
@@ -0,0 +1,268 @@
# 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 typing import TYPE_CHECKING
import paddle
from paddle import _C_ops
from paddle.base.framework import in_dynamic_or_pir_mode
from ..base import framework
from ..base.dygraph import no_grad
from .optimizer import Optimizer
if TYPE_CHECKING:
from collections.abc import Sequence
from typing_extensions import NotRequired
from paddle import Tensor
from paddle.nn.clip import GradientClipBase
from paddle.regularizer import WeightDecayRegularizer
from .lr import LRScheduler
from .optimizer import _ParameterConfig
class _AdadeltaParameterConfig(_ParameterConfig):
epsilon: NotRequired[float]
rho: NotRequired[float]
__all__ = []
class Adadelta(Optimizer):
r"""
**Notes: This API does not support sparse parameter optimization.**
Adadelta Optimizer. Please refer to this for details:
`ADADELTA: AN ADAPTIVE LEARNING RATE METHOD <https://arxiv.org/abs/1212.5701>`_.
The update is done as follows:
.. math::
E(g_t^2) &= \rho * E(g_{t-1}^2) + (1-\rho) * g^2
learning\_rate &= \sqrt{ ( E(dx_{t-1}^2) + \epsilon ) / ( E(g_t^2) + \epsilon ) }
E(dx_t^2) &= \rho * E(dx_{t-1}^2) + (1-\rho) * (-g*learning\_rate)^2
Args:
learning_rate (float|Tensor|LRScheduler, optional): The learning rate used to update ``Parameter``.
It can be a float value, a ``Tensor`` with a float type or a LearningRateDecay. The default value is 0.001.
epsilon (float): a small float number for numeric stability. Default 1.0e-6.
rho (float): a floating point value indicating the decay rate. Default 0.95.
parameters (list|tuple|None, optional): List/Tuple of ``Tensor`` to update to minimize ``loss``. \
This parameter is required in dygraph mode. And you can specify different options for \
different parameter groups such as the learning rate, weight decay, etc, \
then the parameters are list of dict. Note that the learning_rate in parameter groups \
represents the scale of base learning_rate. \
The default value is None in static graph mode, at this time all parameters will be updated.
weight_decay (int|float|WeightDecayRegularizer|None, optional): The strategy of regularization. \
It can be a int or float value as coeff of L2 regularization or \
:ref:`api_paddle_regularizer_L1Decay`, :ref:`api_paddle_regularizer_L2Decay`.
If a parameter has set regularizer using :ref:`api_paddle_ParamAttr` already, \
the regularization setting here in optimizer will be ignored for this parameter. \
Otherwise, the regularization setting here in optimizer will take effect. \
Default None, meaning there is no regularization.
grad_clip (GradientClipBase|None, optional): Gradient clipping strategy, it's an instance of
some derived class of ``GradientClipBase`` . There are three clipping strategies
( :ref:`api_paddle_nn_ClipGradByGlobalNorm` , :ref:`api_paddle_nn_ClipGradByNorm` ,
:ref:`api_paddle_nn_ClipGradByValue` ). Default None, meaning there is no gradient clipping.
name (str|None, optional): The default value is None. Normally there is no need for user
to set this property. For more information, please refer to
:ref:`api_guide_Name` .
Examples:
.. code-block:: pycon
>>> import paddle
>>> inp = paddle.uniform([10, 10], dtype="float32", min=-0.1, max=0.1)
>>> linear = paddle.nn.Linear(10, 10)
>>> out = linear(inp)
>>> loss = paddle.mean(out)
>>> beta1 = paddle.to_tensor([0.9], dtype="float32")
>>> beta2 = paddle.to_tensor([0.99], dtype="float32")
>>> adadelta = paddle.optimizer.Adadelta(learning_rate=0.1, parameters=linear.parameters(), weight_decay=0.01)
>>> back = out.backward()
>>> adadelta.step()
>>> adadelta.clear_grad()
>>> # Note that the learning_rate of linear_2 is 0.01.
>>> linear_1 = paddle.nn.Linear(10, 10)
>>> linear_2 = paddle.nn.Linear(10, 10)
>>> inp = paddle.uniform(shape=[10, 10], min=-0.1, max=0.1)
>>> out = linear_1(inp)
>>> out = linear_2(out)
>>> loss = paddle.mean(out)
>>> adadelta = paddle.optimizer.Adadelta(
... learning_rate=0.1,
... parameters=[{ # type: ignore
... 'params': linear_1.parameters()
... }, {
... 'params': linear_2.parameters(),
... 'weight_decay': 0.001,
... 'learning_rate': 0.1,
... }],
... weight_decay=0.01)
>>> out.backward()
>>> adadelta.step()
>>> adadelta.clear_grad()
"""
type: str
_avg_squared_grad_acc_str = "_avg_squared_grad"
_avg_squared_update_acc_str = "_avg_squared_update"
def __init__(
self,
learning_rate: float | Tensor | LRScheduler = 0.001,
epsilon: float = 1.0e-6,
rho: float = 0.95,
parameters: (
Sequence[Tensor] | Sequence[_AdadeltaParameterConfig] | None
) = None,
weight_decay: float | WeightDecayRegularizer | None = None,
grad_clip: GradientClipBase | None = None,
name: str | None = None,
) -> None:
if learning_rate is None:
raise ValueError("learning_rate is not set.")
if epsilon is None:
raise ValueError("epsilon is not set.")
if rho is None:
raise ValueError("rho is not set.")
super().__init__(
learning_rate=learning_rate,
parameters=parameters,
weight_decay=weight_decay,
grad_clip=grad_clip,
name=name,
)
self._multi_precision = False
self._master_weights = {}
self.type = "adadelta"
self._epsilon = epsilon
self._rho = rho
self._default_dict = {
'epsilon': epsilon,
'rho': rho,
}
def _create_accumulators(self, block, parameters):
if not isinstance(block, (framework.Block, paddle.pir.Block)):
raise TypeError("block is not instance of framework.Block.")
if isinstance(parameters, dict):
parameters = parameters.get('params')
for p in parameters:
if p.name in self._already_create_accumulator:
continue
if self._multi_precision and self._is_dtype_fp16_or_bf16(p.dtype):
master_p = self._create_master_weight(p)
self._add_accumulator(self._avg_squared_grad_acc_str, master_p)
self._add_accumulator(
self._avg_squared_update_acc_str, master_p
)
self._already_create_accumulator.add(p.name)
continue
if (
self._is_dtype_fp16_or_bf16(p.dtype)
and not self._multi_precision
):
warnings.warn(
"Accumulating with FP16/BF16 in optimizer can lead to poor accuracy or slow convergence."
"Consider using multi_precision=True option of the Lars optimizer."
)
self._add_accumulator(self._avg_squared_grad_acc_str, p)
self._add_accumulator(self._avg_squared_update_acc_str, p)
self._already_create_accumulator.add(p.name)
def _append_optimize_op(self, block, param_and_grad):
if isinstance(param_and_grad, dict):
param_and_grad = self._update_param_group(param_and_grad)
avg_squared_grad_acc = self._get_accumulator_master(
self._avg_squared_grad_acc_str, param_and_grad[0]
)
avg_squared_update_acc = self._get_accumulator_master(
self._avg_squared_update_acc_str, param_and_grad[0]
)
find_master = self._multi_precision and self._is_dtype_fp16_or_bf16(
param_and_grad[0].dtype
)
master_weight = (
self._master_weights[param_and_grad[0].name]
if find_master
else None
)
if in_dynamic_or_pir_mode():
with no_grad():
_C_ops.adadelta_(
param_and_grad[0],
param_and_grad[1],
avg_squared_grad_acc,
avg_squared_update_acc,
self._create_param_lr(param_and_grad),
master_weight,
self._rho,
self._epsilon,
find_master,
)
return None
else:
if not isinstance(block, (framework.Block, paddle.pir.Block)):
raise TypeError("block is not instance of framework.Block.")
# Create the adadelta optimizer op
inputs = {
"Param": param_and_grad[0],
"Grad": param_and_grad[1],
"AvgSquaredGrad": avg_squared_grad_acc,
"AvgSquaredUpdate": avg_squared_update_acc,
"LearningRate": self._create_param_lr(param_and_grad),
}
outputs = {
"ParamOut": param_and_grad[0],
"AvgSquaredGradOut": avg_squared_grad_acc,
"AvgSquaredUpdateOut": avg_squared_update_acc,
}
if find_master:
inputs["MasterParam"] = master_weight
outputs["MasterParamOut"] = master_weight
adadelta_op = block.append_op(
type=self.type,
inputs=inputs,
outputs=outputs,
attrs={
"epsilon": self._epsilon,
"rho": self._rho,
"multi_precision": find_master,
},
stop_gradient=True,
)
return adadelta_op
def _update_param_group(self, parameters):
self._epsilon = parameters.get('epsilon', self._default_dict['epsilon'])
self._rho = parameters.get('rho', self._default_dict['rho'])
parameters = parameters.get('params')
return parameters
python/paddle/optimizer/adagrad.py
+273
View File
@@ -0,0 +1,273 @@
# 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 typing import TYPE_CHECKING
from paddle import _C_ops, pir
from paddle.framework import (
in_dynamic_or_pir_mode,
)
from ..base import framework
from .optimizer import Optimizer
if TYPE_CHECKING:
from collections.abc import Sequence
from typing_extensions import NotRequired
from paddle import Tensor
from paddle.nn.clip import GradientClipBase
from paddle.regularizer import WeightDecayRegularizer
from .optimizer import _ParameterConfig
class _AdagradParameterConfig(_ParameterConfig):
epsilon: NotRequired[float]
initial_accumulator_value: NotRequired[float]
__all__ = []
class Adagrad(Optimizer):
r"""
The Adaptive Gradient optimizer (Adagrad for short) use an optimization described
in paper: `Adaptive Subgradient Methods for Online Learning and
Stochastic Optimization <http://www.jmlr.org/papers/volume12/duchi11a/duchi11a.pdf>`_.
The parameter ``param_out`` update rule with gradient ``grad``:
.. math::
moment\_out &= moment + grad * grad
param\_out &= param - \frac{learning\_rate * grad}{\sqrt{moment\_out} + \epsilon}
The original paper does not have the ``epsilon`` attribute. It is added here
in our implementation as also proposed `Per-parameter adaptive learning rate
methods <http://cs231n.github.io/neural-networks-3/#ada>`_
for numerical stability to avoid the division by zero error.
Args:
learning_rate (float|Tensor): The learning rate used to update ``Parameter``.
It can be a float value or a ``Variable`` with a float type.
epsilon (float, optional): A small float value for numerical stability.
The default value is 1e-06.
parameters (list|tuple|None, optional): List/Tuple of ``Tensor`` to update to minimize ``loss``.
This parameter is required in dygraph mode. And you can specify different options for
different parameter groups such as the learning rate, weight decay, etc,
then the parameters are list of dict. Note that the learning_rate in parameter groups
represents the scale of base learning_rate.
The default value is None in static graph mode, at this time all parameters will be updated.
weight_decay (int|float|WeightDecayRegularizer|None, optional): The strategy of regularization.
It canbe a int or float value as coeff of L2 regularization or
:ref:`api_paddle_regularizer_L1Decay`, :ref:`api_paddle_regularizer_L2Decay`.
If a parameter has set regularizer using :ref:`api_paddle_base_param_attr_aramAttr` already,
the regularization setting here in optimizer will be ignored for this parameter.
Otherwise, the regularization setting here in optimizer will take effect.
Default None, meaning there is no regularization.
grad_clip (GradientClipBase|None, optional): Gradient clipping strategy, it's an instance of
some derived class of ``GradientClipBase`` . There are three clipping strategies,
ClipGradByGlobalNorm, ClipGradByNorm and ClipGradByValue. Default None,
meaning there is no gradient clipping.
name (str|None, optional): Normally there is no need for user to set this property.
For more information, please refer to :ref:`api_guide_Name`.
The default value is None.
initial_accumulator_value (float, optional): Initial value for moment accumulator.
The default value is 0.0.
Examples:
.. code-block:: pycon
>>> import paddle
>>> inp = paddle.rand(shape=[10, 10])
>>> linear = paddle.nn.Linear(10, 10)
>>> out = linear(inp)
>>> loss = paddle.mean(out)
>>> adagrad = paddle.optimizer.Adagrad(
... learning_rate=0.1,
... parameters=linear.parameters(),
... )
>>> out.backward()
>>> adagrad.step()
>>> adagrad.clear_grad()
>>> # Note that the learning_rate of linear_2 is 0.01.
>>> linear_1 = paddle.nn.Linear(10, 10)
>>> linear_2 = paddle.nn.Linear(10, 10)
>>> inp = paddle.uniform(shape=[10, 10], min=-0.1, max=0.1)
>>> out = linear_1(inp)
>>> out = linear_2(out)
>>> loss = paddle.mean(out)
>>> adagrad = paddle.optimizer.Adagrad(
... learning_rate=0.1,
... parameters=[ # type: ignore
... {
... 'params': linear_1.parameters(),
... },
... {
... 'params': linear_2.parameters(),
... 'weight_decay': 0.001,
... 'learning_rate': 0.1,
... },
... ],
... weight_decay=0.01,
... )
>>> out.backward()
>>> adagrad.step()
>>> adagrad.clear_grad()
"""
type: str
initial_accumulator_value: float
_moment_acc_str = "moment"
def __init__(
self,
learning_rate: float | Tensor,
epsilon: float = 1.0e-6,
parameters: (
Sequence[Tensor] | Sequence[_AdagradParameterConfig] | None
) = None,
weight_decay: float | WeightDecayRegularizer | None = None,
grad_clip: GradientClipBase | None = None,
name: str | None = None,
initial_accumulator_value: float = 0.0,
) -> None:
assert learning_rate is not None
assert epsilon is not None
super().__init__(
learning_rate=learning_rate,
parameters=parameters,
weight_decay=weight_decay,
grad_clip=grad_clip,
name=name,
)
self.type = "adagrad"
self._epsilon = epsilon
self._multi_precision = False
self._master_weights = {}
self.initial_accumulator_value = initial_accumulator_value
self._default_dict = {
'epsilon': epsilon,
'initial_accumulator_value': initial_accumulator_value,
}
def _create_accumulators(self, block, parameters):
if not isinstance(block, (framework.Block, pir.Block)):
raise TypeError("block is not instance of Block.")
if isinstance(parameters, dict):
parameters = self._update_param_group(parameters)
for p in parameters:
if p.name in self._already_create_accumulator:
continue
if self._multi_precision and self._is_dtype_fp16_or_bf16(p.dtype):
master_p = self._create_master_weight(p)
self._add_accumulator(
self._moment_acc_str,
master_p,
fill_value=self.initial_accumulator_value,
)
self._already_create_accumulator.add(p.name)
continue
if (
self._is_dtype_fp16_or_bf16(p.dtype)
and not self._multi_precision
):
warnings.warn(
"Accumulating with FP16/BF16 in optimizer can lead to poor accuracy or slow convergence."
"Consider using multi_precision=True option of the Momentum optimizer."
)
self._add_accumulator(
self._moment_acc_str,
p,
fill_value=self.initial_accumulator_value,
)
self._already_create_accumulator.add(p.name)
def _append_optimize_op(self, block, param_and_grad):
if not isinstance(block, (framework.Block, pir.Block)):
raise TypeError("block is not instance of Block.")
if isinstance(param_and_grad, dict):
param_and_grad = self._update_param_group(param_and_grad)
moment_acc = self._get_accumulator_master(
self._moment_acc_str, param_and_grad[0]
)
find_master = self._multi_precision and self._is_dtype_fp16_or_bf16(
param_and_grad[0].dtype
)
master_weight = (
self._master_weights[param_and_grad[0].name]
if find_master
else None
)
if in_dynamic_or_pir_mode():
_, _, _ = _C_ops.adagrad_(
param_and_grad[0],
param_and_grad[1],
moment_acc,
self._create_param_lr(param_and_grad),
master_weight if find_master else None,
self._epsilon,
find_master,
)
return None
else:
# Create the adagrad optimizer op
inputs = {
"Param": param_and_grad[0],
"Grad": param_and_grad[1],
"Moment": moment_acc,
"LearningRate": self._create_param_lr(param_and_grad),
}
outputs = {"ParamOut": param_and_grad[0], "MomentOut": moment_acc}
if find_master:
inputs["MasterParam"] = master_weight
outputs["MasterParamOut"] = master_weight
adagrad_op = block.append_op(
type=self.type,
inputs=inputs,
outputs=outputs,
attrs={
"epsilon": self._epsilon,
"multi_precision": find_master,
},
stop_gradient=True,
)
return adagrad_op
def _update_param_group(self, parameters):
self._epsilon = parameters.get('epsilon', self._default_dict['epsilon'])
self.initial_accumulator_value = parameters.get(
'initial_accumulator_value',
self._default_dict['initial_accumulator_value'],
)
parameters = parameters.get('params')
return parameters
python/paddle/optimizer/adam.py
+973
View File
@@ -0,0 +1,973 @@
# 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 typing import TYPE_CHECKING
import paddle
from paddle import _C_ops, pir
from paddle.base.libpaddle import DataType
from paddle.pir import Value
from ..base import core, framework
from ..base.dygraph import base as imperative_base
from ..base.framework import (
Variable,
in_dygraph_mode,
in_dynamic_or_pir_mode,
in_pir_mode,
)
from .optimizer import Optimizer
if TYPE_CHECKING:
from collections.abc import Callable, Sequence
from typing_extensions import NotRequired
from paddle import Tensor
from paddle.nn.clip import GradientClipBase
from paddle.regularizer import WeightDecayRegularizer
from .lr import LRScheduler
from .optimizer import _ParameterConfig
class _AdamParameterConfig(_ParameterConfig):
beta1: NotRequired[float | Tensor]
beta2: NotRequired[float | Tensor]
epsilon: NotRequired[float | Tensor]
lazy_mode: NotRequired[bool]
__all__ = []
class Adam(Optimizer):
r"""
The Adam optimizer uses an optimization described at the end
of section 2 of `Adam paper <https://arxiv.org/abs/1412.6980>`_ ,
it can dynamically adjusts the learning rate of each parameter using
the 1st moment estimates and the 2nd moment estimates of the gradient.
The parameter ``param_out`` update rule with gradient ``grad``:
.. math::
\begin{aligned}
&\hspace{5mm} t = t + 1 \\
&\hspace{5mm} moment\_1\_out = {\beta}_1 * moment\_1 + (1 - {\beta}_1) * grad \\
&\hspace{5mm} moment\_2\_out = {\beta}_2 * moment\_2 + (1 - {\beta}_2) * grad * grad \\
&\hspace{5mm} learning\_rate = learning\_rate * \frac{\sqrt{1 - {\beta}_2^t}}{1 - {\beta}_1^t} \\
&\hspace{5mm}\textbf{if} \: \textit{amsgrad}: \\
&\hspace{15mm} moment\_2\_max\_out = max(moment\_2\_out, moment\_2\_max) \\
&\hspace{15mm} param\_out = param - learning\_rate * \frac{moment\_1\_out}{\sqrt{moment\_2\_max\_out} + \epsilon} \\
&\hspace{5mm}\textbf{else}: \: \\
&\hspace{15mm} param\_out = param - learning\_rate * \frac{moment\_1\_out}{\sqrt{moment\_2\_out} + \epsilon} \\
\end{aligned}
Related paper: `Adam: A Method for Stochastic Optimization <https://arxiv.org/abs/1412.6980>`_
Args:
learning_rate (float|LRScheduler, optional): The learning rate used to update ``Parameter``.
It can be a float value or a LRScheduler. The default value is 0.001.
beta1 (float|Tensor, optional): The exponential decay rate for the 1st moment estimates.
It should be a float number or a 0-D Tensor with shape [] and data type as float32.
The default value is 0.9.
beta2 (float|Tensor, optional): The exponential decay rate for the 2nd moment estimates.
It should be a float number or a 0-D Tensor with shape [] and data type as float32.
The default value is 0.999.
epsilon (float|Tensor, optional): A small float value for numerical stability.
It should be a float number or a 0-D Tensor with shape [] and data type as float32.
The default value is 1e-08.
parameters (list|tuple|None, optional): List/Tuple of ``Tensor`` to update to minimize ``loss``.
This parameter is required in dygraph mode. And you can specify different options for
different parameter groups such as the learning rate, weight decay, etc,
then the parameters are list of dict. Note that the learning_rate in parameter groups
represents the scale of base learning_rate.
The default value is None in static graph mode, at this time all parameters will be updated.
weight_decay (int|float|WeightDecayRegularizer|None, optional): The strategy of regularization.
It canbe a int or float value as coeff of L2 regularization or
:ref:`api_paddle_regularizer_L1Decay`, :ref:`api_paddle_regularizer_L2Decay`.
If a parameter has set regularizer using :ref:`api_paddle_ParamAttr` already,
the regularization setting here in optimizer will be ignored for this parameter.
Otherwise, the regularization setting here in optimizer will take effect.
Default None, meaning there is no regularization.
grad_clip (GradientClipBase|None, optional): Gradient clipping strategy, it's an instance of
some derived class of ``GradientClipBase`` . There are three clipping strategies
( :ref:`api_paddle_nn_ClipGradByGlobalNorm` , :ref:`api_paddle_nn_ClipGradByNorm` ,
:ref:`api_paddle_nn_ClipGradByValue` ). Default None, meaning there is no gradient clipping.
lazy_mode (bool, optional): The official Adam algorithm has two moving-average accumulators.
The accumulators are updated at every step. Every element of the two moving-average
is updated in both dense mode and sparse mode. If the size of parameter is very large,
then the update may be very slow. The lazy mode only update the element that has
gradient in current mini-batch, so it will be much more faster. But this mode has
different semantics with the original Adam algorithm and may lead to different result.
The default value is False.
multi_precision (bool, optional): Whether to use multi-precision during weight updating. Default is false.
use_multi_tensor (bool, optional): Whether to use multi-tensor strategy to update all parameters at once . Default is false.
amsgrad (bool, optional): Whether to use the AMSGrad variant of this algorithm from the paper
`On the Convergence of Adam and Beyond <https://openreview.net/forum?id=ryQu7f-RZ>`_. Default is false.
name (str|None, optional): Normally there is no need for user to set this property.
For more information, please refer to :ref:`api_guide_Name`.
The default value is None.
Examples:
.. code-block:: pycon
:name: code-example1
>>> import paddle
>>> linear = paddle.nn.Linear(10, 10)
>>> inp = paddle.rand([10,10], dtype="float32")
>>> out = linear(inp)
>>> loss = paddle.mean(out)
>>> adam = paddle.optimizer.Adam(
... learning_rate=0.1,
... parameters=linear.parameters()
... )
>>> loss.backward()
>>> adam.step()
>>> adam.clear_grad()
.. code-block:: pycon
:name: code-example2
>>> # Adam with beta1/beta2 as Tensor and weight_decay as float
>>> import paddle
>>> linear = paddle.nn.Linear(10, 10)
>>> inp = paddle.rand([10,10], dtype="float32")
>>> out = linear(inp)
>>> loss = paddle.mean(out)
>>> beta1 = paddle.to_tensor([0.9], dtype="float32")
>>> beta2 = paddle.to_tensor([0.99], dtype="float32")
>>> adam = paddle.optimizer.Adam(
... learning_rate=0.1,
... parameters=linear.parameters(),
... beta1=beta1,
... beta2=beta2,
... weight_decay=0.01
... )
>>> loss.backward()
>>> adam.step()
>>> adam.clear_grad()
>>> # Note that the learning_rate of linear_2 is 0.01.
>>> linear_1 = paddle.nn.Linear(10, 10)
>>> linear_2 = paddle.nn.Linear(10, 10)
>>> inp = paddle.uniform(shape=[10, 10], min=-0.1, max=0.1)
>>> out = linear_1(inp)
>>> out = linear_2(out)
>>> loss = paddle.mean(out)
>>> adam = paddle.optimizer.Adam(
... learning_rate=0.1,
... parameters=[{ # type: ignore
... 'params': linear_1.parameters()
... }, {
... 'params': linear_2.parameters(),
... 'weight_decay': 0.001,
... 'learning_rate': 0.1,
... 'beta1': 0.8
... }],
... weight_decay=0.01,
... beta1=0.9
... )
>>> loss.backward()
>>> adam.step()
>>> adam.clear_grad()
"""
type: str
_moment1_acc_str = "moment1"
_moment2_acc_str = "moment2"
_moment2_acc_max_str = "moment2_max"
_beta1_pow_acc_str = "beta1_pow_acc"
_beta2_pow_acc_str = "beta2_pow_acc"
def __init__(
self,
learning_rate: float | LRScheduler = 0.001,
beta1: float | Tensor = 0.9,
beta2: float | Tensor = 0.999,
epsilon: float | Tensor = 1e-8,
parameters: (
Sequence[Tensor] | Sequence[_AdamParameterConfig] | None
) = None,
weight_decay: float | WeightDecayRegularizer | None = None,
grad_clip: GradientClipBase | None = None,
lazy_mode: bool = False,
multi_precision: bool = False,
use_multi_tensor: bool = False,
amsgrad: bool = False,
name: str | None = None,
) -> None:
assert learning_rate is not None
assert beta1 is not None
assert beta2 is not None
assert epsilon is not None
if not isinstance(beta1, (Variable, Value)):
if not 0 <= beta1 < 1:
raise ValueError(
"Invalid value of beta1, expect beta1 in [0,1)."
)
if not isinstance(beta2, (Variable, Value)):
if not 0 <= beta2 < 1:
raise ValueError(
"Invalid value of beta2, expect beta2 in [0,1)."
)
if not isinstance(epsilon, (Variable, Value)):
if not 0 <= epsilon:
raise ValueError(
"Invalid value of epsilon, expect epsilon >= 0."
)
super().__init__(
learning_rate=learning_rate,
parameters=parameters,
weight_decay=weight_decay,
grad_clip=grad_clip,
name=name,
)
self.type = "adam"
self._beta1 = beta1
self._beta2 = beta2
self._epsilon = epsilon
self._lazy_mode = lazy_mode
self._multi_precision = multi_precision
self._master_weights = {}
self._default_dict = {
'beta1': beta1,
'beta2': beta2,
'epsilon': epsilon,
'lazy_mode': lazy_mode,
}
self._use_multi_tensor = use_multi_tensor
if self._use_multi_tensor:
self._param_dict = self._create_multi_tensor_dict()
self._moment1_dict = self._create_multi_tensor_dict()
self._moment2_dict = self._create_multi_tensor_dict()
self._moment2_max_dict = (
self._create_multi_tensor_dict() if amsgrad else None
)
self._beta1_pow_acc_dict = self._create_multi_tensor_dict()
self._beta2_pow_acc_dict = self._create_multi_tensor_dict()
self._master_weight_dict = self._create_multi_tensor_dict()
self._master_weight_dict['FP32_DenseTensor'] = None
# whether to use AMSGrad
self._amsgrad = amsgrad
def get_lr_dtype(self) -> paddle.dtype:
return paddle.float64
def _create_regularization_of_grad(self, param, grad, regularization=None):
from paddle.regularizer import L2Decay
if (
regularization is not None
and isinstance(regularization, L2Decay)
and paddle.get_flags(['FLAGS_use_accuracy_compatible_kernel']).get(
'FLAGS_use_accuracy_compatible_kernel', False
)
):
# PyTorch fused Adam: grad += param * weight_decay in the kernel
# where weight_decay is double. The effective grad is:
# float32(float64(grad) + float64(param) * float64(wd))
# Replicate without intermediate float32 truncation.
wd = float(regularization._coeff) # Python float (float64)
return (grad.cast('float64') + param.cast('float64') * wd).cast(
'float32'
)
return super()._create_regularization_of_grad(
param, grad, regularization
)
def _add_moments_pows(self, p):
acc_dtype = p.dtype
if self._is_dtype_fp16_or_bf16(acc_dtype):
if in_pir_mode():
acc_dtype = DataType.FLOAT32
else:
acc_dtype = core.VarDesc.VarType.FP32
self._add_accumulator(self._moment1_acc_str, p, dtype=acc_dtype)
self._add_accumulator(self._moment2_acc_str, p, dtype=acc_dtype)
if self._amsgrad:
self._add_accumulator(self._moment2_acc_max_str, p, dtype=acc_dtype)
self._add_accumulator(
name=self._beta1_pow_acc_str,
param=p,
dtype=acc_dtype,
fill_value=(
0.9
if isinstance(self._beta1, (Variable, Value))
else self._beta1
),
shape=[1],
type=core.VarDesc.VarType.DENSE_TENSOR,
device='cpu',
)
self._add_accumulator(
name=self._beta2_pow_acc_str,
param=p,
dtype=acc_dtype,
fill_value=(
0.999
if isinstance(self._beta2, (Variable, Value))
else self._beta2
),
shape=[1],
type=core.VarDesc.VarType.DENSE_TENSOR,
device='cpu',
)
def _create_accumulators(self, block, parameters):
assert isinstance(block, (framework.Block, paddle.pir.Block))
if isinstance(parameters, dict):
parameters = self._update_param_group(parameters)
# Create accumulator tensors for first and second moments
for p in parameters:
if p.name in self._already_create_accumulator:
continue
if self._multi_precision and self._is_dtype_fp16_or_bf16(p.dtype):
master_p = self._create_master_weight(p)
self._add_moments_pows(master_p)
self._already_create_accumulator.add(p.name)
continue
if (
self._is_dtype_fp16_or_bf16(p.dtype)
and not self._multi_precision
):
warnings.warn(
"Accumulating with FP16 or BF16 in optimizer can lead to poor accuracy or slow convergence."
"Consider using multi_precision=True option of the Adam optimizer."
)
self._add_moments_pows(p)
self._already_create_accumulator.add(p.name)
def _append_optimize_op(self, block, param_and_grad):
assert isinstance(block, (framework.Block, paddle.pir.Block))
if isinstance(param_and_grad, dict):
param_and_grad = self._update_param_group(param_and_grad)
moment1 = self._get_accumulator_master(
self._moment1_acc_str, param_and_grad[0]
)
moment2 = self._get_accumulator_master(
self._moment2_acc_str, param_and_grad[0]
)
moment2_max = (
self._get_accumulator_master(
self._moment2_acc_max_str, param_and_grad[0]
)
if self._amsgrad
else None
)
beta1_pow_acc = self._get_accumulator_master(
self._beta1_pow_acc_str, param_and_grad[0]
)
beta2_pow_acc = self._get_accumulator_master(
self._beta2_pow_acc_str, param_and_grad[0]
)
find_master = self._multi_precision and self._is_dtype_fp16_or_bf16(
param_and_grad[0].dtype
)
master_weight = (
self._master_weights[param_and_grad[0].name]
if find_master
else None
)
lr = self._create_param_lr(param_and_grad)
# create the adam optimize op
if in_dynamic_or_pir_mode():
_beta1 = (
self._beta1
if not isinstance(self._beta1, Variable)
else self._beta1.item(0)
)
_beta2 = (
self._beta2
if not isinstance(self._beta2, Variable)
else self._beta2.item(0)
)
found_inf = (
self._get_auxiliary_var('found_inf') if in_pir_mode() else None
)
_, _, _, _, _, _, _ = _C_ops.adam_(
param_and_grad[0],
param_and_grad[1],
lr,
moment1,
moment2,
moment2_max,
beta1_pow_acc,
beta2_pow_acc,
master_weight,
found_inf,
_beta1,
_beta2,
self._epsilon,
self._lazy_mode,
1000,
find_master,
False,
self._amsgrad,
)
return None
else:
inputs = {
"Param": [param_and_grad[0]],
"Grad": [param_and_grad[1]],
"LearningRate": [lr],
"Moment1": [moment1],
"Moment2": [moment2],
"Beta1Pow": [beta1_pow_acc],
"Beta2Pow": [beta2_pow_acc],
}
# Pass found_inf to adam, to skip update for not only param, but also momentum and beta_pow
found_inf = self._get_auxiliary_var('found_inf')
if found_inf:
inputs['SkipUpdate'] = found_inf
outputs = {
"ParamOut": [param_and_grad[0]],
"Moment1Out": [moment1],
"Moment2Out": [moment2],
"Beta1PowOut": [beta1_pow_acc],
"Beta2PowOut": [beta2_pow_acc],
}
attrs = {
"lazy_mode": self._lazy_mode,
"min_row_size_to_use_multithread": 1000,
"multi_precision": find_master,
"amsgrad": self._amsgrad,
}
if isinstance(self._beta1, Variable):
inputs['Beta1Tensor'] = self._beta1
else:
attrs['beta1'] = self._beta1
if isinstance(self._beta2, Variable):
inputs['Beta2Tensor'] = self._beta2
else:
attrs['beta2'] = self._beta2
if isinstance(self._epsilon, Variable):
inputs['EpsilonTensor'] = self._epsilon
else:
attrs['epsilon'] = self._epsilon
if self._amsgrad:
inputs['Moment2Max'] = [moment2_max]
outputs["Moment2MaxOut"] = [moment2_max]
if find_master:
inputs["MasterParam"] = master_weight
outputs["MasterParamOut"] = master_weight
adam_op = block.append_op(
type=self.type,
inputs=inputs,
outputs=outputs,
attrs=attrs,
stop_gradient=True,
)
return adam_op
@imperative_base.no_grad
@framework.non_static_only
def step(
self, closure: Callable[[], Tensor] | None = None
) -> Tensor | None:
"""
Execute the optimizer and update parameters once.
Args:
closure (Callable|None, optional): A closure that reevaluates the model
and returns the loss. It should be a callable that takes no arguments
and returns a Tensor. This is useful for optimizers that need to
evaluate the loss multiple times (e.g., line search). Default is None.
Returns:
Tensor|None: If closure is provided, returns the loss value computed by
the closure. Otherwise returns None.
Examples:
.. code-block:: pycon
>>> import paddle
>>> x = paddle.rand([2, 13], dtype="float32")
>>> linear = paddle.nn.Linear(13, 5)
>>> # This can be any optimizer supported by dygraph.
>>> adam = paddle.optimizer.Adam(
... learning_rate=0.01,
... parameters=linear.parameters(),
... )
>>> out = linear(x)
>>> out.backward()
>>> adam.step()
>>> adam.clear_grad()
>>> # usage 1: not use closure
>>> adam.zero_grad()
>>> output = linear(x)
>>> loss = paddle.mean(output)
>>> loss.backward()
>>> adam.step()
>>> # usage 2: use closure
>>> def closure():
... adam.zero_grad()
... output = linear(x)
... loss = paddle.mean(output)
... loss.backward()
... return loss
>>> step_loss = adam.step(closure)
"""
loss = None
if closure is not None:
with imperative_base.enable_grad():
loss = closure()
if paddle.base.dygraph.base.in_to_static_mode():
self._declarative_step()
return loss
if not isinstance(self._parameter_list[0], dict):
params_grads = []
for param in self._parameter_list:
if param.stop_gradient:
continue
if param._grad_ivar() is not None:
grad_var = param._grad_ivar()
if in_dygraph_mode():
if (
hasattr(grad_var, "is_selected_rows")
and grad_var.is_selected_rows()
and self.regularization is not None
):
raise RuntimeError(
"Adam don't support weight_decay with sparse parameters, please set it to None."
)
else:
if (
hasattr(grad_var, "_is_sparse")
and grad_var._is_sparse()
and self.regularization is not None
):
raise RuntimeError(
"Adam don't support weight_decay with sparse parameters, please set it to None."
)
params_grads.append((param, grad_var))
optimize_ops = self._apply_optimize(
loss=None,
startup_program=None,
params_grads=params_grads,
param_group_idx=0,
)
else:
# optimize parameters in groups
for idx, param_group in enumerate(self._param_groups):
params_grads = defaultdict(lambda: [])
for param in param_group['params']:
if param.stop_gradient:
continue
if param._grad_ivar() is not None:
grad_var = param._grad_ivar()
params_grads['params'].append((param, grad_var))
params_grads.update(
{k: v for k, v in param_group.items() if k != 'params'}
)
self._apply_optimize(
loss=None,
startup_program=None,
params_grads=params_grads,
param_group_idx=idx,
)
return loss
def _multi_tensor_init(self, target_block, parameters, param_group_idx):
"""
All parameters used for optimizer (such as: parameters, master_weight, velocity_acc for momentum) calculations are grouped into a python list by data type (bfloat16, float16, float32).
This function will be overridden in the corresponding optimizer file.
Args:
target_block: the block in which the loss tensor is present
parameters: list of parameter tensors for the optimizer
"""
self._create_accumulators(target_block, parameters)
for param in parameters:
moment1 = self._get_accumulator_master(self._moment1_acc_str, param)
moment2 = self._get_accumulator_master(self._moment2_acc_str, param)
moment2_max = (
self._get_accumulator_master(self._moment2_acc_max_str, param)
if self._amsgrad
else None
)
beta1_pow_acc = self._get_accumulator_master(
self._beta1_pow_acc_str, param
)
beta2_pow_acc = self._get_accumulator_master(
self._beta2_pow_acc_str, param
)
if param.dtype == paddle.float32:
self._param_dict['FP32_DenseTensor'][param_group_idx].append(
param
)
self._moment1_dict['FP32_DenseTensor'][param_group_idx].append(
moment1
)
self._moment2_dict['FP32_DenseTensor'][param_group_idx].append(
moment2
)
if self._amsgrad:
self._moment2_max_dict['FP32_DenseTensor'][
param_group_idx
].append(moment2_max)
self._beta1_pow_acc_dict['FP32_DenseTensor'][
param_group_idx
].append(beta1_pow_acc)
self._beta2_pow_acc_dict['FP32_DenseTensor'][
param_group_idx
].append(beta2_pow_acc)
elif self._is_dtype_fp16_or_bf16(param.dtype):
self._param_dict['FP16_DenseTensor'][param_group_idx].append(
param
)
self._moment1_dict['FP16_DenseTensor'][param_group_idx].append(
moment1
)
self._moment2_dict['FP16_DenseTensor'][param_group_idx].append(
moment2
)
if self._amsgrad:
self._moment2_max_dict['FP16_DenseTensor'][
param_group_idx
].append(moment2_max)
self._beta1_pow_acc_dict['FP16_DenseTensor'][
param_group_idx
].append(beta1_pow_acc)
self._beta2_pow_acc_dict['FP16_DenseTensor'][
param_group_idx
].append(beta2_pow_acc)
if self._multi_precision:
self._master_weight_dict['FP16_DenseTensor'][
param_group_idx
].append(self._master_weights[param.name])
else:
self._master_weight_dict['FP16_DenseTensor'] = None
else:
raise ValueError(
"Now multi_tensor_momentum only support fp32, fp16 or bf16 parameters and grad is DENSE_TENSOR."
)
def _append_optimize_multi_tensor_op(
self,
target_block,
parameters_and_grads,
param_group_idx,
):
"""
For Multi Tensor, append optimize merged_operator to block.
"""
assert isinstance(target_block, (framework.Block, pir.Block))
grad_dict = {'FP32_DenseTensor': [], 'FP16_DenseTensor': []}
lr_dict = {'FP32_DenseTensor': [], 'FP16_DenseTensor': []}
if isinstance(parameters_and_grads, list):
if framework.in_dygraph_mode():
params = [pair[0] for pair in parameters_and_grads]
grads_types = core.eager.get_grads_types(params)
for index, tp in enumerate(grads_types):
if tp == core.DataType.FLOAT32:
grad_dict['FP32_DenseTensor'].append(
parameters_and_grads[index][1]
)
lr = self._create_param_lr(parameters_and_grads[index])
lr_dict['FP32_DenseTensor'].append(lr)
elif (
tp == core.DataType.FLOAT16
or tp == core.DataType.BFLOAT16
):
grad_dict['FP16_DenseTensor'].append(
parameters_and_grads[index][1]
)
lr = self._create_param_lr(parameters_and_grads[index])
lr_dict['FP16_DenseTensor'].append(lr)
elif in_pir_mode():
for param_and_grad in parameters_and_grads:
if param_and_grad[1] is None:
continue
if param_and_grad[0].stop_gradient is False:
if (
param_and_grad[0].dtype == DataType.FLOAT32
and param_and_grad[1].is_dense_tensor_type()
):
grad_dict['FP32_DenseTensor'].append(
param_and_grad[1]
)
lr = self._create_param_lr(param_and_grad)
lr_dict['FP32_DenseTensor'].append(lr)
elif (
self._is_dtype_fp16_or_bf16(param_and_grad[0].dtype)
and param_and_grad[1].is_dense_tensor_type()
):
grad_dict['FP16_DenseTensor'].append(
param_and_grad[1]
)
lr = self._create_param_lr(param_and_grad)
lr_dict['FP16_DenseTensor'].append(lr)
else:
for param_and_grad in parameters_and_grads:
if param_and_grad[1] is None:
continue
if param_and_grad[0].stop_gradient is False:
if (
param_and_grad[0].dtype == paddle.float32
and param_and_grad[1].type
== core.VarDesc.VarType.DENSE_TENSOR
):
grad_dict['FP32_DenseTensor'].append(
param_and_grad[1]
)
lr = self._create_param_lr(param_and_grad)
lr_dict['FP32_DenseTensor'].append(lr)
elif (
self._is_dtype_fp16_or_bf16(param_and_grad[0].dtype)
and param_and_grad[1].type
== core.VarDesc.VarType.DENSE_TENSOR
):
grad_dict['FP16_DenseTensor'].append(
param_and_grad[1]
)
lr = self._create_param_lr(param_and_grad)
lr_dict['FP16_DenseTensor'].append(lr)
else:
for param_and_grad in parameters_and_grads['params']:
if param_and_grad[1] is None:
continue
if param_and_grad[0].stop_gradient is False:
param_grad_dict = {}
param_grad_dict['params'] = param_and_grad
param_grad_dict.update(
{
k: v
for k, v in parameters_and_grads.items()
if k != 'params'
}
)
param_and_grad = self._update_param_group(param_grad_dict)
if in_pir_mode():
if (
param_and_grad[0].dtype == DataType.FLOAT32
and param_and_grad[1].is_dense_tensor_type()
):
grad_dict['FP32_DenseTensor'].append(
param_and_grad[1]
)
lr = self._create_param_lr(param_and_grad)
lr_dict['FP32_DenseTensor'].append(lr)
elif (
self._is_dtype_fp16_or_bf16(param_and_grad[0].dtype)
and param_and_grad[1].is_dense_tensor_type()
):
grad_dict['FP16_DenseTensor'].append(
param_and_grad[1]
)
lr = self._create_param_lr(param_and_grad)
lr_dict['FP16_DenseTensor'].append(lr)
else:
if (
param_and_grad[0].dtype == paddle.float32
and param_and_grad[1].type
== core.VarDesc.VarType.DENSE_TENSOR
):
grad_dict['FP32_DenseTensor'].append(
param_and_grad[1]
)
lr = self._create_param_lr(param_and_grad)
lr_dict['FP32_DenseTensor'].append(lr)
elif (
self._is_dtype_fp16_or_bf16(param_and_grad[0].dtype)
and param_and_grad[1].type
== core.VarDesc.VarType.DENSE_TENSOR
):
grad_dict['FP16_DenseTensor'].append(
param_and_grad[1]
)
lr = self._create_param_lr(param_and_grad)
lr_dict['FP16_DenseTensor'].append(lr)
multi_tensor_list = ['FP32_DenseTensor', 'FP16_DenseTensor']
for key in multi_tensor_list:
if len(self._param_dict[key][param_group_idx]) > 0:
find_master = (
self._multi_precision and key == 'FP16_DenseTensor'
)
_beta1 = (
self._beta1
if not isinstance(self._beta1, Variable)
else self._beta1.item(0)
)
_beta2 = (
self._beta2
if not isinstance(self._beta2, Variable)
else self._beta2.item(0)
)
if in_dygraph_mode():
master_weight = self._master_weight_dict[key]
master_weight = (
master_weight[param_group_idx]
if master_weight is not None
else None
)
found_inf = self._get_auxiliary_var('found_inf')
if found_inf:
if isinstance(
found_inf, (core.eager.Tensor, pir.Value)
):
self._set_auxiliary_var('found_inf', True)
else:
if isinstance(
found_inf, (core.eager.Tensor, pir.Value)
):
self._set_auxiliary_var('found_inf', False)
_, _, _, _, _, _, _ = _C_ops.merged_adam_(
self._param_dict[key][param_group_idx],
grad_dict[key],
lr_dict[key],
self._moment1_dict[key][param_group_idx],
self._moment2_dict[key][param_group_idx],
(
self._moment2_max_dict[key][param_group_idx]
if self._amsgrad
else None
),
self._beta1_pow_acc_dict[key][param_group_idx],
self._beta2_pow_acc_dict[key][param_group_idx],
master_weight,
_beta1,
_beta2,
self._epsilon,
find_master,
False,
self._amsgrad,
)
elif in_pir_mode():
master_weight = self._master_weight_dict[key]
master_weight = (
master_weight[param_group_idx]
if master_weight is not None
else None
)
_, _, _, _, _, _, _ = _C_ops.merged_adam_(
self._param_dict[key][param_group_idx],
grad_dict[key],
lr_dict[key],
self._moment1_dict[key][param_group_idx],
self._moment2_dict[key][param_group_idx],
(
self._moment2_max_dict[key][param_group_idx]
if self._amsgrad
else None
),
self._beta1_pow_acc_dict[key][param_group_idx],
self._beta2_pow_acc_dict[key][param_group_idx],
master_weight,
_beta1,
_beta2,
self._epsilon,
find_master,
False,
self._amsgrad,
)
else:
inputs = {
"Param": self._param_dict[key][param_group_idx],
"Grad": grad_dict[key],
"LearningRate": lr_dict[key],
"Moment1": self._moment1_dict[key][param_group_idx],
"Moment2": self._moment2_dict[key][param_group_idx],
"Beta1Pow": self._beta1_pow_acc_dict[key][
param_group_idx
],
"Beta2Pow": self._beta2_pow_acc_dict[key][
param_group_idx
],
}
outputs = {
"ParamOut": self._param_dict[key][param_group_idx],
"Moment1Out": self._moment1_dict[key][param_group_idx],
"Moment2Out": self._moment2_dict[key][param_group_idx],
"Beta1PowOut": self._beta1_pow_acc_dict[key][
param_group_idx
],
"Beta2PowOut": self._beta2_pow_acc_dict[key][
param_group_idx
],
}
attrs = {
"epsilon": self._epsilon,
"beta1": _beta1,
"beta2": _beta2,
"amsgrad": self._amsgrad,
}
if self._amsgrad:
inputs["Moment2Max"] = self._moment2_max_dict[key][
param_group_idx
]
outputs["Moment2MaxOut"] = self._moment2_max_dict[key][
param_group_idx
]
if find_master:
inputs["MasterParam"] = self._master_weight_dict[key][
param_group_idx
]
outputs["MasterParamOut"] = self._master_weight_dict[
key
][param_group_idx]
attrs["multi_precision"] = find_master
target_block.append_op(
type="merged_adam",
inputs=inputs,
outputs=outputs,
attrs=attrs,
stop_gradient=True,
)
def _update_param_group(self, parameters):
self._beta1 = parameters.get('beta1', self._default_dict['beta1'])
self._beta2 = parameters.get('beta2', self._default_dict['beta2'])
self._epsilon = parameters.get('epsilon', self._default_dict['epsilon'])
self._lazy_mode = parameters.get(
'lazy_mode', self._default_dict['lazy_mode']
)
parameters = parameters.get('params')
return parameters
python/paddle/optimizer/adamax.py
+398
View File
@@ -0,0 +1,398 @@
# 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 typing import TYPE_CHECKING
from paddle import _C_ops, pir
from ..base import core, framework
from ..base.dygraph import no_grad
from ..base.framework import name_scope
from .optimizer import Optimizer
if TYPE_CHECKING:
from collections.abc import Sequence
from typing_extensions import NotRequired
from paddle import Tensor
from paddle.nn.clip import GradientClipBase
from paddle.regularizer import WeightDecayRegularizer
from .lr import LRScheduler
from .optimizer import _ParameterConfig
class _AdamaxParameterConfig(_ParameterConfig):
beta1: NotRequired[float | Tensor]
beta2: NotRequired[float | Tensor]
epsilon: NotRequired[float | Tensor]
__all__ = []
class Adamax(Optimizer):
r"""
The Adamax optimizer is implemented based on the Adamax Optimization
in Section 7 of `Adam paper <https://arxiv.org/abs/1412.6980>`_.
The Adamax algorithm is a variant of the Adam algorithm based on the infinite norm,
which makes the learning rate update algorithm more stable and simple.
The parameter ``param_out`` update rule with gradient ``grad``:
.. math::
t & = t + 1
moment\_out & = {\beta}_1 * moment + (1 - {\beta}_1) * grad
inf\_norm\_out & = max({\beta}_2 * inf\_norm + \epsilon, |grad|)
learning\_rate & = \frac{learning\_rate}{1 - {\beta}_1^t}
param\_out & = param - learning\_rate * \frac{moment\_out}{inf\_norm\_out}
Related paper: `Adam: A Method for Stochastic Optimization <https://arxiv.org/abs/1412.6980>`_
The original paper does not have an ``epsilon`` attribute,
it is added here for numerical stability to prevent the division by 0 error.
Args:
learning_rate (float|LRScheduler, optional): The learning rate used to update ``Parameter``.
It can be a float value or a LRScheduler. The default value is 0.001.
beta1 (float|Tensor, optional): The exponential decay rate for the 1st moment estimates.
It should be a float number or a 0-D Tensor with shape [] and data type as float32.
The default value is 0.9.
beta2 (float|Tensor, optional): The exponential decay rate for the 2nd moment estimates.
It should be a float number or a 0-D Tensor with shape [] and data type as float32.
The default value is 0.999.
epsilon (float|Tensor, optional): A small float value for numerical stability.
The default value is 1e-08.
parameters (list|tuple|None, optional): List/Tuple of ``Tensor`` to update to minimize ``loss``.
This parameter is required in dygraph mode. And you can specify different options for
different parameter groups such as the learning rate, weight decay, etc,
then the parameters are list of dict. Note that the learning_rate in parameter groups
represents the scale of base learning_rate.
The default value is None in static graph mode, at this time all parameters will be updated.
weight_decay (int|float|WeightDecayRegularizer|None, optional): The strategy of regularization.
It can be a int or float value as coeff of L2 regularization or
:ref:`api_paddle_regularizer_L1Decay`, :ref:`api_paddle_regularizer_L2Decay`.
If a parameter has set regularizer using :ref:`api_paddle_ParamAttr` already,
the regularization setting here in optimizer will be ignored for this parameter.
Otherwise, the regularization setting here in optimizer will take effect.
Default None, meaning there is no regularization.
grad_clip (GradientClipBase|None, optional): Gradient clipping strategy, it's an instance of
some derived class of ``GradientClipBase`` . There are three clipping strategies
( :ref:`api_paddle_nn_ClipGradByGlobalNorm` , :ref:`api_paddle_nn_ClipGradByNorm` ,
:ref:`api_paddle_nn_ClipGradByValue` ). Default None, meaning there is no gradient clipping.
name (str|None, optional): Normally there is no need for user to set this property.
For more information, please refer to :ref:`api_guide_Name`.
The default value is None.
**Notes**:
**Currently, Adamax doesn't support sparse parameter optimization.**
Examples:
.. code-block:: pycon
>>> import paddle
>>> inp = paddle.uniform([10, 10], dtype="float32", min=-0.1, max=0.1)
>>> linear = paddle.nn.Linear(10, 10)
>>> inp = paddle.to_tensor(inp)
>>> out = linear(inp)
>>> loss = paddle.mean(out)
>>> beta1 = paddle.to_tensor([0.9], dtype="float32")
>>> beta2 = paddle.to_tensor([0.99], dtype="float32")
>>> adamax = paddle.optimizer.Adamax(
... learning_rate=0.1,
... parameters=linear.parameters(),
... beta1=beta1,
... beta2=beta2,
... weight_decay=0.01,
... )
>>> out.backward()
>>> adamax.step()
>>> adamax.clear_grad()
>>> # Note that the learning_rate of linear_2 is 0.01.
>>> linear_1 = paddle.nn.Linear(10, 10)
>>> linear_2 = paddle.nn.Linear(10, 10)
>>> inp = paddle.uniform(shape=[10, 10], min=-0.1, max=0.1)
>>> out = linear_1(inp)
>>> out = linear_2(out)
>>> loss = paddle.mean(out)
>>> adamax = paddle.optimizer.Adamax(
... learning_rate=0.1,
... parameters=[ # type: ignore
... {
... 'params': linear_1.parameters(),
... },
... {
... 'params': linear_2.parameters(),
... 'weight_decay': 0.001,
... 'learning_rate': 0.1,
... 'beta1': 0.8,
... },
... ],
... weight_decay=0.01,
... beta1=0.9,
... )
>>> out.backward()
>>> adamax.step()
>>> adamax.clear_grad()
"""
type: str
_moment_acc_str = "moment"
_inf_norm_acc_str = "inf_norm"
_beta1_pow_acc_str = "beta1_pow_acc"
def __init__(
self,
learning_rate: float | LRScheduler = 0.001,
beta1: float | Tensor = 0.9,
beta2: float | Tensor = 0.999,
epsilon: float | Tensor = 1e-8,
parameters: (
Sequence[Tensor] | Sequence[_AdamaxParameterConfig] | None
) = None,
weight_decay: float | WeightDecayRegularizer | None = None,
grad_clip: GradientClipBase | None = None,
name: str | None = None,
) -> None:
assert learning_rate is not None
assert beta1 is not None
assert beta2 is not None
assert epsilon is not None
if not 0 <= beta1 < 1:
raise ValueError("Invalid value of beta1, expect beta1 in [0,1).")
if not 0 <= beta2 < 1:
raise ValueError("Invalid value of beta2, expect beta2 in [0,1).")
if not 0 <= epsilon:
raise ValueError("Invalid value of epsilon, expect epsilon >= 0.")
super().__init__(
learning_rate=learning_rate,
parameters=parameters,
weight_decay=weight_decay,
grad_clip=grad_clip,
name=name,
)
self.type = "adamax"
self._beta1 = beta1
self._beta2 = beta2
self._epsilon = epsilon
self._multi_precision = False
self._master_weights = {}
self._default_dict = {
'beta1': beta1,
'beta2': beta2,
'epsilon': epsilon,
}
def _add_moments_pows(self, p):
acc_dtype = p.dtype
if self._is_dtype_fp16_or_bf16(acc_dtype):
acc_dtype = core.VarDesc.VarType.FP32
self._add_accumulator(self._moment_acc_str, p, dtype=acc_dtype)
self._add_accumulator(self._inf_norm_acc_str, p, dtype=acc_dtype)
self._add_accumulator(
name=self._beta1_pow_acc_str,
param=p,
fill_value=self._beta1,
shape=[1],
)
def _create_accumulators(self, block, parameters):
if isinstance(parameters, dict):
parameters = self._update_param_group(parameters)
# Create accumulator tensors for first moment and infinity norm
for p in parameters:
if p.name in self._already_create_accumulator:
continue
if self._multi_precision and self._is_dtype_fp16_or_bf16(p.dtype):
master_p = self._create_master_weight(p)
self._add_moments_pows(master_p)
self._already_create_accumulator.add(p.name)
continue
if (
self._is_dtype_fp16_or_bf16(p.dtype)
and not self._multi_precision
):
warnings.warn(
"Accumulating with FP16/BF16 in optimizer can lead to poor accuracy or slow convergence."
"Consider using multi_precision=True option of the Adam optimizer."
)
self._add_moments_pows(p)
self._already_create_accumulator.add(p.name)
def _append_optimize_op(self, block, param_and_grad):
assert isinstance(block, (framework.Block, pir.Block))
if isinstance(param_and_grad, dict):
param_and_grad = self._update_param_group(param_and_grad)
moment = self._get_accumulator_master(
self._moment_acc_str, param_and_grad[0]
)
inf_norm = self._get_accumulator_master(
self._inf_norm_acc_str, param_and_grad[0]
)
find_master = self._multi_precision and self._is_dtype_fp16_or_bf16(
param_and_grad[0].dtype
)
master_weight = (
self._master_weights[param_and_grad[0].name]
if find_master
else None
)
beta1_pow_acc = self._get_accumulator_master(
self._beta1_pow_acc_str, param_and_grad[0]
)
if framework.in_dynamic_or_pir_mode():
_C_ops.adamax_(
param_and_grad[0],
param_and_grad[1],
self._create_param_lr(param_and_grad),
moment,
inf_norm,
beta1_pow_acc,
master_weight,
self._beta1,
self._beta2,
self._epsilon,
find_master,
)
else:
# create the adamax optimize op
inputs = {
"Param": param_and_grad[0],
"Grad": param_and_grad[1],
"LearningRate": self._create_param_lr(param_and_grad),
"Moment": moment,
"InfNorm": inf_norm,
"Beta1Pow": beta1_pow_acc,
}
outputs = {
"ParamOut": param_and_grad[0],
"MomentOut": moment,
"InfNormOut": inf_norm,
}
if find_master:
inputs["MasterParam"] = master_weight
outputs["MasterParamOut"] = master_weight
attrs = {
"beta1": self._beta1,
"beta2": self._beta2,
"epsilon": self._epsilon,
"multi_precision": find_master,
}
adamax_op = block.append_op(
type=self.type,
inputs=inputs,
outputs=outputs,
attrs=attrs,
stop_gradient=True,
)
return adamax_op
def _finish_update(self, block, parameters_and_grads):
"""Update Beta1 Power accumulator"""
assert isinstance(block, (framework.Block, pir.Block))
if isinstance(parameters_and_grads, list):
for param, grad in parameters_and_grads:
if grad is None or param.stop_gradient is True:
continue
if framework.in_dygraph_mode():
beta1_pow_acc = self._get_accumulator_master(
self._beta1_pow_acc_str, param
)
with no_grad():
tmp = _C_ops.scale(
beta1_pow_acc, self._beta1, 0.0, True
)
beta1_pow_acc.copy_(tmp, False)
elif framework.in_pir_mode():
with param.block.program._optimized_guard([param, grad]):
beta1_pow_acc = self._get_accumulator_master(
self._beta1_pow_acc_str, param
)
_C_ops.scale_(beta1_pow_acc, self._beta1, 0.0, True)
else:
with (
param.block.program._optimized_guard([param, grad]),
name_scope('adamax'),
):
beta1_pow_acc = self._get_accumulator_master(
self._beta1_pow_acc_str, param
)
block.append_op(
type="scale",
inputs={"X": beta1_pow_acc},
outputs={"Out": beta1_pow_acc},
attrs={"scale": self._beta1},
stop_gradient=True,
)
else:
for param, grad in parameters_and_grads['params']:
if grad is None or param.stop_gradient is True:
continue
if framework.in_dygraph_mode():
beta1_pow_acc = self._get_accumulator_master(
self._beta1_pow_acc_str, param
)
self._beta1 = parameters_and_grads.get(
'beta1', self._default_dict['beta1']
)
with no_grad():
tmp = _C_ops.scale(
beta1_pow_acc, self._beta1, 0.0, True
)
beta1_pow_acc.copy_(tmp, False)
else:
with (
param.block.program._optimized_guard([param, grad]),
name_scope('adamax'),
):
beta1_pow_acc = self._get_accumulator_master(
self._beta1_pow_acc_str, param
)
self._beta1 = parameters_and_grads.get(
'beta1', self._default_dict['beta1']
)
block.append_op(
type="scale",
inputs={"X": beta1_pow_acc},
outputs={"Out": beta1_pow_acc},
attrs={"scale": self._beta1},
stop_gradient=True,
)
def _update_param_group(self, parameters):
self._beta1 = parameters.get('beta1', self._default_dict['beta1'])
self._beta2 = parameters.get('beta2', self._default_dict['beta2'])
self._epsilon = parameters.get('epsilon', self._default_dict['epsilon'])
parameters = parameters.get('params')
return parameters
python/paddle/optimizer/adamw.py
+880
View File
@@ -0,0 +1,880 @@
# Copyright (c) 2021 PaddlePaddle Authors. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
from __future__ import annotations
import warnings
from collections import defaultdict
from collections.abc import Callable
from typing import TYPE_CHECKING
import paddle
from paddle import pir
from paddle.base.libpaddle import DataType
from paddle.distributed.flex_checkpoint.dcp.sharded_weight import (
ShardedStateDict,
ShardedWeight,
create_sharded_weight_with_new_local,
)
from paddle.pir import Value
from .. import _C_ops
from ..base import core, framework
from ..base.dygraph import base as imperative_base
from ..base.framework import (
Parameter,
Variable,
in_dynamic_or_pir_mode,
in_pir_mode,
)
from ..nn.clip import GradientClipBase
from .lr import LRScheduler
from .optimizer import Optimizer
if TYPE_CHECKING:
from collections.abc import Sequence
from paddle import Tensor
from .adam import _AdamParameterConfig
__all__ = []
class AdamW(Optimizer):
r"""
The AdamW optimizer is implemented based on the AdamW Optimization
in paper `DECOUPLED WEIGHT DECAY REGULARIZATION <https://arxiv.org/pdf/1711.05101.pdf>`_.
it can resolves the problem of L2 regularization failure in the Adam optimizer.
.. math::
\begin{aligned}
&\hspace{5mm} t = t + 1 \\
&\hspace{5mm} moment\_1\_out = {\beta}_1 * moment\_1 + (1 - {\beta}_1) * grad \\
&\hspace{5mm} moment\_2\_out = {\beta}_2 * moment\_2 + (1 - {\beta}_2) * grad * grad \\
&\hspace{5mm} learning\_rate = learning\_rate * \frac{\sqrt{1 - {\beta}_2^t}}{1 - {\beta}_1^t} \\
&\hspace{5mm}\textbf{if} \: \textit{amsgrad}: \\
&\hspace{15mm} moment\_2\_max\_out = max(moment\_2\_out, moment\_2\_max) \\
&\hspace{15mm} param\_out = param - learning\_rate * (\frac{moment\_1\_out}{\sqrt{moment\_2\_max\_out} + \epsilon} + \lambda * param) \\
&\hspace{5mm}\textbf{else}: \: \\
&\hspace{15mm} param\_out = param - learning\_rate * (\frac{moment\_1\_out}{\sqrt{moment\_2\_out} + \epsilon} + \lambda * param) \\
\end{aligned}
Args:
learning_rate (float|LRScheduler, optional): The learning rate used to update ``Parameter``.
It can be a float value or a LRScheduler. The default value is 0.001.
beta1 (float|Tensor, optional): The exponential decay rate for the 1st moment estimates.
It should be a float number or a 0-D Tensor with shape [] and data type as float32.
The default value is 0.9.
beta2 (float|Tensor, optional): The exponential decay rate for the 2nd moment estimates.
It should be a float number or a 0-D Tensor with shape [] and data type as float32.
The default value is 0.999.
epsilon (float|Tensor, optional): A small float value for numerical stability.
The default value is 1e-08.
parameters (list|tuple|None, optional): List/Tuple of ``Tensor`` names to update to minimize ``loss``.
This parameter is required in dygraph mode. And you can specify different options for
different parameter groups such as the learning rate, weight decay, etc,
then the parameters are list of dict. Note that the learning_rate in parameter groups
represents the scale of base learning_rate.
The default value is None in static graph mode, at this time all parameters will be updated.
weight_decay (int|float|Tensor, optional): The weight decay coefficient, it can be int, float or Tensor. The default value is 0.01.
lr_ratio (Callable|None, optional): If it is not None,
the learning rate will be updated with layer-wise learning rate ratio.
Otherwise, the learning rate is the original.
Default: None.
apply_decay_param_fun (Callable|None, optional): If it is not None,
only tensors that makes apply_decay_param_fun(Tensor.name)==True
will be updated with weight decay. It only works when we want to specify tensors.
Default: None.
grad_clip (GradientClipBase|None, optional): Gradient clipping strategy, it's an instance of
some derived class of ``GradientClipBase`` . There are three clipping strategies
( :ref:`api_paddle_nn_ClipGradByGlobalNorm` , :ref:`api_paddle_nn_ClipGradByNorm` ,
:ref:`api_paddle_nn_ClipGradByValue` ). Default None, meaning there is no gradient clipping.
lazy_mode (bool, optional): The official Adam algorithm has two moving-average accumulators.
The accumulators are updated at every step. Every element of the two moving-average
is updated in both dense mode and sparse mode. If the size of parameter is very large,
then the update may be very slow. The lazy mode only update the element that has
gradient in current mini-batch, so it will be much more faster. But this mode has
different semantics with the original Adam algorithm and may lead to different result.
The default value is False.
multi_precision (bool, optional): Whether to use multi-precision during weight updating. Default is false.
amsgrad (bool, optional): Whether to use the AMSGrad variant of this algorithm from the paper
`On the Convergence of Adam and Beyond <https://openreview.net/forum?id=ryQu7f-RZ>`_. Default is false.
name (str|None, optional): Normally there is no need for user to set this property.
For more information, please refer to :ref:`api_guide_Name`.
The default value is None.
Notes:
**Currently, AdamW doesn't support sparse parameter optimization.**
Examples:
.. code-block:: pycon
>>> import paddle
>>> linear = paddle.nn.Linear(10, 10)
>>> inp = paddle.rand([10,10], dtype="float32")
>>> out = linear(inp)
>>> loss = paddle.mean(out)
>>> beta1 = paddle.to_tensor([0.9], dtype="float32")
>>> beta2 = paddle.to_tensor([0.99], dtype="float32")
>>> opt = paddle.optimizer.AdamW(
... learning_rate=0.1,
... parameters=linear.parameters(),
... beta1=beta1,
... beta2=beta2,
... weight_decay=0.01
... )
>>> loss.backward()
>>> opt.step()
>>> opt.clear_grad()
>>> # Note that the learning_rate of linear_2 is 0.01.
>>> linear_1 = paddle.nn.Linear(10, 10)
>>> linear_2 = paddle.nn.Linear(10, 10)
>>> inp = paddle.uniform(shape=[10, 10], min=-0.1, max=0.1)
>>> out = linear_1(inp)
>>> out = linear_2(out)
>>> loss = paddle.mean(out)
>>> opt = paddle.optimizer.AdamW(
... learning_rate=0.1,
... parameters=[{ # type: ignore
... 'params': linear_1.parameters()
... }, {
... 'params': linear_2.parameters(),
... 'weight_decay': 0.001,
... 'learning_rate': 0.1,
... 'beta1': 0.8
... }],
... weight_decay=0.01,
... beta1=0.9
... )
>>> loss.backward()
>>> opt.step()
>>> opt.clear_grad()
"""
helper: None
type: str
_moment1_acc_str = "moment1"
_moment2_acc_str = "moment2"
_moment2_acc_max_str = "moment2_max"
_beta1_pow_acc_str = "beta1_pow_acc"
_beta2_pow_acc_str = "beta2_pow_acc"
def __init__(
self,
learning_rate: float | LRScheduler = 0.001,
beta1: float | Tensor = 0.9,
beta2: float | Tensor = 0.999,
epsilon: float | Tensor = 1e-8,
parameters: (
Sequence[Tensor] | Sequence[_AdamParameterConfig] | None
) = None,
weight_decay: float | Tensor = 0.01,
use_lowprecision_moment: bool = False,
lr_ratio: Callable[[Tensor], float] | None = None,
apply_decay_param_fun: Callable[[str], bool] | None = None,
grad_clip: GradientClipBase | None = None,
lazy_mode: bool = False,
multi_precision: bool = False,
amsgrad: bool = False,
name: str | None = None,
) -> None:
assert learning_rate is not None
assert beta1 is not None
assert beta2 is not None
assert epsilon is not None
if not isinstance(beta1, Value) and not 0 <= beta1 < 1:
raise ValueError("Invalid value of beta1, expect beta1 in [0,1).")
if not isinstance(beta2, Value) and not 0 <= beta2 < 1:
raise ValueError("Invalid value of beta2, expect beta2 in [0,1).")
if not isinstance(epsilon, Value) and not 0 <= epsilon:
raise ValueError("Invalid value of epsilon, expect epsilon >= 0.")
if not isinstance(weight_decay, (int, float)) and not isinstance(
weight_decay, (framework.Variable, Value)
):
raise TypeError("weight_decay should be int, float or Tensor.")
if lr_ratio is not None:
assert isinstance(lr_ratio, Callable)
if (
not core.is_compiled_with_cuda()
and not core.is_compiled_with_xpu()
and paddle.device.get_device().split(":")[0]
not in paddle.device.get_all_custom_device_type()
):
raise NotImplementedError("'lr_ratio' is unimplemented in CPU.")
if parameters is not None:
# paddle.Tensor is also iterable, so here we don't check whether
# the input is iterable, if the input is paddle.Tensor, the
# list(paddle.Tensor) will be a error value
if isinstance(parameters, paddle.Tensor):
raise TypeError(
"`parameters` argument given to the optimizer should be "
f"an iterable of paddle Tensors, but got argument type is `{type(parameters)}`."
)
if isinstance(parameters, dict):
raise TypeError(
"`parameters` argument should not get dict type, "
"if parameter groups is needed, please set `parameters`"
" as list of dict"
)
self._parameter_list = list(parameters)
else:
self._parameter_list = None
self._name = name
if framework.in_dygraph_mode():
if self._parameter_list is None:
raise AttributeError(
"parameters argument given to the Optimizer should not be None in dygraph mode."
)
if not isinstance(learning_rate, (float, LRScheduler)):
raise TypeError(
f"learning rate should be float or LRScheduler, got {type(learning_rate)} here"
)
if grad_clip is not None:
if not isinstance(grad_clip, GradientClipBase):
raise TypeError(
"'grad_clip' should be an instance of GradientClipBase's derived class"
)
self._dtype = None
# Infer the dtype form parameter
if self._parameter_list:
if isinstance(self._parameter_list[0], dict):
for param_group in self._parameter_list:
assert 'params' in param_group, (
'params should be set in parameters if parameter groups are optimized in different options'
)
self._dtype = self._parameter_list[0]['params'][0].dtype
else:
self._dtype = self._parameter_list[0].dtype
# each program should have a independent learning rate
# program -> tensor(learning_rate)
self._learning_rate_map = {}
# Dictionary of accumulators. Some optimizer subclasses need to
# allocate and manage extra tensors associated with the parameters
# to train. These tensors are called accumulators.
# {accum_name : { parameter_name : accumulator_for_parameter, ...}, ...}
self._accumulators = defaultdict(lambda: {})
self.helper = None
self._opti_name_list = []
self._accumulators_holder = {}
self._param_device_map = {}
self.clear_gradients = self.clear_grad
self.type = "adamw"
self._learning_rate = learning_rate
self._params_name = set()
self._apply_decay_param_fun = apply_decay_param_fun
self._weight_decay = float(weight_decay)
self._use_lowprecision_moment = use_lowprecision_moment
self._grad_clip = grad_clip
self._lr_ratio = lr_ratio
self._beta1 = beta1
self._beta2 = beta2
self._epsilon = epsilon
self._lazy_mode = lazy_mode
self._multi_precision = multi_precision
self._master_weights = {}
# whether to use AMSGrad
self._amsgrad = amsgrad
self._default_dict = {
'weight_decay': float(weight_decay),
'beta1': beta1,
'beta2': beta2,
'epsilon': epsilon,
'lazy_mode': lazy_mode,
'grad_clip': grad_clip,
}
self._param_groups = []
if self._parameter_list and isinstance(self._parameter_list[0], dict):
for param_group in self._parameter_list:
self._add_param_group(param_group.copy())
else:
self._param_groups = self._parameter_list
self._use_multi_tensor = None
self.regularization = None
self._auxiliary_vars = {}
self._already_create_accumulator = set()
self._create_master_grad_states()
self._use_fusion_storage = False
self._need_refuse = False
self.fusion_storage = None
self._fuse_buffer_version = 0
self.merged_model_params = None
def _set_auxiliary_var(self, key, val):
self._auxiliary_vars[key] = val
def _get_auxiliary_var(self, key):
if key in self._auxiliary_vars:
return self._auxiliary_vars[key]
else:
return None
def get_lr_dtype(self) -> paddle.dtype:
return paddle.float64
def _add_param_group(self, param_group):
"""
Add a param group to parameter_list.
Args:
param_group (dict): The group of Tensors to be optimized with
different optimization options.
"""
params = param_group['params']
if isinstance(params, (Parameter, pir.core.ParameterMeta)):
param_group['params'] = [params]
elif isinstance(params, set):
raise TypeError(
"optimizer parameters should be in ordered collections,"
"but received set, please use list instead."
)
else:
param_group['params'] = list(params)
# Update optimization options for each groups
for k, v in self._default_dict.items():
param_group.setdefault(k, v)
param_set = set()
for group in self._param_groups:
param_set.update(set(group['params']))
if not param_set.isdisjoint(set(param_group['params'])):
raise ValueError(
"some parameters appear in more than one parameter group"
)
for param in param_group['params']:
param.optimize_attr['learning_rate'] = param_group.get(
'learning_rate', 1.0
)
self._param_groups.append(param_group)
def _add_moments_pows(self, p):
acc_dtype = p.dtype
if (
self._is_dtype_fp16_or_bf16(acc_dtype)
and not self._use_lowprecision_moment
):
acc_dtype = (
DataType.FLOAT32 if in_pir_mode() else core.VarDesc.VarType.FP32
)
if core.is_compiled_with_xpu():
import os
xpu_adamw_moment_dtype = os.getenv(
"xpu_adamw_moment_dtype", default="fp32"
)
if xpu_adamw_moment_dtype == "fp16":
self._add_accumulator(
self._moment1_acc_str, p, dtype=core.VarDesc.VarType.FP16
)
self._add_accumulator(
self._moment2_acc_str, p, dtype=core.VarDesc.VarType.FP16
)
if self._amsgrad:
self._add_accumulator(
self._moment2_acc_max_str,
p,
dtype=core.VarDesc.VarType.FP16,
)
else:
self._add_accumulator(self._moment1_acc_str, p, dtype=acc_dtype)
self._add_accumulator(self._moment2_acc_str, p, dtype=acc_dtype)
if self._amsgrad:
self._add_accumulator(
self._moment2_acc_max_str, p, dtype=acc_dtype
)
else:
self._add_accumulator(self._moment1_acc_str, p, dtype=acc_dtype)
self._add_accumulator(self._moment2_acc_str, p, dtype=acc_dtype)
if self._amsgrad:
self._add_accumulator(
self._moment2_acc_max_str, p, dtype=acc_dtype
)
self._add_accumulator(
name=self._beta1_pow_acc_str,
param=p,
dtype=acc_dtype,
fill_value=(
0.9
if isinstance(self._beta1, (Variable, Value))
else self._beta1
),
shape=[1],
type=core.VarDesc.VarType.DENSE_TENSOR,
device='cpu',
)
self._add_accumulator(
name=self._beta2_pow_acc_str,
param=p,
dtype=acc_dtype,
fill_value=(
0.999
if isinstance(self._beta2, (Variable, Value))
else self._beta2
),
shape=[1],
type=core.VarDesc.VarType.DENSE_TENSOR,
device='cpu',
)
def _create_accumulators(self, block, parameters):
assert isinstance(block, (framework.Block, pir.Block))
if isinstance(parameters, dict):
parameters = self._update_param_group(parameters)
# Create accumulator tensors for first and second moments
for p in parameters:
if p.name in self._already_create_accumulator:
continue
if self._multi_precision and self._is_dtype_fp16_or_bf16(p.dtype):
master_p = self._create_master_weight(p)
self._add_moments_pows(master_p)
self._already_create_accumulator.add(p.name)
continue
if (
self._is_dtype_fp16_or_bf16(p.dtype)
and not self._multi_precision
):
warnings.warn(
"Accumulating with FP16 or BF16 in optimizer can lead to poor accuracy or slow convergence."
"Consider using multi_precision=True option of the Adam optimizer."
)
self._add_moments_pows(p)
self._already_create_accumulator.add(p.name)
def _append_optimize_op(self, block, param_and_grad):
assert isinstance(block, (framework.Block, pir.Block))
if isinstance(param_and_grad, dict):
param_and_grad = self._update_param_group(param_and_grad)
param, grad = param_and_grad
# Whether we should do weight decay for the parameter.
with_decay = True
if (
self._apply_decay_param_fun is not None
and not self._apply_decay_param_fun(param.name)
):
with_decay = False
moment1 = self._get_accumulator_master(
self._moment1_acc_str, param_and_grad[0]
)
moment2 = self._get_accumulator_master(
self._moment2_acc_str, param_and_grad[0]
)
moment2_max = (
self._get_accumulator_master(
self._moment2_acc_max_str, param_and_grad[0]
)
if self._amsgrad
else None
)
beta1_pow_acc = self._get_accumulator_master(
self._beta1_pow_acc_str, param_and_grad[0]
)
beta2_pow_acc = self._get_accumulator_master(
self._beta2_pow_acc_str, param_and_grad[0]
)
find_master = self._multi_precision and self._is_dtype_fp16_or_bf16(
param_and_grad[0].dtype
)
master_weight = (
self._master_weights[param_and_grad[0].name]
if find_master
else None
)
lr = self._create_param_lr(param_and_grad)
# create the adamw optimize op
if in_dynamic_or_pir_mode():
lr_ratio_ = (
1.0
if self._lr_ratio is None
else self._lr_ratio(param_and_grad[0])
)
_beta1 = (
self._beta1
if not isinstance(self._beta1, Variable)
else self._beta1.item(0)
)
_beta2 = (
self._beta2
if not isinstance(self._beta2, Variable)
else self._beta2.item(0)
)
found_inf = (
self._get_auxiliary_var('found_inf') if in_pir_mode() else None
)
_, _, _, _, _, _, _ = _C_ops.adamw_(
param_and_grad[0],
param_and_grad[1],
lr,
moment1,
moment2,
moment2_max,
beta1_pow_acc,
beta2_pow_acc,
master_weight,
found_inf,
_beta1,
_beta2,
self._epsilon,
lr_ratio_,
self._weight_decay,
with_decay,
self._lazy_mode,
1000,
find_master,
False,
self._amsgrad,
)
return None
else:
inputs = {
"Param": [param_and_grad[0]],
"Grad": [param_and_grad[1]],
"LearningRate": [lr],
"Moment1": [moment1],
"Moment2": [moment2],
"Beta1Pow": [beta1_pow_acc],
"Beta2Pow": [beta2_pow_acc],
}
# Pass found_inf to adamw, to skip update for not only param, but also momentum and beta_pow
found_inf = self._get_auxiliary_var('found_inf')
if found_inf:
inputs['SkipUpdate'] = found_inf
outputs = {
"ParamOut": [param_and_grad[0]],
"Moment1Out": [moment1],
"Moment2Out": [moment2],
"Beta1PowOut": [beta1_pow_acc],
"Beta2PowOut": [beta2_pow_acc],
}
attrs = {
"lazy_mode": self._lazy_mode,
"min_row_size_to_use_multithread": 1000,
"multi_precision": find_master,
"with_decay": with_decay,
"coeff": self._weight_decay,
"lr_ratio": (
1.0
if self._lr_ratio is None
else self._lr_ratio(param_and_grad[0])
),
"amsgrad": self._amsgrad,
}
if isinstance(self._beta1, Variable):
inputs['Beta1Tensor'] = self._beta1
else:
attrs['beta1'] = self._beta1
if isinstance(self._beta2, Variable):
inputs['Beta2Tensor'] = self._beta2
else:
attrs['beta2'] = self._beta2
if isinstance(self._epsilon, Variable):
inputs['EpsilonTensor'] = self._epsilon
else:
attrs['epsilon'] = self._epsilon
if self._amsgrad:
inputs["Moment2Max"] = [moment2_max]
outputs["Moment2MaxOut"] = [moment2_max]
if find_master:
inputs["MasterParam"] = master_weight
outputs["MasterParamOut"] = master_weight
adamw_op = block.append_op(
type=self.type,
inputs=inputs,
outputs=outputs,
attrs=attrs,
stop_gradient=True,
)
return adamw_op
def __str__(self):
return " ".join(["Weight Decay, params:", ",".join(self._params_name)])
@imperative_base.no_grad
@framework.non_static_only
def step(
self, closure: Callable[[], Tensor] | None = None
) -> Tensor | None:
"""
Execute the optimizer and update parameters once.
Args:
closure (Callable|None, optional): A closure that reevaluates the model
and returns the loss. It should be a callable that takes no arguments
and returns a Tensor. This is useful for optimizers that need to
evaluate the loss multiple times (e.g., line search). Default is None.
Returns:
Tensor|None: If closure is provided, returns the loss value computed by
the closure. Otherwise returns None.
Examples:
.. code-block:: pycon
>>> import paddle
>>> x = paddle.rand([2, 13], dtype="float32")
>>> linear = paddle.nn.Linear(13, 5)
>>> # This can be any optimizer supported by dygraph.
>>> opt = paddle.optimizer.AdamW(
... learning_rate=0.01,
... parameters=linear.parameters(),
... )
>>> out = linear(x)
>>> out.backward()
>>> opt.step()
>>> opt.clear_grad()
>>> # usage 1: not use closure
>>> opt.zero_grad()
>>> output = linear(x)
>>> loss = paddle.mean(output)
>>> loss.backward()
>>> opt.step()
>>> # usage 2: use closure
>>> def closure():
... opt.zero_grad()
... output = linear(x)
... loss = paddle.mean(output)
... loss.backward()
... return loss
>>> step_loss = opt.step(closure)
"""
loss = None
if closure is not None:
with imperative_base.enable_grad():
loss = closure()
if paddle.base.dygraph.base.in_to_static_mode():
self._declarative_step()
return loss
if not isinstance(self._parameter_list[0], dict):
params_grads = []
for param in self._parameter_list:
if param.stop_gradient:
continue
if param._grad_ivar() is not None:
grad_var = param._grad_ivar()
if framework.in_dygraph_mode():
if (
hasattr(grad_var, "is_selected_rows")
and grad_var.is_selected_rows()
and self.regularization is not None
):
raise RuntimeError(
"AdamW don't support weight_decay with sparse parameters, please set it to None."
)
else:
if (
hasattr(grad_var, "_is_sparse")
and grad_var._is_sparse()
and self.regularization is not None
):
raise RuntimeError(
"AdamW don't support weight_decay with sparse parameters, please set it to None."
)
params_grads.append((param, grad_var))
optimize_ops = self._apply_optimize(
loss=None, startup_program=None, params_grads=params_grads
)
else:
# optimize parameters in groups
for param_group in self._param_groups:
params_grads = defaultdict(lambda: [])
for param in param_group['params']:
if param.stop_gradient:
continue
if param._grad_ivar() is not None:
grad_var = param._grad_ivar()
if framework.in_dygraph_mode():
if (
hasattr(grad_var, "is_selected_rows")
and grad_var.is_selected_rows()
and self.regularization is not None
):
raise RuntimeError(
"AdamW don't support weight_decay with sparse parameters, please set it to None."
)
else:
if (
hasattr(grad_var, "_is_sparse")
and grad_var._is_sparse()
and self.regularization is not None
):
raise RuntimeError(
"AdamW don't support weight_decay with sparse parameters, please set it to None."
)
params_grads['params'].append((param, grad_var))
params_grads.update(
{k: v for k, v in param_group.items() if k != 'params'}
)
self._apply_optimize(
loss=None, startup_program=None, params_grads=params_grads
)
return loss
def _update_param_group(self, parameters):
self._beta1 = parameters.get('beta1', self._default_dict['beta1'])
self._beta2 = parameters.get('beta2', self._default_dict['beta2'])
self._epsilon = parameters.get('epsilon', self._default_dict['epsilon'])
self._lazy_mode = parameters.get(
'lazy_mode', self._default_dict['lazy_mode']
)
self._weight_decay = parameters.get(
'weight_decay', self._default_dict['weight_decay']
)
parameters = parameters.get('params')
return parameters
def sharded_state_dict(
self,
model_sharded_state_dict: ShardedStateDict,
) -> ShardedStateDict:
"""
Convert optimizer state dict to a sharded state dict based on model sharding information.
Args:
model_sharded_state_dict (dict): Sharded state dict of the model, containing tensor metadata.
Returns:
dict: A new optimizer state dict where weights are wrapped as ShardedWeight.
"""
_FP32_MASTER = "fp32_master_0"
_MOMENT_NAME = "moment"
_optimizer_scalar_name = [
"beta1_pow_acc_0",
"beta2_pow_acc_0",
]
_optimizer_non_scaler_name = [
"moment1_0",
"moment2_0",
"velocity_0",
]
def _generate_base_static_name(vname):
if _FP32_MASTER in vname:
return tuple(vname.split("_" + _FP32_MASTER + "_", 1))
for name in _optimizer_scalar_name + _optimizer_non_scaler_name:
if vname.endswith(name):
return vname[: -(len(name) + 1)], name
raise ValueError(f"Cannot split variable name: {vname}.")
optimizer_sharded_state_dict = {}
optimizer_state_dict = self.state_dict()
# Build name mapping and remove non-tensor entries from optimizer state
static_to_struct_mapping = {}
model_sharded_state_dict = dict(
sorted(model_sharded_state_dict.items())
)
for k, v in model_sharded_state_dict.items():
# When shared weights exist, the v.local_tensor.name of shared parameters are identical, but only the first parameter has optimizer states. Therefore, only the key-value pairs of the first occurrence in the shared parameter group need to be retained.
if v.local_tensor.name not in static_to_struct_mapping:
static_to_struct_mapping[v.local_tensor.name] = k
master_weights = optimizer_state_dict.pop("master_weights", None)
optimizer_state_dict.pop("LR_Scheduler", None)
# Process main optimizer states
for key, tensor in optimizer_state_dict.items():
static_name, optim_state_type = _generate_base_static_name(key)
struct_name = static_to_struct_mapping[static_name]
sharded_weight = model_sharded_state_dict[struct_name]
unified_name = f"{struct_name}.{optim_state_type}"
# Determine tensor partitioning scheme
if _MOMENT_NAME in optim_state_type:
if tensor.is_dist():
optimizer_sharded_state_dict[unified_name] = ShardedWeight(
key=unified_name,
local_tensor=tensor,
local_shape=tensor.shape,
global_shape=tensor.shape,
global_offset=sharded_weight.global_offset,
)
else:
optimizer_sharded_state_dict[unified_name] = (
create_sharded_weight_with_new_local(
unified_name, tensor, sharded_weight
)
)
else: # Non-momentum parameters
optimizer_sharded_state_dict[unified_name] = ShardedWeight(
key=unified_name,
local_tensor=tensor,
local_shape=(1,),
global_shape=(1,),
global_offset=(0,),
)
# Process master weights if using mixed precision
if master_weights is not None:
for key, tensor in master_weights.items():
struct_name = static_to_struct_mapping[key]
sharded_weight = model_sharded_state_dict[struct_name]
unified_name = f"{struct_name}.w_0"
if tensor.is_dist():
optimizer_sharded_state_dict[unified_name] = ShardedWeight(
key=unified_name,
local_tensor=tensor,
local_shape=tensor.shape,
global_shape=tensor.shape,
global_offset=sharded_weight.global_offset,
)
else:
optimizer_sharded_state_dict[unified_name] = (
create_sharded_weight_with_new_local(
unified_name, tensor, sharded_weight
)
)
return optimizer_sharded_state_dict
python/paddle/optimizer/asgd.py
+379
View File
@@ -0,0 +1,379 @@
# 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 typing import TYPE_CHECKING
import paddle
from paddle import _C_ops, pir
from paddle.tensor.creation import to_tensor
from ..base import framework
from ..base.dygraph import no_grad
from ..base.framework import in_dygraph_mode, in_pir_mode
from .optimizer import Optimizer
if TYPE_CHECKING:
from collections.abc import Sequence
from paddle import Tensor
from paddle.nn.clip import GradientClipBase
from paddle.regularizer import WeightDecayRegularizer
from .lr import LRScheduler
from .optimizer import _ParameterConfig
__all__ = []
class ASGD(Optimizer):
r"""
Optimizer of the ASGD algorithm.Please refer to this for details:
`Minimizing Finite Sums with the Stochastic Average Gradient <https://hal.science/hal-00860051v2>`_.
.. math::
\begin{aligned}
&\hspace{0mm} d=0,\ y_i=0\ \textbf{for}\ i=1,2,...,n \\
&\hspace{0mm} \textbf{for}\ \: m=0,1,...\ \textbf{do} \: \\
&\hspace{5mm} i=m\ \%\ n \\
&\hspace{5mm} d=d-y_i+f_i{}'(x) \\
&\hspace{5mm} y_i=f_i{}'(x) \\
&\hspace{5mm} x=x-learning\_rate(\frac{d}{\mathrm{min}(m+1,\ n)}+\lambda x) \\
&\hspace{0mm} \textbf{end for} \\
\end{aligned}
Parameters:
learning_rate (float|Tensor|LRScheduler, optional): The learning rate used to update ``Parameter``.
It can be a float value, a ``Tensor`` with a float type or a LRScheduler. The default value is 0.001.
batch_num (int, optional): The number of batches needed to complete one epoch.
Assuming the total number of samples is ``all``,
it is recommended to set ``batch_num`` to ``all`` / ``batch_size``.
In situations where the graphics memory is tight,
it is possible to reduce the batch_num appropriately.
The default value is 1.
parameters (list|tuple|None, optional): List/Tuple of ``Tensor`` to update to minimize ``loss``.
This parameter is required in dygraph mode.
The default value is None in static graph mode, at this time all parameters will be updated.
weight_decay (int|float|WeightDecayRegularizer|None, optional): The strategy of regularization.
It can be a int or float value as coeff of L2 regularization or :ref:`api_paddle_regularizer_L1Decay`, :ref:`api_paddle_regularizer_L2Decay`.
If a parameter has set regularizer using :ref:`api_paddle_ParamAttr` already,
the regularization setting here in optimizer will be ignored for this parameter.
Otherwise, the regularization setting here in optimizer will take effect.
Default None, meaning there is no regularization.
grad_clip (GradientClipBase|None, optional): Gradient clipping strategy, it's an instance of some derived class of ``GradientClipBase`` .
There are three clipping strategies ( :ref:`api_paddle_nn_ClipGradByGlobalNorm` , :ref:`api_paddle_nn_ClipGradByNorm` , :ref:`api_paddle_nn_ClipGradByValue` ).
Default None, meaning there is no gradient clipping.
multi_precision (bool, optional): In mixed precision training scenarios based on GPU,
this parameter is mainly used to ensure the numerical stability of gradient updates.
When it is set to True, the optimizer will save a backup of FP32 type parameters with an equal value for FP16 type parameters.
When updating gradients, first increase the gradient type to FP32, and then assign it to the FP32 type parameter backup.
Finally, the updated FP32 type value will be converted to FP16 type first,
and then assigned to the actual FP16 type parameters participating in the calculation.
The default value is False.
name (str|None, optional): The default value is None. Normally there is no need for user to set this property.
For more information, please refer to :ref:`api_guide_Name` .
Examples:
.. code-block:: pycon
>>> import paddle
>>> inp = paddle.uniform(min=-0.1, max=0.1, shape=[10, 10], dtype='float32')
>>> linear = paddle.nn.Linear(10, 10)
>>> inp = paddle.to_tensor(inp)
>>> out = linear(inp)
>>> loss = paddle.mean(out)
>>> asgd = paddle.optimizer.ASGD(
... learning_rate=0.001,
... batch_num=10,
... parameters=linear.parameters(),
... weight_decay=0.01
... )
>>> out.backward()
>>> asgd.step()
>>> asgd.clear_grad()
"""
type: str
_d_acc_str = "d"
_y_acc_str = "y"
_m_acc_str = "m"
def __init__(
self,
learning_rate: float | Tensor | LRScheduler = 0.001,
batch_num: int = 1,
parameters: Sequence[Tensor] | Sequence[_ParameterConfig] | None = None,
weight_decay: float | WeightDecayRegularizer | None = None,
grad_clip: GradientClipBase | None = None,
multi_precision: bool = False,
name: str | None = None,
) -> None:
if learning_rate is None:
raise ValueError("learning_rate should not be none")
if batch_num is None:
raise ValueError("batch_num should not be none")
if not 0 < batch_num:
raise ValueError("batch_num should be greater than 0")
super().__init__(
learning_rate=learning_rate,
parameters=parameters,
weight_decay=weight_decay,
grad_clip=grad_clip,
name=name,
)
self.type = "asgd"
self._multi_precision = multi_precision
self._master_weights = {}
self._n = batch_num
self._n_tensor = None
def _create_accumulators(self, block, parameters):
if not isinstance(block, (framework.Block, pir.Block)):
raise TypeError("block is not instance of Block.")
if isinstance(parameters, dict):
parameters = self._update_param_group(parameters)
for p in parameters:
if p.name in self._already_create_accumulator:
continue
p_new = p
if self._multi_precision and self._is_dtype_fp16_or_bf16(p.dtype):
master_p = self._create_master_weight(p)
p_new = master_p
if (
self._is_dtype_fp16_or_bf16(p.dtype)
and not self._multi_precision
):
warnings.warn(
"Accumulating with FP16/BF16 in optimizer can lead to poor accuracy or slow convergence."
"Consider using multi_precision=True option of the Adam optimizer."
)
self._add_accumulator(
self._d_acc_str,
p_new,
p.dtype,
0,
)
# Sometimes p.shape is a tuple, so we need to change it to a list
self._add_accumulator(
self._y_acc_str,
p_new,
p.dtype,
0,
[self._n, *list(p.shape)],
)
self._add_accumulator(
self._m_acc_str,
p_new,
"int64",
0,
[1],
)
self._already_create_accumulator.add(p.name)
def _assign_accumulator_master(
self, block, name, param, assign_value, index
):
if self._name is not None:
name = self._name + "_" + name
find_master = self._multi_precision and self._is_dtype_fp16_or_bf16(
param.dtype
)
target_param = (
self._master_weights[param.name] if find_master else param
)
target_name = target_param.name
if (
name not in self._accumulators
or target_name not in self._accumulators[name]
):
raise Exception(
f"Accumulator {name} does not exist for parameter {target_name}"
)
if in_pir_mode():
if index is None:
self._accumulators[name][target_name] = paddle.assign(
assign_value
)
else:
self._accumulators[name][target_name][index] = paddle.assign(
assign_value
)
else:
assert isinstance(block, framework.Block)
assign_inputs = {
"X": assign_value,
}
assign_outputs = {
"Out": self._accumulators[name][target_name],
}
block.append_op(
type="assign",
inputs=assign_inputs,
outputs=assign_outputs,
)
@no_grad
def _append_optimize_op(self, block, param_and_grad):
if isinstance(param_and_grad, dict):
param_and_grad = self._update_param_group(param_and_grad)
if self._n_tensor is None:
self._n_tensor = to_tensor(
[self._n],
)
d = self._get_accumulator_master(self._d_acc_str, param_and_grad[0])
m = self._get_accumulator_master(self._m_acc_str, param_and_grad[0])
ys = self._get_accumulator_master(self._y_acc_str, param_and_grad[0])
index = paddle.mod(m, self._n_tensor).item()
y = paddle.assign(ys[index])
find_master = self._multi_precision and self._is_dtype_fp16_or_bf16(
param_and_grad[0].dtype
)
master_weight = (
self._master_weights[param_and_grad[0].name]
if find_master
else None
)
lr = self._create_param_lr(param_and_grad)
if in_dygraph_mode():
m.add_(to_tensor([1], dtype=m.dtype))
_C_ops.asgd_(
param_and_grad[0],
param_and_grad[1],
lr,
d,
ys[index],
paddle.fmin(m, self._n_tensor),
master_weight,
find_master,
)
return None
elif in_pir_mode():
m = paddle.assign(paddle.add(m, to_tensor([1], dtype=m.dtype)))
self._assign_accumulator_master(
block, self._m_acc_str, param_and_grad[0], m, None
)
# The y in the static graph has one more dimension than the y in the dynamic graph.
# So we should unify the shape of y in both dynamic and static graph.
# eg:
# dynamic graph: y.shape is [2, 2]
# static graph: y.shape is [1, 2, 2]
# so we should do
# static graph: y = y[0]
y = y[0]
_C_ops.asgd_(
param_and_grad[0],
param_and_grad[1],
lr,
d,
y,
paddle.fmin(m, self._n_tensor),
master_weight,
find_master,
)
self._assign_accumulator_master(
block, self._y_acc_str, param_and_grad[0], y, index
)
return None
else:
assert isinstance(block, framework.Block)
# create the optimize op
add_inputs = {
"X": m,
"Y": to_tensor([1], dtype=m.dtype),
}
add_outputs = {
"Out": m,
}
block.append_op(
type="elementwise_add",
inputs=add_inputs,
outputs=add_outputs,
)
# The y in the static graph has one more dimension than the y in the dynamic graph.
# So we should unify the shape of y in both dynamic and static graph.
# eg:
# dynamic graph: y.shape is [2, 2]
# static graph: y.shape is [1, 2, 2]
# so we should do
# static graph: y = y[0]
y = y[0]
asgd_inputs = {
"param": param_and_grad[0],
"grad": param_and_grad[1],
"learning_rate": lr,
"d": d,
"y": y,
"n": paddle.fmin(m, self._n_tensor),
}
asgd_outputs = {
"param_out": param_and_grad[0],
"d_out": d,
"y_out": y,
}
asgd_attrs = {"multi_precision": find_master}
if find_master:
asgd_inputs["master_param"] = master_weight
asgd_outputs["master_param_out"] = master_weight
asgd_op = block.append_op(
type=self.type,
inputs=asgd_inputs,
outputs=asgd_outputs,
attrs=asgd_attrs,
stop_gradient=True,
)
ys = paddle.static.setitem(ys, index, y)
self._assign_accumulator_master(
block, self._y_acc_str, param_and_grad[0], ys, None
)
return asgd_op
def _update_param_group(self, parameters):
parameters = parameters.get('params')
return parameters
python/paddle/optimizer/fusion_utils.py
+321
View File
@@ -0,0 +1,321 @@
# Copyright (c) 2024 PaddlePaddle Authors. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
import numpy as np
import paddle
import paddle.autograd as imperative_base
from paddle.framework import (
_current_expected_place_,
)
from paddle.incubate.tensor.manipulation import (
async_offload_with_offset,
create_async_load,
)
alignment = {
"gpu": 256,
"npu": 256,
"xpu": 256,
}
align = {
paddle.float16: 2,
paddle.bfloat16: 2,
paddle.float32: 4,
}
__current_device_type__ = None
def _share_tensor_ipc_meta(tensor):
if tensor is None:
return None
if paddle.is_compiled_with_xpu():
return tensor.value().get_tensor()._share_xpu()
if paddle.is_compiled_with_cuda() and not paddle.is_compiled_with_rocm():
return tensor.value().get_tensor()._share_cuda()
return None
def get_current_device_type():
global __current_device_type__
if __current_device_type__ is None:
if paddle.is_compiled_with_cuda():
device_type = "gpu"
elif paddle.is_compiled_with_xpu():
device_type = "xpu"
else:
current_device = _current_expected_place_()
try:
device_type = current_device.get_device_type()
except:
device_type = "unknown"
assert device_type in alignment.keys(), (
f"tensor fusion helper now only support {alignment.keys()}, but got device {device_type} instead."
)
__current_device_type__ = device_type
return __current_device_type__
def get_align(t):
size = np.prod(t.shape) * align[t.dtype]
remaining = size % alignment[get_current_device_type()]
ali = (
0
if remaining == 0
else alignment[get_current_device_type()] - remaining
)
align_ = ali // align[t.dtype]
return align_
class FusionStorage:
def __init__(
self,
accumulators,
master_weights,
merged_model_params=None,
dtype=paddle.float32,
):
assert isinstance(accumulators, dict), "accumulators must be a dict"
assert isinstance(master_weights, dict), "master_weights must be a dict"
assert (
isinstance(merged_model_params, dict) or merged_model_params is None
), "merged_model_params must be a dict or None"
self.accumulators = accumulators
self.master_weights = master_weights
self.merged_model_params = merged_model_params
self.accumulators_meta = {}
self.master_weights_meta = {}
self.merged_model_params_meta = {}
self.dtype = dtype
self.buffer = None
self.offset = 0
self.build_buffer()
self.mapping_tensor()
@imperative_base.no_grad()
def build_buffer(self):
self.offset = 0
for k, v in self.accumulators.items():
if k not in self.accumulators_meta:
self.accumulators_meta[k] = {}
for para_name, var_tmp in v.items():
assert var_tmp.dtype == self.dtype
src_len = var_tmp._numel() + get_align(var_tmp)
self.accumulators_meta[k][para_name] = {
"start": self.offset,
"end": self.offset + src_len,
"name": var_tmp.name,
"shape": var_tmp.shape,
}
self.offset += src_len
for k, v in self.master_weights.items():
assert v.dtype == self.dtype
src_len = v._numel() + get_align(v)
self.master_weights_meta[k] = {
"start": self.offset,
"end": self.offset + src_len,
"name": v.name,
"shape": v.shape,
}
self.offset += src_len
if self.merged_model_params is not None:
for k, v in self.merged_model_params.items():
assert v.dtype == self.dtype
src_len = v._numel() + get_align(v)
self.merged_model_params_meta[k] = {
"start": self.offset,
"end": self.offset + src_len,
"name": v.name,
"shape": v.shape,
}
self.offset += src_len
self.buffer = paddle.zeros((self.offset,), dtype=self.dtype)
@imperative_base.no_grad()
def mapping_tensor(self):
for k, v in self.accumulators_meta.items():
for para_name, meta in v.items():
self.mapping_tensor_impl(
src=self.accumulators[k][para_name],
start=meta["start"],
end=meta["end"],
)
for k, v in self.master_weights_meta.items():
self.mapping_tensor_impl(
src=self.master_weights[k], start=v["start"], end=v["end"]
)
for k, v in self.merged_model_params_meta.items():
self.mapping_tensor_impl(
src=self.merged_model_params[k],
start=v["start"],
end=v["end"],
)
@imperative_base.no_grad()
def mapping_tensor_impl(self, src, start, end):
tensor_shape = src.shape
stop_gradient = src.stop_gradient
src.stop_gradient = True
src.flatten_()
paddle.assign(
src,
self.buffer._slice(start, end),
)
src.get_tensor()._set_dims(tensor_shape)
src.stop_gradient = stop_gradient
self.buffer._slice(start, end)._share_buffer_to(src)
def _refresh_buffer_ipc_meta(self):
return _share_tensor_ipc_meta(self.buffer)
@property
def buffer_ipc_meta(self):
return self._refresh_buffer_ipc_meta()
class FusionStorageHelper:
def __init__(
self,
accumulators_meta,
master_weights_meta,
merged_model_params_meta,
buffer_ipc_meta,
):
self.async_loader = create_async_load()
self.accumulators_meta = None
self.master_weights_meta = None
self.merged_model_params_meta = None
self.buffer = None
self.cpu_buffer = None
self.buffer_length = None
self.tasks = []
self.reset_meta(
accumulators_meta,
master_weights_meta,
merged_model_params_meta,
buffer_ipc_meta,
)
@imperative_base.no_grad()
def reset_meta(
self,
accumulators_meta,
master_weights_meta,
merged_model_params_meta,
buffer_ipc_meta,
):
assert isinstance(accumulators_meta, dict), (
"accumulators_meta must be a dict"
)
self.accumulators_meta = accumulators_meta
assert isinstance(master_weights_meta, dict), (
"master_weights_meta must be a dict"
)
self.master_weights_meta = master_weights_meta
assert (
isinstance(merged_model_params_meta, dict)
or merged_model_params_meta is None
), "merged_model_params_meta must be a dict or None"
self.merged_model_params_meta = merged_model_params_meta
assert isinstance(buffer_ipc_meta, tuple), (
"buffer_ipc_meta must be a tuple"
)
assert len(buffer_ipc_meta) in (5, 7), (
"buffer_ipc_meta must be a tuple with length 5 when FLAGS_use_virtual_memory_auto_growth is True or 7 when FLAGS_use_virtual_memory_auto_growth is False."
)
if paddle.is_compiled_with_xpu():
new_tensor = paddle.base.core.DenseTensor._new_shared_xpu(
buffer_ipc_meta
)
else:
new_tensor = paddle.base.core.DenseTensor._new_shared_cuda(
buffer_ipc_meta
)
self.buffer = paddle.to_tensor(new_tensor)
self.cpu_buffer = self.buffer.pin_memory()
self.buffer_length = self.buffer._numel()
def sync_param(self):
self.sync_partial_param(0, self.buffer_length)
@imperative_base.no_grad()
def sync_partial_param(self, start, end):
assert isinstance(start, int), "start must be an integer"
assert isinstance(end, int), "end must be an integer"
assert start >= 0, "start must be non-negative"
assert end <= self.buffer_length, (
"end must be less than or equal to the total buffer length"
)
task = async_offload_with_offset(
src_tensor=self.buffer,
dst_tensor=self.cpu_buffer,
src_offset=start,
dst_offset=start,
offload_size=(end - start),
async_loader=self.async_loader,
)
self.tasks.append(task)
def wait_all(self):
if len(self.tasks) == 0:
return
last_task = self.tasks.pop(-1)
while len(self.tasks) > 0:
task = self.tasks.pop(0)
if paddle.is_compiled_with_xpu():
task.xpu_wait()
else:
task.cuda_wait()
last_task.cpu_wait()
def state_dict(self):
state_dict = {"master_weights": {}}
for k, v in self.accumulators_meta.items():
for para_name, tensor_meta in v.items():
var_tmp = self.restore_tensor_from_meta(tensor_meta)
state_dict[var_tmp.name] = var_tmp
for k, v in self.master_weights_meta.items():
var_tmp = self.restore_tensor_from_meta(v)
state_dict["master_weights"][k] = var_tmp
if self.merged_model_params_meta:
state_dict["merged_model_params"] = {}
for k, v in self.merged_model_params_meta.items():
var_tmp = self.restore_tensor_from_meta(v)
state_dict["merged_model_params"][k] = var_tmp
return state_dict
@imperative_base.no_grad()
def restore_tensor_from_meta(self, tensor_meta):
shape = tensor_meta["shape"]
name = tensor_meta["name"]
start = tensor_meta["start"]
end = tensor_meta["end"]
tensor = self.cpu_buffer._slice(start, end)
tensor.get_tensor()._set_dims(shape)
tensor.name = name
return tensor
python/paddle/optimizer/lamb.py
+358
View File
@@ -0,0 +1,358 @@
# Copyright (c) 2022 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
from typing import TYPE_CHECKING
from paddle import _C_ops, pir
from paddle.base.executor import global_scope
from ..base import core, framework
from ..base.framework import Variable
from .optimizer import Optimizer
if TYPE_CHECKING:
from collections.abc import Callable, Sequence
from typing_extensions import NotRequired
from paddle import Tensor
from paddle.nn.clip import GradientClipBase
from .optimizer import _ParameterConfig
class _LambParameterConfig(_ParameterConfig):
beta1: NotRequired[float | Tensor]
beta2: NotRequired[float | Tensor]
epsilon: NotRequired[float | Tensor]
lamb_weight_decay: NotRequired[float]
exclude_from_weight_decay_fn: NotRequired[
Callable[[Tensor], bool] | None
]
__all__ = []
class Lamb(Optimizer):
r"""
LAMB (Layer-wise Adaptive Moments optimizer for Batching training) Optimizer.
LAMB Optimizer is designed to scale up the batch size of training without losing
accuracy, which supports adaptive element-wise updating and accurate layer-wise
correction. For more information, please refer to `Large Batch Optimization for
Deep Learning: Training BERT in 76 minutes <https://arxiv.org/abs/1904.00962>`_ .
The updating of parameters follows:
.. math::
m_t &= \beta_1 m_{t - 1}+ (1 - \beta_1)g_t
v_t &= \beta_2 v_{t - 1} + (1 - \beta_2)g_t^2
m_t &= \frac{m_t}{\beta_1^t}
v_t &= \frac{v_t}{\beta_2^t}
r_t &= \frac{m_t}{\sqrt{v_t}+\epsilon}
w_t &= w_{t-1} -\eta_t \frac{\left \| w_{t-1}\right \|}{\left \| r_t + \lambda w_{t-1}\right \|} (r_t + \lambda w_{t-1})
where :math:`m` is the 1st moment, and :math:`v` the 2nd moment, :math:`\\eta` the
learning rate, :math:`\\lambda` the LAMB weight decay rate.
Args:
learning_rate (float|Tensor, optional): the learning rate used to update parameters. \
Can be a float value or a Variable with data type float32. Default 0.001.
lamb_weight_decay (float, optional): The LAMB weight decay rate. Default 0.01. Remind that weight_decay should be None.
beta1 (float|Tensor, optional): The exponential decay rate for the 1st moment estimates.
Default 0.9.
beta2 (float|Tensor, optional): The exponential decay rate for the 2nd moment estimates.
Default 0.999.
epsilon (float|Tensor, optional): A small float value for numerical stability. Default 1e-6.
parameters (list|tuple|None, optional): Iterable of ``Variable`` names to update to minimize ``loss``. \
This parameter is required in dygraph mode. And you can specify different options for \
different parameter groups such as the learning rate, weight decay, etc, \
then the parameters are list of dict. Note that the learning_rate in parameter groups \
represents the scale of base learning_rate. \
The default value is None in static graph mode, at this time all parameters will be updated.
grad_clip (GradientClipBase, optional): Gradient clipping strategy, it's an instance of
some derived class of ``GradientClipBase`` . There are three clipping strategies
( :ref:`api_paddle_base_clip_ClipGradByGlobalNorm` , :ref:`api_paddle_base_clip_ClipGradByNorm` ,
:ref:`api_paddle_base_clip_ClipGradByValue` ). If you want better convergence, it is recommended
to use :ref:`api_paddle_base_clip_ClipGradByGlobalNorm` . Default None, meaning there is no gradient clipping.
exclude_from_weight_decay_fn (Callable|None, optional): whether to skip weight decay for a parameter when this function returns True while take the parameter as input.
multi_precision (bool, optional) - Whether to use it during weight updates multi-precision, Default False。
always_adapt (bool, optional): whether to use Layer-wise LR adaptation. By default, skip adaptation on parameters that are
excluded from weight decay, unless always_adapt == True, then always enable LR adaptation.
name(str|None, optional): For detailed information, please refer to
:ref:`api_guide_Name` . Usually name is no need to set and None by default.
Examples:
.. code-block:: pycon
>>> import paddle
>>> inp = paddle.uniform(shape=[10, 10], dtype='float32', min=-0.1, max=0.1)
>>> linear = paddle.nn.Linear(10, 10)
>>> out = linear(inp)
>>> loss = paddle.mean(out)
>>> beta1 = paddle.to_tensor([0.9], dtype="float32")
>>> beta2 = paddle.to_tensor([0.85], dtype="float32")
>>> lamb = paddle.optimizer.Lamb(
... learning_rate=0.002,
... beta1=beta1,
... beta2=beta2,
... parameters=linear.parameters(),
... lamb_weight_decay=0.01
... )
>>> back = out.backward()
>>> lamb.step()
>>> lamb.clear_grad()
"""
_moment1_acc_str = "moment1"
_moment2_acc_str = "moment2"
_beta1_pow_acc_str = "beta1_pow_acc"
_beta2_pow_acc_str = "beta2_pow_acc"
def __init__(
self,
learning_rate: float | Tensor = 0.001,
lamb_weight_decay: float = 0.01,
beta1: float | Tensor = 0.9,
beta2: float | Tensor = 0.999,
epsilon: float | Tensor = 1e-6,
parameters: (
Sequence[Tensor] | Sequence[_LambParameterConfig] | None
) = None,
grad_clip: GradientClipBase | None = None,
exclude_from_weight_decay_fn: Callable[[Tensor], bool] | None = None,
multi_precision: bool = False,
always_adapt: bool = False,
name: str | None = None,
) -> None:
assert learning_rate is not None
assert beta1 is not None
assert beta2 is not None
assert epsilon is not None
super().__init__(
learning_rate=learning_rate,
parameters=parameters,
weight_decay=None,
grad_clip=grad_clip,
name=name,
)
self.type = "lamb"
self._beta1 = beta1
self._beta2 = beta2
self._epsilon = epsilon
self._lamb_weight_decay = lamb_weight_decay
self._exclude_from_weight_decay_fn = exclude_from_weight_decay_fn
self._default_dict = {
'beta1': beta1,
'beta2': beta2,
'epsilon': epsilon,
'lamb_weight_decay': lamb_weight_decay,
'exclude_from_weight_decay_fn': exclude_from_weight_decay_fn,
}
self._master_weights = {}
self._used_master_weights = {}
# TODO(zengjinle): expose API as soon as possible
self._multi_precision = multi_precision
self.always_adapt = always_adapt
def _get_parameter(self, name, scope=None):
if scope is None:
scope = global_scope()
p_t = scope.find_var(name).get_tensor()
master_name = self._used_master_weights.get(name)
if master_name is not None:
master_p_t = scope.find_var(master_name).get_tensor()
assert master_p_t._dtype() != p_t._dtype()
assert master_p_t.shape() == p_t.shape()
else:
master_p_t = None
return p_t, master_p_t
def _create_accumulators(self, block, parameters):
if not isinstance(block, (framework.Block, pir.Block)):
raise TypeError("block is not instance of Block.")
if isinstance(parameters, dict):
parameters = self._update_param_group(parameters)
# Create accumulator tensors for first and second moments
for p in parameters:
if p.name in self._already_create_accumulator:
continue
if self._multi_precision and self._is_dtype_fp16_or_bf16(p.dtype):
master_p = self._create_master_weight(p)
self._add_moments_pows(master_p)
self._already_create_accumulator.add(p.name)
else:
self._add_moments_pows(p)
self._already_create_accumulator.add(p.name)
def _add_moments_pows(self, p):
acc_dtype = p.dtype
if self._is_dtype_fp16_or_bf16(acc_dtype):
acc_dtype = core.VarDesc.VarType.FP32
self._add_accumulator(self._moment1_acc_str, p, dtype=acc_dtype)
self._add_accumulator(self._moment2_acc_str, p, dtype=acc_dtype)
self._add_accumulator(
name=self._beta1_pow_acc_str,
param=p,
dtype=acc_dtype,
fill_value=(
0.9 if isinstance(self._beta1, Variable) else self._beta1
),
shape=[1],
type=core.VarDesc.VarType.DENSE_TENSOR,
device='cpu',
)
self._add_accumulator(
name=self._beta2_pow_acc_str,
param=p,
dtype=acc_dtype,
fill_value=(
0.999 if isinstance(self._beta2, Variable) else self._beta2
),
shape=[1],
type=core.VarDesc.VarType.DENSE_TENSOR,
device='cpu',
)
def _append_optimize_op(self, block, param_and_grad):
if not isinstance(block, (framework.Block, pir.Block)):
raise TypeError("block is not instance of Block.")
if isinstance(param_and_grad, dict):
param_and_grad = self._update_param_group(param_and_grad)
block.program._use_lamb = True
moment1 = self._get_accumulator_master(
self._moment1_acc_str, param_and_grad[0]
)
moment2 = self._get_accumulator_master(
self._moment2_acc_str, param_and_grad[0]
)
beta1_pow_acc = self._get_accumulator_master(
self._beta1_pow_acc_str, param_and_grad[0]
)
beta2_pow_acc = self._get_accumulator_master(
self._beta2_pow_acc_str, param_and_grad[0]
)
if (
self._exclude_from_weight_decay_fn is not None
and self._exclude_from_weight_decay_fn(param_and_grad[0])
):
weight_decay = 0.0
else:
weight_decay = self._lamb_weight_decay
lr = self._create_param_lr(param_and_grad)
find_master = self._multi_precision and self._is_dtype_fp16_or_bf16(
param_and_grad[0].dtype
)
p_name = param_and_grad[0].name
if find_master:
master_weight = self._master_weights[p_name]
self._used_master_weights[p_name] = master_weight.name
else:
master_weight = None
if framework.in_dynamic_or_pir_mode():
_C_ops.lamb_(
param_and_grad[0],
param_and_grad[1],
lr,
moment1,
moment2,
beta1_pow_acc,
beta2_pow_acc,
master_weight,
None,
weight_decay,
self._beta1,
self._beta2,
self._epsilon,
self.always_adapt,
find_master,
)
return None
else:
# create the lamb optimize op
inputs = {
"Param": param_and_grad[0],
"Grad": param_and_grad[1],
"LearningRate": lr,
"Moment1": moment1,
"Moment2": moment2,
"Beta1Pow": beta1_pow_acc,
"Beta2Pow": beta2_pow_acc,
}
outputs = {
"ParamOut": param_and_grad[0],
"Moment1Out": moment1,
"Moment2Out": moment2,
"Beta1PowOut": beta1_pow_acc,
"Beta2PowOut": beta2_pow_acc,
}
attrs = {
"beta1": self._beta1,
"beta2": self._beta2,
"epsilon": self._epsilon,
"weight_decay": weight_decay,
"always_adapt": self.always_adapt,
"multi_precision": find_master,
}
if find_master:
inputs["MasterParam"] = master_weight
outputs["MasterParamOut"] = master_weight
found_inf = self._get_auxiliary_var('found_inf')
if found_inf:
inputs["SkipUpdate"] = found_inf
lamb_op = block.append_op(
type=self.type,
inputs=inputs,
outputs=outputs,
attrs=attrs,
stop_gradient=True,
)
return lamb_op
def _update_param_group(self, parameters):
self._beta1 = parameters.get('beta1', self._default_dict['beta1'])
self._beta2 = parameters.get('beta2', self._default_dict['beta2'])
self._epsilon = parameters.get('epsilon', self._default_dict['epsilon'])
self._lamb_weight_decay = parameters.get(
'lamb_weight_decay', self._default_dict['lamb_weight_decay']
)
self._exclude_from_weight_decay_fn = parameters.get(
'exclude_from_weight_decay_fn',
self._default_dict['exclude_from_weight_decay_fn'],
)
parameters = parameters.get('params')
return parameters
python/paddle/optimizer/lbfgs.py
+826
View File
@@ -0,0 +1,826 @@
# Copyright (c) 2022 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 os
import warnings
from collections import defaultdict
from functools import reduce
from typing import TYPE_CHECKING, NoReturn, TypedDict
from typing_extensions import NotRequired
import paddle
from ..base import framework
from .optimizer import Optimizer
if TYPE_CHECKING:
from collections.abc import Sequence
from paddle import Tensor
from paddle.nn.clip import GradientClipBase
from paddle.regularizer import WeightDecayRegularizer
from .optimizer import _ParameterConfig
__all__ = []
class _LbfgsState(TypedDict):
func_evals: int
n_iter: int
d: Tensor
alpha: Tensor
old_yk: list[Tensor]
old_sk: list[Tensor]
ro: list[Tensor]
H_diag: Tensor
prev_flat_grad: Tensor
prev_loss: float
al: NotRequired[list[Tensor]]
class _LbfgsStateDict(TypedDict):
state: _LbfgsState
def check_tf32_override():
"""Check and warn about TF32 acceleration status"""
if (
paddle.device.is_compiled_with_cuda()
and os.getenv("NVIDIA_TF32_OVERRIDE") != "0"
): # None or "1"
warnings.warn(
"Warning! TF32 Tensor Cores are enabled by default on some NVIDIA GPUs for faster computation, "
"but may compromise numerical precision in specific cases, particularly with the L-BFGS optimizer."
"To disable it, set: NVIDIA_TF32_OVERRIDE=0"
)
def dot(x, y):
r"""
NOTE: This is a temporary workaround for unstable result computed by `paddle.dot`,
which will be reverted when the problem is fixed."
"""
return (x * y).sum(axis=-1)
def _cubic_interpolate(x1, f1, g1, x2, f2, g2, bounds=None):
r"""Cubic interpolation between (x1, f1, g1) and (x2, f2, g2).
Use two points and their gradient to determine a cubic function and get the minimum point
between them in the cubic curve.
Reference:
Jorge Nocedal, Stephen J. Wright, Numerical Optimization, Second Edition, 2006.
pp59: formula 3.59
Args:
x1, f1, g1: point1's position, value and gradient.
x2, f2, g2: point2's position, value and gradient.
bounds: bounds of interpolation area
Returns:
min_pos: the minimum point between the specified points in the cubic curve.
"""
# Compute bounds of interpolation area
if bounds is not None:
xmin_bound, xmax_bound = bounds
else:
xmin_bound, xmax_bound = (x1, x2) if x1 <= x2 else (x2, x1)
d1 = g1 + g2 - 3 * (f1 - f2) / (x1 - x2)
d2_square = d1**2 - g1 * g2
if d2_square >= 0:
d2 = d2_square.sqrt()
if x1 <= x2:
min_pos = x2 - (x2 - x1) * ((g2 + d2 - d1) / (g2 - g1 + 2 * d2))
else:
min_pos = x1 - (x1 - x2) * ((g1 + d2 - d1) / (g1 - g2 + 2 * d2))
return min(max(min_pos, xmin_bound), xmax_bound)
else:
return (xmin_bound + xmax_bound) / 2.0
def _strong_wolfe(
obj_func,
xk,
alpha,
d,
loss,
grad,
gtd,
c1=1e-4,
c2=0.9,
tolerance_change=1e-9,
max_ls=25,
):
r"""Implements of line search algorithm that satisfies the strong Wolfe conditions using double zoom.
Reference:
Jorge Nocedal, Stephen J. Wright, Numerical Optimization, Second Edition, 2006.
pp60: Algorithm 3.5 (Line Search Algorithm).
Args:
obj_func: the objective function to minimize. ```` accepts a multivariate input and returns a scalar.
xk (Tensor): the starting point of the iterates.
alpha (Scalar): the initial step size.
d (Tensor): search direction.
loss (scalar): the initial loss
grad (Tensor): the initial grad
c1 (Scalar): parameter for sufficient decrease condition.
c2 (Scalar): parameter for curvature condition.
tolerance_change (Scalar): terminates if the change of function value/position/parameter between
two iterations is smaller than this value.
max_ls(int): max iteration of line search.
alpha_max (float): max step length.
Returns:
loss_new (Scaler): loss of obj_func at final alpha.
grad_new, (Tensor): derivative of obj_func at final alpha.
alpha(Tensor): optimal step length, or 0. if the line search algorithm did not converge.
ls_func_evals (Scaler): number of objective function called in line search process.
Following summarizes the essentials of the strong Wolfe line search algorithm.
Some notations used in the description:
- `func` denotes the objective function.
- `obi_func` is a function of step size alpha, restricting `obj_func` on a line.
obi_func = func(xk + alpha * d),
where xk is the position of k'th iterate, d is the line search direction(decent direction),
and a is the step size.
- alpha : substitute of alpha
- a1 is alpha of last iteration, which is alpha_(i-1).
- a2 is alpha of current iteration, which is alpha_i.
- a_lo is alpha in left position when calls zoom, which is alpha_low.
- a_hi is alpha in right position when calls zoom, which is alpha_high.
Line Search Algorithm:
repeat
Compute obi_func(a2) and derphi(a2).
1. If obi_func(a2) > obi_func(0) + c_1 * a2 * obi_func'(0) or [obi_func(a2) >= obi_func(a1) and i > 1],
alpha= zoom(a1, a2) and stop;
2. If |obi_func'(a2)| <= -c_2 * obi_func'(0),
alpha= a2 and stop;
3. If obi_func'(a2) >= 0,
alpha= zoom(a2, a1) and stop;
a1 = a2
a2 = min(2 * a2, a2)
i = i + 1
end(repeat)
zoom(a_lo, a_hi) Algorithm:
repeat
aj = cubic_interpolation(a_lo, a_hi)
Compute obi_func(aj) and derphi(aj).
1. If obi_func(aj) > obi_func(0) + c_1 * aj * obi_func'(0) or obi_func(aj) >= obi_func(a_lo),
then a_hi <- aj;
2.
2.1. If |obi_func'(aj)| <= -c_2 * obi_func'(0), then alpha= a2 and stop;
2.2. If obi_func'(aj) * (a2 - a1) >= 0, then a_hi = a_lo
a_lo = aj;
end(repeat)
reference: https://github.com/pytorch/pytorch
"""
d_norm = d.abs().max()
grad = grad.clone()
# evaluate objective and gradient using initial step
loss_new, grad_new = obj_func(xk, alpha, d)
ls_func_evals = 1
gtd_new = dot(grad_new, d)
# bracket an interval containing a point satisfying the Wolfe criteria
t_prev, f_prev, g_prev, gtd_prev = (0, loss, grad, gtd)
done = False
ls_iter = 0
while ls_iter < max_ls:
# check conditions
if loss_new > (loss + c1 * alpha * gtd) or (
ls_iter > 1 and loss_new >= f_prev
):
bracket = [t_prev, alpha]
bracket_f = [f_prev, loss_new]
bracket_g = [g_prev, grad_new.clone()]
bracket_gtd = [gtd_prev, gtd_new]
break
if abs(gtd_new) <= -c2 * gtd:
bracket = [alpha]
bracket_f = [loss_new]
bracket_g = [grad_new]
done = True
break
if gtd_new >= 0:
bracket = [t_prev, alpha]
bracket_f = [f_prev, loss_new]
bracket_g = [g_prev, grad_new.clone()]
bracket_gtd = [gtd_prev, gtd_new]
break
# interpolate
min_step = alpha + 0.01 * (alpha - t_prev)
max_step = alpha * 10
tmp = alpha
alpha = _cubic_interpolate(
t_prev,
f_prev,
gtd_prev,
alpha,
loss_new,
gtd_new,
bounds=(min_step, max_step),
)
# next step
t_prev = tmp
f_prev = loss_new
g_prev = grad_new.clone()
gtd_prev = gtd_new
loss_new, grad_new = obj_func(xk, alpha, d)
ls_func_evals += 1
gtd_new = dot(grad_new, d)
ls_iter += 1
# reached max number of iterations?
if ls_iter == max_ls:
bracket = [0, alpha]
bracket_f = [loss, loss_new]
bracket_g = [grad, grad_new]
# zoom phase: we now have a point satisfying the criteria, or
# a bracket around it. We refine the bracket until we find the
# exact point satisfying the criteria
insuf_progress = False
# find high and low points in bracket
low_pos, high_pos = (0, 1) if bracket_f[0] <= bracket_f[-1] else (1, 0)
while not done and ls_iter < max_ls:
# line-search bracket is so small
if abs(bracket[1] - bracket[0]) * d_norm < tolerance_change:
break
# compute new trial value
alpha = _cubic_interpolate(
bracket[0],
bracket_f[0],
bracket_gtd[0],
bracket[1],
bracket_f[1],
bracket_gtd[1],
)
# test that we are making sufficient progress:
# in case `alpha` is so close to boundary, we mark that we are making
# insufficient progress, and if
# + we have made insufficient progress in the last step, or
# + `alpha` is at one of the boundary,
# we will move `alpha` to a position which is `0.1 * len(bracket)`
# away from the nearest boundary point.
eps = 0.1 * (max(bracket) - min(bracket))
if min(max(bracket) - alpha, alpha - min(bracket)) < eps:
# interpolation close to boundary
if insuf_progress or alpha >= max(bracket) or alpha <= min(bracket):
# evaluate at 0.1 away from boundary
if abs(alpha - max(bracket)) < abs(alpha - min(bracket)):
alpha = max(bracket) - eps
else:
alpha = min(bracket) + eps
insuf_progress = False
else:
insuf_progress = True
else:
insuf_progress = False
# Evaluate new point
loss_new, grad_new = obj_func(xk, alpha, d)
ls_func_evals += 1
gtd_new = dot(grad_new, d)
ls_iter += 1
if (
loss_new > (loss + c1 * alpha * gtd)
or loss_new >= bracket_f[low_pos]
):
# Armijo condition not satisfied or not lower than lowest point
bracket[high_pos] = alpha
bracket_f[high_pos] = loss_new
bracket_g[high_pos] = grad_new.clone()
bracket_gtd[high_pos] = gtd_new
low_pos, high_pos = (
(0, 1) if bracket_f[0] <= bracket_f[1] else (1, 0)
)
else:
if abs(gtd_new) <= -c2 * gtd:
# Wolfe conditions satisfied
done = True
elif gtd_new * (bracket[high_pos] - bracket[low_pos]) >= 0:
# old high becomes new low
bracket[high_pos] = bracket[low_pos]
bracket_f[high_pos] = bracket_f[low_pos]
bracket_g[high_pos] = bracket_g[low_pos]
bracket_gtd[high_pos] = bracket_gtd[low_pos]
# new point becomes new low
bracket[low_pos] = alpha
bracket_f[low_pos] = loss_new
bracket_g[low_pos] = grad_new.clone()
bracket_gtd[low_pos] = gtd_new
# return stuff
alpha = bracket[low_pos]
loss_new = bracket_f[low_pos]
grad_new = bracket_g[low_pos]
return loss_new, grad_new, alpha, ls_func_evals
class LBFGS(Optimizer):
r"""
The L-BFGS is a quasi-Newton method for solving an unconstrained optimization problem over a differentiable function.
Closely related is the Newton method for minimization. Consider the iterate update formula:
.. math::
x_{k+1} = x_{k} + H_k \nabla{f_k}
If :math:`H_k` is the inverse Hessian of :math:`f` at :math:`x_k`, then it's the Newton method.
If :math:`H_k` is symmetric and positive definite, used as an approximation of the inverse Hessian, then
it's a quasi-Newton. In practice, the approximated Hessians are obtained
by only using the gradients, over either whole or part of the search
history, the former is BFGS, the latter is L-BFGS.
Reference:
Jorge Nocedal, Stephen J. Wright, Numerical Optimization, Second Edition, 2006. pp179: Algorithm 7.5 (L-BFGS).
Args:
learning_rate (float, optional): learning rate .The default value is 1.
max_iter (int, optional): maximal number of iterations per optimization step.
The default value is 20.
max_eval (int|None, optional): maximal number of function evaluations per optimization
step. The default value is max_iter * 1.25.
tolerance_grad (float, optional): termination tolerance on first order optimality
The default value is 1e-5.
tolerance_change (float, optional): termination tolerance on function
value/parameter changes. The default value is 1e-9.
history_size (int, optional): update history size. The default value is 100.
line_search_fn (string|None, optional): either 'strong_wolfe' or None. The default value is strong_wolfe.
parameters (list|tuple|None, optional): List/Tuple of ``Tensor`` names to update to minimize ``loss``. \
This parameter is required in dygraph mode. The default value is None.
weight_decay (int|float|WeightDecayRegularizer|None, optional): The strategy of regularization. \
It can be a int or float value as coeff of L2 regularization or \
:ref:`api_paddle_regularizer_L1Decay`, :ref:`api_paddle_regularizer_L2Decay`.
If a parameter has set regularizer using :ref:`api_paddle_ParamAttr` already, \
the regularization setting here in optimizer will be ignored for this parameter. \
Otherwise, the regularization setting here in optimizer will take effect. \
Default None, meaning there is no regularization.
grad_clip (GradientClipBase|None, optional): Gradient clipping strategy, it's an instance of \
some derived class of ``GradientClipBase`` . There are three clipping strategies \
( :ref:`api_paddle_nn_ClipGradByGlobalNorm` , :ref:`api_paddle_nn_ClipGradByNorm` , \
:ref:`api_paddle_nn_ClipGradByValue` ). Default None, meaning there is no gradient clipping.
name (str|None, optional): Normally there is no need for user to set this property.
For more information, please refer to :ref:`api_guide_Name`.
The default value is None.
Return:
loss (Tensor): the final loss of closure.
Examples:
.. code-block:: pycon
>>> import paddle
>>> import numpy as np
>>> paddle.disable_static()
>>> np.random.seed(0)
>>> np_w = np.random.rand(1).astype(np.float32)
>>> np_x = np.random.rand(1).astype(np.float32)
>>> inputs = [np.random.rand(1).astype(np.float32) for i in range(10)]
>>> # y = 2x
>>> targets = [2 * x for x in inputs]
>>> class Net(paddle.nn.Layer):
... def __init__(self):
... super().__init__()
... w = paddle.to_tensor(np_w)
... self.w = paddle.create_parameter(shape=w.shape, dtype=w.dtype, default_initializer=paddle.nn.initializer.Assign(w))
...
... def forward(self, x):
... return self.w * x
...
>>> net = Net()
>>> opt = paddle.optimizer.LBFGS(learning_rate=1, max_iter=1, max_eval=None, tolerance_grad=1e-07, tolerance_change=1e-09, history_size=100, line_search_fn='strong_wolfe', parameters=net.parameters())
>>> def train_step(inputs, targets):
... def closure():
... outputs = net(inputs)
... loss = paddle.nn.functional.mse_loss(outputs, targets)
... print('loss: ', loss.item())
... opt.clear_grad()
... loss.backward()
... return loss
... opt.step(closure)
...
>>> for input_np, target_np in zip(inputs, targets):
... input = paddle.to_tensor(input_np)
... target = paddle.to_tensor(target_np)
... train_step(input, target)
"""
def __init__(
self,
learning_rate: float = 1.0,
max_iter: int = 20,
max_eval: int | None = None,
tolerance_grad: float = 1e-7,
tolerance_change: float = 1e-9,
history_size: int = 100,
line_search_fn: str | None = None,
parameters: Sequence[Tensor] | Sequence[_ParameterConfig] | None = None,
weight_decay: float | WeightDecayRegularizer | None = None,
grad_clip: GradientClipBase | None = None,
name: str | None = None,
) -> None:
check_tf32_override()
if max_eval is None:
max_eval = max_iter * 5 // 4
self.learning_rate = learning_rate
self.max_iter = max_iter
self.max_eval = max_eval
self.tolerance_grad = tolerance_grad
self.tolerance_change = tolerance_change
self.history_size = history_size
self.line_search_fn = line_search_fn
if isinstance(parameters, paddle.Tensor):
raise TypeError(
"parameters argument given to the optimizer should be "
"an iterable of Tensors or dicts, but got " + type(parameters)
)
self.state = defaultdict(dict)
super().__init__(
learning_rate=1.0,
parameters=parameters,
weight_decay=weight_decay,
grad_clip=grad_clip,
name=name,
)
if not isinstance(self._parameter_list[0], dict):
self._params = self._parameter_list
else:
for idx, param_group in enumerate(self._param_groups):
self._params = param_group['params']
self._numel_cache = None
def state_dict(self) -> _LbfgsStateDict:
r"""Returns the state of the optimizer as a :class:`dict`.
Return:
state, a dict holding current optimization state. Its content
differs between optimizer classes.
Examples:
.. code-block:: pycon
>>> import paddle
>>> paddle.disable_static()
>>> net = paddle.nn.Linear(10, 10)
>>> opt = paddle.optimizer.LBFGS(
... learning_rate=1,
... max_iter=1,
... max_eval=None,
... tolerance_grad=1e-07,
... tolerance_change=1e-09,
... history_size=100,
... line_search_fn='strong_wolfe',
... parameters=net.parameters(),
>>> )
>>> def train_step(inputs, targets):
... def closure():
... outputs = net(inputs)
... loss = paddle.nn.functional.mse_loss(outputs, targets)
... opt.clear_grad()
... loss.backward()
... return loss
...
... opt.step(closure)
>>> inputs = paddle.rand([10, 10], dtype="float32")
>>> targets = paddle.to_tensor([2 * x for x in inputs])
>>> n_iter = 0
>>> while n_iter < 20:
... loss = train_step(inputs, targets)
... n_iter = opt.state_dict()["state"]["func_evals"]
... print("n_iter:", n_iter)
"""
packed_state = {}
for k, v in self.state.items():
packed_state.update({k: v})
return {'state': packed_state}
def _numel(self) -> int:
# compute the number of all parameters
if self._numel_cache is None:
self._numel_cache = reduce(
lambda total, p: total + p.numel(), self._params, 0
)
return self._numel_cache
# flatten grad of all parameters
def _gather_flat_grad(self):
views = []
for p in self._params:
if p.grad is None:
view = paddle.zeros_like(p).reshape([-1])
else:
view = p.grad.reshape([-1])
views.append(view)
return paddle.concat(views, axis=0)
# compute xk = xk + alpha * direction
def _add_grad(self, alpha, direction):
offset = 0
for p in self._params:
numel = reduce(lambda x, y: x * y, p.shape) if p.shape != [] else 1
p = paddle.assign(
p.add(
direction[offset : offset + numel].reshape(p.shape) * alpha
),
p,
)
offset += numel
assert offset == self._numel()
def _clone_param(self):
return [p.clone() for p in self._params]
def _set_param(self, params_data):
for p, pdata in zip(self._params, params_data):
paddle.assign(pdata, p)
def _directional_evaluate(self, closure, x, alpha, d):
self._add_grad(alpha, d)
loss = float(closure())
flat_grad = self._gather_flat_grad()
self._set_param(x)
return loss, flat_grad
@framework.non_static_only
def step(self, closure) -> Tensor:
"""Performs a single optimization step.
Args:
closure (callable): A closure that reevaluates the model
and returns the loss.
Examples:
.. code-block:: pycon
>>> import paddle
>>> paddle.disable_static()
>>> inputs = paddle.rand([10, 10], dtype="float32")
>>> targets = paddle.to_tensor([2 * x for x in inputs])
>>> net = paddle.nn.Linear(10, 10)
>>> opt = paddle.optimizer.LBFGS(
... learning_rate=1,
... max_iter=1,
... max_eval=None,
... tolerance_grad=1e-07,
... tolerance_change=1e-09,
... history_size=100,
... line_search_fn='strong_wolfe',
... parameters=net.parameters(),
>>> )
>>> def closure():
... outputs = net(inputs)
... loss = paddle.nn.functional.mse_loss(outputs, targets)
... print("loss:", loss.item())
... opt.clear_grad()
... loss.backward()
... return loss
>>> opt.step(closure)
"""
with paddle.no_grad():
# Make sure the closure is always called with grad enabled
closure = paddle.enable_grad()(closure)
learning_rate = self.learning_rate
max_iter = self.max_iter
max_eval = self.max_eval
tolerance_grad = self.tolerance_grad
tolerance_change = self.tolerance_change
line_search_fn = self.line_search_fn
history_size = self.history_size
state = self.state
state.setdefault('func_evals', 0)
state.setdefault('n_iter', 0)
# evaluate initial f(x) and df/dx
orig_loss = closure()
loss = float(orig_loss)
current_evals = 1
state['func_evals'] += 1
flat_grad = self._gather_flat_grad()
opt_cond = flat_grad.abs().max() <= tolerance_grad
# optimal condition
if opt_cond:
return orig_loss
# tensors cached in state (for tracing)
d = state.get('d')
alpha = state.get('alpha')
old_yk = state.get('old_yk')
old_sk = state.get('old_sk')
ro = state.get('ro')
H_diag = state.get('H_diag')
prev_flat_grad = state.get('prev_flat_grad')
prev_loss = state.get('prev_loss')
n_iter = 0
# optimize for a max of max_iter iterations
while n_iter < max_iter:
# keep track of nb of iterations
n_iter += 1
state['n_iter'] += 1
############################################################
# compute gradient descent direction
############################################################
if state['n_iter'] == 1:
d = flat_grad.neg()
old_yk = []
old_sk = []
ro = []
H_diag = paddle.to_tensor(1.0, dtype=orig_loss.dtype)
else:
# do lbfgs update (update memory)
y = flat_grad.subtract(prev_flat_grad)
s = d.multiply(paddle.to_tensor(alpha, dtype=d.dtype))
ys = dot(y, s)
if ys > 1e-10:
# updating memory
if len(old_yk) == history_size:
# shift history by one (limited-memory)
old_yk.pop(0)
old_sk.pop(0)
ro.pop(0)
# store new direction/step
old_yk.append(y)
old_sk.append(s)
ro.append(1.0 / ys)
# update scale of initial Hessian approximation
H_diag = ys / dot(y, y) # (y*y)
# compute the approximate (L-BFGS) inverse Hessian
# multiplied by the gradient
num_old = len(old_yk)
if 'al' not in state:
state['al'] = [None] * history_size
al = state['al']
# iteration in L-BFGS loop collapsed to use just one buffer
q = flat_grad.neg()
for i in range(num_old - 1, -1, -1):
al[i] = dot(old_sk[i], q) * ro[i]
paddle.assign(q.add(old_yk[i] * (-al[i])), q)
# multiply by initial Hessian
# r/d is the final direction
d = r = paddle.multiply(q, H_diag)
for i in range(num_old):
be_i = dot(old_yk[i], r) * ro[i]
paddle.assign(r.add(old_sk[i] * (al[i] - be_i)), r)
if prev_flat_grad is None:
prev_flat_grad = flat_grad.clone()
else:
paddle.assign(flat_grad, prev_flat_grad)
prev_loss = loss
############################################################
# compute step length
############################################################
# reset initial guess for step size
if state['n_iter'] == 1:
alpha = (
min(1.0, 1.0 / flat_grad.abs().sum()) * learning_rate
)
else:
alpha = learning_rate
# directional derivative
gtd = dot(flat_grad, d)
# directional derivative is below tolerance
if gtd > -tolerance_change:
break
# optional line search: user function
ls_func_evals = 0
if line_search_fn is not None:
# perform line search, using user function
if line_search_fn != "strong_wolfe":
raise RuntimeError("only 'strong_wolfe' is supported")
else:
x_init = self._clone_param()
def obj_func(x, alpha, d):
return self._directional_evaluate(
closure, x, alpha, d
)
loss, flat_grad, alpha, ls_func_evals = _strong_wolfe(
obj_func, x_init, alpha, d, loss, flat_grad, gtd
)
self._add_grad(alpha, d)
opt_cond = flat_grad.abs().max() <= tolerance_grad
else:
# no line search, simply move with fixed-step
self._add_grad(alpha, d)
if n_iter != max_iter:
with paddle.enable_grad():
loss = float(closure())
flat_grad = self._gather_flat_grad()
opt_cond = flat_grad.abs().max() <= tolerance_grad
ls_func_evals = 1
# update func eval
current_evals += ls_func_evals
state['func_evals'] += ls_func_evals
# optimal condition
if opt_cond:
break
# lack of progress
if (d * alpha).abs().max() <= tolerance_change:
break
if abs(loss - prev_loss) < tolerance_change:
break
# check conditions
if current_evals >= max_eval:
break
if n_iter == max_iter:
break
state['d'] = d
state['alpha'] = alpha
state['old_yk'] = old_yk
state['old_sk'] = old_sk
state['ro'] = ro
state['H_diag'] = H_diag
state['prev_flat_grad'] = prev_flat_grad
state['prev_loss'] = prev_loss
return orig_loss
def minimize(
self, loss, startup_program=None, parameters=None, no_grad_set=None
) -> NoReturn:
"""Empty method. LBFGS optimizer does not use this way to minimize ``loss``. Please refer 'Examples' of LBFGS() above for usage."""
raise NotImplementedError(
"LBFGS optimizer does not use this way to minimize loss. Please refer 'Examples' of LBFGS() for usage."
)
python/paddle/optimizer/lr.py
+3437
View File
File diff suppressed because it is too large Load Diff
python/paddle/optimizer/momentum.py
+594
View File
@@ -0,0 +1,594 @@
# 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 typing import TYPE_CHECKING
import paddle
from paddle import _C_ops, pir
from paddle.framework import in_dynamic_or_pir_mode
from paddle.regularizer import L2Decay
from ..base import core, framework
from .optimizer import Optimizer
if TYPE_CHECKING:
from collections.abc import Sequence
from typing_extensions import NotRequired
from paddle import Tensor
from paddle.nn.clip import GradientClipBase
from paddle.regularizer import WeightDecayRegularizer
from .lr import LRScheduler
from .optimizer import _ParameterConfig
class _MomentumParameterConfig(_ParameterConfig):
momentum: NotRequired[float]
use_nesterov: NotRequired[bool]
rescale_grad: NotRequired[float]
regularization_method: NotRequired[str]
regularization_coeff: NotRequired[float]
__all__ = []
class Momentum(Optimizer):
r"""
Simple Momentum optimizer with velocity state
This optimizer has a flag for Nestrov Momentum.
The update equations are as follows:
.. math::
& velocity = mu * velocity + gradient
& if (use\_nesterov):
&\quad param = param - (gradient + mu * velocity) * learning\_rate
& else:
&\quad param = param - learning\_rate * velocity
Parameters:
learning_rate (float|Tensor|LRScheduler, optional): The learning rate used to update ``Parameter``.
It can be a float value, a ``Tensor`` with a float type or a LRScheduler. The default value is 0.001.
momentum (float): Momentum factor. The default value is 0.9.
parameters (list|tuple|None, optional): List|Tuple of ``Tensor`` to update to minimize ``loss``. \
This parameter is required in dygraph mode. And you can specify different options for \
different parameter groups such as the learning rate, weight decay, etc, \
then the parameters are list of dict. Note that the learning_rate in parameter groups \
represents the scale of base learning_rate. \
The default value is None in static graph mode, at this time all parameters will be updated.
use_nesterov(bool, optional): Enables Nesterov momentum. The default value is False.
weight_decay (int|float|WeightDecayRegularizer|None, optional): The strategy of regularization. \
It can be a int or float value as coeff of L2 regularization or \
:ref:`api_paddle_regularizer_L1Decay`, :ref:`api_paddle_regularizer_L2Decay`.
If a parameter has set regularizer using :ref:`api_paddle_ParamAttr` already, \
the regularization setting here in optimizer will be ignored for this parameter. \
Otherwise, the regularization setting here in optimizer will take effect. \
Default None, meaning there is no regularization.
grad_clip (GradientClipBase|None, optional): Gradient clipping strategy, it's an instance of
some derived class of ``GradientClipBase`` . There are three clipping strategies
( :ref:`api_paddle_nn_ClipGradByGlobalNorm` , :ref:`api_paddle_nn_ClipGradByNorm` ,
:ref:`api_paddle_nn_ClipGradByValue` ). Default None, meaning there is no gradient clipping.
multi_precision (bool, optional): Whether to use multi-precision during weight updating. Default is false.
rescale_grad (float, optional): Multiply the gradient with `rescale_grad` before updating. \
Often choose to be ``1.0/batch_size``.
use_multi_tensor (bool, optional): Whether to use multi-tensor strategy to update all parameters at once . Default is false.
name (str|None, optional): The default value is None. Normally there is no need for user
to set this property. For more information, please refer to
:ref:`api_guide_Name` .
Examples:
.. code-block:: pycon
>>> import paddle
>>> inp = paddle.uniform([10, 10], dtype="float32", min=-0.1, max=0.1)
>>> linear = paddle.nn.Linear(10, 10)
>>> inp = paddle.to_tensor(inp)
>>> out = linear(inp)
>>> loss = paddle.mean(out)
>>> momentum = paddle.optimizer.Momentum(
... learning_rate=0.1,
... parameters=linear.parameters(),
... weight_decay=0.01
... )
>>> back = out.backward()
>>> momentum.step()
>>> momentum.clear_grad()
>>> # Note that the learning_rate of linear_2 is 0.01.
>>> linear_1 = paddle.nn.Linear(10, 10)
>>> linear_2 = paddle.nn.Linear(10, 10)
>>> inp = paddle.uniform(shape=[10, 10], min=-0.1, max=0.1)
>>> out = linear_1(inp)
>>> out = linear_2(out)
>>> loss = paddle.mean(out)
>>> momentum = paddle.optimizer.Momentum(
... learning_rate=0.1,
... parameters=[{ # type: ignore
... 'params': linear_1.parameters()
... }, {
... 'params': linear_2.parameters(),
... 'weight_decay': 0.001,
... 'learning_rate': 0.1
... }],
... weight_decay=0.01,
... momentum=0.9
... )
>>> out.backward()
>>> momentum.step()
>>> momentum.clear_grad()
"""
_velocity_acc_str = "velocity"
def __init__(
self,
learning_rate: float | Tensor | LRScheduler = 0.001,
momentum: float = 0.9,
parameters: (
Sequence[Tensor] | Sequence[_MomentumParameterConfig] | None
) = None,
use_nesterov: bool = False,
weight_decay: float | WeightDecayRegularizer | None = None,
grad_clip: GradientClipBase | None = None,
multi_precision: bool = False,
rescale_grad: float = 1.0,
use_multi_tensor: bool = False,
name: str | None = None,
) -> None:
if learning_rate is None:
raise ValueError("learning_rate is not set")
if momentum is None:
raise ValueError("momentum is not set")
if isinstance(weight_decay, int):
weight_decay = float(weight_decay)
predicate = lambda regular: isinstance(regular, (L2Decay, float))
if isinstance(parameters, list):
if isinstance(parameters[0], dict):
for param_group in parameters:
decay = (
param_group['weight_decay']
if 'weight_decay' in param_group
else weight_decay
)
reg_method, reg_coeff = self._update_regularization(decay)
param_group['regularization_method'] = reg_method
param_group['regularization_coeff'] = reg_coeff
py_regular = None if predicate(decay) else decay
param_group['weight_decay'] = py_regular
py_regular = None if predicate(weight_decay) else weight_decay
super().__init__(
learning_rate=learning_rate,
parameters=parameters,
weight_decay=py_regular,
grad_clip=grad_clip,
name=name,
)
self.type = "momentum"
self._momentum = momentum
self._use_nesterov = bool(use_nesterov)
(
self._regularization_method,
self._regularization_coeff,
) = self._update_regularization(weight_decay)
self._multi_precision = multi_precision
self._rescale_grad = rescale_grad
self._master_weights = {}
self._default_dict = {
'momentum': momentum,
'use_nesterov': use_nesterov,
'rescale_grad': rescale_grad,
'regularization_method': self._regularization_method,
'regularization_coeff': self._regularization_coeff,
}
self._use_multi_tensor = use_multi_tensor
if self._use_multi_tensor:
self._param_dict = self._create_multi_tensor_dict()
self._velocity_dict = self._create_multi_tensor_dict()
self._master_weight_dict = self._create_multi_tensor_dict()
self._master_weight_dict['FP32_DenseTensor'] = None
self._regularization_method_dict = self._create_multi_tensor_dict()
self._regularization_coeff_dict = self._create_multi_tensor_dict()
def _update_regularization(self, weight_decay):
reg_method = ""
reg_coeff = 0.0
if isinstance(weight_decay, L2Decay):
reg_method = "l2_decay"
reg_coeff = weight_decay._coeff
if isinstance(weight_decay, float):
reg_method = "l2_decay"
reg_coeff = weight_decay
return reg_method, reg_coeff
def _create_accumulators(self, block, parameters):
'''
if framework.in_dynamic_mode():
return
'''
assert isinstance(block, (framework.Block, paddle.pir.Block))
if isinstance(parameters, dict):
parameters = self._update_param_group(parameters)
for p in parameters:
if p.name in self._already_create_accumulator:
continue
if self._multi_precision and self._is_dtype_fp16_or_bf16(p.dtype):
master_p = self._create_master_weight(p)
self._add_accumulator(self._velocity_acc_str, master_p)
self._already_create_accumulator.add(p.name)
continue
if (
self._is_dtype_fp16_or_bf16(p.dtype)
and not self._multi_precision
):
warnings.warn(
"Accumulating with FP16/BF16 in optimizer can lead to poor accuracy or slow convergence."
"Consider using multi_precision=True option of the Momentum optimizer."
)
self._add_accumulator(self._velocity_acc_str, p)
self._already_create_accumulator.add(p.name)
def _create_regularization_of_grad(self, param, grad, regularization=None):
"""Create and add backward regularization Operators
Function helper of append_regularization_ops.
"""
# If ParamAttr is set to L2Decay, we skip doing regularization here. And then we fused
# L2Decay with momentum which can refer to _append_optimize_op below.
if hasattr(param, 'regularizer') and isinstance(
param.regularizer, L2Decay
):
return grad
return super()._create_regularization_of_grad(
param, grad, regularization
)
def _append_optimize_op(self, block, param_and_grad):
if not isinstance(block, (framework.Block, pir.Block)):
raise TypeError("block is not instance of Block.")
if isinstance(param_and_grad, dict):
param_and_grad = self._update_param_group(param_and_grad)
velocity_acc = self._get_accumulator_master(
self._velocity_acc_str, param_and_grad[0]
)
lr = self._create_param_lr(param_and_grad)
# For fusion of momentum and l2decay
param = param_and_grad[0]
regularization_method = self._regularization_method
regularization_coeff = self._regularization_coeff
if hasattr(param, 'regularizer'):
# we skip param's l2decay before, so fuse it with momentum here.
if isinstance(param.regularizer, L2Decay):
regularization_method = "l2_decay"
regularization_coeff = param.regularizer._coeff
# the param's regularization has been done before, we avoid do l2decay in momentum.
elif param.regularizer is not None:
regularization_method = ""
regularization_coeff = 0.0
find_master = self._multi_precision and self._is_dtype_fp16_or_bf16(
param_and_grad[0].dtype
)
master_weight = (
self._master_weights[param_and_grad[0].name]
if find_master
else None
)
if in_dynamic_or_pir_mode():
if isinstance(param_and_grad, dict):
self._update_regularization(param_and_grad['weight_decay'])
return _C_ops.momentum_(
param_and_grad[0],
param_and_grad[1],
velocity_acc,
lr,
master_weight,
self._momentum,
self._use_nesterov,
regularization_method,
regularization_coeff,
find_master,
self._rescale_grad,
)
else:
attrs = {
"mu": self._momentum,
"use_nesterov": self._use_nesterov,
"regularization_method": regularization_method,
"regularization_coeff": regularization_coeff,
"multi_precision": find_master,
"rescale_grad": self._rescale_grad,
}
inputs = {
"Param": [param_and_grad[0]],
"Grad": [param_and_grad[1]],
"Velocity": [velocity_acc],
"LearningRate": [lr],
}
outputs = {
"ParamOut": [param_and_grad[0]],
"VelocityOut": [velocity_acc],
}
if find_master:
inputs["MasterParam"] = master_weight
outputs["MasterParamOut"] = master_weight
# create the momentum optimize op
momentum_op = block.append_op(
type=self.type,
inputs=inputs,
outputs=outputs,
attrs=attrs,
stop_gradient=True,
)
return momentum_op
def _multi_tensor_init(self, target_block, parameters, param_group_idx):
"""
All parameters used for optimizer (such as: parameters, master_weight, velocity_acc for momentum) calculations are grouped into a python list by data type (float16, bf16, float32).
This function will be overridden in the corresponding optimizer file.
Args:
target_block: the block in which the loss tensor is present
parameters: list of parameter tensors for the optimizer
"""
self._create_accumulators(target_block, parameters)
for param in parameters:
velocity_acc = self._get_accumulator_master(
self._velocity_acc_str, param
)
regularization_method = self._regularization_method
regularization_coeff = self._regularization_coeff
if hasattr(param, 'regularizer'):
# we skip param's l2decay before, so fuse it with momentum here.
if isinstance(param.regularizer, L2Decay):
regularization_method = "l2_decay"
regularization_coeff = param.regularizer._coeff
elif param.regularizer is not None:
regularization_method = ""
regularization_coeff = 0.0
if param.dtype == paddle.float32:
self._param_dict['FP32_DenseTensor'][param_group_idx].append(
param
)
self._velocity_dict['FP32_DenseTensor'][param_group_idx].append(
velocity_acc
)
# fp32 no master weight
self._regularization_method_dict['FP32_DenseTensor'][
param_group_idx
].append(regularization_method)
self._regularization_coeff_dict['FP32_DenseTensor'][
param_group_idx
].append(regularization_coeff)
elif self._is_dtype_fp16_or_bf16(param.dtype):
self._param_dict['FP16_DenseTensor'][param_group_idx].append(
param
)
self._velocity_dict['FP16_DenseTensor'][param_group_idx].append(
velocity_acc
)
if self._multi_precision:
self._master_weight_dict['FP16_DenseTensor'][
param_group_idx
].append(self._master_weights[param.name])
else:
self._master_weight_dict['FP16_DenseTensor'][
param_group_idx
] = None
self._regularization_method_dict['FP16_DenseTensor'][
param_group_idx
].append(regularization_method)
self._regularization_coeff_dict['FP16_DenseTensor'][
param_group_idx
].append(regularization_coeff)
else:
raise ValueError(
"Now multi_tensor_momentum only support fp32, fp16 or bf16 parameters and grad is DENSE_TENSOR."
)
def _append_optimize_multi_tensor_op(
self,
target_block,
parameters_and_grads,
param_group_idx,
):
"""
For Multi Tensor, append optimize merged_operator to block.
"""
assert isinstance(target_block, framework.Block)
grad_dict = {'FP32_DenseTensor': [], 'FP16_DenseTensor': []}
lr_dict = {'FP32_DenseTensor': [], 'FP16_DenseTensor': []}
if isinstance(parameters_and_grads, list):
for param_and_grad in parameters_and_grads:
if param_and_grad[1] is None:
continue
if param_and_grad[0].stop_gradient is False:
if (
param_and_grad[0].dtype == paddle.float32
and param_and_grad[1].type
== core.VarDesc.VarType.DENSE_TENSOR
):
grad_dict['FP32_DenseTensor'].append(param_and_grad[1])
lr = self._create_param_lr(param_and_grad)
lr_dict['FP32_DenseTensor'].append(lr)
elif (
self._is_dtype_fp16_or_bf16(param_and_grad[0].dtype)
and param_and_grad[1].type
== core.VarDesc.VarType.DENSE_TENSOR
):
grad_dict['FP16_DenseTensor'].append(param_and_grad[1])
lr = self._create_param_lr(param_and_grad)
lr_dict['FP16_DenseTensor'].append(lr)
else:
for param_and_grad in parameters_and_grads['params']:
if param_and_grad[1] is None:
continue
if param_and_grad[0].stop_gradient is False:
param_grad_dict = {}
param_grad_dict['params'] = param_and_grad
param_grad_dict.update(
{
k: v
for k, v in parameters_and_grads.items()
if k != 'params'
}
)
param_and_grad = self._update_param_group(param_grad_dict)
if (
param_and_grad[0].dtype == paddle.float32
and param_and_grad[1].type
== core.VarDesc.VarType.DENSE_TENSOR
):
grad_dict['FP32_DenseTensor'].append(param_and_grad[1])
lr = self._create_param_lr(param_and_grad)
lr_dict['FP32_DenseTensor'].append(lr)
elif (
self._is_dtype_fp16_or_bf16(param_and_grad[0].dtype)
and param_and_grad[1].type
== core.VarDesc.VarType.DENSE_TENSOR
):
grad_dict['FP16_DenseTensor'].append(param_and_grad[1])
lr = self._create_param_lr(param_and_grad)
lr_dict['FP16_DenseTensor'].append(lr)
multi_tensor_list = ['FP32_DenseTensor', 'FP16_DenseTensor']
for key in multi_tensor_list:
if len(self._param_dict[key][param_group_idx]) > 0:
find_master = (
self._multi_precision and key == 'FP16_DenseTensor'
)
master_weight = self._master_weight_dict[key]
master_weight = (
master_weight[param_group_idx]
if master_weight is not None
else None
)
if in_dynamic_or_pir_mode():
found_inf = self._get_auxiliary_var('found_inf')
if found_inf:
if isinstance(
found_inf, (core.eager.Tensor, paddle.pir.Value)
):
self._set_auxiliary_var('found_inf', True)
else:
if isinstance(
found_inf, (core.eager.Tensor, paddle.pir.Value)
):
self._set_auxiliary_var('found_inf', False)
_, _, _ = _C_ops.merged_momentum_(
self._param_dict[key][param_group_idx],
grad_dict[key],
self._velocity_dict[key][param_group_idx],
lr_dict[key],
master_weight,
self._momentum,
self._use_nesterov,
self._regularization_method_dict[key][
param_group_idx
],
self._regularization_coeff_dict[key][
param_group_idx
],
find_master,
self._rescale_grad,
)
else:
inputs = {
"Param": self._param_dict[key][param_group_idx],
"Grad": grad_dict[key],
"Velocity": self._velocity_dict[key][param_group_idx],
"LearningRate": lr_dict[key],
}
outputs = {
"ParamOut": self._param_dict[key][param_group_idx],
"VelocityOut": self._velocity_dict[key][
param_group_idx
],
}
attrs = {
"mu": self._momentum,
"use_nesterov": self._use_nesterov,
"regularization_method": self._regularization_method_dict[
key
][param_group_idx],
"regularization_coeff": self._regularization_coeff_dict[
key
][param_group_idx],
}
if find_master:
inputs["MasterParam"] = self._master_weight_dict[key][
param_group_idx
]
outputs["MasterParamOut"] = self._master_weight_dict[
key
][param_group_idx]
attrs["multi_precision"] = find_master
target_block.append_op(
type="merged_momentum",
inputs=inputs,
outputs=outputs,
attrs=attrs,
stop_gradient=True,
)
def _update_param_group(self, parameters):
self._momentum = parameters.get(
'momentum', self._default_dict['momentum']
)
self._use_nesterov = parameters.get(
'use_nesterov', self._default_dict['use_nesterov']
)
self._rescale_grad = parameters.get(
'rescale_grad', self._default_dict['rescale_grad']
)
self._regularization_method = parameters.get(
'regularization_method', self._default_dict['regularization_method']
)
self._regularization_coeff = parameters.get(
'regularization_coeff', self._default_dict['regularization_coeff']
)
parameters = parameters.get('params')
return parameters
python/paddle/optimizer/muon.py
+837
View File
@@ -0,0 +1,837 @@
# Copyright (c) 2024 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 logging
import os
from typing import TYPE_CHECKING
if TYPE_CHECKING:
from collections.abc import Callable
from dataclasses import dataclass
from typing import TYPE_CHECKING
if TYPE_CHECKING:
from paddle import Tensor
import paddle
from paddle import _C_ops
from paddle.base import framework
from paddle.distributed.flex_checkpoint.dcp.sharded_weight import (
ShardedStateDict,
ShardedWeight,
create_sharded_weight_with_new_local,
)
from ..nn.clip import GradientClipBase
from .optimizer import Optimizer
# Debug logging for Muon optimizer
_logger = logging.getLogger(__name__)
MUON_DEBUG = os.environ.get("MUON_DEBUG", "0") == "1"
__all__ = []
# ------------------------------------------------------------------
# Parameter metadata dataclasses
# ------------------------------------------------------------------
@dataclass
class MuonParamInfo:
"""Muon update metadata for a single parameter.
This replaces the previous approach of setting dynamic attributes
directly on param objects.
Attributes:
use_muon: If True, use Muon (orthogonal) updates; otherwise AdamW.
split_concat_func: Optional callable that implements the slice strategy.
Signature: split_concat_func(matrix, ortho_fn, **kwargs) -> sliced_matrix
If None, whole-matrix orthogonalisation is used.
"""
use_muon: bool = True
split_concat_func: Callable | None = None
# Type alias for the parameter info mapping
MuonParamInfoMap = dict[str, MuonParamInfo]
# ------------------------------------------------------------------
# Newton-Schulz coefficient sets
# ------------------------------------------------------------------
_NS_COEFFICIENT_SETS = {
# Simple coefficient set (original)
"simple": [
(3.4445, -4.7750, 2.0315),
],
# Quintic iteration with optimized coefficients
# Source: https://leloykun.github.io/ponder/muon-opt-coeffs/
"quintic": [
(4.0848, -6.8946, 2.9270),
(3.9505, -6.3029, 2.6377),
(3.7418, -5.5913, 2.3037),
(2.8769, -3.1427, 1.2046),
(2.8366, -3.0525, 1.2012),
],
# Polar Express iteration from https://arxiv.org/abs/2505.16932
"polar_express": [
(8.2051, -22.9019, 16.4607),
(4.0664, -2.8612, 0.5184),
(3.9096, -2.8234, 0.5250),
(3.2856, -2.4153, 0.4853),
(2.2779, -1.6198, 0.3985),
(1.8726, -1.2307, 0.3585),
(1.8564, -1.2132, 0.3568),
(1.8750, -1.2500, 0.3750),
],
# AOL coefficients from https://github.com/thib-s/flash-newton-schulz
"aol": [
(4.0098, -7.0585, 2.4635),
(3.4585, -5.5479, 2.5959),
(2.7573, -3.2939, 1.4254),
(2.7215, -3.0494, 1.3169),
],
"deepseekv4":
# From DeepSeekV4: https://huggingface.co/deepseek-ai/DeepSeek-V4-Pro/resolve/main/DeepSeek_V4.pdf
[(3.4445, -4.7750, 2.0315)] * 8 + [(2.0, -1.5, 0.5)] * 2,
}
# ------------------------------------------------------------------
# Default parameter classification
# ------------------------------------------------------------------
def _default_should_use_muon(name, shape, exclude_patterns):
"""Default fallback logic for determining if a parameter should use Muon.
This is only used when param.is_muon is not set. The actual exclusion
patterns must be configured via training_args.muon_exclude_patterns in yaml.
Args:
name: Parameter name.
shape: Parameter shape.
exclude_patterns: List of substrings to exclude from Muon updates.
Must be provided (e.g., ['embed', 'bias', 'lm_head', 'mlp.gate']).
Returns:
True if the parameter should use Muon (orthogonal) updates.
Raises:
ValueError: If exclude_patterns is None.
"""
if exclude_patterns is None:
raise ValueError(
"muon_exclude_patterns must be set in yaml config. "
"Example: muon_exclude_patterns: ['embed', 'bias', 'lm_head', 'mlp.gate']"
)
if len(shape) not in (2, 3):
return False
name_lower = name.lower()
for pattern in exclude_patterns:
if pattern.lower() in name_lower:
return False
return True
class Muon(Optimizer):
r"""
Muon optimizer for MuonShardingOptimizer (Sharding Stage1 V3) usage.
For 2-D weight matrices (identified by :func:`_default_should_use_muon`), Muon
applies orthogonal gradient updates via Newton-Schulz iteration. For all
other parameters (embeddings, biases, expert weights, …) it falls back to
a standard AdamW update.
Designed for ``MuonShardingOptimizer`` (Sharding Stage1 V3), where 2D parameters are
assigned as whole tensors to ranks. Currently we do not support TP=1, no sharding gather
or TP communication is needed during the optimizer step.
Args:
learning_rate (float | LRScheduler): Learning rate. Default: ``0.02``.
parameters (list[Tensor]): Flat list of parameters to optimize.
momentum (float): Momentum coefficient for the Muon update. Default: ``0.95``.
adam_beta1 (float): β₁ for the AdamW fallback. Default: ``0.9``.
adam_beta2 (float): β₂ for the AdamW fallback. Default: ``0.95``.
weight_decay (float): Decoupled weight decay. Default: ``0.01``.
ns_steps (int): Newton-Schulz iteration steps. Default: ``5``.
ns_coeff_type (str): Preset name for Newton-Schulz coefficients.
Options: ``"simple"``, ``"quintic"``, ``"polar_express"``,
``"aol"``, ``"deepseekv4"``, ``"custom"``. Default: ``"simple"``.
ns_coeffs (list[tuple[float, float, float]] | None): Custom
Newton-Schulz coefficient set. Each tuple is ``(a, b, c)``
for one iteration step. Default: ``None``.
Only used when ns_coeff_type=``custom``.
nesterov (bool): Use Nesterov momentum in Muon. Default: ``True``.
adam_epsilon (float): ε for numerical stability in AdamW. Default: ``1e-9``.
grad_clip (GradientClipBase | None): Gradient clipping. Default: ``None``.
apply_decay_param_fun (callable | None): Function to select which
parameters receive weight decay. Default: ``None``.
muon_version (int): Scaling-function version (1/2/3). Default: ``1``.
muon_exclude_patterns (list[str] | None): Parameter names containing
any of these substrings will use AdamW instead of Muon.
Example: ``['embed', 'bias', 'lm_head', 'mlp.gate']``.
Default: ``None``.
muon_extra_scale_factor (float): Extra multiplicative scale applied
after the dimension-dependent scaling in ``_scaling_fn``.
Default: ``0.2``.
muon_param_info_map (MuonParamInfoMap | None): Per-parameter metadata
dict mapping param name to :class:`MuonParamInfo` (use_muon,
split_concat_func). Built by Trainer and passed in.
Default: ``None``.
ns_matmul_dtype (paddle.dtype | None): Dtype for Newton-Schulz matmul
iterations. ``None`` = auto-detect: bfloat16 on Ampere+ (capability
>= 8.0), float32 on V100 and older. Pass ``paddle.float32``
explicitly to force float32. Default: ``None``.
multi_precision (bool): Maintain FP32 master weights when training in
BF16/FP16. Default: ``False``.
name (str | None): Optional name for the optimizer instance.
"""
_moment_acc_str = "moment1"
_moment2_acc_str = "moment2"
_beta1_pow_acc_str = "beta1_pow_acc"
_beta2_pow_acc_str = "beta2_pow_acc"
def __init__(
self,
learning_rate=0.02,
parameters=None,
momentum=0.95,
adam_beta1=0.9,
adam_beta2=0.95,
weight_decay=0.01,
ns_steps=5,
ns_coeff_type="simple",
ns_coeffs=None,
nesterov=True,
adam_epsilon=1e-9,
grad_clip=None,
lr_ratio: Callable[[Tensor], float] | None = None,
apply_decay_param_fun: Callable[[str], bool] | None = None,
muon_version=1,
muon_exclude_patterns=None,
muon_extra_scale_factor=0.2,
muon_param_info_map: MuonParamInfoMap | None = None,
ns_matmul_dtype=None,
multi_precision=False,
name=None,
**kwargs,
):
if parameters is None:
raise ValueError(
"parameters argument given to the Optimizer should not be None."
)
if not isinstance(parameters, list):
raise TypeError("parameters must be a list.")
if len(parameters) > 0 and isinstance(parameters[0], dict):
raise TypeError(
"Muon optimizer only supports a flat list of parameters, "
"not a list of parameter groups."
)
if grad_clip is not None and not isinstance(
grad_clip, GradientClipBase
):
raise TypeError(
"'grad_clip' should be an instance of GradientClipBase's derived class"
)
defaults = {
"momentum": momentum,
"adam_beta1": adam_beta1,
"adam_beta2": adam_beta2,
"weight_decay": weight_decay,
"ns_steps": ns_steps,
"nesterov": nesterov,
"epsilon": adam_epsilon,
"muon_version": muon_version,
"ns_coeff_type": ns_coeff_type,
}
super().__init__(
learning_rate=learning_rate,
parameters=parameters,
weight_decay=weight_decay,
grad_clip=grad_clip,
name=name,
)
self._multi_precision = multi_precision
self._master_weights = {}
self._lr_ratio = lr_ratio
self._apply_decay_param_fun = apply_decay_param_fun
self._muon_split_logged = False
self._muon_exclude_patterns = muon_exclude_patterns
self._muon_extra_scale_factor = muon_extra_scale_factor
self._ns_coeff_type = ns_coeff_type
if ns_coeff_type == "custom":
assert ns_coeffs is not None, (
"ns_coeffs must be provided when ns_coeff_type is 'custom'."
)
self._ns_coeffs = ns_coeffs
else:
assert ns_coeff_type in _NS_COEFFICIENT_SETS, (
f"Invalid ns_coeff_type: {ns_coeff_type}"
)
self._ns_coeffs = _NS_COEFFICIENT_SETS[ns_coeff_type]
self._muon_param_info_map = muon_param_info_map or {}
# Dtype for Newton-Schulz matmul.
# None = auto: bfloat16 on Ampere+ (capability >= 8.0), float32 on older.
if ns_matmul_dtype is None:
cap = (
paddle.device.cuda.get_device_capability()
if paddle.is_compiled_with_cuda()
else (0, 0)
)
self._ns_matmul_dtype = (
paddle.bfloat16 if cap[0] >= 8 else paddle.float32
)
else:
self._ns_matmul_dtype = ns_matmul_dtype
self._default_dict.update(defaults)
# ------------------------------------------------------------------
# Accumulator management
# ------------------------------------------------------------------
def _ensure_accumulators(self, param, use_muon, group):
"""Create optimizer accumulators for *param* if they do not exist yet."""
if (
self._moment_acc_str in self._accumulators
and param.name in self._accumulators[self._moment_acc_str]
):
return
# FP32 master weight for mixed-precision training
if self._multi_precision and self._is_dtype_fp16_or_bf16(param.dtype):
if param.name not in self._master_weights:
self._create_master_weight(param)
self._add_accumulator(
self._moment_acc_str,
param,
dtype=paddle.float32,
fill_value=0.0,
shape=param.shape,
type=framework.core.VarDesc.VarType.DENSE_TENSOR,
)
if not use_muon:
# AdamW-specific states
self._add_accumulator(
self._moment2_acc_str,
param,
dtype=paddle.float32,
fill_value=0.0,
shape=param.shape,
type=framework.core.VarDesc.VarType.DENSE_TENSOR,
)
for acc_name, init_val in [
(self._beta1_pow_acc_str, group.get("adam_beta1", 0.9)),
(self._beta2_pow_acc_str, group.get("adam_beta2", 0.95)),
]:
self._add_accumulator(
acc_name,
param,
dtype=paddle.float32,
fill_value=init_val,
shape=[1],
type=framework.core.VarDesc.VarType.DENSE_TENSOR,
)
def _create_accumulators(self, block, parameters):
"""Standard entry-point used by checkpoint-resume infrastructure."""
if isinstance(parameters, dict):
parameters = self._update_param_group(parameters)
for p in parameters:
param_info = self._muon_param_info_map.get(p.name)
if param_info is not None:
use_muon = param_info.use_muon
else:
use_muon = _default_should_use_muon(
p.name,
getattr(p, "original_shape", p.shape),
self._muon_exclude_patterns,
)
self._ensure_accumulators(p, use_muon, self._default_dict)
# ------------------------------------------------------------------
# Newton-Schulz orthogonalisation
# ------------------------------------------------------------------
@staticmethod
def _zeropower_via_newtonschulz5(
X,
steps=5,
eps=1e-9,
ns_coeffs=None,
ns_matmul_dtype=paddle.bfloat16,
):
"""Approximate the matrix sign function via Newton-Schulz iteration.
Args:
X: Input tensor to orthogonalize. Must be 2D (M, N) or
3D (B, M, N) for batched operation.
steps: Number of Newton-Schulz iterations.
eps: Small constant for numerical stability.
ns_coeffs: List of (a, b, c) coefficient tuples for iteration.
If None, uses the "simple" preset.
ns_matmul_dtype: Dtype for matmul iterations. Defaults to
bfloat16. Pass paddle.float32 for V100 compatibility.
"""
if X.ndim < 2 or X.ndim > 3:
raise ValueError(
f"Input tensor X must be 2D or 3D (batched), got {X.ndim}D"
)
coeff_sets = (
ns_coeffs
if ns_coeffs is not None
else _NS_COEFFICIENT_SETS["simple"]
)
if X.shape[-2] > X.shape[-1]:
X = paddle.transpose(
X,
perm=[1, 0] if X.ndim == 2 else [0, 2, 1],
)
transpose = True
else:
transpose = False
orig_shape = X.shape
X_flat = X.flatten(start_axis=-2)
X_flat = paddle.nn.functional.normalize(
X_flat, p=2, axis=-1, epsilon=eps
)
X = X_flat.reshape(orig_shape).astype(ns_matmul_dtype)
if X.ndim == 3:
ns_step_fn = Muon._batched_newton_schulz_step
else:
ns_step_fn = Muon._newton_schulz_step
for i in range(steps):
a, b, c = coeff_sets[i % len(coeff_sets)]
X = ns_step_fn(X, a, b, c)
if transpose:
X = paddle.transpose(X, perm=[1, 0] if X.ndim == 2 else [0, 2, 1])
return X
@staticmethod
def _newton_schulz_step(X, a, b, c):
"""Single Newton-Schulz iteration step for 2D input."""
A = paddle.matmul(X, X, transpose_y=True)
B = paddle.addmm(input=A, x=A, y=A, beta=b, alpha=c)
X = paddle.addmm(input=X, x=B, y=X, beta=a, alpha=1.0)
return X
@staticmethod
def _batched_newton_schulz_step(X, a, b, c):
"""Single Newton-Schulz iteration step for 3D batched input."""
A = paddle.matmul(X, X, transpose_y=True)
B = paddle.baddbmm(A, A, A, beta=b, alpha=c)
X = paddle.baddbmm(X, B, X, beta=a, alpha=1.0)
return X
@staticmethod
def _scaling_fn(orthogonal_update, version, extra_scale_factor=1.0):
"""Apply dimension-dependent scaling to the orthogonal update."""
din, dout = orthogonal_update.shape[-2], orthogonal_update.shape[-1]
if version == 1:
scale = max(1, dout / din) ** 0.5
elif version == 2:
scale = (dout / din) ** 0.5
else: # version == 3 (default)
scale = max(dout, din) ** 0.5
return orthogonal_update * scale * extra_scale_factor
# ------------------------------------------------------------------
# Per-parameter update rules
# ------------------------------------------------------------------
def _adamw_update(
self,
param,
grad,
lr,
moment1,
moment2,
beta1_pow,
beta2_pow,
beta1,
beta2,
epsilon,
weight_decay,
):
"""In-place AdamW update for 1-D sharded parameters."""
lr_ratio = 1.0 if self._lr_ratio is None else self._lr_ratio(param)
with_decay = True
if (
self._apply_decay_param_fun is not None
and not self._apply_decay_param_fun(param.name)
):
with_decay = False
find_master = param.name in self._master_weights
master_weight = (
self._master_weights[param.name] if find_master else None
)
_, _, _, _, _, _, _ = _C_ops.adamw_(
param,
grad,
lr,
moment1,
moment2,
None, # moment2_max
beta1_pow,
beta2_pow,
master_weight,
None, # found_inf
beta1,
beta2,
epsilon,
lr_ratio,
weight_decay,
with_decay,
False, # lazy_mode
1000,
find_master,
False,
False, # amsgrad
)
def _muon_update(
self,
param,
grad,
lr,
momentum_buffer,
momentum_beta,
ns_steps,
nesterov,
epsilon,
weight_decay,
version,
):
"""In-place Muon update for a 2D parameter tensor.
Applies Newton-Schulz orthogonalisation to the 2D weight matrix and
updates the parameter in-place. MuonShardingOptimizer assigns whole
2D tensors to ranks, so no sharding gather or TP communication is needed.
"""
param_shape = getattr(param, "original_shape", param.shape)
param_info = self._muon_param_info_map.get(param.name)
with paddle.no_grad():
grad_f32 = (
grad.astype(momentum_buffer.dtype)
if grad.dtype != momentum_buffer.dtype
else grad
)
# Step 1: Momentum update
new_momentum = paddle.lerp(
momentum_buffer, grad_f32, 1.0 - momentum_beta
)
paddle.assign(new_momentum, momentum_buffer)
update_buffer = (
paddle.lerp(grad_f32, momentum_buffer, momentum_beta)
if nesterov
else momentum_buffer
)
# Step 2: Reshape update buffer to 2D matrix.
# MuonShardingOptimizer assigns whole 2D tensors to ranks, so params
# are already 2D/3D (no sharding gather needed).
matrix_2d_global = update_buffer.reshape(param_shape)
# Shared NS + scaling closure (captures ns_steps, epsilon, version, ns_coeffs)
def ortho_fn(m):
ns_out = Muon._zeropower_via_newtonschulz5(
m,
steps=ns_steps,
eps=epsilon,
ns_coeffs=self._ns_coeffs,
ns_matmul_dtype=self._ns_matmul_dtype,
)
scaled = Muon._scaling_fn(
ns_out, version, self._muon_extra_scale_factor
)
return scaled
# Step 3: Newton-Schulz orthogonalisation
# Use split_concat_func from param_info if provided, otherwise default to whole matrix
if (
param_info is not None
and param_info.split_concat_func is not None
):
# Use slice function defined in model configuration
orthogonal_update = param_info.split_concat_func(
matrix_2d_global, ortho_fn
)
if MUON_DEBUG:
_global_rank = paddle.distributed.get_rank()
if _global_rank == 0:
_sf = param_info.split_concat_func
_logger.info(
f"[Muon] Using split_concat_func: param={param.name}, "
f"split_concat_func={_sf.func.__name__}, "
f"args={_sf.args}, kwargs={_sf.keywords}"
)
else:
# Default: whole matrix orthogonalisation
orthogonal_update = ortho_fn(matrix_2d_global)
find_master = param.name in self._master_weights
master_weight = (
self._master_weights[param.name] if find_master else None
)
with_decay = True
if (
self._apply_decay_param_fun is not None
and not self._apply_decay_param_fun(param.name)
):
with_decay = False
if with_decay and weight_decay > 0:
if find_master:
master_weight.scale_(1.0 - lr * weight_decay)
else:
param.scale_(1.0 - lr * weight_decay)
final_step = orthogonal_update * lr
if find_master:
master_weight.subtract_(final_step)
paddle.assign(master_weight.astype(param.dtype), param)
else:
param.subtract_(final_step.astype(param.dtype))
# ------------------------------------------------------------------
# Core optimization step
# ------------------------------------------------------------------
def _apply_optimize(self, loss, startup_program, params_grads):
if not framework.in_dygraph_mode():
raise NotImplementedError(
"Muon optimizer only supports dygraph mode."
)
if self._grad_clip is not None:
params_grads = self._grad_clip(params_grads)
# apply for zcc
self._maybe_refuse()
group = self._default_dict
lr = self._learning_rate
if isinstance(lr, paddle.optimizer.lr.LRScheduler):
lr = lr()
wd = group.get("weight_decay", 0.0)
muon_params = []
adamw_params = []
for param, grad in params_grads:
if grad is None:
continue
param_info = self._muon_param_info_map.get(param.name)
assert param_info is not None, (
f"muon_param_info_map does not have {param.name}"
)
use_muon = param_info.use_muon
self._ensure_accumulators(param, use_muon, group)
if use_muon:
muon_params.append((param, grad))
else:
adamw_params.append((param, grad))
# --- Pass 1: Muon updates (large temporary tensors) ---
lr_tensor = paddle.to_tensor(lr, dtype=paddle.float32)
lr_tensor_f64 = paddle.to_tensor(lr, dtype=paddle.float64)
for param, grad in muon_params:
self._muon_update(
param,
grad,
lr_tensor,
self._get_accumulator(self._moment_acc_str, param),
group.get("momentum", 0.95),
group.get("ns_steps", 5),
group.get("nesterov", True),
group.get("epsilon", 1e-9),
wd,
version=group.get("muon_version", 3),
)
# --- Pass 2: AdamW updates ---
for param, grad in adamw_params:
self._adamw_update(
param,
grad,
lr_tensor_f64,
self._get_accumulator(self._moment_acc_str, param),
self._get_accumulator(self._moment2_acc_str, param),
self._get_accumulator(self._beta1_pow_acc_str, param),
self._get_accumulator(self._beta2_pow_acc_str, param),
group.get("adam_beta1", 0.9),
group.get("adam_beta2", 0.95),
group.get("epsilon", 1e-9),
wd,
)
@framework.dygraph_only
def step(self) -> None:
params_grads = [
(param, param._grad_ivar())
for param in self._parameter_list
if not param.stop_gradient and param._grad_ivar() is not None
]
self._apply_optimize(
loss=None, startup_program=None, params_grads=params_grads
)
def sharded_state_dict(
self,
model_sharded_state_dict: ShardedStateDict,
) -> ShardedStateDict:
"""Build a sharded optimizer state dict for flex checkpoint save/load.
The layout mirrors :class:`paddle.optimizer.AdamW`'s implementation so
that the same ``dist.save_state_dict`` / ``dist.load_state_dict`` path
works for Muon checkpoints.
Args:
model_sharded_state_dict: Sharded model state dict produced by
``model.sharded_state_dict()``.
Returns:
A dict mapping ``"<struct_name>.<state_type>"`` keys to
:class:`ShardedWeight` objects.
"""
_FP32_MASTER = "fp32_master_0"
_optimizer_scalar_names = [
"beta1_pow_acc_0",
"beta2_pow_acc_0",
]
_optimizer_vector_names = [
"moment1_0",
"moment2_0",
]
def _split_state_name(vname):
if _FP32_MASTER in vname:
return tuple(vname.split("_" + _FP32_MASTER + "_", 1))
for suffix in _optimizer_scalar_names + _optimizer_vector_names:
if vname.endswith(suffix):
return vname[: -(len(suffix) + 1)], suffix
raise ValueError(
f"Cannot parse optimizer state variable name: {vname!r}"
)
model_sharded_state_dict = dict(
sorted(model_sharded_state_dict.items())
)
# Build static-name → struct-name mapping (handles shared weights)
static_to_struct = {}
for struct_name, sw in model_sharded_state_dict.items():
local_name = sw.local_tensor.name
if local_name not in static_to_struct:
static_to_struct[local_name] = struct_name
optimizer_state_dict = self.state_dict()
master_weights = optimizer_state_dict.pop("master_weights", None)
optimizer_state_dict.pop("LR_Scheduler", None)
sharded_state: ShardedStateDict = {}
# Optimizer states (moment1, moment2, beta_pow scalars)
for key, tensor in optimizer_state_dict.items():
static_name, state_type = _split_state_name(key)
struct_name = static_to_struct[static_name]
sharded_param = model_sharded_state_dict[struct_name]
unified_name = f"{struct_name}.{state_type}"
if state_type in _optimizer_vector_names:
# Vector states share the same sharding layout as the parameter
if tensor.is_dist():
sharded_state[unified_name] = ShardedWeight(
key=unified_name,
local_tensor=tensor,
local_shape=tensor.shape,
global_shape=tensor.shape,
global_offset=sharded_param.global_offset,
)
else:
# Reshape accumulator if numel matches but shape differs.
# MoE: grouped_gemm_experts param.shape is 3D
# [n_experts, H, I] but model.state_dict() returns actual
# C++ storage shape 2D [n_experts*H, I]. moment1 was
# created with 3D shape, so we need to reshape here.
# V2 is unaffected: its moments are always 1D shards,
# so shape always matches and reshape is never triggered.
target_shape = sharded_param.local_shape
if (
tuple(tensor.shape) != tuple(target_shape)
and tensor.numel()
== paddle.to_tensor(list(target_shape)).prod().item()
):
tensor = tensor.reshape(target_shape)
sharded_state[unified_name] = (
create_sharded_weight_with_new_local(
unified_name, tensor, sharded_param
)
)
else:
# Scalar states (beta_pow) are replicated – save as-is
sharded_state[unified_name] = ShardedWeight(
key=unified_name,
local_tensor=tensor,
local_shape=(1,),
global_shape=(1,),
global_offset=(0,),
)
# FP32 master weights
if master_weights:
for weight_key, tensor in master_weights.items():
struct_name = static_to_struct[weight_key]
sharded_param = model_sharded_state_dict[struct_name]
unified_name = f"{struct_name}.w_0"
if tensor.is_dist():
sharded_state[unified_name] = ShardedWeight(
key=unified_name,
local_tensor=tensor,
local_shape=tensor.shape,
global_shape=tensor.shape,
global_offset=sharded_param.global_offset,
)
else:
sharded_state[unified_name] = (
create_sharded_weight_with_new_local(
unified_name, tensor, sharded_param
)
)
return sharded_state
python/paddle/optimizer/nadam.py
+363
View File
@@ -0,0 +1,363 @@
# Copyright (c) 2024 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 typing import TYPE_CHECKING
from paddle import _C_ops, pir
from paddle.base.libpaddle import DataType
from ..base import core, framework
from ..base.framework import (
in_dynamic_or_pir_mode,
in_pir_mode,
)
from .optimizer import Optimizer
if TYPE_CHECKING:
from collections.abc import Sequence
from typing_extensions import NotRequired
from paddle import Tensor
from paddle.nn.clip import GradientClipBase
from .lr import LRScheduler
from .optimizer import _ParameterConfig
class _NAdamParameterConfig(_ParameterConfig):
beta1: NotRequired[float | Tensor]
beta2: NotRequired[float | Tensor]
epsilon: NotRequired[float]
momentum_decay: NotRequired[float]
__all__ = []
class NAdam(Optimizer):
r"""
The NAdam optimizer is implemented based on the Adam Optimization
in paper `Incorporating Nesterov Momentum into Adam <https://openreview.net/forum?id=OM0jvwB8jIp57ZJjtNEZ>`_.
The main improvement is to combine the advantages of Nesterov momentum and Adam adaptive learning rate.
.. math::
\begin{aligned}
&\textbf{for} \: t=1 \: \textbf{to} \: \ldots \: \textbf{do} \\
&\hspace{5mm} \mu_t \leftarrow \beta_1 \big(1 - \frac{1}{2} \rho ^{t \psi} \big) \\
&\hspace{5mm} \mu_{t+1} \leftarrow \beta_1 \big(1 - \frac{1}{2} 0.96 ^{(t+1)\psi}\big)\\
&\hspace{5mm}m_t \leftarrow \beta_1 m_{t-1} + (1 - \beta_1) g_t \\
&\hspace{5mm}v_t \leftarrow \beta_2 v_{t-1} + (1-\beta_2) g^2_t \\
&\hspace{5mm}\widehat{m_t} \leftarrow \mu_{t+1} m_t/(1-\prod_{i=1}^{t+1}\mu_i) + (1-\mu_t) g_t /(1-\prod_{i=1}^{t} \mu_{i}) \\
&\hspace{5mm}\widehat{v_t} \leftarrow v_t/\big(1-\beta_2^t \big) \\
&\hspace{5mm}\theta_t \leftarrow \theta_t - \gamma \widehat{m_t}/
\big(\sqrt{\widehat{v_t}} + \epsilon \big) \\
&\hspace{0mm} \text{ with: } \gamma_t \text{ (lr)}, \: \beta_1,\beta_2 \text{ (betas)}, \: \theta_0 \text{ (params)}, \: f(\theta) \text{ (objective)} \\
&\hspace{0mm} \: \lambda \text{ (weight decay)}, \:\psi \text{ (momentum decay)} \\
\end{aligned}
Args:
learning_rate (float|LRScheduler, optional): The learning rate used to update ``Parameter``.
It can be a float value or a LRScheduler. The default value is 0.002.
parameters (list|tuple|None, optional): List/Tuple of ``Tensor`` names to update to minimize ``loss``.
This parameter is required in dygraph mode. And you can specify different options for
different parameter groups such as the learning rate, weight decay, etc,
then the parameters are list of dict. Note that the learning_rate in parameter groups
represents the scale of base learning_rate.
The default value is None in static graph mode, at this time all parameters will be updated.
beta1 (float|Tensor, optional): The exponential decay rate for the 1st moment estimates.
It should be a float number or a 0-D Tensor with shape [] and data type as float32.
The default value is 0.9.
beta2 (float|Tensor, optional): The exponential decay rate for the 2nd moment estimates.
It should be a float number or a 0-D Tensor with shape [] and data type as float32.
The default value is 0.999.
epsilon (float, optional): A small float value for numerical stability.
The default value is 1e-08.
weight_decay (int|float|Tensor|None, optional): The weight decay coefficient, it can be int, float or Tensor.
Default None, meaning there is no regularization.
momentum_decay (float, optional): momentum momentum_decay. The default value is 0.004.
grad_clip (GradientClipBase|None, optional): Gradient clipping strategy, it's an instance of
some derived class of ``GradientClipBase`` . There are three clipping strategies
( :ref:`api_paddle_nn_ClipGradByGlobalNorm` , :ref:`api_paddle_nn_ClipGradByNorm` ,
:ref:`api_paddle_nn_ClipGradByValue` ). Default None, meaning there is no gradient clipping.
name (str|None, optional): Normally there is no need for user to set this property.
For more information, please refer to :ref:`api_guide_Name`.
The default value is None.
Notes:
Currently, NAdam doesn't support sparse parameter optimization.
Examples:
.. code-block:: pycon
>>> import paddle
>>> inp = paddle.rand([10,10], dtype="float32")
>>> linear = paddle.nn.Linear(10, 10)
>>> out = linear(inp)
>>> loss = paddle.mean(out)
>>> nadam = paddle.optimizer.NAdam(
... learning_rate=0.1,
... parameters=linear.parameters()
... )
>>> out.backward()
>>> nadam.step()
>>> nadam.clear_grad()
>>> # Note that the learning_rate of linear_2 is 0.01.
>>> linear_1 = paddle.nn.Linear(10, 10)
>>> linear_2 = paddle.nn.Linear(10, 10)
>>> inp = paddle.uniform(shape=[10, 10], min=-0.1, max=0.1)
>>> out = linear_1(inp)
>>> out = linear_2(out)
>>> loss = paddle.mean(out)
>>> opt = paddle.optimizer.NAdam(
... learning_rate=0.1,
... parameters=[{ # type: ignore
... 'params': linear_1.parameters()
... }, {
... 'params': linear_2.parameters(),
... 'weight_decay': 0.001,
... 'learning_rate': 0.1,
... 'beta1': 0.8
... }],
... weight_decay=0.01,
... beta1=0.9
... )
>>> loss.backward()
>>> opt.step()
>>> opt.clear_grad()
"""
_momentum_decay_pow_acc_str = "momentum_decay_pow"
_beta2_pow_acc_str = "beta2_pow"
_mu_product_acc_str = "mu_product"
_moment1_acc_str = "moment1"
_moment2_acc_str = "moment2"
def __init__(
self,
learning_rate: float | LRScheduler = 0.002,
beta1: float | Tensor = 0.9,
beta2: float | Tensor = 0.999,
epsilon: float = 1.0e-8,
momentum_decay: float = 0.004,
parameters: (
Sequence[Tensor] | Sequence[_NAdamParameterConfig] | None
) = None,
weight_decay: float | Tensor | None = None,
grad_clip: GradientClipBase | None = None,
name: str | None = None,
) -> None:
if isinstance(learning_rate, (float, int)) and not 0.0 <= learning_rate:
raise ValueError(
f"Invalid learning rate: {learning_rate}, expect learning_rate >= 0."
)
if not 0.0 <= epsilon:
raise ValueError(
f"Invalid epsilon value: {epsilon}, expect epsilon >= 0."
)
if not 0.0 <= beta1 < 1.0:
raise ValueError(
f"Invalid beta1: {beta1}, expect 0. <= beta1 < 1.0."
)
if not 0.0 <= beta2 < 1.0:
raise ValueError(
f"Invalid beta2: {beta2}, expect 0. <= beta2 < 1.0."
)
if not 0.0 <= momentum_decay:
raise ValueError(
f"Invalid momentum_decay value: {momentum_decay}, expect momentum_decay >= 0."
)
super().__init__(
learning_rate=learning_rate,
parameters=parameters,
weight_decay=weight_decay,
grad_clip=grad_clip,
name=name,
)
self.type = "nadam"
self._beta1 = beta1
self._beta2 = beta2
self._epsilon = epsilon
self._momentum_decay = momentum_decay
self._multi_precision = False
self._master_weights = {}
self._default_dict = {
'beta1': beta1,
'beta2': beta2,
'epsilon': epsilon,
'momentum_decay': momentum_decay,
}
def _add_moments_pows(self, p):
acc_dtype = p.dtype
if self._is_dtype_fp16_or_bf16(acc_dtype):
acc_dtype = (
DataType.FLOAT32 if in_pir_mode() else core.VarDesc.VarType.FP32
)
self._add_accumulator(
name=self._momentum_decay_pow_acc_str,
param=p,
dtype=acc_dtype,
fill_value=1.0,
)
self._add_accumulator(
name=self._beta2_pow_acc_str,
param=p,
dtype=acc_dtype,
fill_value=1.0,
)
self._add_accumulator(
name=self._mu_product_acc_str,
param=p,
dtype=acc_dtype,
fill_value=1.0,
)
self._add_accumulator(self._moment1_acc_str, p, dtype=acc_dtype)
self._add_accumulator(self._moment2_acc_str, p, dtype=acc_dtype)
def _create_accumulators(self, block, parameters):
if not isinstance(block, (framework.Block, pir.Block)):
raise TypeError("block is not instance of Block.")
if isinstance(parameters, dict):
parameters = parameters.get('params')
for p in parameters:
if p.name in self._already_create_accumulator:
continue
if self._multi_precision and self._is_dtype_fp16_or_bf16(p.dtype):
master_p = self._create_master_weight(p)
self._add_moments_pows(master_p)
self._already_create_accumulator.add(p.name)
continue
if (
self._is_dtype_fp16_or_bf16(p.dtype)
and not self._multi_precision
):
warnings.warn(
"Accumulating with FP16 in optimizer can lead to poor accuracy or slow convergence."
"Consider using multi_precision=True option of the Lars optimizer."
)
self._add_moments_pows(p)
self._already_create_accumulator.add(p.name)
def _append_optimize_op(self, block, param_and_grad):
if not isinstance(block, (framework.Block, pir.Block)):
raise TypeError("block is not instance of Block.")
if isinstance(param_and_grad, dict):
param_and_grad = self._update_param_group(param_and_grad)
momentum_decay_pow_acc = self._get_accumulator_master(
self._momentum_decay_pow_acc_str, param_and_grad[0]
)
beta2_pow_acc = self._get_accumulator_master(
self._beta2_pow_acc_str, param_and_grad[0]
)
mu_product_acc = self._get_accumulator_master(
self._mu_product_acc_str, param_and_grad[0]
)
moment1_acc = self._get_accumulator_master(
self._moment1_acc_str, param_and_grad[0]
)
moment2_acc = self._get_accumulator_master(
self._moment2_acc_str, param_and_grad[0]
)
find_master = self._multi_precision and self._is_dtype_fp16_or_bf16(
param_and_grad[0].dtype
)
master_weight = (
self._master_weights[param_and_grad[0].name]
if find_master
else None
)
if in_dynamic_or_pir_mode():
_C_ops.nadam_(
param_and_grad[0],
param_and_grad[1],
self._create_param_lr(param_and_grad),
momentum_decay_pow_acc,
beta2_pow_acc,
mu_product_acc,
moment1_acc,
moment2_acc,
master_weight,
self._beta1,
self._beta2,
self._epsilon,
self._momentum_decay,
find_master,
)
return None
else:
inputs = {
"param": param_and_grad[0],
"grad": param_and_grad[1],
"momentum_decay_pow": momentum_decay_pow_acc,
"beta2_pow": beta2_pow_acc,
"mu_product": mu_product_acc,
"moment1": moment1_acc,
"moment2": moment2_acc,
"learning_rate": self._create_param_lr(param_and_grad),
}
outputs = {
"param_out": param_and_grad[0],
"momentum_decay_pow_out": momentum_decay_pow_acc,
"beta2_pow_out": beta2_pow_acc,
"mu_product_out": mu_product_acc,
"moment1_out": moment1_acc,
"moment2_out": moment2_acc,
}
if find_master:
inputs["master_param"] = master_weight
outputs["master_param_out"] = master_weight
nadam_op = block.append_op(
type=self.type,
inputs=inputs,
outputs=outputs,
attrs={
"epsilon": self._epsilon,
"beta1": self._beta1,
"beta2": self._beta2,
"momentum_decay": self._momentum_decay,
},
stop_gradient=True,
)
return nadam_op
def _update_param_group(self, parameters):
self._epsilon = parameters.get('epsilon', self._default_dict['epsilon'])
self._beta1 = parameters.get('beta1', self._default_dict['beta1'])
self._beta2 = parameters.get('beta2', self._default_dict['beta2'])
self._momentum_decay = parameters.get(
'momentum_decay', self._default_dict['momentum_decay']
)
parameters = parameters.get('params')
return parameters
python/paddle/optimizer/optimizer.py
+2282
View File
File diff suppressed because it is too large Load Diff
python/paddle/optimizer/radam.py
+354
View File
@@ -0,0 +1,354 @@
# Copyright (c) 2024 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 typing import TYPE_CHECKING
from paddle import _C_ops, pir
from paddle.base.libpaddle import DataType
from ..base import core, framework
from ..base.framework import (
in_dynamic_or_pir_mode,
in_pir_mode,
)
from .optimizer import Optimizer
if TYPE_CHECKING:
from collections.abc import Sequence
from typing_extensions import NotRequired
from paddle import Tensor
from paddle.nn.clip import GradientClipBase
from paddle.optimizer.lr import LRScheduler
from paddle.regularizer import WeightDecayRegularizer
from .optimizer import _ParameterConfig
class _RAdamParameterConfig(_ParameterConfig):
beta1: NotRequired[float | Tensor]
beta2: NotRequired[float | Tensor]
epsilon: NotRequired[float]
__all__ = []
class RAdam(Optimizer):
r"""
The RAdam optimizer is implemented based on the Adam Optimization
in paper `On the Variance of the Adaptive Learning Rate and Beyond <https://arxiv.org/abs/1908.03265>`_.
RAdam improved the initial stability of training by modifying Adam's momentum term.
.. math::
\begin{aligned}
&\textbf{for} \: t=1 \: \textbf{to} \: \ldots \: \textbf{do} \\
&\hspace{6mm}m_t \leftarrow \beta_1 m_{t-1} + (1 - \beta_1) g_t \\
&\hspace{6mm}v_t \leftarrow \beta_2 v_{t-1} + (1-\beta_2) g^2_t \\
&\hspace{6mm}\widehat{m_t} \leftarrow m_t/\big(1-\beta_1^t \big) \\
&\hspace{6mm}\rho_t \leftarrow \rho_{\infty} -
2 t \beta^t_2 /\big(1-\beta_2^t \big) \\
&\hspace{6mm}\textbf{if} \: \rho_t > 5 \\
&\hspace{12mm} l_t \leftarrow \frac{\sqrt{ (1-\beta^t_2) }}{ \sqrt{v_t} +\epsilon } \\
&\hspace{12mm} r_t \leftarrow
\sqrt{\frac{(\rho_t-4)(\rho_t-2)\rho_{\infty}}{(\rho_{\infty}-4)(\rho_{\infty}-2) \rho_t}} \\
&\hspace{12mm}\theta_t \leftarrow \theta_t - \gamma \widehat{m_t} r_t l_t \\
&\hspace{6mm}\textbf{else} \\
&\hspace{12mm}\theta_t \leftarrow \theta_t - \gamma \widehat{m_t} \\
&\hspace{0mm} \text{ with: } \gamma_t \text{ (lr)}, \: \beta_1,\beta_2 \text{ (betas)}, \: \theta_t \text{ (params)} \\
&\hspace{0mm} \rho_{\infty} \leftarrow 2/(1-\beta_2) -1
\end{aligned}
Args:
learning_rate (float|LRScheduler, optional): The learning rate used to update ``Parameter``.
It can be a float value or a LRScheduler. The default value is 0.001.
parameters (list|tuple|None, optional): List/Tuple of ``Tensor`` names to update to minimize ``loss``.
This parameter is required in dygraph mode. And you can specify different options for
different parameter groups such as the learning rate, weight decay, etc,
then the parameters are list of dict. Note that the learning_rate in parameter groups
represents the scale of base learning_rate.
The default value is None in static graph mode, at this time all parameters will be updated.
beta1 (float|Tensor, optional): The exponential decay rate for the 1st moment estimates.
It should be a float number or a 0-D Tensor with shape [] and data type as float32.
The default value is 0.9.
beta2 (float|Tensor, optional): The exponential decay rate for the 2nd moment estimates.
It should be a float number or a 0-D Tensor with shape [] and data type as float32.
The default value is 0.999.
epsilon (float, optional): A small float value for numerical stability.
The default value is 1e-08.
weight_decay (int|float|Tensor|WeightDecayRegularizer|None, optional): The weight decay coefficient, it can be int, float or Tensor.
Default None, meaning there is no regularization.
grad_clip (GradientClipBase|None, optional): Gradient clipping strategy, it's an instance of
some derived class of ``GradientClipBase`` . There are three clipping strategies
( :ref:`api_paddle_nn_ClipGradByGlobalNorm` , :ref:`api_paddle_nn_ClipGradByNorm` ,
:ref:`api_paddle_nn_ClipGradByValue` ). Default None, meaning there is no gradient clipping.
name (str|None, optional): Normally there is no need for user to set this property.
For more information, please refer to :ref:`api_guide_Name`.
The default value is None.
Note:
Currently, RAdam doesn't support sparse parameter optimization.
Examples:
.. code-block:: pycon
>>> import paddle
>>> inp = paddle.rand([10,10], dtype="float32")
>>> linear = paddle.nn.Linear(10, 10)
>>> out = linear(inp)
>>> loss = paddle.mean(out)
>>> radam = paddle.optimizer.RAdam(
... learning_rate=0.1,
... parameters=linear.parameters()
... )
>>> out.backward()
>>> radam.step()
>>> radam.clear_grad()
>>> # Note that the learning_rate of linear_2 is 0.01.
>>> linear_1 = paddle.nn.Linear(10, 10)
>>> linear_2 = paddle.nn.Linear(10, 10)
>>> inp = paddle.uniform(shape=[10, 10], min=-0.1, max=0.1)
>>> out = linear_1(inp)
>>> out = linear_2(out)
>>> loss = paddle.mean(out)
>>> opt = paddle.optimizer.RAdam(
... learning_rate=0.1,
... parameters=[{ # type: ignore
... 'params': linear_1.parameters()
... }, {
... 'params': linear_2.parameters(),
... 'weight_decay': 0.001,
... 'learning_rate': 0.1,
... 'beta1': 0.8
... }],
... weight_decay=0.01,
... beta1=0.9
... )
>>> loss.backward()
>>> opt.step()
>>> opt.clear_grad()
"""
_beta1_pow_acc_str = "beta1_pow"
_beta2_pow_acc_str = "beta2_pow"
_rho_acc_str = "rho"
_moment1_acc_str = "moment1"
_moment2_acc_str = "moment2"
def __init__(
self,
learning_rate: float | LRScheduler = 0.001,
beta1: float | Tensor = 0.9,
beta2: float | Tensor = 0.999,
epsilon: float = 1.0e-8,
parameters: (
Sequence[Tensor] | Sequence[_RAdamParameterConfig] | None
) = None,
weight_decay: float | Tensor | WeightDecayRegularizer | None = None,
grad_clip: GradientClipBase | None = None,
name: str | None = None,
) -> None:
if isinstance(learning_rate, (float, int)) and not 0.0 <= learning_rate:
raise ValueError(
f"Invalid learning rate: {learning_rate}, expect learning_rate >= 0."
)
if not 0.0 <= epsilon:
raise ValueError(
f"Invalid epsilon value: {epsilon}, expect epsilon >= 0."
)
if not 0.0 <= beta1 < 1.0:
raise ValueError(
f"Invalid beta1: {beta1}, expect 0. <= beta1 < 1.0."
)
if not 0.0 <= beta2 < 1.0:
raise ValueError(
f"Invalid beta2: {beta2}, expect 0. <= beta2 < 1.0."
)
super().__init__(
learning_rate=learning_rate,
parameters=parameters,
weight_decay=weight_decay,
grad_clip=grad_clip,
name=name,
)
self.type = "radam"
self._beta1 = beta1
self._beta2 = beta2
self._epsilon = epsilon
self._multi_precision = False
self._master_weights = {}
self._default_dict = {
'beta1': beta1,
'beta2': beta2,
'epsilon': epsilon,
}
def _add_moments_pows(self, p):
acc_dtype = p.dtype
if self._is_dtype_fp16_or_bf16(acc_dtype):
acc_dtype = (
DataType.FLOAT32 if in_pir_mode() else core.VarDesc.VarType.FP32
)
self._add_accumulator(
name=self._beta1_pow_acc_str,
param=p,
dtype=acc_dtype,
fill_value=1.0,
)
self._add_accumulator(
name=self._beta2_pow_acc_str,
param=p,
dtype=acc_dtype,
fill_value=1.0,
)
self._add_accumulator(
name=self._rho_acc_str,
param=p,
dtype=acc_dtype,
fill_value=1.0,
)
self._add_accumulator(self._moment1_acc_str, p, dtype=acc_dtype)
self._add_accumulator(self._moment2_acc_str, p, dtype=acc_dtype)
def _create_accumulators(self, block, parameters):
if not isinstance(block, (framework.Block, pir.Block)):
raise TypeError("block is not instance of Block.")
if isinstance(parameters, dict):
parameters = parameters.get('params')
for p in parameters:
if p.name in self._already_create_accumulator:
continue
if self._multi_precision and self._is_dtype_fp16_or_bf16(p.dtype):
master_p = self._create_master_weight(p)
self._add_moments_pows(master_p)
self._already_create_accumulator.add(p.name)
continue
if (
self._is_dtype_fp16_or_bf16(p.dtype)
and not self._multi_precision
):
warnings.warn(
"Accumulating with FP16 in optimizer can lead to poor accuracy or slow convergence."
"Consider using multi_precision=True option of the Lars optimizer."
)
self._add_moments_pows(p)
self._already_create_accumulator.add(p.name)
def _append_optimize_op(self, block, param_and_grad):
if not isinstance(block, (framework.Block, pir.Block)):
raise TypeError("block is not instance of Block.")
if isinstance(param_and_grad, dict):
param_and_grad = self._update_param_group(param_and_grad)
beta1_pow_acc = self._get_accumulator_master(
self._beta1_pow_acc_str, param_and_grad[0]
)
beta2_pow_acc = self._get_accumulator_master(
self._beta2_pow_acc_str, param_and_grad[0]
)
rho_acc = self._get_accumulator_master(
self._rho_acc_str, param_and_grad[0]
)
moment1_acc = self._get_accumulator_master(
self._moment1_acc_str, param_and_grad[0]
)
moment2_acc = self._get_accumulator_master(
self._moment2_acc_str, param_and_grad[0]
)
find_master = self._multi_precision and self._is_dtype_fp16_or_bf16(
param_and_grad[0].dtype
)
master_weight = (
self._master_weights[param_and_grad[0].name]
if find_master
else None
)
if in_dynamic_or_pir_mode():
_C_ops.radam_(
param_and_grad[0],
param_and_grad[1],
self._create_param_lr(param_and_grad),
beta1_pow_acc,
beta2_pow_acc,
rho_acc,
moment1_acc,
moment2_acc,
master_weight,
self._beta1,
self._beta2,
self._epsilon,
find_master,
)
return None
else:
inputs = {
"param": param_and_grad[0],
"grad": param_and_grad[1],
"beta1_pow": beta1_pow_acc,
"beta2_pow": beta2_pow_acc,
"rho": rho_acc,
"moment1": moment1_acc,
"moment2": moment2_acc,
"learning_rate": self._create_param_lr(param_and_grad),
}
outputs = {
"param_out": param_and_grad[0],
"beta1_pow_out": beta1_pow_acc,
"beta2_pow_out": beta2_pow_acc,
"rho_out": rho_acc,
"moment1_out": moment1_acc,
"moment2_out": moment2_acc,
}
if find_master:
inputs["master_param"] = master_weight
outputs["master_param_out"] = master_weight
radam_op = block.append_op(
type=self.type,
inputs=inputs,
outputs=outputs,
attrs={
"epsilon": self._epsilon,
"beta1": self._beta1,
"beta2": self._beta2,
},
stop_gradient=True,
)
return radam_op
def _update_param_group(self, parameters):
self._epsilon = parameters.get('epsilon', self._default_dict['epsilon'])
self._beta1 = parameters.get('beta1', self._default_dict['beta1'])
self._beta2 = parameters.get('beta2', self._default_dict['beta2'])
parameters = parameters.get('params')
return parameters
python/paddle/optimizer/rmsprop.py
+345
View File
@@ -0,0 +1,345 @@
# 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 typing import TYPE_CHECKING
from typing_extensions import NotRequired
from paddle import _C_ops, pir
from ..base import framework
from ..base.framework import in_dynamic_or_pir_mode
from .optimizer import Optimizer
if TYPE_CHECKING:
from collections.abc import Sequence
from typing_extensions import NotRequired
from paddle import Tensor
from paddle.nn.clip import GradientClipBase
from paddle.optimizer.lr import LRScheduler
from paddle.regularizer import WeightDecayRegularizer
from .optimizer import _ParameterConfig
class _RMSPropParameterConfig(_ParameterConfig):
epsilon: NotRequired[float]
momentum: NotRequired[float]
rho: NotRequired[float]
centered: NotRequired[bool]
__all__ = []
class RMSProp(Optimizer):
r"""
Root Mean Squared Propagation (RMSProp) is an unpublished, adaptive learning
rate method. The original slides proposed RMSProp: Slide 29 of
http://www.cs.toronto.edu/~tijmen/csc321/slides/lecture_slides_lec6.pdf .
The original equation is as follows:
.. math::
r(w, t) & = \rho r(w, t-1) + (1 - \rho)(\nabla Q_{i}(w))^2
w & = w - \frac{\eta} {\sqrt{r(w,t) + \epsilon}} \nabla Q_{i}(w)
The first equation calculates moving average of the squared gradient for
each weight. Then dividing the gradient by :math:`sqrt{v(w,t)}`.
In some cases, adding a momentum term :math: `\\beta` is beneficial.
In our implementation, Nesterov momentum is used:
.. math::
r(w, t) & = \rho r(w, t-1) + (1 - \rho)(\nabla Q_{i}(w))^2
v(w, t) & = \beta v(w, t-1) + \frac{\eta} {\sqrt{r(w,t) +
\epsilon}} \nabla Q_{i}(w)
w & = w - v(w, t)
if centered is True:
.. math::
r(w, t) & = \rho r(w, t-1) + (1 - \rho)(\nabla Q_{i}(w))^2
g(w, t) & = \rho g(w, t-1) + (1 - \rho)\nabla Q_{i}(w)
v(w, t) & = \beta v(w, t-1) + \frac{\eta} {\sqrt{r(w,t) - (g(w, t))^2 +
\epsilon}} \nabla Q_{i}(w)
w & = w - v(w, t)
where, :math:`\rho` is a hyperparameter and typical values are 0.9, 0.95
and so on. :math:`\beta` is the momentum term. :math:`\epsilon` is a
smoothing term to avoid division by zero, usually set somewhere in range
from 1e-4 to 1e-8.
Parameters:
learning_rate (float|LRScheduler): The learning rate used to update ``Parameter``.
It can be a float value or a LRScheduler.
rho(float, optional): rho is :math:`\rho` in equation, default is 0.95.
epsilon(float, optional): :math:`\epsilon` in equation is smoothing term to
avoid division by zero, default is 1e-6.
momentum(float, optional): :math:`\beta` in equation is the momentum term,
default is 0.0.
centered(bool, optional): If True, gradients are normalized by the estimated variance of
the gradient; if False, by the uncentered second moment. Setting this to
True may help with training, but is slightly more expensive in terms of
computation and memory. Defaults to False.
parameters (list|tuple|None, optional): List/Tuple of ``Tensor`` to update to minimize ``loss``.
This parameter is required in dygraph mode. And you can specify different options for
different parameter groups such as the learning rate, weight decay, etc,
then the parameters are list of dict. Note that the learning_rate in parameter groups
represents the scale of base learning_rate.
The default value is None in static graph mode, at this time all parameters will be updated.
weight_decay (int|float|WeightDecayRegularizer|None, optional): The strategy of regularization.
It can be a int or float value as coeff of L2 regularization or \
:ref:`api_paddle_regularizer_L1Decay`, :ref:`api_paddle_regularizer_L2Decay`.
If a parameter has set regularizer using :ref:`api_paddle_ParamAttr` already,
the regularization setting here in optimizer will be ignored for this parameter.
Otherwise, the regularization setting here in optimizer will take effect.
Default None, meaning there is no regularization.
grad_clip (GradientClipBase|None, optional): Gradient clipping strategy, it's an instance of
some derived class of ``GradientClipBase`` . There are three clipping strategies
( :ref:`api_paddle_nn_ClipGradByGlobalNorm` , :ref:`api_paddle_nn_ClipGradByNorm` ,
:ref:`api_paddle_nn_ClipGradByValue` ). Default None, meaning there is no gradient clipping.
name (str|None, optional): Normally there is no need for user to set this property.
For more information, please refer to :ref:`api_guide_Name`.
The default value is None.
Examples:
.. code-block:: pycon
>>> import paddle
>>> inp = paddle.rand([10,10], dtype="float32")
>>> linear = paddle.nn.Linear(10, 10)
>>> out = linear(inp)
>>> loss = paddle.mean(out)
>>> rmsprop = paddle.optimizer.RMSProp(
... learning_rate=0.1,
... parameters=linear.parameters(),
... weight_decay=0.01
... )
>>> out.backward()
>>> rmsprop.step()
>>> rmsprop.clear_grad()
>>> # Note that the learning_rate of linear_2 is 0.01.
>>> linear_1 = paddle.nn.Linear(10, 10)
>>> linear_2 = paddle.nn.Linear(10, 10)
>>> inp = paddle.uniform(shape=[10, 10], min=-0.1, max=0.1)
>>> out = linear_1(inp)
>>> out = linear_2(out)
>>> loss = paddle.mean(out)
>>> rmsprop = paddle.optimizer.RMSProp(
... learning_rate=0.1,
... parameters=[{ # type: ignore
... 'params': linear_1.parameters()
... }, {
... 'params': linear_2.parameters(),
... 'weight_decay': 0.001,
... 'learning_rate': 0.1
... }],
... weight_decay=0.01
... )
>>> out.backward()
>>> rmsprop.step()
>>> rmsprop.clear_grad()
"""
_momentum_acc_str = "momentum"
_mean_square_acc_str = "mean_square"
_mean_grad_acc_str = "mean_grad"
def __init__(
self,
learning_rate: float | LRScheduler,
rho: float = 0.95,
epsilon: float = 1.0e-6,
momentum: float = 0.0,
centered: bool = False,
parameters: (
Sequence[Tensor] | Sequence[_RMSPropParameterConfig] | None
) = None,
weight_decay: float | WeightDecayRegularizer | None = None,
grad_clip: GradientClipBase | None = None,
name: str | None = None,
) -> None:
if learning_rate is None:
raise ValueError("learning_rate is not set.")
if rho is None:
raise ValueError("rho is not set.")
if epsilon is None:
raise ValueError("epsilon is not set.")
if momentum is None:
raise ValueError("momentum is not set.")
if not 0.0 <= epsilon:
raise ValueError("Invalid value of epsilon, expect epsilon >= 0.")
if not 0.0 <= momentum:
raise ValueError("Invalid value of momentum, expect momentum >= 0.")
if not 0.0 <= rho:
raise ValueError("Invalid value of rho, expect rho >= 0.")
super().__init__(
learning_rate=learning_rate,
parameters=parameters,
weight_decay=weight_decay,
grad_clip=grad_clip,
name=name,
)
self.type = "rmsprop"
self._rho = rho
self._epsilon = epsilon
self._momentum = momentum
self._centered = centered
self._multi_precision = False
self._master_weights = {}
self._default_dict = {
'rho': rho,
'epsilon': epsilon,
'momentum': momentum,
'centered': centered,
}
def _create_accumulators(self, block, parameters):
if not isinstance(block, (framework.Block, pir.Block)):
raise TypeError("block is not instance of Block.")
if isinstance(parameters, dict):
parameters = parameters.get('params')
for p in parameters:
if p.name in self._already_create_accumulator:
continue
if self._multi_precision and self._is_dtype_fp16_or_bf16(p.dtype):
master_p = self._create_master_weight(p)
self._add_accumulator(self._momentum_acc_str, master_p)
self._add_accumulator(self._mean_square_acc_str, master_p)
self._add_accumulator(self._mean_grad_acc_str, master_p)
self._already_create_accumulator.add(p.name)
continue
if (
self._is_dtype_fp16_or_bf16(p.dtype)
and not self._multi_precision
):
warnings.warn(
"Accumulating with FP16 in optimizer can lead to poor accuracy or slow convergence."
"Consider using multi_precision=True option of the Lars optimizer."
)
self._add_accumulator(self._momentum_acc_str, p)
self._add_accumulator(self._mean_square_acc_str, p)
self._add_accumulator(self._mean_grad_acc_str, p)
self._already_create_accumulator.add(p.name)
def _append_optimize_op(self, block, param_and_grad):
if not isinstance(block, (framework.Block, pir.Block)):
raise TypeError("block is not instance of Block.")
if isinstance(param_and_grad, dict):
param_and_grad = self._update_param_group(param_and_grad)
momentum_acc = self._get_accumulator_master(
self._momentum_acc_str, param_and_grad[0]
)
mean_square_acc = self._get_accumulator_master(
self._mean_square_acc_str, param_and_grad[0]
)
mean_grad_acc = self._get_accumulator_master(
self._mean_grad_acc_str, param_and_grad[0]
)
find_master = self._multi_precision and self._is_dtype_fp16_or_bf16(
param_and_grad[0].dtype
)
master_weight = (
self._master_weights[param_and_grad[0].name]
if find_master
else None
)
if in_dynamic_or_pir_mode():
_C_ops.rmsprop_(
param_and_grad[0],
mean_square_acc,
param_and_grad[1],
momentum_acc,
self._create_param_lr(param_and_grad),
mean_grad_acc,
master_weight,
self._epsilon,
self._rho,
self._momentum,
self._centered,
find_master,
)
return None
else:
inputs = {
"Param": param_and_grad[0],
"Grad": param_and_grad[1],
"Moment": momentum_acc,
"MeanSquare": mean_square_acc,
"MeanGrad": mean_grad_acc,
"LearningRate": self._create_param_lr(param_and_grad),
}
outputs = {
"ParamOut": param_and_grad[0],
"MomentOut": momentum_acc,
"MeanSquareOut": mean_square_acc,
"MeanGradOut": mean_grad_acc,
}
if find_master:
inputs["MasterParam"] = master_weight
outputs["MasterParamOut"] = master_weight
rmsprop_op = block.append_op(
type=self.type,
inputs=inputs,
outputs=outputs,
attrs={
"epsilon": self._epsilon,
"decay": self._rho,
"momentum": self._momentum,
"centered": self._centered,
},
stop_gradient=True,
)
return rmsprop_op
def _update_param_group(self, parameters):
self._epsilon = parameters.get('epsilon', self._default_dict['epsilon'])
self._rho = parameters.get('rho', self._default_dict['rho'])
self._momentum = parameters.get(
'momentum', self._default_dict['momentum']
)
self._centered = parameters.get(
'centered', self._default_dict['centered']
)
parameters = parameters.get('params')
return parameters
python/paddle/optimizer/rprop.py
+286
View File
@@ -0,0 +1,286 @@
# 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 typing import TYPE_CHECKING
from paddle import _C_ops, pir
from paddle.tensor.creation import to_tensor
from ..base import framework
from ..base.dygraph import no_grad
from ..base.framework import in_dynamic_or_pir_mode
from .optimizer import Optimizer
if TYPE_CHECKING:
from collections.abc import Sequence
from paddle import Tensor
from paddle.nn.clip import GradientClipBase
from paddle.optimizer.lr import LRScheduler
from .optimizer import _ParameterConfig
__all__ = []
class Rprop(Optimizer):
r"""
**Notes: This optimizer is only applicable to full-batch training.**
Optimizer of the Rprop algorithm.Please refer to this for details:
`A direct adaptive method for faster backpropagation learning : The RPROP algorithm <https://ieeexplore.ieee.org/document/298623>`_.
.. math::
\begin{aligned}
&\hspace{0mm} For\ all\ weights\ and\ biases\{ \\
&\hspace{5mm} \textbf{if} \: (\frac{\partial E}{\partial w_{ij}}(t-1)*\frac{\partial E}{\partial w_{ij}}(t)> 0)\ \textbf{then} \: \{ \\
&\hspace{10mm} learning\_rate_{ij}(t)=\mathrm{minimum}(learning\_rate_{ij}(t-1)*\eta^{+},learning\_rate_{max}) \\
&\hspace{10mm} \Delta w_{ij}(t)=-sign(\frac{\partial E}{\partial w_{ij}}(t))*learning\_rate_{ij}(t) \\
&\hspace{10mm} w_{ij}(t+1)=w_{ij}(t)+\Delta w_{ij}(t) \\
&\hspace{5mm} \} \\
&\hspace{5mm} \textbf{else if} \: (\frac{\partial E}{\partial w_{ij}}(t-1)*\frac{\partial E}{\partial w_{ij}}(t)< 0)\ \textbf{then} \: \{ \\
&\hspace{10mm} learning\_rate_{ij}(t)=\mathrm{maximum}(learning\_rate_{ij}(t-1)*\eta^{-},learning\_rate_{min}) \\
&\hspace{10mm} w_{ij}(t+1)=w_{ij}(t) \\
&\hspace{10mm} \frac{\partial E}{\partial w_{ij}}(t)=0 \\
&\hspace{5mm} \} \\
&\hspace{5mm} \textbf{else if} \: (\frac{\partial E}{\partial w_{ij}}(t-1)*\frac{\partial E}{\partial w_{ij}}(t)= 0)\ \textbf{then} \: \{ \\
&\hspace{10mm} \Delta w_{ij}(t)=-sign(\frac{\partial E}{\partial w_{ij}}(t))*learning\_rate_{ij}(t) \\
&\hspace{10mm} w_{ij}(t+1)=w_{ij}(t)+\Delta w_{ij}(t) \\
&\hspace{5mm} \} \\
&\hspace{0mm} \} \\
\end{aligned}
Parameters:
learning_rate (float|Tensor|LRScheduler, optional): The initial learning rate used to update ``Parameter``.
It can be a float value, a ``Tensor`` with a float type or a LRScheduler. The default value is 0.001.
learning_rate_range (tuple, optional): The range of learning rate.
Learning rate cannot be smaller than the first element of the tuple;
learning rate cannot be larger than the second element of the tuple.
The default value is (1e-5, 50).
parameters (list|tuple|None, optional): List/Tuple of ``Tensor`` to update to minimize ``loss``.
This parameter is required in dygraph mode.
The default value is None in static graph mode, at this time all parameters will be updated.
etas (tuple, optional): Tuple used to update learning rate.
The first element of the tuple is the multiplicative decrease factor;
the second element of the tuple is the multiplicative increase factor.
The default value is (0.5, 1.2).
grad_clip (GradientClipBase|None, optional): Gradient clipping strategy, it's an instance of some derived class of ``GradientClipBase`` .
There are three clipping strategies ( :ref:`api_paddle_nn_ClipGradByGlobalNorm` , :ref:`api_paddle_nn_ClipGradByNorm` , :ref:`api_paddle_nn_ClipGradByValue` ).
Default None, meaning there is no gradient clipping.
multi_precision (bool, optional): In mixed precision training scenarios based on GPU,
this parameter is mainly used to ensure the numerical stability of gradient updates.
When it is set to True, the optimizer will save a backup of FP32 type parameters with an equal value for FP16 type parameters.
When updating gradients, first increase the gradient type to FP32, and then assign it to the FP32 type parameter backup.
Finally, the updated FP32 type value will be converted to FP16 type first,
and then assigned to the actual FP16 type parameters participating in the calculation.
The default value is False.
name (str|None, optional): The default value is None. Normally there is no need for user to set this property.
For more information, please refer to :ref:`api_guide_Name` .
Examples:
.. code-block:: pycon
>>> import paddle
>>> inp = paddle.uniform(min=-0.1, max=0.1, shape=[1, 100], dtype='float32')
>>> linear = paddle.nn.Linear(100, 10)
>>> inp = paddle.to_tensor(inp)
>>> out = linear(inp)
>>> loss = paddle.mean(out)
>>> rprop = paddle.optimizer.Rprop(
... learning_rate=0.001,
... learning_rate_range=(0.0001, 0.1),
... parameters=linear.parameters(),
... etas=(0.5, 1.2)
... )
>>> out.backward()
>>> rprop.step()
>>> rprop.clear_grad()
"""
_prevs_acc_str = "prevs"
_learning_rates_acc_str = "learning_rates"
def __init__(
self,
learning_rate: float | Tensor | LRScheduler = 0.001,
learning_rate_range: tuple[float, float] = (1e-5, 50),
parameters: Sequence[Tensor] | Sequence[_ParameterConfig] | None = None,
etas: tuple[float, float] = (0.5, 1.2),
grad_clip: GradientClipBase | None = None,
multi_precision: bool = False,
name: str | None = None,
) -> None:
if learning_rate is None:
raise ValueError("learning_rate is not set")
if (
not 0.0
< learning_rate_range[0]
<= learning_rate
<= learning_rate_range[1]
):
raise ValueError(
"'0.0 < learning_rate_range[0] <= learning_rate <= learning_rate_range[1]' must be true"
)
if not 0.0 < etas[0] < 1.0 < etas[1]:
raise ValueError("'0.0 < etas[0] < 1.0 < etas[1]' must be true")
super().__init__(
learning_rate=learning_rate,
parameters=parameters,
weight_decay=0.0,
grad_clip=grad_clip,
name=name,
)
self.type = "rprop"
self._initial_learning_rate = learning_rate
self._multi_precision = multi_precision
self._master_weights = {}
self._learning_rate_range = [learning_rate_range]
self._etas = [etas]
self._sign = True
def _to_tensor(self, block, dtype):
assert isinstance(block, (framework.Block, pir.Block))
self._learning_rate_range = to_tensor(
self._learning_rate_range, dtype=dtype
)
self._etas = to_tensor(self._etas, dtype=dtype)
def _create_accumulators(self, block, parameters):
if not isinstance(block, (framework.Block, pir.Block)):
raise TypeError("block is not instance of Block.")
if isinstance(parameters, dict):
parameters = self._update_param_group(parameters)
# Create accumulator tensors for first and second moments
for p in parameters:
if p.name in self._already_create_accumulator:
continue
if self._multi_precision and self._is_dtype_fp16_or_bf16(p.dtype):
master_p = self._create_master_weight(p)
self._add_accumulator(
self._prevs_acc_str,
master_p,
p.dtype,
0,
)
self._add_accumulator(
self._learning_rates_acc_str,
master_p,
p.dtype,
self._initial_learning_rate,
)
self._already_create_accumulator.add(p.name)
continue
if (
self._is_dtype_fp16_or_bf16(p.dtype)
and not self._multi_precision
):
warnings.warn(
"Accumulating with FP16/BF16 in optimizer can lead to poor accuracy or slow convergence."
"Consider using multi_precision=True option of the Adam optimizer."
)
self._add_accumulator(
self._prevs_acc_str,
p,
p.dtype,
0,
)
self._add_accumulator(
self._learning_rates_acc_str,
p,
p.dtype,
fill_value=self._initial_learning_rate,
)
self._already_create_accumulator.add(p.name)
@no_grad
def _append_optimize_op(self, block, param_and_grad):
if isinstance(param_and_grad, dict):
param_and_grad = self._update_param_group(param_and_grad)
if self._sign:
self._to_tensor(block, param_and_grad[0][0].dtype)
self._sign = False
prevs = self._get_accumulator_master(
self._prevs_acc_str, param_and_grad[0]
)
learning_rates = self._get_accumulator_master(
self._learning_rates_acc_str, param_and_grad[0]
)
find_master = self._multi_precision and self._is_dtype_fp16_or_bf16(
param_and_grad[0].dtype
)
master_weight = (
self._master_weights[param_and_grad[0].name]
if find_master
else None
)
if in_dynamic_or_pir_mode():
_C_ops.rprop_(
param_and_grad[0],
param_and_grad[1],
prevs,
learning_rates,
master_weight,
self._learning_rate_range,
self._etas,
find_master,
)
return None
else:
assert isinstance(block, framework.Block)
# create the optimize op
inputs = {
"param": param_and_grad[0],
"grad": param_and_grad[1],
"prev": prevs,
"learning_rate": learning_rates,
"learning_rate_range": self._learning_rate_range,
"etas": self._etas,
}
outputs = {
"param_out": param_and_grad[0],
"prev_out": prevs,
"learning_rate_out": learning_rates,
}
attrs = {"multi_precision": find_master}
if find_master:
inputs["master_param"] = master_weight
outputs["master_param_out"] = master_weight
rprop_op = block.append_op(
type=self.type,
inputs=inputs,
outputs=outputs,
attrs=attrs,
stop_gradient=True,
)
return rprop_op
def _update_param_group(self, parameters):
parameters = parameters.get('params')
return parameters
python/paddle/optimizer/sgd.py
+190
View File
@@ -0,0 +1,190 @@
# 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 typing import TYPE_CHECKING
from paddle import _C_ops, pir
from ..base import framework
from ..base.dygraph import no_grad
from ..base.framework import in_dynamic_or_pir_mode
from .optimizer import Optimizer
if TYPE_CHECKING:
from collections.abc import Sequence
from paddle import Tensor
from paddle.nn.clip import GradientClipBase
from paddle.regularizer import WeightDecayRegularizer
from .lr import LRScheduler
from .optimizer import _ParameterConfig
__all__ = []
class SGD(Optimizer):
r"""
Optimizer of the stochastic gradient descent algorithm.
.. math::
param\_out = param - learning\_rate * grad
Parameters:
learning_rate (float|LRScheduler, optional): The learning rate used to update ``Parameter``.
It can be a float value or a LRScheduler. The default value is 0.001.
parameters (list|tuple|None, optional): List/Tuple of ``Tensor`` to update to minimize ``loss``. \
This parameter is required in dygraph mode. \
The default value is None in static graph mode, at this time all parameters will be updated.
weight_decay (int|float|WeightDecayRegularizer|None, optional): The strategy of regularization. \
It can be a int or float value as coeff of L2 regularization or \
:ref:`api_paddle_regularizer_L1Decay`, :ref:`api_paddle_regularizer_L2Decay`.
If a parameter has set regularizer using :ref:`api_paddle_ParamAttr` already, \
the regularization setting here in optimizer will be ignored for this parameter. \
Otherwise, the regularization setting here in optimizer will take effect. \
Default None, meaning there is no regularization.
grad_clip (GradientClipBase|None, optional): Gradient clipping strategy, it's an instance of
some derived class of ``GradientClipBase`` . There are three clipping strategies
( :ref:`api_paddle_nn_ClipGradByGlobalNorm` , :ref:`api_paddle_nn_ClipGradByNorm` ,
:ref:`api_paddle_nn_ClipGradByValue` ). Default None, meaning there is no gradient clipping.
multi_precision (bool, optional): Whether to use multi-precision during weight updating.
name (str|None, optional): The default value is None. Normally there is no need for user
to set this property. For more information, please refer to
:ref:`api_guide_Name` .
Examples:
.. code-block:: pycon
>>> import paddle
>>> inp = paddle.uniform(min=-0.1, max=0.1, shape=[10, 10], dtype='float32')
>>> linear = paddle.nn.Linear(10, 10)
>>> inp = paddle.to_tensor(inp)
>>> out = linear(inp)
>>> loss = paddle.mean(out)
>>> sgd = paddle.optimizer.SGD(
... learning_rate=0.1,
... parameters=linear.parameters(),
... weight_decay=0.01
... )
>>> out.backward()
>>> sgd.step()
>>> sgd.clear_grad()
"""
type: str
def __init__(
self,
learning_rate: float | LRScheduler = 0.001,
parameters: Sequence[Tensor] | Sequence[_ParameterConfig] | None = None,
weight_decay: float | WeightDecayRegularizer | None = None,
grad_clip: GradientClipBase | None = None,
multi_precision: bool = False,
name: str | None = None,
) -> None:
if learning_rate is None:
raise ValueError("learning_rate is not set")
super().__init__(
learning_rate=learning_rate,
parameters=parameters,
weight_decay=weight_decay,
grad_clip=grad_clip,
name=name,
)
self.type = "sgd"
self._multi_precision = multi_precision
self._master_weights = {}
def _create_accumulators(self, block, parameters):
assert isinstance(block, (framework.Block, pir.Block))
if isinstance(parameters, dict):
parameters = self._update_param_group(parameters)
# Create accumulator tensors for first and second moments
for p in parameters:
if p.name in self._already_create_accumulator:
continue
if self._multi_precision and self._is_dtype_fp16_or_bf16(p.dtype):
master_p = self._create_master_weight(p)
self._already_create_accumulator.add(p.name)
continue
if (
self._is_dtype_fp16_or_bf16(p.dtype)
and not self._multi_precision
):
warnings.warn(
"Accumulating with FP16/BF16 in optimizer can lead to poor accuracy or slow convergence."
"Consider using multi_precision=True option of the Adam optimizer."
)
@no_grad
def _append_optimize_op(self, block, param_and_grad):
if isinstance(param_and_grad, dict):
param_and_grad = self._update_param_group(param_and_grad)
find_master = self._multi_precision and self._is_dtype_fp16_or_bf16(
param_and_grad[0].dtype
)
master_weight = (
self._master_weights[param_and_grad[0].name]
if find_master
else None
)
lr = self._create_param_lr(param_and_grad)
if in_dynamic_or_pir_mode():
_C_ops.sgd_(
param_and_grad[0],
lr,
param_and_grad[1],
master_weight,
find_master,
)
return None
else:
assert isinstance(block, framework.Block)
# create the optimize op
inputs = {
"Param": param_and_grad[0],
"Grad": param_and_grad[1],
"LearningRate": lr,
}
outputs = {"ParamOut": param_and_grad[0]}
attrs = {"multi_precision": find_master}
if find_master:
inputs["MasterParam"] = master_weight
outputs["MasterParamOut"] = master_weight
sgd_op = block.append_op(
type=self.type,
inputs=inputs,
outputs=outputs,
attrs=attrs,
stop_gradient=True,
)
return sgd_op
def _update_param_group(self, parameters):
parameters = parameters.get('params')
return parameters