Add vLLM v0.18.1 source tree with KV transfer abort fix
third_party/vllm/ now tracked in git for direct patch management.
Based on vLLM v0.18.1 release with one patch applied:
vllm/v1/core/sched/scheduler.py:
Replace fatal assert with graceful skip when KV transfer callback
arrives for an already-aborted request during PD disaggregated serving.
Future vLLM modifications should be made directly in third_party/vllm/
and committed normally. The patches/ directory is kept as documentation
of what changed from upstream.
Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
This commit is contained in:
217
third_party/vllm/tests/lora/test_olmoe_tp.py
vendored
Normal file
217
third_party/vllm/tests/lora/test_olmoe_tp.py
vendored
Normal file
@@ -0,0 +1,217 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
|
||||
import shutil
|
||||
from collections.abc import Sequence
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
from safetensors.torch import load_file, save_file
|
||||
|
||||
import vllm
|
||||
from vllm.lora.request import LoRARequest
|
||||
|
||||
from ..utils import multi_gpu_test
|
||||
|
||||
MODEL_PATH = "allenai/OLMoE-1B-7B-0125-Instruct"
|
||||
|
||||
PROMPT_TEMPLATE = """I want you to act as a SQL terminal in front of an example database, you need only to return the sql command to me. Do not return any additional explanation. Below is an instruction that describes a task, Write a response that appropriately completes the request.
|
||||
"
|
||||
##Instruction:
|
||||
candidate_poll contains tables such as candidate, people. Table candidate has columns such as Candidate_ID, People_ID, Poll_Source, Date, Support_rate, Consider_rate, Oppose_rate, Unsure_rate. Candidate_ID is the primary key.
|
||||
Table people has columns such as People_ID, Sex, Name, Date_of_Birth, Height, Weight. People_ID is the primary key.
|
||||
The People_ID of candidate is the foreign key of People_ID of people.
|
||||
|
||||
|
||||
###Input:
|
||||
{context}
|
||||
|
||||
###Response:""" # noqa: E501
|
||||
|
||||
EXPECTED_LORA_OUTPUT = [
|
||||
"SELECT count(*) FROM candidate",
|
||||
"SELECT count(*) FROM candidate",
|
||||
"SELECT poll_source FROM candidate GROUP BY poll_source ORDER BY count(*) DESC LIMIT 1", # noqa: E501
|
||||
"SELECT poll_source FROM candidate GROUP BY poll_source ORDER BY count(*) DESC LIMIT 1", # noqa: E501
|
||||
]
|
||||
|
||||
EXPECTED_BASE_MODEL_OUTPUT = [
|
||||
"SELECT COUNT(Candidate_ID) FROM candidate",
|
||||
"SELECT COUNT(Candidate_ID) FROM candidate",
|
||||
"SELECT Candidate_ID, COUNT(*) as Total_Candidates\nFROM candidate\nINNER JOIN people ON candidate.People_ID = people.People_ID", # noqa: E501
|
||||
# There are multiple acceptable responses
|
||||
(
|
||||
"SELECT Candidate_ID, Poll_Source FROM candidate WHERE People_ID IN (SELECT People_ID FROM people) ORDER BY COUNT(*) DESC LIMIT 1", # noqa: E501
|
||||
"SELECT Candidate_ID, Poll_Source FROM candidate WHERE COUNT(People_ID) = (SELECT COUNT(People_ID) FROM people) ORDER BY Candidate_ID DESC LIMIT 1", # noqa: E501
|
||||
),
|
||||
]
|
||||
|
||||
|
||||
def _output_matches(generated: str, accepted: str | Sequence[str]) -> bool:
|
||||
if isinstance(accepted, str):
|
||||
accepted = (accepted,)
|
||||
return any(generated.startswith(s) for s in accepted)
|
||||
|
||||
|
||||
def generate_and_test(
|
||||
llm: vllm.LLM,
|
||||
lora_path: str,
|
||||
lora_id: list[int | None] | int | None,
|
||||
compare_lower: bool = False,
|
||||
) -> None:
|
||||
prompts = [
|
||||
PROMPT_TEMPLATE.format(context="How many candidates are there?"),
|
||||
PROMPT_TEMPLATE.format(context="Count the number of candidates."),
|
||||
PROMPT_TEMPLATE.format(
|
||||
context="Which poll resource provided the most number of candidate information?" # noqa: E501
|
||||
),
|
||||
PROMPT_TEMPLATE.format(
|
||||
context="Return the poll resource associated with the most candidates."
|
||||
),
|
||||
]
|
||||
|
||||
lora_request = None
|
||||
if isinstance(lora_id, int):
|
||||
lora_request = LoRARequest(str(lora_id), lora_id, lora_path)
|
||||
elif isinstance(lora_id, list):
|
||||
lora_request = [
|
||||
LoRARequest(str(i), i, lora_path) if i is not None else None
|
||||
for i in lora_id
|
||||
]
|
||||
|
||||
sampling_params = vllm.SamplingParams(temperature=0, max_tokens=64)
|
||||
outputs = llm.generate(prompts, sampling_params, lora_request=lora_request)
|
||||
# Print the outputs.
|
||||
generated_texts: list[str] = []
|
||||
for output in outputs:
|
||||
prompt = output.prompt
|
||||
generated_text = output.outputs[0].text.strip()
|
||||
generated_texts.append(generated_text)
|
||||
print(f"Prompt: {prompt!r}, Generated text: {generated_text!r}")
|
||||
|
||||
for i in range(len(EXPECTED_LORA_OUTPUT)):
|
||||
req_lora_id = lora_id[i] if isinstance(lora_id, list) else lora_id
|
||||
generated_text = generated_texts[i]
|
||||
expected_output = (
|
||||
EXPECTED_LORA_OUTPUT[i]
|
||||
if req_lora_id is not None
|
||||
else EXPECTED_BASE_MODEL_OUTPUT[i]
|
||||
)
|
||||
|
||||
if compare_lower:
|
||||
generated_text = generated_text.lower()
|
||||
if isinstance(expected_output, str):
|
||||
expected_output = (expected_output.lower(),)
|
||||
else:
|
||||
expected_output = tuple(s.lower() for s in expected_output)
|
||||
assert _output_matches(generated_text, expected_output), (
|
||||
f"Output {i}: {generated_text!r} does not match any of {expected_output!r}"
|
||||
)
|
||||
|
||||
|
||||
def test_olmoe_lora(olmoe_lora_files):
|
||||
# We enable enforce_eager=True here to reduce VRAM usage for lora-test CI,
|
||||
# Otherwise, the lora-test will fail due to CUDA OOM.
|
||||
llm = vllm.LLM(
|
||||
MODEL_PATH,
|
||||
max_model_len=1024,
|
||||
enable_lora=True,
|
||||
max_loras=4,
|
||||
enforce_eager=True,
|
||||
trust_remote_code=True,
|
||||
enable_chunked_prefill=True,
|
||||
)
|
||||
|
||||
generate_and_test(llm, olmoe_lora_files, lora_id=1)
|
||||
generate_and_test(llm, olmoe_lora_files, lora_id=2)
|
||||
|
||||
|
||||
def test_olmoe_lora_mixed(olmoe_lora_files):
|
||||
llm = vllm.LLM(
|
||||
MODEL_PATH,
|
||||
max_model_len=1024,
|
||||
enable_lora=True,
|
||||
max_loras=4,
|
||||
enforce_eager=True,
|
||||
trust_remote_code=True,
|
||||
enable_chunked_prefill=True,
|
||||
)
|
||||
|
||||
generate_and_test(llm, olmoe_lora_files, lora_id=[1, None, 3, None])
|
||||
|
||||
|
||||
def test_olmoe_lora_mixed_random(olmoe_lora_files, tmp_path):
|
||||
# Create a dummy LoRA with random weights based on the real one
|
||||
random_lora_path = tmp_path / "random_lora"
|
||||
shutil.copytree(olmoe_lora_files, random_lora_path)
|
||||
|
||||
weights_path = random_lora_path / "adapter_model.safetensors"
|
||||
weights = load_file(str(weights_path))
|
||||
random_weights = {k: torch.randn_like(v) for k, v in weights.items()}
|
||||
save_file(random_weights, str(weights_path))
|
||||
|
||||
llm = vllm.LLM(
|
||||
MODEL_PATH,
|
||||
max_model_len=1024,
|
||||
enable_lora=True,
|
||||
max_loras=4,
|
||||
enforce_eager=True,
|
||||
trust_remote_code=True,
|
||||
enable_chunked_prefill=True,
|
||||
)
|
||||
|
||||
prompts = [
|
||||
PROMPT_TEMPLATE.format(context="How many candidates are there?"),
|
||||
PROMPT_TEMPLATE.format(context="Count the number of candidates."),
|
||||
]
|
||||
|
||||
lora_requests = [
|
||||
LoRARequest("real", 1, olmoe_lora_files),
|
||||
LoRARequest("random", 2, str(random_lora_path)),
|
||||
]
|
||||
|
||||
sampling_params = vllm.SamplingParams(temperature=0, max_tokens=64)
|
||||
outputs = llm.generate(prompts, sampling_params, lora_request=lora_requests)
|
||||
assert outputs[0].outputs[0].text.strip().startswith(EXPECTED_LORA_OUTPUT[0])
|
||||
|
||||
|
||||
@pytest.mark.parametrize("fully_sharded_loras", [False, True])
|
||||
@multi_gpu_test(num_gpus=2)
|
||||
def test_olmoe_lora_tp2(olmoe_lora_files, fully_sharded_loras):
|
||||
llm = vllm.LLM(
|
||||
MODEL_PATH,
|
||||
max_model_len=1024,
|
||||
enable_lora=True,
|
||||
max_loras=4,
|
||||
enforce_eager=True,
|
||||
trust_remote_code=True,
|
||||
enable_chunked_prefill=True,
|
||||
tensor_parallel_size=2,
|
||||
fully_sharded_loras=fully_sharded_loras,
|
||||
)
|
||||
|
||||
generate_and_test(llm, olmoe_lora_files, lora_id=1)
|
||||
generate_and_test(llm, olmoe_lora_files, lora_id=2)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("fully_sharded_loras", [False, True])
|
||||
@multi_gpu_test(num_gpus=4)
|
||||
def test_olmoe_lora_tp4(olmoe_lora_files, fully_sharded_loras):
|
||||
llm = vllm.LLM(
|
||||
MODEL_PATH,
|
||||
max_model_len=1024,
|
||||
enable_lora=True,
|
||||
max_loras=4,
|
||||
enforce_eager=True,
|
||||
trust_remote_code=True,
|
||||
enable_chunked_prefill=True,
|
||||
tensor_parallel_size=4,
|
||||
fully_sharded_loras=fully_sharded_loras,
|
||||
)
|
||||
generate_and_test(
|
||||
llm, olmoe_lora_files, lora_id=1, compare_lower=fully_sharded_loras
|
||||
)
|
||||
generate_and_test(
|
||||
llm, olmoe_lora_files, lora_id=2, compare_lower=fully_sharded_loras
|
||||
)
|
||||
Reference in New Issue
Block a user