Coverage for / home / jenkins / .local / lib / python3.10 / site-packages / hyper_parallel / data / indexed / indexed_split_builder.py: 84%
165 statements
« prev ^ index » next coverage.py v7.13.1, created at 2026-11-15 07:18 +0800
1# Copyright 2026 Huawei Technologies Co., Ltd
2#
3# Licensed under the Apache License, Version 2.0 (the "License");
4# you may not use this file except in compliance with the License.
5# You may obtain a copy of the License at
6#
7# http://www.apache.org/licenses/LICENSE-2.0
8#
9# Unless required by applicable law or agreed to in writing, software
10# distributed under the License is distributed on an "AS IS" BASIS,
11# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12# See the License for the specific language governing permissions and
13# limitations under the License.
14# ============================================================================
15"""Build train, validation, and test Datasets from indexed data blends."""
17from __future__ import annotations
19import math
20from collections.abc import Callable, Sequence
21from enum import Enum
22from functools import partial
23from typing import Any, Literal, TypeAlias, cast
25import numpy as np
27from hyper_parallel.data.dataset_logging import get_dataset_logger
28from hyper_parallel.data.indexed.indexed_blended_dataset import BlendedDataset
29from hyper_parallel.data.indexed.indexed_data_config import GPTDatasetConfig
30from hyper_parallel.data.indexed.indexed_lazy_dataset import LazyDatasetProxy
31from hyper_parallel.data.indexed.indexed_pretrain_dataset import GPTDataset, MockGPTDataset
32from hyper_parallel.data.indexed.indexed_simple_blended_dataset import SimpleBlendedDataset
33from hyper_parallel.data.parallel import (
34 DataLoaderParallelContext,
35 build_dataset_for_dataloader,
36)
38logger = get_dataset_logger(__name__)
40BlendMode: TypeAlias = Literal["no", "inter", "intra"]
41DatasetSplits: TypeAlias = tuple[Any | None, Any | None, Any | None]
42SplitRange: TypeAlias = tuple[float, float]
43SplitMatrix: TypeAlias = Sequence[SplitRange | None]
46class _Split(Enum):
47 """Stable train, validation, and test identifiers used by Dataset caches."""
49 TRAIN = 0
50 VALID = 1
51 TEST = 2
54def _get_prefixes_weights_and_sizes_for_blend(
55 blend: Sequence[str],
56 target_num_samples_per_split: Sequence[int],
57 blend_mode: BlendMode,
58 dataset_margin: float,
59) -> tuple[list[str], list[float], list[list[int]]]:
60 """Determine each source's contribution to the blended Dataset splits.
62 Args:
63 blend: Alternating ``weight, prefix`` pairs for each data source.
64 target_num_samples_per_split: Requested ``[train, validation, test]`` sample counts.
65 blend_mode: ``"no"`` for normalized weighted blending; ``"inter"`` or ``"intra"`` for simple blending.
66 dataset_margin: Extra sample margin applied only to normalized weighted blending.
68 Returns:
69 ``(prefixes, weights, sizes_per_dataset)``, where ``sizes_per_dataset[source][split]`` is the number of
70 samples requested from one source for the train, validation, or test split.
72 Example:
73 Regular: ``["30", "/data/a", "70", "/data/b"]`` with sizes ``[1000, 100, 0]`` and margin ``1.005``
74 -> ``(["/data/a", "/data/b"], [0.3, 0.7], [[302, 31, 0], [704, 71, 0]])``.
75 Simple: ``["1", "/data/a", "1", "/data/b"]`` with sizes ``[1000, 100, 0]``
76 -> ``(["/data/a", "/data/b"], [1, 1], [[500, 50, 0], [500, 50, 0]])``.
77 """
78 if len(blend) % 2 != 0:
79 raise ValueError("Dataset blend must contain alternating weight/prefix pairs")
81 # standard blend
82 if blend_mode == "no":
83 weights, prefixes = zip(*[(float(blend[index]), blend[index + 1].strip()) for index in range(0, len(blend), 2)])
84 weight_sum = float(sum(weights))
85 weights = [weight / weight_sum for weight in weights]
86 sizes_per_dataset = [
87 [
88 int(math.ceil(target_num_samples * weight * dataset_margin))
89 for target_num_samples in target_num_samples_per_split
90 ]
91 for weight in weights
92 ]
93 return prefixes, weights, sizes_per_dataset
95 # inter/intra blend
96 weights, prefixes = zip(*[(int(blend[index]), blend[index + 1].strip()) for index in range(0, len(blend), 2)])
97 sizes_per_dataset = [
98 [int(math.ceil(target_num_samples / len(weights))) for target_num_samples in target_num_samples_per_split]
99 for _ in weights
100 ]
101 return prefixes, weights, sizes_per_dataset
104def _build_low_level_dataset_and_split_indices(
105 dataset_type: type,
106 dataset_path: str,
107 split_matrix: SplitMatrix,
108 config: GPTDatasetConfig,
109) -> tuple[Any, list[np.ndarray | None]]:
110 """Build one low-level Dataset and calculate its requested split indices."""
111 low_level_dataset = dataset_type.build_low_level_dataset(dataset_path, config)
112 num_elements = dataset_type.numel_low_level_dataset(low_level_dataset)
114 split_indices = []
115 for split_range in split_matrix:
116 if split_range is None:
117 split_indices.append(None)
118 continue
120 begin = int(round(split_range[0] * num_elements))
121 end = int(round(split_range[1] * num_elements))
122 split_indice = np.arange(start=begin, stop=end, step=1, dtype=np.int32)
123 split_indices.append(split_indice)
125 # (low_level_dataset, [train_indices, valid_indices, test_indices]).
126 dataset_and_split_indices = (low_level_dataset, split_indices)
127 return dataset_and_split_indices
130class IndexedDatasetSplitBuilder:
131 """Build indexed Dataset splits with one shared distributed context."""
133 def __init__(self, dataloader_context: DataLoaderParallelContext) -> None:
134 """Store the DataLoader ownership policy once for all build steps."""
135 self.dataloader_context = dataloader_context
137 def build(
138 self,
139 dataset_type: type,
140 train_valid_test_num_samples: Sequence[int] | None,
141 config: GPTDatasetConfig,
142 blend_mode: BlendMode,
143 ) -> DatasetSplits:
144 """Build mock, shared-blend, or independent-blend Dataset splits."""
145 if train_valid_test_num_samples is None:
146 raise ValueError("train_valid_test_num_samples must be configured for indexed Datasets")
147 if len(train_valid_test_num_samples) != len(_Split):
148 raise ValueError("train_valid_test_num_samples must contain exactly three values")
149 if any(size < 0 for size in train_valid_test_num_samples):
150 raise ValueError("train_valid_test_num_samples must be non-negative")
151 if blend_mode not in {"no", "inter", "intra"}:
152 raise ValueError(f"Unsupported blend mode: {blend_mode!r}")
153 if config.mock != (dataset_type is MockGPTDataset):
154 raise ValueError("mock_data must match the selected MockGPTDataset class")
155 if not config.mock and (blend_mode == "no") != (dataset_type is GPTDataset):
156 raise ValueError("simple_blend='no' is only for GPTDataset; other Datasets require inter/intra")
158 logger.debug(
159 "Building indexed splits: dataset=%s, num_samples=%s, mock=%s, blend_mode=%s, lazy=%s",
160 dataset_type.__name__,
161 tuple(train_valid_test_num_samples),
162 config.mock,
163 blend_mode,
164 config.data_lazy_load,
165 )
166 if config.mock:
167 return self._build_mock_dataset_splits(dataset_type, train_valid_test_num_samples, config)
169 if config.blend:
170 if config.split_matrix is None:
171 raise ValueError("A shared blend requires blend and split_matrix")
172 return self._build_splits_from_blend(
173 dataset_type, config.blend, config.split_matrix, train_valid_test_num_samples, config, blend_mode
174 )
176 if config.blend_per_split:
177 return self._build_splits_from_independent_blends(
178 dataset_type, train_valid_test_num_samples, config, blend_mode
179 )
180 raise ValueError("One of blend or blend_per_split must be configured")
182 def _build_mock_dataset_splits(
183 self,
184 dataset_type: type,
185 sizes: Sequence[int],
186 config: GPTDatasetConfig,
187 ) -> DatasetSplits:
188 """Build generated train, validation, and test Datasets."""
189 datasets = []
190 for split, size in zip(_Split, sizes):
191 if size == 0:
192 datasets.append(None)
193 continue
195 dataset_factory = partial(dataset_type, None, None, None, int(size), split, config)
196 mock_dataset = self._build_dataset_for_dataloader(dataset_factory, barrier_needed=False)
197 datasets.append(mock_dataset)
198 return cast(DatasetSplits, tuple(datasets))
200 def _build_splits_from_blend(
201 self,
202 dataset_type: type,
203 blend: Sequence[str],
204 split_matrix: SplitMatrix,
205 sizes: Sequence[int],
206 config: GPTDatasetConfig,
207 blend_mode: BlendMode,
208 ) -> DatasetSplits:
209 """
210 Parse one blend and dispatch it to the single- or multiple-source path.
211 Low level: Index -> Mid level: GPTDataset -> Final: Blend.
212 """
213 if len(blend) == 1:
214 return self._build_mid_level_dataset_splits(dataset_type, blend[0], split_matrix, sizes, config)
216 prefixes, weights, sizes_per_dataset = _get_prefixes_weights_and_sizes_for_blend(
217 blend, sizes, blend_mode, config.dataset_margin
218 )
219 logger.debug(
220 "Resolved blend: mode=%s, sources=%d, first_weight=%s, first_source_sizes=%s",
221 blend_mode,
222 len(prefixes),
223 weights[0],
224 sizes_per_dataset[0],
225 )
226 dataset_splits = self._build_multiple_source_splits(
227 dataset_type, prefixes, weights, sizes_per_dataset, split_matrix, config, blend_mode
228 )
229 return dataset_splits
231 def _build_splits_from_independent_blends(
232 self,
233 dataset_type: type,
234 sizes: Sequence[int],
235 config: GPTDatasetConfig,
236 blend_mode: BlendMode,
237 ) -> DatasetSplits:
238 """Build each split from its own data-source distribution."""
239 if config.blend_per_split is None or len(config.blend_per_split) != len(_Split):
240 raise ValueError("Independent blends require one blend per split")
242 datasets = []
243 for split_index, blend in enumerate(config.blend_per_split):
244 if not blend:
245 datasets.append(None)
246 continue
248 split_matrix: list[SplitRange | None] = [None] * len(_Split)
249 split_matrix[split_index] = (0.0, 1.0)
250 split_sizes = [0] * len(_Split)
251 split_sizes[split_index] = int(sizes[split_index])
252 split_datasets = self._build_splits_from_blend(
253 dataset_type, blend, split_matrix, split_sizes, config, blend_mode
254 )
255 datasets.append(split_datasets[split_index])
256 return cast(DatasetSplits, tuple(datasets))
258 def _build_multiple_source_splits(
259 self,
260 dataset_type: type,
261 prefixes: Sequence[str],
262 weights: Sequence[float],
263 sizes_per_dataset: Sequence[Sequence[int]],
264 split_matrix: SplitMatrix,
265 config: GPTDatasetConfig,
266 blend_mode: BlendMode,
267 ) -> DatasetSplits:
268 """Build each source's mid-level splits and assemble top-level blends.
270 Example with two sources::
272 prefixes = ["/data/books", "/data/code"]
273 sizes_per_dataset = [[80, 10, 0], [20, 5, 0]]
275 Offline preparation, tokenize and optionally pre-pack raw records:
276 books.jsonl -> tokenizer -> token IDs -> pack to 2049-token records
277 -> /data/books.idx + /data/books.bin
278 code.jsonl -> tokenizer -> token IDs -> pack to 2049-token records
279 -> /data/code.idx + /data/code.bin
281 Low-level, open one IndexedDataReader per prefix and read one sequence (uint16 example):
282 books.idx: sequence_lengths[1] = 2049, sequence_pointers[1] = 4098
283 books_reader[1] -> books.bin[offset=4098, count=2049] -> text token IDs
285 Mid-level, organize tokens into training records (input_ids, labels, and masks),
286 with one GPTDataset/GPTFromMRDataset per source and split:
287 books_train[0] -> split indices[0] = 1 -> books_reader[1]
288 -> input_ids = text[:-1], labels = text[1:]
289 datasets_per_source = [
290 (books_train[80], books_valid[10], None),
291 (code_train[20], code_valid[5], None),
292 ]
294 Transpose source-major tuples into split-major components:
295 train components = [books_train, code_train], total size = 100
296 valid components = [books_valid, code_valid], total size = 15
297 test components = [], size = 0
299 Blend-level, select the source and source-local sample for each train index, for example:
300 dataset_index = [0, 1, ...]
301 dataset_sample_index = [0, 0, ...]
302 global sample 0 -> (dataset_id=0, sample_id=0) -> books_train[0]
303 global sample 1 -> (dataset_id=1, sample_id=0) -> code_train[0]
305 ``BlendedDataset[i]`` reads ``dataset_index[i]`` and ``dataset_sample_index[i]``;
306 ``SimpleBlendedDataset[i]`` reads ``_locations[i] = (dataset_id, sample_id)``.
307 Both return ``datasets[dataset_id][sample_id]``.
309 final result = (blended_train[100], blended_valid[15], None)
310 """
311 # pylint: disable=too-many-locals
312 datasets_per_source = []
313 logger.debug("Building %d indexed Dataset sources", len(prefixes))
314 # TODO: Add a separate cache-prebuild phase that assigns source_index % world_size to each rank,
315 # synchronizes failures, and atomically writes a completion manifest before Dataset construction.
317 # books -> (books_train, books_valid, None), then code -> (code_train, code_valid, None).
318 for prefix, sizes_for_source in zip(prefixes, sizes_per_dataset):
319 source_datasets = self._build_mid_level_dataset_splits(
320 dataset_type, prefix, split_matrix, sizes_for_source, config
321 )
322 datasets_per_source.append(source_datasets)
324 blended_sizes = [sum(split_sizes) for split_sizes in zip(*sizes_per_dataset)]
325 blended_splits = []
326 for split_index, (split_range, size) in enumerate(zip(split_matrix, blended_sizes)):
327 if split_range is None or size == 0:
328 blended_splits.append(None)
329 continue
331 component_datasets = [datasets[split_index] for datasets in datasets_per_source]
332 missing_datasets = [dataset is None for dataset in component_datasets]
333 if any(missing_datasets) and not all(missing_datasets):
334 raise ValueError("Enabled sources must either all build or all remain None on a rank")
336 component_datasets = [dataset for dataset in component_datasets if dataset is not None]
337 if blend_mode == "no":
338 dataset_factory = partial(BlendedDataset, component_datasets, weights, int(size), config)
339 else:
340 dataset_factory = partial(SimpleBlendedDataset, component_datasets, int(size), blend_mode)
341 blended_dataset = self._build_dataset_for_dataloader(dataset_factory, barrier_needed=True)
342 blended_splits.append(blended_dataset)
344 dataset_splits = cast(DatasetSplits, tuple(blended_splits))
345 split_types = tuple(type(split).__name__ if split is not None else None for split in dataset_splits)
346 logger.debug("Built blended splits=%s", split_types)
347 return dataset_splits
349 def _build_mid_level_dataset_splits(
350 self,
351 dataset_type: type,
352 dataset_path: str,
353 split_matrix: SplitMatrix,
354 sizes_for_source: Sequence[int],
355 config: GPTDatasetConfig,
356 ) -> DatasetSplits:
357 """Build each mid-level split from one low-level indexed Dataset."""
358 # pylint: disable=too-many-locals
359 if (
360 config.data_lazy_load
361 and self.dataloader_context.distributed_enabled
362 and not (
363 self.dataloader_context.data_index_cache
364 or self.dataloader_context.build_on_rank()
365 )
366 ):
367 return cast(DatasetSplits, (None, None, None))
369 low_level_dataset = None
370 split_indices: list[np.ndarray | None] = [None] * len(split_matrix)
371 if not config.data_lazy_load:
372 low_level_dataset, split_indices = _build_low_level_dataset_and_split_indices(
373 dataset_type, dataset_path, split_matrix, config
374 )
376 datasets = []
377 for split, split_range, indices, size in zip(_Split, split_matrix, split_indices, sizes_for_source):
378 if split_range is None or size == 0:
379 datasets.append(None)
380 continue
382 size = int(size)
383 if config.data_lazy_load:
384 lazy_builder = self._build_lazy_mid_level_dataset_split
385 dataset_factory = partial(lazy_builder, dataset_type, dataset_path, split, split_range, size, config)
386 unique_identifiers = dataset_type.build_unique_identifiers(dataset_path, size, split, config)
387 mid_level_dataset = LazyDatasetProxy(dataset_factory, unique_identifiers=unique_identifiers)
388 else:
389 dataset_factory = partial(dataset_type, low_level_dataset, dataset_path, indices, size, split, config)
390 mid_level_dataset = self._build_dataset_for_dataloader(
391 dataset_factory, barrier_needed=dataset_type is GPTDataset
392 )
393 datasets.append(mid_level_dataset)
394 return cast(DatasetSplits, tuple(datasets))
396 def _build_lazy_mid_level_dataset_split(
397 self,
398 dataset_type: type,
399 dataset_path: str,
400 split: _Split,
401 split_range: SplitRange,
402 size: int,
403 config: GPTDatasetConfig,
404 ) -> Any | None:
405 """Apply distributed construction when a lazy split is first accessed."""
407 def _dataset_factory() -> Any:
408 """Build the low-level Dataset and one lazy split on first access."""
409 low_level_dataset, split_indices = _build_low_level_dataset_and_split_indices(
410 dataset_type, dataset_path, [split_range], config
411 )
412 mid_level_dataset = dataset_type(low_level_dataset, dataset_path, split_indices[0], size, split, config)
413 return mid_level_dataset
415 local_dataset = self._build_dataset_for_dataloader(_dataset_factory, barrier_needed=False)
416 return local_dataset
418 def _build_dataset_for_dataloader(
419 self,
420 dataset_factory: Callable[[], Any],
421 *,
422 barrier_needed: bool,
423 ) -> Any | None:
424 """Build one Dataset with the shared DataLoader context."""
425 local_dataset = build_dataset_for_dataloader(
426 dataset_factory,
427 self.dataloader_context,
428 barrier_needed=barrier_needed,
429 )
430 return local_dataset