chore: import upstream snapshot with attribution

python/paddle/incubate/optimizer/__init__.py

@@ -0,0 +1,25 @@
+#   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 functional  # noqa: F401
+from .distributed_fused_lamb import DistributedFusedLamb  # noqa: F401
+from .gradient_merge import GradientMergeOptimizer  # noqa: F401
+from .lars_momentum import LarsMomentumOptimizer  # noqa: F401
+from .lbfgs import LBFGS
+from .lookahead import LookAhead  # noqa: F401
+from .modelaverage import ModelAverage  # noqa: F401
+from .pipeline import PipelineOptimizer  # noqa: F401
+from .recompute import RecomputeOptimizer  # noqa: F401
+
+__all__ = ['LBFGS']

python/paddle/incubate/optimizer/distributed_fused_lamb.py

@@ -0,0 +1,516 @@
+# Copyright (c) 2021 PaddlePaddle Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import os
+
+import paddle
+from paddle.base import core, unique_name
+from paddle.base.executor import global_scope
+from paddle.base.framework import Variable, name_scope
+from paddle.base.layer_helper import LayerHelper
+from paddle.nn import ClipGradByGlobalNorm
+from paddle.optimizer import Optimizer
+
+
+def init_communicator(block, rank, ranks, ring_id):
+    eps = os.environ['PADDLE_TRAINER_ENDPOINTS']
+    eps = [ep.strip() for ep in eps.split(",") if ep.strip()]
+    cur_ep = eps[rank]
+    other_eps = [eps[r] for r in ranks if r != rank]
+
+    local_rank = ranks.index(rank)
+    comm_var_name = unique_name.generate('comm_id')
+    comm_id_var = block.create_var(
+        name=comm_var_name, persistable=True, type=core.VarDesc.VarType.RAW
+    )
+    if core.is_compiled_with_cuda():
+        block.append_op(
+            type='c_gen_nccl_id',
+            inputs={},
+            outputs={'Out': comm_id_var},
+            attrs={
+                'rank': local_rank,
+                'endpoint': cur_ep,
+                'other_endpoints': other_eps,
+                'ring_id': ring_id,
+            },
+        )
+    elif core.is_compiled_with_xpu():
+        block.append_op(
+            type='c_gen_bkcl_id',
+            inputs={},
+            outputs={'Out': comm_id_var},
+            attrs={
+                'rank': local_rank,
+                'endpoint': cur_ep,
+                'other_endpoints': other_eps,
+                'ring_id': ring_id,
+            },
+        )
+    elif (
+        paddle.distributed.ParallelEnv().device_type
+        in paddle.device.get_all_custom_device_type()
+    ):
+        block.append_op(
+            type='c_gen_xccl_id',
+            inputs={},
+            outputs={'Out': comm_id_var},
+            attrs={
+                'rank': local_rank,
+                'endpoint': cur_ep,
+                'other_endpoints': other_eps,
+                'ring_id': ring_id,
+            },
+        )
+    block.append_op(
+        type='c_comm_init',
+        inputs={'X': comm_id_var},
+        outputs={},
+        attrs={
+            'nranks': len(ranks),
+            'rank': local_rank,
+            'ring_id': ring_id,
+            'endpoints': ','.join(eps),
+        },
+    )
+    tmp_var = block.create_var(name=unique_name.generate('tmp'))
+    block.append_op(
+        type='fill_constant', outputs={'Out': tmp_var}, attrs={'value': 1}
+    )
+    block.append_op(
+        type='all_reduce',
+        inputs={'x': tmp_var},
+        outputs={'out': tmp_var},
+        attrs={
+            'ring_id': ring_id,
+            'reduce_type': paddle.distributed.ReduceOp.SUM,
+        },
+    )
+    block.append_op(
+        type='c_sync_calc_stream',
+        inputs={'X': tmp_var},
+        outputs={'Out': tmp_var},
+    )
+    return ring_id
+
+
+def broadcast_parameters(block, parameters, ring_id):
+    for p in parameters:
+        block.append_op(
+            type='broadcast',
+            inputs={'x': p},
+            outputs={'out': p},
+            attrs={
+                'ring_id': ring_id,
+            },
+        )
+
+
+class DistributedFusedLamb(Optimizer):
+    def __init__(
+        self,
+        learning_rate=0.001,
+        lamb_weight_decay=0.01,
+        beta1=0.9,
+        beta2=0.999,
+        epsilon=1e-6,
+        parameters=None,
+        grad_clip=None,
+        exclude_from_weight_decay_fn=None,
+        clip_after_allreduce=True,
+        is_grad_scaled_by_nranks=True,
+        alignment=128,
+        use_master_param_norm=True,
+        gradient_accumulation_steps=1,
+        use_master_acc_grad=True,
+        nproc_per_node=None,
+        use_hierarchical_allreduce=False,
+        name=None,
+    ):
+        assert not paddle.in_dynamic_mode(), (
+            "DistributedFusedLamb does not support dygraph mode"
+        )
+        super().__init__(learning_rate=learning_rate, grad_clip=None, name=name)
+
+        self._beta1 = beta1
+        self._beta2 = beta2
+        self._epsilon = epsilon
+        self._weight_decay = (
+            lamb_weight_decay if lamb_weight_decay is not None else 0.0
+        )
+        if grad_clip is not None:
+            assert isinstance(grad_clip, ClipGradByGlobalNorm), (
+                "Only ClipGradByGlobalNorm is supported in DistributedFusedLamb"
+            )
+            max_global_grad_norm = grad_clip.clip_norm
+        else:
+            max_global_grad_norm = -1.0
+        self._max_global_grad_norm = max_global_grad_norm
+        self._alignment = alignment if alignment is not None else -1
+        self._clip_after_allreduce = clip_after_allreduce
+        self._is_grad_scaled_by_nranks = is_grad_scaled_by_nranks
+        self._exclude_from_weight_decay_fn = exclude_from_weight_decay_fn
+        self._scale = None
+        self._use_master_param_norm = use_master_param_norm
+        self._gradient_accumulation_steps = gradient_accumulation_steps
+        self._use_master_acc_grad = use_master_acc_grad
+        self._nproc_per_node = nproc_per_node
+        self._use_hierarchical_allreduce = use_hierarchical_allreduce
+        assert self._gradient_accumulation_steps >= 1
+
+        self.helper = LayerHelper('distributed_fused_lamb')
+        self._supports_check_nan_inf = True  # very import flag for AMP
+
+        main_block = self.helper.main_program.global_block()
+        self._found_inf = main_block.create_var(
+            name=unique_name.generate('found_inf'),
+            shape=[1],
+            dtype=core.VarDesc.VarType.BOOL,
+        )
+        self._step = None
+
+        if self._gradient_accumulation_steps > 1:
+            self._stop_update = main_block.create_var(
+                name=unique_name.generate('stop_update'),
+                shape=[1],
+                dtype=core.VarDesc.VarType.BOOL,
+            )
+        else:
+            self._stop_update = None
+
+        self._param_to_master_param = {}
+
+    def _get_stop_update_var(self):
+        return self._stop_update if self._stop_update is not None else False
+
+    def _set_step(self, step):
+        self._step = step
+
+    def _get_or_create_step(self):
+        if self._step is None:
+            self._step = self._create_persistable_var('step', dtype='int64')
+        return self._step
+
+    def _set_scale(self, scale):
+        assert scale is not None
+        if not isinstance(scale, Variable):
+            scale = self._create_scale_from_constant(scale)
+        self._scale = scale
+
+    def _create_scale_from_constant(self, value):
+        name = unique_name.generate('global_scale')
+        return paddle.static.create_global_var(
+            name=name,
+            shape=[1],
+            dtype='float32',
+            value=float(value),
+            persistable=True,
+        )
+
+    def _get_or_create_scale(self):
+        if self._scale is None:
+            self._scale = self._create_scale_from_constant(1.0)
+        return self._scale
+
+    def _create_persistable_var(self, name=None, shape=[-1], dtype='float32'):
+        startup_block = self.helper.startup_program.global_block()
+        if name is not None:
+            name = unique_name.generate(name)
+        startup_var = startup_block.create_var(
+            name=name,
+            shape=shape,
+            dtype=dtype,
+            persistable=True,
+            stop_gradient=True,
+        )
+        main_block = self.helper.main_program.global_block()
+        main_var = main_block.create_var(
+            name=startup_var.name,
+            shape=startup_var.shape,
+            dtype=startup_var.dtype,
+            persistable=True,
+            stop_gradient=True,
+        )
+        return main_var
+
+    def _get_parameter(self, name, scope=None):
+        if scope is None:
+            scope = global_scope()
+
+        master_param = self._param_to_master_param.get(name)
+        assert master_param is not None
+
+        master_param_t = scope.find_var(master_param).get_tensor()
+        assert master_param_t._dtype() == paddle.float32
+
+        param_t = scope.find_var(name).get_tensor()
+        if param_t._dtype() == paddle.float32:
+            assert param_t._ptr() == master_param_t._ptr()
+            return param_t, None
+        else:
+            assert param_t._dtype() == paddle.float16
+            assert param_t.shape() == master_param_t.shape()
+            return param_t, master_param_t
+
+    def apply_optimize(self, params_grads):
+        self.apply_gradients(params_grads)
+
+    def apply_gradients(self, params_grads):
+        flattened = []
+        for p, g in params_grads:
+            flattened.extend([p, g])
+        with (
+            flattened[0].block.program._optimized_guard(flattened),
+            name_scope("optimizer"),
+        ):
+            self._apply_gradients_impl(params_grads)
+
+    def _apply_gradients_impl(self, params_grads):
+        for p, g in params_grads:
+            assert g.type == core.VarDesc.VarType.DENSE_TENSOR, (
+                "Only support dense gradient"
+            )
+            g.persistable = True  # the gradient must be persistable for fusion
+
+        fp32_fused_param = self._create_persistable_var('fp32_fused_param')
+        fp32_fused_grad = self._create_persistable_var('fp32_fused_grad')
+        fp16_fused_param = self._create_persistable_var(
+            'fp16_fused_param', dtype='float16'
+        )
+        fp16_fused_grad = self._create_persistable_var(
+            'fp16_fused_grad', dtype='float16'
+        )
+
+        master_params = []
+        for p, g in params_grads:
+            master_p = self._create_persistable_var('master_weight')
+            self._param_to_master_param[p.name] = master_p.name
+            master_params.append(master_p)
+
+        moment1 = self._create_persistable_var('moment1')
+        moment1.is_distributed = True
+        moment2 = self._create_persistable_var('moment2')
+        moment2.is_distributed = True
+        beta1pow = self._create_persistable_var('beta1pow')
+        beta2pow = self._create_persistable_var('beta2pow')
+
+        param_info = self._create_persistable_var('param_info', dtype='int32')
+        param_info.is_distributed = True
+
+        fused_offsets = self._create_persistable_var(
+            'fused_offsets', dtype='int32'
+        )
+
+        fp32_partial_fused_offsets = self._create_persistable_var(
+            'fp32_partial_fused_offsets', dtype='int32'
+        )
+        fp32_partial_fused_offsets.is_distributed = True
+
+        fp16_partial_fused_offsets = self._create_persistable_var(
+            'fp16_partial_fused_offsets', dtype='int32'
+        )
+        fp16_partial_fused_offsets.is_distributed = True
+
+        param_order = self._create_persistable_var('param_order', dtype='int32')
+        param_order.is_distributed = True
+
+        if self._gradient_accumulation_steps > 1:
+            fp32_acc_fused_grad = [
+                self._create_persistable_var('fp32_acc_fused_grad')
+            ]
+            fp16_acc_fused_grad = [
+                self._create_persistable_var(
+                    'fp16_acc_fused_grad', dtype='float16'
+                )
+            ]
+            acc_step = [self._create_persistable_var('acc_step', dtype='int64')]
+        else:
+            fp32_acc_fused_grad = []
+            fp16_acc_fused_grad = []
+            acc_step = []
+
+        step = self._get_or_create_step()
+
+        rank = paddle.distributed.get_rank()
+        nranks = paddle.distributed.get_world_size()
+        if self._nproc_per_node is None:
+            nproc_per_node = nranks
+        else:
+            nproc_per_node = self._nproc_per_node
+        assert nranks % nproc_per_node == 0, (
+            "nranks should be exactly divided by nproc_per_node"
+        )
+
+        shard_inside_node = nranks > nproc_per_node
+        local_rank = rank % nproc_per_node
+        node_id = int(rank / nproc_per_node)
+        node_num = int(nranks / nproc_per_node)
+        ring_ids = []
+        startup_block = self.helper.startup_program.global_block()
+        if nranks > 1:
+            ring_id = init_communicator(
+                startup_block, rank, list(range(nranks)), 0
+            )
+            ring_ids.append(ring_id)
+
+        use_hierarchical_allreduce = False
+        if node_num > 1 and len(ring_ids) <= 1 and shard_inside_node:
+            local_group_ranks = list(
+                range(node_id * nproc_per_node, (node_id + 1) * nproc_per_node)
+            )
+            ring_id = init_communicator(
+                startup_block, rank, local_group_ranks, 1
+            )
+            ring_ids.append(ring_id)
+
+            if self._use_hierarchical_allreduce and nranks > nproc_per_node:
+                use_hierarchical_allreduce = True
+                outer_group_ranks = list(
+                    range(rank % nproc_per_node, nranks, nproc_per_node)
+                )
+                ring_id = init_communicator(
+                    startup_block, rank, outer_group_ranks, ring_ids[-1] + 1
+                )
+                ring_ids.append(ring_id)
+
+        scale = self._get_or_create_scale()
+
+        params = [p for p, _ in params_grads]
+        grads = [g for _, g in params_grads]
+        apply_weight_decay = [1] * len(params)
+        if self._exclude_from_weight_decay_fn is not None:
+            for i, p in enumerate(params):
+                if self._exclude_from_weight_decay_fn(p):
+                    apply_weight_decay[i] = 0
+
+        for g in grads:
+            startup_block.create_var(
+                name=g.name,
+                type=g.type,
+                dtype=g.dtype,
+                persistable=g.persistable,
+                shape=g.shape,
+            )
+
+        if nranks > 1:
+            broadcast_parameters(startup_block, params, ring_ids[0])
+
+        startup_block.append_op(
+            type='distributed_fused_lamb_init',
+            inputs={
+                'Param': params,
+                'Grad': grads,
+            },
+            outputs={
+                'FP32FusedParam': [fp32_fused_param],
+                'FP32FusedGrad': [fp32_fused_grad],
+                'FP16FusedParam': [fp16_fused_param],
+                'FP16FusedGrad': [fp16_fused_grad],
+                'Moment1': [moment1],
+                'Moment2': [moment2],
+                'Beta1Pow': [beta1pow],
+                'Beta2Pow': [beta2pow],
+                'GlobalScale': [scale],
+                'ParamInfo': [param_info],
+                'ParamOut': params,
+                'MasterParamOut': master_params,
+                'GradOut': grads,
+                'FP32ShardFusedParamOffsets': [fp32_partial_fused_offsets],
+                'FP16ShardFusedParamOffsets': [fp16_partial_fused_offsets],
+                'FusedParamOffsets': [fused_offsets],
+                'ParamOrder': [param_order],
+                'Step': [step],
+            },
+            attrs={
+                'alignment': self._alignment,
+                'rank': local_rank if shard_inside_node else rank,
+                'nranks': nproc_per_node if shard_inside_node else nranks,
+                'apply_weight_decay': apply_weight_decay,
+                'moment1': 0.0,
+                'moment2': 0.0,
+                'beta1': self._beta1,
+                'beta2': self._beta2,
+            },
+        )
+
+        main_block = self.helper.main_program.global_block()
+        self._create_global_learning_rate()
+        lr = None
+        for p_g in params_grads:
+            if lr is None:
+                lr = self._create_param_lr(p_g)
+            else:
+                new_lr = self._create_param_lr(p_g)
+                assert id(lr) == id(new_lr), (
+                    "The learning rate for each parameter should be the same"
+                )
+        assert lr is not None
+
+        lamb_op = main_block.append_op(
+            type='distributed_fused_lamb',
+            inputs={
+                'FP32FusedParam': [fp32_fused_param],
+                'FP32FusedGrad': [fp32_fused_grad],
+                'FP16FusedParam': [fp16_fused_param],
+                'FP16FusedGrad': [fp16_fused_grad],
+                'LearningRate': [lr],
+                'Moment1': [moment1],
+                'Moment2': [moment2],
+                'Beta1Pow': [beta1pow],
+                'Beta2Pow': [beta2pow],
+                'GlobalScale': [scale],
+                'ParamInfo': [param_info],
+                'Param': params,
+                'Grad': grads,
+                'FusedParamOffsets': [fused_offsets],
+                'FP32ShardFusedParamOffsets': [fp32_partial_fused_offsets],
+                'FP16ShardFusedParamOffsets': [fp16_partial_fused_offsets],
+                'ParamOrder': [param_order],
+            },
+            outputs={
+                'FP32FusedParamOut': [fp32_fused_param],
+                'FP16FusedParamOut': [fp16_fused_param],
+                'Moment1Out': [moment1],
+                'Moment2Out': [moment2],
+                'Beta1PowOut': [beta1pow],
+                'Beta2PowOut': [beta2pow],
+                'ParamOut': params,
+                'GradOut': grads,
+                'FoundInf': [self._found_inf],
+                'FP32AccFusedGrad': fp32_acc_fused_grad,
+                'FP16AccFusedGrad': fp16_acc_fused_grad,
+                'AccStep': acc_step,
+                'StopUpdate': (
+                    self._stop_update if self._stop_update is not None else []
+                ),
+                'Step': [step],
+            },
+            attrs={
+                'weight_decay': self._weight_decay,
+                'beta1': self._beta1,
+                'beta2': self._beta2,
+                'epsilon': self._epsilon,
+                'max_global_grad_norm': self._max_global_grad_norm,
+                'clip_after_allreduce': self._clip_after_allreduce,
+                'rank': rank,
+                'nranks': nranks,
+                'ring_ids': ring_ids,
+                'use_master_param_norm': self._use_master_param_norm,
+                'is_grad_scaled_by_nranks': self._is_grad_scaled_by_nranks,
+                'acc_steps': self._gradient_accumulation_steps,
+                'use_master_acc_grad': self._use_master_acc_grad,
+                'use_hierarchical_allreduce': use_hierarchical_allreduce,
+            },
+        )
+        return [lamb_op]

