Add TP2 prefill serving-path smoke experiment
This commit is contained in:
@@ -0,0 +1,50 @@
|
||||
#!/usr/bin/env python3
|
||||
"""Create one exact-length, prefix-disjoint prefill request."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import argparse
|
||||
import json
|
||||
from pathlib import Path
|
||||
|
||||
|
||||
def parse_args() -> argparse.Namespace:
|
||||
parser = argparse.ArgumentParser()
|
||||
parser.add_argument("--model", type=Path, required=True)
|
||||
parser.add_argument("--input-tokens", type=int, default=8192)
|
||||
parser.add_argument("--output-tokens", type=int, default=2)
|
||||
parser.add_argument("--output", type=Path, required=True)
|
||||
return parser.parse_args()
|
||||
|
||||
|
||||
def main() -> None:
|
||||
args = parse_args()
|
||||
if min(args.input_tokens, args.output_tokens) <= 0:
|
||||
raise ValueError("token counts must be positive")
|
||||
if args.input_tokens + args.output_tokens > 40960:
|
||||
raise ValueError("request exceeds the server max model length")
|
||||
|
||||
from transformers import AutoTokenizer
|
||||
|
||||
tokenizer = AutoTokenizer.from_pretrained(args.model, trust_remote_code=True)
|
||||
special = set(tokenizer.all_special_ids)
|
||||
candidates = [
|
||||
token for token in range(tokenizer.vocab_size) if token not in special
|
||||
]
|
||||
if len(candidates) < 2:
|
||||
raise ValueError("tokenizer has too few non-special token IDs")
|
||||
|
||||
body = {
|
||||
"model": "qwen30-prefill-profile",
|
||||
"prompt": [candidates[1], *([candidates[0]] * (args.input_tokens - 1))],
|
||||
"min_tokens": args.output_tokens,
|
||||
"max_tokens": args.output_tokens,
|
||||
"ignore_eos": True,
|
||||
"temperature": 0,
|
||||
}
|
||||
args.output.parent.mkdir(parents=True, exist_ok=True)
|
||||
args.output.write_text(json.dumps(body, separators=(",", ":")) + "\n")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
Reference in New Issue
Block a user