Update on "[reland] Add autograd hook for python rpc call"

zhaojuanmao · zhaojuanmao · commit 60e332b81025 · 2019-10-18T14:25:56.000-07:00
1. currently if autograd context is valid, even tensors do not require grads and grads function are not attached. it still send rpc with autograd meta. This is not ideal. This diff makes some change to make sure rpc with autograd meta is sent only if autograd context is valid and tensors require grads 2. meanwhile create a utiliy to attach autograd info and functions as needed 3. add autograd send/recv functions for python rpc call 4. make changes to support nested python rpc calls 5. disallow nested dist autograd context (was landed in #27022) Differential Revision: [D18017554](https://our.internmc.facebook.com/intern/diff/D18017554/) [ghstack-poisoned]
diff --git a/test/dist_autograd_test.py b/test/dist_autograd_test.py
@@ -15,6 +15,7 @@
 # sent from prev rank respectively.
 # rpc_done[2] and ctx_ids[2] represents for prev of prev rank.
 # rpc_done[3] and ctx_ids[3] represents for prev of prev of prev rank.
+# rpc_done[0] and ctx_ids[0] represents for current rank, but mostly not used.
 rpc_done = [False, False, False, False]
 ctx_ids = [-1, -1, -1, -1]
 
@@ -163,8 +164,8 @@ def _verify_graph_for_first_rpc_call(self, send_function, recv_function, t1, t2,
         self.assertEqual(ret.grad_fn, recv_function)
 
     # For a context passed from previous nested chain calls, this rank
-    # recevied two tensors t1 and t2, execute torch.add(t1, t2) and send result
-    # tensor t3 back.
+    # receives two tensors t1 and t2, executes torch.add(t1, t2) and sends
+    # result tensor t3 back.
     # For this context in this rank, it expects graph like this:
     #  send and recv functions:
     #       rpcSendBackward
@@ -191,15 +192,15 @@ def _verify_graph_for_rpc_call_exec(self, send_function):
         self.assertEqual(next_funcs[0][0], next_funcs[1][0])
 
     # For a context passed from previous nested chain calls, this rank
-    # recevied two tensors t1 and t2, forwarding t1 and t2 tensors using
+    # receives two tensors t1 and t2, forwards t1 and t2 tensors using
     # nested rpc call to next dst. In return route, receive result tensor t3
     # from next dst and forwarding t3 back to previous calls.
     # For this context in this rank, it expects graph like this:
-    #  send and recv functions while recevive and forward t1 and t2:
+    #  send and recv functions for receving and forwarding t1 and t2:
     #       rpcSendBackward
     #          /          \
     # t1.recvRpcBackward    t2.recvRpcBackward
-    #  send and recv functions while receive and forward t3:
+    #  send and recv functions for receiving and forwarding t3:
     #       rpcSendBackward
     #             |
     #           t3.recvRpcBackward
diff --git a/torch/csrc/distributed/rpc/request_callback.cpp b/torch/csrc/distributed/rpc/request_callback.cpp
@@ -38,7 +38,7 @@ struct ClearAutogradContextGuard {
 } // anonymous namespace
 
 Message RequestCallback::operator()(Message& request) const {
-  // For a rev thread, current context id should be invalid outside
+  // For a recv thread, current context id should be invalid outside
   // processMessage().
   ClearAutogradContextGuard guard;
   try {
