Fix GQA permutation computation and sequential weight initialization / loading when doing TP

[michaelbenayoun]

What does this PR do?

• Fixes the way indices are compute for GQA permutation of the query and output projection, and add a test case to make sure everything works
• Add the possibility to specify the number of concurrent ranks that can initialize or load the model weights at the same time under TP. It can be useful to avoid going out-of-memory.
• Fixes a typo indicies -> indices

[HuggingFaceDocBuilderDev]

The docs for this PR live here. All of your documentation changes will be reflected on that endpoint. The docs are available until 30 days after the last update.

[michaelbenayoun]

Ran the distributed tests locally and they pass. It's just a flaky test that needs to be solved.