Adding RRef as return value for builtin operators #25169

mrshenli · 2019-08-26T02:31:31Z

Stack from ghstack:

Add Python RRef as args and return value #25499 Add Python RRef as args and return value
Adding RRef as return value for builtin operators #25169 Adding RRef as return value for builtin operators

See #23110 for RRef design details. This commit only implements
RRef as return value for builtin operators, and RRef will communicate
between a user and the owner. More specifically, a RRef is first
created on the dist.remote caller, which is a user of the RRef.
Then the RRef user sends and notification to the owner to report
the fork to the owner, and the owner uses a shared_ptr to keep
the RRef alive. When the user RRef is destructed on the caller,
another notification will be sent to the owner, and the owner
can then drop it's RRef as well.

Differential Revision: D17048343

See #23110 for RRef design details. This commit only implements RRef as return value for builtin operators, and RRef will communicate between a user and the owner. More specifically, a RRef is first created on the `dist.remote` caller, which is a user of the RRef. Then the RRef user sends and notification to the owner to report the fork to the owner, and the owner uses a shared_ptr to keep the RRef alive. When the user RRef is destructed on the caller, another notification will be sent to the owner, and the owner can then drop it's RRef as well.

See #23110 for RRef design details. This commit only implements RRef as return value for builtin operators, and RRef will communicate between a user and the owner. More specifically, a RRef is first created on the `dist.remote` caller, which is a user of the RRef. Then the RRef user sends and notification to the owner to report the fork to the owner, and the owner uses a shared_ptr to keep the RRef alive. When the user RRef is destructed on the caller, another notification will be sent to the owner, and the owner can then drop it's RRef as well. ghstack-source-id: 3bbf8e1 Pull Request resolved: #25169

See #23110 for RRef design details. This commit only implements RRef as return value for builtin operators, and RRef will communicate between a user and the owner. More specifically, a RRef is first created on the `dist.remote` caller, which is a user of the RRef. Then the RRef user sends and notification to the owner to report the fork to the owner, and the owner uses a shared_ptr to keep the RRef alive. When the user RRef is destructed on the caller, another notification will be sent to the owner, and the owner can then drop it's RRef as well. Differential Revision: [D17048343](https://our.internmc.facebook.com/intern/diff/D17048343)

See #23110 for RRef design details. This commit only implements RRef as return value for builtin operators, and RRef will communicate between a user and the owner. More specifically, a RRef is first created on the `dist.remote` caller, which is a user of the RRef. Then the RRef user sends and notification to the owner to report the fork to the owner, and the owner uses a shared_ptr to keep the RRef alive. When the user RRef is destructed on the caller, another notification will be sent to the owner, and the owner can then drop it's RRef as well. ghstack-source-id: 3f1eca4 Pull Request resolved: #25169

See #23110 for RRef design details. This commit only implements RRef as return value for builtin operators, and RRef will communicate between a user and the owner. More specifically, a RRef is first created on the `dist.remote` caller, which is a user of the RRef. Then the RRef user sends and notification to the owner to report the fork to the owner, and the owner uses a shared_ptr to keep the RRef alive. When the user RRef is destructed on the caller, another notification will be sent to the owner, and the owner can then drop it's RRef as well. Differential Revision: [D17048343](https://our.internmc.facebook.com/intern/diff/D17048343)

See #23110 for RRef design details. This commit only implements RRef as return value for builtin operators, and RRef will communicate between a user and the owner. More specifically, a RRef is first created on the `dist.remote` caller, which is a user of the RRef. Then the RRef user sends and notification to the owner to report the fork to the owner, and the owner uses a shared_ptr to keep the RRef alive. When the user RRef is destructed on the caller, another notification will be sent to the owner, and the owner can then drop it's RRef as well. ghstack-source-id: 635a250 Pull Request resolved: #25169

See #23110 for RRef design details. This commit only implements RRef as return value for builtin operators, and RRef will communicate between a user and the owner. More specifically, a RRef is first created on the `dist.remote` caller, which is a user of the RRef. Then the RRef user sends and notification to the owner to report the fork to the owner, and the owner uses a shared_ptr to keep the RRef alive. When the user RRef is destructed on the caller, another notification will be sent to the owner, and the owner can then drop it's RRef as well. Differential Revision: [D17048343](https://our.internmc.facebook.com/intern/diff/D17048343)

