Coverage for / home / jenkins / .local / lib / python3.10 / site-packages / hyper_parallel / data / batching / build_dataloader.py: 20%
170 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 DataLoaders and calculate distributed micro-batch sizing."""
17from __future__ import annotations
19import copy
20from abc import ABC, abstractmethod
21from collections.abc import Callable, Iterator, Mapping, Sequence
22from typing import Any
24import torch
25from torch.utils.data import IterableDataset
26from torchdata.stateful_dataloader import StatefulDataLoader
28from hyper_parallel.data.batching.dynamic_batch import (
29 FullSampleCheckpoint,
30 IndexReplayCheckpoint,
31 MappingReplayDataset,
32 _BatchPipeline,
33 _OmniPackingPipeline,
34 _TokenBatchPipeline,
35)
36from hyper_parallel.data.dataset_logging import get_dataset_logger
37from hyper_parallel.data.parallel import build_dataset_batch_sampler
39logger = get_dataset_logger(__name__)
42def calculate_num_micro_batches(
43 global_batch_size: int,
44 micro_batch_size: int,
45 dp_world_size: int,
46) -> int:
47 """Calculate the number of micro-batches in one optimizer step.
49 Args:
50 global_batch_size: Number of samples processed by one optimizer step.
51 micro_batch_size: Number of samples processed by each DP rank and forward pass.
52 dp_world_size: Number of data-parallel ranks.
54 Returns:
55 Number of forward/backward micro-batches executed by each DP rank.
57 Raises:
58 ValueError: If the global batch cannot be divided exactly.
59 """
60 distributed_micro_batch_size = micro_batch_size * dp_world_size
61 if global_batch_size % distributed_micro_batch_size != 0:
62 raise ValueError(
63 f"global_batch_size ({global_batch_size}) must be divisible by "
64 f"micro_batch_size * dp_world_size ({distributed_micro_batch_size})"
65 )
67 num_micro_batches = global_batch_size // distributed_micro_batch_size
68 logger.debug(
69 "Resolved micro-batches=%d from global_batch_size=%d, micro_batch_size=%d, dp_world_size=%d",
70 num_micro_batches,
71 global_batch_size,
72 micro_batch_size,
73 dp_world_size,
74 )
75 return num_micro_batches
78def _is_iterable_dataset(dataset: Any) -> bool:
79 """Return whether a Dataset streams samples without mapping-style access."""
80 if isinstance(dataset, IterableDataset):
81 is_iterable = True
82 else:
83 has_iterator = callable(getattr(dataset, "__iter__", None))
84 has_index_access = callable(getattr(dataset, "__getitem__", None))
85 is_iterable = has_iterator and not has_index_access
86 return is_iterable
89def build_dataloader(
90 dataloader_target: Any,
91 *,
92 datasets: Sequence[Any | None],
93 collate_fn: Any,
94 mesh_context: Any,
95 training_config: Any,
96 max_seq_len: int | None = None,
97 default_seed: int = 1234,
98) -> tuple[tuple[Any | None, ...], tuple[Any | None, ...]]:
99 """Build train, validation, and test DataLoaders.
101 Mapping Datasets receive a distributed batch sampler. Iterable Datasets
102 control their own order. The configured target decides whether selected
103 samples use fixed collation or dynamic token batching.
105 Args:
106 dataloader_target: DataLoader build target.
107 datasets: Train, validation, and test datasets.
108 collate_fn: Collator applied after fixed or dynamic sample selection.
109 mesh_context: Data-parallel mesh context.
110 training_config: Batch size and random seed configuration.
111 max_seq_len: Maximum sample length used to derive dynamic token budget.
112 default_seed: Seed used when no training seed is configured.
114 Returns:
115 DataLoader and batch-sampler tuples for the three Dataset splits.
116 """
117 if len(datasets) != 3:
118 raise ValueError("datasets must contain train, validation, and test entries")
120 if all(dataset is None for dataset in datasets):
121 dataloader_splits = (None, None, None)
122 batch_sampler_splits = (None, None, None)
123 return dataloader_splits, batch_sampler_splits
125 if dataloader_target is None:
126 raise ValueError("dataloader_target must define a build target")
128 dataloaders: list[Any | None] = [None] * len(datasets)
129 batch_samplers: list[Any | None] = [None] * len(datasets)
131 for split_index, (split_name, dataset) in enumerate(zip(("train", "valid", "test"), datasets)):
132 if dataset is None:
133 logger.debug("Skipping empty Dataset split=%s", split_name)
134 continue
136 batch_sampler = None
137 if not _is_iterable_dataset(dataset):
138 batch_sampler = build_dataset_batch_sampler(
139 total_samples=len(dataset),
140 micro_batch_size=training_config.micro_batch_size,
141 global_batch_size=training_config.global_batch_size,
142 dp_world_size=mesh_context.dp_size,
143 dp_rank=mesh_context.dp_rank,
144 drop_last=getattr(dataloader_target, "drop_last", True),
145 data_rearrange_map=getattr(dataloader_target, "data_rearrange_map", None),
146 sampler_type=getattr(dataloader_target, "sampler_type", "single"),
147 data_sharding=getattr(dataloader_target, "data_sharding", False),
148 seed=training_config.seed if training_config.seed is not None else default_seed,
149 )
151 dataloaders[split_index] = dataloader_target.build(
152 dataset=dataset,
153 collate_fn=collate_fn,
154 batch_sampler=batch_sampler,
155 batch_size=training_config.micro_batch_size,
156 sampler_type=getattr(dataloader_target, "sampler_type", "single"),
157 dp_world_size=mesh_context.dp_size,
158 max_seq_len=max_seq_len,
159 seed=training_config.seed if training_config.seed is not None else default_seed,
160 )
161 logger.debug(
162 "Built DataLoader split=%s, dataset=%s, batch_sampler=%s",
163 split_name,
164 type(dataset).__name__,
165 type(batch_sampler).__name__ if batch_sampler is not None else None,
166 )
167 batch_samplers[split_index] = batch_sampler
169 logger.debug("Finished building train/valid/test DataLoaders")
170 return tuple(dataloaders), tuple(batch_samplers)
173class FixedBatchDataLoader(StatefulDataLoader):
174 """Build fixed-sample batches while retaining Trainer iterator policy."""
176 def __init__(
177 self,
178 dataset: Any,
179 batch_sampler: Any = None,
180 collate_fn: Callable[[list[Any]], Any] | None = None,
181 *,
182 batch_size: int | None = None,
183 drop_last: bool = True,
184 sampler_type: str = "single",
185 num_workers: int = 0,
186 seed: int = 1234,
187 pin_memory: bool = False,
188 prefetch_factor: int | None = None,
189 ) -> None:
190 """Initialize the stateful DataLoader."""
191 if sampler_type not in ("single", "cyclic"):
192 raise ValueError("sampler_type must be 'single' or 'cyclic'")
194 self.drop_last = drop_last
195 self.sampler_type = sampler_type
196 generator = torch.Generator().manual_seed(seed)
197 worker_options = {
198 "num_workers": num_workers,
199 "generator": generator,
200 "pin_memory": pin_memory,
201 }
202 if num_workers > 0 and prefetch_factor is not None:
203 worker_options["prefetch_factor"] = prefetch_factor
204 if batch_sampler is None:
205 resolved_batch_size = 1 if batch_size is None else batch_size
206 super().__init__(
207 dataset=dataset,
208 batch_size=resolved_batch_size,
209 collate_fn=collate_fn,
210 drop_last=drop_last,
211 **worker_options,
212 )
213 else:
214 super().__init__(
215 dataset=dataset,
216 batch_sampler=batch_sampler,
217 collate_fn=collate_fn,
218 **worker_options,
219 )
221 def set_epoch(self, epoch: int) -> None:
222 """Forward epoch state to the Dataset and configured batch sampler."""
223 if self.batch_sampler is not None and hasattr(self.batch_sampler, "set_epoch"):
224 self.batch_sampler.set_epoch(epoch)
226 epoch_setter = getattr(self.dataset, "set_epoch", None)
227 if callable(epoch_setter):
228 dataset_epoch = epoch if self.sampler_type == "cyclic" else 0
229 epoch_setter(dataset_epoch)
232class _DynamicBatchLoader(ABC):
233 """Own common dynamic source iteration, buffering, and checkpoint state.
235 Subclasses define how candidates are selected and finalized before the
236 configured collator builds a batch.
238 Args:
239 dataset: Online mapping or iterable Dataset producing ModelSamples.
240 batch_sampler: Optional mapping-style batch sampler.
241 collate_fn: Final batch collator receiving the selected ModelSamples.
242 batch_size: Configured micro batch size used to derive the token budget.
243 dp_world_size: Number of data-parallel source lanes.
244 save_by_idx: Whether to checkpoint buffer entries as output indices.
245 Mapping Datasets default to true and iterable Datasets default to false.
246 max_seq_len: Maximum sample length used to derive the token budget.
247 token_budget: Optional candidate-selection budget. Values smaller than
248 one sample still emit that sample alone instead of truncating it.
249 min_buffered_samples: Minimum candidate samples buffered before batching.
250 drop_last: Compatibility option retained from fixed DataLoader configuration.
251 Dynamic buffers are always drained when the source is exhausted.
252 sampler_type: Whether each Dataset epoch reuses one order or advances
253 the deterministic cyclic order.
254 num_workers: Number of source DataLoader workers.
255 seed: Source DataLoader random seed.
256 pin_memory: Whether source samples use pinned host memory.
257 prefetch_factor: Number of source batches prefetched by each worker.
258 """
260 def __init__(
261 self,
262 dataset: Any,
263 batch_sampler: Any = None,
264 collate_fn: Callable[[Sequence[Mapping[str, Any]]], Mapping[str, Any]] | None = None,
265 *,
266 batch_size: int,
267 dp_world_size: int | None = None,
268 save_by_idx: bool | None = None,
269 max_seq_len: int | None,
270 token_budget: int | None = None,
271 min_buffered_samples: int = 200,
272 drop_last: bool = True,
273 sampler_type: str = "single",
274 num_workers: int = 0,
275 seed: int = 1234,
276 pin_memory: bool = False,
277 prefetch_factor: int | None = None,
278 ) -> None:
279 """Initialize source reading, token selection, and Online packing."""
280 # pylint: disable=too-many-locals
281 if collate_fn is None:
282 raise ValueError(f"{type(self).__name__} requires a collate_fn")
284 if max_seq_len is None:
285 raise ValueError(f"{type(self).__name__} requires data_transform.max_seq_len")
286 if token_budget is not None and token_budget <= 0:
287 raise ValueError("token_budget must be positive")
289 resolved_dp_world_size = dp_world_size
290 if resolved_dp_world_size <= 0:
291 raise ValueError("dp_world_size must be positive")
293 sampler_dp_world_size = getattr(batch_sampler, "dp_world_size", resolved_dp_world_size)
294 if sampler_dp_world_size != resolved_dp_world_size:
295 raise ValueError("batch_sampler.dp_world_size must match dp_world_size")
297 self.drop_last = drop_last
298 self.batch_collate_fn = collate_fn
299 self.dp_world_size = resolved_dp_world_size
301 # build_candidate_buffer
302 self._batch_pipeline = self._build_batch_pipeline(dataset)
303 resolved_token_budget = batch_size * max_seq_len if token_budget is None else token_budget
304 self.candidate_buffer = self._batch_pipeline.build_candidate_buffer(
305 token_budget=resolved_token_budget,
306 min_buffered_samples=min_buffered_samples,
307 )
308 if batch_sampler is not None:
309 enable_source_resume = getattr(batch_sampler, "enable_source_batch_resume", None)
310 if callable(enable_source_resume):
311 enable_source_resume()
313 # Resuming training from a checkpoint
314 self._buffer_checkpoint = self._build_buffer_checkpoint(dataset, save_by_idx)
316 self.source_dataloader = FixedBatchDataLoader(
317 dataset=self._buffer_checkpoint.source_dataset,
318 batch_sampler=batch_sampler,
319 # Keep source samples separate until the candidate buffer selects them.
320 collate_fn=list,
321 batch_size=1 if batch_sampler is None else batch_size,
322 drop_last=False,
323 sampler_type=sampler_type,
324 num_workers=num_workers,
325 seed=seed,
326 pin_memory=pin_memory,
327 prefetch_factor=prefetch_factor,
328 )
329 self.resume_pending = False
331 @staticmethod
332 def _build_buffer_checkpoint(
333 dataset: Any,
334 save_by_idx: bool | None,
335 ) -> IndexReplayCheckpoint | FullSampleCheckpoint:
336 """Build candidate-buffer checkpoint storage for the Dataset access mode."""
337 is_iterable = _is_iterable_dataset(dataset)
338 resolved_save_by_idx = not is_iterable if save_by_idx is None else save_by_idx
340 # Mapping Datasets have stable indices. Iterable Datasets must explicitly
341 # provide replay access before their buffer can be checkpointed by index.
342 replay_dataset = dataset if is_iterable else MappingReplayDataset(dataset)
343 supports_replay = not is_iterable or IndexReplayCheckpoint.supports(dataset)
344 if resolved_save_by_idx and not supports_replay:
345 raise ValueError("save_by_idx=True requires get_item() and output_index_for_resume")
347 if is_iterable and hasattr(dataset, "output_index_for_resume"):
348 dataset.output_index_for_resume = resolved_save_by_idx
350 if resolved_save_by_idx:
351 buffer_checkpoint = IndexReplayCheckpoint(replay_dataset)
352 else:
353 # Keep replay access only for loading an earlier index-buffer checkpoint.
354 checkpoint_replay_dataset = replay_dataset if supports_replay else None
355 buffer_checkpoint = FullSampleCheckpoint(dataset, checkpoint_replay_dataset)
357 return buffer_checkpoint
359 def __iter__(self) -> Iterator[Mapping[str, Any]]:
360 """Start source iteration while retaining restored buffer state."""
361 if not self.resume_pending:
362 self.candidate_buffer.buffer.clear()
363 self.candidate_buffer.buffer_output_indices.clear()
364 self.candidate_buffer.buffer_token_count = 0
366 source_iterator = iter(self.source_dataloader)
367 batch_iterator = self._batch_data_generator(source_iterator)
368 self.resume_pending = False
370 return batch_iterator
372 def _batch_data_generator(
373 self,
374 source_iterator: Iterator[list[Any]],
375 ) -> Iterator[Mapping[str, Any]]:
376 """Lazily yield micro-batches as source samples fill the candidate buffer."""
377 # Stage 1: Use ready checkpointed candidates before advancing the source.
378 while self.candidate_buffer.is_ready_for_micro_batch():
379 yield self._build_micro_batch()
381 # Stage 2: Retain unselected candidates while incrementally reading source batches.
382 for source_samples in source_iterator:
383 for source_sample in source_samples:
384 self._buffer_checkpoint.put_source_item(source_sample, self.candidate_buffer)
386 while self.candidate_buffer.is_ready_for_micro_batch():
387 yield self._build_micro_batch()
389 # Stage 3: At EOF, bypass readiness thresholds so no pulled candidate is lost.
390 while not self.candidate_buffer.empty():
391 yield self._build_micro_batch()
393 @abstractmethod
394 def _build_batch_pipeline(self, dataset: Any) -> _BatchPipeline:
395 """Build the batching pipeline selected by the concrete Loader."""
396 raise NotImplementedError
398 def _build_micro_batch(self) -> Mapping[str, Any]:
399 """Select, finalize, collate, and encode one micro-batch."""
400 selected_samples = self.candidate_buffer.get_micro_batch()
401 finalized_samples = self._batch_pipeline.finalize_selected_samples(selected_samples)
403 # build batch data
404 batch = self.batch_collate_fn(finalized_samples)
405 encoded_batch = self._batch_pipeline.encode_batch(batch)
406 return encoded_batch
408 def state_dict(self) -> dict[str, Any]:
409 """Capture the future source cursor and unconsumed dynamic buffer."""
410 state = {
411 "dp_world_size": self.dp_world_size,
412 "source_dataloader": self.source_dataloader.state_dict(),
413 "save_by_idx": self._buffer_checkpoint.save_by_idx,
414 "buffer": self._buffer_checkpoint.get_buffer_state(self.candidate_buffer),
415 "buffer_token_count": self.candidate_buffer.buffer_token_count,
416 }
417 checkpoint_state = copy.deepcopy(state)
419 return checkpoint_state
421 def load_state_dict(self, state_dict: Mapping[str, Any]) -> None:
422 """Restore the future source cursor and unconsumed dynamic buffer."""
423 checkpoint_state = copy.deepcopy(dict(state_dict))
424 saved_dp_world_size = checkpoint_state.get("dp_world_size", self.dp_world_size)
425 if saved_dp_world_size != self.dp_world_size:
426 raise ValueError(
427 "Online dataloader resume does not support DP world-size changes: "
428 f"saved_dp_world_size={saved_dp_world_size}, current_dp_world_size={self.dp_world_size}"
429 )
431 previous_save_by_idx = bool(checkpoint_state.get("save_by_idx", False))
432 saved_buffer = checkpoint_state["buffer"]
433 restored_buffer, restored_indices = self._buffer_checkpoint.restore_buffer(
434 saved_buffer,
435 previous_save_by_idx,
436 self.candidate_buffer,
437 )
439 restored_token_count = sum(sample_length for _, sample_length in restored_buffer)
440 if restored_token_count != checkpoint_state["buffer_token_count"]:
441 raise ValueError("buffer_token_count does not match the restored dynamic buffer")
443 self.source_dataloader.load_state_dict(checkpoint_state["source_dataloader"])
444 self.candidate_buffer.buffer = restored_buffer
445 self.candidate_buffer.buffer_output_indices = restored_indices
446 self.candidate_buffer.buffer_token_count = restored_token_count
447 self.resume_pending = True
449 def set_epoch(self, epoch: int) -> None:
450 """Forward epoch state to the stateful source DataLoader."""
451 self.source_dataloader.set_epoch(epoch)
454class TokenBatchLoader(_DynamicBatchLoader):
455 """Select fully encoded samples within one token budget."""
457 def _build_batch_pipeline(self, _dataset: Any) -> _BatchPipeline:
458 """Build token-budget batching for fully encoded Dataset samples."""
459 batch_pipeline = _TokenBatchPipeline()
460 return batch_pipeline
463class OmniPackingLoader(_DynamicBatchLoader):
464 """Pack selected samples after applying the Dataset encoding lifecycle."""
466 def _build_batch_pipeline(self, dataset: Any) -> _BatchPipeline:
467 """Build packing around the normalized Dataset encoding lifecycle."""
468 batch_pipeline = _OmniPackingPipeline.from_dataset(dataset)
469 return batch_pipeline