mindspore 1.10.0__cp39-cp39-win_amd64.whl → 2.0.0rc1__cp39-cp39-win_amd64.whl

mindspore/nn/optim/adam.py

@@ -13,31 +13,246 @@
 # limitations under the License.
 # ============================================================================
 """adam"""
-from __future__ import absolute_import
+from __future__ import absolute_import, division
 
 import numpy as np
 
 from mindspore import context
 from mindspore.common import dtype as mstype
 from mindspore.common.initializer import initializer
-from mindspore.common.api import ms_function
+from mindspore.common.api import jit
 from mindspore.ops import operations as P
 from mindspore.ops import composite as C
 from mindspore.ops import functional as F
 from mindspore.common.parameter import Parameter
 from mindspore.common.tensor import Tensor
-from mindspore._checkparam import Validator as validator
-from mindspore._checkparam import Rel
+from mindspore import _checkparam as validator
 from mindspore.nn.optim.optimizer import Optimizer
 from mindspore.nn.optim.optimizer import opt_init_args_register
 from mindspore.nn.optim._dist_optimizer_registry import _register_dist_optimizer
+from mindspore.common._decorator import deprecated
 
 _adam_opt = C.MultitypeFuncGraph("adam_opt")
 _fused_adam_weight_decay = C.MultitypeFuncGraph("fused_adam_weight_decay")
+_lazy_adam_opt = C.MultitypeFuncGraph("lazy_adam_opt")
 _scaler_one = Tensor(1, mstype.int32)
 _scaler_ten = Tensor(10, mstype.float32)
 
 
