biondizzle/vllm
1
0
Fork 0
Code Issues Pull Requests Actions 2 Packages Projects Releases Wiki Activity

Convert formatting to use ruff instead of yapf + isort (#26247)

Browse Source 
Signed-off-by: Harry Mellor <19981378+hmellor@users.noreply.github.com>
This commit is contained in:
Harry Mellor
2025-10-05 15:06:22 +01:00
committed by GitHub
parent 17edd8a807
commit d6953beb91
1508 changed files with 115244 additions and 94146 deletions
Download Patch File Download Diff File Expand all files Collapse all files

70
tests/model_executor/test_model_load_with_params.py
View File

@@ -5,8 +5,12 @@ import os
import pytest
from vllm.model_executor.layers.pooler import (CLSPool, DispatchPooler,
MeanPool, PoolingType)
from vllm.model_executor.layers.pooler import (
CLSPool,
DispatchPooler,
MeanPool,
PoolingType,
)
from vllm.model_executor.models.bert import BertEmbeddingModel
from vllm.model_executor.models.roberta import RobertaEmbeddingModel
from vllm.platforms import current_platform
@@ -15,25 +19,28 @@ MAX_MODEL_LEN = 128
MODEL_NAME = os.environ.get("MODEL_NAME", "BAAI/bge-base-en-v1.5")
REVISION = os.environ.get("REVISION", "main")
MODEL_NAME_ROBERTA = os.environ.get("MODEL_NAME",
"intfloat/multilingual-e5-base")
MODEL_NAME_ROBERTA = os.environ.get("MODEL_NAME", "intfloat/multilingual-e5-base")
REVISION_ROBERTA = os.environ.get("REVISION", "main")
@pytest.mark.skipif(current_platform.is_rocm(),
reason="Xformers backend is not supported on ROCm.")
@pytest.mark.skipif(
current_platform.is_rocm(), reason="Xformers backend is not supported on ROCm."
)
def test_model_loading_with_params(vllm_runner, monkeypatch):
"""
Test parameter weight loading with tp>1.
"""
# to use apply_model
monkeypatch.setenv("VLLM_ALLOW_INSECURE_SERIALIZATION", "1")
with vllm_runner(model_name=MODEL_NAME,
revision=REVISION,
dtype="float16",
max_model_len=MAX_MODEL_LEN) as vllm_model:
output = vllm_model.embed("Write a short story about a robot that"
" dreams for the first time.\n")
with vllm_runner(
model_name=MODEL_NAME,
revision=REVISION,
dtype="float16",
max_model_len=MAX_MODEL_LEN,
) as vllm_model:
output = vllm_model.embed(
"Write a short story about a robot that dreams for the first time.\n"
)
model_config = vllm_model.llm.llm_engine.model_config
model_tokenizer = vllm_model.llm.llm_engine.tokenizer
@@ -60,20 +67,24 @@ def test_model_loading_with_params(vllm_runner, monkeypatch):
assert output
@pytest.mark.skipif(current_platform.is_rocm(),
reason="Xformers backend is not supported on ROCm.")
@pytest.mark.skipif(
current_platform.is_rocm(), reason="Xformers backend is not supported on ROCm."
)
def test_roberta_model_loading_with_params(vllm_runner, monkeypatch):
"""
Test parameter weight loading with tp>1.
"""
# to use apply_model
monkeypatch.setenv("VLLM_ALLOW_INSECURE_SERIALIZATION", "1")
with vllm_runner(model_name=MODEL_NAME_ROBERTA,
revision=REVISION_ROBERTA,
dtype="float16",
max_model_len=MAX_MODEL_LEN) as vllm_model:
output = vllm_model.embed("Write a short story about a robot that"
" dreams for the first time.\n")
with vllm_runner(
model_name=MODEL_NAME_ROBERTA,
revision=REVISION_ROBERTA,
dtype="float16",
max_model_len=MAX_MODEL_LEN,
) as vllm_model:
output = vllm_model.embed(
"Write a short story about a robot that dreams for the first time.\n"
)
model_config = vllm_model.llm.llm_engine.model_config
model_tokenizer = vllm_model.llm.llm_engine.tokenizer
@@ -93,16 +104,16 @@ def test_roberta_model_loading_with_params(vllm_runner, monkeypatch):
def check_model(model):
assert isinstance(model, RobertaEmbeddingModel)
assert isinstance(pooler := model.pooler, DispatchPooler)
assert isinstance(pooler.poolers_by_task["embed"].pooling,
MeanPool)
assert isinstance(pooler.poolers_by_task["embed"].pooling, MeanPool)
vllm_model.apply_model(check_model)
assert output
@pytest.mark.skipif(current_platform.is_rocm(),
reason="Xformers backend is not supported on ROCm.")
@pytest.mark.skipif(
current_platform.is_rocm(), reason="Xformers backend is not supported on ROCm."
)
def test_facebook_roberta_model_loading_with_params(vllm_runner, monkeypatch):
"""
Test loading roberta-base model with no lm_head.
@@ -110,11 +121,12 @@ def test_facebook_roberta_model_loading_with_params(vllm_runner, monkeypatch):
# to use apply_model
monkeypatch.setenv("VLLM_ALLOW_INSECURE_SERIALIZATION", "1")
model_name = "FacebookAI/roberta-base"
with vllm_runner(model_name=model_name,
dtype="float16",
max_model_len=MAX_MODEL_LEN) as vllm_model:
output = vllm_model.embed("Write a short story about a robot that"
" dreams for the first time.\n")
with vllm_runner(
model_name=model_name, dtype="float16", max_model_len=MAX_MODEL_LEN
) as vllm_model:
output = vllm_model.embed(
"Write a short story about a robot that dreams for the first time.\n"
)
assert vllm_model.llm.llm_engine.model_config.tokenizer == model_name