Aligning modling code for GPT2 to work with vLLM (fallback)

[ariG23498]

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 vllm-project/vllm#15290 on the vLLM side to work.

One can use the following snippet to check the implementation:

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}")

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))

[github-actions]

Hi 👋, thank you for opening this pull request! The pull request is converted to draft by default. When it is ready for review, please click the Ready for review button (at the bottom of the PR page).

[ariG23498]

This is done for weight loading.

We need the prefix set for different platforms.

[hmellor]

I'm planning to investigate this on the vLLM side to see if we can remove this requirement

[HuggingFaceDocBuilderDev]

The docs for this PR live here. All of your documentation changes will be reflected on that endpoint. The docs are available until 30 days after the last update.

[ArthurZucker]

nice thanks!

[ArthurZucker]

Just make sure CI is green!

[ariG23498]

@ArthurZucker I don't think the CI errors stem from the modeling changes. Do you want me to investigate further?

[hmellor]

Adding a base_model_tp_plan, which is currently empty. This is a required attribute for vLLM

I could make this optional so that a model which does not have it simply doesn't support TP, rather than not working at all?

[ariG23498]

@ArthurZucker would it be okay to merge? The CI is broken for issues not related to the PR it seems 😅