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, 02:15 UTC
[Benchmark] Parameterization of streaming loading of multimodal datasets (#20528)
Signed-off-by: wangli <wangli858794774@gmail.com>
Li Wang committed
Jul 9, 2025 at 21:35 UTC
9ff2af6d2ba1757c6e59fe803225a73e61fe526f
4 files changed
+24
-2
benchmarks/benchmark_dataset.py
+3
-1
@@ -701,6 +701,7 @@ class HuggingFaceDataset(BenchmarkDataset):
701
self,
702
dataset_path: str,
703
dataset_split: str,
704
+ no_stream: bool = False,
705
dataset_subset: Optional[str] = None,
706
**kwargs,
707
) -> None:
@@ -708,6 +709,7 @@ class HuggingFaceDataset(BenchmarkDataset):
709
710
self.dataset_split = dataset_split
711
self.dataset_subset = dataset_subset
712
+ self.load_stream = not no_stream
713
self.load_data()
714
715
def load_data(self) -> None:
@@ -716,7 +718,7 @@ class HuggingFaceDataset(BenchmarkDataset):
718
self.dataset_path,
719
name=self.dataset_subset,
720
split=self.dataset_split,
719
- streaming=True,
721
+ streaming=self.load_stream,
722
)
723
self.data = self.data.shuffle(seed=self.random_seed)
724
benchmarks/benchmark_serving.py
+6
@@ -825,6 +825,7 @@ def main(args: argparse.Namespace):
825
dataset_subset=args.hf_subset,
826
dataset_split=args.hf_split,
827
random_seed=args.seed,
828
+ no_stream=args.no_stream,
829
).sample(
830
num_requests=args.num_prompts,
831
tokenizer=tokenizer,
@@ -1033,6 +1034,11 @@ def create_argument_parser():
1034
help="Path to the sharegpt/sonnet dataset. "
1035
"Or the huggingface dataset ID if using HF dataset.",
1036
)
1037
+ parser.add_argument(
1038
+ "--no-stream",
1039
+ action="store_true",
1040
+ help="Do not load the dataset in streaming mode.",
1041
+ )
1042
parser.add_argument(
1043
"--max-concurrency",
1044
type=int,
benchmarks/benchmark_throughput.py
+6
@@ -356,6 +356,7 @@ def get_requests(args, tokenizer):
356
elif args.dataset_name == "burstgpt":
357
dataset_cls = BurstGPTDataset
358
elif args.dataset_name == "hf":
359
+ common_kwargs["no_stream"] = args.no_stream
360
if args.dataset_path in VisionArenaDataset.SUPPORTED_DATASET_PATHS:
361
dataset_cls = VisionArenaDataset
362
common_kwargs["dataset_subset"] = None
@@ -610,6 +611,11 @@ def create_argument_parser():
611
help="Name of the dataset to benchmark on.",
612
default="sharegpt",
613
)
614
+ parser.add_argument(
615
+ "--no-stream",
616
+ action="store_true",
617
+ help="Do not load the dataset in streaming mode.",
618
+ )
619
parser.add_argument(
620
"--dataset",
621
type=str,
vllm/benchmarks/datasets.py
+9
-1
@@ -481,6 +481,11 @@ def add_dataset_parser(parser: FlexibleArgumentParser):
481
choices=["sharegpt", "burstgpt", "sonnet", "random", "hf", "custom"],
482
help="Name of the dataset to benchmark on.",
483
)
484
+ parser.add_argument(
485
+ "--no-stream",
486
+ action="store_true",
487
+ help="Do not load the dataset in streaming mode.",
488
+ )
489
parser.add_argument(
490
"--dataset-path",
491
type=str,
@@ -674,6 +679,7 @@ def get_samples(args, tokenizer) -> list[SampleRequest]:
679
dataset_subset=args.hf_subset,
680
dataset_split=args.hf_split,
681
random_seed=args.seed,
682
+ no_stream=args.no_stream,
683
).sample(
684
num_requests=args.num_prompts,
685
tokenizer=tokenizer,
@@ -971,6 +977,7 @@ class HuggingFaceDataset(BenchmarkDataset):
977
self,
978
dataset_path: str,
979
dataset_split: str,
980
+ no_stream: bool = False,
981
dataset_subset: Optional[str] = None,
982
**kwargs,
983
) -> None:
@@ -978,6 +985,7 @@ class HuggingFaceDataset(BenchmarkDataset):
985
986
self.dataset_split = dataset_split
987
self.dataset_subset = dataset_subset
988
+ self.load_stream = not no_stream
989
self.load_data()
990
991
def load_data(self) -> None:
@@ -986,7 +994,7 @@ class HuggingFaceDataset(BenchmarkDataset):
994
self.dataset_path,
995
name=self.dataset_subset,
996
split=self.dataset_split,
989
- streaming=True,
997
+ streaming=self.load_stream,
998
)
999
self.data = self.data.shuffle(seed=self.random_seed)
1000