+@_lazy_adam_opt.register("Function", "Function", "Function", "Function", "Bool", "Bool", "Bool", "Tensor", "Tensor",
+                         "Tensor", "Tensor", "Tensor", "Tensor", "RowTensor", "Tensor", "Tensor", "Tensor", "Bool",
+                         "Bool", "Function", "Bool", "Function", "Bool")
+def _run_lazy_opt_with_sparse_dist(opt, sparse_opt, push, pull, use_locking, use_nesterov, target, beta1_power,
+                                   beta2_power, beta1, beta2, eps, lr, gradient, params, m, v, ps_parameter,
+                                   cache_enable, distributed_opt, use_flag, distributed_sparse_opt, use_sparse_flag):
+    """Apply sparse lazy adam optimizer to the weight parameter when the gradient is sparse."""
+    success = True
+    indices = gradient.indices
+    values = gradient.values
+    if use_sparse_flag:
+        success = F.depend(success, distributed_sparse_opt(params, m, v, beta1_power, beta2_power, lr, beta1, beta2,
+                                                           eps, values, indices))
+        return success
+    if ps_parameter and not cache_enable:
+        op_shape = P.Shape()
+        shapes = (op_shape(params), op_shape(m), op_shape(v),
+                  op_shape(beta1_power), op_shape(beta2_power), op_shape(lr), op_shape(beta1),
+                  op_shape(beta2), op_shape(eps), op_shape(values), op_shape(indices))
+        success = F.depend(success, pull(push((beta1_power, beta2_power, lr, beta1, beta2,
+                                               eps, values, indices), shapes), params))
+        return success
+
+    if not target:
+        success = F.depend(success, sparse_opt(params, m, v, beta1_power, beta2_power, lr, beta1, beta2,
+                                               eps, values, indices))
+    else:
+        op_gather = P.Gather()
+        op_sqrt = P.Sqrt()
+        scatter_add = P.ScatterAdd(use_locking)
+        scatter_update = P.ScatterUpdate(use_locking)
+
+        m_slice = op_gather(m, indices, 0)
+        v_slice = op_gather(v, indices, 0)
+
+        next_m = m_slice * beta1 + values * (1 - beta1)
+        next_v = v_slice * beta2 + values * values * (1 - beta2)
+
+        lr_t = lr * op_sqrt(1 - beta2_power) / (1 - beta1_power)
+
+        if use_nesterov:
+            m_temp = beta1 * next_m + values * (1 - beta1)
+            param_update = m_temp / (op_sqrt(next_v) + eps)
+        else:
+            param_update = next_m / (op_sqrt(next_v) + eps)
+
+        success = F.depend(success, scatter_add(params, indices, - lr_t * param_update))
+        success = F.depend(success, scatter_update(m, indices, next_m))
+        success = F.depend(success, scatter_update(v, indices, next_v))
+
+    return success
+
+
+@_lazy_adam_opt.register("Function", "Function", "Function", "Function", "Bool", "Bool", "Bool", "Tensor", "Tensor",
+                         "Tensor", "Tensor", "Tensor", "Tensor", "MapTensor", "MapTensor", "MapTensor", "MapTensor",
+                         "Bool", "Bool", "Function", "Bool", "Function", "Bool")
+def _run_map_tensor_lazy_opt_with_sparse_dist(opt, sparse_opt, push, pull, use_locking, use_nesterov, target,
+                                              beta1_power, beta2_power, beta1, beta2, eps, lr, gradient, params, m, v,
+                                              ps_parameter, cache_enable, distributed_opt, use_flag,
+                                              distributed_sparse_opt, use_sparse_flag):
+    """Apply sparse lazy adam optimizer to the weight parameter when the gradient is sparse."""
+    success = True
+    indices, values = gradient.get_data()
+    if use_sparse_flag:
+        # PS Mode.
+        success = F.depend(success, distributed_sparse_opt(params, m, v, beta1_power, beta2_power, lr, beta1, beta2,
+                                                           eps, values, indices))
+    else:
+        # PS Cache mode.
+        op_sqrt = P.Sqrt()
+
+        m_slice = m.get(indices)
+        v_slice = v.get(indices)
+
+        next_m = m_slice * beta1 + values * (1 - beta1)
+        next_v = v_slice * beta2 + values * values * (1 - beta2)
+
+        lr_t = lr * op_sqrt(1 - beta2_power) / (1 - beta1_power)
+
+        if use_nesterov:
+            m_temp = beta1 * next_m + values * (1 - beta1)
+            param_update = m_temp / (op_sqrt(next_v) + eps)
+        else:
+            param_update = next_m / (op_sqrt(next_v) + eps)
+
+        params_need_update = params.get(indices)
+        params.put(indices, params_need_update - lr_t * param_update)
+        m.put(indices, next_m)
+        v.put(indices, next_v)
+
+    return success
+
+
+@_lazy_adam_opt.register("Function", "Function", "Function", "Function", "Bool", "Bool", "Bool", "Tensor", "Tensor",
+                         "Tensor", "Tensor", "Tensor", "Tensor", "Tensor", "Tensor", "Tensor", "Tensor", "Bool", "Bool",
+                         "Function", "Bool", "Function", "Bool")
+def _run_lazy_opt_with_one_number_dist(opt, sparse_opt, push, pull, use_locking, use_nesterov, target,
+                                       beta1_power, beta2_power, beta1, beta2, eps, lr, gradient, params, moment1,
+                                       moment2, ps_parameter, cache_enable, distributed_opt, use_flag,
+                                       distributed_sparse_opt, use_sparse_flag):
+    """Apply lazy adam optimizer to the weight parameter using Tensor."""
+    success = True
+    if use_flag:
+        success = F.depend(success, distributed_opt(params, moment1, moment2, beta1_power, beta2_power, lr, beta1,
+                                                    beta2, eps, gradient))
+    elif ps_parameter and not cache_enable:
+        op_shape = P.Shape()
+        success = F.depend(success, pull(push((beta1_power, beta2_power, lr, beta1, beta2, eps, gradient),
+                                              (op_shape(params), op_shape(moment1), op_shape(moment2))), params))
+    else:
+        success = F.depend(success, opt(params, moment1, moment2, beta1_power, beta2_power, lr, beta1, beta2,
+                                        eps, gradient))
+    return success
+
+
+@_lazy_adam_opt.register("Function", "Function", "Function", "Function", "Bool", "Bool", "Bool", "Tensor", "Tensor",
+                         "Tensor", "Tensor", "Tensor", "Tensor", "RowTensor", "Tensor", "Tensor", "Tensor", "Bool",
+                         "Bool")
+def _run_lazy_opt_with_sparse(opt, sparse_opt, push, pull, use_locking, use_nesterov, target, beta1_power, beta2_power,
+                              beta1, beta2, eps, lr, gradient, params, m, v, ps_parameter, cache_enable):
+    """Apply sparse lazy adam optimizer to the weight parameter when the gradient is sparse."""
+    success = True
+    indices = gradient.indices
+    values = gradient.values
+    if ps_parameter and not cache_enable:
+        op_shape = P.Shape()
+        shapes = (op_shape(params), op_shape(m), op_shape(v),
+                  op_shape(beta1_power), op_shape(beta2_power), op_shape(lr), op_shape(beta1),
+                  op_shape(beta2), op_shape(eps), op_shape(values), op_shape(indices))
+        success = F.depend(success, pull(push((beta1_power, beta2_power, lr, beta1, beta2,
+                                               eps, values, indices), shapes), params))
+        return success
+
+    if not target:
+        success = F.depend(success, sparse_opt(params, m, v, beta1_power, beta2_power, lr, beta1, beta2,
+                                               eps, values, indices))
+    else:
+        op_gather = P.Gather()
+        op_sqrt = P.Sqrt()
+        scatter_add = P.ScatterAdd(use_locking)
+        scatter_update = P.ScatterUpdate(use_locking)
+
+        m_slice = op_gather(m, indices, 0)
+        v_slice = op_gather(v, indices, 0)
+
+        next_m = m_slice * beta1 + values * (1 - beta1)
+        next_v = v_slice * beta2 + values * values * (1 - beta2)
+
+        lr_t = lr * op_sqrt(1 - beta2_power) / (1 - beta1_power)
+
+        if use_nesterov:
+            m_temp = beta1 * next_m + values * (1 - beta1)
+            param_update = m_temp / (op_sqrt(next_v) + eps)
+        else:
+            param_update = next_m / (op_sqrt(next_v) + eps)
+
+        success = F.depend(success, scatter_add(params, indices, - lr_t * param_update))
+        success = F.depend(success, scatter_update(m, indices, next_m))
+        success = F.depend(success, scatter_update(v, indices, next_v))
+
+    return success
+
+
+@_lazy_adam_opt.register("Function", "Function", "Function", "Function", "Bool", "Bool", "Bool", "Tensor", "Tensor",
+                         "Tensor", "Tensor", "Tensor", "Tensor", "MapTensor", "MapTensor", "MapTensor", "MapTensor",
+                         "Bool", "Bool")
+def _run_map_tensor_lazy_opt_with_sparse(opt, sparse_opt, push, pull, use_locking, use_nesterov, target, beta1_power,
+                                         beta2_power, beta1, beta2, eps, lr, gradient, params, m, v, ps_parameter,
+                                         cache_enable):
+    """Apply sparse lazy adam optimizer to the weight parameter when the gradient is sparse(MapTensor)."""
+    success = True
+    indices, values = gradient.get_data()
+
+    op_sqrt = P.Sqrt()
+
+    m_slice = m.get(indices)
+    v_slice = v.get(indices)
+
+    next_m = m_slice * beta1 + values * (1 - beta1)
+    next_v = v_slice * beta2 + values * values * (1 - beta2)
+
+    lr_t = lr * op_sqrt(1 - beta2_power) / (1 - beta1_power)
+
+    if use_nesterov:
+        m_temp = beta1 * next_m + values * (1 - beta1)
+        param_update = m_temp / (op_sqrt(next_v) + eps)
+    else:
+        param_update = next_m / (op_sqrt(next_v) + eps)
+
+    params_need_update = params.get(indices)
+    params.put(indices, params_need_update - lr_t * param_update)
+    m.put(indices, next_m)
+    v.put(indices, next_v)
+
+    return success
+
+
+@_lazy_adam_opt.register("Function", "Function", "Function", "Function", "Bool", "Bool", "Bool", "Tensor", "Tensor",
+                         "Tensor", "Tensor", "Tensor", "Tensor", "Tensor", "Tensor", "Tensor", "Tensor", "Bool", "Bool")
+def _run_lazy_opt_with_one_number(opt, sparse_opt, push, pull, use_locking, use_nesterov, target, beta1_power,
+                                  beta2_power, beta1, beta2, eps, lr, gradient, params, moment1, moment2, ps_parameter,
+                                  cache_enable):
+    """Apply lazy adam optimizer to the weight parameter using Tensor."""
+    success = True
+    if ps_parameter and not cache_enable:
+        op_shape = P.Shape()
+        success = F.depend(success, pull(push((beta1_power, beta2_power, lr, beta1, beta2, eps, gradient),
+                                              (op_shape(params), op_shape(moment1), op_shape(moment2))), params))
+    else:
+        success = F.depend(success, opt(params, moment1, moment2, beta1_power, beta2_power, lr, beta1, beta2,
+                                        eps, gradient))
+    return success
+
+
 @_adam_opt.register("Tensor", "Tensor", "Tensor", "Tensor", "Tensor", "Tensor", "Tensor", "Tensor",
                     "Tensor", "Bool", "Bool")
 def _update_run_op(beta1, beta2, eps, lr, weight_decay, param, m, v, gradient, decay_flag, optim_filter):
