Repository: pytorch
Issue: [RFC] Device Placement API for RPC 
Labels: feature,triaged,module: rpc,module: tensorpipe

## Goal

As of v1.6, `torch.distributed.rpc` only accepts CPU tensors. Applications have to first move the tensor to CPU and then pass it to RPC APIs. This is both inconvenient and slow. The original reason for enforcing CPU tensors is to avoid invalid devices on the destination process. The RPC device placement API aims at (partly) addressing this issue by allowing applications to configure the device mapping on caller and callee. Note that this proposal focuses on changes in API and RPC agents. There will be separate efforts in the comm layer and the distributed autograd engine to improve its efficiency. 

## API

We can add one `map_location` argument to [`RpcBackendOptions`](https://pytorch.org/docs/master/rpc.html#torch.distributed.rpc.RpcBackendOptions). The example below shows how to configure this field. 

```python
rpc.init_rpc(
    "worker0",
    rank=0,
    world_size=2,
    backend=rpc.BackendType.TENSORPIPE,
    rpc_backend_options=rpc.TensorPipeRpcBackendOptions(
        num_worker_threads=8,
        rpc_timeout=20, # 20 second timeout
        map_location={
        	"worker1": {
        		"cpu":     	"cuda:0",
        		"cuda:0":  	"cuda:1"
        	},
        	"worker2": {
        		"cpu": 		"cuda:1",
        		"cuda:0":	"cuda:2"
        	}
        }
    )
)

x = torch.zeros(2).to(0)
y = torch.ones(2).to(0)

# x and y will be materialized on "cuda:1" in "worker1".
# the return value z will be materialized on "cuda:0" on "worker0",
# which is defined by the inverse of the mapping specified by "worker0"
z = rpc.rpc_sync("worker1", torch.add, args=(x, y))

# x and y will be materialized on "cuda:2" in "worker2".
# the return value z will be materialized on "cuda:0" on "worker0".
z = rpc.rpc_sync("worker2", torch.add, args=(x, y))

rpc.shutdown()
```

In this proposal, we only aim to add default `map_location`. If necessary, we can add the per-RPC `map_location` later. 

This is only intended to be added to the `TensorPipe` backend. `ProcessGroup` backend will throw an exception is this field is specified.


### Device Placement for Return Value

There are a few options here.

* Option 1: we can use the inverse of the `map_location` to place tensors in the return value.
* Option 2: we can add another `return_map_location` for this purpose. 
* Option 3: use the `map_location` defined on the destination to send return values. 

The downside of Option 3 is that the caller cannot fully control the behavior of its own RPC. Imagine a Parameter-Server
training application with one PS with multiple trainers. Option 3 would requre the PS to define the `map_location` entry
for all trainers, which is cumbersome. So it might be more convenient for applications to use Option 1 or 2.  

As the first version, we can go with Option 1, and when we see requests for Option 2, we can add a `return_map_location` and
use entries defined there to override the inverse of `map_location`, so that there will be no BC issue.


## Implementation

In the `init_rpc` call, all processes will use the initialized RPC to communicate `map_location` to make sure that there 
are no invalid devices. If there is any in any process, all processes will throw an error. 

This `map_location` will be stored in the agent and not necessary to modify `Message` structure. 

As there is no remote device type yet, agent will first use integers to store devices, -1 for CPU. 

cc @pietern @mrshenli @pritamdamania87 @zhaojuanmao @satgera @gqchen @aazzolini @rohan-varma @xush6528 @jjlilley @osalpekar @jiayisuse @lw @beauby