[GitHub] eric-haibin-lin commented on a change in pull request #10183: [MXNET-120] Float16 support for distributed training

[GitBox <gi...@apache.org>]

eric-haibin-lin commented on a change in pull request #10183: [MXNET-120] Float16 support for distributed training
URL: https://github.com/apache/incubator-mxnet/pull/10183#discussion_r178665723
 
 

 ##########
 File path: src/kvstore/kvstore_dist_server.h
 ##########
 @@ -170,43 +216,90 @@ class KVStoreDistServer {
     app->Response(recved);
   }
 
+  /*
+   * For keys already initialized, if necessary create stored_realt.
+   * This will only be used if by some wrong usage of kvstore,
+   * some keys are initialized before optimizer is set.
+   */
+  void CreateMultiPrecisionCopies() {
+    for (auto const& stored_entry : store_) {
+      const int key = stored_entry.first;
+      const NDArray& stored = stored_entry.second;
+      if (stored.dtype() != mshadow::kFloat32) {
+        auto& stored_realt = store_realt_[key];
+        if (stored.storage_type() == kRowSparseStorage) {
+          stored_realt = NDArray(kRowSparseStorage, stored.shape(), stored.ctx(),
+                                          true, mshadow::kFloat32);
+        } else {
+          stored_realt = NDArray(stored.shape(), stored.ctx(), false, mshadow::kFloat32);
+        }
+
+        auto& update = update_buf_[key];
+        if (!update.merged.is_none()) {
+          update.merged = NDArray(update.merged.shape(), update.merged.ctx(), false,
+                                  mshadow::kFloat32);
+        }
+        CHECK(update.request.size() == 0)
+          << ps::MyRank() << "Multiprecision mode can not be set while pushes are underway."
+          << "Please set optimizer before pushing keys." << key << " " << update.request.size();
+
+        CopyFromTo(stored, stored_realt);
+      }
+    }
+    for (auto const& stored_realt_entry : store_realt_) {
+      stored_realt_entry.second.WaitToRead();
+    }
+  }
+
   void DataHandleEx(const ps::KVMeta& req_meta,
-                    const ps::KVPairs<real_t>& req_data,
-                    ps::KVServer<real_t>* server) {
-    DataHandleType recved_type = static_cast<DataHandleType>(req_meta.cmd);
-    if (recved_type == DataHandleType::kRowSparsePushPull) {
-      DataHandleRowSparse(req_meta, req_data, server);
-    } else if (recved_type == DataHandleType::kCompressedPushPull) {
-      DataHandleCompressed(req_meta, req_data, server);
-    } else {
-      DataHandleDefault(req_meta, req_data, server);
+                    const ps::KVPairs<char>& req_data,
+                    ps::KVServer<char>* server) {
+    DataHandleType type = DepairDataHandleType(req_meta.cmd);
+    switch (type.requestType) {
+      case RequestType::kRowSparsePushPull:
+        DataHandleRowSparse(type, req_meta, req_data, server);
+        break;
+      case RequestType::kCompressedPushPull:
+        DataHandleCompressed(type, req_meta, req_data, server);
+        break;
+      case RequestType::kDefaultPushPull:
+        DataHandleDefault(type, req_meta, req_data, server);
+        break;
     }
-    return;
   }
 
-  inline void ApplyUpdates(const int key, MergeBuf *merged, NDArray *stored,
-                           ps::KVServer<real_t>* server) {
-    if (merged->request.size() == (size_t) ps::NumWorkers()) {
+  inline bool has_multi_precision_copy(const DataHandleType type) {
+    return multi_precision_ && type.dtype != mshadow::kFloat32;
+  }
+
+  inline void ApplyUpdates(const DataHandleType type, const int key,
+                           UpdateBuf *updateBuf, ps::KVServer<char>* server) {
+    if (!sync_mode_ || updateBuf->request.size() == (size_t) ps::NumWorkers()) {
 
 Review comment:
   nit: `updateBuf ` -> `update_buf `?

----------------------------------------------------------------
This is an automated message from the Apache Git Service.
To respond to the message, please log on GitHub and use the
URL above to go to the specific comment.
 
For queries about this service, please contact Infrastructure at:
users@infra.apache.org


With regards,
Apache Git Services