cross silo

mindspore/ccsrc/backend/kernel_compiler/CMakeLists.txt

@@ -80,6 +80,9 @@ if(NOT ENABLE_CPU OR WIN32)
     list(REMOVE_ITEM CPU_SRC_LIST "cpu/ps/sparse_apply_lazy_adam_ps_kernel.cc")
     list(REMOVE_ITEM CPU_SRC_LIST "cpu/fl/fused_pull_weight_kernel.cc")
     list(REMOVE_ITEM CPU_SRC_LIST "cpu/fl/fused_push_weight_kernel.cc")
+    list(REMOVE_ITEM CPU_SRC_LIST "cpu/fl/get_model_kernel.cc")
+    list(REMOVE_ITEM CPU_SRC_LIST "cpu/fl/start_fl_job_kernel.cc")
+    list(REMOVE_ITEM CPU_SRC_LIST "cpu/fl/update_model_kernel.cc")
 endif()
 
 if(ENABLE_GPU)

mindspore/ccsrc/backend/kernel_compiler/cpu/fl/get_model_kernel.cc

@@ -0,0 +1,25 @@
+/**
+ * Copyright 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.
+ * 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.
+ */
+
+#include "backend/kernel_compiler/cpu/fl/get_model_kernel.h"
+
+namespace mindspore {
+namespace kernel {
+MS_REG_CPU_KERNEL(GetModel,
+                  KernelAttr().SetAllSameAttr(true).AddInputAttr(kNumberTypeFloat32).AddOutputAttr(kNumberTypeFloat32),
+                  GetModelKernel);
+}  // namespace kernel
+}  // namespace mindspore

mindspore/ccsrc/backend/kernel_compiler/cpu/fl/get_model_kernel.h

