mindspore 2.6.0rc1__cp39-cp39-win_amd64.whl → 2.7.0rc1__cp39-cp39-win_amd64.whl

mindspore/device_manager.py

@@ -15,16 +15,14 @@
 
 """Device manager interfaces."""
 
+__all__ = ['set_device', 'set_deterministic', 'get_current_device']
+
 import os
 from mindspore import log as logger
 from mindspore._c_expression import DeviceManagerConf, DeviceContextManager, MSContext, CollectiveManager
 from mindspore._checkparam import args_type_check
 from mindspore.parallel._ps_context import _need_reset_device_target_for_ps
 
-
-__all__ = ['set_device', 'set_deterministic', 'get_current_device']
-
-
 class DeviceInfo(tuple):
     """
     DeviceInfo class. Store the current device target and the corresponding device id.
@@ -120,14 +118,21 @@ def set_deterministic(deterministic):
     """
     Enables or disables deterministic computing.
 
-    When deterministic computing is enabled, the same output is generated if an operator is executed
-    for multiple times with the same hardware and input.This often slows down operator execution.
-    In distributed scenario, we suggest user to set deterministic mode before
-    calling :func:`mindspore.communication.init` to enable deterministic operation for
-    communication operators in the global communication group.
+    This configuration is a global configuration, and once enabled, subsequent calculation operations
+    will follow the configuration setting. When deterministic computing is enabled, the same output
+    is generated if an operator is executed for multiple times with the same hardware and input. This often
+    slows down operator execution.
 
     The framework not enabled deterministic computation by default.
 
+    Note:
+        - In distributed scenario, we suggest user to set deterministic computing before
+          calling :func:`mindspore.communication.init` to enable deterministic operation for
+          communication operators in the global communication group.
+        - The fixed method for deterministic calculation must be in the same main process as the network,
+          operator, etc. Only one deterministic state can be set in the same thread, and it is not recommended
+          to set deterministic state multiple times in one thread.
+
     Args:
         deterministic (bool): Whether to enable deterministic computing.
 
@@ -139,6 +144,8 @@ def set_deterministic(deterministic):
     if DeviceManagerConf.get_instance().is_deterministic_configured():
         raise RuntimeError("The 'mindspore.set_deterministic' can not be set repeatedly.")
 
+    logger.info(f"Set deterministic setting to '{deterministic}'.")
+
     # Must wait for all async created groups to be initialized so that
     # deterministic feature could be consistent between all processes.
     CollectiveManager.get_instance().wait_all_comm_init()

mindspore/experimental/llm_boost/ascend_native/llama_boost_ascend_native.py

@@ -18,8 +18,6 @@ import os
 import numpy as np
 from mindspore.common import Tensor, dtype
 from mindspore.experimental.llm_boost.ascend_native.llm_boost import LLMBoost
-from mindspore.experimental.llm_boost.register import LlmBoostRegister, LlmBoostType
-
 
 def RoundUp(val: int, align: int) -> int:
     if align == 0:
