Repository: pytorch
Issue: [RFC] RPC Based Distributed Model Parallel
Labels: feature,triaged,module: rpc

with @pritamdamania87 @zhaojuanmao @aazzolini @gqchen @pietern @satgera @ezyang @zdevito @suo @manojkris @gchanan @soumith @dzhulgakov @yifuwang @bddppq @joxu-cn @dwarakrajagopal @jspisak 

PyTorch currently provides simple APIs for single machine data parallel, distributed data parallel, and single machine model parallel. However, when it comes to distributed model parallel, applications have to build their own scaffold to stitch together local autograd graphs into one global graph. This proposal aims to fill in that gap by providing an RPC-Based distributed model parallel API. In short, applications may run RPC to execute code remotely in the forward pass, and autograd will automatically travel across RPC boundaries in the backward pass. 

# API

## Core Concepts


**RRef[T] -** (abbreviation ref)  A *reference *to a value of some type `T` (e.g. Tensor) on a remote worker. This handle keeps the referenced remote tensor value alive on the owner, but there is no implication that the value will be transferred to the local worker in the future. It is valid to have a reference to local value as well, and values of type `T` can be* implicitly converted* to `RRef[T]`. This implicit conversion will be critical later to allow the expression of different types of RPC. Think of it like the implicit conversion from `std::string` to `const std::string &`. See System Design section for more details about `RRef`.

```python
ref.owner() # what is the worker this value lives on

v = ref.local_value() # if ref.owner() is local worker, then
                      # this returns the the underlying value, otherwise error.
# you can create a ref to a local tensor
t = torch.rand(3, 4)         
ref2 = torch.RRef(t)

# in TorchScript, T can be automatically converted to RRef[T]
ref3 : RRef[Tensor] = t
```


**Future[T] -** (abbreviation fut) a guarantee that at some future point in time the value of type `T` will be available locally. The action to create `T` locally is assumed to be scheduled and in-progress.  Future is already supported in TorchScript and we are extending this to remote calls. 

```python
v = fut.wait() # block the current thread until v is ready

# local cpu task creation returns a future to the computed tensors
fut = torch.fork(lambda x, y: x + y, torch.rand(3, 4), torch.rand(3, 4))
```

## Core Functions

```python
# synchronous
result : T = torch.rpc(on : Worker, remote_callable : Callable, *args)
# asynchronous
result : Future[T] = torch.async_rpc(on : Worker, remote_callable : Callable, *args)
# remote reference
result : RRef[T] = torch.remote(on : Worker, remote_callable : Callable, *args)
```

Each function above invokes `remote_callable` on a remote worker. Value types in the `args` list are *copied by value* to the remote worker. `RRef[T]` types in the `args` list are *copied by reference* to the remote worker (again see the analogy between `std::string` and `const std::string&`). 

The *synchronous* variant copies the result value back, blocking the calling thread until the response occurs. The *asynchronous* variant returns immediately with a future. The remote knows that the call *will expect to receive the value* so it will send a message back at some point with the result without further prompting.

The *remote reference* variant returns immediately with an `RRef` of the return value. The remote knows that the caller *does not expect to receive the result value*. 


Below shows how these functions are used:


```python
# make some local tensors
a : Tensor = torch.rand(3, 4)
b : Tensor = torch.rand(3, 4)

# define a remote function, visible to all machines.
# type annotations define expected input/output types.
def remote_function(a : Tensor, b : RRef[Tensor]) -> Tensor:
   # 'b' in the type signature is a remote reference, so we must copy it here
   # to use it locally.
   # to_here() is defined later in the syntax sugar section, it synchronously
   # copies the tensor to this worker.
   b_l  : Tensor = b.to_here()
   return a + b_l

# run remote_function on a different device. 
# a is copied by value since it is a Tensor
# b is copied by reference remote machine due to the RRef[Tensor] 
# type annotation in the signature, which causes an implicit conversion to a
# reference type.

# torch.remote always creates an RRef of the result type. 
# It does not wait for the remote's response. 
# There is no implied copy of the tensor data yet.
c : RRef[Tensor] = torch.remote("worker1", remote_function, a, b)

# we can explicitly request the data to be copied back here:
c_l : Tensor = c.to_here()

# another example:
def remote_function2(a : Tensor, b : Tensor) -> Tensor:
   return a + b

# Here we call torch.rpc which returns the value directly without
# creating a remote reference.
# we synchronously wait for remote_function2 to return.
c : Tensor = torch.rpc("worker2", remote_function2, a, b)

# When the RPC call is returning a non-reference type, we need to wait for 
# a response from the remote host. To avoid synchronously waiting, use the
# async flag to get a future instead.
c_f : Future[Tensor] = torch.async_rpc("worker2", remote_function2, a, b)
# even before calling wait, the remote knows that the data should be sent back
# to the caller as soon as it is ready.

# force the local thread to wait for the remote's response
c = c_f.wait()


# if you omit type annotations in the remote function, the assumption is that 
# arguments are passed without any implicit conversions
def remote_function3(a, b):
   # no annotations mean that a, b will be Tensor since there is no conversion
   return a + b

c: Tensor = torch.rpc("worker2", remote_function3, a, b)

```

