[CICD] [prefill-only] Consolidate prefill-only model E2E tests (#22405)
This commit is contained in:
Sundara Raman Ramachandran
@@ -1,604 +0,0 @@
|
||||
import unittest
|
||||
|
||||
import torch
|
||||
from transformers import AutoModelForCausalLM, AutoTokenizer
|
||||
|
||||
from sglang.srt.entrypoints.engine import Engine
|
||||
from sglang.test.ci.ci_register import register_cuda_ci
|
||||
from sglang.test.test_utils import DEFAULT_SMALL_MODEL_NAME_FOR_TEST, CustomTestCase
|
||||
|
||||
register_cuda_ci(est_time=260, suite="stage-b-test-1-gpu-large")
|
||||
|
||||
TEST_MODEL_NAME = DEFAULT_SMALL_MODEL_NAME_FOR_TEST
|
||||
|
||||
|
||||
class TestScoreAPI(CustomTestCase):
|
||||
"""Test the scoring API functionality."""
|
||||
|
||||
def setUp(self):
|
||||
"""Set up each test case."""
|
||||
self.engine = Engine(model_path=TEST_MODEL_NAME)
|
||||
|
||||
def tearDown(self):
|
||||
"""Clean up after each test case."""
|
||||
if self.engine is not None:
|
||||
self.engine.shutdown()
|
||||
torch.cuda.empty_cache()
|
||||
|
||||
def compute_hf_scores(
|
||||
self, query, items, label_token_ids, apply_softmax=False, item_first=False
|
||||
):
|
||||
"""Compute scores using direct HuggingFace model inference.
|
||||
Returns probabilities for each token ID, optionally normalized with softmax.
|
||||
|
||||
Args:
|
||||
query: The query text
|
||||
items: List of item texts
|
||||
label_token_ids: List of token IDs to compute probabilities for
|
||||
apply_softmax: Whether to normalize probabilities using softmax
|
||||
item_first: If True, prepend items to query. Otherwise append items to query.
|
||||
"""
|
||||
# Initialize HF model and tokenizer
|
||||
tokenizer = AutoTokenizer.from_pretrained(
|
||||
TEST_MODEL_NAME, trust_remote_code=True
|
||||
)
|
||||
model = AutoModelForCausalLM.from_pretrained(
|
||||
TEST_MODEL_NAME, trust_remote_code=True
|
||||
)
|
||||
|
||||
try:
|
||||
scores = []
|
||||
for item in items:
|
||||
# Construct full text based on item_first parameter
|
||||
full_text = f"{item}{query}" if item_first else f"{query}{item}"
|
||||
inputs = tokenizer(full_text, return_tensors="pt").to(model.device)
|
||||
|
||||
# Get logits for the last token
|
||||
with torch.no_grad():
|
||||
outputs = model(**inputs)
|
||||
last_token_logits = outputs.logits[0, -1]
|
||||
|
||||
# Get logits for just our target tokens
|
||||
target_logits = last_token_logits[label_token_ids]
|
||||
|
||||
# Apply softmax over just the target tokens
|
||||
target_probs = torch.softmax(target_logits, dim=-1)
|
||||
|
||||
# Convert to list of probabilities in order of label_token_ids
|
||||
probs = [target_probs[i].item() for i in range(len(label_token_ids))]
|
||||
|
||||
scores.append(probs)
|
||||
|
||||
return scores
|
||||
finally:
|
||||
# Clean up HF resources
|
||||
model.cpu()
|
||||
del model
|
||||
del tokenizer
|
||||
torch.cuda.empty_cache()
|
||||
|
||||
def _get_token_ids(self, tokens):
|
||||
"""Helper method to get token IDs for a list of tokens."""
|
||||
tokenizer = AutoTokenizer.from_pretrained(
|
||||
TEST_MODEL_NAME, trust_remote_code=True
|
||||
)
|
||||
try:
|
||||
label_token_ids = []
|
||||
for token in tokens:
|
||||
encoding = tokenizer(token, add_special_tokens=False)
|
||||
token_ids = encoding["input_ids"]
|
||||
label_token_ids.append(token_ids[0])
|
||||
return label_token_ids
|
||||
finally:
|
||||
del tokenizer
|
||||
|
||||
def _compare_scores(self, hf_scores, sglang_scores, label_token_ids, case_name=""):
|
||||
"""Helper method to compare scores between HF and SGLang using relative tolerance."""
|
||||
self.assertEqual(
|
||||
len(hf_scores),
|
||||
len(sglang_scores),
|
||||
f"Score lengths don't match for {case_name}",
|
||||
)
|
||||
|
||||
# Use a relative tolerance of 1%
|
||||
TOLERANCE = 0.01
|
||||
|
||||
for hf_score_list, sglang_score_list in zip(hf_scores, sglang_scores):
|
||||
self.assertEqual(
|
||||
len(hf_score_list),
|
||||
len(sglang_score_list),
|
||||
f"Score list lengths don't match for {case_name}",
|
||||
)
|
||||
|
||||
for hf_score, sglang_score in zip(hf_score_list, sglang_score_list):
|
||||
diff = abs(hf_score - sglang_score)
|
||||
self.assertLessEqual(
|
||||
diff,
|
||||
TOLERANCE,
|
||||
msg=f"Scores differ by {diff:.2%} ({case_name}): "
|
||||
f"HF={hf_score:.6f}, SGLang={sglang_score:.6f}",
|
||||
)
|
||||
|
||||
self.assertGreaterEqual(
|
||||
sglang_score, 0, f"SGLang score {sglang_score:.6f} not in [0,1]"
|
||||
)
|
||||
self.assertLessEqual(
|
||||
sglang_score, 1, f"SGLang score {sglang_score:.6f} not in [0,1]"
|
||||
)
|
||||
|
||||
self.assertAlmostEqual(
|
||||
sum(sglang_score_list),
|
||||
1.0,
|
||||
places=6,
|
||||
msg=f"SGLang scores don't sum to 1 ({case_name}): {sum(sglang_score_list):.6f}",
|
||||
)
|
||||
|
||||
def test_score_consistency(self):
|
||||
"""Test that SGLang scoring matches direct HuggingFace model scoring."""
|
||||
# Define test cases
|
||||
test_cases = [
|
||||
{
|
||||
"name": "default case",
|
||||
"query": "I pledge allegiance",
|
||||
"items": ["", " to"],
|
||||
"item_first": False,
|
||||
},
|
||||
{
|
||||
"name": "item_first case",
|
||||
"query": " is a city",
|
||||
"items": ["Tokyo", "Japan"],
|
||||
"item_first": True,
|
||||
},
|
||||
]
|
||||
|
||||
# Common tokens to test for all cases
|
||||
tokens = [" to", " the"]
|
||||
label_token_ids = self._get_token_ids(tokens)
|
||||
|
||||
# Run each test case
|
||||
for case in test_cases:
|
||||
# Get scores from SGLang
|
||||
sglang_scores = self.engine.score(
|
||||
query=case["query"],
|
||||
items=case["items"],
|
||||
label_token_ids=label_token_ids,
|
||||
apply_softmax=True,
|
||||
item_first=case["item_first"],
|
||||
).scores
|
||||
|
||||
# Get scores from HuggingFace using the same parameters
|
||||
hf_scores = self.compute_hf_scores(
|
||||
query=case["query"],
|
||||
items=case["items"],
|
||||
label_token_ids=label_token_ids,
|
||||
apply_softmax=True,
|
||||
item_first=case["item_first"],
|
||||
)
|
||||
|
||||
# Compare scores
|
||||
self._compare_scores(
|
||||
hf_scores, sglang_scores, label_token_ids, case["name"]
|
||||
)
|
||||
|
||||
def test_score_batch_handling(self):
|
||||
"""Test that batch scoring works correctly."""
|
||||
# Test with different batch sizes
|
||||
batch_sizes = [1, 2, 4, 8]
|
||||
label_token_ids = [1, 2, 3]
|
||||
|
||||
for batch_size in batch_sizes:
|
||||
texts = [f"test {i}" for i in range(batch_size)]
|
||||
scores = self.engine.score(
|
||||
query="The test was",
|
||||
items=texts,
|
||||
label_token_ids=label_token_ids,
|
||||
apply_softmax=True,
|
||||
).scores
|
||||
|
||||
self.assertEqual(
|
||||
len(scores),
|
||||
batch_size,
|
||||
f"Expected {batch_size} scores, got {len(scores)}",
|
||||
)
|
||||
|
||||
# Verify each score list has the correct length
|
||||
for score_list in scores:
|
||||
self.assertEqual(
|
||||
len(score_list),
|
||||
len(label_token_ids),
|
||||
f"Score list length {len(score_list)} doesn't match label_token_ids length {len(label_token_ids)}",
|
||||
)
|
||||
self.assertTrue(
|
||||
all(isinstance(v, float) for v in score_list),
|
||||
"All scores should be floats",
|
||||
)
|
||||
self.assertAlmostEqual(
|
||||
1.0, sum(score_list), 6, "Scores should sum to 1"
|
||||
)
|
||||
|
||||
def test_score_request_construction(self):
|
||||
"""Test that scoring requests are constructed to avoid decode phase."""
|
||||
from unittest.mock import patch
|
||||
|
||||
# Capture the internal request to verify optimization
|
||||
captured_requests = []
|
||||
original_gen = self.engine.tokenizer_manager.generate_request
|
||||
|
||||
async def mock_generate_request(req, request=None):
|
||||
captured_requests.append(req)
|
||||
async for result in original_gen(req, request):
|
||||
yield result
|
||||
|
||||
# Patch the generate_request method
|
||||
with patch.object(
|
||||
self.engine.tokenizer_manager,
|
||||
"generate_request",
|
||||
side_effect=mock_generate_request,
|
||||
):
|
||||
# Run a scoring request
|
||||
query = "What is the capital of"
|
||||
items = ["France", "Germany"]
|
||||
label_token_ids = [1, 2, 3]
|
||||
|
||||
scores = self.engine.score(
|
||||
query=query,
|
||||
items=items,
|
||||
label_token_ids=label_token_ids,
|
||||
apply_softmax=True,
|
||||
).scores
|
||||
|
||||
# Verify we got results
|
||||
self.assertEqual(len(scores), len(items))
|
||||
|
||||
# Verify the captured request has decode-avoiding properties
|
||||
self.assertEqual(len(captured_requests), 1)
|
||||
request = captured_requests[0]
|
||||
|
||||
# Key assertions for decode phase avoidance:
|
||||
# 1. max_new_tokens should be 0 (prevents token generation)
|
||||
# Handle both single and batch request cases
|
||||
if isinstance(request.sampling_params, dict):
|
||||
max_new_tokens = request.sampling_params.get("max_new_tokens", 0)
|
||||
elif isinstance(request.sampling_params, list):
|
||||
# For batch requests, check the first item
|
||||
max_new_tokens = request.sampling_params[0].get("max_new_tokens", 0)
|
||||
else:
|
||||
max_new_tokens = getattr(request.sampling_params, "max_new_tokens", 0)
|
||||
|
||||
self.assertEqual(
|
||||
max_new_tokens, 0, "max_new_tokens should be 0 to avoid decode phase"
|
||||
)
|
||||
|
||||
# 2. Should have token_ids_logprob for scoring
|
||||
# Handle both single and batch request cases
|
||||
if (
|
||||
isinstance(request.token_ids_logprob, list)
|
||||
and len(request.token_ids_logprob) > 0
|
||||
and isinstance(request.token_ids_logprob[0], list)
|
||||
):
|
||||
# Batch case: token_ids_logprob is a list of lists
|
||||
# Each item in the batch should have the same label_token_ids
|
||||
for item_token_ids in request.token_ids_logprob:
|
||||
self.assertEqual(
|
||||
item_token_ids,
|
||||
label_token_ids,
|
||||
"Each batch item should have label_token_ids for scoring",
|
||||
)
|
||||
else:
|
||||
# Single request case
|
||||
self.assertEqual(
|
||||
request.token_ids_logprob,
|
||||
label_token_ids,
|
||||
"Should have label_token_ids for scoring",
|
||||
)
|
||||
|
||||
# 3. Should request logprobs but not stream
|
||||
self.assertTrue(
|
||||
request.return_logprob, "Should request logprobs for scoring"
|
||||
)
|
||||
self.assertFalse(request.stream, "Scoring requests should not stream")
|
||||
|
||||
def test_multi_item_scoring_basic(self):
|
||||
"""Test basic multi-item scoring functionality."""
|
||||
# Test with a simple query and items
|
||||
query = "What is the capital of California? Answer Yes or No for each of the following options:"
|
||||
items = ["Sacramento", "San Jose", "San Francisco"]
|
||||
label_token_ids = [9454, 2753] # "Yes" and "No" tokens
|
||||
|
||||
# Get scores using SGLang
|
||||
result = self.engine.score(
|
||||
query=query,
|
||||
items=items,
|
||||
label_token_ids=label_token_ids,
|
||||
apply_softmax=True,
|
||||
)
|
||||
scores = result.scores
|
||||
prompt_tokens = result.prompt_tokens
|
||||
|
||||
# Verify we get the expected number of scores
|
||||
self.assertEqual(len(scores), len(items), "Should get one score list per item")
|
||||
self.assertGreater(prompt_tokens, 0, "Should have positive prompt_tokens")
|
||||
|
||||
# Verify each score list has the correct length
|
||||
for i, score_list in enumerate(scores):
|
||||
self.assertEqual(
|
||||
len(score_list),
|
||||
len(label_token_ids),
|
||||
f"Item {i} should have {len(label_token_ids)} scores",
|
||||
)
|
||||
# Verify scores are probabilities (sum to 1)
|
||||
self.assertAlmostEqual(
|
||||
sum(score_list),
|
||||
1.0,
|
||||
places=6,
|
||||
msg=f"Scores for item {i} should sum to 1",
|
||||
)
|
||||
# Verify all scores are non-negative
|
||||
for j, score in enumerate(score_list):
|
||||
self.assertGreaterEqual(
|
||||
score, 0, f"Score {j} for item {i} should be non-negative"
|
||||
)
|
||||
|
||||
def test_multi_item_scoring_consistency(self):
|
||||
"""Test that multi-item scoring gives consistent results."""
|
||||
query = "Choose the best option:"
|
||||
items = ["Option A", "Option B", "Option C"]
|
||||
label_token_ids = [1, 2, 3]
|
||||
|
||||
# Run the same test multiple times
|
||||
scores1 = self.engine.score(
|
||||
query=query,
|
||||
items=items,
|
||||
label_token_ids=label_token_ids,
|
||||
apply_softmax=True,
|
||||
).scores
|
||||
|
||||
scores2 = self.engine.score(
|
||||
query=query,
|
||||
items=items,
|
||||
label_token_ids=label_token_ids,
|
||||
apply_softmax=True,
|
||||
).scores
|
||||
|
||||
# Results should be identical (deterministic)
|
||||
self.assertEqual(len(scores1), len(scores2), "Should get same number of items")
|
||||
for i, (s1, s2) in enumerate(zip(scores1, scores2)):
|
||||
self.assertEqual(
|
||||
len(s1), len(s2), f"Item {i} should have same number of scores"
|
||||
)
|
||||
for j, (score1, score2) in enumerate(zip(s1, s2)):
|
||||
self.assertAlmostEqual(
|
||||
score1,
|
||||
score2,
|
||||
places=6,
|
||||
msg=f"Score {j} for item {i} should be identical",
|
||||
)
|
||||
|
||||
def test_multi_item_scoring_different_sizes(self):
|
||||
"""Test multi-item scoring with different numbers of items."""
|
||||
query = "Rate each option:"
|
||||
label_token_ids = [1, 2, 3, 4, 5]
|
||||
|
||||
# Test with different numbers of items
|
||||
test_cases = [
|
||||
["Single item"],
|
||||
["Item 1", "Item 2"],
|
||||
["A", "B", "C", "D"],
|
||||
["X", "Y", "Z", "W", "V", "U"],
|
||||
]
|
||||
|
||||
for items in test_cases:
|
||||
with self.subTest(items=items):
|
||||
scores = self.engine.score(
|
||||
query=query,
|
||||
items=items,
|
||||
label_token_ids=label_token_ids,
|
||||
apply_softmax=True,
|
||||
).scores
|
||||
|
||||
self.assertEqual(
|
||||
len(scores), len(items), f"Should get {len(items)} score lists"
|
||||
)
|
||||
|
||||
for i, score_list in enumerate(scores):
|
||||
self.assertEqual(
|
||||
len(score_list),
|
||||
len(label_token_ids),
|
||||
f"Item {i} should have {len(label_token_ids)} scores",
|
||||
)
|
||||
self.assertAlmostEqual(sum(score_list), 1.0, places=6)
|
||||
|
||||
def test_multi_item_scoring_empty_items(self):
|
||||
"""Test multi-item scoring with empty items list."""
|
||||
query = "Test query"
|
||||
items = []
|
||||
label_token_ids = [1, 2]
|
||||
|
||||
result = self.engine.score(
|
||||
query=query,
|
||||
items=items,
|
||||
label_token_ids=label_token_ids,
|
||||
apply_softmax=True,
|
||||
)
|
||||
scores = result.scores
|
||||
prompt_tokens = result.prompt_tokens
|
||||
|
||||
self.assertEqual(len(scores), 0, "Should return empty list for empty items")
|
||||
self.assertEqual(
|
||||
prompt_tokens, 0, "Should return 0 prompt_tokens for empty items"
|
||||
)
|
||||
|
||||
def test_multi_item_scoring_single_item(self):
|
||||
"""Test multi-item scoring with single item (should work like regular scoring)."""
|
||||
query = "Complete this sentence: The capital of France is"
|
||||
items = ["Paris"]
|
||||
label_token_ids = [1, 2, 3]
|
||||
|
||||
result = self.engine.score(
|
||||
query=query,
|
||||
items=items,
|
||||
label_token_ids=label_token_ids,
|
||||
apply_softmax=True,
|
||||
)
|
||||
scores = result.scores
|
||||
prompt_tokens = result.prompt_tokens
|
||||
|
||||
self.assertEqual(len(scores), 1, "Should get one score list")
|
||||
self.assertEqual(
|
||||
len(scores[0]), len(label_token_ids), "Should have correct number of scores"
|
||||
)
|
||||
self.assertAlmostEqual(sum(scores[0]), 1.0, places=6)
|
||||
self.assertGreater(prompt_tokens, 0, "Should have positive prompt_tokens")
|
||||
|
||||
def test_multi_item_scoring_different_queries(self):
|
||||
"""Test multi-item scoring with different types of queries."""
|
||||
items = ["Yes", "No"]
|
||||
label_token_ids = [1, 2]
|
||||
|
||||
test_queries = [
|
||||
"Is this true?",
|
||||
"Choose the correct answer:",
|
||||
"What is the best option?",
|
||||
"Select all that apply:",
|
||||
"", # Empty query
|
||||
]
|
||||
|
||||
for query in test_queries:
|
||||
with self.subTest(query=query):
|
||||
scores = self.engine.score(
|
||||
query=query,
|
||||
items=items,
|
||||
label_token_ids=label_token_ids,
|
||||
apply_softmax=True,
|
||||
).scores
|
||||
|
||||
self.assertEqual(
|
||||
len(scores),
|
||||
len(items),
|
||||
f"Should get {len(items)} score lists for query: '{query}'",
|
||||
)
|
||||
|
||||
for i, score_list in enumerate(scores):
|
||||
self.assertEqual(len(score_list), len(label_token_ids))
|
||||
self.assertAlmostEqual(sum(score_list), 1.0, places=6)
|
||||
|
||||
def test_multi_item_scoring_different_label_tokens(self):
|
||||
"""Test multi-item scoring with different label token sets."""
|
||||
query = "Choose the best option:"
|
||||
items = ["Option A", "Option B"]
|
||||
|
||||
test_label_tokens = [
|
||||
[1, 2], # Two tokens
|
||||
[1, 2, 3, 4], # Four tokens
|
||||
[1], # Single token
|
||||
[1, 2, 3, 4, 5, 6, 7, 8], # Many tokens
|
||||
]
|
||||
|
||||
for label_token_ids in test_label_tokens:
|
||||
with self.subTest(label_tokens=label_token_ids):
|
||||
scores = self.engine.score(
|
||||
query=query,
|
||||
items=items,
|
||||
label_token_ids=label_token_ids,
|
||||
apply_softmax=True,
|
||||
).scores
|
||||
|
||||
self.assertEqual(len(scores), len(items))
|
||||
|
||||
for i, score_list in enumerate(scores):
|
||||
self.assertEqual(
|
||||
len(score_list),
|
||||
len(label_token_ids),
|
||||
f"Item {i} should have {len(label_token_ids)} scores",
|
||||
)
|
||||
self.assertAlmostEqual(sum(score_list), 1.0, places=6)
|
||||
|
||||
def test_multi_item_scoring_without_softmax(self):
|
||||
"""Test multi-item scoring without softmax normalization."""
|
||||
query = "Rate each option:"
|
||||
items = ["Good", "Bad", "Neutral"]
|
||||
label_token_ids = [1, 2, 3]
|
||||
|
||||
scores = self.engine.score(
|
||||
query=query,
|
||||
items=items,
|
||||
label_token_ids=label_token_ids,
|
||||
apply_softmax=False, # No softmax
|
||||
).scores
|
||||
|
||||
self.assertEqual(len(scores), len(items))
|
||||
|
||||
for i, score_list in enumerate(scores):
|
||||
self.assertEqual(len(score_list), len(label_token_ids))
|
||||
# Without softmax, scores don't need to sum to 1
|
||||
# But they should still be valid logits/probabilities
|
||||
for j, score in enumerate(score_list):
|
||||
self.assertIsInstance(
|
||||
score, (int, float), f"Score {j} for item {i} should be numeric"
|
||||
)
|
||||
|
||||
def test_multi_item_scoring_large_batch(self):
|
||||
"""Test multi-item scoring with a large number of items."""
|
||||
query = "Classify each item:"
|
||||
items = [f"Item {i}" for i in range(20)] # 20 items
|
||||
label_token_ids = [1, 2, 3]
|
||||
|
||||
scores = self.engine.score(
|
||||
query=query,
|
||||
items=items,
|
||||
label_token_ids=label_token_ids,
|
||||
apply_softmax=True,
|
||||
).scores
|
||||
|
||||
self.assertEqual(len(scores), len(items), "Should handle large batches")
|
||||
|
||||
for i, score_list in enumerate(scores):
|
||||
self.assertEqual(len(score_list), len(label_token_ids))
|
||||
self.assertAlmostEqual(sum(score_list), 1.0, places=6)
|
||||
|
||||
def test_multi_item_scoring_unicode(self):
|
||||
"""Test multi-item scoring with unicode characters."""
|
||||
query = "选择最佳选项："
|
||||
items = ["选项A", "选项B", "选项C"]
|
||||
label_token_ids = [1, 2, 3]
|
||||
|
||||
scores = self.engine.score(
|
||||
query=query,
|
||||
items=items,
|
||||
label_token_ids=label_token_ids,
|
||||
apply_softmax=True,
|
||||
).scores
|
||||
|
||||
self.assertEqual(len(scores), len(items))
|
||||
|
||||
for i, score_list in enumerate(scores):
|
||||
self.assertEqual(len(score_list), len(label_token_ids))
|
||||
self.assertAlmostEqual(sum(score_list), 1.0, places=6)
|
||||
|
||||
def test_multi_item_scoring_error_handling(self):
|
||||
"""Test multi-item scoring error handling."""
|
||||
query = "Test query"
|
||||
items = ["Item 1", "Item 2"]
|
||||
label_token_ids = [1, 2]
|
||||
|
||||
# Test with invalid label_token_ids
|
||||
with self.assertRaises((ValueError, TypeError)):
|
||||
self.engine.score(
|
||||
query=query,
|
||||
items=items,
|
||||
label_token_ids="invalid", # Should be list of ints
|
||||
apply_softmax=True,
|
||||
)
|
||||
|
||||
# Test with None items
|
||||
with self.assertRaises((ValueError, TypeError)):
|
||||
self.engine.score(
|
||||
query=query,
|
||||
items=None,
|
||||
label_token_ids=label_token_ids,
|
||||
apply_softmax=True,
|
||||
)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
@@ -1,351 +0,0 @@
|
||||
"""Tests for Scoring API with SequenceClassification models.
|
||||
|
||||
Covers both single-item and multi-item scoring (MIS) for classification
|
||||
models. Uses json_model_override_args to override a regular Qwen3 model's
|
||||
architecture to Qwen3ForSequenceClassification so we can validate the
|
||||
scoring pipeline without needing a dedicated classification checkpoint
|
||||
(the score head gets randomly initialised, which is fine for shape /
|
||||
pipeline correctness).
|
||||
"""
|
||||
|
||||
import json
|
||||
import unittest
|
||||
|
||||
import torch
|
||||
|
||||
from sglang.srt.entrypoints.engine import Engine
|
||||
from sglang.test.ci.ci_register import register_cuda_ci
|
||||
from sglang.test.test_utils import CustomTestCase
|
||||
|
||||
register_cuda_ci(est_time=120, suite="stage-b-test-1-gpu-small")
|
||||
|
||||
# A lightweight Qwen3 checkpoint whose backbone weights load cleanly into
|
||||
# Qwen3ForSequenceClassification (the classification head is random).
|
||||
TEST_BASE_MODEL = "Qwen/Qwen3-0.6B"
|
||||
# <|endoftext|> for Qwen3 tokenizer — used as MIS delimiter
|
||||
QWEN3_ENDOFTEXT_TOKEN_ID = 151643
|
||||
|
||||
|
||||
class TestScoreClassification(CustomTestCase):
|
||||
"""Single-item scoring with a SequenceClassification model."""
|
||||
|
||||
NUM_LABELS = 2
|
||||
|
||||
@classmethod
|
||||
def setUpClass(cls):
|
||||
override_args = json.dumps(
|
||||
{
|
||||
"architectures": ["Qwen3ForSequenceClassification"],
|
||||
"num_labels": cls.NUM_LABELS,
|
||||
}
|
||||
)
|
||||
cls.engine = Engine(
|
||||
model_path=TEST_BASE_MODEL,
|
||||
disable_radix_cache=True,
|
||||
json_model_override_args=override_args,
|
||||
mem_fraction_static=0.15,
|
||||
)
|
||||
|
||||
@classmethod
|
||||
def tearDownClass(cls):
|
||||
if hasattr(cls, "engine") and cls.engine is not None:
|
||||
cls.engine.shutdown()
|
||||
torch.cuda.empty_cache()
|
||||
|
||||
def test_basic_single_item(self):
|
||||
"""Each item gets a score vector of length num_labels."""
|
||||
scores = self.engine.score(
|
||||
query="Rate each option:",
|
||||
items=["Option A", "Option B"],
|
||||
apply_softmax=True,
|
||||
).scores
|
||||
|
||||
self.assertEqual(len(scores), 2)
|
||||
for i, score_list in enumerate(scores):
|
||||
self.assertEqual(len(score_list), self.NUM_LABELS)
|
||||
self.assertAlmostEqual(
|
||||
sum(score_list),
|
||||
1.0,
|
||||
places=5,
|
||||
msg=f"Softmax scores for item {i} should sum to 1",
|
||||
)
|
||||
for val in score_list:
|
||||
self.assertGreaterEqual(val, 0.0)
|
||||
self.assertLessEqual(val, 1.0)
|
||||
|
||||
def test_single_item_edge_case(self):
|
||||
"""Single item in the list."""
|
||||
scores = self.engine.score(
|
||||
query="Evaluate:",
|
||||
items=["Only item"],
|
||||
apply_softmax=True,
|
||||
).scores
|
||||
|
||||
self.assertEqual(len(scores), 1)
|
||||
self.assertEqual(len(scores[0]), self.NUM_LABELS)
|
||||
self.assertAlmostEqual(sum(scores[0]), 1.0, places=5)
|
||||
|
||||
def test_raw_logits_without_softmax(self):
|
||||
"""Without softmax, returns raw logits (no probability constraints)."""
|
||||
scores = self.engine.score(
|
||||
query="Evaluate:",
|
||||
items=["Alpha", "Beta"],
|
||||
apply_softmax=False,
|
||||
).scores
|
||||
|
||||
self.assertEqual(len(scores), 2)
|
||||
for score_list in scores:
|
||||
self.assertEqual(len(score_list), self.NUM_LABELS)
|
||||
for val in score_list:
|
||||
self.assertTrue(
|
||||
isinstance(val, (int, float)),
|
||||
f"Expected numeric score, got {type(val)}",
|
||||
)
|
||||
|
||||
def test_deterministic(self):
|
||||
"""Identical inputs yield near-identical scores (fp16 non-determinism allowed)."""
|
||||
kwargs = dict(query="Evaluate:", items=["alpha", "beta", "gamma"])
|
||||
|
||||
scores1 = self.engine.score(**kwargs).scores
|
||||
scores2 = self.engine.score(**kwargs).scores
|
||||
|
||||
self.assertEqual(len(scores1), len(scores2))
|
||||
for s1, s2 in zip(scores1, scores2):
|
||||
for v1, v2 in zip(s1, s2):
|
||||
self.assertAlmostEqual(v1, v2, places=1)
|
||||
|
||||
def test_tokenized_inputs(self):
|
||||
"""Pre-tokenized query and items work the same as text inputs."""
|
||||
from transformers import AutoTokenizer
|
||||
|
||||
tokenizer = AutoTokenizer.from_pretrained(TEST_BASE_MODEL)
|
||||
query_text = "Rate this:"
|
||||
items_text = ["Good", "Bad"]
|
||||
|
||||
text_scores = self.engine.score(
|
||||
query=query_text,
|
||||
items=items_text,
|
||||
apply_softmax=True,
|
||||
).scores
|
||||
|
||||
query_ids = tokenizer.encode(query_text)
|
||||
items_ids = [tokenizer.encode(item) for item in items_text]
|
||||
token_scores = self.engine.score(
|
||||
query=query_ids,
|
||||
items=items_ids,
|
||||
apply_softmax=True,
|
||||
).scores
|
||||
|
||||
self.assertEqual(len(text_scores), len(token_scores))
|
||||
for txt_s, tok_s in zip(text_scores, token_scores):
|
||||
for t, k in zip(txt_s, tok_s):
|
||||
self.assertAlmostEqual(t, k, places=4)
|
||||
|
||||
def test_label_token_ids_ignored(self):
|
||||
"""SequenceClassification models ignore label_token_ids (no crash)."""
|
||||
scores = self.engine.score(
|
||||
query="Evaluate:",
|
||||
items=["Test item"],
|
||||
label_token_ids=[1, 2, 3],
|
||||
apply_softmax=True,
|
||||
).scores
|
||||
|
||||
self.assertEqual(len(scores), 1)
|
||||
self.assertEqual(len(scores[0]), self.NUM_LABELS)
|
||||
|
||||
|
||||
class TestScoreClassificationMIS(CustomTestCase):
|
||||
"""Multi-item scoring (MIS) with a SequenceClassification model.
|
||||
|
||||
MIS packs all items into one sequence separated by a delimiter token.
|
||||
The score_and_pool function extracts per-item scores at delimiter
|
||||
positions.
|
||||
"""
|
||||
|
||||
NUM_LABELS = 2
|
||||
|
||||
@classmethod
|
||||
def setUpClass(cls):
|
||||
override_args = json.dumps(
|
||||
{
|
||||
"architectures": ["Qwen3ForSequenceClassification"],
|
||||
"num_labels": cls.NUM_LABELS,
|
||||
}
|
||||
)
|
||||
cls.engine = Engine(
|
||||
model_path=TEST_BASE_MODEL,
|
||||
disable_radix_cache=True,
|
||||
chunked_prefill_size=-1,
|
||||
multi_item_scoring_delimiter=QWEN3_ENDOFTEXT_TOKEN_ID,
|
||||
json_model_override_args=override_args,
|
||||
mem_fraction_static=0.15,
|
||||
)
|
||||
|
||||
@classmethod
|
||||
def tearDownClass(cls):
|
||||
if hasattr(cls, "engine") and cls.engine is not None:
|
||||
cls.engine.shutdown()
|
||||
torch.cuda.empty_cache()
|
||||
|
||||
def test_mis_basic(self):
|
||||
"""MIS produces one score vector per item."""
|
||||
items = ["Option A", "Option B", "Option C"]
|
||||
scores = self.engine.score(
|
||||
query="Rate each option:",
|
||||
items=items,
|
||||
apply_softmax=True,
|
||||
).scores
|
||||
|
||||
self.assertEqual(len(scores), len(items))
|
||||
for i, score_list in enumerate(scores):
|
||||
self.assertEqual(
|
||||
len(score_list),
|
||||
self.NUM_LABELS,
|
||||
f"Item {i} should have {self.NUM_LABELS} scores",
|
||||
)
|
||||
self.assertAlmostEqual(
|
||||
sum(score_list),
|
||||
1.0,
|
||||
places=5,
|
||||
msg=f"Scores for item {i} should sum to 1",
|
||||
)
|
||||
for val in score_list:
|
||||
self.assertGreaterEqual(val, 0.0)
|
||||
self.assertLessEqual(val, 1.0)
|
||||
|
||||
def test_mis_many_items(self):
|
||||
"""Stress test: 10 items."""
|
||||
items = [f"Item {i}" for i in range(10)]
|
||||
scores = self.engine.score(
|
||||
query="Classify each:",
|
||||
items=items,
|
||||
apply_softmax=True,
|
||||
).scores
|
||||
|
||||
self.assertEqual(len(scores), len(items))
|
||||
for score_list in scores:
|
||||
self.assertEqual(len(score_list), self.NUM_LABELS)
|
||||
self.assertAlmostEqual(sum(score_list), 1.0, places=5)
|
||||
|
||||
def test_mis_single_item(self):
|
||||
"""Edge case: single item through MIS path."""
|
||||
scores = self.engine.score(
|
||||
query="Evaluate:",
|
||||
items=["Single item"],
|
||||
apply_softmax=True,
|
||||
).scores
|
||||
|
||||
self.assertEqual(len(scores), 1)
|
||||
self.assertEqual(len(scores[0]), self.NUM_LABELS)
|
||||
self.assertAlmostEqual(sum(scores[0]), 1.0, places=5)
|
||||
|
||||
def test_items_produce_distinct_scores(self):
|
||||
"""Different items must produce different score vectors.
|
||||
|
||||
Even with a randomly initialised classification head, different
|
||||
item texts produce different hidden states, so scores should
|
||||
differ. This catches bugs where all delimiter tokens share the
|
||||
same pooled representation.
|
||||
"""
|
||||
items = [
|
||||
"Option A is about cats",
|
||||
"Option B is about dogs",
|
||||
"Option C is about fish",
|
||||
]
|
||||
scores = self.engine.score(query="Rate each option:", items=items).scores
|
||||
|
||||
self.assertEqual(len(scores), len(items))
|
||||
all_identical = all(scores[0] == s for s in scores[1:])
|
||||
self.assertFalse(
|
||||
all_identical,
|
||||
f"All {len(items)} items returned identical scores — "
|
||||
f"MIS delimiter indexing is likely broken. Scores: {scores[0]}",
|
||||
)
|
||||
|
||||
def test_deterministic(self):
|
||||
"""Identical MIS requests return identical scores."""
|
||||
kwargs = dict(
|
||||
query="Evaluate:",
|
||||
items=["alpha", "beta", "gamma"],
|
||||
)
|
||||
scores1 = self.engine.score(**kwargs).scores
|
||||
scores2 = self.engine.score(**kwargs).scores
|
||||
|
||||
self.assertEqual(scores1, scores2)
|
||||
|
||||
def test_softmax_valid(self):
|
||||
"""With softmax, each item's scores form a valid probability distribution."""
|
||||
items = ["Option A", "Option B", "Option C"]
|
||||
scores = self.engine.score(
|
||||
query="Rate each option:",
|
||||
items=items,
|
||||
apply_softmax=True,
|
||||
).scores
|
||||
|
||||
for i, score_list in enumerate(scores):
|
||||
self.assertEqual(len(score_list), self.NUM_LABELS)
|
||||
for val in score_list:
|
||||
self.assertGreaterEqual(val, 0.0)
|
||||
self.assertLessEqual(val, 1.0)
|
||||
self.assertAlmostEqual(
|
||||
sum(score_list),
|
||||
1.0,
|
||||
places=6,
|
||||
msg=f"Softmax scores for item {i} don't sum to 1: {sum(score_list)}",
|
||||
)
|
||||
|
||||
|
||||
class TestScoreClassificationMISAdvanced(CustomTestCase):
|
||||
"""Advanced MIS tests with more labels to stress tensor shape handling."""
|
||||
|
||||
NUM_LABELS = 12
|
||||
|
||||
@classmethod
|
||||
def setUpClass(cls):
|
||||
override_args = json.dumps(
|
||||
{
|
||||
"architectures": ["Qwen3ForSequenceClassification"],
|
||||
"num_labels": cls.NUM_LABELS,
|
||||
}
|
||||
)
|
||||
cls.engine = Engine(
|
||||
model_path=TEST_BASE_MODEL,
|
||||
disable_radix_cache=True,
|
||||
chunked_prefill_size=-1,
|
||||
multi_item_scoring_delimiter=QWEN3_ENDOFTEXT_TOKEN_ID,
|
||||
json_model_override_args=override_args,
|
||||
mem_fraction_static=0.15,
|
||||
)
|
||||
|
||||
@classmethod
|
||||
def tearDownClass(cls):
|
||||
if hasattr(cls, "engine") and cls.engine is not None:
|
||||
cls.engine.shutdown()
|
||||
torch.cuda.empty_cache()
|
||||
|
||||
def test_many_labels_shape(self):
|
||||
"""Verify correct shape with many labels (catches 2D tensor bugs)."""
|
||||
items = [f"Item {i}" for i in range(5)]
|
||||
scores = self.engine.score(
|
||||
query="Classify:",
|
||||
items=items,
|
||||
apply_softmax=True,
|
||||
).scores
|
||||
|
||||
self.assertEqual(len(scores), len(items))
|
||||
for score_list in scores:
|
||||
self.assertEqual(len(score_list), self.NUM_LABELS)
|
||||
self.assertAlmostEqual(sum(score_list), 1.0, places=5)
|
||||
|
||||
def test_many_items_distinct(self):
|
||||
"""15 items should not all produce identical scores."""
|
||||
items = [f"City {i}" for i in range(15)]
|
||||
scores = self.engine.score(query="Classify each city:", items=items).scores
|
||||
|
||||
self.assertEqual(len(scores), len(items))
|
||||
unique_count = len({tuple(s) for s in scores})
|
||||
self.assertGreater(unique_count, 1, "All 15 items returned identical scores")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
Reference in New Issue
Block a user