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

• 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

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() API, the call method uses an internal variable with the compiled forward instead of the uncompiled forward.

(I raised a related issue in pytorch, 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()