A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 15, 16:00 UTC
[Benchmark] Add AIMO Dataset to Benchmark (#15955)
Signed-off-by: Ziji Shi <shi.ziji.sm@gmail.com> Signed-off-by: StevenShi-23 <shi.ziji.sm@gmail.com>
Ziji Shi (Steven) committed
Apr 2, 2025 at 23:09 UTC
06f21ce7a55222fbda9d1d6c7543c325df741386
2 files changed
+58
-5
benchmarks/benchmark_dataset.py
+49
@@ -752,3 +752,52 @@ class InstructCoderDataset(HuggingFaceDataset):
752
))
753
self.maybe_oversample_requests(sampled_requests, num_requests)
754
return sampled_requests
755
+
756
+
757
+# -----------------------------------------------------------------------------
758
+# AIMO Dataset Implementation
759
+# -----------------------------------------------------------------------------
760
+
761
+
762
+class AIMODataset(HuggingFaceDataset):
763
+ """
764
+ Dataset class for processing a AIMO dataset with reasoning questions.
765
+ """
766
+ SUPPORTED_DATASET_PATHS = {
767
+ "AI-MO/aimo-validation-aime", "AI-MO/NuminaMath-1.5",
768
+ "AI-MO/NuminaMath-CoT"
769
+ }
770
+
771
+ def sample(self,
772
+ tokenizer: PreTrainedTokenizerBase,
773
+ num_requests: int,
774
+ output_len: Optional[int] = None,
775
+ **kwargs) -> list:
776
+ sampled_requests = []
777
+ dynamic_output = output_len is None
778
+
779
+ for item in self.data:
780
+ if len(sampled_requests) >= num_requests:
781
+ break
782
+ prompt, completion = item['problem'], item["solution"]
783
+
784
+ prompt_ids = tokenizer(prompt).input_ids
785
+ completion_ids = tokenizer(completion).input_ids
786
+ prompt_len = len(prompt_ids)
787
+ completion_len = len(completion_ids)
788
+ output_len = completion_len if dynamic_output else output_len
789
+ assert isinstance(output_len, int) and output_len > 0
790
+ if dynamic_output and not is_valid_sequence(prompt_len,
791
+ completion_len,
792
+ max_prompt_len=2048,
793
+ max_total_len=32000):
794
+ continue
795
+ sampled_requests.append(
796
+ SampleRequest(
797
+ prompt=prompt,
798
+ prompt_len=prompt_len,
799
+ expected_output_len=output_len,
800
+ multi_modal_data=None,
801
+ ))
802
+ self.maybe_oversample_requests(sampled_requests, num_requests)
803
+ return sampled_requests
benchmarks/benchmark_serving.py
+9
-5
@@ -49,10 +49,11 @@ try:
49
except ImportError:
50
from argparse import ArgumentParser as FlexibleArgumentParser
51
52
-from benchmark_dataset import (BurstGPTDataset, ConversationDataset,
53
- HuggingFaceDataset, InstructCoderDataset,
54
- RandomDataset, SampleRequest, ShareGPTDataset,
55
- SonnetDataset, VisionArenaDataset)
52
+from benchmark_dataset import (AIMODataset, BurstGPTDataset,
53
+ ConversationDataset, HuggingFaceDataset,
54
+ InstructCoderDataset, RandomDataset,
55
+ SampleRequest, ShareGPTDataset, SonnetDataset,
56
+ VisionArenaDataset)
57
from benchmark_utils import convert_to_pytorch_benchmark_format, write_to_json
58
59
MILLISECONDS_TO_SECONDS_CONVERSION = 1000
@@ -595,6 +596,9 @@ def main(args: argparse.Namespace):
596
args.hf_split = "train"
597
elif args.dataset_path in ConversationDataset.SUPPORTED_DATASET_PATHS:
598
dataset_class = ConversationDataset
599
+ elif args.dataset_path in AIMODataset.SUPPORTED_DATASET_PATHS:
600
+ dataset_class = AIMODataset
601
+ args.hf_split = "train"
602
else:
603
supported_datasets = set([
604
dataset_name for cls in HuggingFaceDataset.__subclasses__()
@@ -610,10 +614,10 @@ def main(args: argparse.Namespace):
614
dataset_path=args.dataset_path,
615
dataset_subset=args.hf_subset,
616
dataset_split=args.hf_split,
617
+ random_seed=args.seed,
618
).sample(
619
num_requests=args.num_prompts,
620
tokenizer=tokenizer,
616
- random_seed=args.seed,
621
output_len=args.hf_output_len,
622
)
623