text: 
  
  TITLE
  [T5 model parallel] implement input auto-relocation + lots of refactoring/cleanup
  
  COMMENTS
  16
  
  REACTIONS
  +1: 1
  -1: 0
  laugh: 0
  hooray: 0
  heart: 0
  rocket: 0
  eyes: 0
  
  BODY
  As I commented on in another incarnation of generalizing t5 model parallelism https://github.com/huggingface/transformers/pull/9316 so that it could be easily ported to other models I realized that it's quite unnecessary to try and remap inputs to specific devices where they will be needed in the future ahead of time. Since we have `forward` where we have access to the device of the parameters of that layer - we can completely automate the relocation of inputs to the correct devices just before `forward` is called. So this PR builds upon https://github.com/huggingface/transformers/pull/9316 and:
  
  * [x] creates `@model_parallel_inputs_to_device` decorator used for `forward`, which automatically takes any inputs and puts them on the same device as the parameters of that layer. This allowed a complete removal of most of the `.to()` juggling logic for inputs, which was quite complex and noisy.
  * [x] a lot of refactoring to make the MP as little invasive and noisy as possible, and fixing some small issues on the way.
  
  I have tested this with:
  ``` 
  pyt -sv tests/test_modeling_t5.py -k parallel
  ```
  Which I'm not sure covers all bases, but the above tests pass.
  
  @alexorona, please let me know what you think. And if you have real applications besides the great tests you wrote please see if it still works correctly. (It was so awesome having those tests in place! Thank you!) If it looks good and others support this proposal we can then look at doing the same for gpt2 and meanwhile I will look at bart.
  
  @patrickvonplaten, @LysandreJik