[WIP] Adding model parallelism for T5 (should work for other models as well)
Author: thomwolf
Pull request: yes
Labels: name: Distributed Training / Models

This PR adds:
- a `get_block_list()` utility method which returns a list of the blocks in a Transformers model (currently only added on T5). Block can be Modules or list/tuple of Modules (if a single transformer block is spread in several ModuleList like in XLM).
- a `spread_on_devices(devices: Optional[List] = None)` method to spread a model on several devices by spreading the transformers blocks (roughly) evenly on the provided device list or all visible CUDA devices if no device list is given. The first device will host the remaining non-block modules in addition (the embeddings usually).

Currently, the code is in the T5 model but should be generic enough to be applied to other models if needed.

To use:
``` python
model = T5ForConditionalGeneration.from_pretrained('...')
model.spread_on_devices()  # Will spread on all visible CUDA devices by default
input = torch.tensor([...]).to('cuda:0')  # Inputs and outputs are on the first device
model(input)  # you should probably use only positional arguments for the forward pass (see spread_on_devices's docstring)
```

TODO:
 - [ ] try it
 - [ ] add tests if possible (on a dummy device list like ['cpu', 'cpu']?)

cc @patrickvonplaten @craffel

Comments:
# [Codecov](https://codecov.io/gh/huggingface/transformers/pull/3578?src=pr&el=h1) Report
> Merging [#3578](https://codecov.io/gh/huggingface/transformers/pull/3578?src=pr&el=desc) into [master](https://codecov.io/gh/huggingface/transformers/commit/a4ee4da18ad659b196582bbdf40785033ee1d26b?el=desc) will **decrease** coverage by `0.10%`.
> The diff coverage is `12.90%`.

[![Impacted file tree graph](https://codecov.io/gh/huggingface/transformers/pull/3578/graphs/tree.svg?width=650&height=150&src=pr&token=9qOlN6Hb1c)](https://codecov.io/gh/huggingface/transformers/pull/3578?src=pr&el=tree)

```diff
@@            Coverage Diff             @@
##           master    #3578      +/-   ##
==========================================
- Coverage   78.05%   77.94%   -0.11%     
==========================================
  Files         100      100              
  Lines       17135    17166      +31     
==========================================
+ Hits        13374    13380       +6     
- Misses       3761     3786      +25     
```


| [Impacted Files](https://codecov.io/gh/huggingface/transformers/pull/3578?src=pr&el=tree) | Coverage Δ | |
|---|---|---|
| [src/transformers/modeling\_t5.py](https://codecov.io/gh/huggingface/transformers/pull/3578/diff?src=pr&el=tree#diff-c3JjL3RyYW5zZm9ybWVycy9tb2RlbGluZ190NS5weQ==) | `77.18% <12.90%> (-4.34%)` | :arrow_down: |
| [src/transformers/modeling\_tf\_utils.py](https://codecov.io/gh/huggingface/transformers/pull/3578/diff?src=pr&el=tree#diff-c3JjL3RyYW5zZm9ybWVycy9tb2RlbGluZ190Zl91dGlscy5weQ==) | `93.62% <0.00%> (+0.32%)` | :arrow_up: |

------

[Continue to review full report at Codecov](https://codecov.io/gh/huggingface/transformers/pull/3578?src=pr&el=continue).
> **Legend** - [Click here to learn more](https://docs.codecov.io/docs/codecov-delta)
> `Δ = absolute <relative> (impact)`, `ø = not affected`, `? = missing data`
> Powered by [Codecov](https://codecov.io/gh/huggingface/transformers/pull/3578?src=pr&el=footer). Last update [a4ee4da...3bfeebe](https://codecov.io/gh/huggingface/transformers/pull/3578?src=pr&el=lastupdated). Read the [comment docs](https://docs.codecov.io/docs/pull-request-comments).


This issue has been automatically marked as stale because it has not had recent activity. It will be closed if no further activity occurs. Thank you for your contributions.


Hello! Do you have plans to merge this feature to master branch? 
I tried to make it locally in clonned repo but I got an error while tried to use it:
<ipython

> -input-22-5591bd8e45c0> in main()
>     143         cache_dir=model_args.cache_dir,
>     144     )
> --> 145     model = model.spread_on_devices(['cpu', 'cpu'])
>     146 
>     147     # Get datasets
> 
> /usr/local/lib/python3.6/dist-packages/transformers/modeling_t5.py in spread_on_devices(self, devices)
>     936             return
>     937 
> --> 938         modules_to_move = set(self.modules)
>     939 
>     940         # Evenly spread the blocks on devices
> 
> TypeError: 'method' object is not iterable

Hey @exelents, 

At the moment I don't think anybody is working on it and I'm not sure what the importance of this PR is at the moment. Feel free to take over the PR and try to make it work. I would be more than happy to help you if you open a PR :-) 

This is very much related: https://github.com/huggingface/transformers/issues/7526

This issue has been automatically marked as stale because it has not had recent activity. If you think this still needs to be addressed please comment on this thread.

Please note that issues that do not follow the [contributing guidelines](https://github.com/huggingface/transformers/blob/master/CONTRIBUTING.md) are likely to be ignored.

This feature was awesome! I think this would be a major improvement to the transformers package!