### RRef Forks

### Implicit Conversions for RRef Arguments


We allow implicit conversion between `T` and `RRef[T]` for arguments of RPC functions.  Both the actual and formal parameter can either be a `T` or an `RRef[T]`, leading to four cases that might occur:

**T → T (passing a T to an rpc that accepts a T):** the value T is copied by value, and send over the wire as part of the message invoking the RPC

**T → RRef[T] (passing a T to an rpc that accepts  RRef[T]):** The caller constructs a remote reference to the argument, and sends the *reference* over the wire to the callee. The data is not sent. The callee can then use the reference as a handle to either request the data later or to make further remote calls.

**RRef[T] → T (passing an RRef[T] to an rpc that accepts T):** The callee expects to get an actual value, so the callee needs to turn the reference into a value. The network behavior depends on where the `RRef[T]` lives.

* If the `RRef[T]` lives on the caller, then the implementation looks up the actual value of `T` locally and pass it by value along the wire similar to the T → T case.
* If the `RRef[T]` lives on the callee, then the implementation just sends the reference and the callee does the lookup locally.
* If the `RRef[T]` lives on some third machine, then the caller sends 2 messages. One to the third machine telling it to send the data in the remote reference directly to the callee, and one to the callee telling it to start the RPC and expect this input to be coming from the third machine. This effectively forward value of the `RRef[T]` to the callee without the caller having to load it or the callee having to request it later


Examples:

```python
def remote_function1() -> Tensor:
    return torch.ones(2)
    
def remote_function2(a : Tensor) → Tensor:
    b = a * 2
    return b

aref : RRef[Tensor] = remote("worker1", remote_function1)

# this local worker will make two RPC calls: one to tell worker1 to send the 
# tensor to worker2, and another one to tell worker2 to expect this Tensor input
# from worker1. remote_function2 will run on worker2 only after it received the 
# tensor from worker1.  
bref : RRef[Tensor] = remote("worker2", remote_function2, aref) 
```



**RRef[T] → RRef[T]** (**passing an RRef[T] to an RPC that accepts RRef[T]): **The callee expects an `RRef[T]`, but we must make sure we correctly keep track of references to the value on a remote. So the actual behavior depends on where the `RRef[T]` lives.

* If `RRef[T]` lives on the caller, then we simply pass it to the remote and record that this remote now has a live reference to the value.
* If the `RRef[T]` lives on the callee, then we pass it to the remote, and it becomes a local reference on the remote.
* If `RRef[T]` lives on some third machine, then we must forward the reference. To do this the caller sends two messages. One to the third machine telling it to create a remote reference and send it to the callee, and one to the callee telling from where to expect the remote. The callee code is not invoked until the remote is transferred to ensure sane reference counting.

Examples: 

```python
def remote_function1() -> Tensor:
    return torch.ones(2)
    
def remote_function2(a : RRef[Tensor]) -> Tensor:
    int delta = 10
    return a.to_here() + delta

aref : RRef[Tensor] = remote("worker1", remote_function1)

# this local worker will make two RPC calls: one to tell worker1 to create a 
# remote reference and send it to worker2, and another one to tell worker2 to 
# expect this remote reference input from worker1. remote_function2 code will 
# not run on worker2 until it receives the remote reference from worker1 to 
# ensure proper reference counting.   
bref : RRef[Tensor] = remote("worker2", remote_function2, aref) 
```


