327 lines
11 KiB
Python
327 lines
11 KiB
Python
import inspect
|
|
import json
|
|
import os
|
|
import random
|
|
|
|
import numpy as np
|
|
import requests
|
|
|
|
from sglang.srt.utils.hf_transformers_utils import get_tokenizer
|
|
|
|
# LongBench V2 dataset configuration
|
|
# Reference: https://github.com/THUDM/LongBench
|
|
LONGBENCH_V2_DATASET = "THUDM/LongBench-v2"
|
|
LONGBENCH_V2_SPLIT = "train"
|
|
DEFAULT_NUM_SAMPLES = 48 # Number of samples to use
|
|
DEFAULT_PROMPT_TOKENS = 3000 # Maximum number of tokens to use
|
|
CACHE_DIR = os.path.join(os.path.dirname(__file__), ".longbench_cache")
|
|
|
|
# In-memory cache for the current session
|
|
_cached_input_ids = {}
|
|
|
|
|
|
def format_longbench_v2_example(example):
|
|
"""Format a LongBench V2 example into a single text string (context + question only)."""
|
|
context = example.get("context", "")
|
|
question = example.get("question", "")
|
|
return f"{context} {question}"
|
|
|
|
|
|
def get_input_ids(
|
|
tokenizer_path, max_prompt_tokens=DEFAULT_PROMPT_TOKENS, num_samples=None
|
|
):
|
|
"""Get input_ids from LongBench V2 dataset with local caching."""
|
|
# Create cache key based on parameters
|
|
if num_samples is None:
|
|
num_samples = DEFAULT_NUM_SAMPLES
|
|
cache_key = f"{tokenizer_path}_{max_prompt_tokens}_{num_samples}"
|
|
|
|
# Check in-memory cache first (fastest)
|
|
if cache_key in _cached_input_ids:
|
|
print(
|
|
f"Using in-memory cached data ({len(_cached_input_ids[cache_key])} prompts)"
|
|
)
|
|
return _cached_input_ids[cache_key]
|
|
|
|
# Check local file cache
|
|
os.makedirs(CACHE_DIR, exist_ok=True)
|
|
# Use a safe filename
|
|
safe_name = tokenizer_path.replace("/", "_").replace("\\", "_")
|
|
cache_file = os.path.join(
|
|
CACHE_DIR, f"input_ids_{safe_name}_{max_prompt_tokens}_{num_samples}.json"
|
|
)
|
|
|
|
if os.path.exists(cache_file):
|
|
print(f"Loading from local cache: {cache_file}")
|
|
with open(cache_file, "r") as f:
|
|
input_ids = json.load(f)
|
|
_cached_input_ids[cache_key] = input_ids
|
|
print(f"Loaded {len(input_ids)} prompts from cache")
|
|
return input_ids
|
|
|
|
# Download from HuggingFace using streaming
|
|
try:
|
|
from datasets import load_dataset
|
|
except ImportError as exc:
|
|
raise ImportError(
|
|
"Please install the 'datasets' package: pip install datasets"
|
|
) from exc
|
|
|
|
tokenizer = get_tokenizer(tokenizer_path)
|
|
|
|
print(f"Downloading {num_samples} samples from LongBench V2 (streaming)...")
|
|
dataset = load_dataset(
|
|
LONGBENCH_V2_DATASET, split=LONGBENCH_V2_SPLIT, streaming=True
|
|
)
|
|
|
|
input_ids = []
|
|
for i, example in enumerate(dataset):
|
|
if len(input_ids) >= num_samples:
|
|
break
|
|
text = format_longbench_v2_example(example)
|
|
tokens = tokenizer.encode(text)
|
|
# Truncate to a random length between 0.5x and 1.5x of max_prompt_tokens
|
|
truncate_len = int(max_prompt_tokens * random.uniform(0.5, 1.5))
|
|
input_ids.append(tokens[:truncate_len])
|
|
|
|
# Save to local cache
|
|
with open(cache_file, "w") as f:
|
|
json.dump(input_ids, f)
|
|
print(f"Saved {len(input_ids)} prompts to cache: {cache_file}")
|
|
|
|
# Also cache in memory
|
|
_cached_input_ids[cache_key] = input_ids
|
|
|
|
return input_ids
|
|
|
|
|
|
def compare_kl_divergence(
|
|
input_logprobs, output_logprobs, ACC_THRESHOLDS, model_name, test_name
|
|
):
|
|
"""Compare the KL divergence between input and output log probabilities."""
|
|
kl_divs = []
|
|
for input_logprob, output_logprob in zip(input_logprobs, output_logprobs):
|
|
input_logprob = np.array(input_logprob)
|
|
output_logprob = np.array(output_logprob)
|
|
logr = input_logprob - output_logprob
|
|
kl_approx = (np.exp(logr) - 1) - logr
|
|
kl_divs.append(np.mean(kl_approx))
|
|
|
|
print(f"kl_divs={kl_divs}")
|
|
avg_kl_div = sum(kl_divs) / len(kl_divs)
|
|
print(f"avg_kl_div={avg_kl_div}")
|
|
print(f"ACC_THRESHOLDS={ACC_THRESHOLDS[model_name]}")
|
|
assert avg_kl_div < ACC_THRESHOLDS[model_name]["kl_div"], (
|
|
f"avg_kl_div={avg_kl_div} > threshold={ACC_THRESHOLDS[model_name]['kl_div']} "
|
|
f"for {model_name} {test_name}"
|
|
)
|
|
|
|
|
|
# Common request helpers
|
|
def _flush_cache(base_url):
|
|
requests.post(base_url + "/flush_cache")
|
|
|
|
|
|
def _generate(
|
|
base_url, input_ids, max_new_tokens, return_logprob=False, logprob_start_len=-1
|
|
):
|
|
"""Send generate request and return results."""
|
|
json_data = {
|
|
"input_ids": input_ids,
|
|
"sampling_params": {
|
|
"temperature": 1,
|
|
"max_new_tokens": max_new_tokens,
|
|
"ignore_eos": True,
|
|
},
|
|
}
|
|
if return_logprob:
|
|
json_data.update(
|
|
{
|
|
"return_logprob": True,
|
|
"return_text_in_logprobs": False,
|
|
"logprob_start_len": logprob_start_len,
|
|
}
|
|
)
|
|
response = requests.post(base_url + "/generate", json=json_data)
|
|
return response.json()
|
|
|
|
|
|
def _get_input_logprobs(base_url, new_input_ids, output_logprobs):
|
|
"""Run prefill to get input logprobs matching output logprobs."""
|
|
_flush_cache(base_url)
|
|
results = _generate(
|
|
base_url,
|
|
new_input_ids,
|
|
max_new_tokens=0,
|
|
return_logprob=True,
|
|
logprob_start_len=0,
|
|
)
|
|
assert len(results) == len(new_input_ids)
|
|
|
|
input_logprobs = []
|
|
for i, result in enumerate(results):
|
|
logprob = result["meta_info"]["input_token_logprobs"]
|
|
logprob = [x[0] for x in logprob][-len(output_logprobs[i]) :]
|
|
input_logprobs.append(logprob)
|
|
return input_logprobs
|
|
|
|
|
|
def _extract_output_logprobs(result):
|
|
"""Extract output logprobs from a result."""
|
|
return [x[0] for x in result["meta_info"]["output_token_logprobs"]]
|
|
|
|
|
|
def test_input_output_logprobs_match_helper(
|
|
base_url, ACC_THRESHOLDS, model_name, max_samples=None, max_new_tokens=16000
|
|
):
|
|
num_samples = DEFAULT_NUM_SAMPLES
|
|
if max_samples is not None and max_samples > num_samples:
|
|
num_samples = max_samples
|
|
input_ids = get_input_ids(tokenizer_path=model_name, num_samples=num_samples)
|
|
if max_samples is not None:
|
|
input_ids = input_ids[:max_samples]
|
|
print(f"Running test_input_output_logprobs_match with {len(input_ids)} prompts")
|
|
|
|
print("Flush Cache and Running generation to get output logprobs ...")
|
|
_flush_cache(base_url)
|
|
results = _generate(base_url, input_ids, max_new_tokens, return_logprob=True)
|
|
assert len(results) == len(input_ids)
|
|
|
|
new_input_ids = []
|
|
output_logprobs = []
|
|
for i, result in enumerate(results):
|
|
new_input_ids.append(input_ids[i] + result["output_ids"])
|
|
output_logprobs.append(_extract_output_logprobs(result))
|
|
|
|
print("Running prefill to get input logprobs ...")
|
|
input_logprobs = _get_input_logprobs(base_url, new_input_ids, output_logprobs)
|
|
|
|
compare_kl_divergence(
|
|
input_logprobs,
|
|
output_logprobs,
|
|
ACC_THRESHOLDS,
|
|
model_name,
|
|
inspect.currentframe().f_code.co_name,
|
|
)
|
|
|
|
|
|
def test_input_output_logprobs_match_prefill_cache_hit_helper(
|
|
base_url, ACC_THRESHOLDS, model_name, max_samples=None, max_new_tokens=8192
|
|
):
|
|
server_info = requests.get(base_url + "/get_server_info").json()
|
|
if server_info["disable_radix_cache"]:
|
|
print("Radix cache is disabled, skipping test")
|
|
return
|
|
|
|
num_samples = DEFAULT_NUM_SAMPLES
|
|
if max_samples is not None and max_samples > num_samples:
|
|
num_samples = max_samples
|
|
input_ids = get_input_ids(tokenizer_path=model_name, num_samples=num_samples)
|
|
if max_samples is not None:
|
|
input_ids = input_ids[:max_samples]
|
|
print(
|
|
f"Running test_input_output_logprobs_match_prefill_cache_hit with {len(input_ids)} prompts"
|
|
)
|
|
|
|
# Prefill to cache the input
|
|
print("Flush Cache and Prefill to cache the input ...")
|
|
_flush_cache(base_url)
|
|
_generate(base_url, input_ids, max_new_tokens=0)
|
|
|
|
# Generate with cache hit
|
|
print("Running generation to get output logprobs ...")
|
|
results = _generate(base_url, input_ids, max_new_tokens, return_logprob=True)
|
|
assert len(results) == len(input_ids)
|
|
|
|
new_input_ids = []
|
|
output_logprobs = []
|
|
for i, result in enumerate(results):
|
|
if result["meta_info"]["cached_tokens"] == 0:
|
|
print(f"Prefill cache miss for prompt {i}, skipping")
|
|
continue
|
|
new_input_ids.append(input_ids[i] + result["output_ids"])
|
|
output_logprobs.append(_extract_output_logprobs(result))
|
|
|
|
assert len(new_input_ids) > 0.5 * len(
|
|
input_ids
|
|
), f"Too few prefill cache hits: {len(new_input_ids)}/{len(input_ids)}"
|
|
|
|
print("Flush Cache and run prefill to get input logprobs ...")
|
|
input_logprobs = _get_input_logprobs(base_url, new_input_ids, output_logprobs)
|
|
|
|
compare_kl_divergence(
|
|
input_logprobs,
|
|
output_logprobs,
|
|
ACC_THRESHOLDS,
|
|
model_name,
|
|
inspect.currentframe().f_code.co_name,
|
|
)
|
|
|
|
|
|
def test_input_output_logprobs_match_decode_cache_hit_helper(
|
|
base_url, ACC_THRESHOLDS, model_name, max_samples=None, max_new_tokens=8192
|
|
):
|
|
server_info = requests.get(base_url + "/get_server_info").json()
|
|
if server_info["disable_radix_cache"]:
|
|
print("Radix cache is disabled, skipping test")
|
|
return
|
|
|
|
num_samples = DEFAULT_NUM_SAMPLES
|
|
if max_samples is not None and max_samples > num_samples:
|
|
num_samples = max_samples
|
|
first_turn_input_ids = get_input_ids(
|
|
tokenizer_path=model_name, num_samples=num_samples
|
|
)
|
|
if max_samples is not None:
|
|
first_turn_input_ids = first_turn_input_ids[:max_samples]
|
|
print(
|
|
f"Running test_input_output_logprobs_match_decode_cache_hit with {len(first_turn_input_ids)} prompts"
|
|
)
|
|
|
|
# First turn: Prefill + Decode to cache
|
|
print("Flush Cache and First turn: Prefill + Decode to cache decode ...")
|
|
_flush_cache(base_url)
|
|
results = _generate(
|
|
base_url, first_turn_input_ids, max_new_tokens, return_logprob=True
|
|
)
|
|
assert len(results) == len(first_turn_input_ids)
|
|
|
|
tokenizer = get_tokenizer(tokenizer_name=model_name)
|
|
comma_token_id = tokenizer.encode(",")
|
|
|
|
second_turn_input_ids = [
|
|
first_turn_input_ids[i] + result["output_ids"] + comma_token_id
|
|
for i, result in enumerate(results)
|
|
]
|
|
|
|
# Second turn: should hit decode cache
|
|
print("Running generation to get output logprobs ...")
|
|
results = _generate(
|
|
base_url, second_turn_input_ids, max_new_tokens, return_logprob=True
|
|
)
|
|
assert len(results) == len(second_turn_input_ids)
|
|
|
|
new_input_ids = []
|
|
output_logprobs = []
|
|
for i, result in enumerate(results):
|
|
if result["meta_info"]["cached_tokens"] <= len(first_turn_input_ids[i]) + 1:
|
|
print(f"Decode cache miss for prompt {i}, skipping")
|
|
continue
|
|
new_input_ids.append(second_turn_input_ids[i] + result["output_ids"])
|
|
output_logprobs.append(_extract_output_logprobs(result))
|
|
|
|
assert len(new_input_ids) > 0.5 * len(
|
|
second_turn_input_ids
|
|
), f"Too few decode cache hits: {len(new_input_ids)}/{len(second_turn_input_ids)}"
|
|
|
|
print("Flush Cache and run prefill to get input logprobs ...")
|
|
input_logprobs = _get_input_logprobs(base_url, new_input_ids, output_logprobs)
|
|
|
|
compare_kl_divergence(
|
|
input_logprobs,
|
|
output_logprobs,
|
|
ACC_THRESHOLDS,
|
|
model_name,
|
|
inspect.currentframe().f_code.co_name,
|
|
)
|