@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / 7c2ec0fe87

Fork 0

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 14, 05:30 UTC

[Benchmarking] Add disable_shuffle option for dataset loading (#26258)

Signed-off-by: Yasmin Moslem <48152713+ymoslem@users.noreply.github.com>

Yasmin Moslem committed Oct 6, 2025 at 08:05 UTC 7c2ec0fe873af65f74090df8bd2d5792dfccca59

1 file changed +43 -13

vllm/benchmarks/datasets.py

+43 -13

@@ -96,6 +96,8 @@ class BenchmarkDataset(ABC):

96 self,

97 dataset_path: Optional[str] = None,

98 random_seed: int = DEFAULT_SEED,

99 + disable_shuffle: bool = False,

100 + **kwargs,

101 ) -> None:

102 """

103 Initialize the BenchmarkDataset with an optional dataset path and random

@@ -111,6 +113,7 @@ class BenchmarkDataset(ABC):

113 # Set the random seed, ensuring that a None value is replaced with the

114 # default seed.

115 self.random_seed = random_seed if random_seed is not None else self.DEFAULT_SEED

116 + self.disable_shuffle = disable_shuffle

117 self.data = None

118

119 def apply_multimodal_chat_transformation(

@@ -1044,7 +1047,8 @@ class ShareGPTDataset(BenchmarkDataset):

1047 if "conversations" in entry and len(entry["conversations"]) >= 2

1048 ]

1049 random.seed(self.random_seed)

1047 - random.shuffle(self.data)

1050 + if not getattr(self, "disable_shuffle", False):

1051 + random.shuffle(self.data)

1052

1053 def sample(

1054 self,

@@ -1175,6 +1179,11 @@ def add_dataset_parser(parser: FlexibleArgumentParser):

1179 action="store_true",

1180 help="Skip applying chat template to prompt for datasets that support it.",

1181 )

1182 + parser.add_argument(

1183 + "--disable-shuffle",

1184 + action="store_true",

1185 + help="Disable shuffling of dataset samples for deterministic ordering.",

1186 + )

1187

1188 # group for dataset specific arguments

1189 custom_group = parser.add_argument_group("custom dataset options")

@@ -1441,7 +1450,9 @@ def get_samples(args, tokenizer) -> list[SampleRequest]:

1450 args.request_id_prefix = ""

1451

1452 if args.dataset_name == "custom":

1444 - dataset = CustomDataset(dataset_path=args.dataset_path)

1453 + dataset = CustomDataset(

1454 + dataset_path=args.dataset_path, disable_shuffle=args.disable_shuffle

1455 + )

1456 input_requests = dataset.sample(

1457 num_requests=args.num_prompts,

1458 tokenizer=tokenizer,

@@ -1452,7 +1463,9 @@ def get_samples(args, tokenizer) -> list[SampleRequest]:

1463 )

1464

1465 elif args.dataset_name == "sonnet":

1455 - dataset = SonnetDataset(dataset_path=args.dataset_path)

1466 + dataset = SonnetDataset(

1467 + dataset_path=args.dataset_path, disable_shuffle=args.disable_shuffle

1468 + )

1469 # For the "sonnet" dataset, formatting depends on the backend.

1470 if args.backend == "openai-chat":

1471 input_requests = dataset.sample(

@@ -1586,6 +1599,7 @@ def get_samples(args, tokenizer) -> list[SampleRequest]:

1599 random_seed=args.seed,

1600 no_stream=args.no_stream,

1601 hf_name=args.hf_name,

1602 + disable_shuffle=args.disable_shuffle,

1603 ).sample(

1604 num_requests=args.num_prompts,

1605 tokenizer=tokenizer,

@@ -1600,7 +1614,9 @@ def get_samples(args, tokenizer) -> list[SampleRequest]:

1614 # For datasets that follow a similar structure, use a mapping.

1615 dataset_mapping = {

1616 "spec_bench": lambda: SpecBench(

1603 - dataset_path=args.dataset_path, category=args.spec_bench_category

1617 + dataset_path=args.dataset_path,

1618 + category=args.spec_bench_category,

1619 + disable_shuffle=args.disable_shuffle,

1620 ).sample(

1621 num_requests=args.num_prompts,

1622 tokenizer=tokenizer,

@@ -1609,7 +1625,9 @@ def get_samples(args, tokenizer) -> list[SampleRequest]:

1625 no_oversample=args.no_oversample,

1626 ),

1627 "sharegpt": lambda: ShareGPTDataset(

1612 - random_seed=args.seed, dataset_path=args.dataset_path

1628 + random_seed=args.seed,

1629 + dataset_path=args.dataset_path,

1630 + disable_shuffle=args.disable_shuffle,

1631 ).sample(

1632 tokenizer=tokenizer,

1633 num_requests=args.num_prompts,

@@ -1618,7 +1636,9 @@ def get_samples(args, tokenizer) -> list[SampleRequest]:

1636 no_oversample=args.no_oversample,

1637 ),

1638 "burstgpt": lambda: BurstGPTDataset(

1621 - random_seed=args.seed, dataset_path=args.dataset_path

1639 + random_seed=args.seed,

1640 + dataset_path=args.dataset_path,

1641 + disable_shuffle=args.disable_shuffle,

1642 ).sample(

1643 tokenizer=tokenizer,

1644 num_requests=args.num_prompts,

@@ -1626,7 +1646,9 @@ def get_samples(args, tokenizer) -> list[SampleRequest]:

1646 no_oversample=args.no_oversample,

1647 ),

1648 "random": lambda: RandomDataset(

1629 - random_seed=args.seed, dataset_path=args.dataset_path

1649 + random_seed=args.seed,

1650 + dataset_path=args.dataset_path,

1651 + disable_shuffle=args.disable_shuffle,

1652 ).sample(

1653 tokenizer=tokenizer,

1654 num_requests=args.num_prompts,

@@ -1639,7 +1661,9 @@ def get_samples(args, tokenizer) -> list[SampleRequest]:

1661 no_oversample=args.no_oversample,

1662 ),

1663 "random-mm": lambda: RandomMultiModalDataset(

1642 - random_seed=args.seed, dataset_path=args.dataset_path

1664 + random_seed=args.seed,

1665 + dataset_path=args.dataset_path,

1666 + disable_shuffle=args.disable_shuffle,

1667 ).sample(

1668 tokenizer=tokenizer,

1669 num_requests=args.num_prompts,

@@ -1655,7 +1679,9 @@ def get_samples(args, tokenizer) -> list[SampleRequest]:

1679 no_oversample=args.no_oversample,

1680 ),

1681 "prefix_repetition": lambda: PrefixRepetitionRandomDataset(

1658 - random_seed=args.seed, dataset_path=args.dataset_path

1682 + random_seed=args.seed,

1683 + dataset_path=args.dataset_path,

1684 + disable_shuffle=args.disable_shuffle,

1685 ).sample(

1686 tokenizer=tokenizer,

1687 num_requests=args.num_prompts,

@@ -1733,7 +1759,8 @@ class CustomDataset(BenchmarkDataset):

1759 )

1760

1761 random.seed(self.random_seed)

1736 - random.shuffle(self.data)

1762 + if not getattr(self, "disable_shuffle", False):

1763 + random.shuffle(self.data)

1764

1765 def sample(

1766 self,

@@ -1825,7 +1852,8 @@ class SpecBench(CustomDataset):

1852 self.data.append({"prompt": prompt})

1853

1854 random.seed(self.random_seed)

1828 - random.shuffle(self.data)

1855 + if not getattr(self, "disable_shuffle", False):

1856 + random.shuffle(self.data)

1857

1858 def sample(self, **kwargs) -> list:

1859 # leverage CustomDataset sample

@@ -2033,7 +2061,8 @@ class HuggingFaceDataset(BenchmarkDataset):

2061 split=self.dataset_split,

2062 streaming=self.load_stream,

2063 )

2036 - self.data = self.data.shuffle(seed=self.random_seed)

2064 + if not getattr(self, "disable_shuffle", False):

2065 + self.data = self.data.shuffle(seed=self.random_seed)

2066

2067

2068 # -----------------------------------------------------------------------------

@@ -2849,7 +2878,8 @@ class PrefixRepetitionRandomDataset(BenchmarkDataset):

2878 abs(token_mismatch_total),

2879 sign,

2880 )

2852 - random.shuffle(requests)

2881 + if not getattr(self, "disable_shuffle", False):

2882 + random.shuffle(requests)

2883 return requests

2884

2885