@@ -184,18 +399,14 @@ def _run_opt_with_one_number_dist(opt, sparse_opt, push, pull, use_locking, use_
 
 
 @_adam_opt.register("Function", "Function", "Function", "Function",
-                    "Bool", "Bool", "Bool", "Bool",
+                    "Bool", "Bool", "Bool",
                     "Tensor", "Tensor", "Tensor", "Tensor", "Tensor", "Tensor",
-                    "RowTensor", "Tensor", "Tensor", "Tensor", "Tensor", "Bool", "Bool")
+                    "RowTensor", "Tensor", "Tensor", "Tensor", "Bool", "Bool")
 def _run_opt_with_sparse(opt, sparse_opt, push, pull,
-                         use_locking, use_nesterov, use_amsgrad, target,
+                         use_locking, use_nesterov, target,
                          beta1_power, beta2_power, beta1, beta2, eps, lr,
-                         gradient, param, m, v, vhat, ps_parameter, cache_enable):
+                         gradient, param, m, v, ps_parameter, cache_enable):
     """Apply sparse adam optimizer to the weight parameter when the gradient is sparse."""
-    if use_amsgrad:
-        raise Exception("""Adam with amsgrad is currently not supported when the gradients are sparse!
-                        Please set use_amsgrad=False for sparse gradients.""")
-
     success = True
     indices = gradient.indices
     values = gradient.values
