diff --git a/mindspore/ccsrc/backend/session/session_basic.cc b/mindspore/ccsrc/backend/session/session_basic.cc index 7a9b44ca66..ed42ad7b18 100644 --- a/mindspore/ccsrc/backend/session/session_basic.cc +++ b/mindspore/ccsrc/backend/session/session_basic.cc @@ -1556,8 +1556,7 @@ void SessionBasic::InitPSParamAndOptim(const KernelGraphPtr &kernel_graph, auto input_node = input_nodes[i]; MS_EXCEPTION_IF_NULL(input_node); if (input_node->isa() && AnfAlgo::OutputAddrExist(input_node, 0)) { - auto pk_node = input_node->cast(); - ps::worker.InitPSParamAndOptim(pk_node->fullname_with_scope(), tensor); + ps::worker.InitPSParamAndOptim(input_node, tensor); } } } diff --git a/mindspore/ccsrc/ps/worker.h b/mindspore/ccsrc/ps/worker.h index f07154e8f8..8c8b3922f6 100644 --- a/mindspore/ccsrc/ps/worker.h +++ b/mindspore/ccsrc/ps/worker.h @@ -52,7 +52,7 @@ class Worker { void SetOptimInputShapes(size_t key, const ShapeVector &shape); void AddEmbeddingTable(const ::ps::Key &key, const size_t &row_count); void InitPSEmbeddingTable(const std::vector &keys, std::vector shapes, const ShapeVector &sizes); - void InitPSParamAndOptim(const std::string ¶m_name, const tensor::TensorPtr &tensor); + void InitPSParamAndOptim(const AnfNodePtr &input_node, const tensor::TensorPtr &tensor); void DoPSEmbeddingLookup(const ::ps::SArray<::ps::Key> &keys, const ::ps::SArray &lookup_ids, const ::ps::SArray &lens, ::ps::SArray *lookup_result, int cmd); void Finalize(); @@ -321,14 +321,17 @@ void Worker::InitPSEmbeddingTable(const std::vector &keys, std::vecto } template -void Worker::InitPSParamAndOptim(const std::string ¶m_name, const tensor::TensorPtr &tensor) { +void Worker::InitPSParamAndOptim(const AnfNodePtr &input_node, const tensor::TensorPtr &tensor) { MS_EXCEPTION_IF_NULL(tensor); + MS_EXCEPTION_IF_NULL(input_node); + auto pk_node = input_node->cast(); + MS_EXCEPTION_IF_NULL(pk_node); + const std::string ¶m_name = pk_node->fullname_with_scope(); void *param_data = tensor->data_c(); size_t param_size = LongToSize(tensor->data().nbytes()); if (param_size > INT_MAX) { MS_LOG(EXCEPTION) << "PS mode max weight size is " << INT_MAX << ", " << param_name << " size is " << param_size; } - ShapeVector param_shape = tensor->shape_c(); size_t param_key = GetParamKey(param_name); if (param_key == kInvalidKey) { @@ -336,8 +339,8 @@ void Worker::InitPSParamAndOptim(const std::string ¶m_name, const tensor: return; } bool init_in_server = false; - ShapeVector shape_init_in_server = {1}; - if (param_shape == shape_init_in_server) { + auto param_info_ptr = pk_node->param_info(); + if (param_info_ptr != nullptr && param_info_ptr->init_in_server()) { init_in_server = true; } SetParamInitInServer(param_name, init_in_server); diff --git a/mindspore/ccsrc/pybind_api/ir/param_info_py.cc b/mindspore/ccsrc/pybind_api/ir/param_info_py.cc index f151462653..74f2730cf3 100644 --- a/mindspore/ccsrc/pybind_api/ir/param_info_py.cc +++ b/mindspore/ccsrc/pybind_api/ir/param_info_py.cc @@ -26,6 +26,7 @@ REGISTER_PYBIND_DEFINE(ParamInfo, ([](const py::module *m) { .def("clone", &ParamInfo::Clone) .def_property("name", &ParamInfo::name, &ParamInfo::set_name) .def_property("requires_grad", &ParamInfo::requires_grad, &ParamInfo::set_requires_grad) + .def_property("init_in_server", &ParamInfo::init_in_server, &ParamInfo::set_init_in_server) .def_property("layerwise_parallel", &ParamInfo::layerwise_parallel, &ParamInfo::set_layerwise_parallel) .def(py::pickle( diff --git a/mindspore/common/parameter.py b/mindspore/common/parameter.py index de57ca0bf3..c398643e93 100644 --- a/mindspore/common/parameter.py +++ b/mindspore/common/parameter.py @@ -144,7 +144,7 @@ class Parameter(MetaTensor_): if isinstance(data, bool): raise ValueError('Parameter data can not be `bool`') if isinstance(data, MetaTensor): - if _is_in_parallel_mode(): + if _is_in_parallel_mode() or _is_role_worker(): # do not init data while in auto parallel. return (MetaTensor_, data.dtype, data.shape) data = data.to_tensor() @@ -174,8 +174,12 @@ class Parameter(MetaTensor_): def set_param_ps(self, init_in_server=False): if _is_role_worker() or _is_role_pserver() or _is_role_sched(): + if init_in_server and (not self.name.endswith("embedding_table")): + raise RuntimeError("Can not initialize parameter '{}' in server, only parameters of \ + sparse operator support initialization in server.".format(self.name)) self.is_param_ps = True self.init_in_server = init_in_server + self._param_info.init_in_server = init_in_server else: raise RuntimeError("Must complete following two steps before calling set_param_ps: \ 1. set_ps_context(enable_ps=True) \ @@ -270,6 +274,8 @@ class Parameter(MetaTensor_): x._param_info = self._param_info.clone() x._param_info.name = prefix + '.' + self._param_info.name x.is_init = False + x.is_param_ps = self.is_param_ps + x.init_in_server = self.init_in_server if init != 'same': shape = self.shape dtype = self.dtype @@ -403,12 +409,18 @@ class Parameter(MetaTensor_): raise ValueError("The length of layout must be larger than 3! layout is {}.".format(layout)) slice_index = int(_get_slice_index(layout[0], layout[1])) if (self.init_in_server and self.is_param_ps and isinstance(self.init_mode, MetaTensor)): - data = self.init_mode.to_tensor(0, [1]) + if _is_role_worker(): + data = self.init_mode.to_tensor(0, [1]) + else: + data = self.init_mode.to_tensor(slice_index, layout[2]) else: data = self.init_mode.to_tensor(slice_index, layout[2]) else: if (self.init_in_server and self.is_param_ps and isinstance(self.init_mode, MetaTensor)): - data = self.init_mode.to_tensor(0, [1]) + if _is_role_worker(): + data = self.init_mode.to_tensor(0, [1]) + else: + data = self.init_mode.to_tensor() else: data = self.init_mode.to_tensor() diff --git a/mindspore/core/ir/param_info.h b/mindspore/core/ir/param_info.h index 2f71e60d61..216b1f025c 100644 --- a/mindspore/core/ir/param_info.h +++ b/mindspore/core/ir/param_info.h @@ -42,6 +42,9 @@ class ParamInfo { bool requires_grad() const { return requires_grad_; } void set_requires_grad(bool requires_grad) { requires_grad_ = requires_grad; } + bool init_in_server() const { return init_in_server_; } + void set_init_in_server(bool init_in_server) { init_in_server_ = init_in_server; } + bool layerwise_parallel() const { return layerwise_parallel_; } void set_layerwise_parallel(bool layerwise_parallel) { layerwise_parallel_ = layerwise_parallel; } @@ -68,12 +71,14 @@ class ParamInfo { clone->cloned_index_ = index; this->be_cloned_ = true; this->be_cloned_index_.push_back(index); + clone->init_in_server_ = this->init_in_server_; return clone; } private: std::string name_{"Parameter"}; bool requires_grad_{true}; + bool init_in_server_{false}; bool layerwise_parallel_{false}; bool be_cloned_{false}; bool cloned_{false}; diff --git a/tests/st/ps/cmp_sparse_embedding/test_cmp_sparse_embedding.py b/tests/st/ps/cmp_sparse_embedding/test_cmp_sparse_embedding.py index d001dbe2a2..96ca152754 100644 --- a/tests/st/ps/cmp_sparse_embedding/test_cmp_sparse_embedding.py +++ b/tests/st/ps/cmp_sparse_embedding/test_cmp_sparse_embedding.py @@ -26,7 +26,6 @@ from mindspore.nn import TrainOneStepCell, WithLossCell from mindspore.nn.optim import Adam from mindspore.common import set_seed from mindspore.ops import operations as P -from mindspore.common.initializer import TruncatedNormal from mindspore.parallel._ps_context import _is_role_pserver, _is_role_worker parser = argparse.ArgumentParser(description="test_sparse_embedding") @@ -39,18 +38,6 @@ context.set_context( context.set_ps_context(enable_ps=True) -def fc_with_initialize(input_channels, out_channels): - """weight initial for fc layer""" - weight = weight_variable() - bias = weight_variable() - return nn.Dense(input_channels, out_channels, weight, bias) - - -def weight_variable(): - """weight initial""" - return TruncatedNormal(0.02) - - class LeNet5(nn.Cell): def __init__(self, num_class=10): super(LeNet5, self).__init__() @@ -58,7 +45,7 @@ class LeNet5(nn.Cell): self.flatten = nn.Flatten() self.embedding = nn.EmbeddingLookup(16, 4) self.relu = nn.ReLU() - self.fc = fc_with_initialize(12, num_class) + self.fc = nn.Dense(12, num_class) def construct(self, x): x = self.cast(x, mstype.int32)