See #23110 for RRef design details. This commit only implements RRef as return value for builtin operators, and RRef will communicate between a user and the owner. More specifically, a RRef is first created on the `dist.remote` caller, which is a user of the RRef. Then the RRef user sends and notification to the owner to report the fork to the owner, and the owner uses a shared_ptr to keep the RRef alive. When the user RRef is destructed on the caller, another notification will be sent to the owner, and the owner can then drop it's RRef as well. ghstack-source-id: 76b4bde Pull Request resolved: #25169

See #23110 for RRef design details. This commit only implements RRef as return value for builtin operators, and RRef will communicate between a user and the owner. More specifically, a RRef is first created on the `dist.remote` caller, which is a user of the RRef. Then the RRef user sends and notification to the owner to report the fork to the owner, and the owner uses a shared_ptr to keep the RRef alive. When the user RRef is destructed on the caller, another notification will be sent to the owner, and the owner can then drop it's RRef as well. Differential Revision: [D17048343](https://our.internmc.facebook.com/intern/diff/D17048343)

See #23110 for RRef design details. This commit only implements RRef as return value for builtin operators, and RRef will communicate between a user and the owner. More specifically, a RRef is first created on the `dist.remote` caller, which is a user of the RRef. Then the RRef user sends and notification to the owner to report the fork to the owner, and the owner uses a shared_ptr to keep the RRef alive. When the user RRef is destructed on the caller, another notification will be sent to the owner, and the owner can then drop it's RRef as well. ghstack-source-id: ce304cc Pull Request resolved: #25169

See #23110 for RRef design details. This commit only implements RRef as return value for builtin operators, and RRef will communicate between a user and the owner. More specifically, a RRef is first created on the `dist.remote` caller, which is a user of the RRef. Then the RRef user sends and notification to the owner to report the fork to the owner, and the owner uses a shared_ptr to keep the RRef alive. When the user RRef is destructed on the caller, another notification will be sent to the owner, and the owner can then drop it's RRef as well. Differential Revision: [D17048343](https://our.internmc.facebook.com/intern/diff/D17048343)

pietern

This is a first pass, need to do another one.

I noticed this is not clang-format'ed. Do you want to do that as part of this PR as well?

test/test_rpc.py

torch/csrc/distributed/rpc/functions.cpp

torch/csrc/distributed/rpc/script_remote_call.cpp

torch/csrc/distributed/rpc/python_functions.cpp

pietern · 2019-09-04T08:14:38Z

torch/csrc/distributed/rpc/script_ret.cpp

  auto payload_size = message.payload().size();
-  auto value = jit::unpickle(payload, payload_size, nullptr, &message.tensors());
+  auto value = jit::unpickle(
+      payload, payload_size, nullptr, &message.tensors());


This will conflict with #25502.

pietern

Another pass. Looking good!

torch/csrc/distributed/rpc/functions.cpp

torch/csrc/distributed/rpc/python_functions.cpp

pietern · 2019-09-04T16:25:40Z

torch/csrc/distributed/rpc/rref.cpp

+          ScriptRRefFetch(id().toIValue()).toMessage()
+      );
+  auto srv = ScriptRRefValue::fromMessage(fm->wait());
+  return srv.value();


Is it possible to enforce move semantics here? Not sure if this will trigger a copy.

Adding move here gives me "warning: moving a local object in a return statement prevents copy elision [-Wpessimizing-move]".

I modified the return type of srv.value() to const& to explicitly avoid additional copy.

It is possible that returning as value implies move here, given that srv is also local.

A ref will be invalid after return, since srv is local.

A ref will be invalid after return, since srv is local.

Oh, I mean srv.value()'s return type is IValue&, but toHere()'s return type is still IValue. So that there will only be one copy.

Copying IValues is actually pretty cheap since we just copy the underlying intrusive_ptr and not the entire data. The only overhead is incrementing the ref count for intrusive_ptr.

