Repository: pytorch/pytorch

Issue:
Problems in TensorPipeRpcBackendOptions device mapping documentation?
## 📚 Documentation

The new release of PyTorch 1.8 introduces CUDA-support in RPC.
I've referred to the RPC documentation, and the only reference for the CUDA-support I could find is under [`TensorPipeRpcBackendOptions`](https://pytorch.org/docs/1.8.0/rpc.html#torch.distributed.rpc.TensorPipeRpcBackendOptions) and [`set_device_map`](https://pytorch.org/docs/1.8.0/rpc.html#torch.distributed.rpc.TensorPipeRpcBackendOptions.set_device_map).
Seems like setting up CUDA-support is simply done by supplying a device mapping in the `TensorPipeRpcBackendOptions`, pretty cool.

However, I find the documentation for the `device_maps`/`device_map` to be unclear. It seems that `TensorPipeRpcBackendOptions`'s `device_maps` is a dictionary where the keys are worker names, but I'm not exactly sure what the structure of the dictionary's values should be like? Supposedly each value should be some sort of dictionary (as indicated by the parameter's type - `Dict[str, Dict]`), yet the example code provides a set:  `device_maps={"worker1": {0, 1}}`. I don't really understand how does this "map worker0's cuda:0 to worker1's cuda:1"?

Same for `set_device_map`'s `device_map`, the parameter's type also indicates it's a dictionary (`(Dict of python:int, str, or torch.device)`), but doesn't quite explain its structure. And again, the example code provides a set: `options.set_device_map("worker1", {1, 2})`.

It is also not explained how to define a GPU->CPU mapping (or vice versa).

Apart for this, there are 2 obvious errors in the example code provided in that documentation:

1. There is a missing comma in the following part:
```python
>>> rpc.init_rpc(
>>>     "worker0",
>>>     rank=0,
>>>     world_size=2  # <-- missing comma
>>>     backend=rpc.BackendType.TENSORPIPE,
>>>     rpc_backend_options=options
>>> )
```
2. I don't see how it is possible that those two `print`s will give different results. I'm guessing that the second line should read `print(rets[1])`?
```python
>>> print(rets[0])  # tensor([2., 2.], device='cuda:0')
>>> print(rets[0])  # tensor([2., 2.], device='cuda:1')
```

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


Discussion:
Hey @rafi-cohen, thanks a lot for spotting the error in the doc, we will fix that asap.

Currently, the CUDA RPC support is a prototype feature and at an early stage. Please use nightly builds to get the most recent implementation. The tests below can serve as a reference of how it can be used:

https://github.com/pytorch/pytorch/blob/d54be1a9467db5075256d229ea1b01f1a4bcba8d/torch/testing/_internal/distributed/rpc/rpc_test.py#L4625-L5359

> It is also not explained how to define a GPU->CPU mapping (or vice versa).

It does not support direct GPU-to-CPU communication yet. We would like to hear your use case. Is GPU-to-CPU/CPU-to-GPU a requirement for your application? Is today's default mapping API sufficient for your use case or do you need per-RPC device map? 




@rafi-cohen please let me know if you would like to see any other information added to #53508.

We are testing this feature on some large models internally. When we are confident there are no major gaps/flaws in the implementation, we will add a tutorial/recipe for it, likely together with when we graduate this to a beta feature in the next release. 

Resolving patch:
diff --git a/torch/distributed/rpc/options.py b/torch/distributed/rpc/options.py
--- a/torch/distributed/rpc/options.py
+++ b/torch/distributed/rpc/options.py
@@ -73,16 +73,16 @@ def set_device_map(self, to: str, device_map: Dict):
             >>> # on worker 0
             >>> options = TensorPipeRpcBackendOptions(
             >>>     num_worker_threads=8,
-            >>>     device_maps={"worker1": {0, 1}}
+            >>>     device_maps={"worker1": {0: 1}}
             >>>     # maps worker0's cuda:0 to worker1's cuda:1
             >>> )
-            >>> options.set_device_map("worker1", {1, 2})
+            >>> options.set_device_map("worker1", {1: 2})
             >>> # maps worker0's cuda:1 to worker1's cuda:2
             >>>
             >>> rpc.init_rpc(
             >>>     "worker0",
             >>>     rank=0,
-            >>>     world_size=2
+            >>>     world_size=2,
             >>>     backend=rpc.BackendType.TENSORPIPE,
             >>>     rpc_backend_options=options
             >>> )
@@ -94,7 +94,7 @@ def set_device_map(self, to: str, device_map: Dict):
             >>> # the device map, and hence will be moved back to cuda:0 and
             >>> # cuda:1 on worker0
             >>> print(rets[0])  # tensor([2., 2.], device='cuda:0')
-            >>> print(rets[0])  # tensor([2., 2.], device='cuda:1')
+            >>> print(rets[1])  # tensor([2., 2.], device='cuda:1')
         """
         device_index_map = {}
         curr_device_maps = super().device_maps


Test patch: