Repository: pytorch
Issue: [RFC] RemoteTensor: Use Remote Devices as Local Devices
Labels: oncall: distributed,feature,triaged,module: rpc

Please comment below if `RemoteTensor` can be useful to your application. Please also let us know if the preliminary API thoughts missed anything. Thx!

## Motivation:

The [RPC package](https://pytorch.org/docs/master/rpc.html) offers a flexible way to allocate data and run computation on remote machines and remote devices. However, it requires users to take an extra step to decompose programs into smaller pieces and wrap each piece with an RPC, which is less elegant compared to how PyTorch handles local devices.  It will be a lot easier if users can access a remote device as if it is a local device. Hence, we proposal to add a `RemoteTensor` that allows users to easily use remote resources without decomposing their models.

Conceptually, a `RemoteTensor` can be viewed as an [`RRef`](https://pytorch.org/docs/master/rpc.html#torch.distributed.rpc.RRef) of a `Tensor`, but it also exposes a `Tensor`-like API. One naive solution could be convert every operation on a `RemoteTensor` as a standalone RPC. However, this can incur unacceptable overhead. Therefore, `RemoteTensor` should also come with RPC batching.




## API

Remote tensor and regular tensors mostly share same APIs, with two distinctions:

1. Applications needs to pass in the `RemoteDevice` in order to create a `RemoteTensor` E.g.:
    ```python
    # assume RPC is already initialized
    
    # on worker0
    x = torch.zeros(2, device="worker1/cuda:2")
    # x is a RemoteTensor instance
    ```
3. It exposes an additional `materialize()` API to explicitly kick off batched computations when needed. But it is not always necessary to explicitly call this API, as consuming any value of that `RemoteTensor` will also automatically trigger materialization. 

Let's use some examples to further explain `RemoteTensor` behaviors.

## Examples

**Remote Tensor**
Ed: batch RPC instead of tracing graph

```python
from torch.distributed import rpc

rpc.init_rpc(...)

# on worker0
x = torch.zeros(2, device="worker1/cuda:0")  # no RPC
y = x + 1  # no RPC
print(y)  # RPC0: run the above two lines in one RPC and fetch 
          # value of y to worker0

# value of x is not fetched to worker0 yet, but will be stored on worker1, and
# the local handle x will keep that value alive as if x it is an RRef

a = x + 2  # no RPC
b = a * 2  # no RPC
fut_b = b.materialize()  # RPC1: run the above two lines in one RPC and 
                         # return a torch.futures.Future of the result 
                         # Tensor.
                         # This RPC will also mark a and b as in-flight 
                         # and block subsequent RPCs reading a or b
                         # until this one is done. 

# We would like to guarantee that local ordering of reading/writing the
# same tensor is also preserved remotely.

c = a + 3  # no RPC
print(c)  # RPC2: a is in c's lineage, which will delay this RPC until 
          # RPCs using a with smaller seq no. is done, i.e., RPC1.

a = x + 4 # no RPC
fut_a = a.materialize()  # RPC3: this mutates a, and will wait for all
                         # preceding operators that read/write a.

```

Not yet planning to support passing `RemoteTensor` through RPC. When users pass a `RemoteTensor` as an RPC argument, it will be treated as the RPC needs to access the value of the `RemoteTensor` and hence will materialize it and fetch its value. If users would like to treat it as an `RRef`, they need to explicitly create an `RRef` (maybe through `RemoteTensor.rref_()`) The reason for this design is that, if there are different `RemoteTensor` pointing to the same `Tensor`, the above protocol for read/write contention won’t work.

If we don’t have `RemoteTensor`, similar logic can also be implemented using existing RPC APIs, but will be in a more verbose and less efficient way. See the code below. 

```python
from torch.distributed import rpc

rpc.init_rpc(...)

def create_x_and_y():
  x = torch.zeros(2).to(0)
  y = x + 1
  return rpc.RRef(x), rpc.RRef(y)
  
# note that this will actually be slower than the above
# RemoteTensor solution, as here will incur two RPC
# round trips
rref_x, rref_y = rpc.rpc_sync("worker1", create_x_and_y)
print(rref_y.to_here())
  
def create_a_and_b(rref_x):
  a = rref_x.local_value() + 2
  b = a * 2
  return rpc.RRef(a), rpc.RRef(b)
  
rref_a, rref_b = rpc.rpc_sync(
    "worker1", 
    create_a_and_b, 
    args=(rref_x,)
)

rref_b.to_here()

def create_c(rref_a):
  c = rref_a.local_value() + 3
  return rpc.RRef(c)
  
rref_c = rpc.rpc_sync("worker1", create_c, args=(rref_a,))
print(rref_c.to_here())

def mutate_a(rref_a):
  rref_a.local_value() += 4

rpc.rpc_sync("worker1", mutate_a, args=(rref_a,))
rref_a.to_here()
```

**`nn.Module` with Remote Tensor**


```python
from torch.distributed import rpc

rpc.init_rpc(...)

# on worker0
with torch.distributed.device("worker1"):
    net1 = nn.Linear(20, 20)  # no RPC: create a module with RemoteTensor 
                              # params on worker1
    net2 = nn.Linear(20, 20).to(1)  # no RPC: create and move the params 
                                    # to worker1/cuda:1
    
with torch.distributed.device("worker2/cuda:1")
    net3 = nn.Linear(20, 20)  # no RPC: create a module with RemoteTensor 
                              # params on worker1/cuda:1
    
x = torch.zeros(10, 20).to(net1.weight.device)  # no RPC
x = net1(x).to(net2.weight.device)  # no RPC
x = net2(x).to(net3.weight.device)  # no RPC
fut_x = net3(x).materialize()  # kick off all above in one RPC to worker1, 
                               # which will contain a nested RPC to worker2
```


**Combination with Distributed Autograd**

```python
from torch.distributed import rpc, autograd


class MyModule(nn.Module):
    def __init__(self):
        self.net1 = torch.nn.Linear(20, 20)  # local layer
        with torch.distributed.device("worker1/cuda:0"):
            self.net2 = torch.nn.Linear(20, 20)  # remote layer
            
    def forward(self, x):
        x = self.net1(x.to(self.net1.weight.device))
        x = self.net2(x.to(self.net2.weight.device))
        return x

rpc.init_rpc(...)

# on worker0
model = MyModule()  # model contains remote tensors
opt = torch.optim.SGD(model.parameters())  # using a local optimizer
with autograd.context() as ctx:
    inp = ...
    out = model(inp)  # No RPC, out is a remote tensor as well
    autograd.backward(ctx, [out.sum()])  # blocking and trigger all pending RPCs
                                         # we probably should also make 
                                         # out.backward() work for this case
    opt.step()  # no RPC
    
# flush pending RPCs of opt.step() in the last iteration. Can create a context 
# manager for this to simplify.
futs = []
for p in model.parameters():
    if isinstance(p, RemoteTensor):
        futs.append(p.materialize())
torch.futures.wait_all(futs)
```

cc @pietern @mrshenli @pritamdamania87 @zhaojuanmao @satgera @gqchen @aazzolini @rohan-varma @jjlilley @osalpekar @jiayisuse @agolynski @SciPioneer @H-Huang @mrzzd