@@ -44,8 +42,6 @@ def ConvertTensor(nd_mat: np.ndarray, transpose: bool = True, nd2nz: bool = True
     nz_mat = nz_mat.reshape(r, c)
     return nz_mat
 
-
-@LlmBoostRegister.register(LlmBoostType.ASCEND_NATIVE, "Llama")
 class LlamaBoostAscendNative(LLMBoost):
     r"""
     Implements an Llama model in a single kernel.
@@ -68,7 +64,7 @@ class LlamaBoostAscendNative(LLMBoost):
     def _prepare_single_layer(self, ckpt, config, id):
         """ prepares the dictionary of weights of a single layer """
         prefix = 'model.layers.' + str(id)
-        is_last = (id == config.num_layers-1)
+        is_last = id == config.num_layers-1
         layer = 'layers.' + str(id) + '.'
         l_dict = {key: value for key, value in ckpt.items() if layer in key}
         if config.n_kv_heads is None:
@@ -96,8 +92,8 @@ class LlamaBoostAscendNative(LLMBoost):
         else:
             raise RuntimeError("hidden size and ffn hidden size must be divided by rank size without remainder.  \
                                 hidden_size: ", hid_size, " ffn_hidden_size: ", ffn_hid, " rank_size: ", rank_size)
-        quant = (self._get_from_dict(l_dict, "_weight_quantizer") is not None)
-        unite_qkv = (config.num_heads == config.n_kv_heads)
+        quant = self._get_from_dict(l_dict, "_weight_quantizer") is not None
+        unite_qkv = config.num_heads == config.n_kv_heads
         self.dictionary[prefix + ".attention_norm.weight"] = \
             Tensor(self._get_from_dict(l_dict, "attention_norm"), dtype=dtype.float16)
         self.dictionary[prefix + ".ffn_norm.weight"] = \

mindspore/experimental/llm_boost/atb/boost_base.py

@@ -18,7 +18,6 @@ import numpy as np
 import mindspore as ms
 from mindspore import ops, Tensor
 from mindspore import log as logger
-from mindspore.ops import operations as P
 import mindspore.common.dtype as mstype
 from mindspore._c_expression import _set_format
 from mindspore.common.parameter import Parameter
@@ -95,8 +94,8 @@ class AtbBoostBase:
             self.max_base_len, dtype=self.dtype, need_nz=self.need_nz
         )
 
-        self.cast = P.Cast()
-        self.reshape = P.Reshape()
+        self.cast = ops.Cast()
+        self.reshape = ops.Reshape()
         self.kv_quant = None
         self.rank_id = get_real_rank()
         self.device_num = get_real_group_size()

mindspore/experimental/optim/adadelta.py

@@ -1,34 +1,24 @@
-# Copyright 2023 Huawei Technologies Co., Ltd
-#
-# 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.
+# The code implementation refers to the following files from pytorch:
+# - https://github.com/pytorch/pytorch/blob/v1.13.0/torch/optim/adadelta.py
+# Additional modifications are made by Huawei Technologies Co., Ltd in 2023.
 # ============================================================================
 """adadelta"""
 from __future__ import absolute_import
 
-from mindspore.ops import functional as F, composite as C, operations as P
+from mindspore import ops
 import mindspore.common.dtype as mstype
 from mindspore.experimental.optim.optimizer import Optimizer, check_not_less_than, check_not_less_than_without_equal
 from mindspore import _checkparam as validator
 from mindspore import jit
 
-_adadelta_opt = C.MultitypeFuncGraph("adadelta_opt")
+_adadelta_opt = ops.MultitypeFuncGraph("adadelta_opt")
 
 
 @_adadelta_opt.register("Function", "Number", "Number", "Tensor", "Tensor", "Tensor", "Tensor", "Tensor")
 def _tensor_run_opt(opt, rho, epsilon, learning_rate, weight, accum, accum_update, gradient):
     """Apply adadelta optimizer to the weight parameter."""
     success = True
-    success = F.depend(success, opt(weight, accum, accum_update, learning_rate, rho, epsilon, gradient))
+    success = ops.depend(success, opt(weight, accum, accum_update, learning_rate, rho, epsilon, gradient))
     return success
 
 
@@ -60,6 +50,9 @@ class Adadelta(Optimizer):
             &\rule{180mm}{0.4pt}                                                          \\[-1.ex]
         \end{aligned}
 
+    For more details about Adadelta algorithm, please refer to `ADADELTA: An Adaptive Learning Rate Method
+    <https://arxiv.org/abs/1212.5701>`_.
+
     .. warning::
         This is an experimental optimizer API that is subject to change.
         This module must be used with lr scheduler module in `LRScheduler Class
@@ -129,18 +122,18 @@ class Adadelta(Optimizer):
 
         self.accum = self.parameters.clone(prefix="accum", init=0)
         self.accum_update = self.parameters.clone(prefix="accum_update", init=0)
-        self.opt = P.ApplyAdadelta()
-        self.op_cast = P.Cast()
+        self.opt = ops.ApplyAdadelta()
+        self.op_cast = ops.Cast()
 
     @jit
     def implementation(self, lr, rho, eps, maximize, weight_decay, start_id, end_id, gradients):
         """Extract the common computing part for acceleration"""
         params = self.parameters[start_id: end_id]
-        grads = tuple([grad if not maximize else F.neg(grad) for grad in gradients[start_id: end_id]])
+        grads = tuple([grad if not maximize else ops.neg(grad) for grad in gradients[start_id: end_id]])
         grads = self._decay_weight(weight_decay, params, grads)
         accum = self.accum[start_id: end_id]
         accum_update = self.accum_update[start_id: end_id]
-        self.hyper_map(F.partial(_adadelta_opt, self.opt, rho, eps, lr),
+        self.hyper_map(ops.partial(_adadelta_opt, self.opt, rho, eps, lr),
                        params, accum, accum_update, grads)
         return True
 

mindspore/experimental/optim/adagrad.py

@@ -1,34 +1,24 @@
-# Copyright 2023 Huawei Technologies Co., Ltd
-#
-# 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.
+# The code implementation refers to the following files from pytorch:
+# - https://github.com/pytorch/pytorch/blob/v1.13.0/torch/optim/adagrad.py
+# Additional modifications are made by Huawei Technologies Co., Ltd in 2023.
 # ============================================================================
 """adagrad"""
 from __future__ import absolute_import
 
-from mindspore.ops import functional as F, composite as C, operations as P
+from mindspore import ops
 from mindspore.common import Tensor, Parameter
 import mindspore.common.dtype as mstype
 from mindspore.experimental.optim.optimizer import Optimizer, check_not_less_than, check_not_less_than_without_equal
 from mindspore import jit
 
-_adagrad_opt = C.MultitypeFuncGraph("adagrad_opt")
+_adagrad_opt = ops.MultitypeFuncGraph("adagrad_opt")
 
 
 @_adagrad_opt.register("Function", "Tensor", "Tensor", "Tensor", "Tensor")
 def _tensor_run_opt(opt, learning_rate, weight, accum, gradient):
     """Apply adagrad optimizer to the weight parameter."""
     success = True
-    success = F.depend(success, opt(weight, accum, learning_rate, gradient))
+    success = ops.depend(success, opt(weight, accum, learning_rate, gradient))
     return success
 
 
@@ -57,6 +47,9 @@ class Adagrad(Optimizer):
             &\rule{160mm}{0.4pt}                                                          \\[-1.ex]
        \end{aligned}
 
+    For more details about Adagrad algorithm, please refer to `Adaptive Subgradient Methods for Online Learning and \
+    Stochastic Optimization <https://jmlr.org/papers/v12/duchi11a.html>`_.
+
     .. warning::
         This is an experimental optimizer API that is subject to change.
         This module must be used with lr scheduler module in `LRScheduler Class
@@ -129,22 +122,22 @@ class Adagrad(Optimizer):
         super(Adagrad, self).__init__(params, defaults)
 
         self.accum = self.parameters.clone(prefix="accum", init=initial_accumulator_value)
-        self.op_cast = P.Cast()
+        self.op_cast = ops.Cast()
         self.step_t = Parameter(Tensor(0, mstype.int32), "step_t")
         self.increase_tensor = Tensor(1, mstype.int32)
-        self.assignadd = P.AssignAdd()
-        self.assign = P.Assign()
+        self.assignadd = ops.AssignAdd()
+        self.assign = ops.Assign()
 
     @jit
     def implementation(self, eps, lr, lr_decay, maximize, weight_decay, start_id, end_id, gradients):
         """Extract the common computing part for acceleration"""
-        opt = P.ApplyAdagradV2(epsilon=eps, update_slots=True)
+        opt = ops.ApplyAdagradV2(epsilon=eps, update_slots=True)
         decay_lr = lr / (1 + self.step_t * lr_decay)
         params = self.parameters[start_id: end_id]
-        grads = tuple([grad if not maximize else F.neg(grad) for grad in gradients[start_id: end_id]])
+        grads = tuple([grad if not maximize else ops.neg(grad) for grad in gradients[start_id: end_id]])
         grads = self._decay_weight(weight_decay, params, grads)
         accum = self.accum[start_id: end_id]
-        self.hyper_map(F.partial(_adagrad_opt, opt, decay_lr), params, accum, grads)
+        self.hyper_map(ops.partial(_adagrad_opt, opt, decay_lr), params, accum, grads)
         return True
 
     def construct(self, gradients):

mindspore/experimental/optim/adam.py

@@ -1,29 +1,19 @@
-# Copyright 2023 Huawei Technologies Co., Ltd
-#
-# 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.
+# The code implementation refers to the following files from pytorch:
+# - https://github.com/pytorch/pytorch/blob/v1.13.0/torch/optim/adam.py
+# Additional modifications are made by Huawei Technologies Co., Ltd in 2023.
 # ============================================================================
 """adam"""
 from __future__ import absolute_import
 
-from mindspore.ops import functional as F, composite as C, operations as P
+from mindspore import ops
 from mindspore.common.parameter import Parameter
 from mindspore.common.tensor import Tensor
 import mindspore.common.dtype as mstype
 from mindspore.experimental.optim.optimizer import Optimizer
 from mindspore.common.api import jit
 
-_adam_opt = C.MultitypeFuncGraph("adam_opt")
-adam_op = P.Adam(False, False)
+_adam_opt = ops.MultitypeFuncGraph("adam_opt")
+adam_op = ops.Adam(False, False)
 
 
 @_adam_opt.register("Tensor", "Tensor", "Float", "Float", "Float", "Tensor",
@@ -81,6 +71,9 @@ class Adam(Optimizer):
             &\rule{180mm}{0.4pt}                                                          \\[-1.ex]
        \end{aligned}
 
+    For more details about Adam algorithm, please refer to `Adam: A Method for Stochastic Optimization
+    <https://arxiv.org/abs/1412.6980>`_.
+
     .. warning::
         The implementation formula of this optimizer interface is not completely consistent with that in the paper.
         If you want to use an interface that is completely consistent, it is recommended to use
@@ -160,12 +153,12 @@ class Adam(Optimizer):
         self.max_exp_avg_sq = self.parameters.clone(prefix="max_exp_avg_sq", init='zeros')
         self.state_step = Parameter(Tensor(0, mstype.int32), "state_step")
         self.increase_tensor = Tensor(1, mstype.int32)
-        self.assignadd = P.AssignAdd()
-        self.op_add = P.AddN()
-        self.op_mul = P.Mul()
-        self.op_pow = P.Pow()
-        self.adam_opt = P.Adam(False, False)
-        self.op_cast = P.Cast()
+        self.assignadd = ops.AssignAdd()
+        self.op_add = ops.AddN()
+        self.op_mul = ops.Mul()
+        self.op_pow = ops.Pow()
+        self.adam_opt = ops.Adam(False, False)
+        self.op_cast = ops.Cast()
 
     @jit
     def implementation(self, beta1, beta2, eps, lr, start_id, end_id, gradients, maximize, weight_decay):
@@ -173,9 +166,9 @@ class Adam(Optimizer):
         beta1_power = self.op_pow(beta1, self.state_step)
         beta2_power = self.op_pow(beta2, self.state_step)
         params = self.parameters[start_id: end_id]
-        grads = tuple([grad if not maximize else F.neg(grad) for grad in gradients[start_id: end_id]])
+        grads = tuple([grad if not maximize else ops.neg(grad) for grad in gradients[start_id: end_id]])
         grads = self._decay_weight(weight_decay, params, grads)
-        self.hyper_map(F.partial(_adam_opt, beta1_power, beta2_power, beta1, beta2, eps, lr),
+        self.hyper_map(ops.partial(_adam_opt, beta1_power, beta2_power, beta1, beta2, eps, lr),
                        grads, params,
                        self.exp_avg[start_id: end_id], self.exp_avg_sq[start_id: end_id])
         return True

mindspore/experimental/optim/adamax.py

@@ -1,21 +1,10 @@
-# Copyright 2023 Huawei Technologies Co., Ltd
-#
-# 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.
+# The code implementation refers to the following files from pytorch:
+# - https://github.com/pytorch/pytorch/blob/v1.13.0/torch/optim/adamax.py
+# Additional modifications are made by Huawei Technologies Co., Ltd in 2023.
 # ============================================================================
 """adamax"""
 from __future__ import absolute_import
 
-from mindspore.ops import functional as F, composite as C, operations as P
 from mindspore.common import Tensor, Parameter
 import mindspore.common.dtype as mstype
 from mindspore import _checkparam as validator
@@ -23,17 +12,17 @@ from mindspore.experimental.optim.optimizer import Optimizer, check_not_less_tha
 from mindspore import ops
 from mindspore import jit
 
-_adamax_opt = C.MultitypeFuncGraph("adamax_opt")
+_adamax_opt = ops.MultitypeFuncGraph("adamax_opt")
 
 
 @_adamax_opt.register("Number", "Number", "Number", "Tensor", "Tensor", "Tensor", "Tensor", "Tensor")
 def _tensor_run_opt(beta1, beta2, eps, clr, param, grad, exp_avg, exp_inf):
     """Apply adamax optimizer to the weight parameter."""
-    F.assign(exp_avg, exp_avg * beta1 + grad * (1-beta1))
+    ops.assign(exp_avg, exp_avg * beta1 + grad * (1-beta1))
     norm_buf = ops.cat([ops.unsqueeze(exp_inf * beta2, 0), ops.unsqueeze(grad.abs().add(eps), 0)], 0)
-    F.assign(exp_inf, ops.amax(norm_buf, 0))
+    ops.assign(exp_inf, ops.amax(norm_buf, 0))
 
-    F.assign(param, param - clr * exp_avg / exp_inf)
+    ops.assign(param, param - clr * exp_avg / exp_inf)
     return True
 
 
@@ -68,6 +57,9 @@ class Adamax(Optimizer):
         This module must be used with lr scheduler module in `LRScheduler Class
         <https://www.mindspore.cn/docs/en/master/api_python/mindspore.experimental.html#lrscheduler-class>`_ .
 
+    For more details about Adamax algorithm, please refer to `Adam: A Method for Stochastic Optimization
+    <https://arxiv.org/abs/1412.6980>`_.
+
     Args:
         params (Union[list(Parameter), list(dict)]): list of parameters to optimize or dicts defining
             parameter groups.
@@ -135,8 +127,8 @@ class Adamax(Optimizer):
         self.exp_avg = self.parameters.clone(prefix="exp_avg", init='zeros')
         self.exp_inf = self.parameters.clone(prefix="exp_inf", init='zeros')
         self.increase_tensor = Tensor(1, mstype.int32)
-        self.assignadd = P.AssignAdd()
-        self.op_cast = P.Cast()
+        self.assignadd = ops.AssignAdd()
+        self.op_cast = ops.Cast()
 
     @jit
     def implementation(self, group_id, lr, gradients, maximize, weight_decay, beta1, beta2, eps):
@@ -144,13 +136,13 @@ class Adamax(Optimizer):
         start_id = self.group_start_id[group_id]
         end_id = self.group_start_id[group_id + 1]
         params = self.parameters[start_id: end_id]
-        grads = tuple([grad if not maximize else F.neg(grad) for grad in gradients[start_id: end_id]])
+        grads = tuple([grad if not maximize else ops.neg(grad) for grad in gradients[start_id: end_id]])
         grads = self._decay_weight(weight_decay, params, grads)
         exp_avg = self.exp_avg[start_id: end_id]
         exp_inf = self.exp_inf[start_id: end_id]
         bias_correction = 1 - beta1 ** self.step_t
         clr = lr / bias_correction
-        self.hyper_map(F.partial(_adamax_opt, beta1, beta2, eps, clr),
+        self.hyper_map(ops.partial(_adamax_opt, beta1, beta2, eps, clr),
                        params, grads, exp_avg, exp_inf)
         return True
 

mindspore/experimental/optim/adamw.py

@@ -1,21 +1,10 @@
-# Copyright 2023 Huawei Technologies Co., Ltd
-#
-# 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.
+# The code implementation refers to the following files from pytorch:
+# - https://github.com/pytorch/pytorch/blob/v1.13.0/torch/optim/adamw.py
+# Additional modifications are made by Huawei Technologies Co., Ltd in 2023.
 # ============================================================================
 """adamw"""
 from __future__ import absolute_import
 
-from mindspore.ops import functional as F, composite as C, operations as P
 from mindspore.common.parameter import Parameter
 from mindspore.common.tensor import Tensor
 import mindspore.common.dtype as mstype
@@ -25,14 +14,14 @@ from mindspore.ops import auto_generate as gen
 from mindspore import ops
 from mindspore import jit
 
-_adamw_opt = C.MultitypeFuncGraph("adamw_opt")
-_speed_adamw_opt = C.MultitypeFuncGraph("speed_adamw_opt")
+_adamw_opt = ops.MultitypeFuncGraph("adamw_opt")
+_speed_adamw_opt = ops.MultitypeFuncGraph("speed_adamw_opt")
 
-op_mul = P.Mul()
-op_pow = P.Pow()
-op_sqrt = P.Sqrt()
-op_maximum = P.Maximum()
-hyper_map = C.HyperMap()
+op_mul = ops.Mul()
+op_pow = ops.Pow()
+op_sqrt = ops.Sqrt()
+op_maximum = ops.Maximum()
+hyper_map = ops.HyperMap()
 
 
 @_speed_adamw_opt.register("Function", "Float", "Float", "Tensor", "Float", "Float", "Bool", "Bool", "Tensor", "Tensor",
@@ -76,18 +65,18 @@ def _run_adamw_opt(weight_decay_new, step_size, amsgrad, eps, bias_correction2_s
     """Apply adamw optimizer to the weight parameter."""
     success = True
     next_param = op_mul(param, weight_decay_new)
-    F.assign(exp_avg, op_mul(exp_avg, beta1) + op_mul(grad, 1 - beta1))
-    F.assign(exp_avg_sq, ops.addcmul(op_mul(exp_avg_sq, beta2), grad, grad, 1 - beta2))
+    ops.assign(exp_avg, op_mul(exp_avg, beta1) + op_mul(grad, 1 - beta1))
+    ops.assign(exp_avg_sq, ops.addcmul(op_mul(exp_avg_sq, beta2), grad, grad, 1 - beta2))
 
     if amsgrad:
         next_max_exp_avg = op_maximum(max_exp_avg_sq, exp_avg_sq)
         denom = op_sqrt(next_max_exp_avg) / bias_correction2_sqrt + eps
-        F.assign(max_exp_avg_sq, next_max_exp_avg)
+        ops.assign(max_exp_avg_sq, next_max_exp_avg)
     else:
         denom = op_sqrt(exp_avg_sq) / bias_correction2_sqrt + eps
 
     return_param = next_param - op_mul(exp_avg / denom, step_size)
-    F.assign(param, return_param)
+    ops.assign(param, return_param)
     return success
 
 
@@ -129,6 +118,10 @@ class AdamW(Optimizer):
             &\rule{180mm}{0.4pt}                                                          \\[-1.ex]
        \end{aligned}
 
+    More details of the AdamW algorithm can be found in the paper `Decoupled Weight Decay Regularization
+    <https://arxiv.org/abs/1711.05101>`_ and `On the Convergence of Adam and Beyond
+    <https://openreview.net/forum?id=ryQu7f-RZ>`_.
+
     .. warning::
         This is an experimental optimizer API that is subject to change.
         This module must be used with lr scheduler module in `LRScheduler Class
@@ -205,16 +198,16 @@ class AdamW(Optimizer):
         self.max_exp_avg_sq = self.parameters.clone(prefix="max_exp_avg_sq", init='zeros')
         self.state_step = Parameter(Tensor(0, mstype.int32), "state_step")
         self.increase_tensor = Tensor(1, mstype.int32)
-        self.assignadd = P.AssignAdd()
-        self.op_cast = P.Cast()
+        self.assignadd = ops.AssignAdd()
+        self.op_cast = ops.Cast()
 
     @jit
     def implementation(self, lr, weight_decay, beta1, beta2, amsgrad, eps, grads, start_id, end_id):
         """Extract the common computing part for acceleration"""
         weight_decay_new, step_size, bias_correction2_sqrt = prepare_func(lr, weight_decay,
                                                                           self.state_step, beta1, beta2)
-        self.hyper_map(F.partial(_adamw_opt, weight_decay_new, step_size, amsgrad,
-                                 eps, bias_correction2_sqrt, beta1, beta2),
+        self.hyper_map(ops.partial(_adamw_opt, weight_decay_new, step_size, amsgrad,
+                                   eps, bias_correction2_sqrt, beta1, beta2),
                        self.parameters[start_id: end_id], grads, self.exp_avg[start_id: end_id],
                        self.exp_avg_sq[start_id: end_id], self.max_exp_avg_sq[start_id: end_id])
         return True
@@ -228,7 +221,8 @@ class AdamW(Optimizer):
             lr = self.lrs[group_id]
             if isinstance(group.get("lr"), float):
                 lr = self.op_cast(group.get("lr"), mstype.float32)
-            grads = tuple([grad if not group.get("maximize") else F.neg(grad) for grad in gradients[start_id: end_id]])
+            grads = tuple([grad if not group.get("maximize") else ops.neg(grad) \
+                           for grad in gradients[start_id:end_id]])
 
             self.implementation(lr, group.get("weight_decay"), beta1, beta2, group.get("amsgrad"), group.get("eps"),
                                 grads, start_id, end_id)
@@ -265,7 +259,7 @@ class SpeedAdamW(Optimizer):
         self.exp_avg_sq = self.parameters.clone(prefix="exp_avg_sq", init='zeros')
         self.state_step = Parameter(Tensor([0], mstype.float32), "state_step")
         self.increase_tensor = Tensor(1, mstype.float32)
-        self.assignadd = P.AssignAdd()
+        self.assignadd = ops.AssignAdd()
         self.adamw_opt = gen.ApplyAdamW()
 
     def construct(self, gradients):
@@ -285,9 +279,9 @@ class SpeedAdamW(Optimizer):
             if group.get("amsgrad"):
                 raise ValueError("For SpeedAdamW, the value of amsgrad can only be False.")
 
-            self.hyper_map(F.partial(_speed_adamw_opt, self.adamw_opt, beta1, beta2, lr,
-                                     group.get("eps"), group.get("weight_decay"),
-                                     group.get("amsgrad"), maximize, bias_correction1, bias_correction2),
+            self.hyper_map(ops.partial(_speed_adamw_opt, self.adamw_opt, beta1, beta2, lr,
+                                       group.get("eps"), group.get("weight_decay"),
+                                       group.get("amsgrad"), maximize, bias_correction1, bias_correction2),
                            self.parameters[start_id: end_id], grads, self.exp_avg[start_id: end_id],
                            self.exp_avg_sq[start_id: end_id])
 

mindspore/experimental/optim/asgd.py

@@ -1,33 +1,23 @@
-# Copyright 2021-2022 Huawei Technologies Co., Ltd
-#
-# 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.
+# The code implementation refers to the following files from pytorch:
+# - https://github.com/pytorch/pytorch/blob/v1.13.0/torch/optim/asgd.py
+# Additional modifications are made by Huawei Technologies Co., Ltd in 2023.
 # ============================================================================
 """asgd"""
 from __future__ import absolute_import
 
-from mindspore.ops import functional as F, composite as C, operations as P
+from mindspore import ops
 from mindspore.common import Tensor, Parameter
 import mindspore.common.dtype as mstype
 from mindspore.experimental.optim.optimizer import Optimizer, check_not_less_than, check_not_less_than_without_equal
 from mindspore.common.api import jit
 
-_asgd_opt = C.MultitypeFuncGraph("asgd_opt")
+_asgd_opt = ops.MultitypeFuncGraph("asgd_opt")
 
-op_cast = P.Cast()
-op_pow = P.Pow()
-op_maximum = P.Maximum()
-op_assign = P.Assign()
-op_assignadd = P.AssignAdd()
+op_cast = ops.Cast()
+op_pow = ops.Pow()
+op_maximum = ops.Maximum()
+op_assign = ops.Assign()
+op_assignadd = ops.AssignAdd()
 
 
 @_asgd_opt.register("Number", "Number", "Number", "Tensor", "Tensor", "Tensor", "Tensor",
@@ -37,7 +27,7 @@ def _run_asgd_opt(lambd, alpha, t0, step, lr, param, grad, eta, mu, ax):
     if step == 1:
         op_assign(eta, lr)
     next_param = op_cast(param * (1. - lambd * eta) - eta * grad, param.dtype)
-    F.assign(param, next_param)
+    ops.assign(param, next_param)
 
     if mu != 1:
         op_assignadd(ax, op_cast((next_param - ax) * mu, ax.dtype))
@@ -121,8 +111,8 @@ class ASGD(Optimizer):
         self.ax = self.parameters.clone(prefix="ax", init='zeros')
         self.step_t = Parameter(Tensor(0, mstype.int32), "step_t")
         self.increase_tensor = Tensor(1, mstype.int32)
-        self.assignadd = P.AssignAdd()
-        self.op_cast = P.Cast()
+        self.assignadd = ops.AssignAdd()
+        self.op_cast = ops.Cast()
 
     @jit(backend="ms_backend")
     def implementation(self, lambd, alpha, t0, lr, group_id, maximize, gradients, weight_decay):
@@ -130,13 +120,13 @@ class ASGD(Optimizer):
         start_id = self.group_start_id[group_id]
         end_id = self.group_start_id[group_id + 1]
         params = self.parameters[start_id: end_id]
-        grads = tuple([grad if not maximize else F.neg(grad) for grad in gradients[start_id: end_id]])
+        grads = tuple([grad if not maximize else ops.neg(grad) for grad in gradients[start_id: end_id]])
         grads = self._decay_weight(weight_decay, params, grads)
 
         ax = self.ax[start_id: end_id]
         eta = self.eta[start_id: end_id]
         mu = self.mu[start_id: end_id]
-        self.hyper_map(F.partial(_asgd_opt, lambd, alpha, t0, self.step_t, lr),
+        self.hyper_map(ops.partial(_asgd_opt, lambd, alpha, t0, self.step_t, lr),
                        params, grads, eta, mu, ax)
         return True