@@ -0,0 +1,162 @@
+/**
+ * Copyright 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.
+ * 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.
+ */
+
+#ifndef MINDSPORE_CCSRC_BACKEND_KERNEL_COMPILER_FL_GET_MODEL_H_
+#define MINDSPORE_CCSRC_BACKEND_KERNEL_COMPILER_FL_GET_MODEL_H_
+
+#include <map>
+#include <vector>
+#include <string>
+#include <memory>
+#include <utility>
+#include <functional>
+#include "backend/kernel_compiler/cpu/cpu_kernel.h"
+#include "backend/kernel_compiler/cpu/cpu_kernel_factory.h"
+#include "fl/worker/fl_worker.h"
+
+namespace mindspore {
+namespace kernel {
+class GetModelKernel : public CPUKernel {
+ public:
+  GetModelKernel() = default;
+  ~GetModelKernel() override = default;
+
+  bool Launch(const std::vector<AddressPtr> &inputs, const std::vector<AddressPtr> &, const std::vector<AddressPtr> &) {
+    MS_LOG(INFO) << "Launching client GetModelKernel";
+    if (!BuildGetModelReq(fbb_, inputs)) {
+      MS_LOG(EXCEPTION) << "Building request for FusedPushWeight failed.";
+      return false;
+    }
+
+    const schema::ResponseGetModel *get_model_rsp = nullptr;
+    int response_code = schema::ResponseCode_SucNotReady;
+    while (response_code == schema::ResponseCode_SucNotReady) {
+      std::shared_ptr<std::vector<unsigned char>> get_model_rsp_msg = nullptr;
+      if (!fl::worker::FLWorker::GetInstance().SendToServer(target_server_rank_, fbb_->GetBufferPointer(),
+                                                            fbb_->GetSize(), ps::core::TcpUserCommand::kGetModel,
+                                                            &get_model_rsp_msg)) {
+        MS_LOG(EXCEPTION) << "Sending request for GetModel to server " << target_server_rank_ << " failed.";
+        return false;
+      }
+      flatbuffers::Verifier verifier(get_model_rsp_msg->data(), get_model_rsp_msg->size());
+      if (!verifier.VerifyBuffer<schema::ResponseGetModel>()) {
+        MS_LOG(EXCEPTION) << "The schema of ResponseGetModel is invalid.";
+        return false;
+      }
+
+      get_model_rsp = flatbuffers::GetRoot<schema::ResponseGetModel>(get_model_rsp_msg->data());
+      MS_EXCEPTION_IF_NULL(get_model_rsp);
+      response_code = get_model_rsp->retcode();
+      if (response_code == schema::ResponseCode_SUCCEED) {
+        break;
+      } else if (response_code == schema::ResponseCode_SucNotReady) {
+        std::this_thread::sleep_for(std::chrono::milliseconds(200));
+        continue;
+      } else {
+        MS_LOG(EXCEPTION) << "Launching get model for worker failed. Reason: " << get_model_rsp->reason();
+      }
+    }
+
+    auto feature_map = get_model_rsp->feature_map();
+    MS_EXCEPTION_IF_NULL(feature_map);
+    if (feature_map->size() == 0) {
+      MS_LOG(EXCEPTION) << "Feature map after GetModel is empty.";
+      return false;
+    }
+    for (size_t i = 0; i < feature_map->size(); i++) {
+      std::string weight_full_name = feature_map->Get(i)->weight_fullname()->str();
+      float *weight_data = const_cast<float *>(feature_map->Get(i)->data()->data());
+      size_t weight_size = feature_map->Get(i)->data()->size() * sizeof(float);
+      if (weight_name_to_input_idx_.count(weight_full_name) == 0) {
+        MS_LOG(EXCEPTION) << "Weight " << weight_full_name << " doesn't exist in FL worker.";
+        return false;
+      }
+      MS_LOG(INFO) << "Cover weight " << weight_full_name << " by the model in server.";
+      size_t index = weight_name_to_input_idx_[weight_full_name];
+      int ret = memcpy_s(inputs[index]->addr, inputs[index]->size, weight_data, weight_size);
+      if (ret != 0) {
+        MS_LOG(EXCEPTION) << "memcpy_s error, errorno(" << ret << ")";
+        return false;
+      }
+    }
+    return true;
+  }
+
+  void Init(const CNodePtr &kernel_node) {
+    MS_LOG(INFO) << "Initializing GetModel kernel";
+    fbb_ = std::make_shared<fl::FBBuilder>();
+    MS_EXCEPTION_IF_NULL(fbb_);
+
+    MS_EXCEPTION_IF_NULL(kernel_node);
+    server_num_ = fl::worker::FLWorker::GetInstance().server_num();
+    rank_id_ = fl::worker::FLWorker::GetInstance().rank_id();
+    if (rank_id_ == UINT32_MAX) {
+      MS_LOG(EXCEPTION) << "Federated worker is not initialized yet.";
+      return;
+    }
+    target_server_rank_ = rank_id_ % server_num_;
+    fl_name_ = fl::worker::FLWorker::GetInstance().fl_name();
+    MS_LOG(INFO) << "Initializing GetModel kernel. fl_name: " << fl_name_ << ". Request will be sent to server "
+                 << target_server_rank_;
+
+    size_t input_num = AnfAlgo::GetInputTensorNum(kernel_node);
+    for (size_t i = 0; i < input_num; i++) {
+      auto input_node = AnfAlgo::VisitKernelWithReturnType(AnfAlgo::GetInputNode(kernel_node, i), 0).first;
+      MS_EXCEPTION_IF_NULL(input_node);
+      auto weight_node = input_node->cast<ParameterPtr>();
+      MS_EXCEPTION_IF_NULL(weight_node);
+      std::string weight_name = weight_node->fullname_with_scope();
+      MS_LOG(INFO) << "Parameter name is " << weight_name;
+      weight_name_to_input_idx_.insert(std::make_pair(weight_name, i));
+
+      auto weight_shape = AnfAlgo::GetPrevNodeOutputInferShape(kernel_node, i);
+      size_t weight_size_ =
+        std::accumulate(weight_shape.begin(), weight_shape.end(), sizeof(float), std::multiplies<float>());
+      input_size_list_.push_back(weight_size_);
+    }
+    output_size_list_.push_back(sizeof(float));
+  }
+
+  void InitKernel(const CNodePtr &kernel_node) { return; }
+
+ protected:
+  void InitSizeLists() { return; }
+
+ private:
+  bool BuildGetModelReq(const std::shared_ptr<fl::FBBuilder> &fbb, const std::vector<AddressPtr> &weights) {
+    MS_EXCEPTION_IF_NULL(fbb_);
+    auto fbs_fl_name = fbb->CreateString(fl_name_);
+    schema::RequestGetModelBuilder req_get_model_builder(*(fbb.get()));
+    req_get_model_builder.add_fl_name(fbs_fl_name);
+    iteration_ = fl::worker::FLWorker::GetInstance().fl_iteration_num();
+    req_get_model_builder.add_iteration(SizeToInt(iteration_));
+    auto req_get_model = req_get_model_builder.Finish();
+    fbb->Finish(req_get_model);
+    return true;
+  }
+
+  std::shared_ptr<fl::FBBuilder> fbb_;
+  uint32_t rank_id_;
+  uint32_t server_num_;
+  uint32_t target_server_rank_;
+  std::string fl_name_;
+  uint64_t iteration_;
+  std::map<std::string, size_t> weight_name_to_input_idx_;
+};
+}  // namespace kernel
+}  // namespace mindspore
+
+#endif  // MINDSPORE_CCSRC_BACKEND_KERNEL_COMPILER_FL_GET_MODEL_H_

