Coverage for / home / jenkins / .local / lib / python3.10 / site-packages / hyper_parallel / data / batching / build_dataloader.py: 62%
291 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 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 dataclasses import replace
23from typing import Any
25import torch # pylint: disable=forbidden-backend-import
26from torch.utils.data import IterableDataset, default_collate # pylint: disable=forbidden-backend-import
27from torchdata.stateful_dataloader import StatefulDataLoader
29from hyper_parallel.data.batching.dynamic_batch import (
30 FullSampleCheckpoint,
31 IndexReplayCheckpoint,
32 MappingReplayDataset,
33 _BatchPipeline,
34 _OmniPackingPipeline,
35 _TokenBatchPipeline,
36)
37from hyper_parallel.data.batching.build_collate_fn import TextPackingCollator
38from hyper_parallel.data.dataset_logging import get_dataset_logger
39from hyper_parallel.data.parallel import build_dataset_batch_sampler
40from hyper_parallel.data.parallel.batch_sampler import (
41 get_sampler_checkpoint_state, has_torchdata_state, sampler_optimizer_step_progress, sampler_state_for_load,
42)
44from hyper_parallel.distributed_data import (
45 BalancingAlgorithm,
46 CostModel,
47 DistributedDatasetConfig,
48 SampleMetadata,
49 TensorFieldSpec,
50 build_distributed_dataloader,
51)
53logger = get_dataset_logger(__name__)
56def calculate_num_micro_batches(
57 global_batch_size: int,
58 micro_batch_size: int,
59 dp_world_size: int,
60) -> int:
61 """Calculate the number of micro-batches in one optimizer step.
63 Args:
64 global_batch_size: Number of samples processed by one optimizer step.
65 micro_batch_size: Number of samples processed by each DP rank and forward pass.
66 dp_world_size: Number of data-parallel ranks.
68 Returns:
69 Number of forward/backward micro-batches executed by each DP rank.
71 Raises:
72 ValueError: If the global batch cannot be divided exactly.
73 """
74 distributed_micro_batch_size = micro_batch_size * dp_world_size
75 if global_batch_size % distributed_micro_batch_size != 0:
76 raise ValueError(
77 f"global_batch_size ({global_batch_size}) must be divisible by "
78 f"micro_batch_size * dp_world_size ({distributed_micro_batch_size})"
79 )
81 num_micro_batches = global_batch_size // distributed_micro_batch_size
82 logger.debug(
83 "Resolved micro-batches=%d from global_batch_size=%d, micro_batch_size=%d, dp_world_size=%d",
84 num_micro_batches,
85 global_batch_size,
86 micro_batch_size,
87 dp_world_size,
88 )
89 return num_micro_batches
92def _is_iterable_dataset(dataset: Any) -> bool:
93 """Return whether a Dataset streams samples without mapping-style access."""
94 if isinstance(dataset, IterableDataset):
95 is_iterable = True
96 else:
97 has_iterator = callable(getattr(dataset, "__iter__", None))
98 has_index_access = callable(getattr(dataset, "__getitem__", None))
99 is_iterable = has_iterator and not has_index_access
100 return is_iterable
103def _resolve_load_balance(
104 dataloader_target: Any, data_config: Mapping[str, Any],
105) -> Mapping[str, Any] | None:
106 """Validate the balancing block and reject obsolete Dataset-side settings."""
107 removed_options = {
108 "load_balance", "distributed_dataloader", "balance_window_size", "device_prefetch_depth",
109 "log_balance", "min_balance_gain", "repack_sequences",
110 }.intersection(data_config)
111 if removed_options:
112 raise ValueError(f"Move data_config balancing options {sorted(removed_options)} to dataloader.load_balance")
113 load_balance = getattr(dataloader_target, "load_balance", None)
114 if load_balance is not None and not isinstance(load_balance, Mapping):
115 raise ValueError("dataloader.load_balance must be a mapping; omit it to disable balancing")
117 return load_balance
120def _build_distributed_packing_config(
121 dataloader_target: Any,
122 data_config: Mapping[str, Any],
123 *,
124 micro_batch_size: int,
125 seed: int,
126 max_seq_len: int | None = None,
127) -> DistributedDatasetConfig:
128 """Derive distributed-packing sizes and retain optional runtime tuning."""
129 try:
130 seq_len = int(data_config.get("seq_length", max_seq_len))
131 except (KeyError, TypeError, ValueError) as exc:
132 raise ValueError("Distributed loading requires data_config.seq_length or transform.max_seq_len") from exc
133 if max_seq_len is not None and seq_len != max_seq_len:
134 raise ValueError("data_config.seq_length must match transform.max_seq_len for native batch balancing")
135 options = dict(dataloader_target.load_balance)
136 trainer_managed_options = {
137 "seq_len",
138 "local_batch_size",
139 "dp_dim_names",
140 "drop_last",
141 "seed",
142 "num_workers",
143 "pin_memory",
144 "prefetch_factor",
145 "persistent_workers",
146 "dataset_already_sharded",
147 }
148 conflicting_options = sorted(set(options) & trainer_managed_options)
149 if conflicting_options:
150 raise ValueError(
151 "dataloader.load_balance cannot override Trainer-managed options "
152 f"{conflicting_options}; configure DataLoader execution options on dataloader "
153 "and sequence length on dataset.data_config"
154 )
155 return DistributedDatasetConfig(
156 seq_len=seq_len,
157 local_batch_size=micro_batch_size,
158 dp_dim_names=("dp",),
159 drop_last=bool(getattr(dataloader_target, "drop_last", True)),
160 seed=seed,
161 dataset_already_sharded=False,
162 **options,
163 )
166def _native_text_metadata(sample: Mapping[str, Any]) -> SampleMetadata:
167 """Describe a complete GPT output without inferring new attention boundaries."""
168 seq_len = int(sample["input_ids"].shape[-1])
169 return SampleMetadata(pack_tokens=seq_len, features={"P": seq_len, "D": 0})
172def _resolve_policy(dataloader_target: Any, name: str, supplied: Any, model_config: Any) -> Any:
173 """Prefer an explicit policy, otherwise instantiate its configured target."""
174 if supplied is not None:
175 return supplied
176 configured = getattr(dataloader_target, name, None)
177 if configured is None:
178 return None
179 return configured.build(model_config=model_config)
182def _native_dataloader_kwargs(dataloader_target: Any) -> dict[str, Any]:
183 """Forward configured execution options without changing distributed packing."""
184 return {
185 name: getattr(dataloader_target, name)
186 for name in (
187 "num_workers", "pin_memory", "prefetch_factor", "persistent_workers",
188 "timeout", "worker_init_fn", "multiprocessing_context", "pin_memory_device", "in_order",
189 )
190 if getattr(dataloader_target, name, None) is not None
191 }
194def _trainer_model_batch_spec(target: Any, collate_fn: Any, data_config: Mapping[str, Any]) -> Any:
195 """Resolve declared model fields without reading samples on model peers."""
196 configured = getattr(target, "model_batch_spec", None)
197 if configured is not None:
198 return configured.build() if callable(getattr(configured, "build", None)) else configured
199 declared = getattr(collate_fn, "model_batch_spec", None)
200 if declared is not None:
201 return declared
202 if isinstance(collate_fn, TextPackingCollator):
203 return {
204 "input_ids": TensorFieldSpec((1, None), torch.int64),
205 "labels": TensorFieldSpec((1, None), torch.int64),
206 "cu_seq_lens": TensorFieldSpec((None,), torch.int32),
207 }
208 # Instruction and multimodal providers may emit additional model-owned fields.
209 # Their layouts must be configured explicitly rather than guessed from a sample.
210 if collate_fn is default_collate and not data_config.get("is_instruction_dataset", False):
211 spec = {name: TensorFieldSpec((None, None), torch.int64) for name in ("input_ids", "labels")}
212 if data_config.get("create_ltor_fields_in_dataloader", False):
213 spec.update({
214 "loss_mask": TensorFieldSpec((None, None), torch.float32),
215 "position_ids": TensorFieldSpec((None, None), torch.int64),
216 "attention_mask": TensorFieldSpec((None, 1, None, None), torch.bool),
217 })
218 return spec
219 return None
222def _build_native_sampler_loader(
223 dataset: Any,
224 batch_sampler: Any,
225 collate_fn: Any,
226 dataloader_target: Any,
227 mesh_context: Any,
228 data_config: Mapping[str, Any],
229 seed: int,
230 *,
231 metadata_fn: Callable[[Any], SampleMetadata] | None = None,
232 max_seq_len: int | None = None,
233 model_config: Any = None,
234 cost_model: CostModel | None = None,
235 balancing_algorithm: BalancingAlgorithm | None = None,
236) -> Any:
237 """Reuse native Dataset outputs and the existing collator in collective loading."""
238 if batch_sampler is None:
239 raise ValueError("Native BatchSampler balancing requires a mapping Dataset")
240 if getattr(mesh_context, "device_mesh", None) is None:
241 raise ValueError("Native BatchSampler balancing requires mesh_context.device_mesh")
242 if int(getattr(mesh_context, "pp_size", 1)) != 1:
243 raise ValueError("Native BatchSampler balancing does not yet support pipeline parallelism")
244 # Target exposes its resolved callable only as _target_; inspecting it
245 # prevents replacing a second-stage dynamic selector without warning.
246 target_callable = getattr(dataloader_target, "_target_", None)
247 if isinstance(target_callable, type) and issubclass(target_callable, _DynamicBatchLoader):
248 raise ValueError("Native BatchSampler balancing must not override native batches with dynamic batching")
249 config = _build_distributed_packing_config(
250 dataloader_target, data_config,
251 micro_batch_size=batch_sampler.micro_batch_size, seed=seed,
252 max_seq_len=max_seq_len,
253 )
254 device_type = getattr(mesh_context.device_mesh, "device_type", "cpu")
255 if device_type != "npu" and "communication_backend" not in dataloader_target.load_balance:
256 config = replace(config, communication_backend="gloo")
257 communication_device = torch.device("cpu")
258 if device_type != "cpu":
259 communication_device = torch.device(device_type, getattr(torch, device_type).current_device())
260 effective_metadata_fn = metadata_fn
261 if not config.metadata_mode and effective_metadata_fn is None:
262 effective_metadata_fn = _native_text_metadata
263 effective_cost_model = _resolve_policy(dataloader_target, "cost_model", cost_model, model_config)
264 return build_distributed_dataloader(
265 dataset, mesh_context.device_mesh, config,
266 batch_sampler=batch_sampler, metadata_fn=effective_metadata_fn,
267 metadata=getattr(dataset, "metadata", None) if config.metadata_mode else None,
268 collate_fn=None if config.repack_sequences else collate_fn,
269 model_batch_spec=_trainer_model_batch_spec(
270 dataloader_target, None if config.repack_sequences else collate_fn, data_config,
271 ),
272 dataloader_kwargs=_native_dataloader_kwargs(dataloader_target),
273 device=communication_device,
274 model_config=model_config,
275 cost_model=effective_cost_model,
276 balancing_algorithm=_resolve_policy(
277 dataloader_target, "balancing_algorithm", balancing_algorithm, model_config,
278 ),
279 )
282def _build_split_batch_sampler(
283 dataset: Any, dataloader_target: Any, training_config: Any, mesh_context: Any, seed: int,
284) -> Any:
285 """Build rank-local sampling for mapping datasets; iterable sources own their order."""
286 if _is_iterable_dataset(dataset):
287 return None
288 return build_dataset_batch_sampler(
289 total_samples=len(dataset),
290 micro_batch_size=training_config.micro_batch_size,
291 global_batch_size=training_config.global_batch_size,
292 dp_world_size=mesh_context.dp_size,
293 dp_rank=mesh_context.dp_rank,
294 drop_last=getattr(dataloader_target, "drop_last", True),
295 data_rearrange_map=getattr(dataloader_target, "data_rearrange_map", None),
296 sampler_type=getattr(dataloader_target, "sampler_type", "single"),
297 data_sharding=getattr(dataloader_target, "data_sharding", False),
298 seed=seed,
299 )
302def build_dataloader(
303 dataloader_target: Any,
304 *,
305 datasets: Sequence[Any | None],
306 collate_fn: Any,
307 mesh_context: Any,
308 training_config: Any,
309 data_config: Mapping[str, Any] | None = None,
310 max_seq_len: int | None = None,
311 default_seed: int = 1234,
312 metadata_fn: Callable[[Any], SampleMetadata] | None = None,
313 model_config: Any = None,
314 cost_model: CostModel | None = None,
315 balancing_algorithm: BalancingAlgorithm | None = None,
316) -> tuple[tuple[Any | None, ...], tuple[Any | None, ...]]:
317 """Build train, validation, and test DataLoaders.
319 Mapping Datasets receive a distributed batch sampler. Iterable Datasets
320 control their own order. The configured target decides whether selected
321 samples use fixed collation or dynamic token batching.
323 Args:
324 dataloader_target: DataLoader build target. A load_balance mapping enables
325 balancing, including an empty mapping; omission disables it. With balancing, its optional
326 model_batch_spec is a tensor layout or a Target that builds one.
327 Standard GPT and Online text layouts are provided automatically;
328 instruction and multimodal collators require an explicit declaration
329 on this target or the collator when model parallelism is enabled.
330 datasets: Train, validation, and test datasets.
331 collate_fn: Collator applied after fixed or dynamic sample selection.
332 mesh_context: Data-parallel mesh context.
333 training_config: Batch size and random seed configuration.
334 data_config: Dataset options, including the source sequence length.
335 max_seq_len: Maximum sample length used to derive dynamic token budget.
336 default_seed: Seed used when no training seed is configured.
337 metadata_fn: Native physical sample metadata for online
338 ``dataloader.load_balance``. Defaults to complete GPT output
339 lengths; VLMTrainer supplies image metadata. Offline mode uses the
340 ``dataset.metadata`` instead: an indexable SampleMetadata sequence
341 aligned one-to-one with that Dataset split's logical outputs.
342 model_config: Effective model dimensions for the default workload estimator.
343 Unsupported architectures require an explicit cost model.
344 cost_model: Optional workload callback overriding the configured cost-model target.
345 balancing_algorithm: Optional assignment policy overriding its configured target.
347 Returns:
348 DataLoader and batch-sampler tuples for the three Dataset splits.
349 """
350 if len(datasets) != 3:
351 raise ValueError("datasets must contain train, validation, and test entries")
353 if all(dataset is None for dataset in datasets):
354 return (None, None, None), (None, None, None)
356 if dataloader_target is None:
357 raise ValueError("dataloader_target must define a build target")
359 micro_batch_size = training_config.micro_batch_size
360 seed = training_config.seed if training_config.seed is not None else default_seed
362 dp_world_size = mesh_context.dp_size
364 sampler_type = getattr(dataloader_target, "sampler_type", "single")
365 data_config = data_config or {}
366 load_balance = _resolve_load_balance(dataloader_target, data_config)
368 dataloaders: list[Any | None] = [None] * len(datasets)
369 batch_samplers: list[Any | None] = [None] * len(datasets)
371 for split_index, (split_name, dataset) in enumerate(zip(("train", "valid", "test"), datasets)):
372 if dataset is None:
373 logger.debug("Skipping empty Dataset split=%s", split_name)
374 continue
376 batch_sampler = _build_split_batch_sampler(dataset, dataloader_target, training_config, mesh_context, seed)
378 if load_balance is not None:
379 dataloader = _build_native_sampler_loader(
380 dataset, batch_sampler, collate_fn, dataloader_target, mesh_context, data_config, seed,
381 metadata_fn=metadata_fn, max_seq_len=max_seq_len,
382 model_config=model_config, cost_model=cost_model, balancing_algorithm=balancing_algorithm,
383 )
384 else:
385 dataloader = dataloader_target.build(
386 dataset=dataset, collate_fn=collate_fn, batch_sampler=batch_sampler,
387 batch_size=micro_batch_size, sampler_type=sampler_type,
388 dp_world_size=dp_world_size, max_seq_len=max_seq_len, seed=seed,
389 )
390 logger.debug(
391 "Built DataLoader split=%s, dataset=%s, batch_sampler=%s",
392 split_name, type(dataset).__name__,
393 type(batch_sampler).__name__ if batch_sampler is not None else None,
394 )
395 dataloaders[split_index] = dataloader
396 # The collective loader owns sampler checkpoint/epoch state; exposing
397 # its speculative cursor to the Trainer would skip prefetched batches.
398 batch_samplers[split_index] = None if load_balance is not None else batch_sampler
400 logger.debug("Finished building train/valid/test DataLoaders")
401 return tuple(dataloaders), tuple(batch_samplers)
404class FixedBatchDataLoader(StatefulDataLoader):
405 """Build fixed-sample batches while retaining Trainer iterator policy."""
407 def __init__(
408 self,
409 dataset: Any,
410 batch_sampler: Any = None,
411 collate_fn: Callable[[list[Any]], Any] | None = None,
412 *,
413 batch_size: int | None = None,
414 drop_last: bool = True,
415 sampler_type: str = "single",
416 num_workers: int = 0,
417 seed: int = 1234,
418 pin_memory: bool = False,
419 prefetch_factor: int | None = None,
420 ) -> None:
421 """Initialize the stateful DataLoader."""
422 if sampler_type not in ("single", "cyclic"):
423 raise ValueError("sampler_type must be 'single' or 'cyclic'")
425 self.drop_last = drop_last
426 self.sampler_type = sampler_type
427 generator = torch.Generator().manual_seed(seed)
428 worker_options = {
429 "num_workers": num_workers,
430 "generator": generator,
431 "pin_memory": pin_memory,
432 }
433 if num_workers > 0 and prefetch_factor is not None:
434 worker_options["prefetch_factor"] = prefetch_factor
435 if batch_sampler is None:
436 resolved_batch_size = 1 if batch_size is None else batch_size
437 super().__init__(
438 dataset=dataset,
439 batch_size=resolved_batch_size,
440 collate_fn=collate_fn,
441 drop_last=drop_last,
442 **worker_options,
443 )
444 else:
445 super().__init__(
446 dataset=dataset,
447 batch_sampler=batch_sampler,
448 collate_fn=collate_fn,
449 **worker_options,
450 )
451 # Native worker snapshots may precede the delivered optimizer boundary.
452 enable_source_resume = getattr(batch_sampler, "enable_source_batch_resume", None)
453 if callable(enable_source_resume):
454 enable_source_resume()
456 def state_dict(self) -> dict[str, Any]:
457 """Save the native delivered cursor alongside torchdata's worker state."""
458 state = super().state_dict()
459 if "_snapshot" not in state:
460 return state
461 snapshot = state["_snapshot"]["_main_snapshot"]
462 sampler = snapshot.get("_index_sampler_state")
463 if isinstance(sampler, dict) and "consumed_samples" in sampler:
464 sampler = dict(sampler)
465 sampler["consumed_samples"] += (
466 state.get("_steps_since_snapshot", 0) * self.batch_sampler.global_micro_batch_size
467 )
468 state["_index_sampler_state"] = sampler
469 return state
471 def state_dict_for_load(self) -> dict[str, Any]:
472 """Request native worker state and shared or legacy sampler cursors."""
473 return sampler_state_for_load(self.state_dict())
475 def load_state_dict(self, state_dict: dict[str, Any]) -> None:
476 """Restore torchdata state, or a balanced loader's complete native cursor.
478 Args:
479 state_dict: State saved by the ordinary or balanced native loader.
480 """
481 sampler = get_sampler_checkpoint_state(state_dict)
482 snapshot = state_dict.get("_snapshot", {}).get("_main_snapshot", {})
483 native_sampler = sampler if sampler is not None else snapshot.get("_index_sampler_state")
484 repartition = isinstance(native_sampler, Mapping) and any(
485 native_sampler.get(name) != getattr(self.batch_sampler, name, native_sampler.get(name))
486 for name in ("global_batch_size", "dp_world_size")
487 )
488 if state_dict.get("global_rank") is None and has_torchdata_state(state_dict) and not repartition:
489 if isinstance(native_sampler, Mapping) and native_sampler.get("epoch") is not None:
490 self.set_epoch(native_sampler["epoch"])
491 super().load_state_dict(state_dict)
492 return
493 if sampler is None:
494 raise ValueError("Balanced checkpoint has no native sampler cursor on this rank")
495 if sampler["epoch"] < 0 or sampler_optimizer_step_progress(sampler):
496 raise ValueError("Switching loaders requires a complete optimizer-step cursor")
497 self.batch_sampler.load_state_dict(sampler)
498 self.set_epoch(sampler["epoch"])
499 super().load_state_dict({})
500 self.next_iter_state = None
502 def set_epoch(self, epoch: int) -> None:
503 """Forward epoch state to the Dataset and configured batch sampler.
505 Args:
506 epoch: Epoch index used by the sampler and Dataset.
507 """
508 if self.batch_sampler is not None and hasattr(self.batch_sampler, "set_epoch"):
509 self.batch_sampler.set_epoch(epoch)
511 epoch_setter = getattr(self.dataset, "set_epoch", None)
512 if callable(epoch_setter):
513 dataset_epoch = epoch if self.sampler_type == "cyclic" else 0
514 epoch_setter(dataset_epoch)
517class _DynamicBatchLoader(ABC):
518 """Own common dynamic source iteration, buffering, and checkpoint state.
520 Subclasses define how candidates are selected and finalized before the
521 configured collator builds a batch.
523 Args:
524 dataset: Online mapping or iterable Dataset producing ModelSamples.
525 batch_sampler: Optional mapping-style batch sampler.
526 collate_fn: Final batch collator receiving the selected ModelSamples.
527 batch_size: Configured micro batch size used to derive the token budget.
528 dp_world_size: Number of data-parallel source lanes.
529 save_by_idx: Whether to checkpoint buffer entries as output indices.
530 Mapping Datasets default to true and iterable Datasets default to false.
531 max_seq_len: Maximum sample length used to derive the token budget.
532 token_budget: Optional candidate-selection budget. Values smaller than
533 one sample still emit that sample alone instead of truncating it.
534 min_buffered_samples: Minimum candidate samples buffered before batching.
535 drop_last: Compatibility option retained from fixed DataLoader configuration.
536 Dynamic buffers are always drained when the source is exhausted.
537 sampler_type: Whether each Dataset epoch reuses one order or advances
538 the deterministic cyclic order.
539 num_workers: Number of source DataLoader workers.
540 seed: Source DataLoader random seed.
541 pin_memory: Whether source samples use pinned host memory.
542 prefetch_factor: Number of source batches prefetched by each worker.
543 """
545 def __init__(
546 self,
547 dataset: Any,
548 batch_sampler: Any = None,
549 collate_fn: Callable[[Sequence[Mapping[str, Any]]], Mapping[str, Any]] | None = None,
550 *,
551 batch_size: int,
552 dp_world_size: int | None = None,
553 save_by_idx: bool | None = None,
554 max_seq_len: int | None,
555 token_budget: int | None = None,
556 min_buffered_samples: int = 200,
557 drop_last: bool = True,
558 sampler_type: str = "single",
559 num_workers: int = 0,
560 seed: int = 1234,
561 pin_memory: bool = False,
562 prefetch_factor: int | None = None,
563 ) -> None:
564 """Initialize source reading, token selection, and Online packing."""
565 # pylint: disable=too-many-locals
566 if collate_fn is None:
567 raise ValueError(f"{type(self).__name__} requires a collate_fn")
569 if max_seq_len is None:
570 raise ValueError(f"{type(self).__name__} requires data_transform.max_seq_len")
571 if token_budget is not None and token_budget <= 0:
572 raise ValueError("token_budget must be positive")
574 resolved_dp_world_size = dp_world_size
575 if resolved_dp_world_size <= 0:
576 raise ValueError("dp_world_size must be positive")
578 sampler_dp_world_size = getattr(batch_sampler, "dp_world_size", resolved_dp_world_size)
579 if sampler_dp_world_size != resolved_dp_world_size:
580 raise ValueError("batch_sampler.dp_world_size must match dp_world_size")
582 self.drop_last = drop_last
583 self.batch_collate_fn = collate_fn
584 self.dp_world_size = resolved_dp_world_size
586 # build_candidate_buffer
587 self._batch_pipeline = self._build_batch_pipeline(dataset)
588 resolved_token_budget = batch_size * max_seq_len if token_budget is None else token_budget
589 self.candidate_buffer = self._batch_pipeline.build_candidate_buffer(
590 token_budget=resolved_token_budget,
591 min_buffered_samples=min_buffered_samples,
592 )
593 if batch_sampler is not None:
594 enable_source_resume = getattr(batch_sampler, "enable_source_batch_resume", None)
595 if callable(enable_source_resume):
596 enable_source_resume()
598 # Resuming training from a checkpoint
599 self._buffer_checkpoint = self._build_buffer_checkpoint(dataset, save_by_idx)
601 self.source_dataloader = FixedBatchDataLoader(
602 dataset=self._buffer_checkpoint.source_dataset,
603 batch_sampler=batch_sampler,
604 # Keep source samples separate until the candidate buffer selects them.
605 collate_fn=list,
606 batch_size=1 if batch_sampler is None else batch_size,
607 drop_last=False,
608 sampler_type=sampler_type,
609 num_workers=num_workers,
610 seed=seed,
611 pin_memory=pin_memory,
612 prefetch_factor=prefetch_factor,
613 )
614 self.resume_pending = False
616 @staticmethod
617 def _build_buffer_checkpoint(
618 dataset: Any,
619 save_by_idx: bool | None,
620 ) -> IndexReplayCheckpoint | FullSampleCheckpoint:
621 """Build candidate-buffer checkpoint storage for the Dataset access mode."""
622 is_iterable = _is_iterable_dataset(dataset)
623 resolved_save_by_idx = not is_iterable if save_by_idx is None else save_by_idx
625 # Mapping Datasets have stable indices. Iterable Datasets must explicitly
626 # provide replay access before their buffer can be checkpointed by index.
627 replay_dataset = dataset if is_iterable else MappingReplayDataset(dataset)
628 supports_replay = not is_iterable or IndexReplayCheckpoint.supports(dataset)
629 if resolved_save_by_idx and not supports_replay:
630 raise ValueError("save_by_idx=True requires get_item() and output_index_for_resume")
632 if is_iterable and hasattr(dataset, "output_index_for_resume"):
633 dataset.output_index_for_resume = resolved_save_by_idx
635 if resolved_save_by_idx:
636 buffer_checkpoint = IndexReplayCheckpoint(replay_dataset)
637 else:
638 # Keep replay access only for loading an earlier index-buffer checkpoint.
639 checkpoint_replay_dataset = replay_dataset if supports_replay else None
640 buffer_checkpoint = FullSampleCheckpoint(dataset, checkpoint_replay_dataset)
642 return buffer_checkpoint
644 def __iter__(self) -> Iterator[Mapping[str, Any]]:
645 """Start source iteration while retaining restored buffer state."""
646 if not self.resume_pending:
647 self.candidate_buffer.buffer.clear()
648 self.candidate_buffer.buffer_output_indices.clear()
649 self.candidate_buffer.buffer_token_count = 0
651 source_iterator = iter(self.source_dataloader)
652 batch_iterator = self._batch_data_generator(source_iterator)
653 self.resume_pending = False
655 return batch_iterator
657 def _batch_data_generator(
658 self,
659 source_iterator: Iterator[list[Any]],
660 ) -> Iterator[Mapping[str, Any]]:
661 """Lazily yield micro-batches as source samples fill the candidate buffer."""
662 # Stage 1: Use ready checkpointed candidates before advancing the source.
663 while self.candidate_buffer.is_ready_for_micro_batch():
664 yield self._build_micro_batch()
666 # Stage 2: Retain unselected candidates while incrementally reading source batches.
667 for source_samples in source_iterator:
668 for source_sample in source_samples:
669 self._buffer_checkpoint.put_source_item(source_sample, self.candidate_buffer)
671 while self.candidate_buffer.is_ready_for_micro_batch():
672 yield self._build_micro_batch()
674 # Stage 3: At EOF, bypass readiness thresholds so no pulled candidate is lost.
675 logger.debug(
676 "Dynamic loader source exhausted: buffered_samples=%d buffered_tokens=%d "
677 "token_budget=%d drop_last=%s; draining remaining candidates",
678 len(self.candidate_buffer.buffer), self.candidate_buffer.buffer_token_count,
679 self.candidate_buffer.token_budget, self.drop_last, enabled=True,
680 )
681 while not self.candidate_buffer.empty():
682 yield self._build_micro_batch()
683 logger.debug("Dynamic loader tail drained; no candidates remain", enabled=True)
685 @abstractmethod
686 def _build_batch_pipeline(self, dataset: Any) -> _BatchPipeline:
687 """Build the batching pipeline selected by the concrete Loader."""
688 raise NotImplementedError
690 def _build_micro_batch(self) -> Mapping[str, Any]:
691 """Select, finalize, collate, and encode one micro-batch."""
692 selected_samples = self.candidate_buffer.get_micro_batch()
693 logger.debug(
694 "Dynamic loader selected: samples=%d token_budget=%d remaining_samples=%d remaining_tokens=%d",
695 len(selected_samples), self.candidate_buffer.token_budget, len(self.candidate_buffer.buffer),
696 self.candidate_buffer.buffer_token_count, enabled=True,
697 )
698 finalized_samples = self._batch_pipeline.finalize_selected_samples(selected_samples)
700 # build batch data
701 batch = self.batch_collate_fn(finalized_samples)
702 encoded_batch = self._batch_pipeline.encode_batch(batch)
703 return encoded_batch
705 def state_dict(self) -> dict[str, Any]:
706 """Capture the future source cursor and unconsumed dynamic buffer."""
707 state = {
708 "dp_world_size": self.dp_world_size,
709 "source_dataloader": self.source_dataloader.state_dict(),
710 "save_by_idx": self._buffer_checkpoint.save_by_idx,
711 "buffer": self._buffer_checkpoint.get_buffer_state(self.candidate_buffer),
712 "buffer_token_count": self.candidate_buffer.buffer_token_count,
713 }
714 checkpoint_state = copy.deepcopy(state)
716 return checkpoint_state
718 def load_state_dict(self, state_dict: Mapping[str, Any]) -> None:
719 """Restore the future source cursor and unconsumed dynamic buffer.
721 Args:
722 state_dict: Saved source progress and unconsumed candidate buffer.
723 """
724 checkpoint_state = copy.deepcopy(dict(state_dict))
725 saved_dp_world_size = checkpoint_state.get("dp_world_size", self.dp_world_size)
726 if saved_dp_world_size != self.dp_world_size:
727 raise ValueError(
728 "Online dataloader resume does not support DP world-size changes: "
729 f"saved_dp_world_size={saved_dp_world_size}, current_dp_world_size={self.dp_world_size}"
730 )
732 previous_save_by_idx = bool(checkpoint_state.get("save_by_idx", False))
733 saved_buffer = checkpoint_state["buffer"]
734 restored_buffer, restored_indices = self._buffer_checkpoint.restore_buffer(
735 saved_buffer,
736 previous_save_by_idx,
737 self.candidate_buffer,
738 )
740 restored_token_count = sum(sample_length for _, sample_length in restored_buffer)
741 if restored_token_count != checkpoint_state["buffer_token_count"]:
742 raise ValueError("buffer_token_count does not match the restored dynamic buffer")
744 self.source_dataloader.load_state_dict(checkpoint_state["source_dataloader"])
745 self.candidate_buffer.buffer = restored_buffer
746 self.candidate_buffer.buffer_output_indices = restored_indices
747 self.candidate_buffer.buffer_token_count = restored_token_count
748 self.resume_pending = True
750 def set_epoch(self, epoch: int) -> None:
751 """Forward epoch state to the stateful source DataLoader.
753 Args:
754 epoch: Epoch index forwarded to the source DataLoader.
755 """
756 self.source_dataloader.set_epoch(epoch)
759class TokenBatchLoader(_DynamicBatchLoader):
760 """Select fully encoded samples within one token budget."""
762 def _build_batch_pipeline(self, dataset: Any) -> _BatchPipeline:
763 """Build token-budget batching for fully encoded Dataset samples."""
764 batch_pipeline = _TokenBatchPipeline()
765 return batch_pipeline
768class OmniPackingLoader(_DynamicBatchLoader):
769 """Pack selected samples after applying the Dataset encoding lifecycle."""
771 def _build_batch_pipeline(self, dataset: Any) -> _BatchPipeline:
772 """Build packing around the normalized Dataset encoding lifecycle."""
773 batch_pipeline = _OmniPackingPipeline.from_dataset(dataset)
774 return batch_pipeline
777# Compatibility name for callers of the pre-refactor dynamic text loader.
778DynamicBatchDataLoader = TokenBatchLoader