torch/csrc/distributed/rpc/rref.h

pietern · 2019-09-04T16:32:13Z

torch/csrc/distributed/rpc/rref_context.h

+      // NB: cannot use make_shared here as the constructor of OwnerRRef is
+      // private.
+      auto rref = std::shared_ptr<OwnerRRef<T>>(
+          new OwnerRRef<T>(getWorkerId(), rrefId));


Can use std::make_shared here.

OwnerRRef's constructor is private (this is intentional, so that all RRef construction goes through RRefContext). If we use make shared, OwnerRRef needs to add the allocator as a friend or we need to create a wrapper class that enables make_shared, no?

aazzolini

This looks good, but it seems that the messaging mechanism will have to change in order to accommodate the remaining use cases, in particular allowing RRefs as arguments to built-ins. I believe we should commit this one soon and make the required changes on top.

torch/csrc/distributed/rpc/script_rref_proto.h

aazzolini · 2019-09-05T01:06:04Z

torch/csrc/distributed/rpc/types.cpp

+}
+
+GloballyUniqueId GloballyUniqueId::fromIValue(at::IValue&& ivalue) {
+  auto ivalues = ivalue.toTuple()->elements();


what happens if ivalue.toTuple() conversion fails? Does it throw an exception or does it return a nullptr?

It throws an exception if it is not a tuple

torch/csrc/distributed/rpc/script_rref_proto.h

aazzolini · 2019-09-05T01:24:39Z

torch/csrc/distributed/rpc/script_rref_proto.h

+namespace rpc {
+
+// Temporary solution of RRef operations.
+// TODO: Remove all these messages and use rpc + registered functions instead.


What registered functions? You mean JIT built-ins?

I think RRef should be built on top of RPC not in RPC. It does not has to be JIT function, we could have our own namespace for RRef. Similar to using torch.ops.aten.* for builtin ops, we could have torch.distributed.rpc.* for RPC related functions. Then all these messages can be unified as a SCRIPT_CALL.

torch/csrc/distributed/rpc/message.h

aazzolini · 2019-09-05T01:37:32Z

torch/csrc/distributed/rpc/process_group_agent.cpp

-      if (message.isRequest()) {
-        auto response = cb_(std::move(message));
-        send(work.from_, std::move(response));
+      if (message.requiresResponse()) {


I believe we should always send the response, at least as an acknowledgment, and the client should not throw away the acknowledgement, at least for tracing purposes. (e.g. if a FORK_DELETE times out, we may want to re-send the request. etc).

Will address this in #25499

See #23110 for RRef design details. This commit only implements RRef as return value for builtin operators, and RRef will communicate between a user and the owner. More specifically, a RRef is first created on the `dist.remote` caller, which is a user of the RRef. Then the RRef user sends and notification to the owner to report the fork to the owner, and the owner uses a shared_ptr to keep the RRef alive. When the user RRef is destructed on the caller, another notification will be sent to the owner, and the owner can then drop it's RRef as well. Differential Revision: [D17048343](https://our.internmc.facebook.com/intern/diff/D17048343)

mrshenli · 2019-09-05T20:51:22Z

As of the latest commit in #25502, the conflicts no longer exist. Landing now.

torch/csrc/distributed/rpc/script_remote_call.h

pritamdamania87 · 2019-09-05T00:23:49Z

torch/csrc/distributed/rpc/script_remote_call.h

+                   at::IValue retRRefId,
+                   at::IValue retForkId);
+
+  at::IValue retRRefId();


Add docs for retRRefId and retForkId.

pritamdamania87 · 2019-09-05T00:24:07Z

torch/csrc/distributed/rpc/script_remote_call.h

+  static ScriptRemoteCall fromMessage(const Message& message);
+
+ private:
+   const at::IValue retRRefId_;


Add documentation about what these ids are.

torch/csrc/distributed/rpc/script_remote_call.h

torch/csrc/distributed/rpc/script_rref_proto.h

torch/csrc/distributed/rpc/types.cpp

torch/csrc/distributed/rpc/types.h

pritamdamania87 · 2019-09-06T00:27:29Z

test/test_rpc.py

+            expected.append(torch.ones(n, n) * 2)
+
+        for i in range(m):
+            self.assertEqual(rrefs[i].to_here(), expected[i])


We need some more tests here:

Verify all RRefs are cleaned up properly from RRefContext in the end.

Validate state of RRefContext on each node after we've passed a bunch of RRef's around. Also, validate forks are created appropriately.

Add test to verify owner of RRef is populated correctly.

Take down the owner of an rref and verify to_here() throws an appropriate exception.

Take down the owner of an rref and verify to_here() throws an appropriate exception.

@pritamdamania87 How do we take down the owner of an RRef? Are we going to expose test-only APIs for this?

we can create a rref which is not referred in owner?

@gqchen could you please elaborate on "create a rref which is not referred in owner"?

pritamdamania87 · 2019-09-06T00:36:33Z

@aazzolini Thanks for reviewing the PR, although I don't think this PR was in any state to be landed. I got a chance to review it recently and have a lot of concerns as you can see inline. At a high level, the PR needed:

More test coverage
Needed much more documentation.
Some of the APIs needed to be re-worked.
Needed a lot more validation in certain places.
We could have thrown exceptions in destructors which is very dangerous.

All of these issues should have been addressed in this PR and not in a PR on top of this.

facebook-github-bot · 2019-09-06T01:11:59Z

@mrshenli merged this pull request in 197fd4f.

aazzolini · 2019-09-06T18:24:33Z

@pritamdamania87 you're right, my accept was hasty. @mrshenli, definitely address request from other reviewers before moving on to #25499.

torch/csrc/distributed/rpc/functions.cpp

gqchen · 2019-09-09T03:36:43Z

torch/csrc/distributed/rpc/functions.cpp

+      // TODO: make this asynchronous
+      // src is only alive within this block, use reference to avoid copy
+      auto& stack = src.stackRef();
+      src.op()->getOperation()(stack);


We need to handle exceptions here? When we indeed have exception, the ownerRRef should save the exception. This exception will be thrown when client tries to access the ref value.

torch/csrc/distributed/rpc/functions.cpp

gqchen · 2019-09-09T06:26:37Z

torch/csrc/distributed/rpc/rref.h

+//
+// In order to preserve correctness of reference counting, each RRefForkData
+// **MUST** be deserialized into a RRef. This means that if RRefForkData is to
+// be transferred across the network, we need the guarantee that the message


how can we guarantee that? why do we need this guarantee?

gqchen · 2019-09-09T06:39:11Z

torch/csrc/distributed/rpc/rref.cpp

+RRef::RRef(worker_id_t ownerId, const RRefId& rrefId)
+    : ownerId_(ownerId), rrefId_(rrefId) {}
+
+worker_id_t RRef::owner() const {


should we make all those get functions inline?

torch/csrc/distributed/rpc/rref.cpp

torch/csrc/distributed/rpc/rref_context.h

gqchen · 2019-09-09T07:36:24Z

torch/csrc/distributed/rpc/rref_context.h

+    if (ownerId == getWorkerId()) {
+      return getOrCreateOwnerRRef<T>(rrefId);
+    } else {
+      return createUserRRef(ownerId, rrefId, forkId);


should we separate out getOrCreateOwnerRRef() and getOrCreateUserRRef()? Do we need getOrCreateUserRRef?

We don't need getOrCreateUserRRef, because we will always create a new UserRRef on every use.

gqchen

a few more issues. I agree with Pritam that we need to get this PR wrapped up before moving on to add new features.

torch/csrc/distributed/rpc/rref.h

gqchen · 2019-09-11T22:07:42Z

torch/csrc/distributed/rpc/rref.h

+    "RRefForkData must be trivially copyable");
+
+// TODO: make RRef an IValue, and edit createStackForSchema accordingly
+class RRef {


We need a description about the meaning of UserRRef, OwnerRRef, RRefId, ForkId (like what you explained in the chat thread). Those are essential to understand the rest of the code, and they are not obvious.

gqchen · 2019-09-11T22:12:12Z

torch/csrc/distributed/rpc/rref.h

+
+  worker_id_t owner() const;
+  const RRefId& id() const;
+  IValue fork() const;


In what scenario we want to call this API for OwnerRRef? Should we consider moving this method to UserRRef?

The owner may forward its RRef to a user, right? And that's when the fork occurs on owner.

gqchen · 2019-09-11T22:35:58Z

torch/csrc/distributed/rpc/rref_context.h

+  }
+
+  std::shared_ptr<UserRRef> createUserRRef(worker_id_t ownerId) {
+    TORCH_CHECK(ownerId != getWorkerId(), "Cannot create UserRRef on owner.");


we should let user create a UserRRef to a local tensor? This seems to be a very natural behavior.

Here is an example where creating a local UserRRef is very helpful:

def foo(a: RRef[Tensor]): pass x = torch.ones(2, 2) remote(on="worker1", foo, RRef(x))

Yes, will enable this in a separate PR. #25499 has grown too large, if this feature is not absolutely necessary to land #25499, I'd prefer to stack it on top.

gqchen · 2019-09-11T23:16:43Z

torch/csrc/distributed/rpc/rref_context.h

+  const std::shared_ptr<RpcAgent> agent_;
+  std::mutex mutex_;
+  // Keep OwnerRRefs alive until there is no living UserRRefs.
+  std::unordered_map<RRefId, std::shared_ptr<RRef>, RRefId::Hash> owners_;


can we be explicit here? use std::share_ptr<OwnerRRef>, instead of std::share_ptr<RRef>

gqchen · 2019-09-11T23:28:12Z

torch/csrc/distributed/rpc/functions.cpp

+      std::shared_ptr<OwnerRRef<IValue>> rref =
+          RRefContext::getInstance()->getOrCreateOwnerRRef<IValue>(
+              RRefId::fromIValue(srf.value()));
+      auto response = ScriptRRefFetchRet(rref->getValue()).toMessage();


if value is not available, this will block forever. We should add a timeout here.

Yes, will add timeout and error-handling in separate PR. Let me create issues for it.

torch/csrc/distributed/rpc/functions.cpp

mrshenli requested review from apaszke and pietern as code owners August 26, 2019 02:31

This was referenced Aug 26, 2019

Assign each RpcAgent a unique ID, and use ID for sending RPC messages. #24195

Closed

Move worker name collection code from Python to C++ #24260

Closed

mrshenli changed the title ~~Adding RRef as return value for builtin operators~~ [WIP] Adding RRef as return value for builtin operators Aug 26, 2019

mrshenli changed the title ~~[WIP] Adding RRef as return value for builtin operators~~ Adding RRef as return value for builtin operators Aug 28, 2019

mrshenli mentioned this pull request Aug 30, 2019

Add Python RRef as args and return value #25499

Closed

mrshenli added 5 commits August 31, 2019 21:23

pietern reviewed Sep 4, 2019

View reviewed changes

aazzolini approved these changes Sep 5, 2019

View reviewed changes

mrshenli added 5 commits September 5, 2019 08:18

facebook-github-bot closed this in 197fd4f Sep 5, 2019

pritamdamania87 suggested changes Sep 6, 2019

View reviewed changes

facebook-github-bot added the merged label Sep 6, 2019

gqchen reviewed Sep 9, 2019

View reviewed changes

pritamdamania87 mentioned this pull request Sep 10, 2019

Use globally unique id consistently across distributed model parallel code. #25953

Closed

gqchen reviewed Sep 11, 2019

View reviewed changes

mrshenli mentioned this pull request Sep 24, 2019

[RFC] RRef Protocol #26759

Open

facebook-github-bot deleted the gh/mrshenli/12/head branch October 28, 2019 22:17

mruberry added the Merged label Oct 28, 2020

Adding RRef as return value for builtin operators #25169

Adding RRef as return value for builtin operators #25169

Uh oh!

Conversation

mrshenli commented Aug 26, 2019 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

pietern left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

pietern left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

pietern Sep 5, 2019 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

aazzolini left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

mrshenli commented Sep 5, 2019

Uh oh!

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

mrshenli commented Aug 26, 2019 •

edited

Loading

pietern Sep 5, 2019 •

edited

Loading