torch.compile: generate should use call instead of forward
Author: SilverSoldier
Pull request: no
Labels: name: bug

### System Info

- `transformers` version: 4.47.0.dev0
- Platform: Linux-5.14.0-284.73.1.el9_2.x86_64-x86_64-with-glibc2.35
- Python version: 3.10.12
- Huggingface_hub version: 0.25.2
- Safetensors version: 0.4.5
- Accelerate version: not installed
- Accelerate config: not found
- PyTorch version (GPU?): 2.6.0.dev20241008+cu124 (False)
- Tensorflow version (GPU?): not installed (NA)
- Flax version (CPU?/GPU?/TPU?): not installed (NA)
- Jax version: not installed
- JaxLib version: not installed
- Using distributed or parallel set-up in script?: NO

### Who can help?

@ArthurZucker @Cyrilvallez 

### Information

- [X] The official example scripts
- [ ] My own modified scripts

### Tasks

- [X] An officially supported task in the `examples` folder (such as GLUE/SQuAD, ...)
- [ ] My own task or dataset (give details below)

### Reproduction

```
import torch
from transformers import AutoModelForCausalLM, AutoTokenizer

model_name = "facebook/opt-125m"
length = 100

prompt_text = 'In a small, bustling cafe nestled in the heart of a vibrant city, a serendipitous event unfolded, leaving a lasting impression on all who witnessed it. As the patrons sat sipping their coffees and engaging in animated conversations, a talented street musician entered the cafe, carrying a weathered guitar and radiating an aura of creativity.'

tokenizer = AutoTokenizer.from_pretrained(model_name, use_fast=True)
model = AutoModelForCausalLM.from_pretrained(model_name)
model.compile()
input_ids = tokenizer(prompt_text, add_special_tokens=False, return_tensors='pt').input_ids
output = model.generate(input_ids, max_new_tokens=length)
```

### Expected behavior

Expected behaviour is that we use the compiled forward function.

When compiling using the [`model.compile()`](https://github.com/pytorch/pytorch/blob/a6344c8bcd22798987087244e961cdc0cbf9e9df/torch/nn/modules/module.py#L2985) API, the call method uses an [internal variable](https://github.com/pytorch/pytorch/blob/a6344c8bcd22798987087244e961cdc0cbf9e9df/torch/nn/modules/module.py#L1736) with the compiled forward instead of the uncompiled forward. 

(I raised a [related issue in pytorch](https://github.com/pytorch/pytorch/issues/141473), this is the Option 2 there)

So generate, should use the call method instead of the forward to use the compiled version of forward (for this particular case of model.compile).
However, recent changes have changed this call to model.forward() instead of model() for the non-first token :

```
def _sample():
  ...
  def model_forward(model, *args, **kwargs):
      return model.forward(*args, **kwargs)
  ...
      if i == 0:
          outputs = self(**model_inputs, return_dict=True)
          i += 1
      else:
          outputs = model_forward(self, return_dict=True, **model_inputs)
```

model_forward should be changed to call model() instead of model.forward()

Comments:
Hi, so if I understand correctly, the goal of this change is to make the future `model(...)` (potentially used outside `generate`) will use the already compiled version (if that is done in `generate` here). It's good, but currently code isn't really a bug, right?

Previous code called `model()` which works with `model.compile()`. Current code (from [this recent commit](https://github.com/huggingface/transformers/commit/597efd21d20fd320a167fef707a0bcf2be205725)) changed this to `model.forward()`. 
So current code does not work as expected with `model.compile()`, it uses compiled version for the first iteration and eager mode for other iterations. This is sort of a bug.

Hi @SilverSoldier, thanks for opening the issue! Indeed I agree that we should use `__call__`, for consistency between all methods. We wanted to refine this part anyway, this was the first shot! I'll take care of it very soon. Curious to see what the pytorch team has to say with regards to the consistency of the different way to `compile` a model, but I'm not sure this can be solved in general.

Also, we introduced the recent changes because it is inefficient to `compile` the forward in all cases (we only want to compile the iterative decoding part, not the prefill), so your code would me more efficient if you don't call `compile` yourself, and only use `cache_implementation=static` in `generate`, and let us only compile the iterative decoding part 🤗

I see, it's another API of compile! Thanks!

@Cyrilvallez interesting. However, if generate takes care of the compile how can we pass the compile arguments, say custom backend and other parameters?

We are adding this via the classic `generate_config` and `generate_kwargs` api! 

This issue has been automatically marked as stale because it has not had recent activity. If you think this still needs to be addressed please comment on this thread.

Please note that issues that do not follow the [contributing guidelines](https://github.com/huggingface/transformers/blob/main/CONTRIBUTING.md) are likely to be ignored.