diff --git a/tensorrt_llm/bench/benchmark/throughput.py b/tensorrt_llm/bench/benchmark/throughput.py index ff72ccb93803..3c0dfd5912d9 100755 --- a/tensorrt_llm/bench/benchmark/throughput.py +++ b/tensorrt_llm/bench/benchmark/throughput.py @@ -434,6 +434,12 @@ def throughput_command( post_proc_params = None # No detokenization + has_multi_turn = any(r.is_multi_turn for r in requests) + multi_turn_tokenizer = tokenizer if has_multi_turn else None + if has_multi_turn: + logger.info("Multi-turn requests detected. Turns will be processed " + "sequentially within each request.") + # Perform warmup if requested. if options.warmup > 0: logger.info("Setting up for warmup...") @@ -446,7 +452,8 @@ def throughput_command( warmup_dataset, False, options.concurrency, - modality=options.modality)) + modality=options.modality, + tokenizer=multi_turn_tokenizer)) # WAR: IterationResult is a singleton tied to the executor. # Since the benchmark calls asyncio.run() multiple times (e.g., during warmup), # we must reset it to ensure it attaches to the correct event loop. @@ -463,7 +470,8 @@ def throughput_command( options.streaming, options.concurrency, iteration_writer.full_address, - modality=options.modality)) + modality=options.modality, + tokenizer=multi_turn_tokenizer)) logger.info("Benchmark done. Reporting results...") if options.modality is not None: diff --git a/tensorrt_llm/bench/benchmark/utils/asynchronous.py b/tensorrt_llm/bench/benchmark/utils/asynchronous.py index cc84013b0d96..9d8c0ca70688 100644 --- a/tensorrt_llm/bench/benchmark/utils/asynchronous.py +++ b/tensorrt_llm/bench/benchmark/utils/asynchronous.py @@ -1,12 +1,14 @@ from __future__ import annotations import asyncio +import copy import time from contextlib import asynccontextmanager from itertools import chain from typing import Dict, List, Optional, Set, Tuple import tqdm +from transformers import PreTrainedTokenizer from zmq import PUSH from zmq.asyncio import Context @@ -28,7 +30,8 @@ def __init__(self, outbox: asyncio.Queue[PerfItemTuple], streaming: bool, concurrency: int = -1, - modality: Optional[str] = None) -> None: + modality: Optional[str] = None, + tokenizer: Optional[PreTrainedTokenizer] = None) -> None: self.llm = llm self._inbox: asyncio.Queue[Tuple[InferenceRequest, SamplingParams]] = asyncio.Queue() @@ -45,18 +48,28 @@ def __init__(self, self.streaming = streaming self.request_seen = asyncio.Event() self.modality = modality + self.tokenizer = tokenizer async def process_request(self, request: InferenceRequest, sampling_params: SamplingParams, post_proc_params: PostprocParams): - # Set up sampling params with inference request + if request.is_multi_turn and self.tokenizer is not None: + await self._process_multi_turn_request(request, sampling_params, + post_proc_params) + else: + await self._process_single_request(request, sampling_params, + post_proc_params) + + async def _process_single_request(self, request: InferenceRequest, + sampling_params: SamplingParams, + post_proc_params: PostprocParams): self.request_seen.set() + sampling_params = copy.copy(sampling_params) sampling_params.max_tokens = request.output_tokens async with semaphore_guard(self._concurrency_semaphore): request_start_timestamp = time.perf_counter_ns() time_on_first_token = None - # Schedule the request in the LLM API (asynchronously) logger.debug(f"request.lora_request: {request.lora_request}") output: RequestOutput = self.llm.generate_async( request.input_ids if self.modality is None else request.prompt, @@ -70,12 +83,10 @@ async def process_request(self, request: InferenceRequest, time_on_first_token = time.perf_counter_ns() response = stream_output else: - # Wait for the response to return to us. response: RequestOutput = await output.aresult() response_end_timestamp = time.perf_counter_ns() - # Mark that the response returned. Construct a record to send to statistics. tokens = list(chain(*(beam.token_ids for beam in response.outputs))) request_perf_item = PerfItemTuple( start_timestamp=request_start_timestamp, @@ -89,7 +100,79 @@ async def process_request(self, request: InferenceRequest, time_on_first_token=time_on_first_token, ) - # Register the new request perf items in the outbound queue for statistics keeping + await self._outbox.put(request_perf_item) + + async def _process_multi_turn_request(self, request: InferenceRequest, + sampling_params: SamplingParams, + post_proc_params: PostprocParams): + """Process a multi-turn request by iterating through turns sequentially. + + Each turn builds on the previous conversation context: the model's + response from turn N is appended to the messages before encoding + turn N+1. All turns within a single request share one concurrency + slot so that the conversation history stays consistent. + """ + self.request_seen.set() + sampling_params = copy.copy(sampling_params) + sampling_params.max_tokens = request.output_tokens + tokenizer = self.tokenizer + loop = asyncio.get_running_loop() + + messages: List[dict] = [] + total_input_tokens = 0 + all_output_tokens: List[int] = [] + + async with semaphore_guard(self._concurrency_semaphore): + request_start_timestamp = time.perf_counter_ns() + time_on_first_token = None + last_response = None + + for turn_id, question in enumerate(request.turns): + messages.append({"role": "user", "content": question}) + + input_ids = await loop.run_in_executor( + None, lambda: tokenizer.apply_chat_template( + messages, add_generation_prompt=True)) + + output: RequestOutput = self.llm.generate_async( + input_ids, + sampling_params=sampling_params, + _postproc_params=post_proc_params, + streaming=False) + response: RequestOutput = await output.aresult() + + if turn_id == 0 and time_on_first_token is None: + time_on_first_token = time.perf_counter_ns() + + turn_tokens = list( + chain(*(beam.token_ids for beam in response.outputs))) + all_output_tokens.extend(turn_tokens) + total_input_tokens += len(input_ids) + + assistant_text = await loop.run_in_executor( + None, lambda: tokenizer.decode(turn_tokens, + skip_special_tokens=True)) + messages.append({ + "role": "assistant", + "content": assistant_text + }) + + last_response = response + + response_end_timestamp = time.perf_counter_ns() + + request_perf_item = PerfItemTuple( + start_timestamp=request_start_timestamp, + end_timestamp=response_end_timestamp, + request_id=last_response.id, + num_input_tokens=total_input_tokens, + response_is_final=last_response.finished, + error=False, + tokens=all_output_tokens, + decoding_iteration=last_response.decoding_iter, + time_on_first_token=time_on_first_token, + ) + await self._outbox.put(request_perf_item) def _raise_for_failed_tasks(self): @@ -256,6 +339,7 @@ async def async_benchmark( concurrency: int = -1, iteration_log_addr: str = None, modality: Optional[str] = None, + tokenizer: Optional[PreTrainedTokenizer] = None, ) -> StatsKeeper: outbox = asyncio.Queue() statistics = StatsKeeper() @@ -266,7 +350,8 @@ async def async_benchmark( outbox, streaming, concurrency=concurrency, - modality=modality) + modality=modality, + tokenizer=tokenizer) enqueue_task: Optional[asyncio.Task] = None try: backend.run(iteration_addr=iteration_log_addr) diff --git a/tensorrt_llm/bench/dataclasses/general.py b/tensorrt_llm/bench/dataclasses/general.py index 8ea7bd2a32d5..33b23836d176 100644 --- a/tensorrt_llm/bench/dataclasses/general.py +++ b/tensorrt_llm/bench/dataclasses/general.py @@ -24,14 +24,22 @@ class InferenceRequest(BaseModel): input_ids: Optional[List[int]] = Field( alias=AliasChoices("input_ids", "logits")) lora_request: Optional[LoRARequest] = None + turns: Optional[List[str]] = None + category: Optional[str] = None + question_id: Optional[int] = None @model_validator(mode="after") def verify_prompt_and_logits(self) -> InferenceRequest: - if self.prompt is None and self.input_ids is None: + if self.prompt is None and self.input_ids is None and self.turns is None: raise ValueError( - f"Both prompt and input_ids for {self.task_id} are both None.") + f"prompt, input_ids, and turns for task {self.task_id} are all None." + ) return self + @property + def is_multi_turn(self) -> bool: + return self.turns is not None and len(self.turns) > 1 + class DatasetMetadata(BaseModel): isl_stats: PercentileStats diff --git a/tensorrt_llm/bench/dataset/prepare_real_data.py b/tensorrt_llm/bench/dataset/prepare_real_data.py index 063650c926f5..641d95a948f5 100644 --- a/tensorrt_llm/bench/dataset/prepare_real_data.py +++ b/tensorrt_llm/bench/dataset/prepare_real_data.py @@ -11,6 +11,7 @@ from pydantic import BaseModel, model_validator from tensorrt_llm.bench.dataset.utils import ( + generate_multi_turn_dataset, generate_multimodal_dataset, generate_text_dataset, get_norm_dist_lengths, @@ -56,8 +57,6 @@ class DatasetConfig(BaseModel): def check_prompt(self) -> "DatasetConfig": if self.prompt_key and self.prompt: raise AssertionError("--prompt-key and --prompt cannot be set at the same time.") - if (not self.prompt_key) and (not self.prompt): - raise AssertionError("Either --prompt-key or --prompt must be set.") return self @property @@ -77,8 +76,10 @@ def get_prompt(self, req): f"{req.keys()}" ) return req[self.prompt_key] - else: + elif self.prompt: return self.prompt + else: + return "" def get_input(self, req): """Get the input sentence from the given request.""" @@ -207,7 +208,19 @@ def load_dataset_from_hf(dataset_config: DatasetConfig): ) @click.pass_obj def real_dataset(root_args, **kwargs): - """Prepare dataset from real dataset.""" + """Prepare dataset from real dataset. + + Supports three input modes based on the shape of the data at + --dataset-input-key: + + 1. **Single-turn text** (default): The value is a string concatenated with + the prompt from --dataset-prompt-key / --dataset-prompt. + 2. **Multi-turn conversation**: The value is a list of strings (e.g. + MT-Bench ``turns`` field). Each list element is one conversation turn. + No --dataset-prompt-key / --dataset-prompt is required. + 3. **Multimodal**: Detected when the row contains ``image`` / ``video`` + keys. + """ dataset_config = DatasetConfig( **{k[8:]: v for k, v in kwargs.items() if k.startswith("dataset_")} ) @@ -218,8 +231,11 @@ def real_dataset(root_args, **kwargs): task_ids = [] req_cnt = 0 modality = None + multi_turn = False multimodal_texts = [] multimodal_image_paths = [] + all_turns = [] + all_metadata = [] for req in load_dataset_from_hf(dataset_config): if any(key in req for key in ["image", "image_1", "video"]): # multimodal input @@ -248,18 +264,40 @@ def real_dataset(root_args, **kwargs): multimodal_texts.append(text) multimodal_image_paths.append(image_paths) else: - # text input - prompt = dataset_config.get_prompt(req) + " " + dataset_config.get_input(req) - logging.debug(f"Input sequence: {prompt}") - line = root_args.tokenizer.encode(prompt) - if kwargs["max_input_len"] and len(line) > kwargs["max_input_len"]: - continue - input_ids.append(line) - input_lens.append(len(line)) + input_value = dataset_config.get_input(req) + + if isinstance(input_value, list): + # Multi-turn: input_value is a list of conversation turns + multi_turn = True + assert kwargs["output_len_dist"] is not None, ( + "Output length distribution must be set for multi-turn " + "requests (no golden output available)." + ) + turns = input_value + prompt_text = turns[0] + line = root_args.tokenizer.encode(prompt_text) + if kwargs["max_input_len"] and len(line) > kwargs["max_input_len"]: + continue + input_ids.append(line) + input_lens.append(len(line)) + all_turns.append(turns) + all_metadata.append( + {k: req.get(k) for k in ("category", "question_id") if k in req} + ) + else: + # Single-turn text input + prompt = dataset_config.get_prompt(req) + " " + input_value + logging.debug(f"Input sequence: {prompt}") + line = root_args.tokenizer.encode(prompt) + if kwargs["max_input_len"] and len(line) > kwargs["max_input_len"]: + continue + input_ids.append(line) + input_lens.append(len(line)) - # output if fetch from golden - if kwargs["output_len_dist"] is None: - output_lens.append(len(root_args.tokenizer.encode(dataset_config.get_output(req)))) + if kwargs["output_len_dist"] is None: + output_lens.append( + len(root_args.tokenizer.encode(dataset_config.get_output(req))) + ) # lora task id task_id = root_args.task_id @@ -272,13 +310,11 @@ def real_dataset(root_args, **kwargs): if kwargs["num_requests"] and req_cnt >= kwargs["num_requests"]: break - if ( - kwargs["num_requests"] - and (len(input_ids) if modality is None else len(multimodal_texts)) < kwargs["num_requests"] - ): + num_collected = len(input_ids) if modality is None else len(multimodal_texts) + if kwargs["num_requests"] and num_collected < kwargs["num_requests"]: logging.warning( - f"Number of requests={len(input_ids) if modality is None else len(multimodal_texts)} is" - f" smaller than the num-requests user set={kwargs['num_requests']}." + f"Number of requests={num_collected} is smaller than the " + f"num-requests user set={kwargs['num_requests']}." ) # output if randomized @@ -287,7 +323,7 @@ def real_dataset(root_args, **kwargs): output_lens = get_norm_dist_lengths( osl_mean, osl_stdev, - len(input_ids) if modality is None else len(multimodal_texts), + num_collected, root_args.random_seed, ) logging.debug(f"Input lengths: {[len(i) for i in input_ids]}") @@ -295,11 +331,12 @@ def real_dataset(root_args, **kwargs): if modality is not None: logging.debug(f"Modality: {modality}") - dataset_generator = None if modality is not None: dataset_generator = partial( generate_multimodal_dataset, multimodal_texts, multimodal_image_paths ) + elif multi_turn: + dataset_generator = partial(generate_multi_turn_dataset, all_turns, input_ids, all_metadata) else: dataset_generator = partial(generate_text_dataset, input_ids) write_dataset_to_file(dataset_generator(output_lens), root_args.output) diff --git a/tensorrt_llm/bench/dataset/utils.py b/tensorrt_llm/bench/dataset/utils.py index 15c917019530..2ecea0320a5b 100644 --- a/tensorrt_llm/bench/dataset/utils.py +++ b/tensorrt_llm/bench/dataset/utils.py @@ -24,6 +24,20 @@ def generate_text_dataset(input_ids, output_lens, task_ids=None, lora_config=Non yield json.dumps(d, separators=(",", ":"), ensure_ascii=False) +def generate_multi_turn_dataset(all_turns, input_ids, all_metadata, output_lens): + for i, turns in enumerate(all_turns): + d = { + "task_id": i, + "turns": turns, + "output_tokens": output_lens[i], + } + if input_ids[i] is not None: + d["input_ids"] = input_ids[i] + if i < len(all_metadata) and all_metadata[i]: + d.update(all_metadata[i]) + yield json.dumps(d, separators=(",", ":"), ensure_ascii=False) + + def generate_multimodal_dataset(multimodal_texts, multimodal_image_paths, output_lens): for i, (text, image_paths) in enumerate(zip(multimodal_texts, multimodal_image_paths)): d = { diff --git a/tensorrt_llm/bench/utils/data.py b/tensorrt_llm/bench/utils/data.py index c8d735a2bfbb..423e8d3a2201 100644 --- a/tensorrt_llm/bench/utils/data.py +++ b/tensorrt_llm/bench/utils/data.py @@ -51,7 +51,7 @@ def create_dataset_from_stream( tokenizer (PreTrainedTokenizer): HuggingFace tokenizer. stream (TextIO): Stream of input requests. max_input_length (int, optional): Maximum input length to cap prompts to. Defaults to 0. - max_output_length (int, optional): Maximum output length to cap prompts to.. Defaults to 0. + max_output_length (int, optional): Maximum output length to cap prompts to. Defaults to 0. num_requests (int, optional): Number of requests to limit to. Defaults to 0. Returns: @@ -86,30 +86,55 @@ def create_dataset_from_stream( all_logits = [] task_ids = [] lora_requests = [] + all_turns = [] + all_categories = [] + all_question_ids = [] while (line := stream.readline()) and len(task_ids) < max_requests: - # We expect the data to come in as a JSON string. - # For example: + # We support two JSONL formats: + # + # 1. Standard single-turn format: # {"task_id": 1, "prompt": "Generate an infinite response to the following: # There once was a man who.", "output_tokens": 1000} # + # 2. Multi-turn format (e.g. MT-Bench question.jsonl): + # {"question_id": 81, "category": "writing", "turns": ["Write a blog post...", "Rewrite..."]} + # When "turns" is present, the first turn is used as the prompt for + # tokenization/metadata, and all turns are stored for sequential + # multi-turn benchmarking. "output_tokens" is required, same as + # single-turn format. + # # For multimodal data, the data should be of the form: - # {"task_id": 1, "prompt": "Generate an infinite response to the following: - # There once was a man who.", "output_tokens": 1000, + # {"task_id": 1, "prompt": "...", "output_tokens": 1000, # "media_paths": ["/path/to/image1.jpg", "/path/to/image2.jpg"]} # # For LoRA data, the data should be of the form: - # {"task_id": 1, "prompt": "Generate an infinite response to the following: - # There once was a man who.", "output_tokens": 1000, + # {"task_id": 1, "prompt": "...", "output_tokens": 1000, # "lora_request": {"lora_name": "my_lora", "lora_int_id": 1, "lora_path": "/path/to/lora"}} # # Each line should be a complete JSON dictionary with no indentation - # or newline characters. The task_id field is required. + # or newline characters. data = json.loads(line) - prompts.append(data.get("prompt")) - media_paths.append(data.get("media_paths", None)) - all_logits.append(data.get("input_ids", data.get("logits", None))) - all_osl.append(data.get("output_tokens")) - task_ids.append(data.get("task_id")) + + turns = data.get("turns") + if turns is not None and isinstance(turns, list): + prompts.append(data.get("prompt") or turns[0]) + media_paths.append(data.get("media_paths", None)) + all_logits.append(data.get("input_ids", data.get("logits", None))) + all_turns.append(turns) + all_categories.append(data.get("category")) + all_question_ids.append(data.get("question_id")) + all_osl.append(data.get("output_tokens")) + task_ids.append( + data.get("task_id", data.get("question_id", len(task_ids)))) + else: + prompts.append(data.get("prompt")) + media_paths.append(data.get("media_paths", None)) + all_logits.append(data.get("input_ids", data.get("logits", None))) + all_turns.append(None) + all_categories.append(None) + all_question_ids.append(None) + all_osl.append(data.get("output_tokens")) + task_ids.append(data.get("task_id")) # Parse LoRA request if present lora_data = data.get("lora_request", None) @@ -138,8 +163,9 @@ def create_dataset_from_stream( all_isl = [] all_seq_len = [] - for prompt, logits, osl, task_id, lora_request in zip( - prompts, all_logits, all_osl, task_ids, lora_requests): + for prompt, logits, osl, task_id, lora_request, turns, category, question_id in zip( + prompts, all_logits, all_osl, task_ids, lora_requests, all_turns, + all_categories, all_question_ids): if modality is not None: # NOTE: we cannot tokenize multi-modal data, handled by preprocessor # so the actual sequence length is unknown until the model is run @@ -151,7 +177,8 @@ def create_dataset_from_stream( logits = tokenize(prompt)["input_ids"] if logits is None else logits cur_isl = len(logits) all_isl.append(cur_isl) - all_seq_len.append(cur_isl + osl) + num_turns = len(turns) if turns is not None else 1 + all_seq_len.append(cur_isl + num_turns * osl) request = InferenceRequest( task_id=task_id, @@ -159,6 +186,9 @@ def create_dataset_from_stream( output_tokens=output_limiter(osl), input_ids=logits, lora_request=lora_request, + turns=turns, + category=category, + question_id=question_id, ) dataset.append(request)