python/paddle/incubate/optimizer/functional/__init__.py

@@ -0,0 +1,18 @@
+# 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 .bfgs import minimize_bfgs
+from .lbfgs import minimize_lbfgs
+
+__all__ = ['minimize_bfgs', 'minimize_lbfgs']

python/paddle/incubate/optimizer/functional/bfgs.py

@@ -0,0 +1,231 @@
+# 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, Literal
+
+import numpy as np
+
+import paddle
+
+from .line_search import strong_wolfe
+from .utils import (
+    _value_and_gradient,
+    check_initial_inverse_hessian_estimate,
+    check_input_type,
+)
+
+if TYPE_CHECKING:
+    from collections.abc import Callable
+
+    from paddle import Tensor
+
+
+def minimize_bfgs(
+    objective_func: Callable[[Tensor], Tensor],
+    initial_position: Tensor,
+    max_iters: int = 50,
+    tolerance_grad: float = 1e-7,
+    tolerance_change: float = 1e-9,
+    initial_inverse_hessian_estimate: Tensor | None = None,
+    line_search_fn: Literal['strong_wolfe'] = 'strong_wolfe',
+    max_line_search_iters: int = 50,
+    initial_step_length: float = 1.0,
+    dtype: Literal['float32', 'float64'] = 'float32',
+    name: str | None = None,
+) -> tuple[bool, int, Tensor, Tensor, Tensor, Tensor]:
+    r"""
+    Minimizes a differentiable function `func` using the BFGS method.
+    The 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. pp140: Algorithm 6.1 (BFGS Method).
+
+    Args:
+        objective_func: the objective function to minimize. ``objective_func`` accepts a 1D Tensor and returns a scalar.
+        initial_position (Tensor): the starting point of the iterates, has the same shape with the input of ``objective_func`` .
+        max_iters (int, optional): the maximum number of minimization iterations. Default value: 50.
+        tolerance_grad (float, optional): terminates if the gradient norm is smaller than this. Currently gradient norm uses inf norm. Default value: 1e-7.
+        tolerance_change (float, optional): terminates if the change of function value/position/parameter between two iterations is smaller than this value. Default value: 1e-9.
+        initial_inverse_hessian_estimate (Tensor, optional): the initial inverse hessian approximation at initial_position. It must be symmetric and positive definite. If not given, will use an identity matrix of order N, which is size of ``initial_position`` . Default value: None.
+        line_search_fn (str, optional): indicate which line search method to use, only support 'strong wolfe' right now. May support 'Hager Zhang' in the future. Default value: 'strong wolfe'.
+        max_line_search_iters (int, optional): the maximum number of line search iterations. Default value: 50.
+        initial_step_length (float, optional): step length used in first iteration of line search. different initial_step_length may cause different optimal result. For methods like Newton and quasi-Newton the initial trial step length should always be 1.0. Default value: 1.0.
+        dtype ('float32' | 'float64', optional): data type used in the algorithm, the data type of the input parameter must be consistent with the dtype. Default value: 'float32'.
+        name (str, optional): Name for the operation. For more information, please refer to :ref:`api_guide_Name`. Default value: None.
+
+    Returns:
+        output(tuple):
+
+            - is_converge (bool): Indicates whether found the minimum within tolerance.
+            - num_func_calls (int): number of objective function called.
+            - position (Tensor): the position of the last iteration. If the search converged, this value is the argmin of the objective function regarding to the initial position.
+            - objective_value (Tensor): objective function value at the `position`.
+            - objective_gradient (Tensor): objective function gradient at the `position`.
+            - inverse_hessian_estimate (Tensor): the estimate of inverse hessian at the `position`.
+
+    Examples:
+        .. code-block:: pycon
+            :name: code-example1
+
+            >>> # Example1: 1D Grid Parameters
+            >>> import paddle
+            >>> # Randomly simulate a batch of input data
+            >>> inputs = paddle.normal(shape=(100, 1))
+            >>> labels = inputs * 2.0
+            >>> # define the loss function
+            >>> def loss(w):
+            ...     y = w * inputs
+            ...     return paddle.nn.functional.square_error_cost(y, labels).mean()
+            >>> # Initialize weight parameters
+            >>> w = paddle.normal(shape=(1,))
+            >>> # Call the bfgs method to solve the weight that makes the loss the smallest, and update the parameters
+            >>> for epoch in range(0, 10):
+            ...     # Call the bfgs method to optimize the loss, note that the third parameter returned represents the weight
+            ...     w_update = paddle.incubate.optimizer.functional.minimize_bfgs(loss, w)[2]
+            ...     # Use paddle.assign to update parameters in place
+            ...     paddle.assign(w_update, w)
+
+        .. code-block:: pycon
+            :name: code-example2
+
+            >>> # Example2: Multidimensional Grid Parameters
+            >>> import paddle
+            >>> def flatten(x):
+            ...     return x.flatten()
+            >>> def unflatten(x):
+            ...     return x.reshape((2, 2))
+            >>> # Assume the network parameters are more than one dimension
+            >>> def net(x):
+            ...     assert len(x.shape) > 1
+            ...     return x.square().mean()
+            >>> # function to be optimized
+            >>> def bfgs_f(flatten_x):
+            ...     return net(unflatten(flatten_x))
+            >>> x = paddle.rand([2, 2])
+            >>> for i in range(0, 10):
+            ...     # Flatten x before using minimize_bfgs
+            ...     x_update = paddle.incubate.optimizer.functional.minimize_bfgs(bfgs_f, flatten(x))[2]
+            ...     # unflatten x_update, then update parameters
+            ...     paddle.assign(unflatten(x_update), x)
+    """
+
+    if dtype not in ['float32', 'float64']:
+        raise ValueError(
+            f"The dtype must be 'float32' or 'float64', but the specified is {dtype}."
+        )
+
+    op_name = 'minimize_bfgs'
+    check_input_type(initial_position, 'initial_position', op_name)
+
+    I = paddle.eye(initial_position.shape[0], dtype=dtype)
+    if initial_inverse_hessian_estimate is None:
+        initial_inverse_hessian_estimate = I
+    else:
+        check_input_type(
+            initial_inverse_hessian_estimate,
+            'initial_inverse_hessian_estimate',
+            op_name,
+        )
+        check_initial_inverse_hessian_estimate(initial_inverse_hessian_estimate)
+
+    Hk = paddle.assign(initial_inverse_hessian_estimate)
+    # use detach and assign to create new tensor rather than =, or xk will share memory and grad with initial_position
+    xk = paddle.assign(initial_position.detach())
+
+    value, g1 = _value_and_gradient(objective_func, xk)
+    num_func_calls = paddle.full(shape=[1], fill_value=1, dtype='int64')
+
+    # when the dim of x is 1000, it needs more than 30 iters to get all element converge to minimum.
+    k = paddle.full(shape=[1], fill_value=0, dtype='int64')
+    done = paddle.full(shape=[1], fill_value=False, dtype='bool')
+    is_converge = paddle.full(shape=[1], fill_value=False, dtype='bool')
+
+    def cond(k, done, is_converge, num_func_calls, xk, value, g1, Hk):
+        return (k < max_iters) & ~done
+
+    def body(k, done, is_converge, num_func_calls, xk, value, g1, Hk):
+        # --------------   compute pk   -------------- #
+        pk = -paddle.matmul(Hk, g1)
+
+        # --------------   compute alpha by line search   -------------- #
+        if line_search_fn == 'strong_wolfe':
+            alpha, value, g2, ls_func_calls = strong_wolfe(
+                f=objective_func,
+                xk=xk,
+                pk=pk,
+                max_iters=max_line_search_iters,
+                initial_step_length=initial_step_length,
+                dtype=dtype,
+            )
+        else:
+            raise NotImplementedError(
+                f"Currently only support line_search_fn = 'strong_wolfe', but the specified is '{line_search_fn}'"
+            )
+        num_func_calls += ls_func_calls
+
+        # --------------   update Hk   -------------- #
+        sk = alpha * pk
+        yk = g2 - g1
+
+        xk = xk + sk
+        g1 = g2
+
+        sk = paddle.unsqueeze(sk, 0)
+        yk = paddle.unsqueeze(yk, 0)
+
+        rhok_inv = paddle.dot(yk, sk)
+        rhok = paddle.static.nn.cond(
+            rhok_inv == 0.0,
+            lambda: paddle.full(shape=[1], fill_value=1000.0, dtype=dtype),
+            lambda: 1.0 / rhok_inv,
+        )
+
+        Vk_transpose = I - rhok * sk * yk.t()
+        Vk = I - rhok * yk * sk.t()
+        Hk = (
+            paddle.matmul(paddle.matmul(Vk_transpose, Hk), Vk)
+            + rhok * sk * sk.t()
+        )
+
+        k += 1
+
+        # --------------   check convergence   -------------- #
+        gnorm = paddle.linalg.norm(g1, p=np.inf)
+        pk_norm = paddle.linalg.norm(pk, p=np.inf)
+        paddle.assign(
+            done | (gnorm < tolerance_grad) | (pk_norm < tolerance_change), done
+        )
+        paddle.assign(done, is_converge)
+        # when alpha=0, there is no chance to get xk change.
+        paddle.assign(done | (alpha == 0.0), done)
+        return [k, done, is_converge, num_func_calls, xk, value, g1, Hk]
+
+    paddle.static.nn.while_loop(
+        cond=cond,
+        body=body,
+        loop_vars=[k, done, is_converge, num_func_calls, xk, value, g1, Hk],
+    )
+    return is_converge, num_func_calls, xk, value, g1, Hk

python/paddle/incubate/optimizer/functional/lbfgs.py

@@ -0,0 +1,341 @@
+# 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, Literal
+
+import numpy as np
+
+import paddle
+
+from .line_search import strong_wolfe
+from .utils import (
+    _value_and_gradient,
+    check_initial_inverse_hessian_estimate,
+    check_input_type,
+)
+
+if TYPE_CHECKING:
+    from collections.abc import Callable
+
+    from paddle import Tensor
+
+
+def minimize_lbfgs(
+    objective_func: Callable[[Tensor], Tensor],
+    initial_position: Tensor,
+    history_size: int = 100,
+    max_iters: int = 50,
+    tolerance_grad: float = 1e-8,
+    tolerance_change: float = 1e-8,
+    initial_inverse_hessian_estimate: Tensor | None = None,
+    line_search_fn: Literal['strong_wolfe'] = 'strong_wolfe',
+    max_line_search_iters: int = 50,
+    initial_step_length: int = 1.0,
+    dtype: Literal['float32', 'float64'] = 'float32',
+    name: str | None = None,
+) -> tuple[bool, int, Tensor, Tensor, Tensor]:
+    r"""
+    Minimizes a differentiable function `func` using the L-BFGS method.
+    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:
+        objective_func: the objective function to minimize. ``objective_func`` accepts a 1D Tensor and returns a scalar.
+        initial_position (Tensor): the starting point of the iterates, has the same shape with the input of ``objective_func`` .
+        history_size (Scalar): the number of stored vector pairs {si,yi}. Default value: 100.
+        max_iters (int, optional): the maximum number of minimization iterations. Default value: 50.
+        tolerance_grad (float, optional): terminates if the gradient norm is smaller than this. Currently gradient norm uses inf norm. Default value: 1e-7.
+        tolerance_change (float, optional): terminates if the change of function value/position/parameter between two iterations is smaller than this value. Default value: 1e-9.
+        initial_inverse_hessian_estimate (Tensor, optional): the initial inverse hessian approximation at initial_position. It must be symmetric and positive definite. If not given, will use an identity matrix of order N, which is size of ``initial_position`` . Default value: None.
+        line_search_fn (str, optional): indicate which line search method to use, only support 'strong wolfe' right now. May support 'Hager Zhang' in the future. Default value: 'strong wolfe'.
+        max_line_search_iters (int, optional): the maximum number of line search iterations. Default value: 50.
+        initial_step_length (float, optional): step length used in first iteration of line search. different initial_step_length may cause different optimal result. For methods like Newton and quasi-Newton the initial trial step length should always be 1.0. Default value: 1.0.
+        dtype ('float32' | 'float64', optional): data type used in the algorithm, the data type of the input parameter must be consistent with the dtype. Default value: 'float32'.
+        name (str, optional): Name for the operation. For more information, please refer to :ref:`api_guide_Name`. Default value: None.
+
+    Returns:
+        output(tuple):
+
+            - is_converge (bool): Indicates whether found the minimum within tolerance.
+            - num_func_calls (int): number of objective function called.
+            - position (Tensor): the position of the last iteration. If the search converged, this value is the argmin of the objective function regarding to the initial position.
+            - objective_value (Tensor): objective function value at the `position`.
+            - objective_gradient (Tensor): objective function gradient at the `position`.
+
+    Examples:
+        .. code-block:: pycon
+            :name: code-example1
+
+            >>> # Example1: 1D Grid Parameters
+            >>> import paddle
+            >>> # Randomly simulate a batch of input data
+            >>> inputs = paddle.normal(shape=(100, 1))
+            >>> labels = inputs * 2.0
+            >>> # define the loss function
+            >>> def loss(w):
+            ...     y = w * inputs
+            ...     return paddle.nn.functional.square_error_cost(y, labels).mean()
+            >>> # Initialize weight parameters
+            >>> w = paddle.normal(shape=(1,))
+            >>> # Call the bfgs method to solve the weight that makes the loss the smallest, and update the parameters
+            >>> for epoch in range(0, 10):
+            ...     # Call the bfgs method to optimize the loss, note that the third parameter returned represents the weight
+            ...     w_update = paddle.incubate.optimizer.functional.minimize_bfgs(loss, w)[2]
+            ...     # Use paddle.assign to update parameters in place
+            ...     paddle.assign(w_update, w)
+
+        .. code-block:: pycon
+            :name: code-example2
+
+            >>> # Example2: Multidimensional Grid Parameters
+            >>> import paddle
+            >>> def flatten(x):
+            ...     return x.flatten()
+            >>> def unflatten(x):
+            ...     return x.reshape((2, 2))
+            >>> # Assume the network parameters are more than one dimension
+            >>> def net(x):
+            ...     assert len(x.shape) > 1
+            ...     return x.square().mean()
+            >>> # function to be optimized
+            >>> def bfgs_f(flatten_x):
+            ...     return net(unflatten(flatten_x))
+            >>> x = paddle.rand([2, 2])
+            >>> for i in range(0, 10):
+            ...     # Flatten x before using minimize_bfgs
+            ...     x_update = paddle.incubate.optimizer.functional.minimize_bfgs(bfgs_f, flatten(x))[2]
+            ...     # unflatten x_update, then update parameters
+            ...     paddle.assign(unflatten(x_update), x)
+
+    """
+    if dtype not in ['float32', 'float64']:
+        raise ValueError(
+            f"The dtype must be 'float32' or 'float64', but the specified is {dtype}."
+        )
+
+    op_name = 'minimize_lbfgs'
+    check_input_type(initial_position, 'initial_position', op_name)
+
+    if initial_inverse_hessian_estimate is None:
+        H0 = paddle.eye(initial_position.shape[0], dtype=dtype)
+    else:
+        check_input_type(
+            initial_inverse_hessian_estimate,
+            'initial_inverse_hessian_estimate',
+            op_name,
+        )
+        check_initial_inverse_hessian_estimate(initial_inverse_hessian_estimate)
+        H0 = initial_inverse_hessian_estimate
+
+    # use detach and assign to create new tensor rather than =, or xk will share memory and grad with initial_position
+    xk = paddle.assign(initial_position.detach())
+    value, g1 = _value_and_gradient(objective_func, xk)
+
+    k = paddle.full(shape=[1], fill_value=0, dtype='int64')
+    done = paddle.full(shape=[1], fill_value=False, dtype='bool')
+    is_converge = paddle.full(shape=[1], fill_value=False, dtype='bool')
+    num_func_calls = paddle.full(shape=[1], fill_value=1, dtype='int64')
+
+    history_size = paddle.full(shape=[], fill_value=history_size, dtype='int64')
+    head = paddle.full(shape=[1], fill_value=1, dtype='int64')
+    tail = paddle.full(shape=[1], fill_value=0, dtype='int64')
+
+    shape = initial_position.shape[0]
+    # Use tensor as array of fixed length, rather than flexible tensor array. Because in static graph mode,
+    # tensor array will produce tensor of shape[-1], which will cause error when calling jacobian. In this way, can not use append
+    # or pop, so we need head and tail to record where is the newest data and where is the oldest.
+    # Totally speaking, realized a stack by array.
+    sk_vec = paddle.zeros((history_size + 1, shape), dtype=dtype)
+    yk_vec = paddle.zeros((history_size + 1, shape), dtype=dtype)
+    rhok_vec = paddle.zeros((history_size + 1, 1), dtype=dtype)
+    ai_vec = paddle.zeros((history_size + 1, 1), dtype=dtype)
+
+    def cond(
+        k,
+        done,
+        is_converge,
+        num_func_calls,
+        value,
+        xk,
+        g1,
+        sk_vec,
+        yk_vec,
+        rhok_vec,
+        head,
+        tail,
+    ):
+        return (k < max_iters) & ~done
+
+    def body(
+        k,
+        done,
+        is_converge,
+        num_func_calls,
+        value,
+        xk,
+        g1,
+        sk_vec,
+        yk_vec,
+        rhok_vec,
+        head,
+        tail,
+    ):
+        # use assign to cut off the relevance between g1 and q, or they will change together.
+
+        # --------------   compute p_k by two-loop recursion    -------------- #
+        q = paddle.assign(g1)
+        # In a array circle, the index may out of range, so must use mod.
+        i = paddle.full(
+            shape=[], fill_value=(head - 1).mod(history_size), dtype='int64'
+        )
+
+        def cond(i, q, ai_vec):
+            return i != tail
+
+        def body(i, q, ai_vec):
+            if paddle.in_dynamic_mode():
+                ai_vec[i] = rhok_vec[i] * paddle.dot(sk_vec[i], q)
+            else:
+                ai_vec = paddle.static.setitem(
+                    ai_vec, i, rhok_vec[i] * paddle.dot(sk_vec[i], q)
+                )
+            q = q - ai_vec[i] * yk_vec[i]
+            i = (i - 1).mod(history_size)
+            return i, q, ai_vec
+
+        paddle.static.nn.while_loop(
+            cond=cond, body=body, loop_vars=[i, q, ai_vec]
+        )
+
+        r = paddle.matmul(H0, q)
+
+        i = paddle.full(shape=[], fill_value=tail + 1, dtype='int64')
+
+        def cond(i, r):
+            return i != head
+
+        def body(i, r):
+            beta = rhok_vec[i] * paddle.dot(yk_vec[i], r)
+            r = r + sk_vec[i] * (ai_vec[i] - beta)
+            i = (i + 1).mod(history_size)
+            return i, r
+
+        paddle.static.nn.while_loop(cond=cond, body=body, loop_vars=[i, r])
+
+        pk = -r
+
+        # --------------   compute alpha by line search    -------------- #
+        if line_search_fn == 'strong_wolfe':
+            alpha, value, g2, ls_func_calls = strong_wolfe(
+                f=objective_func,
+                xk=xk,
+                pk=pk,
+                max_iters=max_line_search_iters,
+                initial_step_length=initial_step_length,
+                dtype=dtype,
+            )
+        else:
+            raise NotImplementedError(
+                f"Currently only support line_search_fn = 'strong_wolfe', but the specified is '{line_search_fn}'"
+            )
+        paddle.assign(num_func_calls + ls_func_calls, num_func_calls)
+
+        # --------------   update sk_vec, yk_vec, rhok_vec    -------------- #
+        sk = alpha * pk
+        yk = g2 - g1
+
+        rhok_inv = paddle.dot(yk, sk)
+        rhok = paddle.static.nn.cond(
+            rhok_inv == 0.0,
+            lambda: paddle.full(shape=[1], fill_value=1000.0, dtype=dtype),
+            lambda: 1.0 / rhok_inv,
+        )
+        if paddle.in_dynamic_mode():
+            sk_vec[head] = sk
+            yk_vec[head] = yk
+            rhok_vec[head] = rhok
+        else:
+            sk_vec = paddle.static.setitem(sk_vec, head, sk)
+            yk_vec = paddle.static.setitem(yk_vec, head, yk)
+            rhok_vec = paddle.static.setitem(rhok_vec, head, rhok)
+        head = (head + 1) % history_size
+
+        def true_fn(tail):
+            paddle.assign(tail + 1, tail)
+
+        # when array is full, the tail should move forward too.
+        paddle.static.nn.cond(head == tail, lambda: true_fn(tail), None)
+
+        xk = xk + sk
+        g1 = g2
+        k += 1
+
+        # --------------   check convergence    -------------- #
+        gnorm = paddle.linalg.norm(g1, p=np.inf)
+        pk_norm = paddle.linalg.norm(pk, p=np.inf)
+        paddle.assign(
+            done | (gnorm < tolerance_grad) | (pk_norm < tolerance_change), done
+        )
+        paddle.assign(done, is_converge)
+        # when alpha=0, there is no chance to get xk change.
+        paddle.assign(done | (alpha == 0.0), done)
+
+        return [
+            k,
+            done,
+            is_converge,
+            num_func_calls,
+            value,
+            xk,
+            g1,
+            sk_vec,
+            yk_vec,
+            rhok_vec,
+            head,
+            tail,
+        ]
+
+    paddle.static.nn.while_loop(
+        cond=cond,
+        body=body,
+        loop_vars=[
+            k,
+            done,
+            is_converge,
+            num_func_calls,
+            value,
+            xk,
+            g1,
+            sk_vec,
+            yk_vec,
+            rhok_vec,
+            head,
+            tail,
+        ],
+    )
+    return is_converge, num_func_calls, xk, value, g1

