[Bugfix] Replace PoolingParams.normalize with use_activation (#32243)

Signed-off-by: DarkLight1337 <tlleungac@connect.ust.hk>
This commit is contained in:
Cyrus Leung
2026-01-13 18:45:42 +08:00
committed by simon-mo
parent 11b6af5280
commit 0aa8c40552
21 changed files with 68 additions and 70 deletions

View File

@@ -53,7 +53,9 @@ def test_token_embed(llm: LLM):
def test_pooling_params(llm: LLM):
def get_outputs(normalize):
outputs = llm.embed(
prompts, pooling_params=PoolingParams(normalize=normalize), use_tqdm=False
prompts,
pooling_params=PoolingParams(use_activation=normalize),
use_tqdm=False,
)
return torch.tensor([x.outputs.embedding for x in outputs])

View File

@@ -216,7 +216,7 @@ def server_with_chunked_processing():
"512", # Set smaller max_model_len to trigger chunking mechanism
"--pooler-config",
(
'{"pooling_type": "MEAN", "normalize": true, '
'{"pooling_type": "MEAN", "use_activation": true, '
'"enable_chunked_processing": true, "max_embed_len": 10000}'
),
"--gpu-memory-utilization",

View File

@@ -236,17 +236,14 @@ class TestModel:
"use_activation": use_activation,
},
)
if response.status_code != 200:
return response
outputs = response.json()
return torch.tensor([x["score"] for x in outputs["data"]])
if model["is_cross_encoder"]:
default = get_outputs(use_activation=None)
w_activation = get_outputs(use_activation=True)
wo_activation = get_outputs(use_activation=False)
default = get_outputs(use_activation=None)
w_activation = get_outputs(use_activation=True)
wo_activation = get_outputs(use_activation=False)
if model["is_cross_encoder"]:
assert torch.allclose(default, w_activation, atol=1e-2), (
"Default should use activation."
)
@@ -256,9 +253,3 @@ class TestModel:
assert torch.allclose(F.sigmoid(wo_activation), w_activation, atol=1e-2), (
"w_activation should be close to activation(wo_activation)."
)
else:
get_outputs(use_activation=None)
# The activation parameter only works for the is_cross_encoder model
response = get_outputs(use_activation=True)
assert response.status_code == 400

View File

@@ -48,7 +48,7 @@ def test_model_loading_with_params(vllm_runner, monkeypatch):
# asserts on the pooling config files
assert model_config.pooler_config.seq_pooling_type == "CLS"
assert model_config.pooler_config.tok_pooling_type == "ALL"
assert model_config.pooler_config.normalize
assert model_config.pooler_config.use_activation
# asserts on the tokenizer loaded
assert model_config.tokenizer == "BAAI/bge-base-en-v1.5"
@@ -93,7 +93,7 @@ def test_roberta_model_loading_with_params(vllm_runner, monkeypatch):
# asserts on the pooling config files
assert model_config.pooler_config.seq_pooling_type == "MEAN"
assert model_config.pooler_config.tok_pooling_type == "ALL"
assert model_config.pooler_config.normalize
assert model_config.pooler_config.use_activation
# asserts on the tokenizer loaded
assert model_config.tokenizer == "intfloat/multilingual-e5-base"

View File

@@ -66,7 +66,7 @@ def test_embed_models_using_normalize(
model,
max_model_len=512,
dtype=dtype,
pooler_config=PoolerConfig(normalize=False),
pooler_config=PoolerConfig(use_activation=False),
) as vllm_model:
wo_normalize = torch.tensor(vllm_model.embed(example_prompts))
@@ -74,7 +74,7 @@ def test_embed_models_using_normalize(
model,
max_model_len=512,
dtype=dtype,
pooler_config=PoolerConfig(normalize=True),
pooler_config=PoolerConfig(use_activation=True),
) as vllm_model:
w_normalize = torch.tensor(vllm_model.embed(example_prompts))
@@ -146,7 +146,7 @@ def test_multi_vector_retrieval_models_using_normalize(
model,
max_model_len=512,
dtype=dtype,
pooler_config=PoolerConfig(normalize=False),
pooler_config=PoolerConfig(use_activation=False),
) as vllm_model:
wo_normalize = vllm_model.token_embed(example_prompts)
@@ -154,7 +154,7 @@ def test_multi_vector_retrieval_models_using_normalize(
model,
max_model_len=512,
dtype=dtype,
pooler_config=PoolerConfig(normalize=True),
pooler_config=PoolerConfig(use_activation=True),
) as vllm_model:
w_normalize = vllm_model.token_embed(example_prompts)

