@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / 78a47f87ce
Fork 0

A high-throughput and memory-efficient inference and serving engine for LLMs

Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git

Synced Nov 14, 05:30 UTC

Test Prompt Embeds/LoRA compatibility and Enable LoRA Support for OPT Models (#25717)

Signed-off-by: Andrew Sansom <andrew@protopia.ai>

Andrew Sansom committed Sep 29, 2025 at 19:10 UTC 78a47f87ce259a48f0391fa9ae15add05ea7432b
5 files changed +40 -11
docs/features/README.md
+1 -1
@@ -52,7 +52,7 @@ th:not(:first-child) {
52 | [mm](multimodal_inputs.md) | ✅ | ✅ | [🟠](gh-pr:4194)<sup>^</sup> | ❔ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ❔ | ✅ | | | |
53 | best-of | ✅ | ✅ | ✅ | [❌](gh-issue:6137) | ✅ | ❌ | ✅ | ✅ | ✅ | ❔ | [❌](gh-issue:7968) | ✅ | ✅ | | |
54 | beam-search | ✅ | ✅ | ✅ | [❌](gh-issue:6137) | ✅ | ❌ | ✅ | ✅ | ✅ | ❔ | [❌](gh-issue:7968) | ❔ | ✅ | ✅ | |
55 -| [prompt-embeds](prompt_embeds.md) | ✅ | [❌](gh-issue:25096) | ? | ❌ | ✅ | ❌ | ❌ | ✅ | ❌ | ? | ? | ❌ | ? | ? | ✅ |
55 +| [prompt-embeds](prompt_embeds.md) | ✅ | [❌](gh-issue:25096) | ✅ | ❌ | ✅ | ❌ | ❌ | ✅ | ❌ | ❔ | ❔ | ❌ | ❔ | ❔ | ✅ |
56
57 \* Chunked prefill and prefix caching are only applicable to last-token pooling.
58 <sup>^</sup> LoRA is only applicable to the language backbone of multimodal models.
docs/models/supported_models.md
+1 -1
@@ -403,7 +403,7 @@ th {
403 | `OLMo2ForCausalLM` | OLMo2 | `allenai/OLMo-2-0425-1B`, etc. | ✅︎ | ✅︎ | ✅︎ |
404 | `OLMo3ForCausalLM` | OLMo3 | TBA | ✅︎ | ✅︎ | ✅︎ |
405 | `OLMoEForCausalLM` | OLMoE | `allenai/OLMoE-1B-7B-0924`, `allenai/OLMoE-1B-7B-0924-Instruct`, etc. | | ✅︎ | ✅︎ |
406 -| `OPTForCausalLM` | OPT, OPT-IML | `facebook/opt-66b`, `facebook/opt-iml-max-30b`, etc. | | ✅︎ | ✅︎ |
406 +| `OPTForCausalLM` | OPT, OPT-IML | `facebook/opt-66b`, `facebook/opt-iml-max-30b`, etc. | ✅︎ | ✅︎ | ✅︎ |
407 | `OrionForCausalLM` | Orion | `OrionStarAI/Orion-14B-Base`, `OrionStarAI/Orion-14B-Chat`, etc. | | ✅︎ | ✅︎ |
408 | `PhiForCausalLM` | Phi | `microsoft/phi-1_5`, `microsoft/phi-2`, etc. | ✅︎ | ✅︎ | ✅︎ |
409 | `Phi3ForCausalLM` | Phi-4, Phi-3 | `microsoft/Phi-4-mini-instruct`, `microsoft/Phi-4`, `microsoft/Phi-3-mini-4k-instruct`, `microsoft/Phi-3-mini-128k-instruct`, `microsoft/Phi-3-medium-128k-instruct`, etc. | ✅︎ | ✅︎ | ✅︎ |
tests/entrypoints/conftest.py
+8
@@ -208,3 +208,11 @@ def zephyr_lora_files():
208 """Download zephyr LoRA files once per test session."""
209 from huggingface_hub import snapshot_download
210 return snapshot_download(repo_id="typeof/zephyr-7b-beta-lora")
211 +
212 +
213 +@pytest.fixture(scope="session")
214 +def opt125_lora_files() -> str:
215 + """Download opt-125m LoRA files once per test session."""
216 + from huggingface_hub import snapshot_download
217 + return snapshot_download(
218 + repo_id="peft-internal-testing/opt-125m-dummy-lora")
tests/entrypoints/openai/test_completion_with_prompt_embeds.py
+28 -6
@@ -3,6 +3,7 @@
3
4 import base64
5 import io
6 +import json
7
8 import openai # use the official client for correctness check
9 import pytest
@@ -16,13 +17,15 @@ from ...utils import RemoteOpenAIServer
17
18 # any model with a chat template should work here
19 MODEL_NAME = "facebook/opt-125m"
20 +LORA_SERVING_MODEL_NAME = "opt125m-lora"
21
22 CONFIG = AutoConfig.from_pretrained(MODEL_NAME)
23
24
23 -@pytest.fixture(scope="module")
24 -def default_server_args() -> list[str]:
25 - return [
25 +@pytest.fixture(scope="module", params=["use-lora"])
26 +def default_server_args(request: pytest.FixtureRequest,
27 + opt125_lora_files: str) -> list[str]:
28 + args = [
29 # use half precision for speed and memory savings in CI environment
30 "--dtype",
31 "bfloat16",
@@ -35,6 +38,25 @@ def default_server_args() -> list[str]:
38 "--enable-prompt-embeds",
39 ]
40
41 + if request.param == "use-lora":
42 + lora_module_1 = {
43 + "name": LORA_SERVING_MODEL_NAME,
44 + "path": opt125_lora_files,
45 + "base_model_name": MODEL_NAME
46 + }
47 +
48 + args.extend([
49 + "--enable-lora",
50 + "--lora-module",
51 + json.dumps(lora_module_1),
52 + "--max-lora-rank",
53 + "64",
54 + "--max-cpu-loras",
55 + "2",
56 + ])
57 +
58 + return args
59 +
60
61 EXAMPLE_PROMPTS = [
62 "Hello, my name is",
@@ -74,7 +96,7 @@ async def client_with_prompt_embeds(server_with_prompt_embeds):
96
97
98 @pytest.mark.asyncio
77 -@pytest.mark.parametrize("model_name", [MODEL_NAME])
99 +@pytest.mark.parametrize("model_name", [MODEL_NAME, LORA_SERVING_MODEL_NAME])
100 async def test_completions_with_prompt_embeds(
101 example_prompt_embeds,
102 client_with_prompt_embeds: openai.AsyncOpenAI,
@@ -179,7 +201,7 @@ async def test_completions_with_prompt_embeds(
201
202
203 @pytest.mark.asyncio
182 -@pytest.mark.parametrize("model_name", [MODEL_NAME])
204 +@pytest.mark.parametrize("model_name", [MODEL_NAME, LORA_SERVING_MODEL_NAME])
205 async def test_completions_errors_with_prompt_embeds(
206 client_with_prompt_embeds: openai.AsyncOpenAI, model_name: str):
207 # Test error case: invalid prompt_embeds
@@ -194,7 +216,7 @@ async def test_completions_errors_with_prompt_embeds(
216
217 @pytest.mark.asyncio
218 @pytest.mark.parametrize("logprobs_arg", [1, 0])
197 -@pytest.mark.parametrize("model_name", [MODEL_NAME])
219 +@pytest.mark.parametrize("model_name", [MODEL_NAME, LORA_SERVING_MODEL_NAME])
220 async def test_completions_with_logprobs_and_prompt_embeds(
221 example_prompt_embeds,
222 client_with_prompt_embeds: openai.AsyncOpenAI,
vllm/model_executor/models/opt.py
+2 -3
@@ -43,7 +43,7 @@ from vllm.model_executor.layers.vocab_parallel_embedding import (
43 from vllm.model_executor.model_loader.weight_utils import default_weight_loader
44 from vllm.sequence import IntermediateTensors
45
46 -from .interfaces import SupportsPP
46 +from .interfaces import SupportsLoRA, SupportsPP
47 from .utils import (AutoWeightsLoader, WeightsMapper, is_pp_missing_parameter,
48 make_empty_intermediate_tensors_factory, make_layers,
49 maybe_prefix)
@@ -352,10 +352,9 @@ class OPTModel(nn.Module):
352 return loaded_params
353
354
355 -class OPTForCausalLM(nn.Module, SupportsPP):
355 +class OPTForCausalLM(nn.Module, SupportsPP, SupportsLoRA):
356 packed_modules_mapping = {
357 "qkv_proj": ["q_proj", "k_proj", "v_proj"],
358 - "gate_up_proj": ["gate_proj", "up_proj"]
358 }
359
360 hf_to_vllm_mapper = WeightsMapper(orig_to_new_prefix={