@@ -253,30 +464,42 @@ def _run_opt_with_sparse(opt, sparse_opt, push, pull,
 
 
 @_adam_opt.register("Function", "Function", "Function", "Function",
-                    "Bool", "Bool", "Bool", "Bool",
+                    "Bool", "Bool", "Bool",
                     "Tensor", "Tensor", "Tensor", "Tensor", "Tensor", "Tensor",
-                    "Tensor", "Tensor", "Tensor", "Tensor", "Tensor", "Bool", "Bool")
+                    "Tensor", "Tensor", "Tensor", "Tensor", "Bool", "Bool")
 def _run_opt_with_one_number(opt, sparse_opt, push, pull,
-                             use_locking, use_nesterov, use_amsgrad, target,
+                             use_locking, use_nesterov, target,
                              beta1_power, beta2_power, beta1, beta2, eps, lr,
-                             gradient, param, moment1, moment2, vhat, ps_parameter, cache_enable):
+                             gradient, param, moment1, moment2, ps_parameter, cache_enable):
     """Apply adam optimizer to the weight parameter using Tensor."""
     success = True
     if ps_parameter and not cache_enable:
         op_shape = P.Shape()
-        if use_amsgrad:
-            success = F.depend(success, pull(push((beta1_power, beta2_power, lr, gradient),
-                                                  (op_shape(param), op_shape(moment1), op_shape(moment2),
-                                                   op_shape(vhat))), param))
-        else:
-            success = F.depend(success, pull(push((beta1_power, beta2_power, lr, beta1, beta2, eps, gradient),
-                                                  (op_shape(param), op_shape(moment1), op_shape(moment2))), param))
+        success = F.depend(success, pull(push((beta1_power, beta2_power, lr, beta1, beta2, eps, gradient),
+                                              (op_shape(param), op_shape(moment1), op_shape(moment2))), param))
     else:
-        if use_amsgrad:
-            success = F.depend(success, opt(param, moment1, moment2, vhat, beta1_power, beta2_power, lr, gradient))
-        else:
-            success = F.depend(success, opt(param, moment1, moment2, beta1_power, beta2_power, lr, beta1, beta2,
-                                            eps, gradient))
+        success = F.depend(success, opt(param, moment1, moment2, beta1_power, beta2_power, lr, beta1, beta2,
+                                        eps, gradient))
+    return success
+
+
+@_adam_opt.register("Function", "Function", "Function", "Function",
+                    "Bool", "Bool", "Bool",
+                    "Tensor", "Tensor", "Tensor", "Tensor", "Tensor", "Tensor",
+                    "Tensor", "Tensor", "Tensor", "Tensor", "Tensor", "Bool", "Bool")
+def _run_opt_with_one_number_use_amsgrad(opt, sparse_opt, push, pull,
+                                         use_locking, use_nesterov, target,
+                                         beta1_power, beta2_power, beta1, beta2, eps, lr,
+                                         gradient, param, moment1, moment2, vhat, ps_parameter, cache_enable):
+    """Apply adam optimizer to the weight parameter using Tensor and use amsgrad."""
+    success = True
+    if ps_parameter and not cache_enable:
+        op_shape = P.Shape()
+        success = F.depend(success, pull(push((beta1_power, beta2_power, lr, gradient),
+                                              (op_shape(param), op_shape(moment1), op_shape(moment2),
+                                               op_shape(vhat))), param))
+    else:
+        success = F.depend(success, opt(param, moment1, moment2, vhat, beta1_power, beta2_power, lr, gradient))
     return success
 
 
@@ -308,8 +531,8 @@ def _check_param_value(beta1, beta2, eps, prim_name):
     validator.check_value_type("beta1", beta1, [float], prim_name)
     validator.check_value_type("beta2", beta2, [float], prim_name)
     validator.check_value_type("eps", eps, [float], prim_name)
-    validator.check_float_range(beta1, 0.0, 1.0, Rel.INC_NEITHER, "beta1", prim_name)
-    validator.check_float_range(beta2, 0.0, 1.0, Rel.INC_NEITHER, "beta2", prim_name)
+    validator.check_float_range(beta1, 0.0, 1.0, validator.INC_NEITHER, "beta1", prim_name)
+    validator.check_float_range(beta2, 0.0, 1.0, validator.INC_NEITHER, "beta2", prim_name)
     validator.check_positive_float(eps, "eps", prim_name)
 
 
@@ -374,6 +597,16 @@ class Adam(Optimizer):
         parameters are grouped, each group can set `weight_decay`. If not, the `weight_decay` in optimizer will be
         applied.
 
+        When using Adam with use_lazy=True:
+
+        Please note, the optimizer only updates the current index position of the network parameters
+        when the gradient is sparse. The sparse behavior is not equivalent to the original Adam algorithm.
+        If you want to execute a sparse policy, target needs to be set to CPU.
+
+        When using Adam with use_offload=True:
+
+        This optimizer only supports `GRAPH_MODE`.
+
     Args:
         params (Union[list[Parameter], list[dict]]): Must be list of `Parameter` or list of `dict`. When the
             `params` is a list of `dict`, the string "params", "lr", "weight_decay", "grad_centralization" and
