llama-2-7b-chat-hf __call__() method throws memory error
Author: kaoutaar
Pull request: no

### System Info

i am using "meta-llama/Llama-2-7b-chat-hf" model loaded in 4bits for inference in both kaggle (GPU T4x2) and colab (GPU T4x1), the model works fine with generate method: 

`llamatokenizer.batch_decode(llamamodel.generate(inputs=s.input_ids, max_new_tokens=60))
`
OUT:

`["
 \n
 
 
\nYou are a helpful,.....`

 but it immediatly crashes when i try to use __call__ : 

`llamamodel(**s)`

OUT:

```
OutOfMemoryError                          Traceback (most recent call last)
[<ipython-input-45-5cc8812cd90c>](https://localhost:8080/#) in <cell line: 1>()
----> 1 llamamodel(**s)

12 frames
[/usr/local/lib/python3.10/dist-packages/torch/nn/functional.py](https://localhost:8080/#) in softmax(input, dim, _stacklevel, dtype)
   1843         ret = input.softmax(dim)
   1844     else:
-> 1845         ret = input.softmax(dim, dtype=dtype)
   1846     return ret
   1847 

OutOfMemoryError: CUDA out of memory. Tried to allocate 154.00 MiB (GPU 0; 14.75 GiB total capacity; 13.07 GiB already allocated; 54.81 MiB free; 13.65 GiB reserved in total by PyTorch) If reserved memory is >> allocated memory try setting max_split_size_mb to avoid fragmentation.  See documentation for Memory Management and PYTORCH_CUDA_ALLOC_CONF
````

isn't generate() method using __call__() under the hood, why is it crashing then?

### Who can help?

@gante @younesbelkada @ArthurZucker 

### Information

- [ ] The official example scripts
- [ ] My own modified scripts

### Tasks

- [ ] An officially supported task in the `examples` folder (such as GLUE/SQuAD, ...)
- [ ] My own task or dataset (give details below)

### Reproduction

see above

### Expected behavior

see above

Comments:
Hey could you share a full reproducer as well as the output of `transformers-cli env`

@ArthurZucker sure, here's the code used in colab:

```
model_id = "meta-llama/Llama-2-7b-chat-hf"
llamamodel = AutoModelForCausalLM.from_pretrained(model_id,
                                                  torch_dtype=torch.float16,
                                                  load_in_4bit=True,
                                                  low_cpu_mem_usage=True, token=token)
llamatokenizer = AutoTokenizer.from_pretrained(model_id)
llamatokenizer.pad_token = llamatokenizer.eos_token

prompt = "some long text"
s = llamatokenizer(prompt, return_tensors='pt')
#s.input_ids.shape ==>  torch.Size([1, 1120])
llamamodel(**s)
```

tried to run `transformer-cli ` but it said "command not found" !


cc @SunMarc when he comes back for big model inference + quantization 

Hi @kaoutaar , `.generate()` uses `torch.no_grad()`. This might explain with you are getting OOM when you `call()`. 

> Hi @kaoutaar , `.generate()` uses `torch.no_grad()`. This might explain with you are getting OOM when you `call()`.

If we back-prop gradients via `torch.set_grad_enabled(True)`, how can we control the memory usage? 

You should use [gradient checkpointing ](https://huggingface.co/docs/transformers/v4.18.0/en/performance#gradient-checkpointing)

> gradient checkpointing

Dear Arthur,

Well received with thanks!

Best regards,

Shuyue
Dec. 4th, 2023