Default `synced_gpus` to `True` when using `FullyShardedDataParallel`
Author: ringohoffman
Pull request: yes



# What does this PR do?

<!-- Remove if not applicable -->

Fixes #30228

Related:

* https://github.com/pytorch/pytorch/issues/100069
* https://github.com/pytorch/pytorch/issues/123962

Similar to DeepSpeed ZeRO Stage 3, when using FSDP with multiple GPUs and differently sized data per rank, the ranks reach different synchronization points at the same time, leading to deadlock.

To avoid this, we can automatically set `synced_gpus` to `True` if we detect that a `PreTrainedModel` is being managed by FSDP using `_is_fsdp_managed_module`, which was added in 2.0.0 for `torch.compile`: https://github.com/pytorch/pytorch/blob/v2.0.0/torch/distributed/fsdp/_dynamo_utils.py

To facilitate this, I created a module called `transformers.integrations.fsdp` containing the function `is_fsdp_managed_module` which returns `True` if a `Module` has `_is_fsdp_managed_module` set to `True` on it or if the `Module` itself is a `FullyShardedDataParallel` instance.

Here is the script I used to test my fix:

```bash
OMP_NUM_THREADS=2 \
TOKENIZERS_PARALLELISM=false \
CUDA_VISIBLE_DEVICES=6,7 \
torchrun \
    --rdzv-backend=c10d \
    --rdzv-endpoint=localhost:0 \
    --nnodes=1 \
    --nproc-per-node=2 \
    fsdp_generate.py
```

**fsdp_generate.py**
```python
import functools

import torch
import torch.distributed
import torch.distributed.fsdp
import torch.distributed.fsdp.wrap
import transformers
import transformers.models.gpt_neo.modeling_gpt_neo


def main() -> None:
    torch.distributed.init_process_group(world_size=2)
    device = torch.device(torch.distributed.get_rank())
    torch.cuda.set_device(device)

    pretrained_model_name_or_path = "EleutherAI/gpt-neo-125m"
    model = transformers.AutoModelForCausalLM.from_pretrained(
        pretrained_model_name_or_path,
        device_map=device,
        attn_implementation="flash_attention_2",  # I'm using flash-attn @ https://github.com/Dao-AILab/flash-attention/releases/download/v2.6.3/flash_attn-2.6.3+cu123torch2.4cxx11abiFALSE-cp310-cp310-linux_x86_64.whl
        torch_dtype=torch.bfloat16,
    )
    assert isinstance(model, transformers.models.gpt_neo.modeling_gpt_neo.GPTNeoForCausalLM)

    tokenizer = transformers.AutoTokenizer.from_pretrained(
        pretrained_model_name_or_path,
    )
    tokenizer.pad_token_id = tokenizer.eos_token_id

    fsdp_model = torch.distributed.fsdp.FullyShardedDataParallel(
        model,
        auto_wrap_policy=functools.partial(
            torch.distributed.fsdp.wrap.transformer_auto_wrap_policy,
            transformer_layer_cls={
                transformers.models.gpt_neo.modeling_gpt_neo.GPTNeoBlock
            },
        ),
        limit_all_gathers=True,
        use_orig_params=True,  # required to overcome the error "The tensor has a non-zero number of elements, but its data is not allocated yet" ... PreTrainedModel.generate is probably using some torch.compile-wrapped function
    )

    data_by_rank = {  # differently sized causes FSDP to hang
        0: "Hello world!",
        1: "The quick brown fox jumps over the lazy dog."
    }

    batch = tokenizer(
        data_by_rank[torch.distributed.get_rank()],
        return_tensors="pt",
        return_attention_mask=True,
    ).to(device)

    with torch.distributed.fsdp.FullyShardedDataParallel.summon_full_params(fsdp_model):  # required to overcome to the error "'weight' must be 2-D"
        generated_text = fsdp_model.module.generate(
            input_ids=batch["input_ids"],
            attention_mask=batch["attention_mask"],
            max_length=20,
            # synced_gpus=True,  # currently, True is required to use differently sized data with FSDP + generate (current default is False)
        )

    torch.distributed.barrier()
    torch.distributed.destroy_process_group()


if __name__ == "__main__":
    main()
```

## 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?
- [x] 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.
- [x] 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?

Anyone in the community is free to review the PR once the tests have passed. Feel free to tag
members/contributors who may be interested in your PR.

@gante
@ArthurZucker

<!-- 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
- vision models: @amyeroberts, @qubvel
- speech models: @ylacombe, @eustlb
- graph models: @clefourrier

Library:

- flax: @sanchit-gandhi
- generate: @zucchini-nlp (visual-language models) or @gante (all others)
- pipelines: @Rocketknight1
- tensorflow: @gante and @Rocketknight1
- tokenizers: @ArthurZucker
- trainer: @muellerzr and @SunMarc
- chat templates: @Rocketknight1

Integrations:

- deepspeed: HF Trainer/Accelerate: @muellerzr
- ray/raytune: @richardliaw, @amogkam
- Big Model Inference: @SunMarc
- quantization (bitsandbytes, autogpt): @SunMarc

Documentation: @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: See Models above and tag the person corresponding to the modality of the example.
- TensorFlow: @Rocketknight1

 -->


Comments:
I should also mention that my script fails when using `attn_implementation="sdpa"` with an off-by-one error. I don't see this error when using `attn_implementation="flash_attention_2"`...