@@ -446,6 +679,17 @@ class Adam(Optimizer):
             `FixedLossScaleManager`. Refer to class :class:`mindspore.amp.FixedLossScaleManager` for more details.
             Default: 1.0.
 
+        kwargs:
+
+            - use_lazy (bool): Whether to use Lazy Adam algorithm. Default: False.
+              If true, apply lazy adam algorithm.
+              If false, apply normal adam algorithm.
+
+            - use_offload (bool): Whether to offload adam optimizer to host CPU and keep parameters being updated on
+              the device in order to minimize the memory cost. Default: False.
+              If true, apply offload adam.
+              If false, apply normal adam.
+
     Inputs:
         - **gradients** (tuple[Tensor]) - The gradients of `params`, the shape is the same as `params`.
 
@@ -457,10 +701,13 @@ class Adam(Optimizer):
         TypeError: If element of `parameters` is neither Parameter nor dict.
         TypeError: If `beta1`, `beta2`, `eps` or `loss_scale` is not a float.
         TypeError: If `weight_decay` is neither float nor int.
-        TypeError: If `use_locking`, `use_nesterov` or `use_amsgrad` is not a bool.
+        TypeError: If `use_locking`, `use_nesterov`, `use_amsgrad`, `use_lazy` or `use_offload` is not a bool.
         ValueError: If `loss_scale` or `eps` is less than or equal to 0.
         ValueError: If `beta1`, `beta2` is not in range (0.0, 1.0).
         ValueError: If `weight_decay` is less than 0.
+        ValueError: If `use_lazy` and `use_offload` are both true.
+        ValueError: If `use_amsgrad` is true and (`use_lazy` or `use_offload` is true).
+        ValueError: If `use_amsgrad` while using distributed training.
 
     Supported Platforms:
         ``Ascend`` ``GPU``  ``CPU``
@@ -479,7 +726,7 @@ class Adam(Optimizer):
         >>> group_params = [{'params': conv_params, 'weight_decay': 0.01, 'grad_centralization':True},
         ...                 {'params': no_conv_params, 'lr': 0.01},
         ...                 {'order_params': net.trainable_params()}]
-        >>> optim = nn.Adam(group_params, learning_rate=0.1, weight_decay=0.0)
+        >>> optim = nn.Adam(group_params, learning_rate=0.1, weight_decay=0.0, use_lazy=False, use_offload=False)
         >>> # The conv_params's parameters will use default learning rate of 0.1 and weight decay of 0.01 and grad
         >>> # centralization of True.
         >>> # The no_conv_params's parameters will use learning rate of 0.01 and default weight decay of 0.0 and grad
@@ -492,50 +739,174 @@ class Adam(Optimizer):
 
     @opt_init_args_register
     def __init__(self, params, learning_rate=1e-3, beta1=0.9, beta2=0.999, eps=1e-8, use_locking=False,
-                 use_nesterov=False, weight_decay=0.0, loss_scale=1.0, use_amsgrad=False):
+                 use_nesterov=False, weight_decay=0.0, loss_scale=1.0, use_amsgrad=False, **kwargs):
         super(Adam, self).__init__(learning_rate, params, weight_decay, loss_scale)
+        use_lazy = kwargs.get('use_lazy', False)
+        use_offload = kwargs.get('use_offload', False)
         _check_param_value(beta1, beta2, eps, self.cls_name)
         validator.check_value_type("use_locking", use_locking, [bool], self.cls_name)
         validator.check_value_type("use_nesterov", use_nesterov, [bool], self.cls_name)
         validator.check_value_type("use_amsgrad", use_amsgrad, [bool], self.cls_name)
+        validator.check_value_type("use_lazy", use_lazy, [bool], self.cls_name)
+        validator.check_value_type("use_offload", use_offload, [bool], self.cls_name)
+
+        if use_lazy and use_offload:
+            raise ValueError(f"For 'Adam', 'use_lazy' and 'use_offload' can not both be True."
+                             f"But got use_lazy={use_lazy}, use_offload={use_offload}.")
+
+        if use_amsgrad and (use_lazy or use_offload):
+            raise ValueError(f"For lazy Adam and Adam with offload, there is no parameter named 'use_amsgrad'."
+                             f"but got 'use_amsgrad'={use_amsgrad}.")
 
         self.beta1 = Tensor(beta1, mstype.float32)
         self.beta2 = Tensor(beta2, mstype.float32)
-        self.beta1_power = Parameter(initializer(1, [1], mstype.float32), name="beta1_power")
-        self.beta2_power = Parameter(initializer(1, [1], mstype.float32), name="beta2_power")
+        self.beta1_power = Parameter(initializer(1, (), mstype.float32), name="beta1_power")
+        self.beta2_power = Parameter(initializer(1, (), mstype.float32), name="beta2_power")
         self.eps = Tensor(eps, mstype.float32)
         self.use_nesterov = use_nesterov
         self.use_locking = use_locking
         self.use_amsgrad = use_amsgrad
