[Speculative Decoding] Add native UNO serving support (#37667)
Co-authored-by: drproduck <drproduck@MacBook-Air-2.local> Co-authored-by: BBuf <1182563586@qq.com>
This commit is contained in:
Yash Akhauri
co-authored by
drproduck
BBuf
drproduck
BBuf
parent
354ed6d66b
commit
2bb25dc18b
@@ -0,0 +1,148 @@
|
||||
# UNO full-dataset math evaluation
|
||||
|
||||
`run_math_eval.py` evaluates AR, UNO, DFLASH, EAGLE, or EAGLE3 with identical
|
||||
datasets, prompts, sampling parameters, and grading. It creates an in-process
|
||||
`sgl.Engine`; there is no separate server process. Engine startup is excluded
|
||||
from the timed interval, and no additional request warmup is run.
|
||||
|
||||
The runner downloads pinned revisions of GSM8K, MATH-500, AIME 2024, AIME
|
||||
2025, and AIME 2026. It applies the same boxed-answer instruction and Qwen
|
||||
reasoning chat template to every engine, then grades with `math_verify`.
|
||||
|
||||
Install SGLang with its evaluation dependencies:
|
||||
|
||||
```bash
|
||||
pip install -e "python[test]"
|
||||
```
|
||||
|
||||
## Reproduce the H200 table
|
||||
|
||||
Run from the SGLang repository root. Each invocation below produces one row of
|
||||
the PR table. GSM8K and MATH-500 use one sample per problem; AIME 2025 uses ten
|
||||
samples per problem, or 300 completions.
|
||||
|
||||
```bash
|
||||
export MODEL_PATH=Qwen/Qwen3-8B
|
||||
export TOKENIZER_PATH=Qwen/Qwen3-8B
|
||||
export UNO_LORA_PATH=s-sahoo/uno-qwen3-8B
|
||||
export DATA_ROOT=/path/to/math-eval-data
|
||||
export RESULT_ROOT=/path/to/math-eval-results
|
||||
|
||||
COMMON_ARGS=(
|
||||
--model-path "$MODEL_PATH"
|
||||
--tokenizer-path "$TOKENIZER_PATH"
|
||||
--data-root "$DATA_ROOT"
|
||||
--context-length 40960
|
||||
--max-tokens 32768
|
||||
--temperature 1
|
||||
--top-k 50
|
||||
--top-p 0.95
|
||||
--random-seed 42
|
||||
)
|
||||
|
||||
run_ar() {
|
||||
local benchmark=$1 samples=$2 requests=$3 output_name=$4
|
||||
PYTHONPATH=python python -m benchmark.uno.run_math_eval \
|
||||
"${COMMON_ARGS[@]}" \
|
||||
--benchmark "$benchmark" \
|
||||
--num-samples "$samples" \
|
||||
--max-running-requests "$requests" \
|
||||
--output-dir "$RESULT_ROOT/$output_name"
|
||||
}
|
||||
|
||||
run_linear_uno() {
|
||||
local benchmark=$1 samples=$2 requests=$3 output_name=$4
|
||||
PYTHONPATH=python python -m benchmark.uno.run_math_eval \
|
||||
"${COMMON_ARGS[@]}" \
|
||||
--benchmark "$benchmark" \
|
||||
--num-samples "$samples" \
|
||||
--max-running-requests "$requests" \
|
||||
--output-dir "$RESULT_ROOT/$output_name" \
|
||||
--speculative-algorithm UNO \
|
||||
--uno-lora-path "$UNO_LORA_PATH" \
|
||||
--speculative-num-steps 1 \
|
||||
--speculative-eagle-topk 1 \
|
||||
--speculative-num-draft-tokens 8
|
||||
}
|
||||
|
||||
run_tree_uno() {
|
||||
local benchmark=$1 samples=$2 requests=$3 output_name=$4
|
||||
PYTHONPATH=python python -m benchmark.uno.run_math_eval \
|
||||
"${COMMON_ARGS[@]}" \
|
||||
--benchmark "$benchmark" \
|
||||
--num-samples "$samples" \
|
||||
--max-running-requests "$requests" \
|
||||
--output-dir "$RESULT_ROOT/$output_name" \
|
||||
--speculative-algorithm UNO \
|
||||
--uno-lora-path "$UNO_LORA_PATH" \
|
||||
--speculative-num-steps 15 \
|
||||
--speculative-eagle-topk 32 \
|
||||
--speculative-num-draft-tokens 32
|
||||
}
|
||||
```
|
||||
|
||||
Run the six batch-64 AR and linear `B/K/V = 8/1/8` rows:
|
||||
|
||||
```bash
|
||||
run_ar gsm8k 1 64 ar-gsm8k-c64
|
||||
run_linear_uno gsm8k 1 64 uno-linear-b8-k1-v8-gsm8k-c64
|
||||
run_ar math500 1 64 ar-math500-c64
|
||||
run_linear_uno math500 1 64 uno-linear-b8-k1-v8-math500-c64
|
||||
run_ar aime25 10 64 ar-aime25-c64
|
||||
run_linear_uno aime25 10 64 uno-linear-b8-k1-v8-aime25-c64
|
||||
```
|
||||
|
||||
Run the six batch-1 AR and tree `B/K/V = 16/32/32` rows:
|
||||
|
||||
```bash
|
||||
run_ar gsm8k 1 1 ar-gsm8k-c1
|
||||
run_tree_uno gsm8k 1 1 uno-tree-b16-k32-v32-gsm8k-c1
|
||||
run_ar math500 1 1 ar-math500-c1
|
||||
run_tree_uno math500 1 1 uno-tree-b16-k32-v32-math500-c1
|
||||
run_ar aime25 10 1 ar-aime25-c1
|
||||
run_tree_uno aime25 10 1 uno-tree-b16-k32-v32-aime25-c1
|
||||
```
|
||||
|
||||
Each output directory contains raw generations, per-answer grades, and
|
||||
`summary.json` and `summary.md`. AR TPF is one. UNO TPF counts both full
|
||||
target-model forwards in each cycle: the diffusion-pathway draft and
|
||||
AR-pathway verification forwards.
|
||||
|
||||
## Other speculative decoders
|
||||
|
||||
The runner uses the same public option names as `sglang serve`. For example,
|
||||
DFLASH can be evaluated with:
|
||||
|
||||
```bash
|
||||
PYTHONPATH=python python -m benchmark.uno.run_math_eval \
|
||||
"${COMMON_ARGS[@]}" \
|
||||
--benchmark math500 \
|
||||
--num-samples 1 \
|
||||
--output-dir "$RESULT_ROOT/dflash-b8-math500-c64" \
|
||||
--max-running-requests 64 \
|
||||
--speculative-algorithm DFLASH \
|
||||
--speculative-draft-model-path z-lab/Qwen3-8B-DFlash-b16 \
|
||||
--speculative-dflash-block-size 8 \
|
||||
--speculative-draft-attention-backend fa3
|
||||
```
|
||||
|
||||
EAGLE or EAGLE3 can be evaluated with the corresponding draft model:
|
||||
|
||||
```bash
|
||||
export EAGLE_DRAFT_MODEL=/path/to/compatible-eagle-draft-model
|
||||
|
||||
PYTHONPATH=python python -m benchmark.uno.run_math_eval \
|
||||
"${COMMON_ARGS[@]}" \
|
||||
--benchmark math500 \
|
||||
--num-samples 1 \
|
||||
--output-dir "$RESULT_ROOT/eagle3-b8-math500-c64" \
|
||||
--max-running-requests 64 \
|
||||
--speculative-algorithm EAGLE3 \
|
||||
--speculative-draft-model-path "$EAGLE_DRAFT_MODEL" \
|
||||
--speculative-num-steps 7 \
|
||||
--speculative-eagle-topk 1 \
|
||||
--speculative-num-draft-tokens 8
|
||||
```
|
||||
|
||||
For EAGLE and DFLASH, TPF follows SGLang's acceptance-length convention and
|
||||
counts generated tokens per target verification forward.
|
||||
@@ -0,0 +1 @@
|
||||
"""Speculative-decoding benchmark utilities."""
|
||||
@@ -0,0 +1,159 @@
|
||||
"""Pinned dataset preparation for speculative-decoding math evaluation."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
from pathlib import Path
|
||||
from typing import Any, NamedTuple
|
||||
|
||||
MATH_INSTRUCTION = "Please reason step by step and put your final answer in \\boxed{}."
|
||||
|
||||
|
||||
class BenchmarkConfig(NamedTuple):
|
||||
name: str
|
||||
expected_rows: int
|
||||
instruction: str
|
||||
chat_template_kwargs: dict[str, object]
|
||||
|
||||
|
||||
BENCHMARKS = {
|
||||
name: BenchmarkConfig(
|
||||
name=name,
|
||||
expected_rows=expected_rows,
|
||||
instruction=MATH_INSTRUCTION,
|
||||
chat_template_kwargs={"reasoning_effort": "high"},
|
||||
)
|
||||
for name, expected_rows in (
|
||||
("gsm8k", 1319),
|
||||
("math500", 500),
|
||||
("aime24", 30),
|
||||
("aime25", 30),
|
||||
("aime26", 30),
|
||||
)
|
||||
}
|
||||
|
||||
DATASET_REVISIONS = {
|
||||
"openai/gsm8k": "740312add88f781978c0658806c59bc2815b9866",
|
||||
"HuggingFaceH4/MATH-500": "6e4ed1a2a79af7d8630a6b768ec859cb5af4d3be",
|
||||
"hypaai/Hypa_AIME2024": "11ab79f0eed5f4fdf3d469b466663ab86bbd77c8",
|
||||
"math-ai/aime25": "563bb8404243c5f09de6ec262f2db674fe5bce9b",
|
||||
"math-ai/aime26": "79037aebdb6580008fb960d17cb21fd3099083e3",
|
||||
}
|
||||
|
||||
|
||||
def get_benchmark(name: str) -> BenchmarkConfig:
|
||||
try:
|
||||
return BENCHMARKS[name]
|
||||
except KeyError as exc:
|
||||
available = ", ".join(BENCHMARKS)
|
||||
raise KeyError(f"Unknown benchmark {name!r}. Available: {available}") from exc
|
||||
|
||||
|
||||
def _load_dataset(
|
||||
repo_id: str,
|
||||
config_name: str | None = None,
|
||||
*,
|
||||
split: str,
|
||||
) -> list[dict[str, Any]]:
|
||||
from datasets import load_dataset
|
||||
|
||||
dataset = load_dataset(
|
||||
repo_id,
|
||||
config_name,
|
||||
split=split,
|
||||
revision=DATASET_REVISIONS[repo_id],
|
||||
)
|
||||
return [dict(row) for row in dataset]
|
||||
|
||||
|
||||
def _prepare_gsm8k() -> list[dict[str, Any]]:
|
||||
rows = _load_dataset("openai/gsm8k", "main", split="test")
|
||||
records = []
|
||||
for index, row in enumerate(rows):
|
||||
prompt = f"Q: {row['question']}\nA: Let's think step by step."
|
||||
records.append(
|
||||
{
|
||||
"row": index,
|
||||
"ground_truth": row["answer"],
|
||||
"chat_input": [{"role": "user", "content": prompt}],
|
||||
}
|
||||
)
|
||||
return records
|
||||
|
||||
|
||||
def _prepare_math500() -> list[dict[str, Any]]:
|
||||
rows = _load_dataset("HuggingFaceH4/MATH-500", split="test")
|
||||
return [
|
||||
{
|
||||
"row": index,
|
||||
"ground_truth": row["answer"],
|
||||
"chat_input": [{"role": "user", "content": row["problem"]}],
|
||||
}
|
||||
for index, row in enumerate(rows)
|
||||
]
|
||||
|
||||
|
||||
def _prepare_aime(rows: list[dict[str, Any]]) -> list[dict[str, Any]]:
|
||||
records = []
|
||||
for index, row in enumerate(rows):
|
||||
prompt = f"Question: {row['problem']}\nAnswer:"
|
||||
records.append(
|
||||
{
|
||||
"row": index,
|
||||
"ground_truth": str(row.get("answer", row.get("solution"))).strip(),
|
||||
"chat_input": [{"role": "user", "content": prompt}],
|
||||
}
|
||||
)
|
||||
return records
|
||||
|
||||
|
||||
def _prepare_aime24() -> list[dict[str, Any]]:
|
||||
return _prepare_aime(_load_dataset("hypaai/Hypa_AIME2024", split="english"))
|
||||
|
||||
|
||||
def _prepare_aime25() -> list[dict[str, Any]]:
|
||||
return _prepare_aime(_load_dataset("math-ai/aime25", split="test"))
|
||||
|
||||
|
||||
def _prepare_aime26() -> list[dict[str, Any]]:
|
||||
return _prepare_aime(_load_dataset("math-ai/aime26", split="test"))
|
||||
|
||||
|
||||
BUILDERS = {
|
||||
"gsm8k": _prepare_gsm8k,
|
||||
"math500": _prepare_math500,
|
||||
"aime24": _prepare_aime24,
|
||||
"aime25": _prepare_aime25,
|
||||
"aime26": _prepare_aime26,
|
||||
}
|
||||
|
||||
|
||||
def _row_count(path: Path) -> int:
|
||||
with path.open("rb") as handle:
|
||||
return sum(bool(line.strip()) for line in handle)
|
||||
|
||||
|
||||
def prepare_benchmark_data(name: str, *, output_dir: Path) -> Path:
|
||||
benchmark = get_benchmark(name)
|
||||
output = output_dir / f"{name}.jsonl"
|
||||
if output.is_file():
|
||||
count = _row_count(output)
|
||||
if count == benchmark.expected_rows:
|
||||
return output
|
||||
raise ValueError(
|
||||
f"{output} has {count} rows; expected {benchmark.expected_rows}"
|
||||
)
|
||||
|
||||
records = BUILDERS[name]()
|
||||
if len(records) != benchmark.expected_rows:
|
||||
raise ValueError(
|
||||
f"Expected {benchmark.expected_rows} rows for {name}, got {len(records)}"
|
||||
)
|
||||
|
||||
output.parent.mkdir(parents=True, exist_ok=True)
|
||||
temporary = output.with_suffix(".jsonl.tmp")
|
||||
with temporary.open("w", encoding="utf-8") as handle:
|
||||
for record in records:
|
||||
handle.write(json.dumps(record, ensure_ascii=False) + "\n")
|
||||
temporary.replace(output)
|
||||
return output
|
||||
@@ -0,0 +1,310 @@
|
||||
"""Minimal math scorer adapted from Nano-vLLM-UNO's Eval360 grader."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import re
|
||||
from typing import Any
|
||||
|
||||
try:
|
||||
from math_verify import parse, verify
|
||||
except ModuleNotFoundError: # Allow CLI help without optional eval dependencies.
|
||||
parse = None
|
||||
verify = None
|
||||
|
||||
|
||||
def _require_math_verify() -> None:
|
||||
if parse is None or verify is None:
|
||||
raise RuntimeError("math scoring requires the 'math_verify' package")
|
||||
|
||||
|
||||
def extract_last_boxed_content(text: str | None) -> str | None:
|
||||
if not text:
|
||||
return None
|
||||
matches = list(re.finditer(r"\\(?:boxed|fbox)\s*\{", text, re.IGNORECASE))
|
||||
for match in reversed(matches):
|
||||
start = match.end()
|
||||
depth = 1
|
||||
index = start
|
||||
while index < len(text) and depth:
|
||||
depth += (text[index] == "{") - (text[index] == "}")
|
||||
index += 1
|
||||
if depth == 0:
|
||||
return text[start : index - 1].strip().strip("$").strip()
|
||||
return None
|
||||
|
||||
|
||||
def _generation_list(row: dict[str, Any]) -> list[str]:
|
||||
if isinstance(row.get("generations"), list):
|
||||
return [str(value or "") for value in row["generations"]]
|
||||
return [str(row.get("generation") or "")]
|
||||
|
||||
|
||||
def _get_accuracy(correct: list[bool]) -> float:
|
||||
return sum(correct) / len(correct) if correct else float("nan")
|
||||
|
||||
|
||||
def _source_key(row: dict[str, Any], fallback: int) -> str:
|
||||
return str(row.get("source_row", row.get("row", fallback)))
|
||||
|
||||
|
||||
def normalize_answer_text(text: str | None) -> str:
|
||||
if text is None:
|
||||
return ""
|
||||
text = str(text).strip().strip("$").strip()
|
||||
gsm8k_answer = re.search(r"####\s*([^\n]+)", text)
|
||||
if gsm8k_answer:
|
||||
text = gsm8k_answer.group(1).strip()
|
||||
text = re.sub(r"\\(?:boxed|fbox)\s*\{(.+)\}\s*$", r"\1", text)
|
||||
text = text.replace("\\dfrac", "\\frac").replace("\\tfrac", "\\frac")
|
||||
text = re.sub(r"\\(?:left|right|bigl|bigr|Bigl|Bigr|big|Big)", "", text)
|
||||
text = text.replace("\\displaystyle", "")
|
||||
text = re.sub(r"\\mathbf\s*\{([^{}]+)\}", r"\1", text)
|
||||
text = re.sub(r"\\mathbf\s+([A-Za-z])", r"\1", text)
|
||||
text = re.sub(r"\\\\\s*\[[^\]]+\]", r"\\\\", text)
|
||||
text = re.sub(r"\\frac\s*\{([^{}]+)\}\s*\{([^{}]+)\}", r"\1/\2", text)
|
||||
text = re.sub(r"\\frac\s*([+-]?\d+)\s*([+-]?\d+)", r"\1/\2", text)
|
||||
text = text.replace("{,}", "")
|
||||
text = re.sub(r"\\(?:,|;|!|:)", "", text)
|
||||
text = text.replace("\\%", "%")
|
||||
text = text.replace("\\$", "").replace("$", "")
|
||||
text = re.sub(r"\\(?:text|mathrm)\s*\{([^{}]*)\}", r"\1", text)
|
||||
text = re.sub(r"\^\{([^{}])\}", r"^\1", text)
|
||||
text = re.sub(r"_\{([^{}])\}", r"_\1", text)
|
||||
text = re.sub(r"\s+", "", text)
|
||||
text = re.sub(r"(?<=\d),(?=\d{3}(?:\D|$))", "", text)
|
||||
return re.sub(r"^[A-Za-z]+=(?=\\begin\{pmatrix\})", "", text)
|
||||
|
||||
|
||||
def _has_plain_variable(text: str) -> bool:
|
||||
text = re.sub(r"\\(?:text|mathrm)\s*\{[^{}]*\}", "", text)
|
||||
text = re.sub(r"\\[A-Za-z]+", "", text)
|
||||
return bool(re.search(r"[A-Za-z]", text))
|
||||
|
||||
|
||||
def _parse_boxed_content(boxed: str) -> Any:
|
||||
_require_math_verify()
|
||||
boxed = normalize_answer_text(boxed)
|
||||
leading_number = re.match(r"\s*([-+]?[0-9][0-9,]*(?:\.[0-9]+)?)", boxed)
|
||||
if leading_number and re.fullmatch(r"[A-Za-z]+", boxed[leading_number.end() :]):
|
||||
answer = leading_number.group(1).replace(",", "")
|
||||
try:
|
||||
return parse(answer)
|
||||
except Exception:
|
||||
return answer
|
||||
if _has_plain_variable(boxed):
|
||||
return boxed
|
||||
try:
|
||||
parsed = parse(f"${boxed}$") or parse(boxed)
|
||||
if parsed:
|
||||
return parsed
|
||||
except Exception:
|
||||
pass
|
||||
if leading_number:
|
||||
answer = leading_number.group(1).replace(",", "")
|
||||
try:
|
||||
return parse(answer)
|
||||
except Exception:
|
||||
return answer
|
||||
return boxed
|
||||
|
||||
|
||||
def _parse_unboxed_answer(text: str) -> Any:
|
||||
_require_math_verify()
|
||||
try:
|
||||
parsed = parse(f"${text}$") or parse(text)
|
||||
if parsed:
|
||||
return parsed
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
patterns = [
|
||||
r"The answer is:?\s*\$?([\-0-9\.,]+)",
|
||||
r"#### ?\$?([\-0-9\.,]+)",
|
||||
r"Therefore,? the answer is:?\s*\$?([\-0-9\.,]+)",
|
||||
r"So,? the answer is:?\s*\$?([\-0-9\.,]+)",
|
||||
r"Thus,? the answer is:?\s*\$?([\-0-9\.,]+)",
|
||||
r"Hence,? the answer is:?\s*\$?([\-0-9\.,]+)",
|
||||
r"Final answer:?\s*\$?([\-0-9\.,]+)",
|
||||
r"The final answer is:?\s*\$?([\-0-9\.,]+)",
|
||||
r"The answer is:?\s*\$?([\-0-9\.,]+)\s*(?:miles?|minutes?|hours?|dollars?|GB)?",
|
||||
(
|
||||
r"=\s*\$?([\-0-9\.,]+)"
|
||||
r"\s*(?:miles?|minutes?|hours?|dollars?|GB)?"
|
||||
r"\.?\s*(?:The answer|$)"
|
||||
),
|
||||
]
|
||||
for pattern in patterns:
|
||||
matches = re.findall(pattern, text, re.IGNORECASE)
|
||||
if matches:
|
||||
answer = matches[-1].replace(",", "").strip().rstrip(".")
|
||||
try:
|
||||
return parse(answer)
|
||||
except Exception:
|
||||
return answer
|
||||
|
||||
sentence_end = (
|
||||
r"(?:is|are|equals?|makes?|has|have|gets?|arrives?|covers?|travels?)"
|
||||
r"\s+\$?([\-0-9\.,]+)(?:\s*(?:miles?|minutes?|hours?|dollars?|GB))?"
|
||||
r"\.?\s*$"
|
||||
)
|
||||
match = re.search(sentence_end, text, re.MULTILINE | re.IGNORECASE)
|
||||
if match:
|
||||
answer = match.group(1).replace(",", "").strip().rstrip(".")
|
||||
try:
|
||||
return parse(answer)
|
||||
except Exception:
|
||||
return answer
|
||||
|
||||
for sentence in reversed(text.split(".")):
|
||||
if "Human:" in sentence or "Assistant:" in sentence:
|
||||
continue
|
||||
numbers = re.findall(r"[-+]?[0-9]*\.?[0-9]+", sentence)
|
||||
if numbers:
|
||||
answer = numbers[-1].lstrip("0") or "0"
|
||||
try:
|
||||
return parse(answer)
|
||||
except Exception:
|
||||
return answer
|
||||
return None
|
||||
|
||||
|
||||
def _parse_answer(text: str) -> Any:
|
||||
boxed = extract_last_boxed_content(text)
|
||||
return _parse_boxed_content(boxed) if boxed else _parse_unboxed_answer(text)
|
||||
|
||||
|
||||
def _answer_candidates(text: str) -> list[Any]:
|
||||
boxed = extract_last_boxed_content(text)
|
||||
candidates = (
|
||||
[_parse_boxed_content(boxed), normalize_answer_text(boxed), boxed]
|
||||
if boxed
|
||||
else [_parse_unboxed_answer(text)]
|
||||
)
|
||||
unique = []
|
||||
seen = set()
|
||||
for candidate in candidates:
|
||||
if candidate is not None and repr(candidate) not in seen:
|
||||
seen.add(repr(candidate))
|
||||
unique.append(candidate)
|
||||
return unique
|
||||
|
||||
|
||||
def _vector_components(text: str) -> tuple[str, ...] | None:
|
||||
matrix = re.fullmatch(r"\\begin\{pmatrix\}(.+)\\end\{pmatrix\}", text)
|
||||
if matrix:
|
||||
return tuple(part for part in matrix.group(1).split(r"\\") if part)
|
||||
tuple_match = re.fullmatch(r"\(([^()]+)\)", text)
|
||||
if tuple_match and "," in tuple_match.group(1):
|
||||
return tuple(part for part in tuple_match.group(1).split(",") if part)
|
||||
return None
|
||||
|
||||
|
||||
def _text_answers_match(answer: str, gold: str) -> bool:
|
||||
answer_norm = normalize_answer_text(answer)
|
||||
gold_norm = normalize_answer_text(gold)
|
||||
if answer_norm == gold_norm:
|
||||
return True
|
||||
try:
|
||||
if float(answer_norm.rstrip("%")) == float(gold_norm.rstrip("%")):
|
||||
return True
|
||||
except ValueError:
|
||||
pass
|
||||
if re.fullmatch(r"\([A-Za-z]\)", gold_norm) and answer_norm == gold_norm[1:-1]:
|
||||
return True
|
||||
answer_parts = _vector_components(answer_norm)
|
||||
gold_parts = _vector_components(gold_norm)
|
||||
return bool(answer_parts and answer_parts == gold_parts)
|
||||
|
||||
|
||||
def _compare_answers(answer: Any, gold: str | None) -> bool:
|
||||
_require_math_verify()
|
||||
if not answer or gold is None:
|
||||
return False
|
||||
if isinstance(answer, str) and _text_answers_match(answer, gold):
|
||||
return True
|
||||
try:
|
||||
if verify(answer, gold):
|
||||
return True
|
||||
gold_answer = _parse_answer(gold)
|
||||
if gold_answer:
|
||||
return bool(verify(gold_answer, answer))
|
||||
except Exception:
|
||||
return isinstance(answer, str) and _text_answers_match(answer, gold)
|
||||
|
||||
|
||||
def grade_math_row(row: dict[str, Any]) -> dict[str, Any]:
|
||||
expected = row.get("ground_truth")
|
||||
expected_values = expected if isinstance(expected, list) else [expected]
|
||||
correct = []
|
||||
parsed_generations = []
|
||||
for generation in _generation_list(row):
|
||||
candidates = _answer_candidates(generation)
|
||||
parsed_generations.append([str(candidate) for candidate in candidates])
|
||||
correct.append(
|
||||
any(
|
||||
_compare_answers(candidate, str(gold))
|
||||
for candidate in candidates
|
||||
for gold in expected_values
|
||||
if gold is not None
|
||||
)
|
||||
)
|
||||
return {
|
||||
**row,
|
||||
"parsed_generations": parsed_generations,
|
||||
"correct": correct,
|
||||
"accuracy": _get_accuracy(correct),
|
||||
"grader": "math",
|
||||
}
|
||||
|
||||
|
||||
def score_math(
|
||||
rows: list[dict[str, Any]],
|
||||
) -> tuple[list[dict[str, Any]], dict[str, Any]]:
|
||||
graded = [grade_math_row(row) for row in rows]
|
||||
correct = [bool(value) for row in graded for value in row["correct"]]
|
||||
by_source: dict[str, list[tuple[int, bool]]] = {}
|
||||
for index, row in enumerate(graded):
|
||||
try:
|
||||
sample_index = int(row.get("sample_index", 0))
|
||||
except (TypeError, ValueError):
|
||||
sample_index = 0
|
||||
samples = by_source.setdefault(_source_key(row, index), [])
|
||||
samples.extend(
|
||||
(sample_index + offset, bool(value))
|
||||
for offset, value in enumerate(row["correct"])
|
||||
)
|
||||
|
||||
per_problem = [
|
||||
[value for _, value in sorted(samples)]
|
||||
for samples in by_source.values()
|
||||
if samples
|
||||
]
|
||||
samples_per_problem = max((len(samples) for samples in per_problem), default=0)
|
||||
sample0 = [samples[0] for samples in per_problem]
|
||||
summary = {
|
||||
"grader": "math",
|
||||
"num_rows": len(graded),
|
||||
"num_problems": len(per_problem),
|
||||
"samples_per_problem": samples_per_problem,
|
||||
"num_correct": sum(correct),
|
||||
"accuracy": _get_accuracy(correct),
|
||||
}
|
||||
if per_problem:
|
||||
summary.update(
|
||||
avg_at_1=_get_accuracy(sample0),
|
||||
pass_at_1=_get_accuracy(sample0),
|
||||
num_correct_at_1=sum(sample0),
|
||||
)
|
||||
if samples_per_problem > 1:
|
||||
summary.update(
|
||||
{
|
||||
f"avg_at_{samples_per_problem}": _get_accuracy(correct),
|
||||
f"pass_at_{samples_per_problem}": _get_accuracy(
|
||||
[any(samples) for samples in per_problem]
|
||||
),
|
||||
f"num_pass_at_{samples_per_problem}": sum(
|
||||
any(samples) for samples in per_problem
|
||||
),
|
||||
}
|
||||
)
|
||||
return graded, summary
|
||||
@@ -0,0 +1,517 @@
|
||||
"""Run full-dataset AR or speculative-decoding math evaluation offline."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import argparse
|
||||
import json
|
||||
from collections import defaultdict
|
||||
from pathlib import Path
|
||||
from time import perf_counter
|
||||
from typing import Any
|
||||
|
||||
from benchmark.uno.math_data import (
|
||||
BENCHMARKS,
|
||||
BenchmarkConfig,
|
||||
get_benchmark,
|
||||
prepare_benchmark_data,
|
||||
)
|
||||
from benchmark.uno.math_grader import score_math
|
||||
|
||||
CONTEXT_LENGTH = 40960
|
||||
MAX_TOKENS = 2**15
|
||||
SPECULATIVE_ALGORITHMS = ("EAGLE", "EAGLE3", "DFLASH", "UNO")
|
||||
SPECULATIVE_OPTION_NAMES = (
|
||||
"speculative_algorithm",
|
||||
"speculative_draft_model_path",
|
||||
"speculative_draft_model_revision",
|
||||
"speculative_num_steps",
|
||||
"speculative_eagle_topk",
|
||||
"speculative_num_draft_tokens",
|
||||
"speculative_dflash_block_size",
|
||||
"speculative_draft_attention_backend",
|
||||
"uno_lora_path",
|
||||
)
|
||||
|
||||
|
||||
def parse_args() -> argparse.Namespace:
|
||||
parser = argparse.ArgumentParser(description=__doc__)
|
||||
parser.add_argument("--model-path", default="Qwen/Qwen3-8B")
|
||||
parser.add_argument("--tokenizer-path")
|
||||
parser.add_argument("--revision")
|
||||
parser.add_argument("--dtype", default="bfloat16")
|
||||
parser.add_argument("--attention-backend", default="fa3")
|
||||
parser.add_argument("--data-root", type=Path, required=True)
|
||||
parser.add_argument("--output-dir", type=Path, required=True)
|
||||
parser.add_argument(
|
||||
"--speculative-algorithm",
|
||||
type=str.upper,
|
||||
choices=SPECULATIVE_ALGORITHMS,
|
||||
)
|
||||
parser.add_argument("--speculative-draft-model-path")
|
||||
parser.add_argument("--speculative-draft-model-revision")
|
||||
parser.add_argument("--speculative-num-steps", type=int)
|
||||
parser.add_argument("--speculative-eagle-topk", type=int)
|
||||
parser.add_argument("--speculative-num-draft-tokens", type=int)
|
||||
parser.add_argument("--speculative-dflash-block-size", type=int)
|
||||
parser.add_argument("--speculative-draft-attention-backend")
|
||||
parser.add_argument("--uno-lora-path")
|
||||
parser.add_argument(
|
||||
"--benchmark",
|
||||
action="append",
|
||||
choices=tuple(BENCHMARKS),
|
||||
help="Benchmark to run; repeat to select multiple (default: all).",
|
||||
)
|
||||
parser.add_argument("--limit", type=int, help="Maximum problems per benchmark.")
|
||||
parser.add_argument("--num-samples", type=int, default=1)
|
||||
parser.add_argument("--max-running-requests", type=int, default=4)
|
||||
parser.add_argument("--context-length", type=int, default=CONTEXT_LENGTH)
|
||||
parser.add_argument("--max-tokens", type=int, default=MAX_TOKENS)
|
||||
parser.add_argument("--temperature", type=float, default=1.0)
|
||||
parser.add_argument("--top-k", type=int, default=50)
|
||||
parser.add_argument("--top-p", type=float, default=0.95)
|
||||
parser.add_argument("--random-seed", type=int, default=42)
|
||||
args = parser.parse_args()
|
||||
_validate_args(parser=parser, args=args)
|
||||
return args
|
||||
|
||||
|
||||
def _validate_args(
|
||||
*, parser: argparse.ArgumentParser, args: argparse.Namespace
|
||||
) -> None:
|
||||
for name in (
|
||||
"num_samples",
|
||||
"max_running_requests",
|
||||
"context_length",
|
||||
"max_tokens",
|
||||
"top_k",
|
||||
):
|
||||
if getattr(args, name) <= 0:
|
||||
parser.error(f"--{name.replace('_', '-')} must be positive")
|
||||
if args.limit is not None and args.limit <= 0:
|
||||
parser.error("--limit must be positive")
|
||||
if not 0 < args.top_p <= 1:
|
||||
parser.error("--top-p must be in (0, 1]")
|
||||
if args.temperature < 0:
|
||||
parser.error("--temperature must be non-negative")
|
||||
|
||||
|
||||
def _context_reserve(args: argparse.Namespace) -> int:
|
||||
width = args.speculative_num_draft_tokens
|
||||
if width is None:
|
||||
width = args.speculative_dflash_block_size or 1
|
||||
if args.speculative_algorithm == "DFLASH":
|
||||
return 2 * width
|
||||
if args.speculative_algorithm == "UNO" and (args.speculative_eagle_topk or 1) > 1:
|
||||
draft_width = (args.speculative_num_steps or 0) + 1
|
||||
return max(width, draft_width) + 1
|
||||
return width
|
||||
|
||||
|
||||
def _model_ref(value: str) -> str:
|
||||
path = Path(value).expanduser()
|
||||
return str(path.resolve()) if path.exists() else value
|
||||
|
||||
|
||||
def _read_jsonl(path: Path) -> list[dict[str, Any]]:
|
||||
with path.open(encoding="utf-8") as handle:
|
||||
return [json.loads(line) for line in handle if line.strip()]
|
||||
|
||||
|
||||
def _first_value(row: dict[str, Any], *names: str) -> Any | None:
|
||||
return next((row[name] for name in names if row.get(name) is not None), None)
|
||||
|
||||
|
||||
def _format_prompt(
|
||||
*,
|
||||
tokenizer: Any,
|
||||
messages: list[dict[str, Any]],
|
||||
benchmark: BenchmarkConfig,
|
||||
) -> tuple[list[int], str]:
|
||||
messages = [dict(message) for message in messages]
|
||||
system_message = next(
|
||||
(message for message in messages if message.get("role") == "system"),
|
||||
None,
|
||||
)
|
||||
if system_message is None:
|
||||
messages.insert(0, {"role": "system", "content": benchmark.instruction})
|
||||
elif benchmark.instruction not in system_message["content"]:
|
||||
existing = system_message["content"].strip()
|
||||
system_message["content"] = f"{benchmark.instruction}\n\n{existing}"
|
||||
|
||||
rendered = tokenizer.apply_chat_template(
|
||||
messages,
|
||||
tokenize=False,
|
||||
add_generation_prompt=True,
|
||||
**benchmark.chat_template_kwargs,
|
||||
)
|
||||
token_ids = tokenizer.apply_chat_template(
|
||||
messages,
|
||||
tokenize=True,
|
||||
add_generation_prompt=True,
|
||||
return_dict=False,
|
||||
**benchmark.chat_template_kwargs,
|
||||
)
|
||||
return list(token_ids), str(rendered)
|
||||
|
||||
|
||||
def _prepare_prompts(
|
||||
*,
|
||||
args: argparse.Namespace,
|
||||
tokenizer: Any,
|
||||
context_reserve: int,
|
||||
) -> list[dict[str, Any]]:
|
||||
prompts = []
|
||||
for benchmark_name in args.benchmark or BENCHMARKS:
|
||||
benchmark = get_benchmark(benchmark_name)
|
||||
data_path = prepare_benchmark_data(
|
||||
benchmark.name,
|
||||
output_dir=args.data_root,
|
||||
)
|
||||
rows = _read_jsonl(data_path)
|
||||
rows = rows[: args.limit] if args.limit is not None else rows
|
||||
for row_index, row in enumerate(rows):
|
||||
input_ids, rendered = _format_prompt(
|
||||
tokenizer=tokenizer,
|
||||
messages=row["chat_input"],
|
||||
benchmark=benchmark,
|
||||
)
|
||||
available = args.context_length - len(input_ids) - context_reserve
|
||||
if available < args.max_tokens:
|
||||
source = _first_value(row, "id", "problem_id", "index", "row")
|
||||
source = row_index if source is None else source
|
||||
raise ValueError(
|
||||
f"{benchmark.name}:{source} has only {max(0, available)} "
|
||||
"completion tokens available; increase --context-length"
|
||||
)
|
||||
source = _first_value(row, "id", "problem_id", "index", "row")
|
||||
source = row_index if source is None else source
|
||||
for sample_index in range(args.num_samples):
|
||||
prompt_id = f"{source}:sample{sample_index}"
|
||||
prompts.append(
|
||||
{
|
||||
"id": prompt_id,
|
||||
"benchmark": benchmark.name,
|
||||
"input_ids": input_ids,
|
||||
"row": {
|
||||
"id": prompt_id,
|
||||
"source_row": source,
|
||||
"sample_index": sample_index,
|
||||
"problem": rendered,
|
||||
"ground_truth": row["ground_truth"],
|
||||
"prompt_token_count": len(input_ids),
|
||||
"resolved_max_tokens": args.max_tokens,
|
||||
},
|
||||
}
|
||||
)
|
||||
if not prompts:
|
||||
raise ValueError("No prompts were prepared")
|
||||
return prompts
|
||||
|
||||
|
||||
def _engine_options(
|
||||
*,
|
||||
args: argparse.Namespace,
|
||||
) -> dict[str, Any]:
|
||||
tokenizer_path = _model_ref(args.tokenizer_path or args.model_path)
|
||||
options: dict[str, Any] = {
|
||||
"model_path": _model_ref(args.model_path),
|
||||
"tokenizer_path": tokenizer_path,
|
||||
"skip_tokenizer_init": True,
|
||||
"context_length": args.context_length,
|
||||
"dtype": args.dtype,
|
||||
"random_seed": args.random_seed,
|
||||
"max_running_requests": args.max_running_requests,
|
||||
"attention_backend": args.attention_backend,
|
||||
"log_level": "info",
|
||||
}
|
||||
if args.revision is not None:
|
||||
options["revision"] = args.revision
|
||||
for name in SPECULATIVE_OPTION_NAMES:
|
||||
value = getattr(args, name)
|
||||
if value is not None:
|
||||
if name in (
|
||||
"speculative_draft_model_path",
|
||||
"uno_lora_path",
|
||||
):
|
||||
value = _model_ref(value)
|
||||
options[name] = value
|
||||
return options
|
||||
|
||||
|
||||
def _generate(
|
||||
*,
|
||||
args: argparse.Namespace,
|
||||
prompts: list[dict[str, Any]],
|
||||
engine_options: dict[str, Any],
|
||||
) -> tuple[list[dict[str, Any]], float]:
|
||||
import sglang as sgl
|
||||
|
||||
engine = sgl.Engine(**engine_options)
|
||||
sampling = [
|
||||
{
|
||||
"max_new_tokens": args.max_tokens,
|
||||
"temperature": args.temperature,
|
||||
"top_k": args.top_k,
|
||||
"top_p": args.top_p,
|
||||
}
|
||||
for _ in prompts
|
||||
]
|
||||
try:
|
||||
start = perf_counter()
|
||||
outputs = engine.generate(
|
||||
input_ids=[prompt["input_ids"] for prompt in prompts],
|
||||
sampling_params=sampling,
|
||||
rid=[f"{prompt['benchmark']}:{prompt['id']}" for prompt in prompts],
|
||||
)
|
||||
elapsed = perf_counter() - start
|
||||
finally:
|
||||
engine.shutdown()
|
||||
return outputs, elapsed
|
||||
|
||||
|
||||
def _build_rows(
|
||||
*,
|
||||
prompts: list[dict[str, Any]],
|
||||
outputs: list[dict[str, Any]],
|
||||
tokenizer: Any,
|
||||
speculative_algorithm: str | None,
|
||||
) -> list[dict[str, Any]]:
|
||||
rows = []
|
||||
for prompt, output in zip(prompts, outputs, strict=True):
|
||||
token_ids = list(output["output_ids"])
|
||||
metadata = output["meta_info"]
|
||||
verify_forwards = int(metadata.get("spec_verify_ct", 0))
|
||||
if speculative_algorithm == "UNO":
|
||||
# Both UNO pathways are full target-model forward passes.
|
||||
num_forwards = 2 * verify_forwards
|
||||
elif speculative_algorithm is not None:
|
||||
# EAGLE and DFLASH TPF conventionally count target verification.
|
||||
num_forwards = verify_forwards
|
||||
else:
|
||||
num_forwards = len(token_ids)
|
||||
rows.append(
|
||||
{
|
||||
"benchmark": prompt["benchmark"],
|
||||
**prompt["row"],
|
||||
"output_ids": token_ids,
|
||||
"generation": tokenizer.decode(token_ids, skip_special_tokens=True),
|
||||
"num_tokens": len(token_ids),
|
||||
"num_forwards": num_forwards,
|
||||
"tokens_per_forward": _divide(len(token_ids), num_forwards),
|
||||
"sglang_meta_info": metadata,
|
||||
}
|
||||
)
|
||||
return rows
|
||||
|
||||
|
||||
def _divide(numerator: int | float, denominator: int | float) -> float | None:
|
||||
return numerator / denominator if denominator else None
|
||||
|
||||
|
||||
def _write_json(path: Path, value: dict[str, Any]) -> None:
|
||||
path.parent.mkdir(parents=True, exist_ok=True)
|
||||
path.write_text(
|
||||
json.dumps(value, indent=2, ensure_ascii=False) + "\n",
|
||||
encoding="utf-8",
|
||||
)
|
||||
|
||||
|
||||
def _write_jsonl(path: Path, rows: list[dict[str, Any]]) -> None:
|
||||
path.parent.mkdir(parents=True, exist_ok=True)
|
||||
with path.open("w", encoding="utf-8") as handle:
|
||||
for row in rows:
|
||||
handle.write(json.dumps(row, ensure_ascii=False) + "\n")
|
||||
|
||||
|
||||
def _benchmark_summary(rows: list[dict[str, Any]]) -> dict[str, Any]:
|
||||
tokens = sum(row["num_tokens"] for row in rows)
|
||||
forwards = sum(row["num_forwards"] for row in rows)
|
||||
tpfs = [
|
||||
row["tokens_per_forward"]
|
||||
for row in rows
|
||||
if row["tokens_per_forward"] is not None
|
||||
]
|
||||
return {
|
||||
"num_tokens": tokens,
|
||||
"num_forwards": forwards,
|
||||
"tokens_per_forward": _divide(tokens, forwards),
|
||||
"unweighted_mean_tokens_per_forward": (sum(tpfs) / len(tpfs) if tpfs else None),
|
||||
}
|
||||
|
||||
|
||||
def _write_markdown(path: Path, summary: dict[str, Any]) -> None:
|
||||
lines = [
|
||||
"| Dataset | Accuracy | TPF | tok/s | tok/s/request |",
|
||||
"| --- | ---: | ---: | ---: | ---: |",
|
||||
]
|
||||
for name, metrics in summary["by_benchmark"].items():
|
||||
tokens_per_second = metrics.get("tokens_per_second")
|
||||
per_request = metrics.get("tokens_per_second_per_request")
|
||||
tokens_per_forward = metrics.get("tokens_per_forward")
|
||||
tps = f"{tokens_per_second:.2f}" if tokens_per_second is not None else "—"
|
||||
request_tps = f"{per_request:.2f}" if per_request is not None else "—"
|
||||
tpf = f"{tokens_per_forward:.3f}" if tokens_per_forward is not None else "—"
|
||||
lines.append(
|
||||
f"| {name} | {metrics['accuracy']:.2%} | {tpf} | {tps} | {request_tps} |"
|
||||
)
|
||||
lines.extend(
|
||||
[
|
||||
"| **Average** | "
|
||||
f"**{summary['macro_average_accuracy']:.2%}** | "
|
||||
f"**{summary['macro_average_tokens_per_forward']:.3f}** | "
|
||||
f"**{summary['tokens_per_second']:.2f}** | "
|
||||
f"**{summary['tokens_per_second_per_request']:.2f}** |",
|
||||
"",
|
||||
"Accuracy and TPF in the Average row are unweighted means across "
|
||||
"datasets. Throughput is aggregate output tokens divided by timed "
|
||||
"generation seconds; tok/s/request divides it by max running requests.",
|
||||
]
|
||||
)
|
||||
path.write_text("\n".join(lines) + "\n", encoding="utf-8")
|
||||
|
||||
|
||||
def _write_results(
|
||||
*,
|
||||
args: argparse.Namespace,
|
||||
rows: list[dict[str, Any]],
|
||||
generation_seconds: float,
|
||||
engine_options: dict[str, Any],
|
||||
) -> None:
|
||||
rows_by_benchmark: dict[str, list[dict[str, Any]]] = defaultdict(list)
|
||||
for row in rows:
|
||||
output_row = {key: value for key, value in row.items() if key != "benchmark"}
|
||||
rows_by_benchmark[row["benchmark"]].append(output_row)
|
||||
|
||||
benchmark_metrics = {}
|
||||
totals = {"prompts": 0, "completions": 0, "correct": 0, "tokens": 0, "forwards": 0}
|
||||
request_tpfs = []
|
||||
for name, benchmark_rows in rows_by_benchmark.items():
|
||||
output_dir = args.output_dir / name
|
||||
_write_jsonl(output_dir / "generations.jsonl", benchmark_rows)
|
||||
graded, scores = score_math(benchmark_rows)
|
||||
_write_jsonl(output_dir / "grades.jsonl", graded)
|
||||
_write_json(output_dir / "scores.json", scores)
|
||||
|
||||
metrics = {
|
||||
"num_prompts": scores["num_problems"],
|
||||
"num_completions": scores["num_rows"],
|
||||
"num_correct": scores["num_correct"],
|
||||
"accuracy": scores["accuracy"],
|
||||
**_benchmark_summary(benchmark_rows),
|
||||
}
|
||||
benchmark_metrics[name] = metrics
|
||||
totals["prompts"] += scores["num_problems"]
|
||||
totals["completions"] += scores["num_rows"]
|
||||
totals["correct"] += scores["num_correct"]
|
||||
totals["tokens"] += metrics["num_tokens"]
|
||||
totals["forwards"] += metrics["num_forwards"]
|
||||
request_tpfs.extend(
|
||||
row["tokens_per_forward"]
|
||||
for row in benchmark_rows
|
||||
if row["tokens_per_forward"] is not None
|
||||
)
|
||||
|
||||
tokens_per_second = totals["tokens"] / generation_seconds
|
||||
tokens_per_second_per_request = tokens_per_second / args.max_running_requests
|
||||
if len(benchmark_metrics) == 1:
|
||||
only_metrics = next(iter(benchmark_metrics.values()))
|
||||
only_metrics["tokens_per_second"] = tokens_per_second
|
||||
only_metrics["tokens_per_second_per_request"] = tokens_per_second_per_request
|
||||
|
||||
dataset_accuracies = [value["accuracy"] for value in benchmark_metrics.values()]
|
||||
dataset_tpfs = [
|
||||
value["tokens_per_forward"]
|
||||
for value in benchmark_metrics.values()
|
||||
if value["tokens_per_forward"] is not None
|
||||
]
|
||||
summary = {
|
||||
"engine": "sglang-offline",
|
||||
"mode": _mode_name(args),
|
||||
"model_path": engine_options["model_path"],
|
||||
"tokenizer_path": engine_options["tokenizer_path"],
|
||||
"revision": args.revision,
|
||||
"dtype": args.dtype,
|
||||
"attention_backend": args.attention_backend,
|
||||
"speculative_parameters": {
|
||||
name: engine_options[name]
|
||||
for name in SPECULATIVE_OPTION_NAMES
|
||||
if name in engine_options
|
||||
},
|
||||
"max_running_requests": args.max_running_requests,
|
||||
"benchmarks": list(rows_by_benchmark),
|
||||
"context_length": args.context_length,
|
||||
"max_tokens": args.max_tokens,
|
||||
"num_samples": args.num_samples,
|
||||
"random_seed": args.random_seed,
|
||||
"sampling": {
|
||||
"temperature": args.temperature,
|
||||
"top_k": args.top_k,
|
||||
"top_p": args.top_p,
|
||||
},
|
||||
"num_prompts": totals["prompts"],
|
||||
"num_completions": totals["completions"],
|
||||
"num_correct": totals["correct"],
|
||||
"accuracy": _divide(totals["correct"], totals["completions"]),
|
||||
"macro_average_accuracy": sum(dataset_accuracies) / len(dataset_accuracies),
|
||||
"generation_seconds": generation_seconds,
|
||||
"num_tokens": totals["tokens"],
|
||||
"tokens_per_second": tokens_per_second,
|
||||
"tokens_per_second_per_request": tokens_per_second_per_request,
|
||||
"num_forwards": totals["forwards"],
|
||||
"tokens_per_forward": _divide(totals["tokens"], totals["forwards"]),
|
||||
"macro_average_tokens_per_forward": sum(dataset_tpfs) / len(dataset_tpfs),
|
||||
"unweighted_mean_tokens_per_forward": (
|
||||
sum(request_tpfs) / len(request_tpfs) if request_tpfs else None
|
||||
),
|
||||
"by_benchmark": benchmark_metrics,
|
||||
}
|
||||
_write_json(args.output_dir / "summary.json", summary)
|
||||
_write_markdown(args.output_dir / "summary.md", summary)
|
||||
print(json.dumps(summary, indent=2))
|
||||
|
||||
|
||||
def _mode_name(args: argparse.Namespace) -> str:
|
||||
algorithm = args.speculative_algorithm
|
||||
if algorithm is None:
|
||||
return "ar"
|
||||
if algorithm == "UNO":
|
||||
return "uno-tree" if (args.speculative_eagle_topk or 1) > 1 else "uno-linear"
|
||||
return algorithm.lower()
|
||||
|
||||
|
||||
def main() -> None:
|
||||
args = parse_args()
|
||||
from transformers import AutoTokenizer
|
||||
|
||||
engine_options = _engine_options(args=args)
|
||||
tokenizer = AutoTokenizer.from_pretrained(
|
||||
engine_options["tokenizer_path"],
|
||||
use_fast=True,
|
||||
trust_remote_code=True,
|
||||
)
|
||||
prompts = _prepare_prompts(
|
||||
args=args,
|
||||
tokenizer=tokenizer,
|
||||
context_reserve=_context_reserve(args),
|
||||
)
|
||||
outputs, elapsed = _generate(
|
||||
args=args,
|
||||
prompts=prompts,
|
||||
engine_options=engine_options,
|
||||
)
|
||||
rows = _build_rows(
|
||||
prompts=prompts,
|
||||
outputs=outputs,
|
||||
tokenizer=tokenizer,
|
||||
speculative_algorithm=args.speculative_algorithm,
|
||||
)
|
||||
_write_results(
|
||||
args=args,
|
||||
rows=rows,
|
||||
generation_seconds=elapsed,
|
||||
engine_options=engine_options,
|
||||
)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
Reference in New Issue
Block a user