Repository: huggingface/transformers
Pull request: Aligning modling code for GPT2 to work with vLLM (fallback)

This PR changes the modeling code for GPT2 to support working on vLLM using the `transformers fallback backend`.

The changes are as follows:
- Introduction of `kwargs` which is used to propagate information about `attention_indices` in vLLM
- Adding a `base_model_tp_plan`, which is currently empty. This is a required attribute for vLLM
- Changing the reshape structure for the `attn_outputs` (took help from the llama code)

This PR is dependent on https://github.com/vllm-project/vllm/pull/15290 on the vLLM side to work.

One can use the following snippet to check the implementation:

```py
from vllm import LLM, SamplingParams
llm = LLM("openai-community/gpt2", model_impl="transformers", tensor_parallel_size=1)
prompts = [
    "Hello, my name is",
    "The president of the United States is",
    "The capital of France is",
    "The future of AI is",
]
sampling_params = SamplingParams(temperature=0.8, top_p=0.95)
outputs = llm.generate(prompts, sampling_params)
for output in outputs:
    prompt = output.prompt
    generated_text = output.outputs[0].text
    print(f"Prompt: {prompt!r}, Generated text: {generated_text!r}")
```

```py
from transformers import AutoModelForCausalLM, AutoTokenizer

tok = AutoTokenizer.from_pretrained("gpt2")
model = AutoModelForCausalLM.from_pretrained("gpt2")

inputs = tok("Hello there ", return_tensors="pt")
outputs = model.generate(**inputs)
print(tok.batch_decode(outputs))
```

Changed files:
filename: docs/source/en/model_doc/gpt2.md
patch: @@ -73,6 +73,12 @@ echo -e "Hello, I'm a language model" | transformers run --task text-generation
 </hfoption>
 </hfoptions>
 
+One can also serve the model using vLLM with the `transformers backend`.
+
+```
+vllm serve openai-community/gpt2 --model-imp transformers
+```
+
 Quantization reduces the memory burden of large models by representing the weights in a lower precision. Refer to the [Quantization](../quantization/overview) overview for more available quantization backends.
 
 The example below uses [bitsandbytes](../quantization/bitsandbytes) to only quantize the weights to 4-bits.

filename: src/transformers/models/decision_transformer/modeling_decision_transformer.py
patch: @@ -396,6 +396,7 @@ def forward(
         encoder_attention_mask: Optional[torch.FloatTensor] = None,
         use_cache: Optional[bool] = False,
         output_attentions: Optional[bool] = False,
+        **kwargs,
     ) -> Union[Tuple[torch.Tensor], Optional[Tuple[torch.Tensor, Tuple[torch.FloatTensor, ...]]]]:
         residual = hidden_states
         hidden_states = self.ln_1(hidden_states)
@@ -407,6 +408,7 @@ def forward(
             head_mask=head_mask,
             use_cache=use_cache,
             output_attentions=output_attentions,
+            **kwargs,
         )
         # residual connection
         hidden_states = attn_output + residual

filename: src/transformers/models/gpt2/modeling_gpt2.py
patch: @@ -401,6 +401,7 @@ def forward(
         encoder_attention_mask: Optional[torch.FloatTensor] = None,
         use_cache: Optional[bool] = False,
         output_attentions: Optional[bool] = False,
+        **kwargs,
     ) -> Union[Tuple[torch.Tensor], Optional[Tuple[torch.Tensor, Tuple[torch.FloatTensor, ...]]]]:
         residual = hidden_states
         hidden_states = self.ln_1(hidden_states)
@@ -412,6 +413,7 @@ def forward(
             head_mask=head_mask,
             use_cache=use_cache,
             output_attentions=output_attentions,
+            **kwargs,
         )
         # residual connection
         hidden_states = attn_output + residual
@@ -567,6 +569,7 @@ class GPT2PreTrainedModel(PreTrainedModel):
     _skip_keys_device_placement = "past_key_values"
     _supports_flash_attn_2 = True
     _supports_sdpa = True
+    _supports_attention_backend = True
     _supports_cache_class = True
     _supports_static_cache = True
 
@@ -903,6 +906,7 @@ def forward(
         output_attentions: Optional[bool] = None,
         output_hidden_states: Optional[bool] = None,
         return_dict: Optional[bool] = None,
+        **kwargs,
     ) -> Union[Tuple, BaseModelOutputWithPastAndCrossAttentions]:
         output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
         output_hidden_states = (
@@ -1046,6 +1050,7 @@ def forward(
                     encoder_attention_mask=encoder_attention_mask,
                     use_cache=use_cache,
                     output_attentions=output_attentions,
+                    **kwargs,
                 )
 
             hidden_states = outputs[0]