Repository: pytorch
Issue: rfc: automating the switching of inputs to the device of the params
Labels: oncall: distributed,feature,module: cuda,module: rpc,needs design,module: ux

## 🚀 Feature

**tldr**: 

Figure out a way for pytorch to automatically switch arguments of pytorch operations to the device of the params involved in the operation. Primarily to ease the facilitation of model parallelism, but it may have other uses.

----------------------

MP = Model Parallelism

1. We have the simple case of having a model run on a single GPU (DDP-included) and needing to switch model inputs to that GPU. We need a single `.to(device)` call to switch the input to the device of the model and we are done, regardless if the model has dozens of layers or just one.

2. Then we have model parallelism where different layers will be scattered across multiple devices. Things become very complicated now because it's far from trivial to know which device the inputs need to be switched to.

The most simple approach that will always work is to just always add: `to(self.get_device)` to all input arguments in `forward`, e.g.:

```
def self.get_device(self): return next(self.parameters(recurse=True)).device
def forward(self, input):
    x = self.linear(input.to(self.get_device))
```
this would work without needing multiple `if MP: ....` which some of the current MP implementations use with a slight overhead for the non-MP case.

3. but this makes the code hard to read and so why not do this automatically before `forward` as in:

```
def model_parallel_inputs_to_device(func):
    """
    This decorator will try to find a at least one parameter to read layer's .device from and then will
    automatically copy any inputs to that device before `forward` is called.
    this will work only if all params of this layer are on the same device
    """

    def _call__mp(self, *input, **kwargs):

        try:
            device = next(self.parameters(recurse=True)).device
        except StopIteration:
            device = None

        if device is not None:
            input = list(input)
            for i, v in enumerate(input):
                if v is not None:
                    input[i] = v.to(device)
            input = tuple(input)

            for k in kwargs.keys():
                if kwargs[k] is not None and torch.is_tensor(kwargs[k]):
                    kwargs[k] = kwargs[k].to(device)

        return func(self, *input, **kwargs)

    return _call__mp

    
@model_parallel_inputs_to_device
def forward(self, input):
    x = self.linear(x)    
```
pass as many inputs as you want and they are already on the right device MP or not.

So now we get an automatic switch of all inputs to the same device as the params of the model. We assume here that it is so - that is that all params of that layer are on the same model. 

But this is not always the case, as for example shown in the trivial example of https://pytorch.org/tutorials/intermediate/model_parallel_tutorial.html#apply-model-parallel-to-existing-modules

4. So my thinking is this - **the next stage is to extend pytorch ops to automatically switch the inputs to the device of the params at the point of use**, so if there is a `matmul` look at the device of the "designated" argument and switch all the other args to the same device, rather than assert that they have to be on the same device.

I totally understand that the devil is in the detail and how could one figure out what that "designated" argument is, so that, say, we remain on gpu1 if the other arg is on cpu and not mistakenly copy data away from gpu. How do we define DWIM (do what I mean) since it's not always the first arg that "leads". What I'm trying to communicate is that each pytorch op then needs a way to know which argument sets the device for the rest of the arguments of that op.

Or perhaps there should be a way of switching some global flag that unless specified otherwise will switch all args to that device before doing any operation.

The goal of this is to be able to implement model parallelism by simply spreading out the params of a single model through multiple devices and having pytorch automagically sort out the rest.

I realize that perhaps my vision is not quite feasible at the moment, but perhaps there are some ways that can ease this process. I think the decorator I proposed here is a small step towards such automation, but surely more can be done.

Thank you for reading and if you have some inspirations please don't hesitate to share.





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