python/paddle/incubate/optimizer/functional/line_search.py

@@ -0,0 +1,368 @@
+# 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.
+
+import paddle
+
+from .utils import _value_and_gradient
+
+
+def cubic_interpolation_(x1, f1, g1, x2, f2, g2):
+    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.
+    Returns:
+        min_pos: the minimum point between the specified points in the cubic curve.
+    """
+    xmin, xmax = paddle.static.nn.cond(
+        x1 <= x2, lambda: (x1, x2), lambda: (x2, x1)
+    )
+    d1 = g1 + g2 - 3 * (f1 - f2) / (x1 - x2)
+    d2_square = d1**2 - g1 * g2
+
+    def true_func1():
+        d2 = d2_square.sqrt()
+
+        def true_fn2():
+            return x2 - (x2 - x1) * ((g2 + d2 - d1) / (g2 - g1 + 2 * d2))
+
+        def false_fn2():
+            return x1 - (x1 - x2) * ((g1 + d2 - d1) / (g1 - g2 + 2 * d2))
+
+        pred = paddle.less_equal(x=x1, y=x2)
+        min_pos = paddle.static.nn.cond(pred, true_fn2, false_fn2)
+
+        return paddle.minimum(paddle.maximum(min_pos, xmin), xmax)
+
+    def false_func1():
+        return (xmin + xmax) / 2.0
+
+    min_pos = paddle.static.nn.cond(d2_square >= 0.0, true_func1, false_func1)
+    return min_pos
+
+
+def strong_wolfe(
+    f,
+    xk,
+    pk,
+    max_iters=20,
+    tolerance_change=1e-8,
+    initial_step_length=1.0,
+    c1=1e-4,
+    c2=0.9,
+    alpha_max=10,
+    dtype='float32',
+):
+    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:
+        f: the objective function to minimize. ``f`` accepts a multivariate input and returns a scalar.
+        xk (Tensor): the starting point of the iterates.
+        pk (Tensor): search direction.
+        max_iters (Scalar): the maximum number of iterations.
+        tolerance_grad (Scalar): terminates if the gradient norm is smaller than
+            this. Currently gradient norm uses inf norm.
+        tolerance_change (Scalar): terminates if the change of function value/position/parameter between
+            two iterations is smaller than this value.
+        initial_step_length (Scalar): step length used in first iteration.
+        c1 (Scalar): parameter for sufficient decrease condition.
+        c2 (Scalar): parameter for curvature condition.
+        alpha_max (float): max step length.
+        dtype ('float32' | 'float64'): the datatype to be used.
+
+    Returns:
+        num_func_calls (float): number of objective function called in line search process.
+        a_star(Tensor): optimal step length, or 0. if the line search algorithm did not converge.
+        phi_star (Tensor): phi at a_star.
+        derphi_star (Tensor): derivative of phi at a_star.
+
+    Following summarizes the essentials of the strong Wolfe line search algorithm.
+    Some notations used in the description:
+
+        - `f` denotes the objective function.
+        - `phi` is a function of step size alpha, restricting `f` on a line.
+
+            phi = f(xk + a * pk),
+            where xk is the position of k'th iterate, pk is the line search direction(decent direction),
+            and a is the step size.
+        - a : substitute of alpha
+        - a1 is a of last iteration, which is alpha_(i-1).
+        - a2 is a of current iteration, which is alpha_i.
+        - a_lo is a in left position when calls zoom, which is alpha_low.
+        - a_hi is a in right position when calls zoom, which is alpha_high.
+
+    Line Search Algorithm:
+        repeat
+            Compute phi(a2) and derphi(a2).
+            1. If phi(a2) > phi(0) + c_1 * a2 * phi'(0) or [phi(a2) >= phi(a1) and i > 1],
+                a_star= zoom(a1, a2) and stop;
+
+            2. If |phi'(a2)| <= -c_2 * phi'(0),
+                a_star= a2 and stop;
+
+            3. If phi'(a2) >= 0,
+                a_star= 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 phi(aj) and derphi(aj).
+            1. If phi(aj) > phi(0) + c_1 * aj * phi'(0) or phi(aj) >= phi(a_lo),
+                then a_hi <- aj;
+            2.
+                2.1. If |phi'(aj)| <= -c_2 * phi'(0), then a_star= a2 and stop;
+
+                2.2. If phi'(aj) * (a2 - a1) >= 0, then a_hi = a_lo
+
+                a_lo = aj;
+        end(repeat)
+    """
+
+    def phi_and_derphi(a):
+        r"""Compute function value and derivative of phi at a.
+        phi = f(xk + a * pk)
+        phi'(a) = f'(xk + a * pk) * pk
+        """
+        phi_value, f_grad = _value_and_gradient(f, xk + a * pk)
+        phi_grad = paddle.dot(f_grad, pk)
+        # return f_grad to be used in bfgs/l-bfgs to compute yk to avoid computint repeatedly.
+        return phi_value, f_grad, phi_grad
+
+    def zoom(
+        a_lo,
+        phi_lo,
+        derphi_lo,
+        derf_lo,
+        a_hi,
+        phi_hi,
+        derphi_hi,
+        phi_0,
+        derphi_0,
+    ):
+        # find the exact a from the bracket [a_lo, a_hi]
+        max_zoom_iters = max_iters
+        j = paddle.full(shape=[1], fill_value=0, dtype='int64')
+        done_zoom = paddle.full(shape=[1], fill_value=False, dtype='bool')
+
+        def cond_zoom(
+            j,
+            done_zoom,
+            a_lo,
+            phi_lo,
+            derphi_lo,
+            derf_lo,
+            a_hi,
+            phi_hi,
+            derphi_hi,
+        ):
+            pred = paddle.abs(a_hi - a_lo) < tolerance_change
+            paddle.assign(done_zoom | pred, done_zoom)
+            return (j < max_zoom_iters) & ~done_zoom
+
+        def body_zoom(
+            j,
+            done_zoom,
+            a_lo,
+            phi_lo,
+            derphi_lo,
+            derf_lo,
+            a_hi,
+            phi_hi,
+            derphi_hi,
+        ):
+            aj = cubic_interpolation_(
+                a_lo, phi_lo, derphi_lo, a_hi, phi_hi, derphi_hi
+            )  # 21
+            min_change = 0.1 * paddle.abs(a_hi - a_lo)
+            pred = (
+                paddle.minimum(paddle.abs(aj - a_lo), paddle.abs(aj - a_hi))
+                < min_change
+            )
+            aj = paddle.static.nn.cond(
+                pred, lambda: 0.5 * (a_lo + a_hi), lambda: aj
+            )
+
+            phi_j, derf_j, derphi_j = phi_and_derphi(aj)
+
+            def true_fn():
+                # use assign to modify the variable in-place
+                paddle.assign(aj, a_hi)
+                paddle.assign(phi_j, phi_hi)
+                paddle.assign(derphi_j, derphi_hi)
+
+            def false_fn(a_lo, done_zoom):
+                pred3 = paddle.abs(derphi_j) <= -c2 * derphi_0
+                paddle.assign(pred3, done_zoom)
+
+                def true_fn():
+                    paddle.assign(a_lo, a_hi)
+                    paddle.assign(phi_lo, phi_hi)
+                    paddle.assign(derphi_lo, derphi_hi)
+
+                pred4 = ~done_zoom & (derphi_j * (a_hi - a_lo) >= 0)
+                paddle.static.nn.cond(pred4, true_fn, None)
+
+                paddle.assign(aj, a_lo)
+                paddle.assign(phi_j, phi_lo)
+                paddle.assign(derphi_j, derphi_lo)
+                paddle.assign(derf_j, derf_lo)
+
+            pred2 = (phi_j > phi_0 + c1 * aj * derphi_0) | (phi_j >= phi_lo)
+            paddle.static.nn.cond(
+                pred2, true_fn, lambda: false_fn(a_lo, done_zoom)
+            )
+            j = paddle.static.nn.cond(done_zoom, lambda: j, lambda: j + 1)
+            return [
+                j,
+                done_zoom,
+                a_lo,
+                phi_lo,
+                derphi_lo,
+                derf_lo,
+                a_hi,
+                phi_hi,
+                derphi_hi,
+            ]
+
+        paddle.static.nn.while_loop(
+            cond=cond_zoom,
+            body=body_zoom,
+            loop_vars=[
+                j,
+                done_zoom,
+                a_lo,
+                phi_lo,
+                derphi_lo,
+                derf_lo,
+                a_hi,
+                phi_hi,
+                derphi_hi,
+            ],
+        )
+        # j is the number of object function called in zoom.
+        return j
+
+    alpha_max = paddle.full(shape=[1], fill_value=alpha_max, dtype=dtype)
+
+    a1 = paddle.full(shape=[1], fill_value=0.0, dtype=dtype)
+    a2 = paddle.full(shape=[1], fill_value=initial_step_length, dtype=dtype)
+
+    phi_1, derf_1, derphi_1 = phi_and_derphi(a1)
+    # use assign to cut off binding between two variables
+    phi_0 = paddle.assign(phi_1)
+    derphi_0 = paddle.assign(derphi_1)
+    ls_func_calls = paddle.full(shape=[1], fill_value=1, dtype='int64')
+
+    # If not found the a_star, will return alpha=0 and f(xk), derf(xk)
+    a_star = paddle.full(shape=[1], fill_value=0, dtype=dtype)
+    phi_star = paddle.assign(phi_1)
+    derf_star = paddle.assign(derf_1)
+
+    i = paddle.full(shape=[1], fill_value=0, dtype='int64')
+    done = paddle.full(shape=[1], fill_value=False, dtype='bool')
+
+    def cond(i, ls_func_calls, a1, a2, phi_1, derf_1, done):
+        return (i < max_iters) & ~done
+
+    def body(i, ls_func_calls, a1, a2, phi_1, derf_1, done):
+        phi_2, derf_2, derphi_2 = phi_and_derphi(a2)
+        paddle.assign(ls_func_calls + 1, ls_func_calls)
+        paddle.assign(done | paddle.any(paddle.isinf(phi_2)), done)
+
+        def true_fn1():
+            j = zoom(
+                a1,
+                phi_1,
+                derphi_1,
+                derf_1,
+                a2,
+                phi_2,
+                derphi_2,
+                phi_0,
+                derphi_0,
+            )
+            paddle.assign(a1, a_star)
+            paddle.assign(phi_1, phi_star)
+            paddle.assign(derf_1, derf_star)
+            paddle.assign(ls_func_calls + j, ls_func_calls)
+
+        pred1 = ~done & (
+            (phi_2 > phi_0 + c1 * a2 * derphi_0) | ((phi_2 >= phi_1) & (i > 1))
+        )
+        paddle.assign(done | pred1, done)
+        paddle.static.nn.cond(pred1, true_fn1, None)
+
+        def true_fn2():
+            paddle.assign(a2, a_star)
+            paddle.assign(phi_2, phi_star)
+            paddle.assign(derf_2, derf_star)
+
+        pred2 = ~done & (paddle.abs(derphi_2) <= -c2 * derphi_0)
+        paddle.assign(done | pred2, done)
+        paddle.static.nn.cond(pred2, true_fn2, None)
+
+        def true_fn3():
+            j = zoom(
+                a2,
+                phi_2,
+                derphi_2,
+                derf_2,
+                a1,
+                phi_1,
+                derphi_1,
+                phi_0,
+                derphi_0,
+            )
+            paddle.assign(a2, a_star)
+            paddle.assign(phi_2, phi_star)
+            paddle.assign(derf_2, derf_star)
+            paddle.assign(ls_func_calls + j, ls_func_calls)
+
+        pred3 = ~done & (derphi_2 >= 0)
+        paddle.assign(done | pred3, done)
+        paddle.static.nn.cond(pred3, true_fn3, None)
+
+        def false_fn():
+            paddle.assign(a2, a1)
+            paddle.assign(phi_2, phi_1)
+            paddle.assign(derf_2, derf_1)
+            paddle.assign(paddle.minimum(2 * a2, alpha_max), a2)
+            paddle.assign(i + 1, i)
+
+        paddle.static.nn.cond(done, None, false_fn)
+        return [i, ls_func_calls, a1, a2, phi_1, derf_1, done]
+
+    paddle.static.nn.while_loop(
+        cond=cond,
+        body=body,
+        loop_vars=[i, ls_func_calls, a1, a2, phi_1, derf_1, done],
+    )
+
+    return a_star, phi_star, derf_star, ls_func_calls

python/paddle/incubate/optimizer/functional/utils.py