+        self.use_lazy = use_lazy
+        self.use_offload = use_offload
         self.moment1 = self._parameters.clone(prefix="moment1", init='zeros')
         self.moment2 = self._parameters.clone(prefix="moment2", init='zeros')
-        self.vhat = self._parameters.clone(prefix="vhat", init='zeros')
-
-        self._is_device = True
         if use_amsgrad:
-            self.opt = P.ApplyAdamWithAmsgrad(beta1, beta2, eps, use_locking)
-        else:
+            self.vhat = self._parameters.clone(prefix="vhat", init='zeros')
+
+        if use_offload:
+            self.opt = P.AdamNoUpdateParam(use_locking, use_nesterov)
+            self.opt.set_device("CPU")
+
+        elif use_lazy:
+            self._is_device = True
             self.opt = P.Adam(use_locking, use_nesterov)
-        self.sparse_opt = P.FusedSparseAdam(use_locking, use_nesterov)
-        self.sparse_opt.add_prim_attr("primitive_target", "CPU")
-        self._ps_pull = P.Pull()
-        if use_amsgrad:
-            self._ps_push = P.Push("ApplyAdamWithAmsgrad", [0, 1, 2, 3])
-        else:
+            self.sparse_opt = P.FusedSparseLazyAdam(use_locking, use_nesterov)
+            self.sparse_opt.set_device("CPU")
+            self._ps_pull = P.Pull()
             self._ps_push = P.Push("Adam", [0, 1, 2])
             self._ps_push.add_prim_attr("use_nesterov", use_nesterov)
+            self._init_distributed_opts(use_locking, use_nesterov)
 
