Updates in Trainer to support new features in SM Model Parallel library
# What does this PR do?


- Calls state dict on rdp_rank==0 to fix hang for SMP when using tensor parallelism
- Reorders wrapping of optimizer to happen after wrapping of model, which is a requirement for the newer version of SMP.
- Uses the right ranks for data loading when using TP based on whether we are using prescaled batch or not.
Prescaled batch is a new construct in SMP. SMP does tensor parallelism within a DP group. What this means is that it shuffles around the samples read by all ranks so that the rank which owns the shard of a layer/parameter can see relevant portion of samples for all samples across the data parallel group. To reduce the communication before layer execution, and to support cases where we can't use a large batch size (i.e. all ranks can't get their own sample), we recommend using prescaled batch, where all ranks in same rdp group read the same batch of data. Hence this usage affects the number of examples seen. From data loading perspective, the number of ranks which read different data changes in prescaled and non-prescaled case. World size and process index in transformers seem to relate to number of times data parallelism happens in the job, hence this change.

More details here
https://sagemaker.readthedocs.io/en/stable/api/training/smd_model_parallel_general.html#prescaled-batch


Merges https://github.com/huggingface/transformers/pull/15804, https://github.com/huggingface/transformers/pull/15796/, https://github.com/huggingface/transformers/pull/15811/

<!--
Congratulations! You've made it this far! You're not quite done yet though.

Once merged, your PR is going to appear in the release notes with the title you set, so make sure it's a great title that fully reflects the extent of your awesome contribution.

Then, please replace this with a description of the change and which issue is fixed (if applicable). Please also include relevant motivation and context. List any dependencies (if any) that are required for this change.

Once you're done, someone will review your PR shortly (see the section "Who can review?" below to tag some potential reviewers). They may suggest changes to make the code even better. If no one reviewed your PR after a week has passed, don't hesitate to post a new comment @-mentioning the same persons---sometimes notifications get lost.
-->

<!-- Remove if not applicable -->

Fixes # (issue)


## Before submitting
- [x] This PR fixes a typo or improves the docs (you can dismiss the other checks if that's the case).
- [x] Did you read the [contributor guideline](https://github.com/huggingface/transformers/blob/master/CONTRIBUTING.md#start-contributing-pull-requests),
      Pull Request section?
- [x] Was discussed in a Slack channel
- [x] Did you make sure to update the documentation with your changes? Here are the
      [documentation guidelines](https://github.com/huggingface/transformers/tree/master/docs), and
      [here are tips on formatting docstrings](https://github.com/huggingface/transformers/tree/master/docs#writing-source-documentation).
- [ ] Did you write any new necessary tests?


## Who can review?

Anyone in the community is free to review the PR once the tests have passed. Feel free to tag
members/contributors who may be interested in your PR.

<!-- Your PR will be replied to more quickly if you can figure out the right person to tag with @

 If you know how to use git blame, that is the easiest way, otherwise, here is a rough guide of **who to tag**.
 Please tag fewer than 3 people.

Models:

- albert, bert, xlm: @LysandreJik
- blenderbot, bart, marian, pegasus, encoderdecoder,  t5: @patrickvonplaten, @patil-suraj
- longformer, reformer, transfoxl, xlnet: @patrickvonplaten
- fsmt: @stas00
- funnel: @sgugger
- gpt2: @patrickvonplaten, @LysandreJik
- rag: @patrickvonplaten, @lhoestq
- tensorflow: @LysandreJik

Library:

- benchmarks: @patrickvonplaten
- deepspeed: @stas00
- ray/raytune: @richardliaw, @amogkam
- text generation: @patrickvonplaten
- tokenizers: @n1t0, @LysandreJik
- trainer: @sgugger
- pipelines: @LysandreJik

Documentation: @sgugger

HF projects:

- datasets: [different repo](https://github.com/huggingface/datasets)
- rust tokenizers: [different repo](https://github.com/huggingface/tokenizers)

Examples:

- maintained examples (not research project or legacy): @sgugger, @patil-suraj
- research_projects/bert-loses-patience: @JetRunner
- research_projects/distillation: @VictorSanh

 -->

The docs for this PR live [here](https://moon-ci-docs.huggingface.co/docs/transformers/pr_15877). All of your documentation changes will be reflected on that endpoint.<|||||>Lot of TF tests seem to fail, and there are a couple of strange errors. Could you take a look?
<|||||>Thanks for your PR! Please rebase on master to remove the dev documentation failure.<|||||>Done<|||||>Thanks again for your PR!