```
[rank1]: Traceback (most recent call last):
[rank1]:   File "/home/matthew/transformers/fsdp_generate.py", line 64, in <module>
[rank1]:     main()
[rank1]:   File "/home/matthew/transformers/fsdp_generate.py", line 52, in main
[rank1]:     generated_text = fsdp_model.module.generate(
[rank1]:   File "/home/matthew/.conda/envs/transformers310/lib/python3.10/site-packages/torch/utils/_contextlib.py", line 116, in decorate_context
[rank1]:     return func(*args, **kwargs)
[rank1]:   File "/home/matthew/transformers/src/transformers/generation/utils.py", line 2048, in generate
[rank1]:     result = self._sample(
[rank1]:   File "/home/matthew/transformers/src/transformers/generation/utils.py", line 3001, in _sample
[rank1]:     outputs = self(**model_inputs, return_dict=True)
[rank1]:   File "/home/matthew/.conda/envs/transformers310/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1553, in _wrapped_call_impl
[rank1]:     return self._call_impl(*args, **kwargs)
[rank1]:   File "/home/matthew/.conda/envs/transformers310/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1562, in _call_impl
[rank1]:     return forward_call(*args, **kwargs)
[rank1]:   File "/home/matthew/transformers/src/transformers/models/gpt_neo/modeling_gpt_neo.py", line 1038, in forward
[rank1]:     transformer_outputs = self.transformer(
[rank1]:   File "/home/matthew/.conda/envs/transformers310/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1553, in _wrapped_call_impl
[rank1]:     return self._call_impl(*args, **kwargs)
[rank1]:   File "/home/matthew/.conda/envs/transformers310/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1562, in _call_impl
[rank1]:     return forward_call(*args, **kwargs)
[rank1]:   File "/home/matthew/transformers/src/transformers/models/gpt_neo/modeling_gpt_neo.py", line 801, in forward
[rank1]:     outputs = block(
[rank1]:   File "/home/matthew/.conda/envs/transformers310/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1553, in _wrapped_call_impl
[rank1]:     return self._call_impl(*args, **kwargs)
[rank1]:   File "/home/matthew/.conda/envs/transformers310/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1562, in _call_impl
[rank1]:     return forward_call(*args, **kwargs)
[rank1]:   File "/home/matthew/.conda/envs/transformers310/lib/python3.10/site-packages/torch/distributed/fsdp/fully_sharded_data_parallel.py", line 863, in forward
[rank1]:     output = self._fsdp_wrapped_module(*args, **kwargs)
[rank1]:   File "/home/matthew/.conda/envs/transformers310/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1553, in _wrapped_call_impl
[rank1]:     return self._call_impl(*args, **kwargs)
[rank1]:   File "/home/matthew/.conda/envs/transformers310/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1562, in _call_impl
[rank1]:     return forward_call(*args, **kwargs)
[rank1]:   File "/home/matthew/transformers/src/transformers/models/gpt_neo/modeling_gpt_neo.py", line 512, in forward
[rank1]:     attn_outputs = self.attn(
[rank1]:   File "/home/matthew/.conda/envs/transformers310/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1553, in _wrapped_call_impl
[rank1]:     return self._call_impl(*args, **kwargs)
[rank1]:   File "/home/matthew/.conda/envs/transformers310/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1562, in _call_impl
[rank1]:     return forward_call(*args, **kwargs)
[rank1]:   File "/home/matthew/transformers/src/transformers/models/gpt_neo/modeling_gpt_neo.py", line 462, in forward
[rank1]:     return self.attention(
[rank1]:   File "/home/matthew/.conda/envs/transformers310/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1553, in _wrapped_call_impl
[rank1]:     return self._call_impl(*args, **kwargs)
[rank1]:   File "/home/matthew/.conda/envs/transformers310/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1562, in _call_impl
[rank1]:     return forward_call(*args, **kwargs)
[rank1]:   File "/home/matthew/transformers/src/transformers/models/gpt_neo/modeling_gpt_neo.py", line 314, in forward
[rank1]:     attn_output, attn_weights = self._attn(query, key, value, attention_mask, head_mask)
[rank1]:   File "/home/matthew/transformers/src/transformers/models/gpt_neo/modeling_gpt_neo.py", line 278, in _attn
[rank1]:     attn_weights = attn_weights + causal_mask
[rank1]: RuntimeError: The size of tensor a (21) must match the size of tensor b (20) at non-singleton dimension 3
```

The CI failures seem to be from flaky tests. This should be good for review!

Hey @gante and @ArthurZucker, just following up. Do you think you will be able to review this?

@SunMarc and @muellerzr could you take a look here please?

Okay, I gave it a stab based on these files + the example in the description. I added a test for FSDP and FSDP2.

https://github.com/huggingface/transformers/blob/main/tests/trainer/test_trainer_distributed.py
https://github.com/huggingface/transformers/blob/main/tests/generation/test_utils.py


Hi, I found that maybe have a bug of trainer fsdp init, when I use pytest to test the test code `tests/trainer/test_trainer_fsdp.py`, I got error.
![image](https://github.com/user-attachments/assets/f3c37d01-ff45-481d-8676-28bee0726670)
 may u git me some advice how can i avoid it? @ringohoffman 
 
 my packages version are
 torch                    2.4.1
 accelerate            0.34.2
 transformers        4.46.0.dev0
 thx a lot

@YeLuoSuiYou Seems like it was due to the logic added in this PR https://github.com/huggingface/transformers/pull/34032

> @YeLuoSuiYou Seems like it was due to the logic added in this PR #34032@YeLuoSuiYou似乎是由于此 PR 中添加的逻辑所致 #34032

thx a lot, maybe I should wait the normal version to fix it? btw, thx much for the function provided by this PR!

> @YeLuoSuiYou Seems like it was due to the logic added in this PR #34032@YeLuoSuiYou似乎是由于此 PR 中添加的逻辑所致 #34032

I commet this function, the test work well done, thx!