@@ -0,0 +1,120 @@
+# 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.
+
+import paddle
+from paddle.base.data_feeder import check_type
+from paddle.base.framework import Variable, in_pir_mode
+
+
+def check_input_type(input, name, op_name):
+    r"""Check whether the input is tensor or variable."""
+    if paddle.in_dynamic_mode():
+        if not isinstance(input, paddle.Tensor):
+            raise ValueError(f"The input: {input} must be tensor.")
+    else:
+        check_type(input, name, (Variable, paddle.pir.Value), op_name)
+
+
+def check_initial_inverse_hessian_estimate(H0):
+    r"""Check whether the specified initial_inverse_hessian_estimate is symmetric and positive definite.
+        Raise errors when precondition not met.
+
+    Note:
+        In static graph can not raise error directly, so use py_func make raise_func as a op,
+        and use paddle.static.nn.cond to decide if put the op in net.
+        cholesky is the fast way to check positive definition, but in static graph can not catch
+        exception to raise value error, so use eigvals rather than cholesky in static graph.
+    """
+    is_symmetric = paddle.all(paddle.equal(H0, H0.t()))
+
+    def raise_func():
+        raise ValueError(
+            "The initial_inverse_hessian_estimate should be symmetric and positive definite, but the specified is not."
+        )
+
+    if paddle.in_dynamic_mode():
+        if not is_symmetric:
+            raise_func()
+        try:
+            paddle.linalg.cholesky(H0)
+        except RuntimeError as error:
+            raise_func()
+    elif in_pir_mode():
+        paddle.static.nn.control_flow.Assert(
+            is_symmetric,
+            None,
+            10,
+            name="The initial_inverse_hessian_estimate should be symmetric and positive definite, but the specified is not.",
+        )
+        eigvals = paddle.linalg.eigvals(H0)
+        is_positive = paddle.bitwise_and(
+            paddle.all(eigvals.real() > 0.0), paddle.all(eigvals.imag() == 0.0)
+        )
+        paddle.static.nn.control_flow.Assert(
+            is_positive,
+            None,
+            10,
+            name="The initial_inverse_hessian_estimate should be symmetric and positive definite, but the specified is not.",
+        )
+
+    else:
+
+        def create_tmp_var(program, name, dtype, shape):
+            return program.current_block().create_var(
+                name=name, dtype=dtype, shape=shape
+            )
+
+        out_var = create_tmp_var(
+            paddle.static.default_main_program(),
+            name='output',
+            dtype='float32',
+            shape=[-1],
+        )
+
+        def false_fn():
+            paddle.static.nn.py_func(
+                func=raise_func, x=is_symmetric, out=out_var
+            )
+
+        paddle.static.nn.cond(is_symmetric, None, false_fn)
+        # eigvals only support cpu
+        paddle.set_device("cpu")
+        eigvals = paddle.linalg.eigvals(H0)
+        is_positive = paddle.all(eigvals.real() > 0.0) and paddle.all(
+            eigvals.imag() == 0.0
+        )
+        paddle.static.nn.cond(is_positive, None, false_fn)
+
+
+def _value_and_gradient(f, x, v=None):
+    r"""Compute function value and gradient of f at x.
+
+    Args:
+        f (Callable): the objective function.
+        x (Tensor): the input tensor.
+    Returns:
+        value: a tensor that holds the function value.
+        gradient: a tensor that holds the function gradients.
+    """
+    # use detach to cut off relation between x and original graph
+    x = x.detach()
+    x.stop_gradient = False
+    value = f(x)
+    if paddle.in_dynamic_mode():
+        # only need to compute first order derivative, and some op dont support high order derivative.
+        gradient = paddle.grad([value], [x], create_graph=False)[0]
+    else:
+        gradient = paddle.static.gradients([value], [x])[0]
+    # use detach to make results real number without grad to avoid assign error
+    return value.detach(), gradient.detach()

python/paddle/incubate/optimizer/gradient_merge.py

@@ -0,0 +1,383 @@
+# Copyright (c) 2019 PaddlePaddle Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+
+import paddle
+from paddle.base import core
+from paddle.base.framework import (
+    Variable,
+    default_main_program,
+    default_startup_program,
+    device_guard,
+    in_dygraph_mode,
+    program_guard,
+)
+
+__all__ = []
+
+
+class GradientMergeOptimizer:
+    """
+    Gradient Merge, also called as Gradient Accumulation,
+    is a training strategy for larger batches. With this strategy,
+    the parameter will not be updated until specific steps.
+
+    For each step, the forward network and the backward network
+    will run to calculate the gradient of the parameters.
+
+    For every k step, the optimization network will run,
+    applying a specific optimization method (such as SGD, Adam)
+    to the parameters.
+
+    Args:
+        inner_optimizer (Optimizer): The specific optimization (such as SGD, Adam)
+            which update the parameters
+        k_steps (int): the update period of the parameters
+        avg (bool): whether to average the gradients of each mini-batch,
+            the default value is `True`
+
+    Examples:
+        .. code-block:: pycon
+
+            >>> import paddle
+            >>> import numpy as np
+            >>> paddle.enable_static()
+
+            >>> def gen_data(batch_size):
+            ...     return {
+            ...         "x": np.random.random(size=(batch_size, 32)).astype('float32'),
+            ...         "y": np.random.random(size=(batch_size, 1)).astype('int64'),
+            ...     }
+
+            >>> def mlp(input_x, input_y, hid_dim=128, label_dim=2):
+            ...     fc_1 = paddle.static.nn.fc(x=input_x, size=hid_dim)
+            ...     prediction = paddle.static.nn.fc(x=[fc_1], size=label_dim, activation='softmax')
+            ...     cost = paddle.nn.functional.cross_entropy(
+            ...         input=prediction,
+            ...         label=input_y,
+            ...         reduction='none',
+            ...         use_softmax=False,
+            ...     )
+            ...     sum_cost = paddle.mean(cost)
+            ...     return sum_cost, fc_1, prediction
+
+            >>> input_x = paddle.static.data(name="x", shape=[-1, 32], dtype='float32')
+            >>> input_y = paddle.static.data(name="y", shape=[-1, 1], dtype='int64')
+            >>> cost, fc_1, pred = mlp(input_x, input_y)
+            >>> sgd = paddle.optimizer.Adam(learning_rate=0.01)
+            >>> sgd = paddle.incubate.optimizer.GradientMergeOptimizer(sgd, k_steps=4, avg=True)
+            >>> sgd.minimize(cost)
+
+            >>> place = paddle.CPUPlace()
+            >>> exe = paddle.static.Executor(place)
+            >>> exe.run(paddle.static.default_startup_program())
+
+            >>> for i in range(10):
+            ...     cost_val = exe.run(
+            ...         feed=gen_data(32),
+            ...         program=paddle.static.default_main_program(),
+            ...         fetch_list=[cost.name],
+            ...     )
+            ...     print("step=%d, cost=%f" % (i, cost_val[0]))
+    """
+
+    GRAD_MERGE_COND_NAME = "grad_merge_cond_name"
+
+    def __init__(self, inner_optimizer, k_steps=1, avg=True):
+        if in_dygraph_mode():
+            raise Exception(
+                "In dygraph, we don't support GradientMergeOptimizer."
+                "You can do Gradient merge by yourself with k-times forward + backward, "
+                "and one-time optimizer.minimize()"
+            )
+
+        assert inner_optimizer is not None, "inner optimizer can not be None"
+        assert isinstance(k_steps, int) and k_steps > 0, (
+            "k_steps should be a positive integer"
+        )
+
+        self.inner_optimizer = inner_optimizer
+        self.k_steps = k_steps
+        self.type = "gradient_merge"
+        self.avg = avg
+        self._optimize_ops = None
+
+    def _set_k_steps(self, k_steps):
+        self.k_steps = k_steps
+
+    def _set_avg(self, avg):
+        self.avg = avg
+
+    def backward(
+        self,
+        loss,
+        startup_program=None,
+        parameter_list=None,
+        no_grad_set=None,
+        callbacks=None,
+    ):
+        assert isinstance(loss, Variable), "The loss should be an Variable."
+        assert parameter_list is None, (
+            "The parameter_list should be None when using GradientMergeOptimizer"
+        )
+        assert no_grad_set is None, (
+            "The no_grad_set should be None when using GradientMergeOptimizer"
+        )
+
+        params_grads = self.inner_optimizer.backward(
+            loss, startup_program=startup_program
+        )
+        return params_grads
+
+    def apply_optimize(self, loss, startup_program, params_grads):
+        program = loss.block.program
+        with program_guard(program, startup_program):
+            optimize_ops = self.apply_gradients(params_grads)
+        return optimize_ops
+
+    def _is_the_backward_op(self, op):
+        op_maker = core.op_proto_and_checker_maker
+        backward = core.op_proto_and_checker_maker.OpRole.Backward
+        if op_maker.kOpRoleVarAttrName() in op.attr_names and int(
+            op.all_attrs()[op_maker.kOpRoleAttrName()]
+        ) == int(backward):
+            return True
+        return False
+
+    def _remove_op_role_var(self, param, grad):
+        op_maker = core.op_proto_and_checker_maker
+        op = grad.op
+        assert self._is_the_backward_op(op), (
+            f'grad.op={op} is not the backward op which produces the grad={grad.name}'
+        )
+
+        block = grad.block
+        var_attr = op.all_attrs()[op_maker.kOpRoleVarAttrName()]
+        assert param.name in var_attr, (
+            f'when using GradientMergeOptimizer, param={param.name} must be in var_attr={var_attr}'
+        )
+        assert grad.name in var_attr, (
+            f'when using GradientMergeOptimizer, grad={param.name} must be in var_attr={var_attr}'
+        )
+
+        # remove (param, grad) from op_role_var
+        var_attr.remove(param.name)
+        var_attr.remove(grad.name)
+        if len(var_attr) > 1:
+            op._set_attr(op_maker.kOpRoleVarAttrName(), var_attr)
+        else:
+            op._remove_attr(op_maker.kOpRoleVarAttrName())
+
+    def _add_gm_op_role_var(self, op, param, grad, cond):
+        grad.op = op
+        op_maker = core.op_proto_and_checker_maker
+        backward = op_maker.OpRole.Backward
+
+        # NOTE(wangxi). When distributed, we will insert grad_merge_all_reduce_op_handle
+        # in multi_devices_graph_pass, which will allreduce(grad) if cond is True, else
+        # do nothing.
+        # In this way, the gradient can be merged first, and then communicate when the
+        # condition is met, reducing the number of communications to increase the
+        # speed.
+        op._set_attr(self.GRAD_MERGE_COND_NAME, cond.name)
+        op._set_attr(op_maker.kOpRoleAttrName(), backward)
+        op._set_attr(op_maker.kOpRoleVarAttrName(), [param.name, grad.name])
+
+    def _get_gm_cond_var(self, main_block):
+        # Add const var
+        k_step_var = paddle.static.create_global_var(
+            name="gradient_merge_k",
+            shape=[1],
+            value=int(self.k_steps),
+            dtype='int32',
+            persistable=True,
+            force_cpu=True,
+        )
+
+        zero_var = paddle.static.create_global_var(
+            name="gradient_merge_zero",
+            shape=[1],
+            value=0,
+            dtype='int32',
+            persistable=True,
+            force_cpu=True,
+        )
+
+        # Add step var & cond var
+        step_var = paddle.static.create_global_var(
+            name="gradient_merge_step",
+            shape=[1],
+            value=0,
+            dtype='int32',
+            persistable=True,
+            force_cpu=True,
+        )
+
+        cond_var = main_block.create_var(
+            name="gradient_merge_cond", shape=[1], dtype='bool'
+        )
+
+        with device_guard("cpu"):
+            # step_var = (step_var + 1) % k_step
+            paddle.increment(x=step_var, value=1.0)
+            main_block.append_op(
+                type='elementwise_mod',
+                inputs={'X': step_var, 'Y': k_step_var},
+                outputs={'Out': step_var},
+                attrs={'axis': -1},
+            )
+
+            # cond_var = (step_var == 0)
+            main_block.append_op(
+                type='equal',
+                inputs={'X': step_var, 'Y': zero_var},
+                outputs={'Out': cond_var},
+            )
+
+        return cond_var
+
+    def apply_gradients(self, params_grads):
+        main_program = default_main_program()
+        startup_program = default_startup_program()
+        main_block = main_program.global_block()
+        startup_block = startup_program.global_block()
+
+        cond = self._get_gm_cond_var(main_block)
+
+        # TODO(mapingshuo) support sparse embedding
+        # step1: remove grad.op's op_role_var
+        for param, grad in params_grads:
+            assert param.type != core.VarDesc.VarType.SELECTED_ROWS, (
+                "SELECTED_ROWS is not supported in GradientMergeOptimizer for now"
+            )
+
+            self._remove_op_role_var(param, grad)
+
+        param_to_grad = {k.name: v for (k, v) in params_grads}
+        param_names = param_to_grad.keys()
+        param_to_gradient_merge = {}
+
+        new_params_grads = []
+        # step2: create gradient_merge var and init with 0
+        # and update op_role_var
+        for param, grad in params_grads:
+            param_name = param.name
+            param_var = main_block.var(param_name)
+            assert param_var is not None
+            gradient_merge_var = main_block.create_var(
+                name=param_name + "@GRAD@GradientMerge",
+                shape=param_var.shape,
+                dtype=param_var.dtype,
+                persistable=True,
+            )
+            param_to_gradient_merge[param_name] = gradient_merge_var
+
+            startup_gradient_merge_var = startup_block.create_var(
+                name=param_name + "@GRAD@GradientMerge",
+                shape=param_var.shape,
+                dtype=param_var.dtype,
+                persistable=True,
+            )
+            startup_block.append_op(
+                type="fill_constant",
+                outputs={"Out": startup_gradient_merge_var},
+                attrs={
+                    "shape": param_var.shape,
+                    "dtype": param_var.dtype,
+                    "value": float(0),
+                },
+            )
+
+            # grad_merge += grad
+            new_grad_op = main_block.append_op(
+                type="elementwise_add",
+                inputs={'X': grad, 'Y': gradient_merge_var},
+                outputs={'Out': gradient_merge_var},
+                attrs={'axis': -1},
+            )
+            self._add_gm_op_role_var(
+                new_grad_op, param, gradient_merge_var, cond
+            )
+            new_params_grads.append([param, gradient_merge_var])
+
+        def true_apply_gradient():
+            cur_block_idx = main_program.current_block_idx
+            cur_block = main_program.current_block()
+
+            # cur_block's forward_block & backward_block is itself
+            cur_block._set_forward_block_idx(cur_block_idx)
+            op_maker = core.op_proto_and_checker_maker
+
+            if self.avg:
+                for param, new_grad in new_params_grads:
+                    # grad /= k_steps
+                    cur_block.append_op(
+                        type='scale',
+                        inputs={'X': new_grad},
+                        outputs={'Out': new_grad},
+                        attrs={
+                            'scale': 1.0 / self.k_steps,
+                            'bias': 0.0,
+                            'bias_after_scale': False,
+                        },
+                    )
+                    new_grad.op._set_attr(
+                        op_maker.kOpRoleAttrName(), op_maker.OpRole.Backward
+                    )
+
+            for param, new_grad in new_params_grads:
+                # NOTE. regularization will append ops to grad.block,
+                # while new_grad's real block is global_block,
+                # but we want append regularization ops to cur_block,
+                # so we set new_grad.block = cur_block
+                new_grad.block = cur_block
+
+            self._optimize_ops = self.inner_optimizer.apply_gradients(
+                new_params_grads
+            )
+
+            # clear gradient_merge_vars
+            for param, new_grad in new_params_grads:
+                paddle.tensor.fill_constant(
+                    shape=new_grad.shape,
+                    dtype=new_grad.dtype,
+                    value=0.0,
+                    out=new_grad,
+                )
+                new_grad.op._set_attr(
+                    op_maker.kOpRoleAttrName(), op_maker.OpRole.Optimize
+                )
+
+        # step3. apply gradient
+        paddle.static.nn.cond(cond, true_fn=true_apply_gradient, false_fn=None)
+
+        return self._optimize_ops
+
+    def minimize(
+        self, loss, startup_program=None, parameter_list=None, no_grad_set=None
+    ):
+        assert isinstance(loss, Variable), "The loss should be an Variable."
+
+        params_grads = self.backward(
+            loss,
+            startup_program=startup_program,
+            parameter_list=parameter_list,
+            no_grad_set=no_grad_set,
+        )
+
+        optimize_ops = self.apply_optimize(
+            loss, startup_program=startup_program, params_grads=params_grads
+        )
+
+        return optimize_ops, params_grads

python/paddle/incubate/optimizer/lars_momentum.py

