[FSDP] redundant additional allgather during backward when using FSDP FULL_SHARD with gradient checkpointing
Author: yundai424
Pull request: no
Labels: name: PyTorch FSDP

### System Info

- `transformers` version: 4.39.3
- Platform: Linux-5.15.138.1-4.cm2-x86_64-with-glibc2.35
- Python version: 3.10.2
- Huggingface_hub version: 0.22.2
- Safetensors version: 0.4.2
- Accelerate version: 0.29.3
- Accelerate config:    not found
- PyTorch version (GPU?): 2.1.2 (True)
- Tensorflow version (GPU?): not installed (NA)
- Flax version (CPU?/GPU?/TPU?): not installed (NA)
- Jax version: not installed
- JaxLib version: not installed
- Using GPU in script?: yes
- Using distributed or parallel set-up in script?: yes

### Who can help?

@pacman100 @muellerzr 

### Information

- [ ] The official example scripts
- [X] My own modified scripts

### Tasks

- [ ] An officially supported task in the `examples` folder (such as GLUE/SQuAD, ...)
- [x] My own task or dataset (give details below)

### Reproduction

**For reproduction script, please refer to this gist https://gist.github.com/yundai424/1b6f0fa9f23796033aaf585ae9e1d4b6**

Input alpaca data can be downloaded via https://huggingface.co/datasets/QingyiSi/Alpaca-CoT/blob/1d971b5591d198b810ccca6feaed420df9f39d28/alpaca/alpaca_data.json


If run the script as-is (which is essentially very basic finetuning, enabling fsdp and grad checkpointing via `--fsdp` and `--gradient_checkpointing` flags to `TrainingArguments`), it'll use the existing way of enabling gradient checkpointing, which is directly through `torch.utils.checkpoint`. By looking at the profiled trace view, it's obvious that there are 2 allgather ops per transformer block backward, and if tracing the call stack a bit you'll find the 1st one corresponds to the allgather for gradient checkpointing forward, and the 2nd one corresponds to the actual `pre_backward` hook for FSDP. This introduces unnecessary redundancy because the time gap between them is very short (i.e. the forward compute in gradient checkpointing) and is resulting in 2 -> 3 collectives per backprop (hence 3 -> 4 collectives per training step):

![Screenshot 2024-04-22 at 1 02 45 PM](https://github.com/huggingface/transformers/assets/43726198/a6b3c333-a88f-4302-82b7-08501d4815fd)


If uncomment [line 67-75](https://gist.github.com/yundai424/1b6f0fa9f23796033aaf585ae9e1d4b6#file-training-py-L67-L75) in the training script, it uses a different torch gradient checkpointing implementation `torch.distributed.algorithm._checkpoint` as also suggested by Meta's official [llama recipe](https://github.com/meta-llama/llama-recipes/blob/main/src/llama_recipes/policies/activation_checkpointing_functions.py). It's another layer wrapping on top of `torch.utils.checkpoint`. This will address the redundant allgather as described above, resulting in trace that looks like follows:

![Screenshot 2024-04-22 at 1 03 21 PM](https://github.com/huggingface/transformers/assets/43726198/7e391fb3-91c5-4e11-996c-252d7b7953e7)


Note that the relative order of wrapping with `FSDP` module and calling `apply_activation_checkpointing` doesn't matter here. The above code example call `apply_activation_checkpointing` before wrapping with FSDP module, but it works in the other way too.


### Expected behavior

When using FSDP with gradient checkpointing, should instead use `torch.distributed.algorithm._checkpoint.apply_activation_checkpointing`


Will be happy to draft a PR.

Comments:
Hi @yundai424, thanks for the detailed report ! If you can submit a PR to fix this issue, that would be awesome 🔥 cc @pacman100 

Hi @yundai424, hope you are doing well. Do you still want to draft the PR ? Otherwise, I will try to fix the issue based on your report. Any advice is appreciated. 

Hi @SunMarc thanks for checking! I later found there's already another flag ([this](https://github.com/huggingface/transformers/blob/main/src/transformers/training_args.py#L530)) under `fsdp_config` that uses the torch distributed ckpt wrapper which addressed the issue. But it would be nice to just merge them 

Hi @yundai424, thanks for answering and i'm glad you manage to fix the issue ! What do you mean by merging them ? I guess you are talking about `gradient_checkpointing` and `activation_checkpointing`. 

@SunMarc yes exactly!

I see [here](https://github.com/huggingface/transformers/blob/1e79eade41c346d1d8c2ad79c9e32e3186f20b54/src/transformers/trainer.py#L4619) that we can't set both args to `True` and that we advise the user to use FSDP `activation_checkpointing`.

 I think that we should not merge the two args for BC but maybe add a warning when `gradient_checkpointing` is set to `True` when using fsdp advising users to use `activation_checkpointing` arg instead since it uses a different torch gradient checkpointing implementation which is better suited for fsdp ! Would you like to open the PR ? However, I can do it ! 

@SunMarc that makes sense, i can open the PR to add the warning : ) thanks!

With this change, how can we still set `use_reetrant=Ture`  using `gradient_checkpointing_kwargs`? If `--gradient_checkpointing=False` and FSDP config has `{activation_checkpointing: true}`, by default the huggingface trainer doesn't load gradient_checkpointing_kwargs.

This is probably something that could be added to the fsdp plugin or config wdyt @muellerzr ? This would affect this specific [code](https://github.com/huggingface/accelerate/blob/bf4572b6ce0a534a9d73537485a0edf1d68144b8/src/accelerate/accelerator.py#L1510). If you are willing to contribute, feel free to open a PR ! 

> With this change, how can we still set `use_reetrant=Ture` using `gradient_checkpointing_kwargs`? If `--gradient_checkpointing=False` and FSDP config has `{activation_checkpointing: true}`, by default the huggingface trainer doesn't load gradient_checkpointing_kwargs.

I have the same issue here. 
Will get following error: 

```
Traceback (most recent call last):
  File "/miniforge/envs/iris/lib/python3.10/site-packages/torch/utils/checkpoint.py", line 1416, in _checkpoint_without_reentrant_generator
    with _checkpoint_hook(new_frame), forward_context:
  File "/miniforge/envs/iris/lib/python3.10/site-packages/torch/autograd/graph.py", line 284, in __exit__
Exception ignored in: <generator object _checkpoint_without_reentrant_generator at 0x7f40341acc80>
```