Enable PyTorch/XLA Fully Sharded Data Parallel (FSDP)
Author: AlexWertheim
Pull request: yes

# What does this PR do?

This PR enables the user to make use of the [PyTorch/XLA implementation of FSDP](https://github.com/pytorch/xla/tree/master/torch_xla/distributed/fsdp), including the newly added [auto-wrap feature](https://github.com/pytorch/xla/pull/4318). Four arguments have been added to `training_args.py` to facilitate this functionality:

- `xla_fsdp`: this flag is a string containing the location of a `.json` file which specifies the FSDP arguments the user wants to use when wrapping their model.
- `xla_fsdp_min_num_params`: this flag is an int which will set a size-based automatic wrapping policy which automatically FSDP wraps any module with at least `xla_fsdp_min_num_params` many parameters.
- `xla_fsdp_transformer_layer_cls_to_wrap`: this flag is a list of (case-sensitive) strings which will set a layer-class-based automatic wrapping policy which automatically FSDP wraps any module whose name matches one of the listed strings.
- `xla_fsdp_grad_ckpt`: this flag is a bool which determines whether gradient checkpointing is enabled for the automatically wrapped layers.

# Design notes and future work

1) This PR is an updated version of [this closed PR](https://github.com/huggingface/transformers/pull/20774), which enabled FSDP for a more restricted class of models. This PR now enables nested FSDP wrapping via two auto-wrap policies, avoiding the restrictions of the previous PR.
2) For very large model sizes (greater than, say, 128B parameters), users may see host-side OOMs on TPUs during initialization. This can be mitigated by initializing layer weights immediately after construction, wrapping with FSDP, and moving onto the XLA device, as can be seen in [this branch](https://github.com/AlexWertheim/transformers/blob/einsum/src/transformers/models/gpt2/modeling_gpt2.py#L690-L723). We opted to enable FSDP wrapping at the trainer level, since it does not necessitate model-specific changes and does not disrupt the existing architecture for model construction and initialization. 
3) Checkpointing support for XLA FSDP is not included as part of this PR. We hope to add it soon via another PR. 
4) We have not included testing for XLA FSDP as part of this PR. We would like to add this in a future PR. 

Thanks to @ronghanghu for his assistance in the preparation of this PR. Among many other contributions, the observations that one must copy the model's forward method and replace the optimizer step are his.