mindspore/ccsrc/backend/kernel_compiler/cpu/fl/start_fl_job_kernel.cc

@@ -0,0 +1,23 @@
+/**
+ * Copyright 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.
+ * 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.
+ */
+
+#include "backend/kernel_compiler/cpu/fl/start_fl_job_kernel.h"
+
+namespace mindspore {
+namespace kernel {
+MS_REG_CPU_KERNEL(StartFLJob, KernelAttr().AddOutputAttr(kNumberTypeFloat32), StartFLJobKernel);
+}  // namespace kernel
+}  // namespace mindspore

mindspore/ccsrc/backend/kernel_compiler/cpu/fl/start_fl_job_kernel.h

@@ -0,0 +1,125 @@
+/**
+ * Copyright 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.
+ * 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.
+ */
+
+#ifndef MINDSPORE_CCSRC_BACKEND_KERNEL_COMPILER_FL_START_FL_JOB_H_
+#define MINDSPORE_CCSRC_BACKEND_KERNEL_COMPILER_FL_START_FL_JOB_H_
+
+#include <vector>
+#include <string>
+#include <memory>
+#include "backend/kernel_compiler/cpu/cpu_kernel.h"
+#include "backend/kernel_compiler/cpu/cpu_kernel_factory.h"
+#include "fl/worker/fl_worker.h"
+
+namespace mindspore {
+namespace kernel {
+class StartFLJobKernel : public CPUKernel {
+ public:
+  StartFLJobKernel() = default;
+  ~StartFLJobKernel() override = default;
+
+  bool Launch(const std::vector<AddressPtr> &inputs, const std::vector<AddressPtr> &, const std::vector<AddressPtr> &) {
+    MS_LOG(INFO) << "Launching client StartFLJobKernel";
+    if (!BuildStartFLJobReq(fbb_)) {
+      MS_LOG(EXCEPTION) << "Building request for StartFLJob failed.";
+      return false;
+    }
+
+    std::shared_ptr<std::vector<unsigned char>> start_fl_job_rsp_msg = nullptr;
+    if (!fl::worker::FLWorker::GetInstance().SendToServer(target_server_rank_, fbb_->GetBufferPointer(),
+                                                          fbb_->GetSize(), ps::core::TcpUserCommand::kStartFLJob,
+                                                          &start_fl_job_rsp_msg)) {
+      MS_LOG(EXCEPTION) << "Sending request for StartFLJob to server " << target_server_rank_ << " failed.";
+      return false;
+    }
+    flatbuffers::Verifier verifier(start_fl_job_rsp_msg->data(), start_fl_job_rsp_msg->size());
+    if (!verifier.VerifyBuffer<schema::ResponseFLJob>()) {
+      MS_LOG(EXCEPTION) << "The schema of ResponseFLJob is invalid.";
+      return false;
+    }
+
+    const schema::ResponseFLJob *start_fl_job_rsp =
+      flatbuffers::GetRoot<schema::ResponseFLJob>(start_fl_job_rsp_msg->data());
+    MS_EXCEPTION_IF_NULL(start_fl_job_rsp);
+    auto response_code = start_fl_job_rsp->retcode();
+    switch (response_code) {
+      case schema::ResponseCode_SUCCEED:
+      case schema::ResponseCode_OutOfTime:
+        break;
+      default:
+        MS_LOG(EXCEPTION) << "Launching start fl job for worker failed. Reason: " << start_fl_job_rsp->reason();
+    }
+
+    uint64_t iteration = IntToSize(start_fl_job_rsp->iteration());
+    fl::worker::FLWorker::GetInstance().set_fl_iteration_num(iteration);
+    MS_LOG(INFO) << "Start fl job for iteration " << iteration;
+    return true;
+  }
+
+  void Init(const CNodePtr &kernel_node) {
+    MS_EXCEPTION_IF_NULL(kernel_node);
+    server_num_ = fl::worker::FLWorker::GetInstance().server_num();
+    rank_id_ = fl::worker::FLWorker::GetInstance().rank_id();
+    if (rank_id_ == UINT32_MAX) {
+      MS_LOG(EXCEPTION) << "Federated worker is not initialized yet.";
+      return;
+    }
+    target_server_rank_ = rank_id_ % server_num_;
+    fl_name_ = fl::worker::FLWorker::GetInstance().fl_name();
+    fl_id_ = fl::worker::FLWorker::GetInstance().fl_id();
+    data_size_ = LongToInt(AnfAlgo::GetNodeAttr<int64_t>(kernel_node, "data_size"));
+    fl::worker::FLWorker::GetInstance().set_data_size(data_size_);
+    MS_LOG(INFO) << "Initializing StartFLJob kernel. fl_name: " << fl_name_ << ", fl_id: " << fl_id_
+                 << ", data_size: " << data_size_ << ". Request will be sent to server " << target_server_rank_;
+
+    fbb_ = std::make_shared<fl::FBBuilder>();
+    MS_EXCEPTION_IF_NULL(fbb_);
+
+    input_size_list_.push_back(sizeof(int));
+    output_size_list_.push_back(sizeof(float));
+  }
+
+  void InitKernel(const CNodePtr &kernel_node) { return; }
+
+ protected:
+  void InitSizeLists() { return; }
+
+ private:
+  bool BuildStartFLJobReq(const std::shared_ptr<fl::FBBuilder> &fbb) {
+    MS_EXCEPTION_IF_NULL(fbb);
+    auto fbs_fl_name = fbb->CreateString(fl_name_);
+    auto fbs_fl_id = fbb->CreateString(fl_id_);
+    schema::RequestFLJobBuilder req_fl_job_builder(*(fbb.get()));
+    req_fl_job_builder.add_fl_name(fbs_fl_name);
+    req_fl_job_builder.add_fl_id(fbs_fl_id);
+    req_fl_job_builder.add_data_size(data_size_);
+    auto req_fl_job = req_fl_job_builder.Finish();
+    fbb->Finish(req_fl_job);
+    return true;
+  }
+
+  uint32_t rank_id_;
+  uint32_t server_num_;
+  uint32_t target_server_rank_;
+  std::string fl_name_;
+  std::string fl_id_;
+  int data_size_;
+  std::shared_ptr<fl::FBBuilder> fbb_;
+};
+}  // namespace kernel
+}  // namespace mindspore
+
+#endif  // MINDSPORE_CCSRC_BACKEND_KERNEL_COMPILER_FL_START_FL_JOB_H_

