solin319 commented on a change in pull request #8373: distribute training in 
fp16
URL: https://github.com/apache/incubator-mxnet/pull/8373#discussion_r168054096
 
 

 ##########
 File path: src/kvstore/kvstore_dist_server.h
 ##########
 @@ -378,25 +379,46 @@ class KVStoreDistServer {
     if (req_meta.push) {
       size_t ds[] = {(size_t)req_data.lens[0]};
       TShape dshape(ds, ds + 1);
-      TBlob recv_blob((real_t*)req_data.vals.data(), // NOLINT(*)
+      TBlob recv_blob((DType*)req_data.vals.data(), // NOLINT(*)
                       dshape, cpu::kDevMask);
       NDArray recved = NDArray(recv_blob, 0);
+      NDArray recved_tmp;
+      if (recved.dtype() != mshadow::DataType<real_t>::kFlag) {
+        recved.WaitToRead();
+        recved_tmp = NDArray(dshape, Context::CPU(0), false, 
mshadow::DataType<real_t>::kFlag);
+        CopyFromTo(recved, &recved_tmp, 0);
+        recved_tmp.WaitToRead();
+      }
       if (stored.is_none()) {
         // initialization
-        stored = NDArray(dshape, Context());
-        CopyFromTo(recved, &stored, 0);
+        if (recved.dtype() != mshadow::DataType<real_t>::kFlag) {
+          stored = NDArray(dshape, Context::CPU(0), false, 
mshadow::DataType<real_t>::kFlag);
 
 Review comment:
   Because the compute in fp16 is slow in cpu.
   When the size of  received data is large, it's better to move it to gpu with 
fp16.
   When the size of  received data is small, it's better to move it to cpu with 
fp32.
   It's better to implement like below.
   ```
   if (total_size > 100000) {
                  stored = NDArray(dshape, Context::GPU(0), false, 
mshadow::DataType<DType>::kFlag);
   } else{
                  stored = NDArray(dshape, Context::CPU(0), false, 
mshadow::DataType<real_t>::kFlag);
   }
   ```

----------------------------------------------------------------
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:
us...@infra.apache.org


With regards,
Apache Git Services

Reply via email to