-        self._init_distributed_opts(use_locking, use_nesterov)
+        else:
+            self._is_device = True
+            if use_amsgrad:
+                self.opt = P.ApplyAdamWithAmsgrad(beta1, beta2, eps, use_locking)
+            else:
+                self.opt = P.Adam(use_locking, use_nesterov)
+            self.sparse_opt = P.FusedSparseAdam(use_locking, use_nesterov)
+            self.sparse_opt.set_device("CPU")
+            self._ps_pull = P.Pull()
+            if use_amsgrad:
+                self._ps_push = P.Push("ApplyAdamWithAmsgrad", [0, 1, 2, 3])
+            else:
+                self._ps_push = P.Push("Adam", [0, 1, 2])
+                self._ps_push.add_prim_attr("use_nesterov", use_nesterov)
+
+            self._init_distributed_opts(use_locking, use_nesterov)
+
+    def _apply_adam(self, params, beta1_power, beta2_power, moment1, moment2, lr, gradients):
+        """Execute Adam optimizer and its variants."""
+        if self.use_offload:
+            if self.is_group_lr:
+                success = self.map_reverse(F.partial(_adam_opt, self.opt, beta1_power, beta2_power, self.beta1,
+                                                     self.beta2, self.eps), lr, gradients, params, moment1, moment2)
+            else:
+                success = self.map_reverse(F.partial(_adam_opt, self.opt, beta1_power, beta2_power, self.beta1,
+                                                     self.beta2, self.eps, lr), gradients, params, moment1, moment2)
+        # Lazy adam or normal adam
+        else:
+            if self.use_dist_optimizer:
+                if self.use_dist_optimizer and self.use_amsgrad:
+                    raise ValueError(f"Adam with amsgrad is currently not supporting distributed training!"
+                                     f"Please set use_amsgrad=False for distributed training.")
+                if self.is_group_lr:
+                    if self.use_lazy:
+                        success = self.map_reverse(F.partial(_lazy_adam_opt, self.opt, self.sparse_opt,
+                                                             self._ps_push, self._ps_pull, self.use_locking,
+                                                             self.use_nesterov,
+                                                             self._is_device, beta1_power, beta2_power,
+                                                             self.beta1, self.beta2, self.eps),
+                                                   lr, gradients, self._parameters, self.moment1, self.moment2,
+                                                   self.ps_parameters, self.cache_enable, self.dense_lazyadam_opts,
+                                                   self.use_dense_opt_flags, self.sparse_lazyadam_opts,
+                                                   self.use_sparse_opt_flags)
+                    # Normal Adam
+                    else:
+                        success = self.map_(F.partial(_adam_opt, self.opt, self.sparse_opt, self._ps_push,
+                                                      self._ps_pull, self.use_locking, self.use_nesterov,
+                                                      self._is_device, beta1_power, beta2_power, self.beta1, self.beta2,
+                                                      self.eps), lr, gradients, params, moment1, moment2,
+                                            self.ps_parameters, self.cache_enable, self.dense_adam_opts,
+                                            self.use_dense_opt_flags, self.sparse_adam_opts, self.use_sparse_opt_flags)
+                else:
+                    if self.use_lazy:
+                        success = self.map_reverse(F.partial(_lazy_adam_opt, self.opt, self.sparse_opt, self._ps_push,
+                                                             self._ps_pull, self.use_locking, self.use_nesterov,
+                                                             self._is_device, beta1_power, beta2_power, self.beta1,
+                                                             self.beta2, self.eps, lr), gradients, self._parameters,
+                                                   self.moment1, self.moment2, self.ps_parameters, self.cache_enable,
+                                                   self.dense_lazyadam_opts, self.use_dense_opt_flags,
+                                                   self.sparse_lazyadam_opts, self.use_sparse_opt_flags)
+                    else:
+                        success = self.map_(F.partial(_adam_opt, self.opt, self.sparse_opt, self._ps_push,
+                                                      self._ps_pull, self.use_locking, self.use_nesterov,
+                                                      self._is_device, beta1_power, beta2_power, self.beta1, self.beta2,
+                                                      self.eps, lr), gradients, params, moment1, moment2,
+                                            self.ps_parameters, self.cache_enable, self.dense_adam_opts,
+                                            self.use_dense_opt_flags, self.sparse_adam_opts, self.use_sparse_opt_flags)
+            else:
+                if self.is_group_lr:
+                    if self.use_lazy:
+                        success = self.map_(F.partial(_lazy_adam_opt, self.opt, self.sparse_opt, self._ps_push,
+                                                      self._ps_pull, self.use_locking, self.use_nesterov,
+                                                      self._is_device, beta1_power, beta2_power, self.beta1, self.beta2,
+                                                      self.eps), lr, gradients, params, moment1, moment2,
+                                            self.ps_parameters, self.cache_enable)
+                    else:
+                        if self.use_amsgrad:
+                            success = self.map_(F.partial(_adam_opt, self.opt, self.sparse_opt, self._ps_push,
+                                                          self._ps_pull, self.use_locking, self.use_nesterov,
+                                                          self._is_device, beta1_power, beta2_power,
+                                                          self.beta1, self.beta2, self.eps), lr, gradients, params,
+                                                moment1, moment2, self.vhat, self.ps_parameters, self.cache_enable)
+                        else:
+                            success = self.map_(F.partial(_adam_opt, self.opt, self.sparse_opt, self._ps_push,
+                                                          self._ps_pull, self.use_locking, self.use_nesterov,
+                                                          self._is_device, beta1_power, beta2_power,
+                                                          self.beta1, self.beta2, self.eps), lr, gradients, params,
+                                                moment1, moment2, self.ps_parameters, self.cache_enable)
+                else:
+                    if self.use_lazy:
+                        success = self.map_(F.partial(_lazy_adam_opt, self.opt, self.sparse_opt, self._ps_push,
+                                                      self._ps_pull, self.use_locking, self.use_nesterov,
+                                                      self._is_device, beta1_power, beta2_power, self.beta1, self.beta2,
+                                                      self.eps, lr), gradients, params, moment1, moment2,
+                                            self.ps_parameters, self.cache_enable)
+                    else:
+                        if self.use_amsgrad:
+                            success = self.map_(F.partial(_adam_opt, self.opt, self.sparse_opt, self._ps_push,
+                                                          self._ps_pull, self.use_locking, self.use_nesterov,
+                                                          self._is_device, beta1_power, beta2_power,
+                                                          self.beta1, self.beta2, self.eps, lr), gradients, params,
+                                                moment1, moment2, self.vhat, self.ps_parameters, self.cache_enable)
+                        else:
+                            success = self.map_(F.partial(_adam_opt, self.opt, self.sparse_opt, self._ps_push,
+                                                          self._ps_pull, self.use_locking, self.use_nesterov,
+                                                          self._is_device, beta1_power, beta2_power,
+                                                          self.beta1, self.beta2, self.eps, lr), gradients, params,
+                                                moment1, moment2, self.ps_parameters, self.cache_enable)
 
-    @ms_function
+        return success
+
+    @jit
     def construct(self, gradients):
         params = self._parameters
         moment1 = self.moment1
         moment2 = self.moment2
-        vhat = self.vhat
         gradients = self.flatten_gradients(gradients)
         gradients = self.decay_weight(gradients)
-        gradients = self.gradients_centralization(gradients)
+        if not self.use_offload:
+            gradients = self.gradients_centralization(gradients)
         gradients = self.scale_grad(gradients)
         gradients = self._grad_sparse_indices_deduplicate(gradients)
         lr = self.get_lr()
@@ -544,38 +915,8 @@ class Adam(Optimizer):
         self.beta1_power = beta1_power
         beta2_power = self.beta2_power * self.beta2
         self.beta2_power = beta2_power
