'LlamaModel' object has no attribute '_gradient_checkpointing_func'.
Author: foreverpiano
Pull request: no

### System Info

[rank3]:   File "/home/dhl/miniconda3/envs/light/lib/python3.10/site-packages/transformers/models/llama/modeling_llama.py", line 1060, in forward
[rank3]:     layer_outputs = self._gradient_checkpointing_func(
[rank3]:   File "/home/dhl/miniconda3/envs/light/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1709, in __getattr__
[rank3]:     raise AttributeError(f"'{type(self).__name__}' object has no attribute '{name}'")
[rank3]: AttributeError: 'LlamaModel' object has no attribute '_gradient_checkpointing_func'. Did you mean: 'gradient_checkpointing'?


https://github.com/huggingface/transformers/blob/main/src/transformers/models/llama/modeling_llama.py#L1007


### Who can help?


 @ArthurZucker and @younesbelkada @Narsil


Comments:
```
accelerate==0.29.3
certifi==2024.2.2
charset-normalizer==3.3.2
einops==0.8.0
exceptiongroup==1.2.1
filelock==3.13.4
flash-attn==2.5.8
fsspec==2024.3.1
huggingface-hub==0.22.2
idna==3.7
iniconfig==2.0.0
Jinja2==3.1.3
MarkupSafe==2.1.5
mpmath==1.3.0
networkx==3.3
ninja==1.11.1.1
numpy==1.26.4
nvidia-cublas-cu12==12.1.3.1
nvidia-cuda-cupti-cu12==12.1.105
nvidia-cuda-nvrtc-cu12==12.1.105
nvidia-cuda-runtime-cu12==12.1.105
nvidia-cudnn-cu12==8.9.2.26
nvidia-cufft-cu12==11.0.2.54
nvidia-curand-cu12==10.3.2.106
nvidia-cusolver-cu12==11.4.5.107
nvidia-cusparse-cu12==12.1.0.106
nvidia-nccl-cu12==2.20.5
nvidia-nvjitlink-cu12==12.4.127
nvidia-nvtx-cu12==12.1.105
packaging==24.0
pillow==10.3.0
pip==23.3.1
pluggy==1.5.0
psutil==5.9.8
pytest==8.2.0
PyYAML==6.0.1
regex==2024.4.16
requests==2.31.0
safetensors==0.4.3
setuptools==68.2.2
sympy==1.12
tokenizers==0.15.2
tomli==2.0.1
torch==2.3.0
torchaudio==2.3.0
torchvision==0.18.0
tqdm==4.66.2
transformers==4.37.2
triton==2.3.0
typing_extensions==4.11.0
urllib3==2.2.1
wheel==0.41.2

```

HI @foreverpiano 
Can you share more details about the script that you are running?

when I try to use Llamamodel
```python
        data = {"input_ids": torch.randint(0, 1000, (1, length,), device="cuda"),
                "labels": torch.randint(0, 1000, (1, length,), device="cuda"), 
                "attention_mask": torch.ones(1, length, device="cuda")}
        for i in tqdm(range(20)):
            if i > 10:
                time_s = time.time()
            outputs = model(**data, use_cache=True) ----> fail
```

log
```
[rank3]:   File "/home/dhl/LongChat-dev/longchat/train/fine_tune/train_no_trainer.py", line 166, in train
[rank3]:     outputs = model(**data, use_cache=True)
[rank3]:   File "/home/dhl/miniconda3/envs/light/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1532, in _wrapped_call_impl
[rank3]:     return self._call_impl(*args, **kwargs)
[rank3]:   File "/home/dhl/miniconda3/envs/light/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1541, in _call_impl
[rank3]:     return forward_call(*args, **kwargs)
[rank3]:   File "/home/dhl/miniconda3/envs/light/lib/python3.10/site-packages/torch/distributed/fsdp/fully_sharded_data_parallel.py", line 857, in forward
[rank3]:     output = self._fsdp_wrapped_module(*args, **kwargs)
[rank3]:   File "/home/dhl/miniconda3/envs/light/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1532, in _wrapped_call_impl
[rank3]:     return self._call_impl(*args, **kwargs)
[rank3]:   File "/home/dhl/miniconda3/envs/light/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1541, in _call_impl
[rank3]:     return forward_call(*args, **kwargs)
[rank3]:   File "/home/dhl/miniconda3/envs/light/lib/python3.10/site-packages/transformers/models/llama/modeling_llama.py", line 1183, in forward
[rank3]:     outputs = self.model(
[rank3]:   File "/home/dhl/miniconda3/envs/light/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1532, in _wrapped_call_impl
[rank3]:     return self._call_impl(*args, **kwargs)
[rank3]:   File "/home/dhl/miniconda3/envs/light/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1541, in _call_impl
[rank3]:     return forward_call(*args, **kwargs)
[rank3]:   File "/home/dhl/miniconda3/envs/light/lib/python3.10/site-packages/transformers/models/llama/modeling_llama.py", line 1060, in forward
[rank3]:     layer_outputs = self._gradient_checkpointing_func(
[rank3]:   File "/home/dhl/miniconda3/envs/light/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1709, in __getattr__
[rank3]:     raise AttributeError(f"'{type(self).__name__}' object has no attribute '{name}'")
[rank3]: AttributeError: 'LlamaModel' object has no attribute '_gradient_checkpointing_func'. Did you mean: 'gradient_checkpointing'?
```
for the model definition
```python
    config = transformers.AutoConfig.from_pretrained(model_args.model_name_or_path)
model = transformers.AutoModelForCausalLM.from_config(config).to(device="cuda", dtype=torch.float16)
    model.model.gradient_checkpointing = True
    model_parameters = filter(lambda p: p.requires_grad, model.parameters())


    model = FSDP(
        model,
        auto_wrap_policy=auto_wrap_policy,
        sharding_strategy=ShardingStrategy.FULL_SHARD
        #sharding_strategy=ShardingStrategy.SHARD_GRAD_OP
        )
```

@younesbelkada 

Thanks @foreverpiano 
Can you try:

```diff
    config = transformers.AutoConfig.from_pretrained(model_args.model_name_or_path)
model = transformers.AutoModelForCausalLM.from_config(config).to(device="cuda", dtype=torch.float16)
-   model.model.gradient_checkpointing = True
+   model.gradient_checkpointing_enable()
    model_parameters = filter(lambda p: p.requires_grad, model.parameters())


    model = FSDP(
        model,
        auto_wrap_policy=auto_wrap_policy,
        sharding_strategy=ShardingStrategy.FULL_SHARD
        #sharding_strategy=ShardingStrategy.SHARD_GRAD_OP
        )
```

@younesbelkada  Thanks for your timely reply. It works for me.

I got the same error with transformers 4.39.3 and facebook/opt-350m, after explicitly calling `gradient_checkpointing_enable `.
![image](https://github.com/huggingface/transformers/assets/87317405/f044d22c-dd9b-41ba-9a94-4bdbf4bb231a)
![image](https://github.com/huggingface/transformers/assets/87317405/193ff0ec-790f-47db-bb0c-6bd223e10f94)


Hi @Edenzzzz, 

`gradient_checkpointing_enable` is meant to be called on the model outside of other methods. Which version of transformers are you running on? The code example does not match the [forward logic of the OPTDecoder](https://github.com/huggingface/transformers/blob/e65502951593a76844e872fee9c56b805598538a/src/transformers/models/opt/modeling_opt.py#L874)

![image](https://github.com/huggingface/transformers/assets/91032945/1bbb199e-49a1-49ed-b690-824d26c63fea)
@amyeroberts I came across the same issue recently, even though everything was normal a few days ago. The method below worked, but the OOM (Out of Memory) issues have appeared. I'm not sure if it is due to a different version of Transformers.

```py
In [1]: from transformers import CLIPVisionModel
   ...: vision_tower = CLIPVisionModel.from_pretrained("openai/clip-vit-large-patch14-336")
   ...: vision_tower.gradient_checkpointing_enable({"use_reentrant": True})
   ...: vision_tower.vision_model.encoder._gradient_checkpointing_func
Out[1]: functools.partial(<function checkpoint at 0x15accd240>, use_reentrant=True)
```

_Originally posted by @amyeroberts in https://github.com/huggingface/transformers/issues/30264#issuecomment-2061082162_
     

> Hi @Edenzzzz,
> 
> `gradient_checkpointing_enable` is meant to be called on the model outside of other methods. Which version of transformers are you running on? The code example does not match the [forward logic of the OPTDecoder](https://github.com/huggingface/transformers/blob/e65502951593a76844e872fee9c56b805598538a/src/transformers/models/opt/modeling_opt.py#L874)

I did call it outside the forward logic. Here I'm modifying the forward function for personal uses, with `gradient_checkpointing_enable` added to show that it's ensured to be called.

I have a funnier example 
```
    # Build OPT model
    config = AutoConfig.from_pretrained(args.model_name_or_path)
    model = OPTForCausalLM.from_pretrained(args.model_name_or_path, config=config)

    # Enable gradient checkpointing
    model.gradient_checkpointing_enable()
    print("model._gradient_checkpointing_func:", model._gradient_checkpointing_func)
```
<img width="966" alt="image" src="https://github.com/huggingface/transformers/assets/87317405/a6f3dfa7-0d19-4940-adcf-86a168c54703">


> ![image](https://private-user-images.githubusercontent.com/91032945/344862397-1bbb199e-49a1-49ed-b690-824d26c63fea.png?jwt=eyJhbGciOiJIUzI1NiIsInR5cCI6IkpXVCJ9.eyJpc3MiOiJnaXRodWIuY29tIiwiYXVkIjoicmF3LmdpdGh1YnVzZXJjb250ZW50LmNvbSIsImtleSI6ImtleTUiLCJleHAiOjE3MTk5NTczNTQsIm5iZiI6MTcxOTk1NzA1NCwicGF0aCI6Ii85MTAzMjk0NS8zNDQ4NjIzOTctMWJiYjE5OWUtNDlhMS00OWVkLWI2OTAtODI0ZDI2YzYzZmVhLnBuZz9YLUFtei1BbGdvcml0aG09QVdTNC1ITUFDLVNIQTI1NiZYLUFtei1DcmVkZW50aWFsPUFLSUFWQ09EWUxTQTUzUFFLNFpBJTJGMjAyNDA3MDIlMkZ1cy1lYXN0LTElMkZzMyUyRmF3czRfcmVxdWVzdCZYLUFtei1EYXRlPTIwMjQwNzAyVDIxNTA1NFomWC1BbXotRXhwaXJlcz0zMDAmWC1BbXotU2lnbmF0dXJlPTA5MWU3NzlmZDRiZmY2M2ZjNmZkMmY1NGVjYjVjYzNiZDkxNmIwNzkzZmI5ODc5MWZhYjA5NDM2NjExNjViN2QmWC1BbXotU2lnbmVkSGVhZGVycz1ob3N0JmFjdG9yX2lkPTAma2V5X2lkPTAmcmVwb19pZD0wIn0.J4AIQOWc9JYYr8sy-mJnk18ouyBUXg5Ul_4oKS5nlAw) @amyeroberts I came across the same issue recently, even though everything was normal a few days ago. The method below worked, but the OOM (Out of Memory) issues have appeared. I'm not sure if it is due to a different version of Transformers.
> 
> ```python
> In [1]: from transformers import CLIPVisionModel
>    ...: vision_tower = CLIPVisionModel.from_pretrained("openai/clip-vit-large-patch14-336")
>    ...: vision_tower.gradient_checkpointing_enable({"use_reentrant": True})
>    ...: vision_tower.vision_model.encoder._gradient_checkpointing_func
> Out[1]: functools.partial(<function checkpoint at 0x15accd240>, use_reentrant=True)
> ```
> 
> _Originally posted by @amyeroberts in [#30264 (comment)](https://github.com/huggingface/transformers/issues/30264#issuecomment-2061082162)_

@nzomi Thanks for sharing this example. Is it related to the screenshot above? I'm able to run the code snippet provided without issue, but the example is using the vision model and the screenshot is referring to the CLIPEncoder 

> > Hi @Edenzzzz,
> > `gradient_checkpointing_enable` is meant to be called on the model outside of other methods. Which version of transformers are you running on? The code example does not match the [forward logic of the OPTDecoder](https://github.com/huggingface/transformers/blob/e65502951593a76844e872fee9c56b805598538a/src/transformers/models/opt/modeling_opt.py#L874)
> 
> I did call it outside the forward logic. Here I'm modifying the forward function for personal uses, with `gradient_checkpointing_enable` added to show that it's ensured to be called.

@Edenzzzz in the screenshot provided, it shows `gradient_checkpointing_enable` being called in a modified version of the forward pass, which is in the forward logic. 


> I have a funnier example
> 
> ```
>     # Build OPT model
>     config = AutoConfig.from_pretrained(args.model_name_or_path)
>     model = OPTForCausalLM.from_pretrained(args.model_name_or_path, config=config)
> 
>     # Enable gradient checkpointing
>     model.gradient_checkpointing_enable()
>     print("model._gradient_checkpointing_func:", model._gradient_checkpointing_func)
> ```
> 
> <img alt="image" width="966" src="https://private-user-images.githubusercontent.com/87317405/344917834-a6f3dfa7-0d19-4940-adcf-86a168c54703.png?jwt=eyJhbGciOiJIUzI1NiIsInR5cCI6IkpXVCJ9.eyJpc3MiOiJnaXRodWIuY29tIiwiYXVkIjoicmF3LmdpdGh1YnVzZXJjb250ZW50LmNvbSIsImtleSI6ImtleTUiLCJleHAiOjE3MTk5NTczNTQsIm5iZiI6MTcxOTk1NzA1NCwicGF0aCI6Ii84NzMxNzQwNS8zNDQ5MTc4MzQtYTZmM2RmYTctMGQxOS00OTQwLWFkY2YtODZhMTY4YzU0NzAzLnBuZz9YLUFtei1BbGdvcml0aG09QVdTNC1ITUFDLVNIQTI1NiZYLUFtei1DcmVkZW50aWFsPUFLSUFWQ09EWUxTQTUzUFFLNFpBJTJGMjAyNDA3MDIlMkZ1cy1lYXN0LTElMkZzMyUyRmF3czRfcmVxdWVzdCZYLUFtei1EYXRlPTIwMjQwNzAyVDIxNTA1NFomWC1BbXotRXhwaXJlcz0zMDAmWC1BbXotU2lnbmF0dXJlPWI0NTMzZmNhNmVkM2M5ZGI0YWUzYjc4YTViY2IxYmJhNzI5MmY3NGNjYjg2NmYzMTU2MGNkYzg2M2U3NGI0MWYmWC1BbXotU2lnbmVkSGVhZGVycz1ob3N0JmFjdG9yX2lkPTAma2V5X2lkPTAmcmVwb19pZD0wIn0.uTG3eXbKX0eMLgzP21KaW-JhMzKYHizA33XNSRls-_g">

It's the decoder which will have the function set:

```py
In [1]: from transformers import AutoModelForCausalLM

In [2]: checkpoint = "facebook/opt-125m"

In [3]: model = AutoModelForCausalLM.from_pretrained(checkpoint)

In [4]: model.gradient_checkpointing_enable()

In [5]: model.model.decoder._gradient_checkpointing_func
Out[5]: functools.partial(<function checkpoint at 0x146f676d0>, use_reentrant=True)
```