eric-haibin-lin commented on a change in pull request #8732: rsp push and rsp
pull for comm device, used in kvstore('device')
URL: https://github.com/apache/incubator-mxnet/pull/8732#discussion_r154187393
##########
File path: src/kvstore/comm.h
##########
@@ -621,7 +659,63 @@ class CommDevice : public Comm {
const std::vector<std::pair<NDArray*, NDArray>>& dst,
const bool use_copy,
const int priority) override {
- LOG(FATAL) << "Not implemented yet";
+ CHECK_EQ(src.storage_type(), kRowSparseStorage)
+ << "BroadcastRowSparse expects row-sparse src NDArray";
+
+ bool is_same_rowid = true;
+ for (size_t i = 1; i < dst.size(); ++i) {
+ if (dst[i].second.var() != dst[0].second.var()) {
+ is_same_rowid = false;
+ }
+ }
+
+ for (size_t i = 0; i < dst.size(); ++i) {
+ if (is_same_rowid && i != 0) {
+ CopyFromTo(*dst[0].first, dst[i].first, priority);
+ continue;
+ }
+
+ NDArray* out = dst[i].first;
+ NDArray row_id = dst[i].second;
+ if (use_copy) {
+ CopyFromTo(src, out, priority);
+ } else {
+ CHECK_EQ(out->storage_type(), kRowSparseStorage)
+ << "BroadcastRowSparse expects row_sparse dst NDArray";
+ const bool is_diff_ctx = out->ctx() != src.ctx();
Review comment:
That's true at the beginning. But as soon as you push some gradients on GPU,
it copies the weight from pinned_ctx to GPU. See
https://github.com/apache/incubator-mxnet/blob/master/src/kvstore/kvstore_local.h#L173
----------------------------------------------------------------
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
