mindspore 2.7.0rc1__cp311-cp311-win_amd64.whl → 2.7.1__cp311-cp311-win_amd64.whl

mindspore/nn/optim/lazyadam.py

@@ -26,34 +26,19 @@ from mindspore.common.tensor import Tensor
 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
 
 _lazy_adam_opt = C.MultitypeFuncGraph("lazy_adam_opt")
 
 
-@_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_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):
+@_lazy_adam_opt.register("Function", "Function", "Bool", "Bool", "Bool", "Tensor", "Tensor",
+                         "Tensor", "Tensor", "Tensor", "Tensor", "RowTensor", "Tensor", "Tensor", "Tensor")
+def _run_opt_with_sparse(opt, sparse_opt, use_locking, use_nesterov, target, beta1_power, beta2_power,
+                         beta1, beta2, eps, lr, gradient, params, m, v):
     """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,
@@ -85,122 +70,10 @@ def _run_opt_with_sparse_dist(opt, sparse_opt, push, pull, use_locking, use_nest
     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_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_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_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_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):
+@_lazy_adam_opt.register("Function", "Function", "Bool", "Bool", "Bool", "Tensor", "Tensor",
+                         "Tensor", "Tensor", "Tensor", "Tensor", "MapTensor", "MapTensor", "MapTensor", "MapTensor")
+def _run_map_tensor_opt_with_sparse(opt, sparse_opt, use_locking, use_nesterov, target, beta1_power,
+                                    beta2_power, beta1, beta2, eps, lr, gradient, params, m, v):
     """Apply sparse lazy adam optimizer to the weight parameter when the gradient is sparse(MapTensor)."""
     success = True
     indices, values = gradient.get_data()
@@ -229,19 +102,14 @@ def _run_map_tensor_opt_with_sparse(opt, sparse_opt, push, pull, use_locking, us
     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_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):
+@_lazy_adam_opt.register("Function", "Function", "Bool", "Bool", "Bool", "Tensor", "Tensor",
+                         "Tensor", "Tensor", "Tensor", "Tensor", "Tensor", "Tensor", "Tensor", "Tensor")
+def _run_opt_with_one_number(opt, sparse_opt, use_locking, use_nesterov, target, beta1_power, beta2_power,
+                             beta1, beta2, eps, lr, gradient, params, moment1, moment2):
     """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))
+    success = F.depend(success, opt(params, moment1, moment2, beta1_power, beta2_power, lr, beta1, beta2,
+                                    eps, gradient))
     return success
 
 
@@ -436,15 +304,9 @@ class LazyAdam(Optimizer):
         self.opt = P.Adam(use_locking, use_nesterov)
         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)
 
     @jit
     def construct(self, gradients):
-        gradients = self.flatten_gradients(gradients)
         gradients = self.decay_weight(gradients)
         gradients = self.gradients_centralization(gradients)
         gradients = self.scale_grad(gradients)
@@ -457,40 +319,18 @@ class LazyAdam(Optimizer):
         beta2_power = self.beta2_power * self.beta2
         self.beta2_power = beta2_power
 