<!--
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
- [ ] 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](https://github.com/huggingface/transformers/blob/main/CONTRIBUTING.md#start-contributing-pull-requests),
      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?

@sgugger @JackCaoG

<!-- 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:

- text models: @ArthurZucker and @younesbelkada
- vision models: @amyeroberts and @NielsRogge
- speech models: @sanchit-gandhi

Library:

- flax: @sanchit-gandhi
- generate: @gante
- pipelines: @Narsil
- tensorflow: @gante and @Rocketknight1
- tokenizers: @ArthurZucker
- trainer: @sgugger

Integrations:

- deepspeed: @stas00
- ray/raytune: @richardliaw, @amogkam

Documentation: @sgugger and @stevhliu

HF projects:

- accelerate: [different repo](https://github.com/huggingface/accelerate)
- datasets: [different repo](https://github.com/huggingface/datasets)
- diffusers: [different repo](https://github.com/huggingface/diffusers)
- rust tokenizers: [different repo](https://github.com/huggingface/tokenizers)

Maintained examples (not research project or legacy):

- Flax: @sanchit-gandhi
- PyTorch: @sgugger
- TensorFlow: @Rocketknight1

 -->


Comments:
_The documentation is not available anymore as the PR was closed or merged._

Thanks so much for the feedback!

> I think this can all be added via the new `fsdp_config` training arguments instead of adding four new training arguments in a class that already has too many of them (and for which users have complained a lot about). Adding an `xla` boolean inside and relying on the existing `fsdp_min_num_params`, then adding new keys for the two other arguments you want to add would work better.

The new `fsdp_config` training argument is great. I envision the strategy would be something like adding four keys to the `fsdp_config` dictionary:
- An `xla` boolean which indicates whether the user is using Fairscale FSDP or XLA FSDP
- An `xla_config` string which points to the location of a JSON file which stores the XLA FSDP configuration parameters
- The arguments `xla_fsdp_transformer_layer_class_to_wrap` and `xla_grad_ckpt` as before

Does this make sense to you? One thing I am wondering about is why `fsdp_transformer_layer_class_to_wrap` hasn't been absorbed into `fsdp_config` the same way `fsdp_min_num_params` has. It would be a bit strange to have `xla_fsdp_transformer_layer_class_to_wrap` as part of the `fsdp_config` dictionary and not `fsdp_transformer_layer_class_to_wrap`.

> The other thing to change is that `self.model` should always be the original model, not the wrapped one, which is one you need the hack on the signature. That attribute should be left as is.

Just to clarify, do you mean that we should set `self.model = model` before we modify `model`'s forward signature, or that we shouldn't set `self.model = model` at all? The former is no problem, but I think we need to do need to set `self.model = model`; among other things, without it, the program hits a segfault between training and evaluation.



Your plan for the config sounds sound, cc @pacman100 on the question around `fsdp_transformer_layer_class_to_wrap`.

Regarding the `self.model` question, it looks like the traditional FSDP also changes `self.model` to use the FSDP model, though it doesn't look like it updates the signature. Is it missing there as well?

> Your plan for the config sounds sound, cc @pacman100 on the question around `fsdp_transformer_layer_class_to_wrap`.

Excellent, thanks. If it helps, I'm happy to add `fsdp_transformer_layer_class_to_wrap` to the`fsdp_config` training argument as part of my PR, and correct a few small typos I noticed.

> Regarding the `self.model` question, it looks like the traditional FSDP also changes `self.model` to use the FSDP model, though it doesn't look like it updates the signature. Is it missing there as well?

Great question. I wasn't sure, so I asked @ronghanghu about it. He pointed out to me that we likely do not need to update the signature for either case anymore, now that XLA FSDP wrapping is occurring in `_wrap_model`. Indeed, `_remove_unused_columns` is called by `get_train_dataloader` before `_wrap_model` is called in`inner_training_loop`, and so it will use the original unwrapped model's signature. Moreover, while the functions `get_eval_dataloader` and `get_test_dataloader` call `_remove_unused_columns` after the model is wrapped, this should be ok because `self.self._signature_columns` is already set in the previous call to `get_train_dataloader`. I still need to test to verify this, but once I do, I can remove the update to the wrapped model signature. In the event that it's necessary for XLA FSDP, it will also be necessary for traditional FSDP. 



> I'm happy to add fsdp_transformer_layer_class_to_wrap to the fsdp_config training argument as part of my PR, and correct a few small typos I noticed.

Hello @AlexWertheim, adding  `fsdp_transformer_layer_class_to_wrap` to `fsdp_config` would be great. I had this change in my backlog but it would be great if this PR does that alongside the conterpart XLA changes.

I feel only boolean `xla` would be enough along with `grad_ckpt`, `compute_dtype` and `buffer_dtype` to the config. The other args `fsdp_min_num_params ` and `fsdp_transformer_layer_class_to_wrap` can be used in either cases and doc mentioning that `grad_ckpt`, `compute_dtype` and `buffer_dtype` is only applicable with `xla` should be enough along with a warning when fsdp is being used without xla with this any of these fields passed.

I feel `xla_config` as a path to file inside `fsdp_config` which itself is a path to file would be lot of load on user







> Hello @AlexWertheim, adding `fsdp_transformer_layer_class_to_wrap` to `fsdp_config` would be great. I had this change in my backlog but it would be great if this PR does that alongside the conterpart XLA changes.

Thanks @pacman100, I'd be happy to do this. One question I had was how we should handle the (unusual) situation when both the deprecated flag `fsdp_transformer_layer_cls_to_wrap` and the option within `fsdp_config` are both specified. With `fsdp_min_num_params`, there is the reasonably simple approach to take the max of the two specified numbers. In the case of `fsdp_transformer_layer_cls_to_wrap`, there are a few approaches that come to mind:
- Raise an error
- Raise a warning and use the `fsdp_config` specified string
- Use both layer classes

What do you think? 

Also, XLA FSDP supports automatic wrapping based on a set of layer names. I believe Fairscale FSDP also does as well, but currently, `fsdp_transformer_layer_cls_to_wrap` only accepts a single string as input. What do you think about expanding `fsdp_transformer_layer_cls_to_wrap` to accept a list of strings? The modifications to the existing Fairscale FSDP logic will be quite small - right now, I think it just passes a singleton set to the auto wrap policy. 

 
> I feel only boolean `xla` would be enough along with `grad_ckpt`, `compute_dtype` and `buffer_dtype` to the config. The other args `fsdp_min_num_params ` and `fsdp_transformer_layer_class_to_wrap` can be used in either cases and doc mentioning that `grad_ckpt`, `compute_dtype` and `buffer_dtype` is only applicable with `xla` should be enough along with a warning when fsdp is being used without xla with this any of these fields passed.
> 
> I feel `xla_config` as a path to file inside `fsdp_config` which itself is a path to file would be lot of load on user

With appreciation for the confusion that this might cause the user, I think it's still important to allow the user to specify XLA FSDP specific configuration parameters via an `xla_fsdp_settings` flag. The [XLA FSDP parameters](https://github.com/pytorch/xla/blob/master/torch_xla/distributed/fsdp/xla_fully_sharded_data_parallel.py#L122-L240) actually differ quite substantially from [Fairscale FSDP parameters](https://github.com/facebookresearch/fairscale/blob/v0.4.13/fairscale/nn/data_parallel/fully_sharded_data_parallel.py#L170-L291), and so I think it would be valuable for the user to be able to specify the XLA FSDP arguments that they want. Adding a separate flag for each non-redundant XLA FSDP parameter seems like it might bloat the `fsdp_config` file with too many XLA specific arguments. If they don't supply an XLA config file, then no harm done - the XLA FSDP default parameters will be used.

Compute dtype and buffer dtype are actually already Fairscale FSDP arguments, and so these could be inferred from existing flags if you'd like, though I think it's best for the user to specify them through `xla_fsdp_settings`. (Right now, the trainer sets a mixed precision policy based on whether the `bf16` or `fp16` training args are specified). 


Still happy to discuss what the final changes look like, but as a proof of concept, I've pushed some changes which (among other things) implement some of the items discussed in prior comments:
- `fsdp_transformer_layer_cls_to_wrap` is now moved inside `fsdp_config`, and is treated as a list of strings instead of a single string. In the event that the user enters `fsdp_transformer_layer_cls_to_wrap` as a string in the JSON file, the program converts this to a list. 
- If the user uses the deprecated version of `fsdp_transformer_layer_cls_to_wrap` outside `fsdp_config`, the program issues a warning and combines it (as a list of strings) with whatever `fsdp_transformer_layer_cls_to_wrap` is set to. 
- The PyTorch FSDP logic has been modified to iterate through the elements of `fsdp_transformer_layer_cls_to_wrap` instead of passing a singleton set to the auto-wrap policy.
- I noted that `fsdp_min_num_params` wasn't getting set correctly. This is because `getattr(dict, key, default)` seems to always return the default. I've replaced instances of `getattr(dict, key, default)` with `dict.get(key, default)`, including with `fsdp_min_num_params`. 
- XLA FSDP parameters are still loaded into their own separate dictionary from the JSON file specified by `xla_fsdp_settings`, though I am happy to continue discussion over whether this should be modified
- Defaults are now correctly set for the flags `xla` and `xla_fsdp_grad_ckpt`

Hello @AlexWertheim, thanks a lot for all the changes and detailed notes:

At overall level, what if the fsdp_config looked like below. i.e., xla_fsdp_settings was a dict instead of path to another json. Also, `xla_fsdp_grad_ckpt` could be pushed into `xla_fsdp_settings` as `grad_ckpt` param or something like that as it is unique to xla fsdp. This way user would need to give path to only one json `fsdp_config` while all xla params are in `xla_fsdp_settings` when `xla` is True.

```json
{
    "fsdp_transformer_layer_cls_to_wrap": "T5Block",
    "xla": true,
    "xla_fsdp_settings": {
        "buffer_dtype": "bfloat16",
        "compute_dtype": "bfloat16",
        "grad_ckpt": true,
    }
}
``` 

> At overall level, what if the fsdp_config looked like below. i.e., xla_fsdp_settings was a dict instead of path to another json. Also, `xla_fsdp_grad_ckpt` could be pushed into `xla_fsdp_settings` as `grad_ckpt` param or something like that as it is unique to xla fsdp. This way user would need to give path to only one json `fsdp_config` while all xla params are in `xla_fsdp_settings` when `xla` is True.
> 
> ```json
> {
>     "fsdp_transformer_layer_cls_to_wrap": "T5Block",
>     "xla": true,
>     "xla_fsdp_settings": {
>         "buffer_dtype": "bfloat16",
>         "compute_dtype": "bfloat16",
>         "grad_ckpt": true,
>     }
> }
> ```

@pacman100 Thanks very much for the feedback. I think the suggestion to turn `xla_fsdp_settings` into a dictionary is a great one, and a good way to resolve the issue of having a separate configuration path within `fsdp_config`. 

As far as absorbing `xla_fsdp_grad_ckpt` into `xla_fsdp_settings` goes, I agree that it is nice to have all of the XLA FSDP related configuration information in one place from a logical point of view. That being said, I do have a concern, namely that `xla_fsdp_settings` was supposed to keep track of the XLA FSDP wrapping parameters, and `xla_fsdp_grad_ckpt` is not part of the XLA FSDP wrapping arguments. I worry that the user will get confused about what this flag is for, and in particular, that users will not realize that they have to add this gradient checkpointing flag (or, for that matter, any future XLA FSDP related flags that are not themselves XLA FSDP wrapping params). What do you think?

Would also love to hear your thoughts, @ronghanghu 

@sgugger @pacman100 Thanks for taking time to review this pr. As Pytorch 2.0 branch cut just happened and will most likely be release in ~ a month, it will be great if we can enable the FSDP for HF on the master(I believe we will also need to make some other changes, but that can happen in the subsequent pr). This way we can share our benchmark using FSDP with HF more broadly when we doing the announcement/blog post for the 2.0 release.

@JackCaoG This should be merged in the next few days :-)
@AlexWertheim could you rebase your branch on main to fix the tests, and run `make style` on your branch to fix the quality jobs?

@pacman100 let me know when you are happy with all the changes.

> @AlexWertheim could you rebase your branch on main to fix the tests, and run `make style` on your branch to fix the quality jobs?

Done! Please note that all commits including and after [77e99c8](https://github.com/huggingface/transformers/pull/21406/commits/77e99c8452aa9fc7b7ca5b03c0d69b2a28e71b0e) are just concerned with rebasing and formatting/style changes. Apologies for the many commits; I had some difficulties running all of the automatic style checks, but I think I've run all of them now, and the CircleCI checks are now passing. 

Commit [e65dc0b](https://github.com/huggingface/transformers/pull/21406/commits/e65dc0b65a21417a15e30498fa3fd6447334d7ed) modified a check which would have prevented XLA FSDP wrapping when `local_rank = -1`, which is not correct for TPUs. As discussed in prior comments, commit [632318a](https://github.com/huggingface/transformers/pull/21406/commits/632318ac88e3d313843023cf0eff8bb9ac45716f) modified the argument `xla_fsdp_settings` to be a dictionary instead of a path to a JSON file; note that `xla_fsdp_grad_ckpt` was not absorbed into `xla_fsdp_settings`, for the reasons mentioned earlier. 

Please let me know if there are any additional questions or concerns upon review. Thanks!

Hi,
I was able to do parallel training flawlessly without FSDP on xlm-roberta-**base** on the 8 cores of the TPU-v3 VM (because the model and batch fit properly within each core) by following the [run_glue example](https://github.com/huggingface/transformers/tree/main/examples/pytorch#running-on-tpus) for TPUs.

Now I'm trying to get the XL version (facebook/xlm-roberta-**xl**) to work with XLA FSDP with the Trainer integration (set as "full_shard") but I'm getting these memory errors when running the second or first batch:
`0%|          | 2/939 [03:05<28:08:11, 108.10s/it]Exception in device=TPU:0: RESOURCE_EXHAUSTED: From /job:localservice/replica:0/task:0: 2 root error(s) found.
(0) RESOURCE_EXHAUSTED: Attempting to reserve 13.20G at the bottom of memory. That was not possible. There are 10.59G free, 0B reserved, and 10.59G reservable.`

From my understanding, the model was supposed to split loaded onto the TPU cores, along with whatever Zero-3 entails, but it doesn't seem to be happening.
I tried setting `per_device_train_batch_size=1`, limited tokenizer `max_length=200` and a played with a bunch of `fsdp_config` parameters but none worked.

Also I'm confused if I'm still supposed to use the xla_spawn.py script to run with FSDP or not. I tried both and got the same error though.

Additionally, trying to debug by myself I found that this [line](https://github.com/huggingface/transformers/blob/main/src/transformers/trainer.py#L617) condition is never True because the conditions on line 601 prevent it.

Maybe we could have an example on how to use this feature with the Trainer?

Thanks in advance!

@AlexWertheim This seems to be a device side oom, any advise you can give?