Coverage for / home / jenkins / .local / lib / python3.10 / site-packages / hyper_parallel / data / indexed / indexed_split_builder.py: 0%
162 statements
« prev ^ index » next coverage.py v7.13.1, created at 2026-10-25 05:45 +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_simple_blended_dataset import SimpleBlendedDataset
32from hyper_parallel.data.parallel import (
33 DataLoaderParallelContext,
34 build_dataset_for_dataloader,
35)
37logger = get_dataset_logger(__name__)
39BlendMode: TypeAlias = Literal["no", "inter", "intra"]
40DatasetSplits: TypeAlias = tuple[Any | None, Any | None, Any | None]
41SplitRange: TypeAlias = tuple[float, float]
42SplitMatrix: TypeAlias = Sequence[SplitRange | None]
45class _Split(Enum):
46 """Stable train, validation, and test identifiers used by Dataset caches."""
48 TRAIN = 0
49 VALID = 1
50 TEST = 2
53def _get_prefixes_weights_and_sizes_for_blend(
54 blend: Sequence[str],
55 target_num_samples_per_split: Sequence[int],
56 blend_mode: BlendMode,
57 dataset_margin: float,
58) -> tuple[list[str], list[float], list[list[int]]]:
59 """Determine each source's contribution to the blended Dataset splits.
61 Args:
62 blend: Alternating ``weight, prefix`` pairs for each data source.
63 target_num_samples_per_split: Requested ``[train, validation, test]`` sample counts.
64 blend_mode: ``"no"`` for normalized weighted blending; ``"inter"`` or ``"intra"`` for simple blending.
65 dataset_margin: Extra sample margin applied only to normalized weighted blending.
67 Returns:
68 ``(prefixes, weights, sizes_per_dataset)``, where ``sizes_per_dataset[source][split]`` is the number of
69 samples requested from one source for the train, validation, or test split.
71 Example:
72 Regular: ``["30", "/data/a", "70", "/data/b"]`` with sizes ``[1000, 100, 0]`` and margin ``1.005``
73 -> ``(["/data/a", "/data/b"], [0.3, 0.7], [[302, 31, 0], [704, 71, 0]])``.
74 Simple: ``["1", "/data/a", "1", "/data/b"]`` with sizes ``[1000, 100, 0]``
75 -> ``(["/data/a", "/data/b"], [1, 1], [[500, 50, 0], [500, 50, 0]])``.
76 """
77 if len(blend) % 2 != 0:
78 raise ValueError("Dataset blend must contain alternating weight/prefix pairs")
80 # standard blend
81 if blend_mode == "no":
82 weights, prefixes = zip(*[(float(blend[index]), blend[index + 1].strip()) for index in range(0, len(blend), 2)])
83 weight_sum = float(sum(weights))
84 weights = [weight / weight_sum for weight in weights]
85 sizes_per_dataset = [
86 [
87 int(math.ceil(target_num_samples * weight * dataset_margin))
88 for target_num_samples in target_num_samples_per_split
89 ]
90 for weight in weights
91 ]
92 return prefixes, weights, sizes_per_dataset
94 # inter/intra blend
95 weights, prefixes = zip(*[(int(blend[index]), blend[index + 1].strip()) for index in range(0, len(blend), 2)])
96 sizes_per_dataset = [
97 [int(math.ceil(target_num_samples / len(weights))) for target_num_samples in target_num_samples_per_split]
98 for _ in weights
99 ]
100 return prefixes, weights, sizes_per_dataset
103def _build_low_level_dataset_and_split_indices(
104 dataset_type: type,
105 dataset_path: str,
106 split_matrix: SplitMatrix,
107 config: GPTDatasetConfig,
108) -> tuple[Any, list[np.ndarray | None]]:
109 """Build one low-level Dataset and calculate its requested split indices."""
110 low_level_dataset = dataset_type.build_low_level_dataset(dataset_path, config)
111 num_elements = dataset_type.numel_low_level_dataset(low_level_dataset)
113 split_indices = []
114 for split_range in split_matrix:
115 if split_range is None:
116 split_indices.append(None)
117 continue
119 begin = int(round(split_range[0] * num_elements))
120 end = int(round(split_range[1] * num_elements))
121 split_indice = np.arange(start=begin, stop=end, step=1, dtype=np.int32)
122 split_indices.append(split_indice)
124 # (low_level_dataset, [train_indices, valid_indices, test_indices]).
125 dataset_and_split_indices = (low_level_dataset, split_indices)
126 return dataset_and_split_indices
129class IndexedDatasetSplitBuilder:
130 """Build indexed Dataset splits with one shared distributed context."""
132 def __init__(self, dataloader_context: DataLoaderParallelContext) -> None:
133 """Store the DataLoader ownership policy once for all build steps."""
134 self.dataloader_context = dataloader_context
136 def build(
137 self,
138 dataset_type: type,
139 train_valid_test_num_samples: Sequence[int] | None,
140 config: GPTDatasetConfig,
141 blend_mode: BlendMode,
142 ) -> DatasetSplits:
143 """Build mock, shared-blend, or independent-blend Dataset splits."""
144 if train_valid_test_num_samples is None:
145 raise ValueError("train_valid_test_num_samples must be configured for indexed Datasets")
146 if len(train_valid_test_num_samples) != len(_Split):
147 raise ValueError("train_valid_test_num_samples must contain exactly three values")
148 if any(size < 0 for size in train_valid_test_num_samples):
149 raise ValueError("train_valid_test_num_samples must be non-negative")
150 if blend_mode not in {"no", "inter", "intra"}:
151 raise ValueError(f"Unsupported blend mode: {blend_mode!r}")
152 if blend_mode != "no" and not config.is_dataset_from_mr:
153 raise ValueError("inter/intra blending requires is_dataset_from_mr=True")
155 logger.debug(
156 "Building indexed splits: dataset=%s, num_samples=%s, mock=%s, blend_mode=%s, lazy=%s",
157 dataset_type.__name__,
158 tuple(train_valid_test_num_samples),
159 config.mock,
160 blend_mode,
161 config.data_lazy_load,
162 )
163 if config.mock:
164 return self._build_mock_dataset_splits(dataset_type, train_valid_test_num_samples, config)
166 if config.blend:
167 if config.split_matrix is None:
168 raise ValueError("A shared blend requires blend and split_matrix")
169 return self._build_splits_from_blend(
170 dataset_type, config.blend, config.split_matrix, train_valid_test_num_samples, config, blend_mode
171 )
173 if config.blend_per_split:
174 return self._build_splits_from_independent_blends(
175 dataset_type, train_valid_test_num_samples, config, blend_mode
176 )
177 raise ValueError("One of blend or blend_per_split must be configured")
179 def _build_mock_dataset_splits(
180 self,
181 dataset_type: type,
182 sizes: Sequence[int],
183 config: GPTDatasetConfig,
184 ) -> DatasetSplits:
185 """Build generated train, validation, and test Datasets."""
186 datasets = []
187 for split, size in zip(_Split, sizes):
188 if size == 0:
189 datasets.append(None)
190 continue
192 dataset_factory = partial(dataset_type, None, None, None, int(size), split, config)
193 mock_dataset = self._build_dataset_for_dataloader(dataset_factory, barrier_needed=False)
194 datasets.append(mock_dataset)
195 return cast(DatasetSplits, tuple(datasets))
197 def _build_splits_from_blend(
198 self,
199 dataset_type: type,
200 blend: Sequence[str],
201 split_matrix: SplitMatrix,
202 sizes: Sequence[int],
203 config: GPTDatasetConfig,
204 blend_mode: BlendMode,
205 ) -> DatasetSplits:
206 """
207 Parse one blend and dispatch it to the single- or multiple-source path.
208 Low level: Index -> Mid level: GPTDataset -> Final: Blend.
209 """
210 if len(blend) == 1:
211 return self._build_mid_level_dataset_splits(dataset_type, blend[0], split_matrix, sizes, config)
213 prefixes, weights, sizes_per_dataset = _get_prefixes_weights_and_sizes_for_blend(
214 blend, sizes, blend_mode, config.dataset_margin
215 )
216 logger.debug(
217 "Resolved blend: mode=%s, sources=%d, first_weight=%s, first_source_sizes=%s",
218 blend_mode,
219 len(prefixes),
220 weights[0],
221 sizes_per_dataset[0],
222 )
223 dataset_splits = self._build_multiple_source_splits(
224 dataset_type, prefixes, weights, sizes_per_dataset, split_matrix, config, blend_mode
225 )
226 return dataset_splits
228 def _build_splits_from_independent_blends(
229 self,
230 dataset_type: type,
231 sizes: Sequence[int],
232 config: GPTDatasetConfig,
233 blend_mode: BlendMode,
234 ) -> DatasetSplits:
235 """Build each split from its own data-source distribution."""
236 if config.blend_per_split is None or len(config.blend_per_split) != len(_Split):
237 raise ValueError("Independent blends require one blend per split")
239 datasets = []
240 for split_index, blend in enumerate(config.blend_per_split):
241 if not blend:
242 datasets.append(None)
243 continue
245 split_matrix: list[SplitRange | None] = [None] * len(_Split)
246 split_matrix[split_index] = (0.0, 1.0)
247 split_sizes = [0] * len(_Split)
248 split_sizes[split_index] = int(sizes[split_index])
249 split_datasets = self._build_splits_from_blend(
250 dataset_type, blend, split_matrix, split_sizes, config, blend_mode
251 )
252 datasets.append(split_datasets[split_index])
253 return cast(DatasetSplits, tuple(datasets))
255 def _build_multiple_source_splits(
256 self,
257 dataset_type: type,
258 prefixes: Sequence[str],
259 weights: Sequence[float],
260 sizes_per_dataset: Sequence[Sequence[int]],
261 split_matrix: SplitMatrix,
262 config: GPTDatasetConfig,
263 blend_mode: BlendMode,
264 ) -> DatasetSplits:
265 """Build each source's mid-level splits and assemble top-level blends.
267 Example with two sources::
269 prefixes = ["/data/books", "/data/code"]
270 sizes_per_dataset = [[80, 10, 0], [20, 5, 0]]
272 Offline preparation, tokenize and optionally pre-pack raw records:
273 books.jsonl -> tokenizer -> token IDs -> pack to 2049-token records
274 -> /data/books.idx + /data/books.bin
275 code.jsonl -> tokenizer -> token IDs -> pack to 2049-token records
276 -> /data/code.idx + /data/code.bin
278 Low-level, open one IndexedDataReader per prefix and read one sequence (uint16 example):
279 books.idx: sequence_lengths[1] = 2049, sequence_pointers[1] = 4098
280 books_reader[1] -> books.bin[offset=4098, count=2049] -> text token IDs
282 Mid-level, organize tokens into training records (tokens, labels, and masks),
283 with one GPTDataset/GPTFromMRDataset per source and split:
284 books_train[0] -> split indices[0] = 1 -> books_reader[1]
285 -> tokens = text[:-1], labels = text[1:]
286 datasets_per_source = [
287 (books_train[80], books_valid[10], None),
288 (code_train[20], code_valid[5], None),
289 ]
291 Transpose source-major tuples into split-major components:
292 train components = [books_train, code_train], total size = 100
293 valid components = [books_valid, code_valid], total size = 15
294 test components = [], size = 0
296 Blend-level, select the source and source-local sample for each train index, for example:
297 dataset_index = [0, 1, ...]
298 dataset_sample_index = [0, 0, ...]
299 global sample 0 -> (dataset_id=0, sample_id=0) -> books_train[0]
300 global sample 1 -> (dataset_id=1, sample_id=0) -> code_train[0]
302 ``BlendedDataset[i]`` reads ``dataset_index[i]`` and ``dataset_sample_index[i]``;
303 ``SimpleBlendedDataset[i]`` reads ``_locations[i] = (dataset_id, sample_id)``.
304 Both return ``datasets[dataset_id][sample_id]``.
306 final result = (blended_train[100], blended_valid[15], None)
307 """
308 # pylint: disable=too-many-locals
309 datasets_per_source = []
310 logger.debug("Building %d indexed Dataset sources", len(prefixes))
311 # TODO: Add a separate cache-prebuild phase that assigns source_index % world_size to each rank,
312 # synchronizes failures, and atomically writes a completion manifest before Dataset construction.
314 # books -> (books_train, books_valid, None), then code -> (code_train, code_valid, None).
315 for prefix, sizes_for_source in zip(prefixes, sizes_per_dataset):
316 source_datasets = self._build_mid_level_dataset_splits(
317 dataset_type, prefix, split_matrix, sizes_for_source, config
318 )
319 datasets_per_source.append(source_datasets)
321 blended_sizes = [sum(split_sizes) for split_sizes in zip(*sizes_per_dataset)]
322 blended_splits = []
323 for split_index, (split_range, size) in enumerate(zip(split_matrix, blended_sizes)):
324 if split_range is None or size == 0:
325 blended_splits.append(None)
326 continue
328 component_datasets = [datasets[split_index] for datasets in datasets_per_source]
329 missing_datasets = [dataset is None for dataset in component_datasets]
330 if any(missing_datasets) and not all(missing_datasets):
331 raise ValueError("Enabled sources must either all build or all remain None on a rank")
333 component_datasets = [dataset for dataset in component_datasets if dataset is not None]
334 if blend_mode == "no":
335 dataset_factory = partial(BlendedDataset, component_datasets, weights, int(size), config)
336 else:
337 dataset_factory = partial(SimpleBlendedDataset, component_datasets, int(size), blend_mode)
338 blended_dataset = self._build_dataset_for_dataloader(dataset_factory, barrier_needed=True)
339 blended_splits.append(blended_dataset)
341 dataset_splits = cast(DatasetSplits, tuple(blended_splits))
342 split_types = tuple(type(split).__name__ if split is not None else None for split in dataset_splits)
343 logger.debug("Built blended splits=%s", split_types)
344 return dataset_splits
346 def _build_mid_level_dataset_splits(
347 self,
348 dataset_type: type,
349 dataset_path: str,
350 split_matrix: SplitMatrix,
351 sizes_for_source: Sequence[int],
352 config: GPTDatasetConfig,
353 ) -> DatasetSplits:
354 """Build each mid-level split from one low-level indexed Dataset."""
355 # pylint: disable=too-many-locals
356 if (
357 config.data_lazy_load
358 and self.dataloader_context.distributed_enabled
359 and not (self.dataloader_context.data_index_cache or self.dataloader_context.build_on_rank())
360 ):
361 return cast(DatasetSplits, (None, None, None))
363 low_level_dataset = None
364 split_indices: list[np.ndarray | None] = [None] * len(split_matrix)
365 if not config.data_lazy_load:
366 low_level_dataset, split_indices = _build_low_level_dataset_and_split_indices(
367 dataset_type, dataset_path, split_matrix, config
368 )
370 datasets = []
371 for split, split_range, indices, size in zip(_Split, split_matrix, split_indices, sizes_for_source):
372 if split_range is None or size == 0:
373 datasets.append(None)
374 continue
376 size = int(size)
377 if config.data_lazy_load:
378 lazy_builder = self._build_lazy_mid_level_dataset_split
379 dataset_factory = partial(lazy_builder, dataset_type, dataset_path, split, split_range, size, config)
380 unique_identifiers = dataset_type.build_unique_identifiers(dataset_path, size, split, config)
381 mid_level_dataset = LazyDatasetProxy(dataset_factory, unique_identifiers=unique_identifiers)
382 else:
383 dataset_factory = partial(dataset_type, low_level_dataset, dataset_path, indices, size, split, config)
384 mid_level_dataset = self._build_dataset_for_dataloader(
385 dataset_factory, barrier_needed=not config.is_dataset_from_mr
386 )
387 datasets.append(mid_level_dataset)
388 return cast(DatasetSplits, tuple(datasets))
390 def _build_lazy_mid_level_dataset_split(
391 self,
392 dataset_type: type,
393 dataset_path: str,
394 split: _Split,
395 split_range: SplitRange,
396 size: int,
397 config: GPTDatasetConfig,
398 ) -> Any | None:
399 """Apply distributed construction when a lazy split is first accessed."""
401 def _dataset_factory() -> Any:
402 """Build the low-level Dataset and one lazy split on first access."""
403 low_level_dataset, split_indices = _build_low_level_dataset_and_split_indices(
404 dataset_type, dataset_path, [split_range], config
405 )
406 mid_level_dataset = dataset_type(low_level_dataset, dataset_path, split_indices[0], size, split, config)
407 return mid_level_dataset
409 local_dataset = self._build_dataset_for_dataloader(_dataset_factory, barrier_needed=False)
410 return local_dataset
412 def _build_dataset_for_dataloader(
413 self,
414 dataset_factory: Callable[[], Any],
415 *,
416 barrier_needed: bool,
417 ) -> Any | None:
418 """Build one Dataset with the shared DataLoader context."""
419 local_dataset = build_dataset_for_dataloader(
420 dataset_factory,
421 self.dataloader_context,
422 barrier_needed=barrier_needed,
423 )
424 return local_dataset