Add DistributedSamplerWithLoop
# What does this PR do?

This PR adds a new distributed sampler that will provide a round multiple of the batch size samples on all processes by looping back at the beginning of the (shuffled) dataset. This is useful:
- for TPUs to avoid triggering a new XLA compilation for the last training batch
- for model parallelism to have batches of the same size on all processes

This PR also refactors some logic regarding the wold_size and process_rank in the `TrainingArguments`, as well as adds a test of the new `DistributedSamplerWithLoop`.

Tested on:
- single-GPU
- multi-GPU
- TPU
- SageMaker MP