ByT5 parallelization
## ByT5 parallelization

Regarding ByT5, we observed the following and it looks like a potential bug:
ByT5 models are handled by HF through the same set of classes that handle T5 models, for instance T5ForConditionalGeneration [link](https://github.com/huggingface/transformers/blob/ad0cba08ea295cd7450484468dd34b53816c85fb/src/transformers/models/t5/modeling_t5.py#L1456)
One issue this creates is the way parallelization is performed by calling the parallelize method [link](https://github.com/huggingface/transformers/blob/ad0cba08ea295cd7450484468dd34b53816c85fb/src/transformers/models/t5/modeling_t5.py#L1494)
The method allows using a single device map to distribute attention blocks across devices. The device map is either inferred from the number of layers of the encoder or can be passed as an argument.
The bug stems from the fact that while T5 model have the same number of layers for encoder and decoder, this is not the case for ByT5 models which have a 3:1 structure (see for instance [here](https://github.com/google-research/byt5/blob/master/byt5/gin/models/byt5.xxl.gin#L4-L5) or Table 1 in the [arXiv paper](https://arxiv.org/pdf/2105.13626.pdf)).
We worked around this by reimplementing parallelize to have different encoder and decoder configurations:

```
encoder_device_map = (
    get_device_map(len(model.encoder.block), range(torch.cuda.device_count()))
    if encoder_device_map is None
    else encoder_device_map
    )
decoder_device_map = (
    get_device_map(len(model.decoder.block), range(torch.cuda.device_count()))
    if decoder_device_map is None
    else decoder_device_map
    )
```

Gently pinging @stas00 here. Do we still use this method to parallelize models such as T5? How do we currently parallelize models like T5?
Until we implement a proper PP (pipeline parallelism) (if we ever do) we only have this naive approach for PP and yes, it assumes a balanced number of layers between encoder and decoder. 

There is absolutely no reason to use the naive PP (pipeline parallelism) as Deepspeed ZeRO solves this issue far more efficiently. but we left these in T5 and GPT2 models for now in case someone still uses those.

The next best alternative should someone not want to use Deepspeed ZeRO is TP (tensor parallelism), which soon will come via OSLO and Deepspeed-Inference projects.

For now you can probably just assert if someone tries to use the naive PP with other than T5, since it was designed for T5 in mind. Or you can fix it to handle unbalanced number of layers but it might be a waste of time if nobody uses it.<|||||>Thanks a lot @stas00 !<|||||>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/main/CONTRIBUTING.md) are likely to be ignored.