mindspore/ccsrc/backend/kernel_compiler/cpu/fl/update_model_kernel.cc

@@ -0,0 +1,25 @@
+/**
+ * Copyright 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.
+ * 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.
+ */
+
+#include "backend/kernel_compiler/cpu/fl/update_model_kernel.h"
+
+namespace mindspore {
+namespace kernel {
+MS_REG_CPU_KERNEL(UpdateModel,
+                  KernelAttr().SetAllSameAttr(true).AddInputAttr(kNumberTypeFloat32).AddOutputAttr(kNumberTypeFloat32),
+                  UpdateModelKernel);
+}  // namespace kernel
+}  // namespace mindspore

mindspore/ccsrc/backend/kernel_compiler/cpu/fl/update_model_kernel.h

@@ -0,0 +1,172 @@
+/**
+ * Copyright 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.
+ * 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.
+ */
+
+#ifndef MINDSPORE_CCSRC_BACKEND_KERNEL_COMPILER_FL_UPDATE_MODEL_H_
+#define MINDSPORE_CCSRC_BACKEND_KERNEL_COMPILER_FL_UPDATE_MODEL_H_
+
+#include <vector>
+#include <string>
+#include <memory>
+#include <functional>
+#include "backend/kernel_compiler/cpu/cpu_kernel.h"
+#include "backend/kernel_compiler/cpu/cpu_kernel_factory.h"
+#include "fl/worker/fl_worker.h"
+
+namespace mindspore {
+namespace kernel {
+class UpdateModelKernel : public CPUKernel {
+ public:
+  UpdateModelKernel() = default;
+  ~UpdateModelKernel() override = default;
+
+  bool Launch(const std::vector<AddressPtr> &inputs, const std::vector<AddressPtr> &, const std::vector<AddressPtr> &) {
+    MS_LOG(INFO) << "Launching client UpdateModelKernel";
+    if (inputs.size() != weight_full_names_.size()) {
+      MS_LOG(EXCEPTION) << "Input number of UpdateModelKernel should be " << weight_full_names_.size() << ", but got "
+                        << inputs.size();
+      return false;
+    }
+
+    if (!WeightingData(inputs)) {
+      MS_LOG(EXCEPTION) << "Weighting data with data_size failed.";
+      return false;
+    }
+
+    if (!BuildUpdateModelReq(fbb_, inputs)) {
+      MS_LOG(EXCEPTION) << "Building request for FusedPushWeight failed.";
+      return false;
+    }
+
+    std::shared_ptr<std::vector<unsigned char>> update_model_rsp_msg = nullptr;
+    if (!fl::worker::FLWorker::GetInstance().SendToServer(target_server_rank_, fbb_->GetBufferPointer(),
+                                                          fbb_->GetSize(), ps::core::TcpUserCommand::kUpdateModel,
+                                                          &update_model_rsp_msg)) {
+      MS_LOG(EXCEPTION) << "Sending request for UpdateModel to server " << target_server_rank_ << " failed.";
+      return false;
+    }
+    flatbuffers::Verifier verifier(update_model_rsp_msg->data(), update_model_rsp_msg->size());
+    if (!verifier.VerifyBuffer<schema::ResponseUpdateModel>()) {
+      MS_LOG(EXCEPTION) << "The schema of ResponseUpdateModel is invalid.";
+      return false;
+    }
+
+    const schema::ResponseFLJob *update_model_rsp =
+      flatbuffers::GetRoot<schema::ResponseFLJob>(update_model_rsp_msg->data());
+    MS_EXCEPTION_IF_NULL(update_model_rsp);
+    auto response_code = update_model_rsp->retcode();
+    switch (response_code) {
+      case schema::ResponseCode_SUCCEED:
+      case schema::ResponseCode_OutOfTime:
+        break;
+      default:
+        MS_LOG(EXCEPTION) << "Launching start fl job for worker failed. Reason: " << update_model_rsp->reason();
+    }
+    return true;
+  }
+
+  void Init(const CNodePtr &kernel_node) {
+    MS_LOG(INFO) << "Initializing UpdateModel kernel";
+    fbb_ = std::make_shared<fl::FBBuilder>();
+    MS_EXCEPTION_IF_NULL(fbb_);
+
+    MS_EXCEPTION_IF_NULL(kernel_node);
+    server_num_ = fl::worker::FLWorker::GetInstance().server_num();
+    rank_id_ = fl::worker::FLWorker::GetInstance().rank_id();
+    if (rank_id_ == UINT32_MAX) {
+      MS_LOG(EXCEPTION) << "Federated worker is not initialized yet.";
+      return;
+    }
+    target_server_rank_ = rank_id_ % server_num_;
+    fl_name_ = fl::worker::FLWorker::GetInstance().fl_name();
+    fl_id_ = fl::worker::FLWorker::GetInstance().fl_id();
+    MS_LOG(INFO) << "Initializing StartFLJob kernel. fl_name: " << fl_name_ << ", fl_id: " << fl_id_
+                 << ". Request will be sent to server " << target_server_rank_;
+
+    size_t input_num = AnfAlgo::GetInputTensorNum(kernel_node);
+    for (size_t i = 0; i < input_num; i++) {
+      auto input_node = AnfAlgo::VisitKernelWithReturnType(AnfAlgo::GetInputNode(kernel_node, i), 0).first;
+      MS_EXCEPTION_IF_NULL(input_node);
+      auto weight_node = input_node->cast<ParameterPtr>();
+      MS_EXCEPTION_IF_NULL(weight_node);
+      std::string weight_name = weight_node->fullname_with_scope();
+      MS_LOG(INFO) << "Parameter name is " << weight_name;
+      weight_full_names_.push_back(weight_name);
+
+      auto weight_shape = AnfAlgo::GetPrevNodeOutputInferShape(kernel_node, i);
+      size_t weight_size_ =
+        std::accumulate(weight_shape.begin(), weight_shape.end(), sizeof(float), std::multiplies<float>());
+      input_size_list_.push_back(weight_size_);
+    }
+    output_size_list_.push_back(sizeof(float));
+  }
+
+  void InitKernel(const CNodePtr &kernel_node) { return; }
+
+ protected:
+  void InitSizeLists() { return; }
+
+ private:
+  bool BuildUpdateModelReq(const std::shared_ptr<fl::FBBuilder> &fbb, const std::vector<AddressPtr> &weights) {
+    MS_EXCEPTION_IF_NULL(fbb_);
+    auto fbs_fl_name = fbb->CreateString(fl_name_);
+    auto fbs_fl_id = fbb->CreateString(fl_id_);
+    std::vector<flatbuffers::Offset<schema::FeatureMap>> fbs_feature_maps;
+    for (size_t i = 0; i < weight_full_names_.size(); i++) {
+      const std::string &weight_name = weight_full_names_[i];
+      auto fbs_weight_fullname = fbb->CreateString(weight_name);
+      auto fbs_weight_data =
+        fbb->CreateVector(reinterpret_cast<const float *>(weights[i]->addr), weights[i]->size / sizeof(float));
+      auto fbs_feature_map = schema::CreateFeatureMap(*(fbb.get()), fbs_weight_fullname, fbs_weight_data);
+      fbs_feature_maps.push_back(fbs_feature_map);
+    }
+    auto fbs_feature_maps_vector = fbb->CreateVector(fbs_feature_maps);
+
+    schema::RequestUpdateModelBuilder req_update_model_builder(*(fbb.get()));
+    req_update_model_builder.add_fl_name(fbs_fl_name);
+    req_update_model_builder.add_fl_id(fbs_fl_id);
+    iteration_ = fl::worker::FLWorker::GetInstance().fl_iteration_num();
+    req_update_model_builder.add_iteration(SizeToInt(iteration_));
+    req_update_model_builder.add_feature_map(fbs_feature_maps_vector);
+    auto req_update_model = req_update_model_builder.Finish();
+    fbb->Finish(req_update_model);
+    return true;
+  }
+
+  bool WeightingData(const std::vector<AddressPtr> &inputs) {
+    data_size_ = fl::worker::FLWorker::GetInstance().data_size();
+    for (auto &input : inputs) {
+      float *data = reinterpret_cast<float *>(input->addr);
+      for (size_t i = 0; i < input->size / sizeof(float); i++) {
+        data[i] *= data_size_;
+      }
+    }
+    return true;
+  }
+
+  std::shared_ptr<fl::FBBuilder> fbb_;
+  uint32_t rank_id_;
+  uint32_t server_num_;
+  uint32_t target_server_rank_;
+  std::string fl_name_;
+  std::string fl_id_;
+  int data_size_;
+  uint64_t iteration_;
+  std::vector<std::string> weight_full_names_;
+};
+}  // namespace kernel
+}  // namespace mindspore
+
+#endif  // MINDSPORE_CCSRC_BACKEND_KERNEL_COMPILER_FL_UPDATE_MODEL_H_

