xinyun/vllm
1
0
Fork 0
mirror of https://git.datalinker.icu/vllm-project/vllm.git synced 2025-12-25 16:40:53 +08:00
Code Issues Packages Projects Releases Wiki Activity

[Bugfix] Use runner_type instead of task in GritLM (#11144)

Browse Source 
Signed-off-by: Pooya Davoodi <pooya.davoodi@parasail.io>
This commit is contained in:
Pooya Davoodi 2024-12-12 20:09:53 -08:00 committed by GitHub
parent 30870b4f66
commit 1efce68605
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
2 changed files with 7 additions and 7 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

6
tests/models/embedding/language/test_gritlm.py
View File

@ -35,7 +35,7 @@ def test_find_array(monkeypatch):
from vllm.model_executor.models.gritlm import GritLMPooler
# Create an LLM object to get the model config.
llm = vllm.LLM(MODEL_NAME, task="embedding", max_model_len=MAX_MODEL_LEN)
llm = vllm.LLM(MODEL_NAME, task="embed", max_model_len=MAX_MODEL_LEN)
pooler = GritLMPooler(model_config=llm.llm_engine.model_config)
arr = _arr([0, 1, 2, 3, 4, 5, 6, 7, 8, 9])
@ -55,7 +55,7 @@ def server_embedding():
with pytest.MonkeyPatch.context() as mp:
mp.setenv("VLLM_ATTENTION_BACKEND", "XFORMERS")
args = ["--task", "embedding", "--max_model_len", str(MAX_MODEL_LEN)]
args = ["--task", "embed", "--max_model_len", str(MAX_MODEL_LEN)]
with RemoteOpenAIServer(MODEL_NAME, args) as remote_server:
yield remote_server
@ -141,7 +141,7 @@ def test_gritlm_offline_embedding(monkeypatch):
queries, q_instruction, documents, d_instruction = get_test_data()
llm = vllm.LLM(MODEL_NAME, task="embedding", max_model_len=MAX_MODEL_LEN)
llm = vllm.LLM(MODEL_NAME, task="embed", max_model_len=MAX_MODEL_LEN)
d_rep = run_llm_encode(
llm,

8
vllm/model_executor/models/gritlm.py
View File

@ -203,12 +203,12 @@ class GritLM(LlamaForCausalLM):
) -> None:
super().__init__(vllm_config=vllm_config, prefix=prefix, **kwargs)
self.task = vllm_config.model_config.task
self.runner_type = vllm_config.model_config.runner_type
self._pooler = GritLMPooler(vllm_config.model_config)
for layer in self.model.layers:
if self.task == "embedding" and hasattr(layer, "self_attn"):
if self.runner_type == "pooling" and hasattr(layer, "self_attn"):
assert isinstance(layer.self_attn.attn.impl, XFormersImpl), (
"GritLM embedding is only supported by XFormers backend, "
"which can be forced by VLLM_ATTENTION_BACKEND=XFORMERS")
@ -222,8 +222,8 @@ class GritLM(LlamaForCausalLM):
**kwargs,
) -> Union[torch.Tensor, IntermediateTensors]:
# Change attention to non-causal for embedding task.
if self.task == "embedding":
# Change attention to non-causal for pooling tasks.
if self.runner_type == "pooling":
assert attn_metadata.prefill_metadata.attn_bias is None
attn_metadata.prefill_metadata.attn_bias = [
BlockDiagonalMask.from_seqlens(attn_metadata.seq_lens)