View File

@@ -160,7 +160,7 @@ def test_get_pooling_config():
model_config = ModelConfig(model_id)
assert model_config.pooler_config is not None
assert model_config.pooler_config.normalize
assert model_config.pooler_config.use_activation
assert model_config.pooler_config.seq_pooling_type == "MEAN"
assert model_config.pooler_config.tok_pooling_type == "ALL"

View File

@@ -18,7 +18,7 @@ EMBEDDING_MODELS = [
]
classify_parameters = ["use_activation"]
embed_parameters = ["dimensions", "normalize"]
embed_parameters = ["dimensions", "use_activation"]
step_pooling_parameters = ["step_tag_id", "returned_token_ids"]
@@ -42,17 +42,17 @@ def test_embed():
task = "embed"
model_config = MockModelConfig(pooler_config=PoolerConfig(seq_pooling_type="CLS"))
pooling_params = PoolingParams(normalize=None)
pooling_params = PoolingParams(use_activation=None)
pooling_params.verify(task=task, model_config=model_config)
pooling_params = PoolingParams(normalize=True)
pooling_params = PoolingParams(use_activation=True)
pooling_params.verify(task=task, model_config=model_config)
pooling_params = PoolingParams(normalize=False)
pooling_params = PoolingParams(use_activation=False)
pooling_params.verify(task=task, model_config=model_config)
invalid_parameters = classify_parameters + step_pooling_parameters
for p in invalid_parameters:
for p in set(invalid_parameters) - set(embed_parameters):
with pytest.raises(ValueError):
pooling_params = PoolingParams(**{p: True})
pooling_params.verify(task=task, model_config=model_config)
@@ -98,7 +98,7 @@ def test_classify(task):
pooling_params.verify(task=task, model_config=model_config)
invalid_parameters = embed_parameters + step_pooling_parameters
for p in invalid_parameters:
for p in set(invalid_parameters) - set(classify_parameters):
with pytest.raises(ValueError):
pooling_params = PoolingParams(**{p: True})
pooling_params.verify(task=task, model_config=model_config)
@@ -111,20 +111,20 @@ def test_token_embed(pooling_type: str):
pooler_config=PoolerConfig(tok_pooling_type=pooling_type)
)
pooling_params = PoolingParams(normalize=None)
pooling_params = PoolingParams(use_activation=None)
pooling_params.verify(task=task, model_config=model_config)
pooling_params = PoolingParams(normalize=True)
pooling_params = PoolingParams(use_activation=True)
pooling_params.verify(task=task, model_config=model_config)
pooling_params = PoolingParams(normalize=False)
pooling_params = PoolingParams(use_activation=False)
pooling_params.verify(task=task, model_config=model_config)
invalid_parameters = classify_parameters
if pooling_type != "STEP":
invalid_parameters = classify_parameters + step_pooling_parameters
for p in invalid_parameters:
for p in set(invalid_parameters) - set(embed_parameters):
with pytest.raises(ValueError):
pooling_params = PoolingParams(**{p: True})
pooling_params.verify(task=task, model_config=model_config)
@@ -150,7 +150,7 @@ def test_token_classify(pooling_type: str):
if pooling_type != "STEP":
invalid_parameters = embed_parameters + step_pooling_parameters
for p in invalid_parameters:
for p in set(invalid_parameters) - set(classify_parameters):
with pytest.raises(ValueError):
pooling_params = PoolingParams(**{p: True})
pooling_params.verify(task=task, model_config=model_config)