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 15, 14:45 UTC
[Bugfix] Fix broken GritLM model and tests (missing pooling_metadata) (#16631)
Signed-off-by: Pooya Davoodi <pooya.davoodi@parasail.io>
Pooya Davoodi committed
Apr 14, 2025 at 23:09 UTC
bc5dd4f669e2f83adec58b38ea11d75c74bc1706
2 files changed
+13
-11
tests/models/embedding/language/test_gritlm.py
+11
-10
@@ -57,24 +57,25 @@ def test_find_array(monkeypatch: pytest.MonkeyPatch):
57
def server_embedding():
58
# GritLM embedding implementation is only supported by XFormers backend.
59
args = ["--task", "embed", "--max_model_len", str(MAX_MODEL_LEN)]
60
- with RemoteOpenAIServer(MODEL_NAME, args) as remote_server:
61
- yield remote_server
60
+ with pytest.MonkeyPatch.context() as m:
61
+ m.setenv(STR_BACKEND_ENV_VAR, "XFORMERS")
62
+ with RemoteOpenAIServer(MODEL_NAME, args) as remote_server:
63
+ yield remote_server
64
65
66
@pytest.fixture(scope="module")
67
def server_generate():
68
args = ["--task", "generate", "--max_model_len", str(MAX_MODEL_LEN)]
67
- with RemoteOpenAIServer(MODEL_NAME, args) as remote_server:
68
- yield remote_server
69
+ with pytest.MonkeyPatch.context() as m:
70
+ m.setenv(STR_BACKEND_ENV_VAR, "XFORMERS")
71
+ with RemoteOpenAIServer(MODEL_NAME, args) as remote_server:
72
+ yield remote_server
73
74
75
@pytest_asyncio.fixture
72
-async def client_embedding(monkeypatch: pytest.MonkeyPatch,
73
- server_embedding: RemoteOpenAIServer):
74
- with monkeypatch.context() as m:
75
- m.setenv("VLLM_ATTENTION_BACKEND", "XFORMERS")
76
- async with server_embedding.get_async_client() as async_client:
77
- yield async_client
76
+async def client_embedding(server_embedding: RemoteOpenAIServer):
77
+ async with server_embedding.get_async_client() as async_client:
78
+ yield async_client
79
80
81
@pytest_asyncio.fixture
vllm/model_executor/models/gritlm.py
+2
-1
@@ -170,7 +170,8 @@ class GritLMPooler(nn.Module):
170
mean_embeddings = sum_embeddings / num_non_instruction_tokens.unsqueeze(
171
1)
172
173
- pooled_data = self.head(mean_embeddings)
173
+ pooled_data = self.head(mean_embeddings,
174
+ pooling_metadata=pooling_metadata)
175
176
pooled_outputs = [
177
PoolingSequenceGroupOutput(data) for data in pooled_data