-        if self.use_dist_optimizer:
-            if self.use_amsgrad:
-                raise Exception("""Adam with amsgrad is currently not supporting distributed training!
-                                Please set use_amsgrad=False for distributed training.""")
-            if self.is_group_lr:
-                success = self.map_(F.partial(_adam_opt, self.opt, self.sparse_opt, self._ps_push, self._ps_pull,
-                                              self.use_locking, self.use_nesterov, self._is_device,
-                                              beta1_power, beta2_power, self.beta1, self.beta2, self.eps),
-                                    lr, gradients, params, moment1, moment2, self.ps_parameters, self.cache_enable,
-                                    self.dense_adam_opts, self.use_dense_opt_flags,
-                                    self.sparse_adam_opts, self.use_sparse_opt_flags)
-            else:
-                success = self.map_(F.partial(_adam_opt, self.opt, self.sparse_opt, self._ps_push, self._ps_pull,
-                                              self.use_locking, self.use_nesterov, self._is_device,
-                                              beta1_power, beta2_power, self.beta1, self.beta2, self.eps, lr),
-                                    gradients, params, moment1, moment2, self.ps_parameters, self.cache_enable,
-                                    self.dense_adam_opts, self.use_dense_opt_flags,
-                                    self.sparse_adam_opts, self.use_sparse_opt_flags)
-        else:
-            if self.is_group_lr:
-                success = self.map_(F.partial(_adam_opt, self.opt, self.sparse_opt, self._ps_push, self._ps_pull,
-                                              self.use_locking, self.use_nesterov, self.use_amsgrad, self._is_device,
-                                              beta1_power, beta2_power, self.beta1, self.beta2, self.eps),
-                                    lr, gradients, params, moment1, moment2, vhat,
-                                    self.ps_parameters, self.cache_enable)
-            else:
-                success = self.map_(F.partial(_adam_opt, self.opt, self.sparse_opt, self._ps_push, self._ps_pull,
-                                              self.use_locking, self.use_nesterov, self.use_amsgrad, self._is_device,
-                                              beta1_power, beta2_power, self.beta1, self.beta2, self.eps, lr),
-                                    gradients, params, moment1, moment2, vhat, self.ps_parameters, self.cache_enable)
 
-        return success
+        return self._apply_adam(params, beta1_power, beta2_power, moment1, moment2, lr, gradients)
 
     @Optimizer.target.setter
     def target(self, value):
@@ -630,13 +971,13 @@ class AdamWeightDecay(Optimizer):
     :math:`m` represents the 1st moment vector `moment1`, :math:`v` represents the 2nd moment vector `moment2`,
     :math:`g` represents `gradients`, :math:`\gamma` represents `learning_rate`,
     :math:`\beta_1, \beta_2` represent `beta1` and `beta2`, :math:`t` represents the current step,
-    :math:`w` represents `params`, :math:`\gamma` represents `weight_decay`.
+    :math:`w` represents `params`, :math:`\lambda` represents `weight_decay`.
 
     Note:
         There is usually no connection between a optimizer and mixed precision. But when `FixedLossScaleManager` is used
         and `drop_overflow_update` in `FixedLossScaleManager` is set to False, optimizer needs to set the 'loss_scale'.
         As this optimizer has no argument of `loss_scale`, so `loss_scale` needs to be processed by other means, refer
-        document `LossScale <https://www.mindspore.cn/tutorials/experts/en/r1.10/others/mixed_precision.html>`_ to
+        document `LossScale <https://www.mindspore.cn/tutorials/en/r2.0/advanced/mixed_precision.html>`_ to
         process `loss_scale` correctly.
 
         If parameters are not grouped, the `weight_decay` in optimizer will be applied on the network parameters without
@@ -753,7 +1094,7 @@ class AdamWeightDecay(Optimizer):
         else:
             self.use_fused_opt = False
 
-    @ms_function
+    @jit
     def construct(self, gradients):
         gradients = self.flatten_gradients(gradients)
         weight_decay = self.get_weight_decay()
@@ -804,7 +1145,7 @@ class AdamWeightDecay(Optimizer):
         """
         self._set_base_target(value)
         if value == 'CPU':
-            self.fused_opt.add_prim_attr("primitive_target", "CPU")
+            self.fused_opt.set_device("CPU")
             self.use_fused_opt = True
         else:
             self.use_fused_opt = False
@@ -949,6 +1290,7 @@ class AdamOffload(Optimizer):
         >>> model = ms.Model(net, loss_fn=loss, optimizer=optim)
     """
 
+    @deprecated("2.0", "Adam", False)
     def __init__(self, params, learning_rate=1e-3, beta1=0.9, beta2=0.999, eps=1e-8, use_locking=False,
                  use_nesterov=False, weight_decay=0.0, loss_scale=1.0):
         super(AdamOffload, self).__init__(learning_rate, params, weight_decay, loss_scale)
@@ -965,9 +1307,9 @@ class AdamOffload(Optimizer):
         self.moment1 = self._parameters.clone(prefix="moment1", init='zeros')
         self.moment2 = self._parameters.clone(prefix="moment2", init='zeros')
         self.opt = P.AdamNoUpdateParam(use_locking, use_nesterov)
-        self.opt.add_prim_attr("primitive_target", "CPU")
+        self.opt.set_device("CPU")
 
-    @ms_function
+    @jit
     def construct(self, gradients):
         params = self._parameters
         moment1 = self.moment1