266 lines
10 KiB
Python
266 lines
10 KiB
Python
import io
|
|
import os
|
|
import pickle
|
|
import random
|
|
import time
|
|
import unittest
|
|
|
|
import numpy as np
|
|
import requests
|
|
import torch
|
|
|
|
import sglang as sgl
|
|
from sglang.test.test_utils import (
|
|
DEFAULT_SMALL_MODEL_NAME_FOR_TEST,
|
|
write_github_step_summary,
|
|
)
|
|
|
|
# Dense model configuration
|
|
DENSE_MODEL_NAME = DEFAULT_SMALL_MODEL_NAME_FOR_TEST
|
|
if torch.version.hip is not None:
|
|
print("Running on AMD ROCm GPU")
|
|
DENSE_INPUT_PKL_URL = "https://huggingface.co/datasets/yushengsu/logprobs/resolve/main/sglang_baseline_2000_amd.pkl"
|
|
DENSE_TOLERANCE_MAX_DIFF = 1.4
|
|
DENSE_TOLERANCE_MEAN_DIFF = 0.1
|
|
elif torch.version.cuda is not None:
|
|
print("Running on NVIDIA CUDA GPU")
|
|
DENSE_INPUT_PKL_URL = "https://huggingface.co/datasets/font-info/logprobs/resolve/main/sglang_baseline_2000.pkl"
|
|
DENSE_TOLERANCE_MAX_DIFF = 1.5
|
|
DENSE_TOLERANCE_MEAN_DIFF = 0.1
|
|
else:
|
|
print("No GPU backend (CPU only)")
|
|
|
|
# Common configuration
|
|
TOP_K = 20
|
|
MAX_RETRIES = 3
|
|
RETRY_DELAY = 2
|
|
NUM_SAMPLES = 1000
|
|
LOGPROB_SAMPLE_RATIO = 0.5
|
|
TEMPERATURE = 1.0
|
|
|
|
|
|
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(
|
|
model_path=DENSE_MODEL_NAME,
|
|
random_seed=42,
|
|
skip_tokenizer_init=True,
|
|
mem_fraction_static=0.80,
|
|
)
|
|
|
|
@classmethod
|
|
def tearDownClass(cls):
|
|
"""Clean up after all tests - shutdown the engine."""
|
|
cls.engine.shutdown()
|
|
torch.cuda.empty_cache()
|
|
|
|
def load_test_data(self):
|
|
"""Load test data from Hugging Face dataset with retry mechanism."""
|
|
print(f"Loading data from {DENSE_INPUT_PKL_URL}...")
|
|
|
|
for attempt in range(MAX_RETRIES):
|
|
try:
|
|
response = requests.get(DENSE_INPUT_PKL_URL, timeout=30)
|
|
response.raise_for_status()
|
|
|
|
with io.BytesIO(response.content) as f:
|
|
records = pickle.load(f)
|
|
|
|
if not records:
|
|
raise ValueError("Empty dataset")
|
|
|
|
print(f"Successfully loaded {len(records)} records")
|
|
return records
|
|
|
|
except Exception as e:
|
|
print(f"Attempt {attempt + 1}/{MAX_RETRIES} failed: {e}")
|
|
if attempt == MAX_RETRIES - 1:
|
|
raise Exception(
|
|
f"Failed to load data after {MAX_RETRIES} attempts: {e}"
|
|
)
|
|
time.sleep(RETRY_DELAY)
|
|
|
|
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 / 2,
|
|
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])
|
|
)
|
|
return max(diffs), float(np.mean(diffs))
|
|
|
|
def test_logprobs_comparison(self):
|
|
"""Test the logprobs comparison functionality with different parameter combinations."""
|
|
# Load test data with retry mechanism
|
|
records = self.load_test_data()
|
|
|
|
with self.subTest(
|
|
config={
|
|
"num_samples": NUM_SAMPLES,
|
|
"logprob_sample_ratio": LOGPROB_SAMPLE_RATIO,
|
|
"temperature": TEMPERATURE,
|
|
}
|
|
):
|
|
|
|
# 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
|
|
]
|
|
|
|
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}",
|
|
)
|
|
|
|
# Write results to GitHub summary
|
|
summary_content = f"""
|
|
- **Configuration**: {{"num_samples": {NUM_SAMPLES}, "logprob_sample_ratio": {LOGPROB_SAMPLE_RATIO}, "temperature": {TEMPERATURE}}}
|
|
- **Max of max Δ**: {max_of_max:.6g}
|
|
- **Mean of mean Δ**: {mean_of_mean:.6g}
|
|
- **Status**: {'✅ Passed' if max_of_max <= DENSE_TOLERANCE_MAX_DIFF and mean_of_mean <= DENSE_TOLERANCE_MEAN_DIFF else '❌ Failed'}
|
|
"""
|
|
write_github_step_summary(summary_content)
|
|
|
|
# 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])
|
|
)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
unittest.main()
|