Add parallelization support for T5EncoderModel

[agemagician]

What does this PR do?

Extend T5EncoderModel to support model parallization across different GPUs.

Before submitting

• This PR fixes a typo or improves the docs (you can dismiss the other checks if that's the case).
• Did you read the contributor guideline,
  Pull Request section?
• Was this discussed/approved via a Github issue or the forum? Please add a link
  to it if that's the case.
• Did you make sure to update the documentation with your changes? Here are the
  documentation guidelines, and
  here are tips on formatting docstrings.
• Did you write any new necessary tests?

Who can review?

T5: @patrickvonplaten

[LysandreJik]

Very cool! Could you also enable the parallelization tests for these models? You can check how it was done in the initial model parallel PR, here's the commit related to the tests. You can just add the T5EncoderModel to the all_parallelizable_model_classes attribute of the T5ModelTester class.

[agemagician]

Very cool! Could you also enable the parallelization tests for these models? You can check how it was done in the initial model parallel PR, here's the commit related to the tests. You can just add the T5EncoderModel to the all_parallelizable_model_classes attribute of the T5ModelTester class.

Thanks for the tip.
Done, please let me know if anything else is needed from my side.

[LysandreJik]

This LGTM. Looking into it it seems we have an error in T5Stask as it is creating the device map with torch.cuda.device_count(), rather than the range of that value like you're doing it here. Since we're always passing the device map to T5Stack (it's never used as a standalone model) we don't see it, but it doesn't seem correct.

What do you think? If you think this is true, do you mind adding a range in T5Stack so that we can merge it together? Thanks!

[LysandreJik]

Also it would be great if you could run make style && make quality or make fixup to solve the quality issues.

[agemagician]

This LGTM. Looking into it it seems we have an error in T5Stask as it is creating the device map with torch.cuda.device_count(), rather than the range of that value like you're doing it here. Since we're always passing the device map to T5Stack (it's never used as a standalone model) we don't see it, but it doesn't seem correct.

What do you think? If you think this is true, do you mind adding a range in T5Stack so that we can merge it together? Thanks!

Yes, you are correct, T5Stack should also use range. Since "get_device_map" function apply len to it .
I have updated T5Stack using range.

[agemagician]

Also it would be great if you could run make style && make quality or make fixup to solve the quality issues.

Done and passed the code quality testing.

[LysandreJik]

Wonderful!