[TPU] Support PyTorch/XLA FSDP via SPMD

[alanwaketan]

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: [RFC] FSDP via SPMD pytorch/xla#6379

Besides the initial implementation of FSDPv2 in r2.2, this change will also requires the following changes in PyTorch/XLA:

1. [FSDPv2] Enable auto-wrapping pytorch/xla#6499
2. [FSDPv2] Use the global mesh API pytorch/xla#6500
3. [SPMD] Introduce global mesh pytorch/xla#6498
4. [FSDPv2] Move the module to xla device pytorch/xla#6525
   Therefore, it will only be compatible with the nightly builds.

Example use cases:

1. Prepare a FSDPv2 config:

{ "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).
• 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?

@ArthurZucker @younesbelkada

[alanwaketan]

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

[alanwaketan]

cc @yeounoh @jonb377

[ArthurZucker]

LGTM overall! We might want to add a small test, it can be done in a followup PR.
Pinging @muellerzr for a second look!

[ArthurZucker]

I am not super fan of super short names but seems common in trainer!

[ArthurZucker]

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

[muellerzr]

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! :)

[muellerzr]

Could we make this easier by importing FSDPv2 as FSDP instead?

[alanwaketan]

May I ask what's the benefits of doing so?

[muellerzr]

And then leave the check for down here on what to do.

[alanwaketan]

shard_output is not used by FSDPv1. Shouldn't we guard that with the flag too?

[alanwaketan]

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?

[alanwaketan]

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?

[alanwaketan]

I fixed a bug here. cc @ArthurZucker @jonb377

[alanwaketan]

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

[alanwaketan]

Thanks @ArthurZucker and @muellerzr for approving the change.

[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.

[alanwaketan]

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

[amyeroberts]

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