mindspore/ccsrc/backend/optimizer/pass/convert_tuple_input_to_dynamic_input.cc

@@ -65,7 +65,7 @@ void ConvertMakeTupleInputToPlantInputs(const FuncGraphPtr &graph, const CNodePt
   std::vector<AnfNodePtr> plant_inputs;
   std::vector<int64_t> dyn_input_sizes;
   plant_inputs.push_back(AnfAlgo::GetCNodePrimitiveNode(cnode_ptr));
-  size_t input_num = AnfAlgo::GetInputTensorNum(cnode_ptr);
+  size_t input_num = cnode_ptr->inputs().size() - 1;
   for (size_t i = 0; i < input_num; ++i) {
     auto input_node = AnfAlgo::GetInputNode(cnode_ptr, i);
     MS_EXCEPTION_IF_NULL(input_node);

mindspore/ccsrc/fl/server/kernel/round/get_model_kernel.cc

@@ -54,6 +54,10 @@ bool GetModelKernel::Launch(const std::vector<AddressPtr> &inputs, const std::ve
   }
 
   const schema::RequestGetModel *get_model_req = flatbuffers::GetRoot<schema::RequestGetModel>(req_data);
+  if (get_model_req == nullptr) {
+    MS_LOG(ERROR) << "RequestGetModel is nullptr.";
+    return false;
+  }
   GetModel(get_model_req, fbb);
   GenerateOutput(outputs, fbb->GetBufferPointer(), fbb->GetSize());
   return true;

mindspore/ccsrc/fl/worker/fl_worker.cc

@@ -25,6 +25,10 @@ namespace mindspore {
 namespace fl {
 namespace worker {
 void FLWorker::Run() {
+  if (running_) {
+    return;
+  }
+  running_ = true;
   worker_num_ = ps::PSContext::instance()->worker_num();
   server_num_ = ps::PSContext::instance()->server_num();
   scheduler_ip_ = ps::PSContext::instance()->scheduler_ip();
@@ -64,6 +68,8 @@ void FLWorker::Run() {
 
   InitializeFollowerScaler();
   worker_node_->Start();
+  rank_id_ = worker_node_->rank_id();
+
   std::this_thread::sleep_for(std::chrono::milliseconds(kWorkerSleepTimeForNetworking));
   return;
 }
@@ -133,6 +139,8 @@ uint32_t FLWorker::server_num() const { return server_num_; }
 
 uint32_t FLWorker::worker_num() const { return worker_num_; }
 
+uint32_t FLWorker::rank_id() const { return rank_id_; }
+
 uint64_t FLWorker::worker_step_num_per_iteration() const { return worker_step_num_per_iteration_; }
 
 void FLWorker::SetIterationRunning() {
@@ -145,6 +153,18 @@ void FLWorker::SetIterationCompleted() {
   worker_iteration_state_ = IterationState::kCompleted;
 }
 
+void FLWorker::set_fl_iteration_num(uint64_t iteration_num) { iteration_num_ = iteration_num; }
+
+uint64_t FLWorker::fl_iteration_num() const { return iteration_num_.load(); }
+
+void FLWorker::set_data_size(int data_size) { data_size_ = data_size; }
+
+int FLWorker::data_size() const { return data_size_; }
+
+std::string FLWorker::fl_name() const { return ps::kServerModeFL; }
+
+std::string FLWorker::fl_id() const { return std::to_string(rank_id_); }
+
 void FLWorker::InitializeFollowerScaler() {
   if (!worker_node_->InitFollowerScaler()) {
     MS_LOG(EXCEPTION) << "Initializing follower elastic scaler failed.";

mindspore/ccsrc/fl/worker/fl_worker.h

@@ -22,6 +22,7 @@
 #include <vector>
 #include "proto/comm.pb.h"
 #include "schema/fl_job_generated.h"
+#include "schema/cipher_generated.h"
 #include "ps/ps_context.h"
 #include "ps/core/worker_node.h"
 #include "ps/core/cluster_metadata.h"
@@ -64,12 +65,22 @@ class FLWorker {
 
   uint32_t server_num() const;
   uint32_t worker_num() const;
+  uint32_t rank_id() const;
   uint64_t worker_step_num_per_iteration() const;
 
   // These methods set the worker's iteration state.
   void SetIterationRunning();
   void SetIterationCompleted();
 
+  void set_fl_iteration_num(uint64_t iteration_num);
+  uint64_t fl_iteration_num() const;
+
+  void set_data_size(int data_size);
+  int data_size() const;
+
+  std::string fl_name() const;
+  std::string fl_id() const;
+
  private:
   FLWorker()
       : server_num_(0),
@@ -77,6 +88,7 @@ class FLWorker {
         scheduler_ip_(""),
         scheduler_port_(0),
         worker_node_(nullptr),
+        rank_id_(UINT32_MAX),
         worker_step_num_per_iteration_(1),
         server_iteration_state_(IterationState::kCompleted),
         worker_iteration_state_(IterationState::kCompleted),
@@ -100,11 +112,19 @@ class FLWorker {
   void ProcessAfterScalingOut();
   void ProcessAfterScalingIn();
 
+  bool running_;
   uint32_t server_num_;
   uint32_t worker_num_;
   std::string scheduler_ip_;
   uint16_t scheduler_port_;
   std::shared_ptr<ps::core::WorkerNode> worker_node_;
+  uint32_t rank_id_;
+
+  // The federated learning iteration number.
+  std::atomic<uint64_t> iteration_num_;
+
+  // Data size for this federated learning job.
+  int data_size_;
 
   // The worker standalone training step number before communicating with server. This used in hybrid training mode.
   uint64_t worker_step_num_per_iteration_;

mindspore/ccsrc/ps/core/communicator/tcp_communicator.h

@@ -51,7 +51,11 @@ enum class TcpUserCommand {
   kNotifyLeaderToNextIter,
   kPrepareForNextIter,
   kProceedToNextIter,
-  kEndLastIter
+  kEndLastIter,
+
+  kStartFLJob,
+  kUpdateModel,
+  kGetModel
 };
 
 const std::unordered_map<TcpUserCommand, std::string> kUserCommandToMsgType = {
@@ -69,7 +73,10 @@ const std::unordered_map<TcpUserCommand, std::string> kUserCommandToMsgType = {
   {TcpUserCommand::kNotifyLeaderToNextIter, "notifyLeaderToNextIter"},
   {TcpUserCommand::kPrepareForNextIter, "prepareForNextIter"},
   {TcpUserCommand::kProceedToNextIter, "proceedToNextIter"},
-  {TcpUserCommand::kEndLastIter, "endLastIter"}};
+  {TcpUserCommand::kEndLastIter, "endLastIter"},
+  {TcpUserCommand::kStartFLJob, "startFLJob"},
+  {TcpUserCommand::kUpdateModel, "updateModel"},
+  {TcpUserCommand::kGetModel, "getModel"}};
 
 class TcpCommunicator : public CommunicatorBase {
  public:

mindspore/ccsrc/ps/ps_context.cc

@@ -243,10 +243,6 @@ void PSContext::set_worker_num(uint32_t worker_num) {
     MS_LOG(EXCEPTION) << "The worker number should be set to 1 in hybrid training mode.";
     return;
   }
-  if (server_mode_ == kServerModeFL && worker_num != 0) {
-    MS_LOG(EXCEPTION) << "The worker number should be 0 in federated learning mode.";
-    return;
-  }
   worker_num_ = worker_num;
 }
 uint32_t PSContext::worker_num() const { return worker_num_; }

mindspore/ccsrc/runtime/device/cpu/kernel_select_cpu.cc

@@ -267,8 +267,7 @@ void SetKernelInfo(const CNodePtr &kernel_node) {
     }
   }
 
-  if (selected_kernel_attr.GetInputSize() > 0 &&
-      (matched.first || input_types.size() == input_not_cnode_indexes.size())) {
+  if (matched.first || input_types.size() == input_not_cnode_indexes.size()) {
     MS_LOG(INFO) << "Input format and dtype is matched";
     GetOutputFormatsAndDtypes(kernel_node, selected_kernel_attr, &selected_output_formats, &selected_output_types);
     UpdatePrevNotCNodeFormatDtype(selected_kernel_attr, input_not_cnode_indexes, kernel_node);

mindspore/ops/operations/__init__.py

@@ -91,7 +91,8 @@ from . import _quant_ops
 from ._quant_ops import *
 from .other_ops import (Assign, InplaceAssign, IOU, BoundingBoxDecode, BoundingBoxEncode,
                         ConfusionMatrix, PopulationCount, UpdateState, Load,
-                        CheckValid, Partial, Depend, identity, CheckBprop, Push, Pull, PullWeight, PushWeight)
+                        CheckValid, Partial, Depend, identity, CheckBprop, Push, Pull, PullWeight, PushWeight,
+                        StartFLJob, UpdateModel, GetModel)
 from ._thor_ops import (CusBatchMatMul, CusCholeskyTrsm, CusFusedAbsMax1, CusImg2Col, CusMatMulCubeDenseLeft,
                         CusMatMulCubeFraczRightMul, CusMatMulCube, CusMatrixCombine, CusTranspose02314,
                         CusMatMulCubeDenseRight,

mindspore/ops/operations/other_ops.py

@@ -770,7 +770,7 @@ class PushWeight(PrimitiveWithInfer):
     def __init__(self):
         """Initialize PushWeight"""
         self.add_prim_attr("primitive_target", "CPU")
-        self.init_prim_io_names(inputs=['weight', "name", "index"], outputs=['output'])
+        self.init_prim_io_names(inputs=["weight", "name", "index"], outputs=["output"])
 
     def infer_shape(self, weight, name, index):
         return [1]
@@ -779,6 +779,57 @@ class PushWeight(PrimitiveWithInfer):
         return mstype.float32
 
 
+class StartFLJob(PrimitiveWithInfer):
+    """
+    StartFLJob for federated learning worker.
+    """
+    @prim_attr_register
+    def __init__(self, data_size):
+        self.add_prim_attr("primitive_target", "CPU")
+        self.add_prim_attr("data_size", data_size)
+        self.init_prim_io_names(inputs=[], outputs=["result"])
+
+    def infer_shape(self):
+        return [1]
+
+    def infer_dtype(self):
+        return mstype.float32
+
+
+class UpdateModel(PrimitiveWithInfer):
+    """
+    UpdateModel for federated learning worker.
+    """
+    @prim_attr_register
+    def __init__(self):
+        self.add_prim_attr("primitive_target", "CPU")
+        self.add_prim_attr('side_effect_mem', True)
+        self.init_prim_io_names(inputs=["weights"], outputs=["result"])
+
+    def infer_shape(self, weights):
+        return [1]
+
+    def infer_dtype(self, weights):
+        return mstype.float32
+
+
+class GetModel(PrimitiveWithInfer):
+    """
+    GetModel for federated learning worker.
+    """
+    @prim_attr_register
+    def __init__(self):
+        self.add_prim_attr("primitive_target", "CPU")
+        self.add_prim_attr('side_effect_mem', True)
+        self.init_prim_io_names(inputs=["weights"], outputs=["result"])
+
+    def infer_shape(self, weights):
+        return [1]
+
+    def infer_dtype(self, weights):
+        return mstype.float32
+
+
 class identity(Primitive):
     """
     Makes a identify primitive, used for pynative mode.