From 86dcd17caa8ba3c6c75e8c1447e767fbf81e97db Mon Sep 17 00:00:00 2001 From: vmehrotra50 Date: Sun, 26 Jan 2025 22:17:48 -0500 Subject: [PATCH 01/11] Add core correctness testing setup --- sarathi/benchmark/benchmark_runner.py | 38 ++++++++++++++++----- sarathi/benchmark/config.py | 9 +++++ sarathi/benchmark/utils/dataset_loader.py | 40 +++++++++++++++++++++++ 3 files changed, 79 insertions(+), 8 deletions(-) create mode 100644 sarathi/benchmark/utils/dataset_loader.py diff --git a/sarathi/benchmark/benchmark_runner.py b/sarathi/benchmark/benchmark_runner.py index 083dcb3..7bcb019 100644 --- a/sarathi/benchmark/benchmark_runner.py +++ b/sarathi/benchmark/benchmark_runner.py @@ -11,6 +11,7 @@ from sarathi.benchmark.entities import Request from sarathi.benchmark.request_generator import RequestGeneratorRegistry from sarathi.benchmark.utils.random import set_seeds +from sarathi.benchmark.utils import dataset_loader from sarathi.config import ReplicaConfig from sarathi.metrics.metrics_store import MetricsStore from sarathi.types import ReplicaResourceMapping, ResourceMapping @@ -38,11 +39,19 @@ def __init__( os.makedirs(replica_config.output_dir, exist_ok=True) set_seeds(self.config.seed) + request_generator = RequestGeneratorRegistry.get( self.config.request_generator_config.get_type(), self.config.request_generator_config, ) - self.requests = request_generator.generate() + + if not config.run_correctness_tests: + self.requests = request_generator.generate() + else: + # NOTE: needs to be trace-based with correctness tests + self.requests = dataset_loader.get_data_loader(dataset) + self.requests[self.replica_id :: self.config.num_replicas] + self.correctness_output = {} # select every nth request for this replica # e.g. if there are 4 replicas, and this is the 2nd replica, then @@ -70,14 +79,24 @@ def _get_input_params( temperature=0, top_p=1.0, ) - prompt_token_ids = [1] * request.num_prefill_tokens - return { - "prompt": None, - "prompt_token_ids": prompt_token_ids, - "sampling_params": sampling_params, - "arrival_time": first_request_time + request.arrived_at, - } + if self.config.run_correctness_tests: + return { + "prompt": request.prompt, + "prompt_token_ids": None, + "sampling_params": sampling_params, + "arrival_time": first_request_time + request.arrived_at, + } + else: + + prompt_token_ids = [1] * request.num_prefill_tokens + + return { + "prompt": None, + "prompt_token_ids": prompt_token_ids, + "sampling_params": sampling_params, + "arrival_time": first_request_time + request.arrived_at, + } def warmup(self) -> None: self.llm_engine.add_request(**self._get_input_params(self.requests[0], 0)) @@ -111,6 +130,9 @@ def _run(self) -> None: num_steps += 1 for output in step_outputs: + if self.config.run_correctness_tests: + self.correctness_outputs[output.seq_id] = self.correctness_outputs.get(output.seq_id, "") + output.text + if output.finished: num_processed_requests += 1 pbar.update(1) diff --git a/sarathi/benchmark/config.py b/sarathi/benchmark/config.py index bbc5d10..080bd6e 100644 --- a/sarathi/benchmark/config.py +++ b/sarathi/benchmark/config.py @@ -240,6 +240,15 @@ class BenchmarkConfig(BaseEndpointConfig): request_generator_config: BaseRequestGeneratorConfig = field( default_factory=SyntheticRequestGeneratorConfig ) + run_correctness_tests: bool = field( + default=False, metadata={"help": "Collect correctness data in this run"} + ) + correctness_test_dataset: Optional[str] = field( + default=None, metadata={"help": "Dataset for correctness tests"} + ) + correctness_test_ground_truth: Optional[str] = field( + default=None, metadata={"help": "Ground truth file to compare LLM output with."} + ) def __post_init__(self): super().__post_init__() diff --git a/sarathi/benchmark/utils/dataset_loader.py b/sarathi/benchmark/utils/dataset_loader.py new file mode 100644 index 0000000..a463779 --- /dev/null +++ b/sarathi/benchmark/utils/dataset_loader.py @@ -0,0 +1,40 @@ +from collections.abc import Iterable +from itertools import islice +from typing import Optional + +from transformers import AutoTokenizer + +from datasets import load_dataset + +DATASET_FIELDS = { + "xsum": "document", + "openai_humaneval": "prompt", + "ccdv/arxiv-summarization": "article", + "lmsys/lmsys-chat-1m": "conversation", + "Fredithefish/ShareGPT-unfiltered-alpaca-lora-format": "instruction" + # 'another_dataset': 'text_field', + # ... add other datasets and their respective fields +} + +def get_data_loader( + dataset_str: Optional[str], max_samples: Optional[int] +) -> Iterable[str]: + + dataset = load_dataset(dataset_str, split='train') + field_name = DATASET_FIELDS[dataset_str] + + if "conversation" not in field_name: + sample_iterator = map( + lambda x: x, map(lambda x: x[field_name], dataset) + ) + else: + # assume llama3 + # don't use meta prompt, this is chat case + tokenizer = AutoTokenizer.from_pretrained("meta-llama/Meta-Llama-3-8B-Instruct") + conversations = dataset[field_name] + sample_iterator = tokenizer.apply_chat_template(conversations, tokenize=False, add_generation_prompt=True,) + + if not max_samples: + return sample_iterator + + return islice(sample_iterator, max_samples) From 94bc5976af387e98deb3be1cc02a23e02c8cbf09 Mon Sep 17 00:00:00 2001 From: vmehrotra50 Date: Sun, 26 Jan 2025 23:14:30 -0500 Subject: [PATCH 02/11] Add ground truth comparison --- sarathi/benchmark/benchmark_runner.py | 13 ++++++++++++- 1 file changed, 12 insertions(+), 1 deletion(-) diff --git a/sarathi/benchmark/benchmark_runner.py b/sarathi/benchmark/benchmark_runner.py index 7bcb019..d4dc1c0 100644 --- a/sarathi/benchmark/benchmark_runner.py +++ b/sarathi/benchmark/benchmark_runner.py @@ -131,7 +131,7 @@ def _run(self) -> None: for output in step_outputs: if self.config.run_correctness_tests: - self.correctness_outputs[output.seq_id] = self.correctness_outputs.get(output.seq_id, "") + output.text + self.correctness_output[output.seq_id] = self.correctness_output.get(output.seq_id, "") + output.text if output.finished: num_processed_requests += 1 @@ -165,6 +165,17 @@ def run(self) -> None: metric_store = self.llm_engine.get_metric_store() return metric_store + def check_correctness(self) -> bool: + assert self.config.run_correctness_checks + + with open(file_path, 'r') as file: + data = yaml.safe_load(file) + + for k, v in data.items(): + if self.correctness_output[k] != v: return False + + return True + class BenchmarkRunnerLauncher: From aa7e2ce43f118af43541186f55028c7a03baba0f Mon Sep 17 00:00:00 2001 From: vmehrotra50 Date: Tue, 28 Jan 2025 22:48:49 -0500 Subject: [PATCH 03/11] Add debug statements --- sarathi/benchmark/benchmark_runner.py | 5 ++++- sarathi/benchmark/main.py | 5 ++++- unit_test/correctness_test.py | 0 3 files changed, 8 insertions(+), 2 deletions(-) create mode 100644 unit_test/correctness_test.py diff --git a/sarathi/benchmark/benchmark_runner.py b/sarathi/benchmark/benchmark_runner.py index d4dc1c0..854670f 100644 --- a/sarathi/benchmark/benchmark_runner.py +++ b/sarathi/benchmark/benchmark_runner.py @@ -49,6 +49,7 @@ def __init__( self.requests = request_generator.generate() else: # NOTE: needs to be trace-based with correctness tests + print('CORRECTNESS TESTS ENABLED') self.requests = dataset_loader.get_data_loader(dataset) self.requests[self.replica_id :: self.config.num_replicas] self.correctness_output = {} @@ -131,6 +132,8 @@ def _run(self) -> None: for output in step_outputs: if self.config.run_correctness_tests: + print("CORRECTNESS OUTPUT") + print(output.text) self.correctness_output[output.seq_id] = self.correctness_output.get(output.seq_id, "") + output.text if output.finished: @@ -165,7 +168,7 @@ def run(self) -> None: metric_store = self.llm_engine.get_metric_store() return metric_store - def check_correctness(self) -> bool: + def check_correctness(self) -> bool: assert self.config.run_correctness_checks with open(file_path, 'r') as file: diff --git a/sarathi/benchmark/main.py b/sarathi/benchmark/main.py index 333033d..9335514 100644 --- a/sarathi/benchmark/main.py +++ b/sarathi/benchmark/main.py @@ -13,7 +13,10 @@ def main() -> None: - config = BenchmarkConfig.create_from_cli_args() + config = BenchmarkConfig( + run_correctness_tests=True, + correctness_test_dataset="openai_humaneval" + ) os.makedirs(config.output_dir, exist_ok=True) with open(os.path.join(config.output_dir, "config.yaml"), "w") as f: diff --git a/unit_test/correctness_test.py b/unit_test/correctness_test.py new file mode 100644 index 0000000..e69de29 From 5487c0861a8f749700ee63ff2747f3c71e2dc7b1 Mon Sep 17 00:00:00 2001 From: vmehrotra50 Date: Tue, 28 Jan 2025 22:52:36 -0500 Subject: [PATCH 04/11] debug --- sarathi/benchmark/main.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/sarathi/benchmark/main.py b/sarathi/benchmark/main.py index 9335514..71b2cb1 100644 --- a/sarathi/benchmark/main.py +++ b/sarathi/benchmark/main.py @@ -18,9 +18,9 @@ def main() -> None: correctness_test_dataset="openai_humaneval" ) - os.makedirs(config.output_dir, exist_ok=True) - with open(os.path.join(config.output_dir, "config.yaml"), "w") as f: - yaml.dump(config.to_dict(), f) + # os.makedirs(config.output_dir, exist_ok=True) + # with open(os.path.join(config.output_dir, "config.yaml"), "w") as f: + # yaml.dump(config.to_dict(), f) logger.info(f"Starting benchmark with config: {config}") From 1f54ea1e5fdf15e58436ea845d1efcb2e88a0d40 Mon Sep 17 00:00:00 2001 From: vmehrotra50 Date: Tue, 28 Jan 2025 22:54:18 -0500 Subject: [PATCH 05/11] Minor bug fix --- sarathi/benchmark/benchmark_runner.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/sarathi/benchmark/benchmark_runner.py b/sarathi/benchmark/benchmark_runner.py index 854670f..6dc566d 100644 --- a/sarathi/benchmark/benchmark_runner.py +++ b/sarathi/benchmark/benchmark_runner.py @@ -50,7 +50,7 @@ def __init__( else: # NOTE: needs to be trace-based with correctness tests print('CORRECTNESS TESTS ENABLED') - self.requests = dataset_loader.get_data_loader(dataset) + self.requests = dataset_loader.get_data_loader(self.config.correctness_test_dataset) self.requests[self.replica_id :: self.config.num_replicas] self.correctness_output = {} From 33c5ea6dba62ac865e876779b816b57fd2700be7 Mon Sep 17 00:00:00 2001 From: vmehrotra50 Date: Tue, 28 Jan 2025 22:55:22 -0500 Subject: [PATCH 06/11] Minor bug fix --- sarathi/benchmark/utils/dataset_loader.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/sarathi/benchmark/utils/dataset_loader.py b/sarathi/benchmark/utils/dataset_loader.py index a463779..00b31ae 100644 --- a/sarathi/benchmark/utils/dataset_loader.py +++ b/sarathi/benchmark/utils/dataset_loader.py @@ -17,7 +17,7 @@ } def get_data_loader( - dataset_str: Optional[str], max_samples: Optional[int] + dataset_str: Optional[str], max_samples: Optional[int]=None ) -> Iterable[str]: dataset = load_dataset(dataset_str, split='train') From 3eadefaca4633478fe3facb5fd99239823e4b576 Mon Sep 17 00:00:00 2001 From: vmehrotra50 Date: Tue, 28 Jan 2025 22:56:23 -0500 Subject: [PATCH 07/11] minor bug fix --- sarathi/benchmark/utils/dataset_loader.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/sarathi/benchmark/utils/dataset_loader.py b/sarathi/benchmark/utils/dataset_loader.py index 00b31ae..e99cd85 100644 --- a/sarathi/benchmark/utils/dataset_loader.py +++ b/sarathi/benchmark/utils/dataset_loader.py @@ -20,7 +20,7 @@ def get_data_loader( dataset_str: Optional[str], max_samples: Optional[int]=None ) -> Iterable[str]: - dataset = load_dataset(dataset_str, split='train') + dataset = load_dataset(dataset_str, split='test') field_name = DATASET_FIELDS[dataset_str] if "conversation" not in field_name: From 2e166af5caba617303aa7ae66cfb7f4cd69ae81d Mon Sep 17 00:00:00 2001 From: vmehrotra50 Date: Tue, 28 Jan 2025 22:57:34 -0500 Subject: [PATCH 08/11] Remove sharding requests logic (temporarily) --- sarathi/benchmark/benchmark_runner.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/sarathi/benchmark/benchmark_runner.py b/sarathi/benchmark/benchmark_runner.py index 6dc566d..0bad4c6 100644 --- a/sarathi/benchmark/benchmark_runner.py +++ b/sarathi/benchmark/benchmark_runner.py @@ -51,7 +51,7 @@ def __init__( # NOTE: needs to be trace-based with correctness tests print('CORRECTNESS TESTS ENABLED') self.requests = dataset_loader.get_data_loader(self.config.correctness_test_dataset) - self.requests[self.replica_id :: self.config.num_replicas] + # self.requests[self.replica_id :: self.config.num_replicas] self.correctness_output = {} # select every nth request for this replica From c3f5a50d6b6f0e6dee428cca07dabe665abfadcb Mon Sep 17 00:00:00 2001 From: vmehrotra50 Date: Tue, 28 Jan 2025 22:58:12 -0500 Subject: [PATCH 09/11] Remove sharding requests logic (temporarily) --- sarathi/benchmark/benchmark_runner.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/sarathi/benchmark/benchmark_runner.py b/sarathi/benchmark/benchmark_runner.py index 0bad4c6..51f0b5a 100644 --- a/sarathi/benchmark/benchmark_runner.py +++ b/sarathi/benchmark/benchmark_runner.py @@ -58,7 +58,7 @@ def __init__( # e.g. if there are 4 replicas, and this is the 2nd replica, then # we will select the 2nd, 6th, 10th, ... requests # round robin scheduling - self.requests = self.requests[self.replica_id :: self.config.num_replicas] + # self.requests = self.requests[self.replica_id :: self.config.num_replicas] if self.config.num_replicas > 1: # disable per-replica wandb logging for multi-replica runs From 1d072f35f98b2b8a4fd6248ae676361e8c96af0f Mon Sep 17 00:00:00 2001 From: apoorvam-mk Date: Tue, 28 Jan 2025 23:15:53 -0500 Subject: [PATCH 10/11] fix prompt processing issues --- sarathi/benchmark/benchmark_runner.py | 7 ++++++- sarathi/benchmark/entities/request.py | 8 +++++++- 2 files changed, 13 insertions(+), 2 deletions(-) diff --git a/sarathi/benchmark/benchmark_runner.py b/sarathi/benchmark/benchmark_runner.py index 51f0b5a..33e6ab5 100644 --- a/sarathi/benchmark/benchmark_runner.py +++ b/sarathi/benchmark/benchmark_runner.py @@ -50,7 +50,12 @@ def __init__( else: # NOTE: needs to be trace-based with correctness tests print('CORRECTNESS TESTS ENABLED') - self.requests = dataset_loader.get_data_loader(self.config.correctness_test_dataset) + self.requests = [Request( + arrived_at=0, + num_prefill_tokens=len(prompt), + num_decode_tokens=1024, + prompt=prompt + ) for prompt in dataset_loader.get_data_loader(self.config.correctness_test_dataset)] # self.requests[self.replica_id :: self.config.num_replicas] self.correctness_output = {} diff --git a/sarathi/benchmark/entities/request.py b/sarathi/benchmark/entities/request.py index 2d74498..d53b7c0 100644 --- a/sarathi/benchmark/entities/request.py +++ b/sarathi/benchmark/entities/request.py @@ -1,5 +1,5 @@ import logging -from typing import Tuple +from typing import Tuple, Optional from sarathi.benchmark.entities.base_entity import BaseEntity @@ -13,11 +13,13 @@ def __init__( arrived_at: float, num_prefill_tokens: int, num_decode_tokens: int, + prompt: Optional[str]=None, ): self._id = Request.generate_id() self._arrived_at = arrived_at self._num_prefill_tokens = num_prefill_tokens self._num_decode_tokens = num_decode_tokens + self._prompt = prompt assert num_prefill_tokens > 0 assert num_decode_tokens > 0 @@ -28,6 +30,10 @@ def size(self) -> Tuple[int, int]: @property def arrived_at(self) -> float: return self._arrived_at + + @property + def prompt(self) -> float: + return self._prompt @property def num_prefill_tokens(self) -> int: From c387d4730c549451e534b59837bb5aa5c466ae3f Mon Sep 17 00:00:00 2001 From: vmehrotra50 Date: Mon, 3 Feb 2025 23:32:26 -0500 Subject: [PATCH 11/11] Refactor correctness tests to work as a unit test --- sarathi/benchmark/benchmark_runner.py | 14 ++-- sarathi/benchmark/config.py | 52 ++++++++++++--- sarathi/benchmark/entities/request.py | 10 ++- .../base_request_length_generator.py | 4 +- .../dataset_request_length_generator.py | 35 ++++++++++ .../request_length_generator_registry.py | 6 ++ .../synthetic_request_generator.py | 7 ++ sarathi/benchmark/utils/dataset_loader.py | 65 +++++++++++++++---- sarathi/types.py | 2 +- unit_test/correctness_test.py | 52 +++++++++++++++ 10 files changed, 215 insertions(+), 32 deletions(-) create mode 100644 sarathi/benchmark/request_generator/dataset_request_length_generator.py diff --git a/sarathi/benchmark/benchmark_runner.py b/sarathi/benchmark/benchmark_runner.py index 51f0b5a..1671547 100644 --- a/sarathi/benchmark/benchmark_runner.py +++ b/sarathi/benchmark/benchmark_runner.py @@ -45,14 +45,10 @@ def __init__( self.config.request_generator_config, ) - if not config.run_correctness_tests: - self.requests = request_generator.generate() - else: - # NOTE: needs to be trace-based with correctness tests - print('CORRECTNESS TESTS ENABLED') - self.requests = dataset_loader.get_data_loader(self.config.correctness_test_dataset) - # self.requests[self.replica_id :: self.config.num_replicas] - self.correctness_output = {} + self.requests = request_generator.generate() + + self.run_correctness_tests = self.config.correctness_test_config is not None \ + and self.config.correctness_test_config.run_correctness_tests # select every nth request for this replica # e.g. if there are 4 replicas, and this is the 2nd replica, then @@ -81,7 +77,7 @@ def _get_input_params( top_p=1.0, ) - if self.config.run_correctness_tests: + if self.run_correctness_tests: return { "prompt": request.prompt, "prompt_token_ids": None, diff --git a/sarathi/benchmark/config.py b/sarathi/benchmark/config.py index 080bd6e..7d4c159 100644 --- a/sarathi/benchmark/config.py +++ b/sarathi/benchmark/config.py @@ -161,6 +161,35 @@ def get_type(): return RequestLengthGeneratorType.FIXED + +@dataclass +class DatasetRequestLengthGeneratorConfig(BaseRequestLengthGeneratorConfig): + dataset: str = field( + default="ccdv/arxiv-summarization", + metadata={"help": "Path to the trace file for request lengths."}, + ) + meta_prompt: str = field( + default=None, + metadata={"help": "Meta prompt for the dataset."}, + ) + max_prompt_len: int = field( + default=4096, metadata={"help": "Maximum prompt length allowed."} + ) + max_num_prompts: int = field( + default=300, metadata={"help": "Maximum number of prompts to use."} + ) + max_decode_tokens: int = field( + default=512, metadata={"help": "Maximum number of decode tokens."} + ) + tokenizer_model: str = field( + default="meta-llama/Meta-Llama-3-8B-Instruct", metadata={"help": "Name or path of the huggingface model to use for the tokenizer."} + ) + + @staticmethod + def get_type(): + return RequestLengthGeneratorType.DATASET + + @dataclass class BaseRequestGeneratorConfig(BasePolyConfig): seed: int = field( @@ -214,6 +243,19 @@ class TraceRequestGeneratorConfig(BaseRequestGeneratorConfig): def get_type(): return RequestGeneratorType.TRACE +@dataclass +class CorrectnessTestConfig(BaseTestConfig): + run_correctness_tests: bool = field( + default=False, metadata={"help": "Collect correctness data in this run"} + ) + run_correctness_baseline: bool = field( + default=False, metadata={"help": "Make this correctness ground truth for correctness tests"} + ) + correctness_test_file: bool = field( + default=None, metadata={"help": "Ground truth file for model output. If run_correctness_baseline is True, then the model output will be saved to \ + this file to be used as a ground truth file. Otherwise, the test will read from this file to be used as ground truth for \ + the correctness test."} + ) @dataclass class BenchmarkConfig(BaseEndpointConfig): @@ -240,14 +282,8 @@ class BenchmarkConfig(BaseEndpointConfig): request_generator_config: BaseRequestGeneratorConfig = field( default_factory=SyntheticRequestGeneratorConfig ) - run_correctness_tests: bool = field( - default=False, metadata={"help": "Collect correctness data in this run"} - ) - correctness_test_dataset: Optional[str] = field( - default=None, metadata={"help": "Dataset for correctness tests"} - ) - correctness_test_ground_truth: Optional[str] = field( - default=None, metadata={"help": "Ground truth file to compare LLM output with."} + correctness_test_config: Optional[CorrectnessTestConfig] = field( + default_factory=CorrectnessTestConfig ) def __post_init__(self): diff --git a/sarathi/benchmark/entities/request.py b/sarathi/benchmark/entities/request.py index 2d74498..097c0fb 100644 --- a/sarathi/benchmark/entities/request.py +++ b/sarathi/benchmark/entities/request.py @@ -1,5 +1,5 @@ import logging -from typing import Tuple +from typing import Tuple, Optional from sarathi.benchmark.entities.base_entity import BaseEntity @@ -13,11 +13,13 @@ def __init__( arrived_at: float, num_prefill_tokens: int, num_decode_tokens: int, + prompt: Optional[str] ): self._id = Request.generate_id() self._arrived_at = arrived_at self._num_prefill_tokens = num_prefill_tokens self._num_decode_tokens = num_decode_tokens + self._prompt = prompt assert num_prefill_tokens > 0 assert num_decode_tokens > 0 @@ -44,6 +46,12 @@ def pd_ratio(self) -> float: @property def total_tokens(self) -> int: return self._num_prefill_tokens + self._num_decode_tokens + + @property + def prompt(self) -> int: + return self._prompt + + def to_dict(self) -> dict: return { diff --git a/sarathi/benchmark/request_generator/base_request_length_generator.py b/sarathi/benchmark/request_generator/base_request_length_generator.py index 1609de3..dfe6b7d 100644 --- a/sarathi/benchmark/request_generator/base_request_length_generator.py +++ b/sarathi/benchmark/request_generator/base_request_length_generator.py @@ -1,5 +1,5 @@ from abc import ABC, abstractmethod -from typing import Tuple +from typing import Tuple, Union from sarathi.benchmark.config import BaseRequestLengthGeneratorConfig @@ -10,5 +10,5 @@ def __init__(self, config: BaseRequestLengthGeneratorConfig): self.config = config @abstractmethod - def get_next_num_tokens(self) -> Tuple[float, float]: + def get_next_num_tokens(self) -> Tuple[Union[str|float], float]: pass diff --git a/sarathi/benchmark/request_generator/dataset_request_length_generator.py b/sarathi/benchmark/request_generator/dataset_request_length_generator.py new file mode 100644 index 0000000..5e398f2 --- /dev/null +++ b/sarathi/benchmark/request_generator/dataset_request_length_generator.py @@ -0,0 +1,35 @@ +import logging +from typing import Tuple, Union + +import numpy as np +import pandas as pd + +from sarathi.benchmark.utils import data_loader +from sarathi.benchmark.config import DatasetRequestLengthGeneratorConfig +from sarathi.benchmark.request_generator.base_request_length_generator import ( + BaseRequestLengthGenerator, +) + +logger = logging.getLogger(__name__) + + +class DatasetRequestLengthGenerator(BaseRequestLengthGenerator): + + def __init__(self, config: DatasetRequestLengthGeneratorConfig): + super().__init__(config) + self.next_request_idx = 0 + prompts = data_loader.get_data_loader(None, config.dataset, config.meta_prompt, None, config.tokenizer_model) + self.requests = [prompt for prompt in prompts if len(prompt) <= config.max_prompt_len][:config.max_num_prompts] + self.decode_tokens = config.max_decode_tokens + + def get_next_num_tokens(self) -> Tuple[Union[str|float], float]: + if self.next_request_idx >= len(self.requests): + return None, None + + row = self.requests[self.next_request_idx] + self.next_request_idx += 1 + + return ( + row, + self.decode_tokens, + ) \ No newline at end of file diff --git a/sarathi/benchmark/request_generator/request_length_generator_registry.py b/sarathi/benchmark/request_generator/request_length_generator_registry.py index da3f12e..ef76896 100644 --- a/sarathi/benchmark/request_generator/request_length_generator_registry.py +++ b/sarathi/benchmark/request_generator/request_length_generator_registry.py @@ -7,6 +7,9 @@ from sarathi.benchmark.request_generator.uniform_request_length_generator import ( UniformRequestLengthGenerator, ) +from sarathi.benchmark.request_generator.dataset_request_length_generator import ( + DatasetRequestLengthGenerator +) from sarathi.benchmark.request_generator.zipf_request_length_generator import ( ZipfRequestLengthGenerator, ) @@ -30,3 +33,6 @@ class RequestLengthGeneratorRegistry(BaseRegistry): RequestLengthGeneratorRegistry.register( RequestLengthGeneratorType.FIXED, FixedRequestLengthGenerator ) +RequestLengthGeneratorRegistry.register( + RequestLengthGeneratorType.DATASET, DatasetRequestLengthGenerator +) diff --git a/sarathi/benchmark/request_generator/synthetic_request_generator.py b/sarathi/benchmark/request_generator/synthetic_request_generator.py index a835791..1b69b1f 100644 --- a/sarathi/benchmark/request_generator/synthetic_request_generator.py +++ b/sarathi/benchmark/request_generator/synthetic_request_generator.py @@ -43,11 +43,18 @@ def _generate_next_request(self, last_arrived_at: float) -> Request: if prefill_tokens is None or decode_tokens is None: return None + + prompt = None + # custom request prompt + if type(prefill_tokens) is str: + prompt = prefill_tokens + prefill_tokens = len(prefill_tokens) return Request( arrived_at=arrived_at, num_prefill_tokens=int(prefill_tokens), num_decode_tokens=int(decode_tokens), + prompt=prompt ) def _generate_requests(self) -> List[Request]: diff --git a/sarathi/benchmark/utils/dataset_loader.py b/sarathi/benchmark/utils/dataset_loader.py index e99cd85..e273fa9 100644 --- a/sarathi/benchmark/utils/dataset_loader.py +++ b/sarathi/benchmark/utils/dataset_loader.py @@ -7,30 +7,46 @@ from datasets import load_dataset DATASET_FIELDS = { - "xsum": "document", - "openai_humaneval": "prompt", - "ccdv/arxiv-summarization": "article", - "lmsys/lmsys-chat-1m": "conversation", - "Fredithefish/ShareGPT-unfiltered-alpaca-lora-format": "instruction" + "xsum": ("document", "train"), + "openai_humaneval": ("prompt", "test"), + "ccdv/arxiv-summarization": ("article", "train"), + "lmsys/lmsys-chat-1m": ("conversation", "train"), + "OpenGVLab/ShareGPT-4o": ("conversations", "image_caption"), + "Fredithefish/ShareGPT-unfiltered-alpaca-lora-format": ("instruction", "train"), + "openai/gsm8k": ("question", "main"), # 'another_dataset': 'text_field', # ... add other datasets and their respective fields } + def get_data_loader( - dataset_str: Optional[str], max_samples: Optional[int]=None + input_string: Optional[str], dataset_str: Optional[str], meta_prompt: Optional[str], max_samples: Optional[int], + model_for_tokenizer_chat_template: Optional[str] ) -> Iterable[str]: + if input_string: + return [input_string] + if dataset_str not in DATASET_FIELDS: + # assume bwb path + return islice(load_bwb(dataset_str, meta_prompt), max_samples) - dataset = load_dataset(dataset_str, split='test') - field_name = DATASET_FIELDS[dataset_str] + field_name = DATASET_FIELDS[dataset_str][0] + split_name = DATASET_FIELDS[dataset_str][1] + if dataset_str != "OpenGVLab/ShareGPT-4o": + dataset = load_dataset(dataset_str, split=split_name) + else: + dataset = load_dataset(dataset_str, 'image_caption') if "conversation" not in field_name: sample_iterator = map( - lambda x: x, map(lambda x: x[field_name], dataset) + lambda x: f"{meta_prompt} {x}", map(lambda x: x[field_name], dataset) ) else: - # assume llama3 + # reformat sharegpt-4 + if "lmsys" not in dataset_str: + dataset['images'] = dataset['images'].map(format_conversations) + dataset = dataset['images'] # don't use meta prompt, this is chat case - tokenizer = AutoTokenizer.from_pretrained("meta-llama/Meta-Llama-3-8B-Instruct") + tokenizer = AutoTokenizer.from_pretrained(model_for_tokenizer_chat_template) conversations = dataset[field_name] sample_iterator = tokenizer.apply_chat_template(conversations, tokenize=False, add_generation_prompt=True,) @@ -38,3 +54,30 @@ def get_data_loader( return sample_iterator return islice(sample_iterator, max_samples) + + +def format_conversations(sample): + sample['conversations'] = [ + { + "content": turn["value"], + "role": "user" if turn["from"] == "human" else "assistant" + } + for turn in sample['conversations'] + ] + return sample + + +def load_bwb(file_path, meta_prompt, chunk_size=1000): + translation_chunks = [] + try: + with open(file_path, 'r', encoding='utf-8') as file: + while True: + chunk = file.read(chunk_size) + if not chunk: + break + chunk = meta_prompt + chunk + translation_chunks.append(chunk) + except FileNotFoundError: + print(f"Error: File '{file_path}' not found.") + + return translation_chunks \ No newline at end of file diff --git a/sarathi/types.py b/sarathi/types.py index fbf2cbb..3fd45c3 100644 --- a/sarathi/types.py +++ b/sarathi/types.py @@ -32,7 +32,7 @@ class RequestLengthGeneratorType(Enum): ZIPF = "ZIPF" TRACE = "TRACE" FIXED = "FIXED" - + DATASET = "DATASET" class AttentionBackend(Enum): FLASHINFER = "FLASHINFER" diff --git a/unit_test/correctness_test.py b/unit_test/correctness_test.py index e69de29..8007359 100644 --- a/unit_test/correctness_test.py +++ b/unit_test/correctness_test.py @@ -0,0 +1,52 @@ +import glob +import shutil +import json +import os + +import pandas as pd +import pytest + +from sarathi.benchmark.benchmark_runner import BenchmarkRunnerLauncher +from sarathi.benchmark.config import BenchmarkConfig, SyntheticRequestGeneratorConfig, \ + DatasetRequestLengthGeneratorConfig, CorrectnessTestConfig +from sarathi.config.config import ModelConfig, ParallelConfig + + +# pytest -k "perf_test" +@pytest.mark.parametrize( + "model, max_model_len, pp_size, tp_size, dataset, request_pattern, baseline_run, test_file", + [ + ("meta-llama/Meta-Llama-3-8B", 4096, 1, 1, "OpenGVLab/ShareGPT-4o", "uniform", False, None), + ] +) +def test_correctness(model: str, max_model_len: int, pp_size: int, tp_size: int, dataset: str, request_pattern: str, baseline_run: bool, baseline_file: str): + # TODO: Test over 3d space + model_config = ModelConfig( + model=model, + max_model_len=max_model_len + ) + parallel_config = ParallelConfig( + pipeline_parallel_size=pp_size, + tensor_parallel_size=tp_size + ) + request_generator_config = SyntheticRequestGeneratorConfig( + length_generator_config=DatasetRequestLengthGeneratorConfig(dataset=dataset) + ) + correctness_test_config = CorrectnessTestConfig( + run_correctness_tests=True, + run_correctness_baseline=baseline_run, + correctness_test_file=baseline_file + ) + cwd = os.getcwd() + output_dir = os.path.join(cwd, "benchmark_output") + benchmark_config = BenchmarkConfig( + log_level="error", + output_dir=output_dir, + model_config=model_config, + parallel_config=parallel_config, + request_generator_config=request_generator_config, + test_config=correctness_test_config, + ) + BenchmarkRunnerLauncher(benchmark_config).run() + + print("correctness_test finished.") \ No newline at end of file