Llama et al. / FSDP : Fix breaking change in 4.40 for FSDP
Author: younesbelkada
Pull request: yes

# What does this PR do?

Fixes: https://github.com/huggingface/transformers/issues/30523

<details><summary>Click to see the snippet (make sure to run `accelerate config` and select FSDP options before hand and run `accelerate launch script.py`) </summary>

```python
from functools import partial
import torch
from transformers import AutoModelForCausalLM, AutoTokenizer
from accelerate import Accelerator

# verify we have FSDP activation support ready by importing:
from torch.distributed.algorithms._checkpoint.checkpoint_wrapper import (
    checkpoint_wrapper,
    CheckpointImpl,
    apply_activation_checkpointing,
)

from transformers.models.llama.modeling_llama import LlamaDecoderLayer

model_id = "HuggingFaceM4/tiny-random-Llama3ForCausalLM"

model = AutoModelForCausalLM.from_pretrained(model_id)

model.train()
model.gradient_checkpointing_enable()

accelerator = Accelerator()
model = accelerator.prepare(model)

check_fn = lambda submodule: isinstance(submodule, LlamaDecoderLayer)

non_reentrant_wrapper = partial(
    checkpoint_wrapper,
    offload_to_cpu=False,
    checkpoint_impl=CheckpointImpl.NO_REENTRANT,
)

apply_activation_checkpointing(
    model, checkpoint_wrapper_fn=non_reentrant_wrapper, check_fn=check_fn
)

print(model)

rand_input = torch.LongTensor([[0, 1, 0, 1]]).to(0)

model(rand_input)
```
</details>

https://github.com/huggingface/transformers/pull/30743 introduced a breaking change for users that use Llama-based models + FSDP + activation checkpointing with FSDP. 

Before https://github.com/huggingface/transformers/pull/30743 - we were able to pass arbitrary kwargs within Llama modules that were silently ignored. When doing FSDP + activation checkpointing, the target gradient checkpointing classes are wrapped in a new class, and additional kwargs are passed along that class forward pass 

The script above used to work for transformers <= 4.40.0 and does not work anymore due to https://github.com/huggingface/transformers/pull/30743 , re-intoducing kwargs in the foward method signature fixes the bug

cc @amyeroberts 

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

Thanks ! 

> Would it make sense to add a test to make sure we don't accidentally break this again?

Yes, I'll add a test in this PR to test this behavior and catch bugs in the future! 

> Having **kwargs in the forward method isn't standard amongst transformers models. Is there something special about these models which need this for FSDP? If not, should we be adding to other models? 

Yes agreed, I think we should add it to all 'most-used' models. FSDP is useful for large models, so I would say we should add it for LLMs (llama, gemma, mistral, mixtral, gpt-neo, etc.) to make things consistent. Happy to do that within this PR !

> Is there an alternative to using this injection? Relying on kwargs being passed isn't ideal

I am not sure, this seems to be something internal to FSDP + CPU offloading, I don't think we can find a workaround to this :/ for me since it used to work before, it should be still working for future transformers versions to ensure BC. What do you think?

> Yes, I'll add a test in this PR to test this behavior and catch bugs in the future!
> Yes agreed, I think we should add it to all 'most-used' models. FSDP is useful for large models, so I would say we should add it for LLMs (llama, gemma, mistral, mixtral, gpt-neo, etc.) to make things consistent. Happy to do that within this PR !

Awesome - thank you! 

> I am not sure, this seems to be something internal to FSDP + CPU offloading, I don't think we can find a workaround to this :/ for me since it used to work before, it should be still working for future transformers versions to ensure BC. What do you think?

Make sense - let's leave as-is :)

@younesbelkada I'm really sorry I missed the rerequest for review. ~I don't have permissions to make changes, so copied the branch here: #31638 and sync with `main`.~ I couldn't push working locally but could change through the editor