add ps mode consistence check

mindspore/ccsrc/backend/session/ascend_session.cc

@@ -45,6 +45,9 @@
 #include "debug/anf_ir_dump.h"
 #include "debug/dump_proto.h"
 #include "toolchain/adx_datadump_server.h"
+#if ENABLE_CPU && ENABLE_D
+#include "ps/util.h"
+#endif
 
 namespace mindspore {
 namespace session {
@@ -166,9 +169,12 @@ GraphId AscendSession::CompileGraphImpl(NotNull<FuncGraphPtr> func_graph) {
   RootGraphExecutorValidate(NOT_NULL(root_graph));
   // adjust kernel
   AdjustKernel(root_graph);
-#if (ENABLE_CPU && (ENABLE_D || ENABLE_GPU))
-  // Assign parameter keys.
-  AssignParamKey(root_graph);
+#if ENABLE_CPU && ENABLE_D
+  if (ps::Util::IsParamServerMode()) {
+    CheckPSModeConsistence(root_graph);
+    // Assign parameter keys.
+    AssignParamKey(root_graph);
+  }
 #endif
   // assign stream
   AssignStream(NOT_NULL(root_graph));
@@ -314,7 +320,7 @@ void AscendSession::RunGraphImpl(const GraphId &graph_id, const std::vector<tens
   if (debugger_) {
     debugger_->PreExecute(kernel_graph);
   }
-#if (ENABLE_CPU && (ENABLE_D || ENABLE_GPU))
+#if ENABLE_CPU && ENABLE_D
   // Initialize parameter server
   InitPSParamAndOptim(kernel_graph, inputs);
 #endif

mindspore/ccsrc/backend/session/cpu_session.cc

@@ -68,9 +68,11 @@ GraphId CPUSession::CompileGraphImpl(const AnfNodePtrList &lst, const AnfNodePtr
   MS_LOG(INFO) << "Set kernel info";
   SetKernelInfo(graph.get());
 #if (ENABLE_CPU && (ENABLE_D || ENABLE_GPU))
-  AssignParamKey(graph);
-  if (ps::Util::IsRoleOfWorker()) {
-    Optimize(graph);
+  if (ps::Util::IsParamServerMode()) {
+    AssignParamKey(graph);
+    if (ps::Util::IsRoleOfWorker()) {
+      Optimize(graph);
+    }
   }
 #endif
   MS_LOG(INFO) << "Build kernel";

mindspore/ccsrc/backend/session/gpu_session.cc

@@ -56,6 +56,9 @@
 #include "utils/ms_utils.h"
 #include "utils/config_manager.h"
 #include "utils/ms_context.h"
+#if ENABLE_CPU && ENABLE_GPU
+#include "ps/util.h"
+#endif
 
 namespace mindspore {
 namespace session {
@@ -255,9 +258,12 @@ GraphId GPUSession::CompileGraphImpl(const AnfNodePtrList &lst, const AnfNodePtr
   // Graph kernel fusion optimization
   GraphKernelOptimize(graph);
 
-#if (ENABLE_CPU && (ENABLE_D || ENABLE_GPU))
-  // Assign parameter keys.
-  AssignParamKey(graph);
+#if ENABLE_CPU && ENABLE_GPU
+  if (ps::Util::IsParamServerMode()) {
+    CheckPSModeConsistence(graph);
+    // Assign parameter keys.
+    AssignParamKey(graph);
+  }
 #endif
   // Start gpu kernel runtime
   StartKernelRT();
@@ -299,7 +305,7 @@ void GPUSession::RunGraphImpl(const GraphId &graph_id, const std::vector<tensor:
   // Load input data from user input
   LoadInputData(kernel_graph, inputs);
   PreIterationDbg(kernel_graph);
-#if (ENABLE_CPU && (ENABLE_D || ENABLE_GPU))
+#if ENABLE_CPU && ENABLE_GPU
   // Initialize parameter server
   InitPSParamAndOptim(kernel_graph, inputs);
 #endif

mindspore/ccsrc/backend/session/session_basic.cc

@@ -1496,6 +1496,23 @@ void SessionBasic::UpdateGraphDynamicShapeAttr(const NotNull<KernelGraphPtr> &ro
 }
 
 #if (ENABLE_CPU && (ENABLE_D || ENABLE_GPU))
+void SessionBasic::CheckPSModeConsistence(const KernelGraphPtr &kernel_graph) {
+  auto input_nodes = kernel_graph->inputs();
+  for (const auto &input_node : input_nodes) {
+    if (!input_node->isa<Parameter>()) {
+      continue;
+    }
+    auto pk_node = input_node->cast<ParameterPtr>();
+    MS_EXCEPTION_IF_NULL(pk_node);
+    auto param_info_ptr = pk_node->param_info();
+    if (param_info_ptr != nullptr && param_info_ptr->init_in_server()) {
+      const std::string &param_name = pk_node->fullname_with_scope();
+      MS_LOG(EXCEPTION) << "Can not initialize the parameter[" << param_name
+                        << "] in server, this parameter is used by kernel which executes in device";
+    }
+  }
+}
+
 void SessionBasic::AssignParamKey(const KernelGraphPtr &kernel_graph) {
   if (!ps::Util::IsRoleOfWorker()) {
     MS_LOG(INFO) << "Not parameter server mode.";

mindspore/ccsrc/backend/session/session_basic.h

@@ -91,6 +91,7 @@ class SessionBasic : public std::enable_shared_from_this<SessionBasic> {
   // get graph id in child graphs by ME front anf node pointer
   virtual GraphId GetGraphIdByNode(const AnfNodePtr &) const { return kInvalidGraphId; }
   virtual GraphId GetFinalRunGraph() const { return kInvalidGraphId; }
+  void CheckPSModeConsistence(const KernelGraphPtr &Kernel_graph);
   void AssignParamKey(const KernelGraphPtr &kernel_graph);
   void InitPSParamAndOptim(const KernelGraphPtr &kernel_graph, const std::vector<tensor::TensorPtr> &inputs_const);
   virtual bool CheckModelInputs(uint32_t graph_id, const std::vector<tensor::TensorPtr> &inputs,

mindspore/common/parameter.py

@@ -169,8 +169,8 @@ 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))
+                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