[Core] Add Lora Support to Beam Search

[alex-jw-brooks]

Adds support for passing lora adapters for beam search (offline and online).

FIX #17205

Sample usage (offline):

from vllm import LLM
from vllm.sampling_params import BeamSearchParams
from vllm.lora.request import LoRARequest
from huggingface_hub import snapshot_download

sql_lora_path = snapshot_download(repo_id="yard1/llama-2-7b-sql-lora-test")
llm = LLM(model="meta-llama/Llama-2-7b-hf", enable_lora=True)

sampling_params = BeamSearchParams(
    beam_width=2,
    temperature=0,
    max_tokens=256,
)

outputs = llm.beam_search(
    [
        {"prompt": "[user] Write a SQL query to answer the question based on the table schema.\n\n context: CREATE TABLE table_name_74 (icao VARCHAR, airport VARCHAR)\n\n question: Name the ICAO for lilongwe international airport [/user] [assistant]"},
        {"prompt": "[user] Write a SQL query to answer the question based on the table schema.\n\n context: CREATE TABLE table_name_11 (nationality VARCHAR, elector VARCHAR)\n\n question: When Anchero Pantaleone was the elector what is under nationality? [/user] [assistant]"},
    ],
    sampling_params,
    # Note - you can also pass a list of LoRARequest/Nones that aligns with the prompt list
    lora_request=LoRARequest("sql_adapter", 1, sql_lora_path),
)

for output in outputs:
    print(f"{output.sequences[0].text}\n\n")

sample output:

<s> [user] Write a SQL query to answer the question based on the table schema.

 context: CREATE TABLE table_name_74 (icao VARCHAR, airport VARCHAR)

 question: Name the ICAO for lilongwe international airport [/user] [assistant] Write a SQL query to answer the question based on the table schema.  context: CREATE TABLE table_name_74 (icao VARCHAR, airport VARCHAR)  question: Name the ICAO for lilongwe international airport     SELECT ICAO FROM table_name_74 WHERE airport = 'lilongwe international airport'     SELECT ICAO FROM table_name_74 WHERE airport = 'lilongwe international airport'    SELECT ICAO FROM table_name_74 WHERE airport = 'lilongwe international airport'   SELECT ICAO FROM table_name_74 WHERE airport = 'lilongwe international airport'   SELECT ICAO FROM table_name_74 WHERE airport = 'lilongwe international airport'  SELECT ICAO FROM table_name_74 WHERE airport = 'lilongwe international airport'  SELECT ICAO FROM table_name_74 WHERE airport = 'lilongwe international airport'  SELECT ICAO FROM table_name_74 WHERE


<s> [user] Write a SQL query to answer the question based on the table schema.

 context: CREATE TABLE table_name_11 (nationality VARCHAR, elector VARCHAR)

 question: When Anchero Pantaleone was the elector what is under nationality? [/user] [assistant] Write a SQL query to answer the question based on the table schema.  context: CREATE TABLE table_name_11 (nationality VARCHAR, elector VARCHAR)  question: When Anchero Pantaleone was the elector what is under nationality?     SELECT nationality FROM table_name_11 WHERE elector = 'anchero pantaleone'    SELECT nationality FROM table_name_11 WHERE elector = 'anchero pantaleone'    SELECT nationality FROM table_name_11 WHERE elector = 'anchero pantaleone'   SELECT nationality FROM table_name_11 WHERE elector = 'anchero pantaleone'   SELECT nationality FROM table_name_11 WHERE elector = 'anchero pantaleone'  SELECT nationality FROM table_name_11 WHERE elector = 'anchero pantaleone'  SELECT nationality FROM table_name_11 WHERE elector = 'anchero pantaleone'  SELECT nationality FROM table_name_11 WHERE elector = 'anchero pantaleone'

Sample usage (online):

Start the server:

python -m vllm.entrypoints.openai.api_server --device cuda --model meta-llama/Llama-2-7b-hf --api-key token-abc123 --dtype=half --enable-lora  --lora-modules test=yard1/llama-2-7b-sql-lora-test

from openai import OpenAI
client = OpenAI(
    base_url="http://localhost:8000/v1",
    api_key="token-abc123",
)

completion = client.completions.create(
    model="test",
    prompt="[user] Write a SQL query to answer the question based on the table schema.\n\n context: CREATE TABLE table_name_74 (icao VARCHAR, airport VARCHAR)\n\n question: Name the ICAO for lilongwe international airport [/user] [assistant]",
    n=2,
    extra_body=dict(use_beam_search=True),
    max_tokens=256,
)

print(completion.choices[0].text)

Sample output:

 SELECT icao FROM table_name_74 WHERE airport = 'lilongwe international airport'     SELECT icao FROM table_name_74 WHERE airport = 'lilongwe international airport'    SELECT icao FROM table_name_74 WHERE airport = 'lilongwe international airport'   SELECT icao FROM table_name_74 WHERE airport = 'lilongwe international airport'   SELECT icao FROM table_name_74 WHERE airport = 'lilongwe international airport'   SELECT icao FROM table_name_74 WHERE airport = 'lilongwe international airport'  SELECT icao FROM table_name_74 WHERE airport = 'lilongwe international airport'  SELECT icao FROM table_name_74 WHERE airport = 'lilongwe international airport'  SELECT icao FROM table_name_74 WHERE airport = 'lilongwe international airport'  SELECT icao FROM table_name_74 WHERE airport

CC @jeejeelee

[github-actions]

👋 Hi! Thank you for contributing to the vLLM project.

💬 Join our developer Slack at https://slack.vllm.ai to discuss your PR in #pr-reviews, coordinate on features in #feat- channels, or join special interest groups in #sig- channels.

Just a reminder: PRs would not trigger full CI run by default. Instead, it would only run fastcheck CI which starts running only a small and essential subset of CI tests to quickly catch errors. You can run other CI tests on top of those by going to your fastcheck build on Buildkite UI (linked in the PR checks section) and unblock them. If you do not have permission to unblock, ping simon-mo or khluu to add you in our Buildkite org.

Once the PR is approved and ready to go, your PR reviewer(s) can run CI to test the changes comprehensively before merging.

To run CI, PR reviewers can either: Add ready label to the PR or enable auto-merge.

🚀

[jeejeelee]

Thank you, will look at this PR ASAP

[jeejeelee]

LGTM, thanks

[jeejeelee]

@alex-jw-brooks Please sync with the main branch to verify if it can fix the CI failure

[alex-jw-brooks]

Awesome, thanks for the review @jeejeelee! Rebased this PR, but it looks like the image build timed out. Could you please retry it?

[DarkLight1337]

Retrying

[alex-jw-brooks]

Awesome, thanks @DarkLight1337! I think the test failures look unrelated