Migrate sampling tests to test/registered/sampling/ (#16455)
This commit is contained in:
Alison Shao
@@ -0,0 +1,200 @@
|
||||
"""Test original log probability alignment between SGLang and Hugging Face.
|
||||
|
||||
This test suite verifies the correctness of the `origin_logprobs` output (temperature=1)
|
||||
and the `logprobs` output (temperature=0.5) in SGLang by comparing it against
|
||||
raw logit-based probabilities computed directly from a reference Hugging Face model.
|
||||
|
||||
The test covers the following scenarios:
|
||||
- Next-token prediction: Verifies that the log probability of the next token from
|
||||
SGLang matches the Hugging Face model.
|
||||
- Top-k logprobs: Ensures that the top-k original logprobs returned by SGLang are
|
||||
consistent with Hugging Face outputs.
|
||||
- Specified token IDs: Confirms that the original logprobs for specific token IDs
|
||||
match the values computed from Hugging Face logits.
|
||||
"""
|
||||
|
||||
import os
|
||||
|
||||
from sglang.test.ci.ci_register import register_amd_ci, register_cuda_ci
|
||||
|
||||
register_cuda_ci(est_time=41, suite="stage-b-test-small-1-gpu")
|
||||
register_amd_ci(est_time=60, suite="stage-b-test-small-1-gpu")
|
||||
import random
|
||||
import unittest
|
||||
|
||||
import torch
|
||||
import torch.nn.functional as F
|
||||
from transformers import AutoModelForCausalLM, AutoTokenizer
|
||||
|
||||
import sglang as sgl
|
||||
from sglang.test.test_utils import DEFAULT_SMALL_MODEL_NAME_FOR_TEST
|
||||
|
||||
# ------------------------- Configurable via env ------------------------- #
|
||||
MODEL_ID = DEFAULT_SMALL_MODEL_NAME_FOR_TEST
|
||||
PROMPTS = [
|
||||
"Hello, my name is",
|
||||
"The future of AI is",
|
||||
"The president of the United States is",
|
||||
"The capital of France is ",
|
||||
]
|
||||
TOP_LOGPROBS_NUM = 50
|
||||
NUM_RANDOM_TOKEN_IDS = 10
|
||||
RTOL = 0.20
|
||||
ATOL = 0.00
|
||||
# ------------------------------------------------
|
||||
|
||||
torch.manual_seed(1234)
|
||||
if torch.cuda.is_available():
|
||||
torch.cuda.manual_seed_all(1234)
|
||||
torch.backends.cuda.matmul.allow_tf32 = False
|
||||
torch.backends.cudnn.allow_tf32 = False
|
||||
|
||||
|
||||
class TestOriginalLogprob(unittest.TestCase):
|
||||
def setUp(self):
|
||||
# ----- HF side (float32 weights) -----
|
||||
self.tokenizer = AutoTokenizer.from_pretrained(MODEL_ID, padding_side="right")
|
||||
self.hf_model = AutoModelForCausalLM.from_pretrained(
|
||||
MODEL_ID, torch_dtype=torch.float32, device_map="auto"
|
||||
)
|
||||
|
||||
# Shared sampling parameters
|
||||
self.sampling_params = {
|
||||
"temperature": 0.5, # SGLang uses 0.5, but original logprobs are used 1.0
|
||||
"top_p": 1.0,
|
||||
"top_k": 10,
|
||||
"max_new_tokens": 1,
|
||||
}
|
||||
|
||||
# ---------------------------------------------------------------------
|
||||
# Helper: compare one SGLang block (token_logprobs / top_logprobs / ids_logprobs)
|
||||
# against a reference HF log‑prob vector.
|
||||
# ---------------------------------------------------------------------
|
||||
def assert_logprobs_block_equal(
|
||||
self,
|
||||
hf_log_probs: torch.Tensor, # [V]
|
||||
token_log_probs: list,
|
||||
top_log_probs: list,
|
||||
ids_log_probs: list,
|
||||
random_token_ids: list,
|
||||
tag: str = "",
|
||||
):
|
||||
vals, idxs, _ = zip(*token_log_probs)
|
||||
sgl_vals = torch.tensor(vals, device=self.hf_model.device, dtype=torch.float32)
|
||||
sgl_idxs = torch.tensor(idxs, device=self.hf_model.device, dtype=torch.long)
|
||||
hf_vals = hf_log_probs[sgl_idxs]
|
||||
|
||||
self.assertTrue(
|
||||
torch.allclose(hf_vals, sgl_vals, rtol=RTOL, atol=ATOL),
|
||||
msg=f"[{tag}] token‑level mismatch at indices {sgl_idxs.tolist()}",
|
||||
)
|
||||
|
||||
hf_topk, _ = torch.topk(hf_log_probs, k=TOP_LOGPROBS_NUM, dim=-1)
|
||||
|
||||
sgl_topk = torch.tensor(
|
||||
[float(t[0]) for t in top_log_probs[0] if t and t[0] is not None][
|
||||
:TOP_LOGPROBS_NUM
|
||||
],
|
||||
dtype=torch.float32,
|
||||
device=self.hf_model.device,
|
||||
)
|
||||
|
||||
k = min(hf_topk.numel(), sgl_topk.numel())
|
||||
self.assertTrue(
|
||||
torch.allclose(hf_topk[:k], sgl_topk[:k], rtol=RTOL, atol=ATOL),
|
||||
msg=f"[{tag}] top‑k mismatch",
|
||||
)
|
||||
|
||||
indices = torch.tensor(
|
||||
random_token_ids, dtype=torch.long, device=hf_log_probs.device
|
||||
)
|
||||
|
||||
hf_token_ids = hf_log_probs[indices]
|
||||
|
||||
sgl_token_ids = torch.tensor(
|
||||
[v for v, _, _ in ids_log_probs[0]],
|
||||
device=self.hf_model.device,
|
||||
dtype=torch.float32,
|
||||
)
|
||||
self.assertTrue(
|
||||
torch.allclose(hf_token_ids, sgl_token_ids, rtol=RTOL, atol=ATOL),
|
||||
msg=f"[{tag}] token‑IDs mismatch",
|
||||
)
|
||||
|
||||
# Optional: print max abs diff for quick diagnostics
|
||||
max_diff = torch.max(torch.abs(hf_vals - sgl_vals)).item()
|
||||
print(f"[{tag}] max|diff| token‑level = {max_diff:.4f}")
|
||||
|
||||
def test_logprob_match(self):
|
||||
vocab_size = self.tokenizer.vocab_size
|
||||
|
||||
for env_val in ["True", "False"]:
|
||||
with self.subTest(SGLANG_RETURN_ORIGINAL_LOGPROB=env_val):
|
||||
os.environ["SGLANG_RETURN_ORIGINAL_LOGPROB"] = env_val
|
||||
|
||||
# ----- SGLang side -----
|
||||
sgl_engine = sgl.Engine(
|
||||
model_path=MODEL_ID,
|
||||
skip_tokenizer_init=True,
|
||||
trust_remote_code=True,
|
||||
mem_fraction_static=0.60,
|
||||
)
|
||||
|
||||
for prompt in PROMPTS:
|
||||
random_token_ids = sorted(
|
||||
random.sample(range(vocab_size), NUM_RANDOM_TOKEN_IDS)
|
||||
)
|
||||
|
||||
enc = self.tokenizer(prompt, return_tensors="pt")
|
||||
input_ids = enc["input_ids"].to(self.hf_model.device)
|
||||
attn_mask = enc["attention_mask"].to(self.hf_model.device)
|
||||
|
||||
with torch.inference_mode():
|
||||
hf_out = self.hf_model(
|
||||
input_ids=input_ids,
|
||||
attention_mask=attn_mask,
|
||||
return_dict=True,
|
||||
)
|
||||
logits = hf_out.logits[:, -1, :] # [1, V]
|
||||
hf_log_probs = F.log_softmax(
|
||||
logits.float() / self.sampling_params["temperature"], dim=-1
|
||||
)[0]
|
||||
hf_original_log_probs = F.log_softmax(logits.float(), dim=-1)[0]
|
||||
|
||||
outputs = sgl_engine.generate(
|
||||
input_ids=input_ids[0].tolist(),
|
||||
sampling_params=self.sampling_params,
|
||||
return_logprob=True,
|
||||
top_logprobs_num=TOP_LOGPROBS_NUM,
|
||||
token_ids_logprob=random_token_ids,
|
||||
)
|
||||
|
||||
if isinstance(outputs, list):
|
||||
outputs = outputs[0]
|
||||
meta = outputs["meta_info"]
|
||||
|
||||
# Check original logprobs only if enabled
|
||||
if env_val.lower() == "true":
|
||||
self.assert_logprobs_block_equal(
|
||||
hf_log_probs=hf_original_log_probs,
|
||||
token_log_probs=meta["output_token_logprobs"],
|
||||
top_log_probs=meta["output_top_logprobs"],
|
||||
ids_log_probs=meta["output_token_ids_logprobs"],
|
||||
random_token_ids=random_token_ids,
|
||||
tag=f"Original logprobs SGLang vs HF: {prompt} ({env_val})",
|
||||
)
|
||||
else:
|
||||
# Always check regular logprobs
|
||||
self.assert_logprobs_block_equal(
|
||||
hf_log_probs=hf_log_probs,
|
||||
token_log_probs=meta["output_token_logprobs"],
|
||||
top_log_probs=meta["output_top_logprobs"],
|
||||
ids_log_probs=meta["output_token_ids_logprobs"],
|
||||
random_token_ids=random_token_ids,
|
||||
tag=f"logprobs SGLang vs HF: {prompt} ({env_val})",
|
||||
)
|
||||
sgl_engine.shutdown()
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
@@ -0,0 +1,250 @@
|
||||
import json
|
||||
import random
|
||||
import re
|
||||
import unittest
|
||||
from concurrent.futures import ThreadPoolExecutor
|
||||
|
||||
import requests
|
||||
|
||||
from sglang.srt.utils import kill_process_tree
|
||||
from sglang.test.ci.ci_register import register_amd_ci, register_cuda_ci
|
||||
|
||||
register_cuda_ci(est_time=82, suite="stage-b-test-small-1-gpu")
|
||||
register_amd_ci(est_time=180, suite="stage-b-test-small-1-gpu")
|
||||
from sglang.test.test_utils import (
|
||||
DEFAULT_SMALL_MODEL_NAME_FOR_TEST,
|
||||
DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
|
||||
DEFAULT_URL_FOR_TEST,
|
||||
CustomTestCase,
|
||||
popen_launch_server,
|
||||
)
|
||||
|
||||
|
||||
class TestPenalty(CustomTestCase):
|
||||
@classmethod
|
||||
def setUpClass(cls):
|
||||
cls.model = DEFAULT_SMALL_MODEL_NAME_FOR_TEST
|
||||
cls.base_url = DEFAULT_URL_FOR_TEST
|
||||
cls.process = popen_launch_server(
|
||||
cls.model,
|
||||
cls.base_url,
|
||||
timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
|
||||
)
|
||||
|
||||
@classmethod
|
||||
def tearDownClass(cls):
|
||||
kill_process_tree(cls.process.pid)
|
||||
|
||||
def run_decode(self, sampling_params):
|
||||
"""Helper method for basic decode tests."""
|
||||
return_logprob = True
|
||||
top_logprobs_num = 5
|
||||
return_text = True
|
||||
n = 1
|
||||
|
||||
response = requests.post(
|
||||
self.base_url + "/generate",
|
||||
json={
|
||||
# prompt that is supposed to generate < 32 tokens
|
||||
"text": "<|start_header_id|>user<|end_header_id|>\n\nWhat is the answer for 1 + 1 = ?<|eot_id|><|start_header_id|>assistant<|end_header_id|>\n\n",
|
||||
"sampling_params": {
|
||||
"max_new_tokens": 48,
|
||||
"n": n,
|
||||
**sampling_params,
|
||||
},
|
||||
"return_logprob": return_logprob,
|
||||
"top_logprobs_num": top_logprobs_num,
|
||||
"return_text_in_logprobs": return_text,
|
||||
"logprob_start_len": 0,
|
||||
},
|
||||
)
|
||||
self.assertEqual(response.status_code, 200)
|
||||
print(json.dumps(response.json()))
|
||||
print("=" * 100)
|
||||
|
||||
def run_generate_with_prompt(self, prompt, sampling_params, max_tokens=100):
|
||||
"""Helper method to generate text with a specific prompt and parameters."""
|
||||
sampling_params.setdefault("temperature", 0.05)
|
||||
sampling_params.setdefault("top_p", 1.0)
|
||||
|
||||
response = requests.post(
|
||||
self.base_url + "/v1/chat/completions",
|
||||
json={
|
||||
"model": self.model,
|
||||
"messages": [{"role": "user", "content": prompt}],
|
||||
"max_tokens": max_tokens,
|
||||
**sampling_params,
|
||||
},
|
||||
)
|
||||
self.assertEqual(response.status_code, 200)
|
||||
result = response.json()
|
||||
content = result["choices"][0]["message"]["content"]
|
||||
return content
|
||||
|
||||
def count_word_repetitions(self, text, word):
|
||||
"""Count how many times a specific word appears in the text."""
|
||||
return len(re.findall(r"\b" + re.escape(word) + r"\b", text.lower()))
|
||||
|
||||
def _test_penalty_effect(
|
||||
self,
|
||||
prompt,
|
||||
baseline_params,
|
||||
penalty_params,
|
||||
target_word,
|
||||
expected_reduction=True,
|
||||
max_tokens=50,
|
||||
):
|
||||
"""Generic test for penalty effects."""
|
||||
# Run multiple iterations to get more reliable results
|
||||
baseline_counts = []
|
||||
penalty_counts = []
|
||||
|
||||
for i in range(5):
|
||||
baseline_output = self.run_generate_with_prompt(
|
||||
prompt, baseline_params, max_tokens
|
||||
)
|
||||
penalty_output = self.run_generate_with_prompt(
|
||||
prompt, penalty_params, max_tokens
|
||||
)
|
||||
|
||||
baseline_count = self.count_word_repetitions(baseline_output, target_word)
|
||||
penalty_count = self.count_word_repetitions(penalty_output, target_word)
|
||||
|
||||
baseline_counts.append(baseline_count)
|
||||
penalty_counts.append(penalty_count)
|
||||
|
||||
# Calculate averages
|
||||
avg_baseline = sum(baseline_counts) / len(baseline_counts)
|
||||
avg_penalty = sum(penalty_counts) / len(penalty_counts)
|
||||
|
||||
if expected_reduction:
|
||||
# Simple check: penalty should reduce repetition
|
||||
self.assertLess(
|
||||
avg_penalty,
|
||||
avg_baseline,
|
||||
f"Penalty should reduce '{target_word}' repetition: {avg_baseline:.1f} → {avg_penalty:.1f}",
|
||||
)
|
||||
else:
|
||||
self.assertGreater(
|
||||
avg_penalty,
|
||||
avg_baseline,
|
||||
f"Negative penalty should increase '{target_word}' repetition",
|
||||
)
|
||||
|
||||
def test_default_values(self):
|
||||
self.run_decode({})
|
||||
|
||||
def test_frequency_penalty(self):
|
||||
self.run_decode({"frequency_penalty": 2})
|
||||
|
||||
def test_min_new_tokens(self):
|
||||
self.run_decode({"min_new_tokens": 16})
|
||||
|
||||
def test_presence_penalty(self):
|
||||
self.run_decode({"presence_penalty": 2})
|
||||
|
||||
def test_penalty_mixed(self):
|
||||
args = [
|
||||
{},
|
||||
{},
|
||||
{},
|
||||
{"frequency_penalty": 2},
|
||||
{"presence_penalty": 1},
|
||||
{"min_new_tokens": 16},
|
||||
{"frequency_penalty": 0.2},
|
||||
{"presence_penalty": 0.4},
|
||||
{"min_new_tokens": 8},
|
||||
{"frequency_penalty": 0.4, "presence_penalty": 0.8},
|
||||
{"frequency_penalty": 0.4, "min_new_tokens": 12},
|
||||
{"presence_penalty": 0.8, "min_new_tokens": 12},
|
||||
{"presence_penalty": -0.3, "frequency_penalty": 1.3, "min_new_tokens": 32},
|
||||
{"presence_penalty": 0.3, "frequency_penalty": -1.3, "min_new_tokens": 32},
|
||||
]
|
||||
random.shuffle(args * 5)
|
||||
with ThreadPoolExecutor(8) as executor:
|
||||
list(executor.map(self.run_decode, args))
|
||||
|
||||
def test_frequency_penalty_reduces_word_repetition(self):
|
||||
"""Test frequency penalty using word repetition."""
|
||||
prompt = "Write exactly 10 very small sentences, each containing the word 'data'. Use the word 'data' as much as possible."
|
||||
baseline_params = {"frequency_penalty": 0.0, "repetition_penalty": 1.0}
|
||||
penalty_params = {"frequency_penalty": 1.99, "repetition_penalty": 1.0}
|
||||
self._test_penalty_effect(prompt, baseline_params, penalty_params, "data")
|
||||
|
||||
def test_presence_penalty_reduces_topic_repetition(self):
|
||||
"""Test presence penalty using topic repetition."""
|
||||
prompt = "Write the word 'machine learning' exactly 20 times in a row, separated by spaces."
|
||||
baseline_params = {"presence_penalty": 0.0, "repetition_penalty": 1.0}
|
||||
penalty_params = {"presence_penalty": 1.99, "repetition_penalty": 1.0}
|
||||
self._test_penalty_effect(
|
||||
prompt, baseline_params, penalty_params, "machine learning"
|
||||
)
|
||||
|
||||
def test_combined_penalties_reduce_repetition(self):
|
||||
"""Test combined penalty effects."""
|
||||
prompt = "Write exactly 10 short sentences, each containing the word 'data'. Use the word 'data' as much as possible."
|
||||
baseline_params = {
|
||||
"frequency_penalty": 0.0,
|
||||
"presence_penalty": 0.0,
|
||||
"repetition_penalty": 1.0,
|
||||
}
|
||||
penalty_params = {
|
||||
"frequency_penalty": 1.99,
|
||||
"presence_penalty": 1.99,
|
||||
"repetition_penalty": 1.99,
|
||||
}
|
||||
self._test_penalty_effect(
|
||||
prompt, baseline_params, penalty_params, "data", max_tokens=100
|
||||
)
|
||||
|
||||
def test_penalty_edge_cases_negative_penalty_values(self):
|
||||
"""Test edge cases with negative penalty values."""
|
||||
prompt = "Write the word 'test' exactly 15 times in a row, separated by spaces."
|
||||
baseline_params = {
|
||||
"frequency_penalty": 0.0,
|
||||
"presence_penalty": 0.0,
|
||||
"repetition_penalty": 1.0,
|
||||
}
|
||||
negative_penalty_params = {
|
||||
"frequency_penalty": -0.5,
|
||||
"presence_penalty": -0.25,
|
||||
"repetition_penalty": 1.0,
|
||||
}
|
||||
# Negative penalties should increase repetition (expected_reduction=False)
|
||||
self._test_penalty_effect(
|
||||
prompt,
|
||||
baseline_params,
|
||||
negative_penalty_params,
|
||||
"test",
|
||||
expected_reduction=False,
|
||||
max_tokens=60,
|
||||
)
|
||||
|
||||
def test_penalty_edge_cases_extreme_penalty_values(self):
|
||||
"""Test edge cases with extreme penalty values."""
|
||||
prompt = (
|
||||
"Write the word 'extreme' exactly 20 times in a row, separated by spaces."
|
||||
)
|
||||
baseline_params = {
|
||||
"frequency_penalty": 0.0,
|
||||
"presence_penalty": 0.0,
|
||||
"repetition_penalty": 1.0,
|
||||
}
|
||||
extreme_penalty_params = {
|
||||
"frequency_penalty": 2.0,
|
||||
"presence_penalty": 2.0,
|
||||
"repetition_penalty": 2.0,
|
||||
}
|
||||
# Extreme penalties should strongly reduce repetition
|
||||
self._test_penalty_effect(
|
||||
prompt,
|
||||
baseline_params,
|
||||
extreme_penalty_params,
|
||||
"extreme",
|
||||
expected_reduction=True,
|
||||
max_tokens=80,
|
||||
)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main(verbosity=3)
|
||||
@@ -0,0 +1,94 @@
|
||||
import unittest
|
||||
from types import SimpleNamespace
|
||||
|
||||
import requests
|
||||
|
||||
from sglang.srt.utils import kill_process_tree
|
||||
from sglang.test.ci.ci_register import register_amd_ci, register_cuda_ci
|
||||
|
||||
register_cuda_ci(est_time=66, suite="stage-b-test-small-1-gpu")
|
||||
register_amd_ci(est_time=66, suite="stage-b-test-small-1-gpu")
|
||||
from sglang.test.run_eval import run_eval
|
||||
from sglang.test.test_utils import (
|
||||
DEFAULT_MODEL_NAME_FOR_TEST,
|
||||
DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
|
||||
DEFAULT_URL_FOR_TEST,
|
||||
CustomTestCase,
|
||||
popen_launch_server,
|
||||
)
|
||||
|
||||
|
||||
class TestPyTorchSamplingBackend(CustomTestCase):
|
||||
@classmethod
|
||||
def setUpClass(cls):
|
||||
cls.model = DEFAULT_MODEL_NAME_FOR_TEST
|
||||
cls.base_url = DEFAULT_URL_FOR_TEST
|
||||
cls.process = popen_launch_server(
|
||||
cls.model,
|
||||
cls.base_url,
|
||||
timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
|
||||
other_args=["--sampling-backend", "pytorch", "--disable-radix-cache"],
|
||||
)
|
||||
|
||||
@classmethod
|
||||
def tearDownClass(cls):
|
||||
kill_process_tree(cls.process.pid)
|
||||
|
||||
def test_mmlu(self):
|
||||
args = SimpleNamespace(
|
||||
base_url=self.base_url,
|
||||
model=self.model,
|
||||
eval_name="mmlu",
|
||||
num_examples=64,
|
||||
num_threads=32,
|
||||
temperature=0.1,
|
||||
)
|
||||
|
||||
metrics = run_eval(args)
|
||||
self.assertGreaterEqual(metrics["score"], 0.65)
|
||||
|
||||
def test_greedy(self):
|
||||
|
||||
first_text = None
|
||||
|
||||
# ensure the answer is identical across single response
|
||||
for _ in range(5):
|
||||
response_single = requests.post(
|
||||
self.base_url + "/generate",
|
||||
json={
|
||||
"text": "The capital of Germany is",
|
||||
"sampling_params": {
|
||||
"temperature": 0,
|
||||
"max_new_tokens": 32,
|
||||
},
|
||||
},
|
||||
).json()
|
||||
text = response_single["text"]
|
||||
if first_text is None:
|
||||
first_text = text
|
||||
|
||||
self.assertEqual(text, first_text)
|
||||
|
||||
first_text = None
|
||||
|
||||
response_batch = requests.post(
|
||||
self.base_url + "/generate",
|
||||
json={
|
||||
"text": ["The capital of Germany is"] * 10,
|
||||
"sampling_params": {
|
||||
"temperature": 0,
|
||||
"max_new_tokens": 32,
|
||||
},
|
||||
},
|
||||
).json()
|
||||
|
||||
# ensure the answer is identical among the batch
|
||||
for i in range(10):
|
||||
text = response_batch[i]["text"]
|
||||
if first_text is None:
|
||||
first_text = text
|
||||
self.assertEqual(text, first_text)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
Reference in New Issue
Block a user