[t5 model parallel] misc fixes
Author: stas00
Pull request: yes
Labels: name: Model Parallel

name: WIP

This PR:

* in 2 places fixes an assumption that devices on the device map are always ` (0, 1, 2, 3)` and:
   1. are ordered by their cuda device id and not `(2, 3, 0, 1)`
   2. have a stride of 1 and not `(0, 2, 3, 5) `
* adds a missing `to()`, removes a redundant `to()`
* removes obvious comments
* removes code that gets run twice
* this PR continues at #9323 -  I branched off from this PR and implemented an automatic remap of inputs and a lot refactoring. 

I will comment on the reasons for changes in the code.

There is one gotcha wrt py36 w/o cython not having its dict ordered. Please see https://github.com/huggingface/transformers/pull/9316#discussion_r549068073

I think sorting out the logic first device/last device/is_this_the_last_layer_of_this_device and such logic should be abstracted away for readability, and not needing to replicate the same logic in each model. Perhaps `self.device_map` should be a smart class that can provide all the answers via its methods.

@alexorona, I'm studying your t5-mp implementation to do the same for bart. Thank you for doing the hard work of putting the foundation in place and porting 2 models!!!

Please have a look and let me know if my tweaks make sense. Your original code is excellent - I'm just trying to think how to make it easier to replicate it in other models and improve readability, hence a gazillion of questions/suggestions.

Also, if you don't mind I have a few design questions:

1. Could you please comment on why you are splitting the encoder between all devices on the device map and the same for the decoder? Won't it be more efficient performance-wise to put the encoder on the first group of devices and decoder on the second? 

2. I also find it confusing that the device map doesn't map out the whole model, but just the encoder and assumes that the decoder has the same config. I'm not familiar with t5 but other models definitely can have encoder and decoder that don't at all match number of layers-wise. And while this is perhaps not the case for t5, I think the device map should be intuitively similar for all models as we slowly progress with porting other models to MP. That is I think it should include all layers of the model and not half of them.

@patrickvonplaten, @LysandreJik



Comments:
Somehow I'm feeling that this approach of having special logic to remap inputs ahead of time is over-complicated. I haven't tried it yet, but won't it be much much simpler to remap inputs once the model layer is visible and just before they are used by that layer - i.e. at the point where one gets:
```
RuntimeError: Expected all tensors to be on the same device, but found at least two devices, cuda:1 and cuda:0!
```
Then we just do `input_foo = input_foo.to(next(self.parameters()).device)` and we are done? No logic required other than just put the inputs on the same device as this layer.

We might be able to even remove all those `if self.model_parallel` in most places in `forward`, and have the same code for w/ or w/o MP. Perhaps with some wrapper that will be noop when not under MP. It could also handle `None` to avoid a gazillion of `if not None` checks. I'd also make it an in-place operation, just like `nn.Module.to` does.


**edit** I branched off from this PR and implemented this - works amazingly simple: https://github.com/huggingface/transformers/pull/9323

I also think #9323 is the way to go here

too long. closing.