When an `RRef[T]` goes dead on machine A, a message is sent to the owner of `T` telling it that the reference from machine A is dead.

### Explicit RRef type for return values

The above implicit `RRef` argument conversion does not apply to return values. If `remote_function` returns `RRef[T]`, calling it remotely using `torch.remote` would return `RRef[RRef[T]]` instead of `RRef[T]`. This is because when the return value `RRef` of `torch.remote` is first created on the caller who does not know the owner of the real data `T`. T could be stored on the callee of `torch.remote`, but it could also be on a different worker as callee may also make another remote call within `remote_function` and return an `RRef[T]` owned by a different worker. Moreover, the caller is allowed to share the returned `RRef` with other workers immediately after `torch.remote` returns.  However, as by then, the caller does not know the real owner of `T` yet, sharing the `RRef` would break the reference count algorithm.

Examples:


```python
def remote_function3() -> RRef[Tensor]:
    return torch.remote("Worker2", torch.ones, 2, 2)
 
cref : RRef[RRef[Tensor]] = remote("worker1", remote_function3) 
```



## Initialization API

Users may choose communication backend for RPC, and users are responsible for setting up the backend properly before calling the `init_rpc` method. 


```python
# backend: specifies the underlying communication implementation
# init_method: contains the information to initialize/connect a name store to 
#              resolve names
# name: is a unique identifier for the current worker
torch.distributed.init_rpc(backend="pg", init_method="file:///...", name="worker1")
```

The `init_rpc` method will create an `RpcAgent` under the hood and will make the current worker ready to send and receive RPC calls. If you call `init_rpc` and use the `ProcessGroup` (`pg`) backend, it acts as a global barrier, where all the node names as collectively synchronized before continuing. This is not the case if you use a peer to peer backend (e.g. tensor pipes), where calling `init_rpc` will register the node name in the specified store and start serving.  

Applications don’t need to explicitly register functions for remote execution, but we do assume same functions are defined on both caller and callee. This is often true as all workers can import the same set of libraries or even share the same Python script. 


## Syntax Sugar

Other operations are now implementable using syntax sugar.

### Retrieving Value From RRef

```python
# helper private RPC functions
def _identity(v : Tensor) -> Tensor:
    # copy the tensor by value to this remote,
    return v

def _to_here(v : RRef[T]) -> T:
    # take a reference, send it to the device that owns it
    # and have that device return the actual tensor by value
    return v.local_value()

class RRef[T]:
    ...
    # copy a remote tensor to the local worker, sync version
    def to_here(self) -> T:
        return torch.rpc(_to_here, self, on=self.owner())

```

### Builtin Operators

```python
# proxy methods for all builtin functions exist on references for 
# existing TorchScript types like Tensors. They always follow a fixed pattern:
def _mm(a : RRef[Tensor], b : RRef[Tensor]) -> RRef[Tensor]:
    return a.local_value() + b.local_value()
     
class RRef[Tensor]:
    def mm(self : RRef[Tensor], other : RRef[Tensor]) -> RRef[Tensor]:
        on = same_worker(self.owner(), other.owner())
        return torch.remote(on, _mm, self, other)


c : Tensor = a.mm(b).to_here()

```

### Callable and RRef

If `RRef[T]` holds a callable object `T`, the application may directly call  the `RRef` which will be translated into `torch.remote` call to the owner of the callable.

```python
# if T is callable for RRef[T], rref(x) will be translated to calling T(x) 
# on the owner of the RRef
def _call_rref(v : RRef[T], *args):
    return v.local_value()(*args)

class RRef[T]:
    def __call__(self, *args):
        return torch.remote(self.on(), _call_rref, self, *args)

net = torch.remote("Worker1", Net)
net(inputs)
```

### Optimizer and RRef

As models might have remote sub-modules (i.e., `RRef[nn.Module]`), we should provide an optimizer sugar to handle it. The optimizer sugar (`torch.optim.remote`) takes a local optimizer constructor, a distributed model parallel model, and an argument list for the local optimizer constructor. The `torch.optim.remote` recursively creates a local optimizer on every remote sub-module owner, and exposes the same step API as a local optimizer which recursively calls every local optimizer. 


