Coverage for / home / jenkins / .local / lib / python3.10 / site-packages / hyper_parallel / data / indexed / indexed_data_config.py: 0%
157 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"""Resolve indexed data paths and normalize GPT Dataset build options."""
17from __future__ import annotations
19import os
20import re
21from collections.abc import Mapping, Sequence
22from dataclasses import dataclass, field
23from typing import Any
25import numpy as np
27import torch.distributed as dist
29from hyper_parallel.data.dataset_logging import get_dataset_logger
31logger = get_dataset_logger(__name__)
34@dataclass
35class BlendedMegatronDatasetConfig:
36 """Configuration shared by blended indexed datasets."""
38 random_seed: int
39 sequence_length: int
40 blend: list[str] | None = None
41 blend_per_split: list[list[str] | None] | None = None
42 split: str | None = None
43 split_matrix: list[tuple[float, float] | None] | None = field(init=False, default=None)
44 path_to_cache: str | None = None
45 mmap_bin_files: bool = True
46 mock: bool = False
47 tokenizer: Any = None
49 def __post_init__(self) -> None:
50 """Validate the configured data distribution and derive split ranges."""
51 if self.mock:
52 return
54 if self.blend_per_split is not None and any(self.blend_per_split):
55 if self.blend is not None:
56 raise ValueError("blend and blend_per_split are incompatible")
57 if self.split is not None:
58 raise ValueError("split and blend_per_split are incompatible")
59 if len(self.blend_per_split) != 3:
60 raise ValueError("blend_per_split must contain train, validation, and test blends")
61 return
63 if self.blend is None:
64 raise ValueError("one of either blend or blend_per_split must be provided")
65 if self.split is None:
66 raise ValueError("both blend and split must be provided")
68 # Parse at most three train/validation/test ratios, padding omitted
69 # trailing splits with zero before normalization.
70 split_values = [float(value) for value in re.findall(r"[.0-9]+", self.split)]
71 split_values.extend([0.0] * (3 - len(split_values)))
72 if len(split_values) != 3:
73 raise ValueError(f"split length {len(split_values)} does not match expected 3")
75 if not all(value >= 0.0 for value in split_values):
76 raise ValueError("all split values must be non-negative")
78 normalized_values = (np.asarray(split_values, dtype=np.float64) / np.sum(split_values)).tolist()
79 canonical_values = []
80 for value in normalized_values:
81 rounded_value = round(value, 12)
82 precision = 2 if rounded_value == round(rounded_value, 2) else 12
83 canonical_values.append(f"{rounded_value:.{precision}f}")
84 self.split = ", ".join(canonical_values)
86 # Convert normalized ratios into non-overlapping ranges. A zero-sized
87 # split is represented by None.
88 boundaries = [0.0]
89 for value in normalized_values:
90 boundaries.append(boundaries[-1] + value)
91 self.split_matrix = []
92 for begin, end in zip(boundaries[:-1], boundaries[1:]):
93 self.split_matrix.append(None if end <= begin else (begin, end))
96@dataclass
97class GPTDatasetConfig(BlendedMegatronDatasetConfig):
98 """Configuration required to build indexed GPT Dataset splits."""
100 reset_position_ids: bool = False
101 reset_attention_mask: bool = False
102 eod_mask_loss: bool = False
103 create_ltor_fields: bool = False
104 drop_last_partial_validation_sequence: bool = True
105 add_extra_token_to_sequence: bool = True
106 allow_ambiguous_pad_tokens: bool = False
107 dataset_margin: float = 1.005
108 skip_data_check: bool = False
109 reuse_idx: bool = False
110 data_lazy_load: bool = False
112 is_dataset_from_mr: bool = False
113 simple_blend: str = "no"
115 def __post_init__(self) -> None:
116 """Validate the blended Dataset fields followed by GPT-specific fields."""
117 super().__post_init__()
119 if self.sequence_length <= 0:
120 raise ValueError("sequence_length must be positive")
121 if self.tokenizer is None and not self.mock:
122 raise ValueError("Attribute 'tokenizer' must not be None")
125def _discover_indexed_paths(
126 current_path: str,
127 rank: int,
128 world_size: int,
129 distributed_walk: bool,
130) -> list[str]:
131 """Discover indexed prefixes under one directory."""
132 if not distributed_walk:
133 return _walk_bin_directories([current_path])
135 root_layout = None
136 if rank == 0:
137 subdirectories = []
138 top_files = []
139 for name in os.listdir(current_path):
140 path = os.path.join(current_path, name)
141 if os.path.isdir(path):
142 subdirectories.append(path)
143 elif name.endswith(".bin"):
144 top_files.append(os.path.splitext(path)[0])
145 root_layout = (subdirectories, top_files)
147 gathered_layouts: list[tuple[list[str], list[str]] | None] = [None] * world_size
148 dist.all_gather_object(gathered_layouts, root_layout)
149 rank_zero_layout = gathered_layouts[0]
150 if rank_zero_layout is None:
151 raise ValueError("Rank 0 did not provide the indexed directory layout")
152 subdirectories, file_paths = rank_zero_layout
153 local_files = _walk_bin_directories(subdirectories[rank::world_size])
154 gathered_files: list[list[str] | None] = [None] * world_size
155 dist.all_gather_object(gathered_files, local_files)
156 for rank_files in gathered_files:
157 if rank_files:
158 file_paths.extend(rank_files)
159 return file_paths
162def resolve_data_paths(
163 data_path: str | Sequence[str],
164 *,
165 distributed_walk: bool = True,
166) -> list[str]:
167 """Find ``.bin`` files and return indexed prefixes with optional weights.
169 Args:
170 data_path: One indexed prefix or a sequence of directories/prefixes.
171 distributed_walk: Whether initialized ranks divide directory walking.
173 Returns:
174 One prefix directly, or alternating weight/prefix values for multiple sources. If ``data_path`` is not a
175 collection of discoverable indexed paths, return it unchanged so an explicit blend can be parsed.
176 """
177 # 1. Normalize one path and multiple paths to the same traversal form.
178 data_paths = [data_path] if isinstance(data_path, str) else list(data_path)
180 # 2. Use distributed walking only after the process group is available.
181 # Before initialization, one process walks all configured directories.
182 try:
183 rank = int(dist.get_rank())
184 world_size = int(dist.get_world_size())
185 except (RuntimeError, ValueError):
186 rank, world_size = 0, 1
187 use_distributed_walk = distributed_walk and world_size > 1
189 # 3. Accept a direct prefix, or recursively discover prefixes in a directory.
190 file_paths: list[str] = []
191 for current_path in data_paths:
192 if os.path.isfile(current_path + ".bin"):
193 file_paths.append(current_path)
194 continue
195 if not os.path.isdir(current_path):
196 # Keep args.data_path unchanged when filesystem discovery
197 # returns no files. This preserves explicit forms such as
198 # ["30", prefix_a, "70", prefix_b] for the blend parser.
199 original_paths = list(data_paths)
200 return original_paths
202 file_paths.extend(_discover_indexed_paths(current_path, rank, world_size, use_distributed_walk))
204 if not file_paths:
205 raise ValueError(f"data_path {data_paths!r} has no data file ending with .bin")
207 if len(file_paths) == 1:
208 return file_paths
210 # 4. Prefer the numeric shard order encoded by the corpus filenames.
211 # Unknown naming schemes remain usable in their original walk order.
212 try:
213 def numeric_shard_key(path: str) -> tuple[int, int]:
214 """Extract numeric corpus and shard identifiers from a path."""
215 path_parts = os.path.basename(path).split("_")
216 return int(path_parts[2]), int(path_parts[-1])
218 file_paths.sort(key=numeric_shard_key)
219 except (IndexError, TypeError, ValueError):
220 logger.warning("Cannot sort indexed files with the numeric filename rule; keeping walk order")
222 # 5. Automatically discovered sources use equal blend weights.
223 weighted_paths = [value for path in file_paths for value in ("1", path)]
224 return weighted_paths
227def build_gpt_dataset_config(
228 data_paths: Sequence[str],
229 data_config: Mapping[str, Any],
230) -> GPTDatasetConfig:
231 """Build a GPT Dataset config from provider fields.
233 Args:
234 data_paths: Weighted prefixes returned by :func:`resolve_data_paths`.
235 data_config: Indexed pretraining configuration supplied by the caller.
237 Returns:
238 Configuration required by GPT Dataset and split construction.
240 Raises:
241 ValueError: If a required build field is missing.
242 """
243 # The source argument parser owns defaults and validation. This stage only
244 # renames fields where the Dataset implementation expects another name.
245 blend_per_split = [
246 data_config.get("train_data_path"),
247 data_config.get("valid_data_path"),
248 data_config.get("test_data_path"),
249 ]
250 has_independent_blends = any(blend_per_split)
251 try:
252 mock_data = bool(data_config["mock_data"])
253 tokenizer = data_config.get("tokenizer")
254 if tokenizer is None and not mock_data:
255 raise ValueError("tokenizer is required when mock_data is false")
257 config = GPTDatasetConfig(
258 random_seed=data_config["random_seed"],
259 sequence_length=data_config["seq_length"],
260 tokenizer=tokenizer,
261 mock=mock_data,
262 blend=None if has_independent_blends else list(data_paths),
263 blend_per_split=blend_per_split if has_independent_blends else None,
264 split=None if has_independent_blends else data_config["split"],
265 path_to_cache=data_config.get("data_cache_path"),
266 mmap_bin_files=data_config.get("mmap_bin_files", True),
267 reuse_idx=data_config.get("reuse_idx", False),
268 data_lazy_load=data_config["data_lazy_load"],
269 reset_position_ids=data_config.get("reset_position_ids", False),
270 reset_attention_mask=data_config.get("reset_attention_mask", False),
271 eod_mask_loss=data_config.get("eod_mask_loss", False),
272 create_ltor_fields=data_config.get("create_ltor_fields_in_dataloader", False),
273 drop_last_partial_validation_sequence=data_config.get("drop_last_partial_validation_sequence", True),
274 add_extra_token_to_sequence=data_config.get("add_extra_token_to_sequence", True),
275 allow_ambiguous_pad_tokens=data_config.get("allow_ambiguous_pad_tokens", False),
276 dataset_margin=data_config.get("dataset_margin", 1.005),
277 skip_data_check=data_config.get("skip_data_check", False),
278 is_dataset_from_mr=data_config["is_dataset_from_mr"],
279 simple_blend=data_config["simple_blend"],
280 )
281 except KeyError as error:
282 raise ValueError(f"Missing indexed Dataset build option: {error.args[0]}") from error
283 source_paths = list(data_paths) if len(data_paths) <= 1 else list(data_paths[1::2])
284 logger.debug("Resolved indexed sources=%d, first four=%s", len(source_paths), source_paths[:4])
285 return config
288def _walk_bin_directories(directories: Sequence[str]) -> list[str]:
289 """Recursively find ``.bin`` files below the assigned directories."""
290 file_paths = []
291 for directory in directories:
292 for root, _, filenames in os.walk(directory):
293 for filename in filenames:
294 if filename.endswith(".bin"):
295 path = os.path.abspath(os.path.join(root, filename))
296 file_paths.append(os.path.splitext(path)[0])
297 return file_paths