minke.yu/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 12 Packages Projects Releases Wiki Activity

Test reorganization: Move tests to manual/ (#13610)

Browse Source
This commit is contained in:
alisonshao
2025-11-20 13:41:58 -08:00
committed by GitHub
parent ada8ce1fd0
commit 6b262ac839
74 changed files with 0 additions and 74 deletions
Download Patch File Download Diff File Expand all files Collapse all files
test/manual/test_logprobs.py
+526
View File
@@ -0,0 +1,526 @@
"""
Logprobs Accuracy Test for SGLang
======================
With deterministic/batch invariant kernels, we can ensure that SGLang produces exactly the same
logprobs results for identical inputs. However, logprobs are highly sensitive to GPU hardware,
kernels, torch versions, and other factors, so we cannot maintain a unified logprobs baseline
across different machines.
This test is designed to be run locally by contributors to verify logprobs accuracy
before making changes to related code.
When submitting changes that affect logprobs computation, please:
1. Generate baseline
2. Run test
3. Submit results
We really appreciate your effort and contribution to SGLang!
======================
What does this test do?
This test fetches 1000 samples from the ShareGPT dataset, generates logprobs for each sample,
and saves them as a baseline. Then, by running the test mode, it validates the accuracy of
logprobs by comparing them against the baseline.
This test ensures that:
- the boundary of log probs requests are correct, eg, the index for tokens that required log probs are strictly followed
- logprobs remain invariant between test runs, and also before and after your code changes;
======================
Usage
Step 1: Generate Baseline (Before Code Changes)
```bash
python test/srt/test_logprobs.py gen
```
Step 2: Test Against Baseline (After Code Changes)
```bash
python test/srt/test_logprobs.py test
```
This tests your changes against the locally generated baseline from Step 1.
The test passes if the maximum and mean differences are within the tolerance thresholds.
======================
"""
import argparse
import json
import os
import pickle
import random
import unittest
import numpy as np
import requests
import torch
from transformers import AutoTokenizer
import sglang as sgl
from sglang.test.test_utils import DEFAULT_SMALL_MODEL_NAME_FOR_TEST
# Configuration
DENSE_MODEL_NAME = DEFAULT_SMALL_MODEL_NAME_FOR_TEST
SHAREGPT_URL = (
"https://huggingface.co/datasets/anon8231489123/"
"ShareGPT_Vicuna_unfiltered/resolve/main/ShareGPT_V3_unfiltered_cleaned_split.json"
)
# Hardware-specific configuration
if torch.version.cuda is not None:
print("Running on NVIDIA CUDA GPU")
DENSE_TOLERANCE_MAX_DIFF = 1e-5
DENSE_TOLERANCE_MEAN_DIFF = 1e-5
else:
print("No GPU backend (CPU only)")
raise ValueError("No GPU backend (CPU only)")
# Common configuration
TOP_K = 20
NUM_SAMPLES = 1000
LOGPROB_SAMPLE_RATIO = 0.5
TEMPERATURE = 1.0
MAX_LEN = 20000
# Default output files
DEFAULT_BASELINE_PKL = "sglang_baseline_local.pkl"
DEFAULT_META_JSON = "baseline_meta_preview.json"
# Default engine configuration
DEFAULT_ENGINE_CONFIG = {
"model_path": DENSE_MODEL_NAME,
"random_seed": 42,
"skip_tokenizer_init": True,
"mem_fraction_static": 0.8,
"enable_deterministic_inference": True,
"attention_backend": "flashinfer",
}
def generate_baseline(
baseline_file=DEFAULT_BASELINE_PKL,
meta_file=DEFAULT_META_JSON,
num_samples=NUM_SAMPLES,
):
"""Generate a local baseline for logprobs testing.
Args:
baseline_file: Path to save the baseline pickle file
meta_file: Path to save the metadata preview JSON file
num_samples: Number of samples to generate
"""
print(f"SGLang version: {sgl.__version__}")
print("Downloading ShareGPT dataset...")
# Download ShareGPT dataset
try:
response = requests.get(SHAREGPT_URL, timeout=30)
response.raise_for_status()
data = response.json()
print(f"Dataset size: {len(data)}")
except requests.exceptions.RequestException as e:
raise Exception(f"Failed to download ShareGPT dataset: {e}") from e
# Filter and prepare texts
texts = []
for s in data:
if "conversations" in s and len(s["conversations"]) > 0:
try:
text = s["conversations"][0]["value"]
if isinstance(text, str) and len(text) <= MAX_LEN and len(text) >= 5500:
texts.append(text)
if len(texts) >= num_samples * 40: # Get more samples for filtering
break
except (KeyError, IndexError, TypeError) as e:
print(f"Warning: Skipping invalid conversation data: {e}")
continue
if not texts:
raise ValueError("No valid texts found in the dataset")
print(f"Loading tokenizer for {DENSE_MODEL_NAME}...")
tokenizer = AutoTokenizer.from_pretrained(DENSE_MODEL_NAME, use_fast=True)
rng = np.random.default_rng(42)
print(f"Launching SGLang Engine with {DENSE_MODEL_NAME}...")
engine = sgl.Engine(
model_path=DENSE_MODEL_NAME,
attention_backend="flashinfer",
enable_deterministic_inference=True,
random_seed=42,
skip_tokenizer_init=True,
mem_fraction_static=0.8,
max_running_requests=1,
)
records = []
prompt_lengths = []
try:
for i, text in enumerate(texts):
if len(records) >= num_samples:
break
try:
ids = tokenizer.encode(text, add_special_tokens=False)
if len(ids) < 5:
continue
start_pos = int(rng.integers(0, max(1, len(ids) - 3)))
outputs = engine.generate(
input_ids=[ids],
sampling_params={
"temperature": 1.0,
"top_p": 1.0,
"top_k": TOP_K,
"max_new_tokens": 1,
},
return_logprob=True,
logprob_start_len=start_pos,
top_logprobs_num=TOP_K,
)
meta = outputs[0]["meta_info"]
records.append(
dict(id=i, text=text, ids=ids, start_pos=start_pos, meta=meta)
)
prompt_lengths.append(len(ids))
if (i + 1) % 50 == 0:
print(f"Processed {len(records)}/{num_samples} samples")
except Exception as e:
print(f"Warning: Failed to process sample {i}: {e}")
continue
if not records:
raise RuntimeError(
"Failed to generate any baseline records. Please check the warnings above for errors."
)
# Save baseline files
with open(baseline_file, "wb") as f:
pickle.dump(records, f)
with open(meta_file, "w", encoding="utf-8") as f:
json.dump(records[:2], f, ensure_ascii=False, indent=2)
print(f"✅ Saved {len(records)} samples to {baseline_file}")
print(f"✅ Meta preview saved to {meta_file}")
if prompt_lengths:
avg_prompt_length = sum(prompt_lengths) / len(prompt_lengths)
print(f"📊 Average prompt length: {avg_prompt_length:.2f} tokens")
finally:
engine.shutdown()
torch.cuda.empty_cache()
class TestLogprobsDense(unittest.TestCase):
@classmethod
def setUpClass(cls):
"""Set up the test class - initialize the engine once for all tests."""
print(f"Launching SGLang Engine with {DENSE_MODEL_NAME}...")
cls.engine = sgl.Engine(**DEFAULT_ENGINE_CONFIG)
@classmethod
def tearDownClass(cls):
"""Clean up after all tests - shutdown the engine."""
cls.engine.shutdown()
torch.cuda.empty_cache()
@classmethod
def restart_engine_with_config(cls, **kwargs):
"""Create engine with custom configuration"""
# Safely shutdown existing engine
cls.engine.shutdown()
torch.cuda.empty_cache()
# Set chunk size
chunk_size = kwargs.pop("chunk_size", None)
if chunk_size is not None:
print(f"Setting chunk size to {chunk_size}")
os.environ["SGLANG_ENABLE_LOGITS_PROCESSER_CHUNK"] = "True"
os.environ["SGLANG_LOGITS_PROCESSER_CHUNK_SIZE"] = str(chunk_size)
else:
os.environ["SGLANG_ENABLE_LOGITS_PROCESSER_CHUNK"] = "False"
# Create engine with merged configuration
engine_config = {**DEFAULT_ENGINE_CONFIG, **kwargs}
cls.engine = sgl.Engine(**engine_config)
def load_test_data(self, baseline_file=None):
"""Load test data from local baseline file. In test mode, only local baseline is supported."""
if not baseline_file:
raise ValueError("baseline_file is required in test mode")
if not os.path.exists(baseline_file):
raise FileNotFoundError(
f"Baseline file not found: {baseline_file}. Please run 'gen' mode first to generate the baseline."
)
print(f"Loading local baseline from {baseline_file}...")
try:
with open(baseline_file, "rb") as f:
records = pickle.load(f)
print(f"Successfully loaded {len(records)} records from local baseline")
return records
except (IOError, pickle.PickleError) as e:
raise Exception(f"Failed to load local baseline: {e}") from e
def compare_meta(self, baseline_meta, sglang_meta):
"""Compare metadata between two outputs and return max and mean differences."""
diffs = []
for key in ["input_top_logprobs", "output_top_logprobs"]:
baseline_logprobs, sglang_logprobs = baseline_meta[key], sglang_meta[key]
self.assertEqual(
len(baseline_logprobs),
len(sglang_logprobs),
f"Length of {key} is not equal, sglang did not return the correct number of log probs(should be top 20)",
)
for baseline_entry, sglang_entry in zip(baseline_logprobs, sglang_logprobs):
if not baseline_entry or not sglang_entry:
continue
baseline_token_map = {tid: lp for lp, tid, _ in baseline_entry}
sglang_token_map = {tid: lp for lp, tid, _ in sglang_entry}
common_tokens = baseline_token_map.keys() & sglang_token_map.keys()
self.assertGreaterEqual(
len(common_tokens),
TOP_K,
f"there are only {len(common_tokens)} common topk tokens that matches",
)
for token_id in common_tokens:
diffs.append(
abs(baseline_token_map[token_id] - sglang_token_map[token_id])
)
if not diffs:
return 0.0, 0.0
return max(diffs), float(np.mean(diffs))
def test_logprobs_comparison(self, baseline_file=None):
"""Test the logprobs comparison functionality with different parameter combinations."""
# Load test data with retry mechanism
records = self.load_test_data(baseline_file)
# Fast configs for CI
test_configs = [
{"num_samples": NUM_SAMPLES},
{"num_samples": 42, "chunk_size": 1, "max_running_requests": 16},
{"num_samples": 42, "chunk_size": 2, "max_running_requests": 16},
{"num_samples": 42, "chunk_size": 3, "max_running_requests": 16},
{"num_samples": NUM_SAMPLES, "chunk_size": 16, "max_running_requests": 128},
{"num_samples": NUM_SAMPLES, "chunk_size": 128, "max_running_requests": 16},
{"num_samples": NUM_SAMPLES, "chunk_size": 128, "max_running_requests": 8},
{"num_samples": NUM_SAMPLES, "chunk_size": 128, "max_running_requests": 32},
{
"num_samples": NUM_SAMPLES,
"chunk_size": 128,
"max_running_requests": 128,
},
{"num_samples": NUM_SAMPLES, "chunk_size": 256, "max_running_requests": 8},
{"num_samples": NUM_SAMPLES, "chunk_size": 256, "max_running_requests": 32},
{
"num_samples": NUM_SAMPLES,
"chunk_size": 256,
"max_running_requests": 128,
},
]
# Run tests
for config in test_configs:
with self.subTest(config=config):
print(f"Testing with config: {config}")
# Sample records for this config
test_records = random.sample(records, k=min(NUM_SAMPLES, len(records)))
random.shuffle(test_records)
# Calculate how many samples should return logprobs
logprob_count = int(len(test_records) * LOGPROB_SAMPLE_RATIO)
print(
f"Testing with {len(test_records)} samples, temperature={TEMPERATURE}"
)
print(
f"Will return logprobs for {logprob_count} samples (ratio: {LOGPROB_SAMPLE_RATIO})"
)
all_max, all_mean = [], []
logprob_returned_count = 0
# Process all records at once
input_ids = [rec["ids"] for rec in test_records]
logprob_start_lens = [rec["start_pos"] for rec in test_records]
# Determine which samples should return logprobs (randomly selected)
logprob_indices = set(
random.sample(range(len(test_records)), logprob_count)
)
return_logprob_array = [
sample_idx in logprob_indices
for sample_idx in range(len(test_records))
]
# Sampling param per request
sampling_params = [
{
"temperature": TEMPERATURE,
"top_p": 1.0,
"top_k": TOP_K,
"max_new_tokens": 1,
}
for _ in test_records
]
# Some configs must restart the engine to take effect
chunk_size = config.get("chunk_size", None)
max_running_requests = config.get("max_running_requests", None)
if chunk_size is not None or max_running_requests is not None:
self.restart_engine_with_config(
chunk_size=chunk_size,
max_running_requests=max_running_requests,
)
outputs = self.engine.generate(
input_ids=input_ids,
sampling_params=sampling_params,
return_logprob=return_logprob_array,
logprob_start_len=logprob_start_lens,
top_logprobs_num=TOP_K,
)
for sample_idx, (rec, output) in enumerate(zip(test_records, outputs)):
# Only compare logprobs for samples that should have them
if sample_idx in logprob_indices:
# Safe access to meta_info and input_top_logprobs
meta_info = output.get("meta_info")
input_top_logprobs = (
meta_info.get("input_top_logprobs") if meta_info else None
)
self.assertIsNotNone(
input_top_logprobs,
f"return_logprob enabled on this sample, but input_top_logprobs is None (length: {len(input_top_logprobs) if input_top_logprobs is not None else 'N/A'})",
)
baseline_meta = rec["meta"]
sglang_meta = meta_info
max_diff, mean_diff = self.compare_meta(
baseline_meta, sglang_meta
)
all_max.append(max_diff)
all_mean.append(mean_diff)
logprob_returned_count += 1
else:
# Verify that logprobs were not returned for this sample
meta_info = output.get("meta_info")
input_top_logprobs = (
meta_info.get("input_top_logprobs") if meta_info else None
)
output_token_ids_logprobs = (
meta_info.get("output_token_ids_logprobs")
if meta_info
else None
)
self.assertFalse(
input_top_logprobs,
f"return_logprob is disabled on this sample, Sample {sample_idx} should not have logprobs, content: {output_token_ids_logprobs}",
)
max_of_max = max(all_max) if all_max else 0.0
mean_of_mean = np.mean(all_mean) if all_mean else 0.0
print(f"max Δ={max_of_max:.6g}")
print(f"mean Δ={mean_of_mean:.6g}")
print(
f"logprobs returned for {logprob_returned_count} samples (expected: {logprob_count})"
)
# Verify correct number of logprobs returned
self.assertEqual(
logprob_returned_count,
logprob_count,
f"Expected {logprob_count} samples with logprobs, got {logprob_returned_count}",
)
# Basic validation
self.assertIsInstance(all_max, list)
self.assertIsInstance(all_mean, list)
self.assertGreater(
len(all_max),
0,
f"No test samples processed for config {{'num_samples': {NUM_SAMPLES}, 'logprob_sample_ratio': {LOGPROB_SAMPLE_RATIO}, 'temperature': {TEMPERATURE}}}",
)
# Tolerance checks with clear error messages
failed_samples = []
for sample_idx, (max_diff, mean_diff) in enumerate(
zip(all_max, all_mean)
):
if max_diff > DENSE_TOLERANCE_MAX_DIFF:
failed_samples.append(
f"Sample {sample_idx}: max_diff={max_diff:.6g} > {DENSE_TOLERANCE_MAX_DIFF}"
)
if mean_diff > DENSE_TOLERANCE_MEAN_DIFF:
failed_samples.append(
f"Sample {sample_idx}: mean_diff={mean_diff:.6g} > {DENSE_TOLERANCE_MEAN_DIFF}"
)
if failed_samples:
self.fail(
f"Config {{'num_samples': {NUM_SAMPLES}, 'logprob_sample_ratio': {LOGPROB_SAMPLE_RATIO}, 'temperature': {TEMPERATURE}}} - Tolerance exceeded in {len(failed_samples)} samples:\n"
+ "\n".join(failed_samples[:5])
)
def main():
"""Main function to handle command line arguments and run either generation or testing."""
parser = argparse.ArgumentParser(
description="SGLang Logprobs Test and Baseline Generation"
)
parser.add_argument(
"mode",
choices=["gen", "test"],
help="Mode to run: 'gen' to generate baseline, 'test' to run tests",
)
args = parser.parse_args()
if args.mode == "gen":
print("🚀 Generating baseline...")
generate_baseline()
print(f"\n✅ Baseline generation complete!")
print(f"📁 Baseline saved to: {DEFAULT_BASELINE_PKL}")
print(f"📁 Metadata preview saved to: {DEFAULT_META_JSON}")
print(f"\n💡 Next steps:")
print(f" 1. Make your code changes")
print(f" 2. Run: python {__file__} test")
elif args.mode == "test":
print("🧪 Running logprobs test...")
if not os.path.exists(DEFAULT_BASELINE_PKL):
print(f"❌ Baseline file not found: {DEFAULT_BASELINE_PKL}")
print(f"💡 Generate baseline first by running:")
print(f" python {__file__} gen")
print(f" This will download ShareGPT data and generate a local baseline.")
return 1
# Set environment variable for testing
os.environ["RETURN_ORIGINAL_LOGPROB"] = "True"
# Create test instance and run
test_instance = TestLogprobsDense()
test_instance.setUpClass()
try:
test_instance.test_logprobs_comparison(baseline_file=DEFAULT_BASELINE_PKL)
print("\n✅ Test completed successfully!")
finally:
test_instance.tearDownClass()
return 0
if __name__ == "__main__":
exit(main())