| 385 | |
| 386 | |
| 387 | def main(args: argparse.Namespace): |
| 388 | if args.tokenizer is None: |
| 389 | args.tokenizer = args.model |
| 390 | validate_args(args) |
| 391 | if args.seed is None: |
| 392 | args.seed = 0 |
| 393 | random.seed(args.seed) |
| 394 | # Sample the requests. |
| 395 | if args.backend == "hf": |
| 396 | tokenizer = AutoTokenizer.from_pretrained(args.tokenizer, trust_remote_code=args.trust_remote_code) |
| 397 | else: |
| 398 | tokenizer = None |
| 399 | requests = get_requests(args, tokenizer) |
| 400 | # is_multi_modal = any(request.multi_modal_data is not None |
| 401 | # for request in requests) |
| 402 | request_outputs: Optional[list[RequestOutput]] = None |
| 403 | if args.backend == "fastdeploy": |
| 404 | elapsed_time, request_outputs = run_fd( |
| 405 | requests, args.n, EngineArgs.from_cli_args(args), args.disable_detokenize |
| 406 | ) |
| 407 | elif args.backend == "hf": |
| 408 | if not TORCH_AVAILABLE: |
| 409 | raise Exception("PyTorch is not available.") |
| 410 | else: |
| 411 | assert args.tensor_parallel_size == 1 |
| 412 | elapsed_time = run_hf( |
| 413 | requests, |
| 414 | args.model, |
| 415 | tokenizer, |
| 416 | args.n, |
| 417 | args.hf_max_batch_size, |
| 418 | args.trust_remote_code, |
| 419 | args.disable_detokenize, |
| 420 | ) |
| 421 | elif args.backend == "fastdeploy-chat": |
| 422 | elapsed_time, request_outputs = run_fd_chat( |
| 423 | requests, args.n, EngineArgs.from_cli_args(args), args.disable_detokenize |
| 424 | ) |
| 425 | else: |
| 426 | raise ValueError(f"Unknown backend: {args.backend}") |
| 427 | |
| 428 | if request_outputs: |
| 429 | # Note: with the vllm and vllm-chat backends, |
| 430 | # we have request_outputs, which we use to count tokens. |
| 431 | total_prompt_tokens = 0 |
| 432 | total_output_tokens = 0 |
| 433 | for ro in request_outputs: |
| 434 | if not isinstance(ro, RequestOutput): |
| 435 | continue |
| 436 | total_prompt_tokens += len(ro.prompt_token_ids) if ro.prompt_token_ids else 0 |
| 437 | if ro.outputs and hasattr(ro.outputs, "token_ids"): |
| 438 | total_output_tokens += len(ro.outputs.token_ids) |
| 439 | total_num_tokens = total_prompt_tokens + total_output_tokens |
| 440 | else: |
| 441 | total_num_tokens = sum(r.prompt_len + r.expected_output_len for r in requests) |
| 442 | total_output_tokens = sum(r.expected_output_len for r in requests) |
| 443 | total_prompt_tokens = total_num_tokens - total_output_tokens |
| 444 | |