@@ -0,0 +1,240 @@
+# Copyright (c) 2019 PaddlePaddle Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+import warnings
+
+from paddle import _C_ops, _legacy_C_ops, pir
+from paddle.base import framework
+from paddle.framework import (
+    in_dynamic_mode,
+    in_pir_mode,
+)
+from paddle.optimizer import Optimizer
+
+
+class LarsMomentumOptimizer(Optimizer):
+    r"""
+    Momentum optimizer with LARS support
+
+    The update equations are as follows:
+
+    .. math::
+
+        & local\_learning\_rate = learning\_rate * lars\_coeff * \\
+          \\frac{||param||}{||gradient|| + lars\_weight\_decay * ||param||}
+
+        & velocity = mu * velocity + local\_learning\_rate * (gradient + lars\_weight\_decay * param + epsilon)
+
+        & param = param - velocity
+
+    Parameters:
+        learning_rate (float|Variable): The learning rate used to update parameters. \
+            Can be a float value or a Variable with one float value as data element. \
+            momentum (float): momentum factor
+        lars_coeff (float): Defines how much we trust the layer to change its weights.
+        lars_weight_decay (float): Weight decay coefficient for decaying using LARS.
+        parameter_list (Iterable, optional):  Iterable of ``Variable`` names 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.
+        regularization (WeightDecayRegularizer, optional): The strategy of regularization. There are two method: \
+             :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, 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, optional): This parameter is used by developers to print debugging information. \
+            For details, please refer to :ref:`api_guide_Name`. Default is None.
+        exclude_from_weight_decay (list[str], optional): Name string of layers which will be exclude from lars weight decay. Default is None.
+        epsilon (float, optional): Epsilon to avoid Division by Zero when calculate local lr. Default is 0.
+        multi_precision (bool, optional): Whether to use multi-precision during weight updating.
+        rescale_grad (float, optional): Multiply the gradient with `rescale_grad` \
+            before updating. Often choose to be `1.0/batch_size`.
+
+    Examples:
+        .. code-block:: pycon
+
+            >>> import paddle
+            >>> import numpy as np
+
+            >>> paddle.enable_static()
+            >>> np_inp = np.array([[1.0, 2.0], [3.0, 4.0]], dtype=np.float32)
+            >>> inp = paddle.static.data(
+            ...     name="inp", shape=[2, 2], dtype='float32')
+            >>> out = paddle.static.nn.fc(inp, size=3)
+            >>> out = paddle.sum(out)
+            >>> optimizer = paddle.incubate.optimizer.LarsMomentumOptimizer(learning_rate=0.001, momentum=0.9)
+            >>> optimizer.minimize(out)
+
+            >>> exe = paddle.static.Executor(paddle.CPUPlace())
+            >>> exe.run(paddle.static.default_startup_program())
+            >>> exe.run(
+            ...     feed={"inp": np_inp},
+            ...     fetch_list=[out.name])
+    """
+
+    _velocity_acc_str = "velocity"
+
+    def __init__(
+        self,
+        learning_rate,
+        momentum,
+        lars_coeff=0.001,
+        lars_weight_decay=0.0005,
+        parameter_list=None,
+        regularization=None,
+        grad_clip=None,
+        name=None,
+        exclude_from_weight_decay=None,
+        epsilon=0,
+        multi_precision=False,
+        rescale_grad=1.0,
+    ):
+        assert learning_rate is not None
+        assert momentum is not None
+        super().__init__(
+            learning_rate=learning_rate,
+            parameters=parameter_list,
+            weight_decay=regularization,
+            grad_clip=grad_clip,
+            name=name,
+        )
+        self.type = "lars_momentum"
+        self._momentum = momentum
+        self._lars_coeff = float(lars_coeff)
+        self._lars_weight_decay = float(lars_weight_decay)
+        self._epsilon = float(epsilon)
+        if exclude_from_weight_decay is None:
+            self._exclude_from_weight_decay = []
+        else:
+            self._exclude_from_weight_decay = exclude_from_weight_decay
+        self._multi_precision = multi_precision
+        self._rescale_grad = float(rescale_grad)
+        self._master_weights = {}
+
+    def _create_accumulators(self, block, parameters):
+        if not isinstance(block, (framework.Block, pir.Block)):
+            raise TypeError("block is not instance of Block.")
+        for p in parameters:
+            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)
+                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._velocity_acc_str, p)
+
+    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.")
+        _lars_weight_decay = self._lars_weight_decay
+        param_name = param_and_grad[0].name
+        if len(self._exclude_from_weight_decay) > 0:
+            for name in self._exclude_from_weight_decay:
+                if name in param_name:
+                    _lars_weight_decay = 0.0
+                    break
+
+        velocity_acc = self._get_accumulator_master(
+            self._velocity_acc_str, param_and_grad[0]
+        )
+        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
+        )
+        master_weight = (
+            self._master_weights[param_and_grad[0].name]
+            if find_master
+            else None
+        )
+
+        attrs = {
+            "mu": self._momentum,
+            "lars_coeff": self._lars_coeff,
+            "lars_weight_decay": [_lars_weight_decay],
+            "multi_precision": find_master,
+            "epsilon": self._epsilon,
+            "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
+
+        if in_dynamic_mode():
+            tmp, tmp2 = _legacy_C_ops.lars_momentum(
+                [param_and_grad[0]],
+                [param_and_grad[1]],
+                [velocity_acc],
+                [lr],
+                [param_and_grad[0]],
+                [velocity_acc],
+                "mu",
+                self._momentum,
+                "lars_coeff",
+                self._lars_coeff,
+                "lars_weight_decay",
+                [_lars_weight_decay],
+                "multi_precision",
+                find_master,
+                "epsilon",
+                self._epsilon,
+                "rescale_grad",
+                self._rescale_grad,
+            )
+        elif in_pir_mode():
+            if isinstance(master_weight, pir.Value):
+                master_weight = [master_weight]
+            _, _, _ = _C_ops.lars_momentum_(
+                [param_and_grad[0]],
+                [param_and_grad[1]],
+                [velocity_acc],
+                [lr],
+                master_weight,
+                self._momentum,
+                self._lars_coeff,
+                [_lars_weight_decay],
+                self._epsilon,
+                find_master,
+                self._rescale_grad,
+            )
+            return None
+        else:
+            # 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

python/paddle/incubate/optimizer/lbfgs.py

@@ -0,0 +1,440 @@
+# 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 collections import defaultdict
+from functools import reduce
+from typing import TYPE_CHECKING, Any, Literal, TypeVar
+
+import paddle
+from paddle.optimizer import Optimizer
+from paddle.utils import deprecated
+
+from .line_search_dygraph import _strong_wolfe
+
+if TYPE_CHECKING:
+    from collections.abc import Callable, Sequence
+
+    from paddle import Tensor
+    from paddle.nn.clip import GradientClipBase
+    from paddle.optimizer.optimizer import _ParameterConfig
+    from paddle.regularizer import WeightDecayRegularizer
+
+    _T_co = TypeVar('_T_co', covariant=True)
+
+
+@deprecated(since="2.5.0", update_to="paddle.optimizer.LBFGS", level=1)
+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, 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, optional): either 'strong_wolfe' or None. The default value is strong_wolfe.
+        parameters (list|tuple, 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 (float|WeightDecayRegularizer, optional): The strategy of regularization. \
+            It canbe a 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, 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, 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
+            >>> from paddle.incubate.optimizer import LBFGS
+
+            >>> 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 = 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, target in zip(inputs, targets):
+            ...     input_tensor = paddle.to_tensor(input)
+            ...     target_tensor = paddle.to_tensor(target)
+            ...     train_step(input_tensor, target_tensor)
+
+    """
+
+    learning_rate: float
+    max_iter: int
+    max_eval: int
+    tolerance_grad: float
+    tolerance_change: float
+    history_size: int
+    line_search_fn: Literal['strong_wolfe'] | None
+    state: dict[str, dict[str, Any]]
+
+    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: Literal['strong_wolfe'] | None = None,
+        parameters: Sequence[Tensor] | Sequence[_ParameterConfig] | None = None,
+        weight_decay: float | WeightDecayRegularizer | None = None,
+        grad_clip: GradientClipBase | None = None,
+        name: str | None = None,
+    ) -> Tensor:
+        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) -> dict[str, dict[str, Any]]:
+        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.
+        """
+
+        packed_state = {}
+        for k, v in self.state.items():
+            packed_state.update({k: v})
+
+        return {'state': packed_state}
+
+    def _numel(self):
+        # 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)
+            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
+
+    def step(self, closure: Callable[[], _T_co]) -> _T_co:
+        """
+        Performs a single optimization step.
+
+        Args:
+            closure (callable): A closure that reevaluates the model
+                and returns the loss.
+
+        """
+
+        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 = y.dot(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 / y.dot(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] = old_sk[i].dot(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 = old_yk[i].dot(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 = flat_grad.dot(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

python/paddle/incubate/optimizer/line_search_dygraph.py

@@ -0,0 +1,297 @@
+# 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.
+
+import paddle
+
+
+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)
+    """
+
+    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 = paddle.dot(grad_new, d)
+
+    # bracket an interval containing a point satisfying the Wolfe criteria
+    t_prev, f_prev, g_prev, gtd_prev = (
+        paddle.to_tensor(0, dtype=grad.dtype),
+        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 paddle.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 = grad_new.dot(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 paddle.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 paddle.abs(alpha - max(bracket)) < paddle.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 = grad_new.dot(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(memory_format=torch.contiguous_format)
+            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 paddle.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

python/paddle/incubate/optimizer/lookahead.py

@@ -0,0 +1,361 @@
+# 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
+
+from typing import TYPE_CHECKING
+
+import paddle
+from paddle.base import framework, unique_name
+from paddle.base.dygraph import base as imperative_base
+from paddle.base.framework import Variable
+from paddle.base.layer_helper import LayerHelper
+from paddle.framework import in_pir_mode
+from paddle.optimizer import Optimizer
+from paddle.pir.core import create_parameter
+
+if TYPE_CHECKING:
+    from paddle import Tensor
+    from paddle.base.framework import Operator
+    from paddle.static import Program
+
+
+__all__ = []
+
+
+class LookAhead(Optimizer):
+    r"""
+    This implements the Lookahead optimizer of the
+    paper : https://arxiv.org/abs/1907.08610.
+
+    Lookahead keeps two sets of params: the fast_params and
+    the slow_params. inner_optimizer update fast_params every
+    training step. Lookahead updates the slow_params and fast_params
+    every k training steps as follows:
+
+    .. math::
+
+        slow\_param_t &= slow\_param_{t-1} + \\alpha * (fast\_param_{t-1} - slow\_param_{t-1})
+
+        fast\_param_t &=  slow\_param_t
+
+    Args:
+        inner_optimizer (Optimizer): The optimizer that update fast params step by step.
+        alpha (float, optional): The learning rate of Lookahead. The default value is 0.5.
+        k (int, optional): The slow params is updated every k steps. The default value is 5.
+        name (str, 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 numpy as np
+            >>> import paddle
+            >>> import paddle.nn as nn
+
+            >>> BATCH_SIZE = 16
+            >>> BATCH_NUM = 4
+            >>> EPOCH_NUM = 4
+
+            >>> IMAGE_SIZE = 784
+            >>> CLASS_NUM = 10
+            >>> # define a random dataset
+            >>> class RandomDataset(paddle.io.Dataset):  # type: ignore[type-arg]
+            ...     def __init__(self, num_samples):
+            ...         self.num_samples = num_samples
+            ...
+            ...     def __getitem__(self, idx):
+            ...         image = np.random.random([IMAGE_SIZE]).astype('float32')
+            ...         label = np.random.randint(0, CLASS_NUM - 1, (1,)).astype('int64')
+            ...         return image, label
+            ...
+            ...     def __len__(self):
+            ...         return self.num_samples
+
+            >>> class LinearNet(nn.Layer):
+            ...     def __init__(self):
+            ...         super().__init__()
+            ...         self._linear = nn.Linear(IMAGE_SIZE, CLASS_NUM)
+            ...         self.bias = self._linear.bias
+            ...
+            ...     @paddle.jit.to_static
+            ...     def forward(self, x):
+            ...         return self._linear(x)
+
+            >>> def train(layer, loader, loss_fn, opt):
+            ...     for epoch_id in range(EPOCH_NUM):
+            ...         for batch_id, (image, label) in enumerate(loader()):
+            ...             out = layer(image)
+            ...             loss = loss_fn(out, label)
+            ...             loss.backward()
+            ...             opt.step()
+            ...             opt.clear_grad()
+            ...             print("Train Epoch {} batch {}: loss = {}".format(epoch_id, batch_id, np.mean(loss.numpy())))
+            >>> layer = LinearNet()
+            >>> loss_fn = nn.CrossEntropyLoss()
+            >>> optimizer = paddle.optimizer.SGD(learning_rate=0.1, parameters=layer.parameters())
+            >>> lookahead = paddle.incubate.LookAhead(optimizer, alpha=0.2, k=5)
+
+            >>> # create data loader
+            >>> dataset = RandomDataset(BATCH_NUM * BATCH_SIZE)
+            >>> loader = paddle.io.DataLoader(
+            ...     dataset,
+            ...     batch_size=BATCH_SIZE,
+            ...     shuffle=True,
+            ...     drop_last=True,
+            ...     num_workers=2,
+            ... )
+
+            >>> # doctest: +SKIP('The run time is too long to pass the CI check.')
+            >>> train(layer, loader, loss_fn, lookahead)
+
+    """
+
+    inner_optimizer: Optimizer
+    alpha: float
+    k: int
+    type: str
+    helper: LayerHelper
+
+    _slow_str = "slow"
+
+    def __init__(
+        self,
+        inner_optimizer: Optimizer,
+        alpha: float = 0.5,
+        k: int = 5,
+        name: str | None = None,
+    ) -> None:
+        assert inner_optimizer is not None, "inner optimizer can not be None"
+        assert 0.0 <= alpha <= 1.0, (
+            "alpha should be larger or equal to 0.0, and less or equal than 1.0"
+        )
+        assert isinstance(k, int) and k > 0, "k should be a positive integer"
+
+        self.inner_optimizer = inner_optimizer
+        if self.inner_optimizer._parameter_list is None:
+            parameters = (
+                paddle.static.default_main_program()
+                .global_block()
+                .all_parameters()
+            )
+        else:
+            parameters = self.inner_optimizer._parameter_list
+
+        super().__init__(
+            learning_rate=alpha,
+            parameters=parameters,
+            weight_decay=None,
+            grad_clip=None,
+            name=name,
+        )
+
+        self.alpha = alpha
+        self.k = k
+        self.type = "lookahead"
+        self.helper = LayerHelper(self.__class__.__name__)
+        self._global_step_var = None
+        self._k_var = None
+
+    def _set_auxiliary_var(self, key, val):
+        super()._set_auxiliary_var(key, val)
+        self.inner_optimizer._set_auxiliary_var(key, val)
+
+    @framework.dygraph_only
+    @imperative_base.no_grad
+    def step(self) -> None:
+        """
+        Execute the optimizer and update parameters once.
+
+        Returns:
+            None
+
+        Examples:
+
+            .. code-block:: pycon
+
+                >>> import paddle
+                >>> inp = paddle.rand([1, 10], dtype="float32")
+                >>> linear = paddle.nn.Linear(10, 1)
+                >>> out = linear(inp)
+                >>> loss = paddle.mean(out)
+                >>> sgd = paddle.optimizer.SGD(learning_rate=0.1, parameters=linear.parameters())
+                >>> lookahead = paddle.incubate.LookAhead(sgd, alpha=0.2, k=5)
+                >>> loss.backward()
+                >>> lookahead.step()
+                >>> lookahead.clear_grad()
+
+        """
+        self.inner_optimizer.step()
+
+        self._increment_global_var()
+        params_grads = []
+        for param in self._parameter_list:
+            if not param.trainable:
+                continue
+            if param._grad_ivar() is not None:
+                grad_var = param._grad_ivar()
+                params_grads.append((param, grad_var))
+
+        self._apply_optimize(
+            loss=None, startup_program=None, params_grads=params_grads
+        )
+
+    def _create_accumulators(self, block, parameters):
+        assert isinstance(block, (framework.Block, paddle.pir.Block))
+
+        for p in parameters:
+            self._add_accumulator(self._slow_str, p)
+
+    def _increment_global_var(self):
+        if in_pir_mode():
+            if self._global_step_var is None:
+                self._global_step_var = create_parameter(
+                    dtype='int32',
+                    shape=[1],
+                    name=unique_name.generate("lookahead_step"),
+                    trainable=False,
+                    initializer=paddle.nn.initializer.ConstantInitializer(
+                        value=0.0, force_cpu=False
+                    ),
+                )
+            self._global_step_var = paddle.increment(self._global_step_var, 1.0)
+        else:
+            if self._global_step_var is None:
+                self._global_step_var = paddle.static.create_global_var(
+                    name=unique_name.generate("lookahead_step"),
+                    shape=[1],
+                    value=0,
+                    dtype='int32',
+                    persistable=True,
+                )
+
+            self.helper.append_op(
+                type='increment',
+                inputs={'X': [self._global_step_var]},
+                outputs={'Out': [self._global_step_var]},
+                attrs={'step': 1.0},
+            )
+
+    def _append_optimize_op(self, block, param_and_grad):
+        one_var = paddle.ones(shape=[1], dtype='int32', name='lookahead_ones')
+        zero_var = paddle.zeros(
+            shape=[1], dtype='int32', name='lookahead_zeros'
+        )
+        if in_pir_mode():
+            k_var = create_parameter(
+                dtype='int32',
+                shape=[1],
+                name=unique_name.generate("lookahead_k"),
+                trainable=False,
+                initializer=paddle.nn.initializer.ConstantInitializer(
+                    value=float(self.k), force_cpu=False
+                ),
+            )
+        else:
+            k_var = paddle.static.create_global_var(
+                name=unique_name.generate("lookahead_k"),
+                shape=[1],
+                value=self.k,
+                dtype='int32',
+                persistable=True,
+            )
+
+        mod = paddle.remainder(self._global_step_var, k_var)
+
+        cond_1 = paddle.equal(self._global_step_var, one_var)
+        cond_1 = paddle.cast(cond_1, dtype='float32')
+
+        cond_2 = paddle.equal(mod, zero_var)
+        cond_2 = paddle.cast(cond_2, dtype='float32')
+
+        slow_var = self._get_accumulator(self._slow_str, param_and_grad[0])
+
+        tmp_var = cond_1 * param_and_grad[0] + (1 - cond_1) * slow_var
+        paddle.assign(tmp_var, slow_var)
+
+        tmp_var = self.alpha * param_and_grad[0] + (1.0 - self.alpha) * slow_var
+        tmp_var_1 = cond_2 * tmp_var + (1 - cond_2) * param_and_grad[0]
+        paddle.assign(tmp_var_1, param_and_grad[0])
+
+        tmp_var_1 = cond_2 * tmp_var + (1 - cond_2) * slow_var
+        paddle.assign(tmp_var_1, slow_var)
+
+    @imperative_base.no_grad
+    def minimize(
+        self,
+        loss: Tensor,
+        startup_program: Program | None = None,
+        parameters: list[Tensor] | list[str] | None = None,
+        no_grad_set: set[Tensor] | set[str] | None = None,
+    ) -> tuple[list[Operator], list[tuple[Tensor, Tensor]]]:
+        """
+        Add operations to minimize ``loss`` by updating ``parameters``.
+
+        Args:
+            loss (Tensor): A ``Tensor`` containing the value to minimize.
+            startup_program (Program, optional): :ref:`api_paddle_static_Program` for
+                initializing parameters in ``parameters``. The default value
+                is None, at this time :ref:`api_paddle_static_default_startup_program` will be used.
+            parameters (list, optional): List of ``Tensor`` or ``Tensor.name`` to update
+                to minimize ``loss``. The default value is None, at this time all parameters
+                will be updated.
+            no_grad_set (set, optional): Set of ``Tensor``  or ``Tensor.name`` that don't need
+                to be updated. The default value is None.
+
+        Returns:
+            tuple: tuple (optimize_ops, params_grads), A list of operators appended
+            by minimize and a list of (param, grad) tensor pairs, param is
+            ``Parameter``, grad is the gradient value corresponding to the parameter.
+            In static graph mode, the returned tuple can be passed to ``fetch_list`` in ``Executor.run()`` to
+            indicate program pruning. If so, the program will be pruned by ``feed`` and
+            ``fetch_list`` before run, see details in ``Executor``.
+
+        Examples:
+
+            .. code-block:: pycon
+
+                >>> import paddle
+
+                >>> inp = paddle.rand([1, 10], dtype="float32")
+                >>> linear = paddle.nn.Linear(10, 1)
+                >>> out = linear(inp)
+                >>> loss = paddle.mean(out)
+                >>> sgd = paddle.optimizer.SGD(learning_rate=0.1, parameters=linear.parameters())
+                >>> lookahead = paddle.incubate.LookAhead(sgd, alpha=0.2, k=5)
+                >>> loss.backward()
+                >>> lookahead.minimize(loss)
+                >>> lookahead.clear_grad()
+
+        """
+        assert isinstance(loss, (Variable, paddle.pir.Value)), (
+            "The loss should be an Tensor."
+        )
+
+        # Apply inner optimizer to the main_program
+        optimize_ops, params_grads = self.inner_optimizer.minimize(
+            loss,
+            startup_program=startup_program,
+            parameters=parameters,
+            no_grad_set=no_grad_set,
+        )
+
+        self._increment_global_var()
+
+        _ = self._apply_optimize(
+            loss, startup_program=startup_program, params_grads=params_grads
+        )
+
+        return optimize_ops, params_grads

python/paddle/incubate/optimizer/modelaverage.py

@@ -0,0 +1,625 @@
+# 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
+
+from typing import TYPE_CHECKING
+
+import paddle
+from paddle import _C_ops
+from paddle.base import framework
+from paddle.base.dygraph import base as imperative_base
+from paddle.base.layer_helper import LayerHelper
+from paddle.base.wrapped_decorator import signature_safe_contextmanager
+from paddle.framework import (
+    in_dynamic_mode,
+    in_dynamic_or_pir_mode,
+    in_pir_mode,
+)
+from paddle.optimizer import Optimizer
+
+if TYPE_CHECKING:
+    from collections.abc import Generator, Sequence
+
+    from paddle import Tensor
+    from paddle.optimizer.optimizer import _ParameterConfig
+    from paddle.static import Executor, Program
+
+
+__all__ = []
+
+
+class ModelAverage(Optimizer):
+    r"""
+    The ModelAverage optimizer accumulates specific continuous historical
+    parameters during training. The accumulated historical range can be controlled
+    by the passed ``average_window_rate`` argument. The averaged ``Parameter`` are
+    used in the prediction, which usually can improve the accuracy of the prediction.
+
+    Accumulate the average of the ``Parameter`` in the sliding window, the result will be saved
+    in a temporary variable, can be applied to the current model's ``Parameter`` by calling
+    the ``apply()`` method, and the current model ``Parameter`` can be restored by calling
+    the ``restore()`` method.
+
+    The window size for calculating the average is determined by ``average_window_rate``,
+    ``min_average_window``, ``max_average_window`` and the current ``Parameter`` update times (num_updates).
+
+    When the cumulative times (num_accumulates) is greater than the specific window
+    threshold (average_window), the accumulated ``Parameter`` temporary variable is set to 0.0.
+    The following example will help to understand the role of these arguments:
+
+    ::
+
+        if num_accumulates >= min_average_window and num_accumulates >= min(max_average_window, num_updates * average_window_rate):
+            num_accumulates = 0
+
+    In the above conditional judgment statement, ``num_accumulates`` indicates the current
+    accumulated number, which can be abstractly understood as the length of the cumulative window.
+    The length of the window must be at least the length set by the ``min_average_window`` argument,
+    and cannot exceed the length specified by the ``max_average_window`` argument or
+    ``num_updates * average_window_rate``, where ``num_updates`` indicates the current ``Parameter``
+    update times, ``average_window_rate`` is a coefficient that calculates the length of the window.
+
+    Args:
+        average_window_rate (float): The calculate ratio of the window length relative to ``Parameter`` update times.
+        parameters (list, optional): List of ``Tensor`` names 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.
+        min_average_window (int, optional): the minimum size of average window length. The default value is 10000.
+        max_average_window (int, optional): The maximum size of average window length. The default value is 10000.
+        name (str, 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
+
+            >>> # doctest: +SKIP("Cannot get source code by to_static in REPL")
+            >>> import numpy as np
+            >>> import paddle
+            >>> import paddle.nn as nn
+            >>> import paddle.optimizer as opt
+
+            >>> BATCH_SIZE = 16
+            >>> BATCH_NUM = 4
+            >>> EPOCH_NUM = 4
+
+            >>> IMAGE_SIZE = 784
+            >>> CLASS_NUM = 10
+
+            >>> # define a random dataset
+            >>> class RandomDataset(paddle.io.Dataset): # type: ignore[type-arg]
+            ...     def __init__(self, num_samples):
+            ...         self.num_samples = num_samples
+            ...     def __getitem__(self, idx):
+            ...         image = np.random.random([IMAGE_SIZE]).astype('float32')
+            ...         label = np.random.randint(0, CLASS_NUM - 1, (1, )).astype('int64')
+            ...         return image, label
+            ...     def __len__(self):
+            ...         return self.num_samples
+            ...
+            >>> class LinearNet(nn.Layer):
+            ...     def __init__(self):
+            ...         super().__init__()
+            ...         self._linear = nn.Linear(IMAGE_SIZE, CLASS_NUM)
+            ...         self.bias = self._linear.bias
+            ...
+            ...     @paddle.jit.to_static
+            ...     def forward(self, x):
+            ...         return self._linear(x)
+            ...
+            >>> def train(layer, loader, loss_fn, opt, model_average):
+            ...     for epoch_id in range(EPOCH_NUM):
+            ...         for batch_id, (image, label) in enumerate(loader()):
+            ...             out = layer(image)
+            ...             loss = loss_fn(out, label)
+            ...             loss.backward()
+            ...             opt.step()
+            ...             model_average.step()
+            ...             opt.clear_grad()
+            ...             model_average.clear_grad()
+            ...             print("Train Epoch {} batch {}: loss = {}, bias = {}".format(
+            ...                 epoch_id, batch_id, np.mean(loss.numpy()), layer.bias.numpy()))
+            ...
+            >>> def evaluate(layer, loader, loss_fn):
+            ...     for batch_id, (image, label) in enumerate(loader()):
+            ...         out = layer(image)
+            ...         loss = loss_fn(out, label)
+            ...         loss.backward()
+            ...         print("Evaluate batch {}: loss = {}, bias = {}".format(
+            ...             batch_id, np.mean(loss.numpy()), layer.bias.numpy()))
+            ...
+            >>> # create network
+            >>> layer = LinearNet()
+            >>> loss_fn = nn.CrossEntropyLoss()
+            >>> optimizer = opt.Momentum(learning_rate=0.2, momentum=0.1, parameters=layer.parameters())
+            >>> model_average = paddle.incubate.ModelAverage(
+            ...     0.15,
+            ...     parameters=layer.parameters(),
+            ...     min_average_window=2,
+            ...     max_average_window=10
+            ... )
+            ...
+            >>> # create data loader
+            >>> dataset = RandomDataset(BATCH_NUM * BATCH_SIZE)
+            >>> loader = paddle.io.DataLoader(dataset,
+            ...     batch_size=BATCH_SIZE,
+            ...     shuffle=True,
+            ...     drop_last=True,
+            ...     num_workers=2)
+            ...
+            >>> # create data loader
+            >>> eval_loader = paddle.io.DataLoader(dataset,
+            ...     batch_size=BATCH_SIZE,
+            ...     shuffle=True,
+            ...     drop_last=True,
+            ...     num_workers=1
+            ... )
+            ...
+            >>> # train
+            >>> train(layer, loader, loss_fn, optimizer, model_average)
+
+            >>> print("\nEvaluate With ModelAverage")
+            >>> with model_average.apply(need_restore=False):
+            ...     evaluate(layer, eval_loader, loss_fn)
+
+            >>> print("\nEvaluate With Restored Parameters")
+            >>> model_average.restore()
+            >>> evaluate(layer, eval_loader, loss_fn)
+
+    """
+
+    helper: LayerHelper
+    average_window: float
+    min_average_window: int
+    max_average_window: int
+    type: str
+    apply_program: Program
+    restore_program: Program
+
+    def __init__(
+        self,
+        average_window_rate: float,
+        parameters: Sequence[Tensor] | Sequence[_ParameterConfig] | None = None,
+        min_average_window: int = 10000,
+        max_average_window: int = 10000,
+        name: str | None = None,
+    ) -> None:
+        super().__init__(
+            learning_rate=0.0,
+            parameters=parameters,
+            weight_decay=None,
+            grad_clip=None,
+            name=name,
+        )
+
+        self.helper = LayerHelper(self.__class__.__name__)
+        self.average_window = average_window_rate
+        self.min_average_window = min_average_window
+        self.max_average_window = max_average_window
+        self.type = "average_accumulates"
+
+        if not in_dynamic_mode():
+            global_block = paddle.static.default_main_program().global_block()
+            all_parameters = (
+                parameters if parameters else global_block.all_parameters()
+            )
+
+            self._create_accumulators(global_block, all_parameters)
+            for param in all_parameters:
+                self._append_optimize_op(global_block, [param, None])
+            self.apply_program = paddle.static.Program()
+            block = self.apply_program.global_block()
+            with paddle.static.program_guard(main_program=self.apply_program):
+                for param in all_parameters:
+                    self._add_average_apply_op(block, param)
+            self.restore_program = paddle.static.Program()
+            block = self.restore_program.global_block()
+            with paddle.static.program_guard(main_program=self.restore_program):
+                for param in all_parameters:
+                    self._add_average_restore_op(block, param)
+
+    def _create_accumulators(self, block, parameters):
+        assert isinstance(block, (framework.Block, paddle.pir.Block))
+
+        for param in parameters:
+            self._add_accumulator('sum_1', param)
+            self._add_accumulator('sum_2', param)
+            self._add_accumulator('sum_3', param)
+            self._add_accumulator('restore', param)
+            self._add_accumulator(
+                'num_accumulates', param, dtype='int64', shape=[1]
+            )
+            self._add_accumulator(
+                'old_num_accumulates', param, dtype='int64', shape=[1]
+            )
+            self._add_accumulator(
+                'num_updates', param, dtype='int64', shape=[1]
+            )
+
+    def _append_optimize_op(self, block, param_and_grad):
+        assert isinstance(block, (framework.Block, paddle.pir.Block))
+
+        sum_1 = self._get_accumulator('sum_1', param_and_grad[0])
+        sum_2 = self._get_accumulator('sum_2', param_and_grad[0])
+        sum_3 = self._get_accumulator('sum_3', param_and_grad[0])
+        num_accumulates = self._get_accumulator(
+            'num_accumulates', param_and_grad[0]
+        )
+        old_num_accumulates = self._get_accumulator(
+            'old_num_accumulates', param_and_grad[0]
+        )
+        num_updates = self._get_accumulator('num_updates', param_and_grad[0])
+
+        if in_dynamic_or_pir_mode():
+            _, _, _, _, _, _ = _C_ops.average_accumulates_(
+                param_and_grad[0],
+                sum_1,
+                sum_2,
+                sum_3,
+                num_accumulates,
+                old_num_accumulates,
+                num_updates,
+                self.average_window,
+                self.max_average_window,
+                self.min_average_window,
+            )
+            return None
+
+        block = framework.default_main_program().global_block()
+        attrs = {
+            "average_window": self.average_window,
+            "min_average_window": self.min_average_window,
+            "max_average_window": self.max_average_window,
+        }
+
+        inputs = {
+            "param": param_and_grad[0],
+            "in_sum_1": sum_1,
+            "in_sum_2": sum_2,
+            "in_sum_3": sum_3,
+            "in_num_accumulates": num_accumulates,
+            "in_old_num_accumulates": old_num_accumulates,
+            "in_num_updates": num_updates,
+        }
+
+        outputs = {
+            "out_sum_1": sum_1,
+            "out_sum_2": sum_2,
+            "out_sum_3": sum_3,
+            "out_num_accumulates": num_accumulates,
+            "out_old_num_accumulates": old_num_accumulates,
+            "out_num_updates": num_updates,
+        }
+
+        average_accumulates_op = block.append_op(
+            type=self.type,
+            inputs=inputs,
+            outputs=outputs,
+            attrs=attrs,
+            stop_gradient=True,
+        )
+
+        return average_accumulates_op
+
+    @imperative_base.no_grad
+    def minimize(
+        self,
+        loss: Tensor,
+        startup_program: Program | None = None,
+        parameters: list[Tensor] | None = None,
+        no_grad_set: set[Tensor] | set[str] | None = None,
+    ) -> None:
+        """
+        Add operations to minimize ``loss`` by updating ``parameters``.
+
+        Args:
+            loss (Tensor): A ``Tensor`` containing the value to minimize.
+            startup_program (Program, optional): :ref:`api_paddle_static_Program` for
+                initializing parameters in ``parameters``. The default value
+                is None, at this time :ref:`api_paddle_static_default_startup_program` will be used.
+            parameters (list, optional): List of ``Tensor`` or ``Tensor.name`` to update
+                to minimize ``loss``. The default value is None, at this time all parameters
+                will be updated.
+            no_grad_set (set, optional): Set of ``Tensor``  or ``Tensor.name`` that don't need
+                to be updated. The default value is None.
+
+        Returns:
+            tuple: tuple (optimize_ops, params_grads), A list of operators appended
+            by minimize and a list of (param, grad) tensor pairs, param is
+            ``Parameter``, grad is the gradient value corresponding to the parameter.
+            In static graph mode, the returned tuple can be passed to ``fetch_list`` in ``Executor.run()`` to
+            indicate program pruning. If so, the program will be pruned by ``feed`` and
+            ``fetch_list`` before run, see details in ``Executor``.
+
+        Examples:
+
+            .. code-block:: pycon
+
+                >>> import paddle
+                >>> inp = paddle.rand([1, 10], dtype="float32")
+                >>> linear = paddle.nn.Linear(10, 1)
+                >>> out = linear(inp)
+                >>> loss = paddle.mean(out)
+                >>> loss.backward()
+
+                >>> sgd = paddle.optimizer.SGD(learning_rate=0.1, parameters=linear.parameters())
+                >>> sgd.minimize(loss)
+
+                >>> modelaverage = paddle.incubate.ModelAverage(
+                ...     0.15,
+                ...     parameters=linear.parameters(),
+                ...     min_average_window=2,
+                ...     max_average_window=4,
+                ... )
+                >>> modelaverage.minimize(loss)
+                >>> sgd.clear_grad()
+                >>> modelaverage.clear_grad()
+
+        """
+        if in_dynamic_mode():
+            self.step()
+
+    @framework.dygraph_only
+    @imperative_base.no_grad
+    def step(self) -> None:
+        """
+        Execute the optimizer and update parameters once.
+
+        Returns:
+            None
+
+        Examples:
+
+            .. code-block:: pycon
+
+                >>> import paddle
+                >>> inp = paddle.rand([1, 10], dtype="float32")
+                >>> linear = paddle.nn.Linear(10, 1)
+                >>> out = linear(inp)
+                >>> loss = paddle.mean(out)
+                >>> sgd = paddle.optimizer.SGD(learning_rate=0.1, parameters=linear.parameters())
+                >>> modelaverage = paddle.incubate.ModelAverage(
+                ...     0.15,
+                ...     parameters=linear.parameters(),
+                ...     min_average_window=2,
+                ...     max_average_window=4,
+                ... )
+                >>> loss.backward()
+                >>> sgd.step()
+                >>> modelaverage.step()
+                >>> sgd.clear_grad()
+                >>> modelaverage.clear_grad()
+        """
+
+        params_grads = []
+        for param in self._parameter_list:
+            if not param.trainable:
+                continue
+            if param._grad_ivar() is not None:
+                grad_var = param._grad_ivar()
+                params_grads.append((param, grad_var))
+
+        block = framework.default_main_program().global_block()
+        self._create_accumulators(block, self._parameter_list)
+        for param_and_grad in params_grads:
+            self._append_optimize_op(block, param_and_grad)
+
+    @signature_safe_contextmanager
+    @imperative_base.no_grad
+    def apply(
+        self, executor: Executor | None = None, need_restore: bool = True
+    ) -> Generator[None, None, None]:
+        """
+        Apply the average of the cumulative ``Parameter`` to the parameters of the current model.
+
+        Args:
+            executor(Executor): The network executor in static-graph mode. The default value is None in dygraph mode.
+            need_restore(bool): Restore flag variable, if set to True, the network will restore
+                the parameters of the network to the default value, if set to False,
+                it will not be restored. The default value is True.
+
+        Examples:
+
+            .. code-block:: pycon
+
+                >>> import paddle
+                >>> inp = paddle.rand([1, 10], dtype="float32")
+                >>> linear = paddle.nn.Linear(10, 1)
+                >>> out = linear(inp)
+                >>> loss = paddle.mean(out)
+                >>> loss.backward()
+
+                >>> sgd = paddle.optimizer.SGD(learning_rate=0.1, parameters=linear.parameters())
+
+                >>> modelaverage = paddle.incubate.ModelAverage(
+                ...     0.15,
+                ...     parameters=linear.parameters(),
+                ...     min_average_window=2,
+                ...     max_average_window=4,
+                ... )
+                >>> sgd.step()
+                >>> modelaverage.step()
+
+                >>> with modelaverage.apply():
+                ...     for param in linear.parameters():
+                ...         print(param)
+
+                >>> for param in linear.parameters():
+                ...     print(param)
+        """
+        if in_dynamic_mode():
+            for param in self._parameter_list:
+                num_accumulates = self._get_accumulator(
+                    'num_accumulates', param
+                )
+                old_num_accumulates = self._get_accumulator(
+                    'old_num_accumulates', param
+                )
+                sum_1 = self._get_accumulator('sum_1', param)
+                sum_2 = self._get_accumulator('sum_2', param)
+                sum_3 = self._get_accumulator('sum_3', param)
+                param_restore = self._get_accumulator('restore', param)
+
+                paddle.assign(param, param_restore)
+                total_param = sum_1 + sum_2 + sum_3
+                total_accumulates = num_accumulates + old_num_accumulates
+                total_param = paddle.cast(total_param, dtype='float32')
+                total_accumulates = paddle.cast(
+                    total_accumulates, dtype='float32'
+                )
+                average_param = total_param / total_accumulates
+                paddle.assign(average_param, param)
+            try:
+                yield
+            finally:
+                if need_restore:
+                    self.restore()
+            return
+        if executor is None:
+            raise RuntimeError(
+                "Executor should not be None in static graph mode."
+            )
+        executor.run(self.apply_program)
+        try:
+            yield
+        finally:
+            if need_restore:
+                self.restore(executor)
+
+    @imperative_base.no_grad
+    def restore(self, executor: Executor | None = None) -> None:
+        """
+        Restore ``Parameter`` values of current model.
+
+        Args:
+            executor(Executor): The network executor in static-graph mode. The default value is None in dygraph mode
+
+        Examples:
+
+            .. code-block:: pycon
+
+                >>> import paddle
+                >>> inp = paddle.rand([1, 10], dtype="float32")
+                >>> linear = paddle.nn.Linear(10, 1)
+                >>> out = linear(inp)
+                >>> loss = paddle.mean(out)
+                >>> loss.backward()
+
+                >>> sgd = paddle.optimizer.SGD(learning_rate=0.1, parameters=linear.parameters())
+
+                >>> modelaverage = paddle.incubate.ModelAverage(
+                ...     0.15,
+                ...     parameters=linear.parameters(),
+                ...     min_average_window=2,
+                ...     max_average_window=4,
+                ... )
+                >>> sgd.step()
+                >>> modelaverage.step()
+
+                >>> with modelaverage.apply(need_restore=False):
+                ...     for param in linear.parameters():
+                ...         print(param)
+
+                >>> for param in linear.parameters():
+                ...     print(param)
+
+                >>> modelaverage.restore()
+
+                >>> for param in linear.parameters():
+                ...     print(param)
+        """
+        if in_dynamic_mode():
+            for param in self._parameter_list:
+                param_restore = self._get_accumulator('restore', param)
+                paddle.assign(param_restore, param)
+            return
+        if executor is None:
+            raise RuntimeError(
+                "Executor should not be None in static graph mode."
+            )
+        executor.run(self.restore_program)
+
+    def _add_average_apply_op(self, block, param):
+        if in_pir_mode():
+            target_program = paddle.static.default_main_program()
+            param = paddle.pir.core._get_parameter(target_program, param)
+            restore_value = self._get_accumulator('restore', param)
+            grad = paddle.pir.core._get_persistable_value(
+                target_program, restore_value
+            )
+            sum_1 = self._get_accumulator('sum_1', param)
+            sum_1 = paddle.pir.core._get_persistable_value(
+                target_program, sum_1
+            )
+            sum_2 = self._get_accumulator('sum_2', param)
+            sum_2 = paddle.pir.core._get_persistable_value(
+                target_program, sum_2
+            )
+            sum_3 = self._get_accumulator('sum_3', param)
+            sum_3 = paddle.pir.core._get_persistable_value(
+                target_program, sum_3
+            )
+            num_accumulates = self._get_accumulator('num_accumulates', param)
+            num_accumulates = paddle.pir.core._get_persistable_value(
+                target_program, num_accumulates
+            )
+            old_num_accumulates = self._get_accumulator(
+                'old_num_accumulates', param
+            )
+            old_num_accumulates = paddle.pir.core._get_persistable_value(
+                target_program, old_num_accumulates
+            )
+        else:
+            param = block._clone_variable(param)
+            grad = block._clone_variable(
+                self._get_accumulator('restore', param)
+            )
+            sum_1 = block._clone_variable(self._get_accumulator('sum_1', param))
+            sum_2 = block._clone_variable(self._get_accumulator('sum_2', param))
+            sum_3 = block._clone_variable(self._get_accumulator('sum_3', param))
+            num_accumulates = block._clone_variable(
+                self._get_accumulator('num_accumulates', param)
+            )
+            old_num_accumulates = block._clone_variable(
+                self._get_accumulator('old_num_accumulates', param)
+            )
+        # backup param value to grad
+        paddle.assign(param, output=grad)
+        # param = (sum_1 + sum_2 + sum_3) / (num_accumulates + old_num_accumulates)
+        tmp = paddle.add_n([num_accumulates, old_num_accumulates])
+        sum = paddle.add_n([sum_1, sum_2, sum_3])
+        tmp = paddle.cast(
+            x=tmp, dtype='float32' if self._dtype is None else self._dtype
+        )
+        sum = paddle.cast(
+            x=sum, dtype='float32' if self._dtype is None else self._dtype
+        )
+        divide_out = paddle.divide(x=sum, y=tmp)
+        paddle.assign(divide_out, output=param)
+
+    def _add_average_restore_op(self, block, param):
+        if in_pir_mode():
+            target_program = paddle.static.default_main_program()
+            param = paddle.pir.core._get_parameter(target_program, param)
+            restore_value = self._get_accumulator('restore', param)
+            grad = paddle.pir.core._get_persistable_value(
+                target_program, restore_value
+            )
+        else:
+            param = block._clone_variable(param)
+            grad = block._clone_variable(
+                self._get_accumulator('restore', param)
+            )
+        paddle.assign(grad, output=param)

python/paddle/incubate/optimizer/recompute.py

@@ -0,0 +1,811 @@
+# Copyright (c) 2019 PaddlePaddle Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import logging
+
+import paddle
+from paddle.base import core, framework, unique_name
+from paddle.base.backward import append_backward
+from paddle.base.framework import Variable, in_dygraph_mode, program_guard
+from paddle.optimizer import Optimizer
+
+
+class RecomputeOptimizer(Optimizer):
+    """
+        :api_attr: Static Graph
+
+    Recompute Optimizer Wrapper
+
+    Normally, a training step contains three sub-steps: first, run forward
+    Operators to calculate the loss; second, run backward Operators to
+    calculate gradient of the parameters; third, apply optimization method
+    to update the value of the parameters.
+
+    In the forward computation process, all variables that are needed by
+    backward computation process will be kept in memory, which occupy a great
+    amount of memory when the network becomes very deep.
+
+    Recompute split the network to k segments. In each segment, It will
+    recompute the forward Operators, before running backward operators. It is
+    very helpful for saving memory.
+
+    The Variables that separate a network to segments are called as checkpoints,
+    and users should set it manually. The usage is very simple:
+
+    Args:
+        optimizer (Optimizer): The optimizer that is applied to parameters.
+
+    Examples:
+        .. code-block:: pycon
+
+            >>> import paddle
+            >>> import numpy as np
+
+            >>> paddle.enable_static()
+
+            >>> def gen_data():
+            ...     return {
+            ...         "x": np.random.random(size=(32, 32)).astype('float32'),
+            ...         "y": np.random.randint(2, size=(32, 1)).astype('int64'),
+            ...     }
+            >>> def mlp(input_x, input_y, hid_dim=128, label_dim=2):
+            ...     print(input_x)
+            ...     fc_1 = paddle.static.nn.fc(x=input_x, size=hid_dim)
+            ...     prediction = paddle.static.nn.fc(x=[fc_1], size=label_dim, activation='softmax')
+            ...     cost = paddle.nn.functional.cross_entropy(
+            ...         input=prediction,
+            ...         label=input_y,
+            ...         reduction='none',
+            ...         use_softmax=False,
+            ...     )
+            ...     sum_cost = paddle.mean(cost)
+            ...     return sum_cost, fc_1, prediction
+            >>> input_x = paddle.static.data(name="x", shape=[-1, 32], dtype='float32')
+            >>> input_y = paddle.static.data(name="y", shape=[-1, 1], dtype='int64')
+            >>> cost, fc_1, pred = mlp(input_x, input_y)
+
+            >>> sgd = paddle.optimizer.Adam(learning_rate=0.01)
+            >>> sgd = paddle.incubate.optimizer.RecomputeOptimizer(sgd)
+            >>> sgd._set_checkpoints([fc_1, pred])
+            >>> sgd.minimize(cost)
+
+            >>> print("Finished optimize")
+            Finished optimize
+            >>> place = paddle.CPUPlace()
+            >>> exe = paddle.static.Executor(place)
+            >>> exe.run(paddle.static.default_startup_program())
+            >>> step = 10
+
+            >>> for i in range(step):
+            ...     cost_val = exe.run(
+            ...         feed=gen_data(),
+            ...         program=paddle.static.default_main_program(),
+            ...         fetch_list=[cost.name],
+            ...     )
+            ...     print("step=%d cost=%f" % (i, cost_val[0]))
+            var x : DENSE_TENSOR.shape(-1, 32).dtype(float32).stop_gradient(True)
+            Finished optimize
+            step=0 cost=0.737203
+            step=1 cost=1.308077
+            step=2 cost=0.768422
+            step=3 cost=1.239475
+            step=4 cost=0.882643
+            step=5 cost=0.738027
+            step=6 cost=0.819374
+            step=7 cost=0.818534
+            step=8 cost=0.753692
+            step=9 cost=0.787448
+
+    """
+
+    def __init__(self, optimizer):
+        if in_dygraph_mode():
+            raise Exception("In dygraph, don't support RecomputeOptimizer.")
+        self._optimizer = optimizer
+        self._checkpoints = None
+        self._learning_rate = self._optimizer._learning_rate
+        self._learning_rate_map = self._optimizer._learning_rate_map
+        self.enable_offload = False
+
+    def _set_checkpoints(self, checkpoints):
+        """
+        Args:
+            checkpoints (list): List of Variable or string
+        """
+        assert isinstance(checkpoints, list), (
+            "_checkpoints should be a list of Variable or a list of String"
+        )
+        for ckpt in checkpoints:
+            assert isinstance(ckpt, (Variable, str)), (
+                "_checkpoints should be a list of Variable or a list of String"
+            )
+        self._checkpoints = checkpoints
+
+    # should enable offload before calling backward
+    def _enable_offload(self):
+        self.enable_offload = True
+
+    @framework.deprecate_stat_dict
+    def load(self, state_dict):
+        """
+            :api_attr: Static Graph
+
+        load function is not supported by Recompute Optimizer for now.
+        :return: None
+
+        Args:
+            state_dict: the dict load by load_persistable method
+
+        Examples:
+            .. code-block:: pycon
+
+                >>> import paddle
+
+                >>> paddle.enable_static()
+                >>> def mlp(input_x, input_y, hid_dim=128, label_dim=2):
+                ...     fc_1 = paddle.static.nn.fc(x=input_x, size=hid_dim)
+                ...     prediction = paddle.static.nn.fc(x=[fc_1], size=label_dim, activation='softmax')
+                ...     cost = paddle.nn.functional.cross_entropy(
+                ...         input=prediction,
+                ...         label=input_y,
+                ...         reduction='none',
+                ...         use_softmax=False,
+                ...     )
+                ...     sum_cost = paddle.mean(cost)
+                ...     return sum_cost, fc_1, prediction
+
+                >>> input_x = paddle.static.data(name="x", shape=[-1, 32], dtype='float32')
+                >>> input_y = paddle.static.data(name="y", shape=[-1, 1], dtype='int64')
+                >>> cost, fc_1, pred = mlp(input_x, input_y)
+                >>> print("Finished FF")
+                Finished FF
+
+                >>> sgd = paddle.optimizer.Adam(learning_rate=0.01)
+                >>> sgd = paddle.incubate.optimizer.RecomputeOptimizer(sgd)
+                >>> sgd._set_checkpoints([fc_1, pred])
+                >>> try:
+                ...     state_dict = {}
+                ...     sgd.load(state_dict)
+                >>> except NotImplementedError as e:
+                ...     print(e)
+                load function is not supported by Recompute Optimizer for now
+        """
+        raise NotImplementedError(
+            "load function is not supported by Recompute Optimizer for now"
+        )
+
+    def apply_gradients(self, params_grads):
+        """
+        call apply_gradients function of self._optimizer.
+
+        Args:
+            params_grads (list): list of (param, grad) pair to do optimization.
+
+        Returns:
+            list: A list of operators appended to the current program.
+
+        Examples:
+            .. code-block:: pycon
+
+                >>> import paddle
+                >>> import paddle.base.framework as framework
+
+                >>> paddle.enable_static()
+
+                >>> def mlp(input_x, input_y, hid_dim=128, label_dim=2):
+                ...     fc_1 = paddle.static.nn.fc(x=input_x, size=hid_dim)
+                ...     prediction = paddle.static.nn.fc(x=[fc_1], size=label_dim, activation='softmax')
+                ...     cost = paddle.nn.functional.cross_entropy(
+                ...         input=prediction,
+                ...         label=input_y,
+                ...         reduction='none',
+                ...         use_softmax=False,
+                ...     )
+                ...     sum_cost = paddle.mean(cost)
+                ...     return sum_cost, fc_1, prediction
+
+                >>> input_x = paddle.static.data(name="x", shape=[-1, 32], dtype='float32')
+                >>> input_y = paddle.static.data(name="y", shape=[-1, 1], dtype='int64')
+                >>> cost, fc_1, pred = mlp(input_x, input_y)
+                >>> print("Finished FF")
+                Finished FF
+
+                >>> sgd = paddle.optimizer.Adam(learning_rate=0.01)
+                >>> sgd = paddle.incubate.optimizer.RecomputeOptimizer(sgd)
+                >>> sgd._set_checkpoints([fc_1, pred])
+                >>> params_grads = sgd.backward(
+                ...     cost,
+                ...     startup_program=None,
+                ...     parameter_list=None,
+                ...     no_grad_set=None,
+                ... )
+
+                >>> program = cost.block.program
+                >>> with framework.program_guard(program, None):
+                ...     optimize_ops = sgd.apply_gradients(params_grads)
+
+                >>> print("Finished apply gradients")
+                Finished apply gradients
+        """
+
+        return self._optimizer.apply_gradients(params_grads=params_grads)
+
+    def _create_vars(self, varname):
+        pinned_var_name = unique_name.generate(varname + "@Pinned")
+        fetched_var_name = unique_name.generate(varname + "@Fetch")
+
+        pinned_var = self._main_program.global_block().create_var(
+            name=pinned_var_name,
+            shape=self.checkpoint_shape,
+            dtype=self._main_program.global_block().var(varname).dtype,
+            persistable=False,
+            stop_gradient=True,
+        )
+
+        fetch_var = self._main_program.global_block().create_var(
+            name=fetched_var_name,
+            shape=self.checkpoint_shape,
+            dtype=self._main_program.global_block().var(varname).dtype,
+            persistable=False,
+            stop_gradient=False,
+        )
+
+        return pinned_var_name, fetched_var_name
+
+    def _append_fill_constant_ops(self, startup_program):
+        """
+        add fill_constant_ops to the end of the prog
+
+        we should fill the pinned vars before running the main_prog
+        to instantiate their tensor hold_, which could tell us whether
+        the host memory could hold all the checkpoints from all the
+        GPU devices in this node.
+        """
+        op_role = 0
+        block = startup_program.global_block()
+        fill_constant_vars = self.checkpoint_name2pinned_name.values()
+        OP_ROLE_KEY = core.op_proto_and_checker_maker.kOpRoleAttrName()
+        for varname in fill_constant_vars:
+            var = self._main_program.global_block().var(varname)
+            # NOTE (JZ-LIANG) to pre-allocate the CUDAPinned MEM
+            pinned_var = block.create_var(
+                name=varname,
+                shape=self.checkpoint_shape,
+                dtype=self._main_program.global_block().var(var.name).dtype,
+                persistable=False,
+                stop_gradient=True,
+            )
+            block.append_op(
+                type='fill_constant',
+                outputs={'Out': varname},
+                attrs={
+                    "shape": var.shape,
+                    "dtype": var.dtype,
+                    "value": 0.0,
+                    "place_type": 2,
+                    OP_ROLE_KEY: op_role,
+                },
+            )
+
+    def _insert_async_memcpy_op(
+        self, insert_idx, src_varname, dst_varname, op_role, dst_place_type
+    ):
+        OP_ROLE_KEY = core.op_proto_and_checker_maker.kOpRoleAttrName()
+        self.block._insert_op_without_sync(
+            insert_idx,
+            type='memcpy',
+            inputs={'X': [self._main_program.global_block().var(src_varname)]},
+            outputs={
+                'Out': [self._main_program.global_block().var(dst_varname)]
+            },
+            attrs={"dst_place_type": int(dst_place_type), OP_ROLE_KEY: op_role},
+        )
+
+    def _insert_fetch_op(self, idx, varname):
+        assert varname in self.checkpoint_name2pinned_name, (
+            f"Try to fetch {varname} from Pinned Memory, but it is NOT a checkpoint"
+        )
+
+        pinned_varname = self.checkpoint_name2pinned_name[varname]
+        fetch_varname = self.checkpoint_name2fetch_name[varname]
+        self._insert_async_memcpy_op(idx, pinned_varname, fetch_varname, 1, 1)
+
+    def _insert_offload_op(self, idx, varname):
+        assert varname in self.checkpoint_name2pinned_name, (
+            f"Try to offload {varname} to Pinned Memory, but it is NOT a checkpoint"
+        )
+        pinned_varname = self.checkpoint_name2pinned_name[varname]
+        self._insert_async_memcpy_op(idx, varname, pinned_varname, 0, 2)
+
+    def _insert_sync_op(self, op_idx, checkpoint_name):
+        # single stream offload no need sync
+        pass
+
+    def _record_fetch_op(self, idx):
+        assert len(self.un_fetch_checkpoint_names) > 0, (
+            "Could NOT found checkpoint to fetch"
+        )
+        checkpoint_name = self.un_fetch_checkpoint_names.pop(-1)
+        logging.debug(f"Record fetch [{checkpoint_name}]")
+        self.idx2insertions[idx] = ("fetch", checkpoint_name)
+
+        return checkpoint_name
+
+    def _record_offload_op(self, idx, checkpoint_name):
+        expected_checkpoint_name = self.un_offload_checkpoint_names.pop(0)
+        assert checkpoint_name == expected_checkpoint_name, (
+            f"expected to offload [{expected_checkpoint_name}] but got [{checkpoint_name}]"
+        )
+        logging.debug(f"Record offload [{checkpoint_name}]")
+        self.idx2insertions[idx] = ("offload", checkpoint_name)
+
+    def _record_sync_op(self, idx, checkpoint_name):
+        assert checkpoint_name not in self.synced_checkpoints, (
+            f"Try to sync the checkpoint [{checkpoint_name}] twice"
+        )
+        self.synced_checkpoints.add(checkpoint_name)
+        logging.debug(f"Record offload sync [{checkpoint_name}]")
+        self.idx2insertions[idx] = ("sync", checkpoint_name)
+
+    def _parse_backward(self):
+        self.idx2insertions = {}
+        # don't offload the last checkpoints, to favor throughput
+        self.un_fetch_checkpoint_names = self.sorted_checkpoint_names[:]
+        self.un_fetch_checkpoint_names.pop(-1)
+        need_fetch_checkpoint_names = self.un_fetch_checkpoint_names[:]
+        self.checkpoint_usage_count = {}
+        for checkpoint_name in self.un_fetch_checkpoint_names:
+            self.checkpoint_usage_count[checkpoint_name] = 0
+
+        self.bw_start_op_idx = len(self.block.ops)
+        for idx, op in enumerate(self.block.ops):
+            if int(op.desc.attr("op_role")) == 1:
+                self.bw_start_op_idx = idx
+                break
+
+        assert self.bw_start_op_idx < len(self.block.ops), (
+            "Could NOT found backward op in prog"
+        )
+
+        # fetch second to last checkpoint at the beginning of BW
+        fetched_checkpoint_varname = self._record_fetch_op(self.bw_start_op_idx)
+        last_last_fetch_checkpoint = None
+
+        for i, op in enumerate(self.block.ops[self.bw_start_op_idx :]):
+            idx = self.bw_start_op_idx + i
+            input_vars = op.desc.input_arg_names()
+
+            for input_var in input_vars:
+                if input_var in need_fetch_checkpoint_names:
+                    if input_var not in self.un_fetch_checkpoint_names:
+                        # fetch the offload checkpoint when the first usage of its previous one
+                        if self.checkpoint_usage_count[input_var] == 0:
+                            # TODO (JZ-LIANG) sync memcpy_stream if extra stream for memcpy
+                            second_to_last_fetch_checkpoint = (
+                                fetched_checkpoint_varname
+                            )
+                            # there is NO fetch ahead the first checkpoint
+                            if input_var != self.sorted_checkpoint_names[0]:
+                                fetched_checkpoint_varname = (
+                                    self._record_fetch_op(idx)
+                                )
+
+                        # should check the current used checkpoint is the last fetch one
+                        assert second_to_last_fetch_checkpoint == input_var, (
+                            f"Current recompute segment should use [{second_to_last_fetch_checkpoint}] BUT got [{input_var}]"
+                        )
+                        # rename
+                        self.block.ops[idx]._rename_input(
+                            input_var,
+                            self.checkpoint_name2fetch_name[input_var],
+                        )
+                        self.checkpoint_usage_count[input_var] += 1
+                    else:
+                        raise ValueError(
+                            f"use checkpoint [{input_var}] before fetch in BW"
+                        )
+
+        assert len(self.un_fetch_checkpoint_names) == 0, (
+            f"{self.un_fetch_checkpoint_names} checkpoints have NOT been Recorded"
+        )
+
+    def _update_backward(self):
+        if len(self.idx2insertions) == 0:
+            return
+        total_op = len(self.block.ops)
+        for op_idx in reversed(range(self.bw_start_op_idx, total_op)):
+            if op_idx in self.idx2insertions:
+                operation, checkpoint_name = self.idx2insertions[op_idx]
+                if operation == "fetch":
+                    self._insert_fetch_op(op_idx, checkpoint_name)
+                    logging.debug(f"Insert [{checkpoint_name}] fetch op.")
+                    del self.idx2insertions[op_idx]
+                elif operation == "sync":
+                    self._insert_sync_op(op_idx, checkpoint_name)
+                    logging.debug(f"Sync [{checkpoint_name}] fetch op.")
+        self.block._sync_with_cpp()
+        assert len(self.idx2insertions) == 0, (
+            f"{[ele[1] for ele in self.idx2insertions.values()]} checkpoints left un-Fetched"
+        )
+
+    def _parse_forward(self):
+        self.idx2insertions = {}
+        # don't offload the last checkpoints, faster, less memory saving
+        self.un_offload_checkpoint_names = self.sorted_checkpoint_names[:]
+        last_checkpoint = self.un_offload_checkpoint_names.pop(-1)
+        need_offload_checkpoint_names = self.un_offload_checkpoint_names[:]
+        self.checkpoint_usage_count_and_idx = {}
+        for checkpoint_name in self.un_offload_checkpoint_names:
+            self.checkpoint_usage_count_and_idx[checkpoint_name] = {
+                'count': 0,
+                'idx': -1,
+            }
+        self.synced_checkpoints = set()
+        self.fw_start_op_idx = len(self.block.ops)
+        for idx, op in enumerate(self.block.ops):
+            if int(op.desc.attr("op_role")) == 0:
+                self.fw_start_op_idx = idx
+                break
+
+        assert self.fw_start_op_idx < len(self.block.ops), (
+            "Could NOT found Forward op in prog"
+        )
+        last_offload_checkpoint = None
+
+        for i, op in enumerate(
+            self.block.ops[self.fw_start_op_idx : self.bw_start_op_idx]
+        ):
+            idx = self.fw_start_op_idx + i
+            output_vars = op.desc.output_arg_names()
+            input_vars = op.desc.input_arg_names()
+
+            for output_var in output_vars:
+                if output_var in need_offload_checkpoint_names:
+                    assert len(output_vars) == 1, (
+                        f"checkpoint should be the only Output of a certain op, but [{output_var}] is from [{op}]"
+                    )
+
+                    if output_var in self.un_offload_checkpoint_names:
+                        # insert sync op if last checkpoint has not been sync
+                        if last_offload_checkpoint is not None:
+                            if (
+                                self.checkpoint_usage_count_and_idx[
+                                    last_offload_checkpoint
+                                ]['count']
+                                == 0
+                            ):
+                                self._record_sync_op(
+                                    idx, last_offload_checkpoint
+                                )
+                            else:
+                                last_usage_idx = (
+                                    self.checkpoint_usage_count_and_idx[
+                                        last_offload_checkpoint
+                                    ]['idx']
+                                )
+                                assert last_usage_idx > 0, (
+                                    f"last_usage_idx of checkpoint [{last_offload_checkpoint}] should large than 0"
+                                )
+                                self._record_sync_op(
+                                    last_usage_idx + 1, last_offload_checkpoint
+                                )
+                        # insert offload op after the checkpoint's generation op
+                        self._record_offload_op(idx + 1, output_var)
+                        last_offload_checkpoint = output_var
+                    else:
+                        raise ValueError(
+                            f"There should be just ONE op that output checkpoint [{output_var}]"
+                        )
+                # need to sync the last need to offload checkpoint before the last checkpoint as output op
+                if output_var == last_checkpoint:
+                    assert len(output_vars) == 1, (
+                        f"checkpoint should be the only Output of a certain op, but [{output_var}] is from [{op}]"
+                    )
+                    assert (
+                        last_offload_checkpoint
+                        == self.sorted_checkpoint_names[-2]
+                    ), (
+                        f"the last offload checkpoint before [{last_checkpoint}] is suppose to be [{self.sorted_checkpoint_names[-2]}], but got [{last_offload_checkpoint}]"
+                    )
+                    # sync if last checkpoint has not been sync
+                    if (
+                        self.checkpoint_usage_count_and_idx[
+                            last_offload_checkpoint
+                        ]['idx']
+                        == 0
+                    ):
+                        self._record_sync_op(idx, last_offload_checkpoint)
+                    else:
+                        last_usage_idx = self.checkpoint_usage_count_and_idx[
+                            last_offload_checkpoint
+                        ]['idx']
+                        assert last_usage_idx > 0, (
+                            f"last_usage_idx of checkpoint [{last_offload_checkpoint}] should large than 0"
+                        )
+                        self._record_sync_op(
+                            last_usage_idx + 1, last_offload_checkpoint
+                        )
+            # record checkpoint usage
+            for input_var in input_vars:
+                if input_var in need_offload_checkpoint_names:
+                    assert input_var not in self.synced_checkpoints, (
+                        f"checkpoint [{input_var}] used after sync"
+                    )
+                    self.checkpoint_usage_count_and_idx[input_var]['count'] += 1
+                    self.checkpoint_usage_count_and_idx[input_var]['idx'] = idx
+
+        assert len(self.un_offload_checkpoint_names) == 0, (
+            f"{self.un_fetch_checkpoint_names} checkpoints have NOT been Recorded"
+        )
+        assert len(self.synced_checkpoints) == len(
+            need_offload_checkpoint_names
+        ), (
+            f"{set(need_offload_checkpoint_names) - set(self.synced_checkpoints)} checkpoints have NOT been Recorded"
+        )
+
+    def _update_forward(self):
+        if len(self.idx2insertions) == 0:
+            return
+        for op_idx in reversed(
+            range(self.fw_start_op_idx, self.bw_start_op_idx)
+        ):
+            if op_idx in self.idx2insertions:
+                operation, checkpoint_name = self.idx2insertions[op_idx]
+                if operation == "offload":
+                    self._insert_offload_op(op_idx, checkpoint_name)
+                    logging.debug(f"Insert [{checkpoint_name}] offload op.")
+                    del self.idx2insertions[op_idx]
+                elif operation == "sync":
+                    self._insert_sync_op(op_idx, checkpoint_name)
+                    logging.debug(
+                        f"Insert [{checkpoint_name}] offload_sync op."
+                    )
+                    del self.idx2insertions[op_idx]
+
+        self.block._sync_with_cpp()
+        assert len(self.idx2insertions) == 0, (
+            f"{[ele[1] for ele in self.idx2insertions.values()]} checkpoints left un-Offloaded"
+        )
+
+    def _check_offload_fetch(self):
+        # TODO(JZ-LIANG) the single stream offload need no sync
+        pass
+
+    def _offload(self, loss, startup_program=None):
+        """
+        core steps for recompute offload
+        1. create pinned vars and temp vars
+        2. parse & update Forward pass: offload, sync
+        3. parse & update Backward pass: rename, fetch, sync
+        4. verify the correctness
+        """
+        self._main_program = loss.block.program
+        self.block = loss.block
+        if startup_program is None:
+            startup_program = paddle.static.default_startup_program()
+
+        with program_guard(self._main_program, startup_program):
+            assert len(self.checkpoint_shape) > 0, (
+                f"checkpoints shape {self.checkpoint_shape} should be an non empty list like: [12, 512, 1024]"
+            )
+            assert all(ele > 0 for ele in self.checkpoint_shape), (
+                f"all ele in checkpoints shape {self.checkpoint_shape} should be a determined integer larger than 0"
+            )
+            self.checkpoint_name2pinned_name = {}
+            self.checkpoint_name2fetch_name = {}
+            for checkpoint_varname in self.sorted_checkpoint_names:
+                pinned_var_name, fetch_var_name = self._create_vars(
+                    checkpoint_varname
+                )
+                self.checkpoint_name2pinned_name[checkpoint_varname] = (
+                    pinned_var_name
+                )
+                self.checkpoint_name2fetch_name[checkpoint_varname] = (
+                    fetch_var_name
+                )
+            self._append_fill_constant_ops(startup_program)
+            # TODO (JZ-LIANG) to provide two offload strategy in future
+            # step 2. parse & update FW: rename, offload, sync
+            self._parse_backward()
+            self._update_backward()
+            # step 3. parse & update BW: rename, offload, sync
+            self._parse_forward()
+            self._update_forward()
+            # step 4. verify the correctness
+            self._check_offload_fetch()
+
+    def backward(
+        self,
+        loss,
+        startup_program=None,
+        parameter_list=None,
+        no_grad_set=None,
+        callbacks=None,
+    ):
+        """
+        call append_backward with checkpoints.
+
+        Args:
+            loss (Variable): loss variable to run optimizations.
+            startup_program (Program): startup_program for initializing parameters
+                in `parameter_list`.
+            parameter_list (list): list of Variables or Variable.names to update.
+            no_grad_set (set|None): set of Variables or Variables.names should be ignored.
+            callbacks (list|None): list of callables to run when appending backward
+                operator for one parameter.
+            checkpoints (list): list of Variables as checkpoints
+
+        Examples:
+            .. code-block:: pycon
+
+                >>> import paddle
+
+                >>> paddle.enable_static()
+
+                >>> def mlp(input_x, input_y, hid_dim=128, label_dim=2):
+                ...     fc_1 = paddle.static.nn.fc(x=input_x, size=hid_dim)
+                ...     prediction = paddle.static.nn.fc(x=[fc_1], size=label_dim, activation='softmax')
+                ...     cost = paddle.nn.functional.cross_entropy(
+                ...         input=prediction,
+                ...         label=input_y,
+                ...         reduction='none',
+                ...         use_softmax=False,
+                ...     )
+                ...     sum_cost = paddle.mean(cost)
+                ...     return sum_cost, fc_1, prediction
+
+                >>> input_x = paddle.static.data(name="x", shape=[-1, 32], dtype='float32')
+                >>> input_y = paddle.static.data(name="y", shape=[-1, 1], dtype='int64')
+                >>> cost, fc_1, pred = mlp(input_x, input_y)
+                >>> print("Finished FF")
+                Finished FF
+
+                >>> sgd = paddle.optimizer.Adam(learning_rate=0.01)
+                >>> sgd = paddle.incubate.optimizer.RecomputeOptimizer(sgd)
+                >>> sgd._set_checkpoints([fc_1, pred])
+                >>> params_grads = sgd.backward(
+                ...     cost,
+                ...     startup_program=None,
+                ...     parameter_list=None,
+                ...     no_grad_set=None,
+                ... )
+                >>> print("Finished backward")
+                Finished backward
+        """
+        assert self._checkpoints is not None, (
+            "You should call _set_checkpoints first"
+        )
+
+        if in_dygraph_mode():
+            raise NotImplementedError(
+                "DyGraph current does not support recompute"
+            )
+
+        self._dtype = loss.dtype
+        program = loss.block.program
+        with program_guard(program, startup_program):
+            checkpoint_vars = []
+            for ckpt in self._checkpoints:
+                if isinstance(ckpt, Variable):
+                    checkpoint_vars.append(ckpt)
+                else:
+                    checkpoint_vars.append(loss.block.var(ckpt))
+
+            # allow return to non-recompute when checkpoints is empty
+            if len(checkpoint_vars) > 0:
+                params_grads, sorted_checkpoint_names = append_backward(
+                    loss,
+                    parameter_list,
+                    no_grad_set,
+                    checkpoints=checkpoint_vars,
+                )
+            else:
+                params_grads = append_backward(
+                    loss,
+                    parameter_list,
+                    no_grad_set,
+                    checkpoints=checkpoint_vars,
+                )
+
+        if self.enable_offload:
+            self.sorted_checkpoint_names = sorted_checkpoint_names
+            self._offload(loss, startup_program=startup_program)
+
+        return params_grads
+
+    def apply_optimize(self, loss, startup_program, params_grads):
+        """
+        call the apply_optimize function of self._optimizer
+        Args:
+            loss (Variable): loss variable to run optimizations.
+            startup_program (Program): startup_program for initializing parameters
+                in `parameter_list`.
+            params_grads (list): list of (param, grad) pair to do optimization.
+        Examples:
+            .. code-block:: pycon
+
+                >>> import paddle
+
+                >>> paddle.enable_static()
+
+                >>> def mlp(input_x, input_y, hid_dim=128, label_dim=2):
+                ...     fc_1 = paddle.static.nn.fc(x=input_x, size=hid_dim)
+                ...     prediction = paddle.static.nn.fc(x=[fc_1], size=label_dim, activation='softmax')
+                ...     cost = paddle.nn.functional.cross_entropy(
+                ...         input=prediction,
+                ...         label=input_y,
+                ...         reduction='none',
+                ...         use_softmax=False,
+                ...     )
+                ...     sum_cost = paddle.mean(cost)
+                ...     return sum_cost, fc_1, prediction
+
+                >>> input_x = paddle.static.data(name="x", shape=[-1, 32], dtype='float32')
+                >>> input_y = paddle.static.data(name="y", shape=[-1, 1], dtype='int64')
+                >>> cost, fc_1, pred = mlp(input_x, input_y)
+                >>> print("Finished FF")
+                Finished FF
+
+                >>> sgd = paddle.optimizer.Adam(learning_rate=0.01)
+                >>> sgd = paddle.incubate.optimizer.RecomputeOptimizer(sgd)
+                >>> sgd._set_checkpoints([fc_1, pred])
+                >>> params_grads = sgd.backward(
+                ...     cost,
+                ...     startup_program=None,
+                ...     parameter_list=None,
+                ...     no_grad_set=None,
+                ... )
+
+                >>> optimize_ops = sgd.apply_optimize(
+                ...     cost,
+                ...     startup_program=None,
+                ...     params_grads=params_grads,
+                ... )
+
+                >>> print("Finished apply_optimize")
+                Finished apply_optimize
+        """
+
+        func = (
+            self._optimizer.apply_optimize
+            if hasattr(self._optimizer, 'apply_optimize')
+            else self._optimizer._apply_optimize
+        )
+        return func(
+            loss, startup_program=startup_program, params_grads=params_grads
+        )
+
+    def minimize(
+        self, loss, startup_program=None, parameter_list=None, no_grad_set=None
+    ):
+        assert isinstance(loss, Variable), "The loss should be an Variable."
+        assert self._checkpoints is not None, (
+            "You should call _set_checkpoints first"
+        )
+        if in_dygraph_mode():
+            raise NotImplementedError(
+                "DyGraph current does not support recompute"
+            )
+        params_grads = self.backward(
+            loss,
+            startup_program=startup_program,
+            parameter_list=parameter_list,
+            no_grad_set=no_grad_set,
+        )
+
+        optimize_ops = self.apply_optimize(
+            loss, startup_program=startup_program, params_grads=params_grads
+        )
+
+        return optimize_ops, params_grads