-        if self.use_dist_optimizer:
-            if self.is_group_lr:
-                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_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)
+        if self.is_group_lr:
+            success = self.map_reverse(F.partial(_lazy_adam_opt, self.opt, self.sparse_opt,
+                                                 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)
         else:
-            if self.is_group_lr:
-                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)
-            else:
-                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)
+            success = self.map_reverse(F.partial(_lazy_adam_opt, self.opt, self.sparse_opt,
+                                                 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)
         return success
 
     @Optimizer.target.setter
@@ -500,36 +340,3 @@ class LazyAdam(Optimizer):
         optimizer operation.
         """
         self._set_base_target(value)
-
-    def _init_distributed_opts(self, use_locking, use_nesterov):
-        self.use_dist_optimizer = self._use_distibuted_optimizer()
-        self.dense_lazyadam_opts, self.use_dense_opt_flags =\
-        self._get_distributed_optimizer_list("adam", use_locking, use_nesterov)
-        self.sparse_lazyadam_opts, self.use_sparse_opt_flags =\
-        self._get_distributed_optimizer_list("fused_sparse_lazy_adam", use_locking, use_nesterov)
-
-
-def create_distributed_adam(*args, **kwargs):
-    """
-    Create the distributed Adam op.
-    """
-    adam = P.Adam(*args, **kwargs)
-    adam.add_prim_attr("gradient_type", "dense_gradient")
-    adam.add_prim_attr("parameter_input_index", 0)
-    adam.add_prim_attr("gradient_input_index", 9)
-    return adam
-
-
-def create_distributed_fused_sparse_lazy_adam(*args, **kwargs):
-    """
-    Create the distributed FusedSparseLazyAdam op.
-    """
-    sparse_lazy_adam = P.FusedSparseLazyAdam(*args, **kwargs)
-    sparse_lazy_adam.add_prim_attr("gradient_type", "sparse_gradient")
-    sparse_lazy_adam.add_prim_attr("parameter_input_index", 0)
-    sparse_lazy_adam.add_prim_attr("gradient_input_index", 9)
-    sparse_lazy_adam.add_prim_attr("indices_input_index", 10)
-    return sparse_lazy_adam
-
-_register_dist_optimizer("adam", create_distributed_adam)
-_register_dist_optimizer("fused_sparse_lazy_adam", create_distributed_fused_sparse_lazy_adam)

mindspore/nn/optim/momentum.py

@@ -23,7 +23,6 @@ import mindspore.common.dtype as mstype
 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
 
 
 _momentum_opt = C.MultitypeFuncGraph("momentum_opt")
@@ -35,18 +34,6 @@ def _tensor_run_opt_ext(opt, momentum, learning_rate, gradient, weight, moment):
     success = F.depend(True, opt(weight, moment, learning_rate, gradient, momentum))
     return success
 
-
-@_momentum_opt.register("Function", "Tensor", "Tensor", "Tensor", "Tensor", "Tensor", "Function", "Bool")
-def _tensor_run_opt_ext_dist(opt, momentum, learning_rate, gradient, weight, moment,
-                             distributed_opt, use_flag):
-    """Apply momentum optimizer to the weight parameter using Tensor."""
-    if use_flag:
-        success = F.depend(True, distributed_opt(weight, moment, learning_rate, gradient, momentum))
-    else:
-        success = F.depend(True, opt(weight, moment, learning_rate, gradient, momentum))
-    return success
-
-
 class Momentum(Optimizer):
     r"""
     Implements the Momentum algorithm.
@@ -196,45 +183,20 @@ class Momentum(Optimizer):
         self.moments = self.params.clone(prefix="moments", init='zeros')
         self.opt = P.ApplyMomentum(use_nesterov=self.use_nesterov)
 
-        self.distributed_opts, self.use_distributed_opt_flags =\
-        self._get_distributed_optimizer_list("momentum", use_nesterov=self.use_nesterov)
-        self.use_dist_optimizer = self._use_distibuted_optimizer()
 
     @jit(backend="ms_backend")
     def construct(self, gradients):
         params = self.params
         moments = self.moments
-        gradients = self.flatten_gradients(gradients)
         gradients = self.decay_weight(gradients)
         gradients = self.gradients_centralization(gradients)
         gradients = self.scale_grad(gradients)
         lr = self.get_lr()
         self.assignadd(self.global_step, self.global_step_increase_tensor)
-        if self.use_dist_optimizer:
-            if self.is_group_lr:
-                success = self.hyper_map_reverse(F.partial(_momentum_opt, self.opt, self.momentum),
-                                                 lr, gradients, params, moments,
-                                                 self.distributed_opts, self.use_distributed_opt_flags)
-            else:
-                success = self.hyper_map_reverse(F.partial(_momentum_opt, self.opt, self.momentum, lr),
-                                                 gradients, params, moments,
-                                                 self.distributed_opts, self.use_distributed_opt_flags)
+        if self.is_group_lr:
+            success = self.hyper_map_reverse(F.partial(_momentum_opt, self.opt, self.momentum),
+                                             lr, gradients, params, moments)
         else:
-            if self.is_group_lr:
-                success = self.hyper_map_reverse(F.partial(_momentum_opt, self.opt, self.momentum),
-                                                 lr, gradients, params, moments)
-            else:
-                success = self.hyper_map_reverse(F.partial(_momentum_opt, self.opt, self.momentum, lr),
-                                                 gradients, params, moments)
+            success = self.hyper_map_reverse(F.partial(_momentum_opt, self.opt, self.momentum, lr),
+                                             gradients, params, moments)
         return success
-
-
-def _create_distributed_momentum(*args, **kwargs):
-    momentum = P.ApplyMomentum(*args, **kwargs)
-    momentum.add_prim_attr("gradient_type", "dense_gradient")
-    momentum.add_prim_attr("parameter_input_index", 0)
-    momentum.add_prim_attr("gradient_input_index", 3)
-    return momentum
-
-
-_register_dist_optimizer("momentum", _create_distributed_momentum)

mindspore/nn/optim/optimizer.py

@@ -1,4 +1,4 @@
-# Copyright 2020-2022 Huawei Technologies Co., Ltd
+# Copyright 2020-2021 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.
@@ -32,11 +32,9 @@ import mindspore.common.dtype as mstype
 from mindspore import _checkparam as validator
 from mindspore import log as logger
 from mindspore.parallel._utils import _get_global_rank, _get_device_num, _get_parallel_mode
-from mindspore.parallel._ps_context import _is_ps_mode
 from mindspore.context import ParallelMode
 from mindspore import context
 from mindspore.nn.learning_rate_schedule import LearningRateSchedule
-from mindspore.nn.optim._dist_optimizer_registry import generate_dist_optimizer_list
 
 __all__ = ['Optimizer', 'opt_init_args_register']
 
@@ -161,7 +159,6 @@ class Optimizer(Cell):
         ...     def construct(self, gradients):
         ...         params = self.parameters
         ...         lr = self.get_lr()
-        ...         gradients = self.flatten_gradients(gradients)
         ...         gradients = self.decay_weight(gradients)
         ...         gradients = self.gradients_centralization(gradients)
         ...         gradients = self.scale_grad(gradients)
@@ -218,15 +215,12 @@ class Optimizer(Cell):
 
         self._unique = True
         self._target = context.get_context("device_target")
-        self._use_flattened_params = False
-        self._grad_fusion_size = 0
         self.dynamic_lr = False
         self.assignadd = P.AssignAdd()
         self.global_step = Parameter(initializer(0, [1], mindspore.int32), name='global_step')
         self.is_group = False
         self.is_group_lr = False
         self.is_group_params_ordered = False
-        self.use_parallel = False
         learning_rate = self._preprocess_single_lr(learning_rate)
         if isinstance(parameters[0], dict):
             self.is_group = True
@@ -263,11 +257,7 @@ class Optimizer(Cell):
             self.grad_centralization_flags = tuple(self.group_grad_centralization)
         else:
             self.parameters = ParameterTuple(parameters)
-            flat_params = self._get_flattened_params(parameters)
-            if self._use_flattened_params:
-                self._parameters = ParameterTuple(flat_params)
-            else:
-                self._parameters = self.parameters
+            self._parameters = self.parameters
             decay_filter = lambda x: 'beta' not in x.name and 'gamma' not in x.name
             self.decay_flags = tuple(decay_filter(x) for x in self._parameters)
             self.dynamic_decay_flags = isinstance(weight_decay, Cell)
@@ -281,10 +271,6 @@ class Optimizer(Cell):
         # set user's parameters as local parameters
         for param in self._parameters:
             self._user_parameters.append(param.name)
-        ps_filter = lambda x: x.is_param_ps
-        self.ps_parameters = tuple(ps_filter(x) for x in self._parameters)
-        cache_filter = lambda x: x.cache_enable
-        self.cache_enable = tuple(cache_filter(x) for x in self._parameters)
         self.reciprocal_scale = Tensor(1.0 / self.loss_scale, mstype.float32)
         self.need_scale = self.loss_scale != 1.0
         self.global_step_increase_tensor = Tensor([1], mstype.int32)
@@ -296,28 +282,6 @@ class Optimizer(Cell):
         self._use_parallel_optimizer()
         self.enable_tuple_broaden = True
 
-    def _get_flattened_params(self, parameters):
-        """Get parameters for each contiguous memory chunks used by input parameters if they are flattened."""
-        if self.is_group:
-            # We don't use flattened parameters when parameters are grouped.
-            return parameters
-        # Check whether parameters are flattened.
-        flattened = Tensor._is_flattened(parameters)  # pylint: disable=W0212
-        if not flattened:
-            # Parameters are not flattened.
-            return parameters
-        # Try to get chunk tensors from flattened parameters.
-        chunk_tensors = Tensor._get_flattened_tensors(parameters)  # pylint: disable=W0212
-        if not chunk_tensors:
-            # Failed to get chunk tensors.
-            logger.warning("Parameters are not properly flattened, fallback to not flattened parameters.")
-            return parameters
-        # Convert chunk tensors to parameters.
-        self._use_flattened_params = True
-        self._grad_fusion_size = Tensor._get_fusion_size(chunk_tensors)  # pylint: disable=W0212
-        return [Parameter._from_tensor(t, name='_chunk_param' + str(i) + '_' + str(t.dtype))  # pylint: disable=W0212
-                for i, t in enumerate(chunk_tensors)]
-
     def _use_parallel_optimizer(self):
         """Indicates whether to use automatic parallelism."""
         if context.get_auto_parallel_context("enable_parallel_optimizer"):
@@ -331,7 +295,10 @@ class Optimizer(Cell):
                 raise RuntimeError("For 'Optimizer', parallel optimizer is not supported in {}, you should set "
                                    "parallel mode to 'data_parallel', 'semi_auto_parallel' or 'auto_parallel'."
                                    .format(_get_parallel_mode()))
-
+            else:
+                self.use_parallel = False
+        else:
+            self.use_parallel = False
         if self.use_parallel:
             if not self._support_parallel_optimizer:
                 raise RuntimeError("For 'Optimizer', parallel optimizer only support optimizer 'Lamb' and "
@@ -403,13 +370,6 @@ class Optimizer(Cell):
             raise ValueError(f"For 'Optimizer', the argument {param_info} must not be empty.")
         return parameters
 
-    @staticmethod
-    def _use_distibuted_optimizer():
-        """
-        Whether use distributed optimizers.
-        """
-        return _is_ps_mode()
-
     def flatten_gradients(self, gradients):
         """
         Flatten gradients into several chunk tensors grouped by data type if network parameters are flattened.
@@ -424,9 +384,6 @@ class Optimizer(Cell):
         Returns:
             tuple[Tensor], The gradients after flattened, or the original gradients if parameters are not flattened.
         """
-        if self._use_flattened_params:
-            flatten_concat = inner.FlattenConcat(fusion_size=self._grad_fusion_size)
-            return flatten_concat(gradients)
         return gradients
 
     def decay_weight(self, gradients):
@@ -869,12 +826,6 @@ class Optimizer(Cell):
                 F.assign(param_group[root][i], next_params[i])
         return new_param_group
 
-    def _get_distributed_optimizer_list(self, optimizer_type, *args, **kwargs):
-        """
-        Get the distributed optimizers list in distributed training mode.
-        """
-        return generate_dist_optimizer_list(optimizer_type, self._parameters, *args, **kwargs)
-
     def construct(self, *hyper_params):
         raise NotImplementedError
 

mindspore/nn/optim/proximal_ada_grad.py

@@ -205,7 +205,6 @@ class ProximalAdagrad(Optimizer):
     def construct(self, grads):
         params = self._parameters
         accum = self.accum
-        grads = self.flatten_gradients(grads)
         grads = self.decay_weight(grads)
         grads = self.gradients_centralization(grads)
         grads = self.scale_grad(grads)

mindspore/nn/optim/rmsprop.py

@@ -238,7 +238,6 @@ class RMSProp(Optimizer):
     @jit
     def construct(self, gradients):
         params = self._parameters
-        gradients = self.flatten_gradients(gradients)
         gradients = self.decay_weight(gradients)
         gradients = self.gradients_centralization(gradients)
         gradients = self.scale_grad(gradients)

mindspore/nn/optim/rprop.py

@@ -203,7 +203,6 @@ class Rprop(Optimizer):
 
     @jit(backend="ms_backend")
     def construct(self, gradients):
-        gradients = self.flatten_gradients(gradients)
         gradients = self.decay_weight(gradients)
         gradients = self.gradients_centralization(gradients)
         gradients = self.scale_grad(gradients)

mindspore/nn/optim/sgd.py

@@ -226,7 +226,6 @@ class SGD(Optimizer):
         accum = self.accum
         stat = self.stat
         gradients = self.decay_weight(gradients)
-        gradients = self.flatten_gradients(gradients)
         gradients = self.gradients_centralization(gradients)
         gradients = self.scale_grad(gradients)
         lr = self.get_lr()

mindspore/nn/optim/tft_wrapper.py

@@ -69,10 +69,9 @@ class OptTFTWrapper(Optimizer):
         tft_env = os.getenv("MS_ENABLE_TFT", "")
         if ("TTP:1" not in tft_env) and ("UCE:1" not in tft_env) and ("ARF:1" not in tft_env):
             raise ValueError("MindIO TFT regitster need custom switch on[MS_ENABLE_TFT='{TTP:1,UCE:1,ARF:1}']!")
-        mode = context.get_context("mode")
         device_target = context.get_context("device_target")
-        if device_target != "Ascend" or mode != context.GRAPH_MODE:
-            raise ValueError("MindIO adataper only support on Ascend device with GRAPH Mode!")
+        if device_target != "Ascend":
+            raise ValueError("MindIO adataper only support on Ascend device!")
         self.opt = opt
         self.report = TensorReport()
         self.report_end = TensorReport()
@@ -109,7 +108,6 @@ class OptTFTWrapper(Optimizer):
         self.dynamic_decay_flags = opt.dynamic_decay_flags
         self.weight_decay = opt.weight_decay
         self.exec_weight_decay = opt.exec_weight_decay
-        self.ps_parameters = opt.ps_parameters
         self.cache_enable = opt.cache_enable
         self.reciprocal_scale = opt.reciprocal_scale
         self.need_scale = opt.need_scale

mindspore/nn/optim/thor.py

@@ -585,7 +585,6 @@ class ThorGpu(Optimizer):
     def construct(self, gradients):
         params = self.params
         moments = self.moments
-        gradients = self.flatten_gradients(gradients)
         gradients = self.scale_grad(gradients)
         damping_step = self.gather(self.damping, self.cov_step, self.axis)
         damping_step = self.cast(damping_step, mstype.float32)
@@ -1247,7 +1246,6 @@ class ThorAscend(Optimizer):
     def construct(self, gradients):
         params = self.params
         moments = self.moments
-        gradients = self.flatten_gradients(gradients)
         gradients = self.scale_grad(gradients)
         damping_step = self.gather(self.damping, self.cov_step, self.axis)
         damping_step = self.cast(damping_step, mstype.float32)

mindspore/nn/probability/bijector/bijector.py

@@ -15,8 +15,7 @@
 """Bijector"""
 from mindspore import context
 from mindspore.nn.cell import Cell
-from mindspore.ops import operations as P
-from mindspore.ops import functional as F
+import mindspore.ops as ops
 from mindspore.ops.operations import _inner_ops as inner
 from mindspore.common import dtype as mstype
 from mindspore.common.tensor import Tensor
@@ -99,9 +98,9 @@ class Bijector(Cell):
         self.checktensor = CheckTensor()
 
         # ops needed for the base class
-        self.cast_base = P.Cast()
-        self.dtype_base = P.DType()
-        self.shape_base = P.Shape()
+        self.cast_base = ops.Cast()
+        self.dtype_base = ops.DType()
+        self.shape_base = ops.Shape()
         self.sametypeshape_base = inner.SameTypeShape()
         self.issubclass_base = inner.IsSubClass()
 
@@ -145,13 +144,13 @@ class Bijector(Cell):
             if self.issubclass_base(value_type, mstype.float_):
                 return value
             return raise_type_error('input value of bijector', value_type, mstype.float_)
-        dtype_tensor = F.fill(self.dtype, self.shape_base(value), 0.0)
+        dtype_tensor = ops.fill(self.dtype, self.shape_base(value), 0.0)
         self.sametypeshape_base(value, dtype_tensor)
         return value
 
     def _shape_mapping(self, shape):
-        shape_tensor = F.fill(self.parameter_type, shape, 0.0)
-        dist_shape_tensor = F.fill(
+        shape_tensor = ops.fill(self.parameter_type, shape, 0.0)
+        dist_shape_tensor = ops.fill(
             self.parameter_type, self.batch_shape, 0.0)
         return (shape_tensor + dist_shape_tensor).shape
 

mindspore/nn/probability/bijector/gumbel_cdf.py

@@ -13,7 +13,7 @@
 # limitations under the License.
 # ============================================================================
 """GumbelCDF Bijector"""
-from mindspore.ops import operations as P
+import mindspore.ops as ops
 from ..distribution._utils.utils import check_greater_zero
 from ..distribution._utils.custom_ops import exp_generic, log_generic
 from .bijector import Bijector
@@ -86,7 +86,7 @@ class GumbelCDF(Bijector):
         self._scale = self._add_parameter(scale, 'scale')
         check_greater_zero(self._scale, "scale")
 
-        self.cast = P.Cast()
+        self.cast = ops.Cast()
         self.exp = exp_generic
         self.log = log_generic