Distributed Autograd Design
This note will present the detailed design for distributed autograd
and walk through the internals of the same. Make sure you're familiar
with Autograd mechanics and the Distributed RPC Framework before
proceeding.
Background
Let's say you have two nodes and a very simple model partitioned
across two nodes. This can be implemented using
"torch.distributed.rpc" as follows:
import torch
   import torch.distributed.rpc as rpc
def my_add(t1, t2):
     return torch.add(t1, t2)
# On worker 0:
   t1 = torch.rand((3, 3), requires_grad=True)
   t2 = torch.rand((3, 3), requires_grad=True)
# Perform some computation remotely.
   t3 = rpc.rpc_sync("worker1", my_add, args=(t1, t2))
# Perform some computation locally based on remote result.
   t4 = torch.rand((3, 3), requires_grad=True)
   t5 = torch.mul(t3, t4)
# Compute some loss.
   loss = t5.sum()
The main motivation behind distributed autograd is to enable running a