[TPU] Support PyTorch/XLA FSDP via SPMD
Author: alanwaketan
Pull request: yes

# What does this PR do?

Summary:
This is the first attempt to enable FSDP via SPMD (FSDPv2) on PyTorch/XLA model.

More information about FSDPv2 can be found here:
1. A user guide: https://github.com/pytorch/xla/blob/master/docs/fsdpv2.md
2. A RFC: https://github.com/pytorch/xla/issues/6379

Besides the initial implementation of FSDPv2 in r2.2, this change will also requires the following changes in PyTorch/XLA:
1. https://github.com/pytorch/xla/pull/6499
2. https://github.com/pytorch/xla/pull/6500
3. https://github.com/pytorch/xla/pull/6498
4. https://github.com/pytorch/xla/pull/6525
Therefore, it will only be compatible with the nightly builds.

Example use cases:
1. Prepare a FSDPv2 config:
```json
{
    "fsdp_transformer_layer_cls_to_wrap": [
        "LlamaDecoderLayer"
    ],
    "xla": true,
    "xla_fsdp_v2": true,
    "xla_fsdp_grad_ckpt": true
}
```
2. Invoke the trainer using the following command:
```
XLA_USE_SPMD=1 XLA_USE_BF16=1 python3 examples/pytorch/language-modeling/run_clm.py     --num_train_epochs 1     --dataset_name wikitext     --dataset_config_name wikitext-2-raw-v1 --per_device_train_batch_size 128     --do_train     --output_dir /tmp/test-clm     --overwrite_output_dir     --config_name ../transformers_pt/2B.config     --cache_dir /tmp     --tokenizer_name hf-internal-testing/llama-tokenizer      --block_size 1024     --optim adafactor     --save_strategy no     --logging_strategy no --fsdp "full_shard" --fsdp_config fsdp_config.json --torch_dtype bfloat16 --dataloader_drop_last yes
```

## Before submitting
- [ ] 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/main/CONTRIBUTING.md#create-a-pull-request),
      Pull Request section?
- [ ] Was this discussed/approved via a Github issue or the [forum](https://discuss.huggingface.co/)? 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](https://github.com/huggingface/transformers/tree/main/docs), and
      [here are tips on formatting docstrings](https://github.com/huggingface/transformers/tree/main/docs#writing-source-documentation).
- [ ] Did you write any new necessary tests?


## Who can review?

@ArthurZucker @younesbelkada 

Comments:
Can HF folks point me on how to add test case in this case and also how to update the documentation?

cc @yeounoh @jonb377

Tests should be added in the `tests/trainer/test_trainer.py` file. You should find similar tests! 

> As @ArthurZucker hinted at, we now don't handle things like this in the trainer directly. I would rather see this code over in accelerate which we can then bring into Trainer automatically since it relies on it for preparation. Especially as this deals with the dataloaders. Would that be possible please! :)

Can you elaborate it a bit more? I can move the `model = model.to(xm.xla_device())` logic. But for the dataloader logic, i.e., tpu_spmd_dataloader, where do you suggest me to move it to?  

> Tests should be added in the `tests/trainer/test_trainer.py` file. You should find similar tests!

Speaking of adding tests, what should I test? I mean do you have TPU CI?

The test failures don't seem to be related. I tried rebasing as well.

Thanks @ArthurZucker and @muellerzr for approving the change.

The docs for this PR live [here](https://moon-ci-docs.huggingface.co/docs/transformers/pr_28949). All of your documentation changes will be reflected on that endpoint. The docs are available until 30 days after the last update.

It's all green. Can HF folks help with landing the PR? Appreciate it.

I can merge :) Thanks for adding this support @alanwaketan!