mirror of
https://github.com/deepseek-ai/DeepInfer.git
synced 2026-10-26 16:35:05 +00:00
310 lines
12 KiB
Python
310 lines
12 KiB
Python
from __future__ import annotations
|
|
|
|
import argparse
|
|
import json
|
|
import os
|
|
import sys
|
|
from pathlib import Path
|
|
from types import ModuleType
|
|
from typing import Any
|
|
|
|
EVAL_DIR = Path(__file__).resolve().parent
|
|
if str(EVAL_DIR) not in sys.path:
|
|
sys.path.insert(0, str(EVAL_DIR))
|
|
|
|
from common.io import write_jsonl # noqa: E402
|
|
from common.loader import SUPPORTED, load_bench # noqa: E402
|
|
from common.paths import ensure_sys_path, resolve_out_dir # noqa: E402
|
|
|
|
REPO = ensure_sys_path() # put repo root on sys.path so examples / deepinfer can be imported
|
|
|
|
from examples._config import DEFAULT_MODEL_NAME # noqa: E402
|
|
|
|
BENCH_DEFAULTS = {
|
|
"aime": {"n_sample": 32, "temperature": 1.0, "top_p": 1.0},
|
|
"hle": {"n_sample": 1, "temperature": 1.0, "top_p": 1.0},
|
|
"gpqa": {"n_sample": 8, "temperature": 1.0, "top_p": 1.0},
|
|
}
|
|
|
|
|
|
_REASONING_EFFORT_LEVELS = ("low", "high", "max")
|
|
|
|
|
|
def _reasoning_effort(value: str) -> int | str:
|
|
"""Accept an integer 1-100, or a named level low/high/max."""
|
|
if value in _REASONING_EFFORT_LEVELS:
|
|
return value
|
|
try:
|
|
effort = int(value)
|
|
except ValueError:
|
|
raise argparse.ArgumentTypeError(f"expected an integer 1-100 or one of {_REASONING_EFFORT_LEVELS}, got {value!r}") from None
|
|
if not 1 <= effort <= 100:
|
|
raise argparse.ArgumentTypeError(f"expected 1-100, got {effort}")
|
|
return effort
|
|
|
|
|
|
parser = argparse.ArgumentParser(description="eval benchmark")
|
|
# benchmark setting
|
|
parser.add_argument("--bench", default="aime", choices=list(SUPPORTED))
|
|
parser.add_argument("--model-name", default=DEFAULT_MODEL_NAME, help="HF model id or HF ckpt directory")
|
|
parser.add_argument("--data-file", required=True, help="problem jsonl")
|
|
parser.add_argument("--limit", type=int, default=None, help="run only the first N problems")
|
|
parser.add_argument("--out-dir", default=None)
|
|
parser.add_argument("--tag", default=None, help="output dir suffix, to compare settings of the same model")
|
|
# sample param
|
|
parser.add_argument("--n-sample", type=int, default=None, help="samples per problem (default: aime 32 / hle 1 / gpqa 8)")
|
|
parser.add_argument("--temperature", type=float, default=None, help="default 1.0 for every benchmark")
|
|
parser.add_argument("--top-p", type=float, default=None, help="default 1.0")
|
|
parser.add_argument(
|
|
"--reasoning-effort", type=_reasoning_effort, default=100, help="integer 1-100, or low/high/max; higher means more thinking"
|
|
)
|
|
parser.add_argument("--max-len", type=int, default=1024 * 144, help="max prompt + output length per request")
|
|
# engine
|
|
parser.add_argument("--ep", type=int, default=8, help="expert parallel size = total GPUs in this job")
|
|
parser.add_argument("--cp", type=int, default=1)
|
|
parser.add_argument("--dp", type=int, default=None, help="defaults to ep // cp")
|
|
parser.add_argument("--max-num-seqs", type=int, default=384)
|
|
parser.add_argument("--max-num-batched-tokens", type=int, default=32768)
|
|
parser.add_argument("--gpu-memory-utilization", type=float, default=0.9)
|
|
parser.add_argument(
|
|
"--env", action="append", default=[], metavar="KEY=VALUE", help="extra env vars, repeatable; applied before the engine is built"
|
|
)
|
|
parser.add_argument("--decode-blacklist", default="0,2,129262..129264", help="tokens blocked during generation; empty string disables it")
|
|
parser.add_argument(
|
|
"--batch-invariant", action="store_true", help="DEEPINFER_BATCH_INVARIANT=1: results independent of batch composition (slower)"
|
|
)
|
|
|
|
args = parser.parse_args()
|
|
|
|
# fill in per-benchmark defaults (only for args not passed explicitly)
|
|
for _key, _default in BENCH_DEFAULTS[args.bench].items():
|
|
if getattr(args, _key) is None:
|
|
setattr(args, _key, _default)
|
|
|
|
REASONING_EFFORT = args.reasoning_effort
|
|
|
|
OUT_DIR = resolve_out_dir(
|
|
args.out_dir,
|
|
bench=args.bench,
|
|
tag=args.tag or f"{args.model_name.replace('/', '_')}-n{args.n_sample}-j{args.reasoning_effort}",
|
|
)
|
|
META_PATH = os.path.join(OUT_DIR, "meta.json")
|
|
|
|
|
|
# --------------------------------------------------------------------------- #
|
|
# generation
|
|
# --------------------------------------------------------------------------- #
|
|
|
|
|
|
def _render_prompts(encoder: ModuleType, records: list[dict]) -> list[str]:
|
|
return [
|
|
encoder.encode_messages(
|
|
rec["messages"],
|
|
thinking_mode="thinking",
|
|
reasoning_effort=REASONING_EFFORT,
|
|
)
|
|
for rec in records
|
|
]
|
|
|
|
|
|
def _multi_modal_entry(record: dict) -> dict:
|
|
"""One multi_modal_data entry: image bytes plus dimensions."""
|
|
return {
|
|
"image_bytes": [Path(path).read_bytes() for path in record.get("image_paths") or []],
|
|
"image_sizes": [tuple(size) for size in record.get("image_sizes") or []],
|
|
}
|
|
|
|
|
|
def _setup_env() -> None:
|
|
# MTP / DFlash speculative decoding, set before setup_common_env (which only setdefaults it).
|
|
os.environ["DEEPINFER_USE_NEXTN_DECODE"] = "5"
|
|
os.environ["DEEPINFER_ENABLE_MTP_VARLEN"] = "1"
|
|
|
|
from examples._config import setup_common_env
|
|
|
|
setup_common_env(model_name=args.model_name, sampler_force_topk=0, set_override_experts=True, enable_eplb=True)
|
|
for item in args.env:
|
|
key, _, value = item.partition("=")
|
|
assert key and "=" in item, f"--env expects KEY=VALUE, got {item!r}"
|
|
os.environ[key] = value
|
|
if args.batch_invariant:
|
|
os.environ["DEEPINFER_BATCH_INVARIANT"] = "1"
|
|
if args.decode_blacklist:
|
|
os.environ["DEEPINFER_DECODE_TOKEN_BLACKLIST"] = args.decode_blacklist
|
|
|
|
|
|
def _seed(global_index: int, sample_index: int) -> int:
|
|
"""Per-sample seed: identical seeds produce identical output, so every sample needs its own."""
|
|
return (global_index * args.n_sample + sample_index) % (2**64)
|
|
|
|
|
|
def _rows_from_outputs(tokenizer: Any, plan: list, records: list[dict], outputs: list) -> list[dict]:
|
|
rows = []
|
|
for (problem, sample), generated in zip(plan, outputs, strict=True):
|
|
assert generated, f"{records[problem]['task_id']}#{sample}: generated no tokens"
|
|
text = tokenizer.decode([token.token_id for token in generated], skip_special_tokens=False, clean_up_tokenization_spaces=False)
|
|
rows.append(
|
|
{
|
|
"task_id": records[problem]["task_id"],
|
|
"sample_idx": sample,
|
|
"seq_seed": _seed(problem, sample),
|
|
"prompt_len": generated[0].prompt_len,
|
|
"num_output_tokens": len(generated),
|
|
"finish_reason": generated[-1].finish_reason,
|
|
"raw_text": text,
|
|
}
|
|
)
|
|
return rows
|
|
|
|
|
|
def run_rank(dp_size: int, dp_rank: int) -> None:
|
|
_setup_env()
|
|
|
|
from common.encoder import load_encoder
|
|
|
|
from examples._config import create_llm
|
|
|
|
encoder = load_encoder(args.model_name)
|
|
|
|
bench = load_bench(args.bench)
|
|
problems = bench.load(args.data_file)
|
|
records = bench.build_records(problems, limit=args.limit)
|
|
|
|
# Shard by sample so every rank gets an equal share and no rank inherits a problem's long tail.
|
|
plan = [(problem, sample) for problem in range(len(records)) for sample in range(args.n_sample)][dp_rank::dp_size]
|
|
|
|
llm = create_llm(
|
|
args.model_name,
|
|
dp_rank,
|
|
ep_size=args.ep,
|
|
cp_size=args.cp,
|
|
data_parallel_size=dp_size,
|
|
max_num_batched_tokens=args.max_num_batched_tokens,
|
|
# The engine's runtime_max_model_len is max_model_len + 1, and every request's
|
|
# prompt + max_tokens is already clamped to <= max_len, so a tiny margin is enough.
|
|
max_model_len=args.max_len + 64,
|
|
max_num_seqs=args.max_num_seqs,
|
|
gpu_memory_utilization=args.gpu_memory_utilization,
|
|
lookahead_tokens=min(1, int(os.environ.get("DEEPINFER_USE_NEXTN_DECODE", "0"))),
|
|
decode_sequence=False,
|
|
)
|
|
|
|
try:
|
|
# Import after create_llm: deepinfer imports torch, which must not happen before the engine
|
|
# has set up the accelerator devices.
|
|
from deepinfer.sampling_params import SamplingParams
|
|
|
|
tokenizer = llm.get_tokenizer()
|
|
|
|
if not plan:
|
|
# dp_size > number of requests leaves some ranks with an empty shard: write an empty
|
|
# file and stop, do not build a batch.
|
|
os.makedirs(OUT_DIR, exist_ok=True)
|
|
write_jsonl(os.path.join(OUT_DIR, f"raw_rank{dp_rank:03d}.jsonl"), [])
|
|
print(f"[rank {dp_rank}] empty shard (dp_size={dp_size} > request count), skipping", flush=True)
|
|
return
|
|
|
|
# Render and tokenize each needed prompt once: a rank may use the same prompt for several
|
|
# samples of the same problem.
|
|
needed = sorted({problem for problem, _ in plan})
|
|
texts = dict(zip(needed, _render_prompts(encoder, [records[problem] for problem in needed]), strict=True))
|
|
has_images = any(records[problem].get("image_paths") for problem in needed)
|
|
multi_modal: dict[int, dict] = {}
|
|
if has_images:
|
|
# Replace the <image> placeholder with the image token the engine understands. Image
|
|
# questions then take the string path and the engine expands the token.
|
|
from examples._prompts import get_engine_image_token, replace_image_tokens
|
|
|
|
image_token_str, _ = get_engine_image_token(llm)
|
|
texts = dict(zip(needed, replace_image_tokens([texts[problem] for problem in needed], image_token_str), strict=True))
|
|
# The engine consumes multi_modal_data in place, so every request gets its own copy.
|
|
multi_modal = {problem: _multi_modal_entry(records[problem]) for problem in needed}
|
|
token_ids = {problem: tokenizer.encode(text, add_special_tokens=False) for problem, text in texts.items()}
|
|
|
|
# max_tokens is clamped to the remaining budget: max_len - prompt_len, at least 1.
|
|
sampling_params = [
|
|
SamplingParams(
|
|
temperature=args.temperature,
|
|
top_p=args.top_p,
|
|
max_tokens=max(args.max_len - len(token_ids[problem]), 1),
|
|
ignore_eos=False,
|
|
seq_seed=_seed(problem, sample),
|
|
)
|
|
for problem, sample in plan
|
|
]
|
|
|
|
lens = [len(token_ids[problem]) for problem, _ in plan]
|
|
print(
|
|
f"[rank {dp_rank}] problems={len(needed)} requests={len(plan)} "
|
|
f"prompt_tokens[min/avg/max]={min(lens)}/{sum(lens) // len(lens)}/{max(lens)}"
|
|
+ (f" multimodal={len(plan)}" if has_images else ""),
|
|
flush=True,
|
|
)
|
|
|
|
if has_images:
|
|
outputs = llm.generate(
|
|
prompts=[texts[problem] for problem, _ in plan],
|
|
multi_modal_data=[dict(multi_modal[problem]) for problem, _ in plan],
|
|
sampling_params=sampling_params,
|
|
use_tqdm=False,
|
|
)
|
|
else:
|
|
outputs = llm.generate(
|
|
prompts=None,
|
|
prompt_token_ids=[token_ids[problem] for problem, _ in plan],
|
|
sampling_params=sampling_params,
|
|
use_tqdm=False,
|
|
)
|
|
|
|
rows = _rows_from_outputs(tokenizer, plan, records, outputs)
|
|
|
|
os.makedirs(OUT_DIR, exist_ok=True)
|
|
path = os.path.join(OUT_DIR, f"raw_rank{dp_rank:03d}.jsonl")
|
|
write_jsonl(path, rows)
|
|
print(f"[rank {dp_rank}] wrote {path} rows={len(rows)}", flush=True)
|
|
finally:
|
|
llm.close()
|
|
|
|
|
|
def main() -> int:
|
|
dp_size = max(args.ep // args.cp, 1) if args.dp is None else args.dp
|
|
os.makedirs(OUT_DIR, exist_ok=True)
|
|
|
|
meta = {
|
|
"bench": args.bench,
|
|
"model_name": args.model_name,
|
|
"ep": args.ep,
|
|
"cp": args.cp,
|
|
"dp": dp_size,
|
|
"data_file": args.data_file,
|
|
"limit": args.limit,
|
|
"n_sample": args.n_sample,
|
|
"temperature": args.temperature,
|
|
"top_p": args.top_p,
|
|
"reasoning_effort": REASONING_EFFORT,
|
|
"max_len": args.max_len,
|
|
"mtp": 5,
|
|
"mtp_varlen": True,
|
|
"batch_invariant": args.batch_invariant,
|
|
"decode_blacklist": args.decode_blacklist,
|
|
"max_num_seqs": args.max_num_seqs,
|
|
"max_num_batched_tokens": args.max_num_batched_tokens,
|
|
"env": args.env,
|
|
}
|
|
|
|
for old in Path(OUT_DIR).glob("raw_rank*.jsonl"):
|
|
old.unlink()
|
|
with open(META_PATH, "w") as f:
|
|
json.dump(meta, f, indent=2, ensure_ascii=False)
|
|
|
|
print(f"REPO={REPO}\nOUT_DIR={OUT_DIR}\nmeta={json.dumps(meta, ensure_ascii=False)}", flush=True)
|
|
|
|
from examples._launch import launch_dp
|
|
|
|
launch_dp(run_rank, dp_size)
|
|
return 0
|
|
|
|
|
|
if __name__ == "__main__":
|
|
raise SystemExit(main())
|