```python
class Net1(nn.Module):
    ...

class Net2(nn.Module):
    ...

class DMP(nn.Module):
    def __init__(self):
        self.net1 = dist.remote("worker1", Net1)
        self.net2 = dist.remote("worker2", Net2)
        
dmp = dist.remote("worker0", DMP)
# dist.optimizer creates an optimizer on all RRef owners
optimizer = dist.optimizer(torch.optim.SGD, dmp, lr=0.1)

with dist.autograd.context():
  loss = dmp(inputs)
  dist.autograd.backward(loss)
  optimizer.step()
```

## Model Parallel Training Examples

### Multi-Machine Model Training

```python
# 1. load data
inputs_rref = torch.remote("worker1"， load_inputs, path_to_inputs) 
labels_rref = torch.remote("worker2"， load_labels, path_to_inputs)

# 2. define model
class Net1(nn.Module):
    ...

class Net2(nn.Module):
    ...

class DMP(nn.Module):
    def __init__(self):
        self.net1 = torch.remote("worker1", Net1)
        self.net2 = torch.remote("worker2", Net2)
        
    def forward(self, inputs_rref):
        # RRef[T].__call__(args) is a sugar that translates to 
        # dist.remote(T, RRef.on(), args)
        outputs1_rref = self.net1(inputs_rref)
        outputs2_rref = self.net2(outputs1_rref)
        return outputs2_rref
        
# 3. training, run it where you want to call autograd
def train(inputs_rref, labels_rref):
    dmp = DMP()
    # torch.optim.remote creates an optimizer on every RRef destination
    optimizer = dist.optimizer(torch.optim.SGD, dmp, lr=0.1)
    outputs_rref = dmp(inputs_rref)
    loss = loss_func(outputs_rref.to_here(), labels_rref.to_here())
    autograd_ctx_id = dist.autograd.backward(loss)
    optimizer.step(autograd_ctx_id)
    
dist.rpc(dev2, train, args=(inputs_rref, labels_rref))

```

### Parameter Server Training


```python
class ParameterServer:
    def __init__(self):
        self.params = torch.zeros(100, 100).to(0)
        
    def get_params(self) -> Tensor:
        return self.params
        
    def add_grads(self, grad: Tensor):
        return self.params += grad.to(0)
        
def train(ps)
    for _ in range(10):
        params = torch.rpc("ps", ParameterServer.get_params, args=(ps, ))
        # run forward and backward
        torch.rpc("ps", ParameterServer.add_grads, args=(ps, params.grad))
        torch.distributed.barrier(group=TRAINER_GROUP)
        
ps = torch.remote("worker1"，ParameterServer)
torch.remote("worker2", train, args=(ps,))
torch.remote("worker3", train, args=(ps,))
            
```

# System Design

## Distributed Autograd

### Basic Idea

In the first version, `dist.autograd.backward` does not support `RRef` arguments, but `RRef` can still help build the autograd graph. The overall idea is as follows.

* When calling  `torch.rpc` or `RRef.to_here()`, `send` and `recv` autograd functions will be inserted to connect local autograd graphs on multiple workers into one distributed autograd graph. 
* Every distributed backward pass is assigned a globally unique id (***`autograd_context_id`***), and every participating worker will keep a dedicate context for it. 
* When the backward computation reaches a `recv` function, it packs the gradient and the `autograd_context_id` in the message, and pass it to its `send` counterpart. 
* Upon receiving a message for a `send` function in the backward pass, it uses the `autograd_context_id` in the message to identify which backward pass it belongs to, and uses the gradient in the message to continue autograd computation locally.



### Send and Recv Autograd Functions 

Let’s start with a simple example where there is just one synchronized RPC call and there is only one tensor passed across worker boundaries. Code is on the left and the autograd graph is on the right where `AccumulateGrad` autograd functions for leaf nodes are omitted for simplicity.

```python
# the add function should be 
# defined on both workers
def add() -> Tensor:
    a = torch.rand(2, 2)
    b = torch.rand(2, 2)
    c = a + b
    return c
    
# make RPC call from worker0
# to execute add on worker1
c1 = dist.rpc(add, on="worker1")
d = torch.ones_like(c1)
e = c1 * d
e.sum().backward()
```

![Screen Shot 2019-07-16 at 11 19 31 AM](https://user-images.githubusercontent.com/16999635/61566619-d4712f00-aa4a-11e9-922d-8124f17fbef7.png)


The `send` and `recv` autograd functions are inserted during the forward pass, which connect two local graphs into one distributed graph. In the backward pass, the gradient will be passed to the `recv` autograd function on `worker0`, and the `recv` autograd function will then transmit the gradient tensor to `worker1`’s `send` autograd function. Then, `worker1` can kick off the local autograd engine to resume the backward pass. There are a few more details need to be clarified in this simple example:

* On `worker1`, how do we keep the autograd graph alive after the RPC call returns?
    * In short, the distributed autograd engine on `worker1` will keep a reference to the `send` function which can keep the graph alive.
    * Reasoning: The graph can be kept alive by keeping a reference to either tensor `C` or the `send` autograd function, as both of them hold a reference to the `add` autograd function. We choose to keep a reference to the `send` function instead of tensor `C`, because `C` as a non-leaf node produced by `add` is not needed in the backward pass. It should be freed as soon as possible. It is not memory efficient to hold C alive just because we want to have an entrance point to the autograd graph. 
* In the backward pass, how does `recv` on `worker0` find the correct `send` on `worker1` to talk to?
    * This can be done by assigning a globally unique ID (**worker*****_id + local send/recv id***) for each `send` / `recv` function pair. 
* When can `worker1` delete its local autograd graph?
    * `send` should have the same lifetime as its corresponding `recv` function. This can be done by sending a message from `worker0` to `worker1` when `recv` is destructed on `worker0`. The `recv` function is kept alive by the `loss` tensor. So, conceptually, the global autograd graph will be deleted when the final loss tensor is gone. 

### Hidden Autograd Path and Circular Dependency 

Things can become complicated when an autograd graph contains multiple send/recv pairs. Consider the following example.

```python
# all functions shoud be defined on all workers
def worker0_func(c2: Tensor) -> Tensor:
    g = torch.rand(2, 2)
    h = g + c2
    return h

def worker1_func_top() -> Tensor:
    a = torch.rand(2, 2)
    b = torch.rand(2, 2)
    c = a + b
    return c

def worker1_func_bottom(c: Tensor, e1: Tensor) -> Tensor:
    f = c + e1
    return f

def worker2_func(c1: Tensor) -> Tensor:
    d = torch.rand(2, 2)
    e = c1 + d
    return e

# on Worker3
c_ref = torch.remote(worker1_func_top, on="Worker1")
h1 = torch.rpc(worker0_func, c_ref, on="Worker0")
e_ref = torch.remote(worker2_func, c_ref, on="Worker2")
f1 = torch.rpc(worker1_funct_bottom, c_ref, e_ref, on="Worker1")
i = h1 + f1
i.sum().backward()
```

![Screen Shot 2019-07-16 at 6 20 52 PM](https://user-images.githubusercontent.com/16999635/61566634-e94dc280-aa4a-11e9-9514-d8a51e54e4a2.png)



This example highlights two problems that we need to address:

* **Hidden Autograd Path:** Existing local autograd engine starts from loss (or all outputs), and do a discovery/marking phase to identify all participating functions before executing the real autograd computation. So that all paths in the autograd graph are known upfront. However, we don’t have this luxury in distributed autograd because some parts of the autograd graph reside on remote workers. For example, when grad arrives at `send5`, worker1 cannot tell whether `send3` will be in the backward pass if it only looks at local information. More specifically, `i.sum().backward()` will be the same as `f1.sum().backward()` from worker1’s perspective, but the former involves `send3` and the latter does not. 
    * To address this problem, we propose to record all globally upstream (upstream in the forward pass, downstream in the autograd graph) `send` / `recv` pairs in the forward pass, so that we know exactly which `send` / `recv` to wait for in the backward pass. 
* **Circular Dependency:** there are circular dependencies between worker1 and worker2, i.e., it is impossible to finish autograd computation on one worker before kicking off on another worker. One option is to start autograd computation on `worker1` first, and having an autograd thread blocking there waiting for grads for `send1`, but this is less ideal.
    * To address this problem, we propose to only create the `send` autograd function and put it in the ready queue when the grad is received. Note that, when computing dependency count for `add1`, the autograd engine still takes `send1` into account, so that the engine will only start computing grads for add1 after both `add2` and `send1` finish. 


Note that we need to record information in the forward pass and do the discovery in the backward pass because we don’t know which `send` function will be participating in the autograd computation. However, if the application can guarantee that all `send` functions will receive grad in the backward pass, we can skip all these complexity and have a more efficient version. Both scenarios are useful, so we propose to have two modes:

* **Smart Mode** supports running backward on a subgraph of the global autograd graph, but there will be extra overhead in both forward and backward pass. 
* **Fast Mode** skips dependency recording in the forward pass and graph discovery in the backward pass, but the application needs to guarantee that *_all send autograd function will receive grad in the backward pass_*. 

The two sections below describe the two algorithms in more details. 

### Distributed Autograd Algorithm Smart mode


**Forward pass:** 

For every `send`  **x**:

1. Find `send` functions in **x**’s lineage, by:
    1. Finds all locally reachable `recv` functions from `send` **x** in the autograd graph. In the example above, `send2` finds `recv1`, `send4` finds `recv3`, and `send5` finds `recv2`. 
    2. Use those found `recv` functions to find globally reachable `recv` functions in `send` **x**’s lineage. Note that this can be done, because in step 2 we send enough information from `send` to `recv`. In the example above `send4` knows `send3`, and `send5` knows `send1` and `send2`. 
2. Then, `send` **x** includes ids of its lineage `send` functions in the message. Intuitively, it means that if there is a grad received for `send` **x**, the backward pass must reach all `send` functions in its lineage as well. It helps a node to determine whether it should wait for a `send` grad.

```python
# pseudo code to demonstrate how send works in forward
def find_global_lineage(tensor):
    # find local lineage
    recvs = find_recvs(tensor.grad_fn)
    dep_ids = {recv.id for recv in recvs}
    # find global lineage
    dep_ids.update({dep_id for recv in recvs for dep_id in recv.dep_ids})
    return dep_ids

def send(func, tensors, on):
    msg = Message(func)
    for tensor in tensors:
        lineage = find_global_lineage(tensor)
        # connect send to autograd graph
        send = SendFunc()
        send.next = tensor.grad_fn
        # remember the send by its id
        RpcAgent.send_map[send.id] = send
        # coalesce data
        msg.data.append((tensor, send.id, lineage))
    send_msg(msg, on)
    
def recv(func, data, from):
    tensors = []
    for tensor, send_id, lineage in data:
        # use send_id as recv_id, and remember global lineage
        recv = RecvFunc(send_id, lineage)
        tensor.grad_fn = recv
        tensors.append(tensor)
        
    return func(tensors)
```


**Backward pass**: 

On the node that calls `torch.distributed.backward`:

1. Find all `send` functions in the lineage of the loss tensor. In the above example, it will be all 5 `send` functions. These ids will be propagated to the `recv` functions and will be passed to the counterpart `send` functions accordingly. 
    1. Optimizations can be added, e.g., drop unnecessary ids in backward pass to reduce message size. 

On every node:

1. Upon receiving the first message (be it a dedicated discovery message or grad of a send), record its `autograd_context_id`, and retrieve all participating `send` ids from the message. Compute dependency count from those `send` functions (and also from loss `grad_fn` if loss is on this node). Set dependency count for `send` functions as 1. If there is any autograd function has dependency count 0, put them into the ready queue. 
2. Upon receiving a `send` grad, decrement the dependency count of that `send` by 1, and add it to the ready queue. Note this is done on an `RpcAgent` thread, and some autograd engine thread will pick up the autograd function for execution. 



```python
# pseudo code to demonstrate backward
graph_tasks = {}
def backward(loss):
    global graph_tasks
    
    autograd_context_id = gen_autograd_id()
    lineage = find_global_lineage(loss)
    # these send will participate in the autograd pass
    roots = local_sends.intersection(lineage)
        
    # propagate the autograd_id and deps info to all
    # participating workers. This is non-blocking and can
    # run concurrently with the real backward computation. 
    # This step is not absolutely necessary, but can help other
    # workers to kick off autograd earlier.
    disseminate(autograd_context_id, lineage)

    # below is a handwaving impl to show how it works with local autograd engine
    graph_task = GraphTask()
    graph_tasks[autograd_context_id] = graph_task
    roots.append(loss.grad_fn)
    # setup dependency count properly
    compute_dependencies(GraphRoot(roots), graph_task)
    # insert the task to local engine ready queue. Only the FunctionTask
    # for loss is inserted now, send FunctionTasks will be inserted later
    # when their grad becomes available.
    ready_queue.push_back(FunctionTask(graph_task, loss.grad_fn, ...))
    return autograd_context_id
    
    
def on_grad_send(send_id, grad, autograd_id):
    global graph_tasks
    graph_task = graph_tasks[autograd_id]
    send_func = RpcAgent.send_map[send_id]
    ready_queue.push_back(FunctionTask(graph_task, send_func, grad))
```



### Distributed Autograd Algorithm Fast mode

The problem with the above approach is that including ids in  `send` / `recv` messages incurs overhead, especially when there are a lot of tensors communicated across multiple workers. And this discovery phase is only necessary when running autograd on subgraph. For example,  `f1.sum().loss()` requires the discovery phase to avoid waiting for `send3`, but it is easier for `i.sum().loss()` as all `send` are involved in the backward. So, we propose to have one additional mode for distributed autograd to bypass `send` / `recv` dependency discovery in both forward and backward **if all send for non-leaf or `requires_grad` tensors will receive grad in the backward pass**. The mode can be toggled when initializing RPC agents:


```python
# all_requires_grad (bool): If True, the application guarantees that all 
# send functions on non-leaf or requires_grad tensors will receive grad 
# in the backward pass. Hence, we can skip the distributed dependency 
# discovery algorithm (fast mode). If False, run smart mode, where 
# messages beween send/recv will contain dependency ids in both forward
# and backward pass. (default False)
torch.distributed.init_rpc(name, backend="pg", all_requires_grad=False)
```


Internally, `RpcAgent` will create a thread-local driver ID, where a driver is the worker that pieces together the autograd graph. In the above example, `Worker3` is the driver. In the forward pass, every `send` function originated from this driver will be tagged with its thread-local driver ID, and this applies to all downstream (upstream in the autograd graph) `send` functions as well. This can be done by either propagating this driver ID to RPC calls recursively, or do an active driver ID discovery by walking the autograd graph before sending a tensor. If this information is ambiguous, e.g., one `send` function traces back to two upstream (downstream in the autograd graph) `recv` functions from two different drivers, it will throw an error. In the backward pass, the thread-local driver id of the loss will be included in the entire autograd execution to identify participating `send` functions. Note that, in this mode, the application cannot keep two disjoint autograd graphs alive at the same time, as that would break the assumption that all send (originated from the driver) will receive grad in the backward pass.


### Concurrent distributed Backward passes  

```python
A = torch.rand(2, 2)
B = torch.rand(2, 2)
    
# on all workers
def add() -> Tensor:
    global A, B
    return A + B
    
# on worker0
C = torch.remote(add, on="worker2").to_here()
C.sum().backward()

# on worker1
C = torch.remote(add, on="worker2").to_here()
C.sum().backward()

```


In the above example, there are two concurrent backward passes triggered by `worker0` and `worker1` respectively, and both will reach `worker2`. To avoid race, the distributed autograd engine will use the globally unique `autograd_context_id` to create a dedicated context on every participating worker. Later, pass this `autograd_context_id` to optimizer to apply gradients.  More concretely, this would work as follows:

1. Compute all the leaf nodes in the autograd graph.
2. As part of running distributed backwards, use the outputs parameter of the autograd engine to avoid executing `AccumulateGrad` for the leaf nodes we have and instead return the appropriate `output_edges` to execute for accumulating gradients.
3. Store the `output_edges` with the `autograd_context_id`. This would ensure multiple backward passes won't accumulate gradients in the same context. 
4. This completes the backward pass and gradients are accumulated in the autograd engine per `autograd_context_id.`
5. Now we run the optimizer on each of the worker nodes and pass the `autograd_context_id` to the optimizer.
6. The optimizer applies all the gradients to the leaf nodes that we computed originally.
7. The context and enclosing gradients should be destroyed when the `autograd_context_id` is destructed on the caller of `backward()`.


Some pseudo-code to illustrate this:

```python
optimizer = dist.optimizer(model)
loss = model(inputs)
bw_ctx_id = dist.autograd.backward(loss, timeout=60) # timeout of 60s
optimizer.step(bw_ctx_id)
```




## RRef 

(more details are described in #26759)

`RRef` is an important concept for building a distributed autograd graph. Each `RRef` is owned by a single worker (i.e., owner) and can be used by multiple users. The owner stores the real data referenced by its `RRef`s, and keeps track of the global reference counts for its `RRef`s. Every `RRef` can be uniquely identified by a global id `ref_id`, which is assigned at the time it is first created either on a user or on the owner. 

The owner only keeps one `RRef` instance for each data object, while users can fork as many `RRef` instances as necessary. All usage on the owner should retrieve the `RRef` instance using the globally unique `ref_id`. A fork of `RRef` will be created when it is used as an argument or return value in a RPC call, but users don't need to worry about forking/forwarding and reference counting (RC) `RRef`s. These will be handled transparently, and every fork will also have its own `fork_id`, which is guaranteed to be unique across all `RRef` instances for the same data object. 

`RRef` needs to support fast and scalable RPC. Hence, in the RC design, we avoid using any global master to keep `RRef` states. Besides, when worker X invokes RPC on worker Y, Y should be able to start immediately after receiving the RPC request, without waiting for any third-party owner Z (unless Y needs to pull real data from Z), even if neither X nor Y owns the `RRef`. We propose the following algorithm: 


1. If the owner is the RPC caller, the owner will update RC for the `RRef` accordingly. 
2. If the owner is the RPC callee, the owner will drop the new fork, and use the unique `RRef` id in the fork to access its singleton local `RRef` instance.
3. If the RPC is between two users:
    1. The caller sends an RPC message to the callee, and also notifies the owner on the new fork.
    2. The owner, upon receiving the notification, updates its local RC and then tells the callee the new fork is now known by the owner.
    3. The callee can starts executing the RPC as soon as it receives the RPC message from the caller, and does not need to wait for the message from the owner. However, it cannot delete its local `RRef` fork until owner's message arrives.

### Reference Count

The right time to delete an `RRef` on owner is when there are no living forks on any user and Python GC also agrees to delete the `RRef` instance on the owner. The tricky part is to determine if there are any living forks.

A user can get a fork in three situations:


1. Receiving a fork from the owner.
2. Receiving a fork from another user.
3. Creating a new `RRef` fork owned by another worker.


`#1` is the simplest case where the owner initiates the fork, and hence it can easily increase local RC. The only requirement is that any fork must notify the owner before destruction. Hence, we need the first guarantee:

* G1. The owner will be notified when any fork is deleted.*

Note that the notification might come delayed or out-of-order. 

With `#2` and `#3`, it is possible that the owner only partially knows the `RRef` fork graph or not even knowing it at all. For example, the `RRef` could be constructed on a user, and before the owner receives the RPC call, the creator user might have already shared the `RRef` with other users, and those users could further share the `RRef`. One invariant is that the fork graph of any `RRef` is a tree rooted at the owner, because forking an `RRef` always creates a new `RRef` instance, and hence every `RRef` has a parent. One nasty detail is that when an `RRef` is created on a user, technically the owner is not its parent but we still  consider it that way and it does not break the argument below.

The owner's view on any node (fork) in the tree has three stages 1) **unknown** →  2) **known** →  3) **deleted**, and the owner's view on the entire tree keeps changing. The owner deletes its `RRef` instance when it thinks there is no living forks, i.e., all the forks could be either indeed deleted or unknown. Therefore, the dangerous case is when some forks are unknown and others are deleted. We only need a simple guarantee to prevent this situation:

*G2. No fork x can be deleted on a user before the owner knows x’s parent fork.
* 
This works because owner's view on x can only change from **known** to **deleted** when x's parent is **known** or **deleted**. If the parent is **known**, owner will not delete local `RRef`. If the parent is **deleted**, this rule recursively applies to the parent's parent, until it reaches the root (owner). To implement the guarantee, we only need to make the caller include its own `fork_id` when notifying the owner on a new fork.

G1 and G2 guarantee correct RC, but does not prevent a user deleting before finishes its own prior RPC calls using that `RRef` fork. This should be OK, because when the caller deserializes the RPC message, it would hold a reference () to that `RRef`, preventing it from been deleted.


cc @pietern @mrshenli @pritamdamania87 @zhaojuanmao @satgera