Coverage for / home / jenkins / .local / lib / python3.10 / site-packages / hyper_parallel / distributed_data / distributed_dataloader.py: 95%
570 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"""Dataset Reader, Planner, and Data Constructor orchestration."""
17from __future__ import annotations
19import copy
20import logging
21import math
22from collections import deque
23from collections.abc import Iterator, Mapping
24from concurrent.futures import Future
25from contextlib import nullcontext
26from dataclasses import dataclass, field
27from queue import Queue
28from threading import Thread
29from typing import Any, Union
31import torch # pylint: disable=forbidden-backend-import
33from hyper_parallel.data.parallel.batch_sampler import (
34 get_sampler_checkpoint_state, has_torchdata_state, sampler_optimizer_step_progress, sampler_state_for_load,
35)
36from hyper_parallel.distributed_data.balance_logging import format_balance_stats, logger as balance_logger
37from hyper_parallel.distributed_data.data_constructor import PackingDataConstructor
38from hyper_parallel.distributed_data.batch_sampler import BatchSamplerReader, _ReaderPlanInput
39from hyper_parallel.distributed_data.device_prefetch import DevicePrefetchedStep, DeviceStepPrefetcher, _record_stream
40from hyper_parallel.distributed_data.planner import DynamicPackingPlanner
41from hyper_parallel.distributed_data.schema import (
42 DistributedPackingPlan,
43 SampleKey,
44 SampleMetadata,
45)
46from hyper_parallel.distributed_data.metadata import PlannedSampleLoader
47from hyper_parallel.distributed_data.topology import DataTopology
48from hyper_parallel.distributed_data.transport import (
49 DataPlaneTransport,
50 ModelParallelTransport,
51 PreparedPayloadExchange,
52)
55_PlanControl = Union[DistributedPackingPlan, tuple[DistributedPackingPlan, ...], None]
58def _validate_loader_components(
59 *,
60 topology: DataTopology,
61 dataset_reader_ranks: tuple[int, ...],
62 dataset_reader: Any | None,
63 metadata_reader: BatchSamplerReader | None,
64 direct_sample_loader: PlannedSampleLoader | None,
65 metadata_mode: bool,
66) -> None:
67 """Check Reader ownership and require mode-compatible loading components."""
68 if not isinstance(metadata_mode, bool):
69 raise ValueError("metadata_mode must be boolean.")
70 if metadata_mode and dataset_reader is not None:
71 raise ValueError("Metadata mode must not configure an online Dataset Reader.")
72 if not metadata_mode and (metadata_reader is not None or direct_sample_loader is not None):
73 raise ValueError("Online mode must not configure metadata loading components.")
74 is_reader = topology.global_rank in dataset_reader_ranks
75 planning_reader = metadata_reader if metadata_mode else dataset_reader
76 if is_reader != (planning_reader is not None):
77 raise ValueError("Dataset Reader ownership does not match dataset_reader_ranks.")
78 if not metadata_mode and planning_reader is not None and not callable(
79 getattr(planning_reader, "prepare_next_step", None)
80 ):
81 raise ValueError("Online Dataset Reader must provide prepare_next_step().")
82 _validate_metadata_loader_owner(
83 topology,
84 metadata_mode=metadata_mode,
85 direct_sample_loader=direct_sample_loader,
86 )
89def _validate_metadata_loader_owner(
90 topology: DataTopology,
91 *,
92 metadata_mode: bool,
93 direct_sample_loader: PlannedSampleLoader | None,
94) -> None:
95 if not metadata_mode:
96 return
97 if topology.is_constructor != (direct_sample_loader is not None):
98 raise ValueError("Every metadata Data Constructor must own one plan-aware sample loader.")
101@dataclass
102class _PreparedWindow:
103 """Retain one Host balancing window and an optional complete-step checkpoint."""
105 batches: deque[Any]
106 exhausted: bool
107 checkpoint: dict[str, Any] | None
108 balance_logs: list[str] = field(default_factory=list)
111@dataclass
112class _PreparedBatch:
113 """Carry one batch and its delivery state through caller-side staging."""
115 batch: Any
116 checkpoint: dict[str, Any] | None = None
117 balance_logs: list[str] = field(default_factory=list)
118 ready_event: Any = None
119 error: BaseException | None = None
122class DistributedDataLoader(Iterator[Any]):
123 """Prepare balancing windows and deliver one local micro-batch at a time.
125 One Host producer owns reads, planning, data-plane collectives and pinning.
126 The caller prefetches H2D and broadcasts only the batch being delivered.
127 Host windows and device inputs have independent buffering limits.
128 Checkpoints advance at complete optimizer-step delivery boundaries.
129 """
131 def __init__(
132 self,
133 *,
134 topology: DataTopology,
135 dataset_reader_ranks: tuple[int, ...],
136 dataset_reader: Any | None,
137 metadata_reader: BatchSamplerReader | None,
138 direct_sample_loader: PlannedSampleLoader | None,
139 metadata_mode: bool,
140 planner: DynamicPackingPlanner,
141 data_constructor: PackingDataConstructor,
142 data_plane: DataPlaneTransport,
143 model_transport: ModelParallelTransport,
144 device_prefetch: DeviceStepPrefetcher | None = None,
145 initial_epoch: int = 0,
146 initial_consumed_samples: int = 0,
147 broadcast_plan: bool = True,
148 log_balance: bool = False,
149 accumulation_steps: int = 1,
150 balance_window_size: int = 1,
151 device_prefetch_depth: int = 1,
152 device: Any = "cpu",
153 ) -> None:
154 """Store the fully validated runtime components."""
155 _validate_loader_components(
156 topology=topology,
157 dataset_reader_ranks=dataset_reader_ranks,
158 dataset_reader=dataset_reader,
159 metadata_reader=metadata_reader,
160 direct_sample_loader=direct_sample_loader,
161 metadata_mode=metadata_mode,
162 )
163 self._topology = topology
164 self._dataset_reader_ranks = frozenset(dataset_reader_ranks)
165 self._dataset_reader = dataset_reader
166 self._metadata_reader = metadata_reader
167 self._direct_sample_loader = direct_sample_loader
168 self._metadata_mode = metadata_mode
169 self._broadcast_plan = broadcast_plan
170 self._log_balance = log_balance
171 self._planner = planner
172 self._data_constructor = data_constructor
173 self._data_plane = data_plane
174 self._model_transport = model_transport
175 self._device_prefetch = device_prefetch
176 self._epoch = initial_epoch
177 self._step = 0
178 self._initial_consumed_samples: int | None = initial_consumed_samples
179 self._stopped = False
180 self._pending_local_keys: set[SampleKey] = set()
181 self._window_plans: deque[DistributedPackingPlan] = deque()
182 for name, value in (("accumulation_steps", accumulation_steps), ("balance_window_size", balance_window_size)):
183 if isinstance(value, bool) or not isinstance(value, int) or value < 1:
184 raise ValueError(f"{name} must be a positive integer.")
185 self._accumulation_steps = accumulation_steps
186 self._balance_window_size = balance_window_size
187 if (isinstance(device_prefetch_depth, bool) or not isinstance(device_prefetch_depth, int)
188 or device_prefetch_depth < 0):
189 raise ValueError("device_prefetch_depth must be a non-negative integer.")
190 self._device_prefetch_depth = device_prefetch_depth
191 self._device = torch.device(device)
192 self._current: _PreparedWindow | None = None
193 self._future: Future[_PreparedWindow] | None = None
194 self._requests: Queue[Future[_PreparedWindow] | None] = Queue(maxsize=1)
195 self._prefetch_thread: Thread | None = None
196 self._device_batches: deque[_PreparedBatch] = deque()
197 self._prefetch_target = 0
198 self._consumer_stream: Any = None
199 # Each staged batch spends one input slot; advancing returns it with the
200 # previous consumer's completion event, even if that work is still running.
201 # Include the active batch so device lookahead can overlap its computation.
202 self._prefetch_slots: deque[Any] = deque([None] * (device_prefetch_depth + 1))
203 self._device_stopped = False
204 self._batch_error: BaseException | None = None
205 self._last_ready_event: Any = None
206 self._has_delivered = False
207 self._closed = False
208 self._balance_logs: list[str] = []
209 self._delivered_state = self._snapshot_state()
211 def __iter__(self) -> DistributedDataLoader:
212 """Return the rank-local iterator."""
213 return self
215 def __next__(self) -> Any:
216 """Deliver one batch and refill the configured device lookahead."""
217 if self._closed or self._delivered_state["stopped"]:
218 raise StopIteration
219 if self._batch_error is not None:
220 raise self._batch_error
221 self._record_consumption()
222 if self._device.type in ("npu", "cuda"):
223 prepared = self._next_device_batch()
224 else:
225 prepared = self._next_host_batch()
226 if prepared.batch is None:
227 self._delivered_state["stopped"] = True
228 raise StopIteration
229 batch = prepared.batch
230 if prepared.ready_event is not None:
231 module = getattr(torch, self._device.type)
232 self._consumer_stream = module.current_stream(self._device)
233 _record_stream(batch, self._consumer_stream)
234 if prepared.checkpoint is not None:
235 self._delivered_state = prepared.checkpoint
236 for message in prepared.balance_logs:
237 balance_logger.info("%s", message)
238 self._has_delivered = True
239 self.prefetch()
240 return batch
242 def _record_consumption(self) -> None:
243 # Advancing the iterator retires the previous batch. Callers must finish
244 # submitting its uses and release input references before advancing.
245 if self._consumer_stream is not None:
246 prefetch_wait_event = getattr(torch, self._device.type).Event()
247 prefetch_wait_event.record(self._consumer_stream)
248 self._prefetch_slots.append(prefetch_wait_event)
249 self._consumer_stream = None
251 def _next_host_batch(self) -> _PreparedBatch:
252 """Pop one Host batch and its delivery state, prefetching the following window."""
253 if self._current is None or not self._current.batches:
254 if self._current is not None and self._current.exhausted:
255 return _PreparedBatch(None)
256 self._prefetch_host_window()
257 self._current = self._future.result()
258 self._future = None
259 if not self._current.batches:
260 return _PreparedBatch(None)
261 self._prefetch_host_window()
262 batch = self._current.batches.popleft()
263 logs, self._current.balance_logs = self._current.balance_logs, []
264 checkpoint = self._current.checkpoint if not self._current.batches else None
265 return _PreparedBatch(batch, checkpoint, logs)
267 def prefetch(self) -> None:
268 """Prefetch Constructor H2D without any model-peer communication."""
269 if self._closed or self._delivered_state["stopped"] or self._batch_error is not None:
270 return
271 if (self._device.type not in ("npu", "cuda") or self._device_prefetch_depth == 0
272 or not self._topology.is_constructor):
273 self._prefetch_target = 0
274 self._prefetch_host_window()
275 return
276 self._prefetch_device_batches(self._device_prefetch_depth)
278 def _host_batch_ready(self) -> bool:
279 """Check Host readiness without waiting for speculative window work."""
280 if self._current is not None and (self._current.batches or self._current.exhausted):
281 return True
282 self._prefetch_host_window()
283 return self._future.done()
285 def _prefetch_device_batches(self, count: int, *, blocking: bool = False) -> None:
286 """Fill an active Constructor's H2D queue after caller-side lifecycle checks."""
287 self._prefetch_target = count
288 while len(self._device_batches) < count and self._prefetch_slots and not self._device_stopped:
289 prefetch_wait_event = self._prefetch_slots[0]
290 if not blocking and (
291 not self._host_batch_ready()
292 or (prefetch_wait_event is not None and not prefetch_wait_event.query())
293 ):
294 return
295 try:
296 prepared = self._stage_device_batch(prefetch_wait_event)
297 except BaseException as exc: # pylint: disable=broad-exception-caught
298 # Preserve FIFO delivery: a speculative failure must not prevent
299 # earlier batches from training or enter a peer collective early.
300 prepared = _PreparedBatch(None, error=exc)
301 self._device_stopped = True
302 self._prefetch_slots.popleft()
303 self._device_batches.append(prepared)
305 def _prefetch_host_window(self) -> None:
306 """Submit one Host window for an active loader, preserving double buffering."""
307 if self._future is not None:
308 return
309 if self._current is not None and self._current.exhausted:
310 return
311 self._future = Future()
312 if self._prefetch_thread is None:
313 self._prefetch_thread = Thread(
314 target=self._run_prefetch, name=f"hp-data-prefetch-rank-{self._topology.global_rank}", daemon=True,
315 )
316 self._prefetch_thread.start()
317 self._requests.put(self._future)
319 def _stage_device_batch(self, prefetch_wait_event: Any) -> _PreparedBatch:
320 """Keep copy events local until the corresponding batch is consumed."""
321 prepared = self._next_host_batch()
322 if prefetch_wait_event is not None:
323 prefetch_wait_event.synchronize()
324 prepared.batch = self._stage_batch(prepared.batch)
325 if isinstance(prepared.batch, DevicePrefetchedStep):
326 prepared.ready_event = prepared.batch.ready_event
327 self._device_stopped = prepared.batch is None
328 return prepared
330 def _next_device_batch(self) -> _PreparedBatch:
331 """Broadcast on the training stream after waiting for only the current copy."""
332 try:
333 if self._topology.is_constructor:
334 self._prefetch_device_batches(1, blocking=True)
335 prepared = self._device_batches.popleft()
336 else:
337 prepared = self._next_host_batch()
338 prefetch_wait_event = self._prefetch_slots.popleft()
339 if prefetch_wait_event is not None:
340 prefetch_wait_event.synchronize()
341 except BaseException as exc: # pylint: disable=broad-exception-caught
342 # All preparation failures must reach peers through the broadcast error header.
343 prepared = _PreparedBatch(None, error=exc)
344 module = getattr(torch, self._device.type)
345 consumer_stream = module.current_stream(self._device)
346 try:
347 batch = prepared.batch
348 if isinstance(batch, DevicePrefetchedStep):
349 batch = batch.take_batch()
350 prepared.batch = self._broadcast_batch(batch) if prepared.error is None else (
351 self._model_transport.broadcast(None, error=prepared.error)
352 )
353 except BaseException as exc:
354 self._batch_error = exc
355 self._device_batches.clear()
356 raise
357 finally:
358 # Explicit drains must cover partially submitted broadcasts on errors
359 # without synchronizing the caller's later training work.
360 prepared.ready_event = module.Event()
361 prepared.ready_event.record(consumer_stream)
362 self._last_ready_event = prepared.ready_event
363 return prepared
365 def _run_prefetch(self) -> None:
366 """Serve window requests on a persistent producer and propagate failures via futures."""
367 stream = None
368 try:
369 while True:
370 future = self._requests.get()
371 if future is None:
372 return
373 try:
374 if stream is None and self._device.type in ("npu", "cuda"):
375 module = getattr(torch, self._device.type)
376 module.set_device(self._device)
377 stream = module.Stream(device=self._device)
378 prefetch_stream_context = (
379 getattr(torch, self._device.type).stream(stream) if stream is not None else nullcontext()
380 )
381 with prefetch_stream_context:
382 window = self._prepare_window()
383 except BaseException as exc: # pylint: disable=broad-exception-caught
384 # A worker interrupt must complete its future instead of stranding result().
385 future.set_exception(exc)
386 else:
387 future.set_result(window)
388 del window
389 # An idle producer must not retain a discarded epoch's inputs.
390 del future
391 finally:
392 if stream is not None:
393 stream.synchronize()
395 def _prepare_window(self) -> _PreparedWindow:
396 """Prepare Host batches up to the window, optimizer-step or epoch boundary."""
397 window_micro_batch_num = min(
398 self._balance_window_size, self._accumulation_steps - self._step % self._accumulation_steps,
399 )
400 self._balance_logs = []
401 self._prepare_local_metadata(window_micro_batch_num)
402 batches = deque()
403 exhausted = False
404 for _ in range(window_micro_batch_num):
405 try:
406 batch = self._prepare_batch()
407 except StopIteration:
408 exhausted = True
409 break
410 batches.append(batch)
411 # Intermediate windows must replay from the optimizer step's beginning.
412 checkpoint = self._snapshot_state() if not exhausted and self._step % self._accumulation_steps == 0 else None
413 return _PreparedWindow(batches, exhausted, checkpoint, balance_logs=self._balance_logs)
415 def _prepare_batch(self) -> Any:
416 """Construct one Host round; accelerator peers receive it at device delivery."""
417 if self._stopped:
418 raise StopIteration
419 batch = self._prepare_next_batch()
420 if batch is not None and self._device_prefetch is not None:
421 batch = self._device_prefetch.prepare_host(batch)
422 if self._device.type == "cpu":
423 batch = self._broadcast_batch(self._stage_batch(batch))
424 self._stopped = batch is None
425 if self._stopped:
426 raise StopIteration
427 return self._consume_batch(batch)
429 def _stage_batch(self, batch: Any) -> Any:
430 if batch is not None and self._device_prefetch is not None:
431 return self._device_prefetch.stage_prepared(batch)
432 return batch
434 def _prepare_next_batch(self) -> Any:
435 """Select, balance, route, and construct samples without model-group broadcast."""
436 if self._data_plane.is_member:
437 return self._produce_on_data_plane()
438 return None
440 def _broadcast_batch(self, constructed_batch: Any) -> Any:
441 """Broadcast one batch on the caller stream, or on the Host producer for CPU."""
442 return self._model_transport.broadcast(constructed_batch)
444 def _consume_batch(self, received: Any) -> Any:
445 """Commit Host preparation; delivery publishes checkpoints only after broadcast."""
446 # Non-Constructors hold placeholders; the caller receives both
447 # real batches and the EOF sentinel from their Constructor's broadcast.
448 planning_reader = self._planning_reader()
449 if planning_reader is not None:
450 planning_reader.commit(self._pending_local_keys)
451 self._pending_local_keys.clear()
452 self._step += 1
453 return received
455 def _fill_requested_device_prefetch(self) -> None:
456 """Complete requested Constructor lookahead before waiting or closing."""
457 # Only accelerator Constructor staging sets a positive target.
458 if (not self._prefetch_target or self._closed or self._delivered_state["stopped"]
459 or self._batch_error is not None):
460 return
461 self._prefetch_device_batches(self._prefetch_target, blocking=True)
463 def wait_for_prefetch(self) -> None:
464 """Drain buffered window communication before checkpoint teardown or shutdown."""
465 self._fill_requested_device_prefetch()
466 if self._last_ready_event is not None:
467 self._last_ready_event.synchronize()
468 ready_event = None
469 for prepared in self._device_batches:
470 if prepared.ready_event is not None:
471 ready_event = prepared.ready_event
472 # Queued events belong to the copy stream, independently of the last
473 # delivered batch's training-stream marker. Neither drain broadcasts ahead.
474 if ready_event is not None:
475 ready_event.synchronize()
476 if self._future is not None:
477 self._future.result()
479 def close(self) -> None:
480 """Drain the producer before releasing buffered inputs and reader resources."""
481 if self._closed:
482 return
483 # Constructors may have skipped optional staging at different times.
484 # Complete the requested lookahead before any producer leaves its collectives;
485 # discarded input errors still belong to next()/wait_for_prefetch(), not close().
486 self._fill_requested_device_prefetch()
487 if self._last_ready_event is not None:
488 self._last_ready_event.synchronize()
489 if self._prefetch_thread is not None:
490 if self._prefetch_thread.is_alive():
491 self._requests.put(None)
492 self._prefetch_thread.join()
493 self._prefetch_thread = None
494 try:
495 if self._device_prefetch is not None:
496 self._device_prefetch.close()
497 finally:
498 reader = self._planning_reader()
499 if reader is not None:
500 reader.close()
501 self._current = None
502 self._future = None
503 self._device_batches.clear()
504 self._consumer_stream = None
505 self._prefetch_slots.clear()
506 self._batch_error = None
507 self._last_ready_event = None
508 self._balance_logs = []
509 self._closed = True
511 def __enter__(self) -> "DistributedDataLoader":
512 """Support deterministic cleanup on early loop exit."""
513 return self
515 def __exit__(self, exc_type: Any, exc_value: Any, traceback: Any) -> None:
516 """Release prefetch resources before process-group teardown."""
517 self.close()
519 def state_dict(self) -> dict[str, Any]:
520 """Snapshot the last fully delivered optimizer step without waiting.
522 Save training checkpoints only after the corresponding optimizer update.
523 An incomplete step leaves this cursor at that step's beginning.
524 """
525 state = self._delivered_state
526 if "_index_sampler_state" not in state:
527 # Old non-Reader checkpoints may lack an absolute source cursor.
528 # Keep their same-mode progress without inventing a portable position.
529 return copy.deepcopy(dict(state, global_rank=self._topology.global_rank))
530 return copy.deepcopy({
531 "_index_sampler_state": state["_index_sampler_state"],
532 "_iterator_finished": state["stopped"],
533 "dataset_state": state["dataset_state"],
534 })
536 def state_dict_for_load(self) -> dict[str, Any]:
537 """Request the shared cursor and legacy Reader aliases without defaulting to zero."""
538 return sampler_state_for_load(self.state_dict())
540 def _snapshot_state(self) -> dict[str, Any]:
541 """Capture producer progress at an optimizer-step boundary."""
542 reader = self._planning_reader()
543 reader_state = reader.state_dict() if reader is not None else None
544 payload_state = (
545 self._direct_sample_loader.state_dict() if self._direct_sample_loader is not None
546 else (reader_state or {}).get("sample_loader")
547 )
548 state = {
549 "epoch": self._epoch,
550 "step": self._step,
551 "stopped": self._stopped,
552 "dataset_state": (
553 {key: payload_state[key] for key in ("dataset_size", "seed")} if payload_state is not None else None
554 ),
555 }
556 if reader_state is not None:
557 state["_index_sampler_state"] = reader_state["sampler"]
558 elif self._initial_consumed_samples is not None:
559 # Model peers have no Reader; advance their cursor from the known initial position.
560 state["_index_sampler_state"] = {
561 "consumed_samples": self._initial_consumed_samples + (
562 self._step * self._planner.local_batch_size * self._topology.data_parallel_size
563 ),
564 "epoch": self._epoch,
565 "global_batch_size": (
566 self._accumulation_steps * self._planner.local_batch_size * self._topology.data_parallel_size
567 ),
568 "dp_world_size": self._topology.data_parallel_size,
569 }
570 sampler = state["_index_sampler_state"]
571 offset = self._initial_consumed_samples % sampler["global_batch_size"]
572 if offset:
573 sampler["optimizer_step_offset"] = offset
574 try:
575 copied_state = copy.deepcopy(state)
576 except Exception as exc:
577 raise ValueError(f"Distributed DataLoader state is not checkpointable: {exc}") from exc
578 return copied_state
580 def load_state_dict(self, state_dict: Mapping[str, Any]) -> None:
581 """Restore a shared native cursor or a legacy rank-local checkpoint.
583 Args:
584 state_dict: A complete-step cursor; legacy rank-local state must match this rank.
585 """
586 self._validate_restore_boundary()
587 try:
588 state = copy.deepcopy(dict(state_dict))
589 except Exception as exc:
590 raise ValueError(f"Distributed DataLoader state is not copyable: {exc}") from exc
591 if state.get("global_rank") not in (None, self._topology.global_rank):
592 raise ValueError(
593 f"Distributed DataLoader checkpoint global_rank={state.get('global_rank')!r} "
594 f"does not match {self._topology.global_rank!r}."
595 )
596 epoch, step, stopped = self._validate_checkpoint_values(state)
597 sampler = get_sampler_checkpoint_state(state)
598 # Legacy step counters use the source accumulation size. The sampler
599 # cursor validates complete steps independently of the target GBS/DP.
600 layout_changed = sampler is not None and (
601 sampler["global_batch_size"] != (
602 self._accumulation_steps * self._planner.local_batch_size * self._topology.data_parallel_size
603 ) or sampler["dp_world_size"] != self._topology.data_parallel_size
604 )
605 if (
606 (not layout_changed and step % self._accumulation_steps) or state.get("window_plans")
607 or state.get("optimizer_step_remaining", self._accumulation_steps) != self._accumulation_steps
608 or any(isinstance(state.get(key), Mapping) and "window" in state[key]
609 for key in ("dataset_reader", "metadata_reader", "sidecar_reader"))
610 ):
611 raise ValueError("Distributed DataLoader checkpoint must align to a complete optimizer step.")
612 if sampler is not None:
613 epoch, step = self._progress_from_sampler(sampler)
614 elif self._dataset_reader is not None or self._metadata_reader is not None or state.get("global_rank") is None:
615 raise ValueError("Distributed DataLoader checkpoint has no native sampler cursor.")
616 self._restore_component_states(state, sampler, epoch)
617 self._epoch = epoch
618 self._step = step
619 self._stopped = stopped
620 self._initial_consumed_samples = None
621 if sampler is not None:
622 self._initial_consumed_samples = sampler["consumed_samples"] - (
623 step * self._planner.local_batch_size * self._topology.data_parallel_size
624 )
626 self._delivered_state = self._snapshot_state()
628 def _progress_from_sampler(self, sampler: Mapping[str, Any]) -> tuple[int, int]:
629 """Validate the shared cursor and derive the epoch and micro-batch position."""
630 global_micro_batch_size = self._planner.local_batch_size * self._topology.data_parallel_size
631 global_batch_size = self._accumulation_steps * global_micro_batch_size
632 consumed, epoch = sampler["consumed_samples"], sampler["epoch"]
633 if any(isinstance(value, bool) or not isinstance(value, int) for value in (consumed, epoch)):
634 raise ValueError("Sampler consumed_samples and epoch must be integers")
635 if sampler_optimizer_step_progress(sampler) or epoch < 0:
636 raise ValueError("Sampler cursor must align to a complete optimizer step")
637 if sampler["global_batch_size"] != global_batch_size:
638 return epoch, 0
639 offset = sampler.get("optimizer_step_offset") or 0
640 return epoch, (consumed - offset) // global_micro_batch_size
642 def _validate_restore_boundary(self) -> None:
643 if self._last_ready_event is not None or self._prefetch_thread is not None or self._step != 0:
644 raise ValueError("load_state_dict must run before distributed iteration starts.")
646 @staticmethod
647 def _validate_checkpoint_values(state: Mapping[str, Any]) -> tuple[int, int, bool]:
648 """Validate and return the saved epoch, micro-batch index and stopped flag."""
649 epoch = state.get("epoch", 0)
650 step = state.get("step", 0)
651 stopped = state.get("_iterator_finished")
652 if stopped is None:
653 stopped = state.get("stopped", False)
654 if not isinstance(epoch, int) or isinstance(epoch, bool) or epoch < 0:
655 raise ValueError("Distributed DataLoader checkpoint epoch must be non-negative.")
656 if not isinstance(step, int) or isinstance(step, bool) or step < 0:
657 raise ValueError("Distributed DataLoader checkpoint step must be non-negative.")
658 if not isinstance(stopped, bool):
659 raise ValueError("Distributed DataLoader checkpoint stopped must be boolean.")
660 return epoch, step, stopped
662 def _restore_component_states(
663 self,
664 state: Mapping[str, Any],
665 sampler_state: Mapping[str, Any] | None,
666 epoch: int,
667 ) -> None:
668 """Rebuild current Readers from shared progress and optional Dataset checks."""
669 reader = self._planning_reader()
670 reader_state = reader.state_dict() if reader is not None else None
671 payload_state = (
672 self._direct_sample_loader.state_dict() if self._direct_sample_loader is not None
673 else (reader_state or {}).get("sample_loader")
674 )
675 # Native torchdata Dataset state belongs to the user's Dataset. Only
676 # cursor-only and legacy balanced states carry our Dataset compatibility fields.
677 if not has_torchdata_state(state):
678 candidates = [state.get("dataset_state"), state.get("direct_sample_loader")]
679 for key in ("dataset_reader", "metadata_reader", "sidecar_reader"):
680 legacy_reader = state.get(key)
681 if isinstance(legacy_reader, Mapping):
682 candidates.append(legacy_reader.get("sample_loader"))
683 for saved in candidates:
684 if saved is None or payload_state is None:
685 continue
686 if not isinstance(saved, Mapping):
687 raise ValueError("Dataset checkpoint state must be a mapping.")
688 for key in ("dataset_size", "seed"):
689 if saved.get(key) is not None and saved[key] != payload_state[key]:
690 raise ValueError(f"Dataset checkpoint {key} does not match the current loader.")
691 if payload_state is not None:
692 payload_state["epoch"] = epoch
693 if reader is not None:
694 reader_state.update(epoch=epoch, sampler=sampler_state)
695 if reader_state["sample_loader"] is not None:
696 reader_state["sample_loader"] = payload_state
697 reader.load_state_dict(reader_state)
698 if self._direct_sample_loader is not None:
699 self._direct_sample_loader.load_state_dict(payload_state)
701 def set_epoch(self, epoch: int) -> None:
702 """Reset the native source while retaining the Host thread and device streams.
704 Args:
705 epoch: Non-negative Dataset epoch.
706 """
707 if isinstance(epoch, bool) or not isinstance(epoch, int) or epoch < 0:
708 raise ValueError("epoch must be a non-negative integer.")
709 if self._closed:
710 raise ValueError("set_epoch cannot restart a closed loader.")
711 if (self._has_delivered or self._delivered_state["step"] != 0) and not self._delivered_state["stopped"]:
712 raise ValueError("set_epoch requires a fresh or exhausted optimizer-step loader.")
713 if epoch == self._delivered_state["epoch"] and not self._delivered_state["stopped"]:
714 return
715 self.wait_for_prefetch()
716 if self._device_prefetch is not None:
717 self._device_prefetch.close()
718 if self._dataset_reader is not None:
719 self._dataset_reader.set_epoch(epoch)
720 if self._metadata_reader is not None:
721 self._metadata_reader.set_epoch(epoch)
722 if self._direct_sample_loader is not None:
723 self._direct_sample_loader.set_epoch(epoch)
724 self._epoch = epoch
725 self._step = 0
726 self._initial_consumed_samples = 0
727 self._stopped = False
728 self._pending_local_keys.clear()
729 self._window_plans.clear()
730 self._current = None
731 self._future = None
732 self._device_batches.clear()
733 self._consumer_stream = None
734 self._prefetch_slots = deque([None] * (self._device_prefetch_depth + 1))
735 self._device_stopped = False
736 self._prefetch_target = 0
737 self._batch_error = None
738 self._last_ready_event = None
739 self._has_delivered = False
740 self._balance_logs = []
741 self._delivered_state = self._snapshot_state()
743 def _produce_on_data_plane(self) -> Any:
744 """Execute one plan through direct reads or payload exchange, then construct the batch."""
745 plan = self._next_plan_control()
746 if plan is None:
747 self._stopped = True
748 return None
750 self._pending_local_keys = {
751 key for key in plan.selected_keys if key.reader_rank == self._topology.global_rank
752 }
753 if self._metadata_mode:
754 if not self._topology.is_constructor:
755 return None
756 received_payloads = self._direct_sample_loader.fetch_samples(
757 plan.local_sample_keys(self._topology.data_rank)
758 )
759 else:
760 prepared_exchange = self._prepare_outgoing(plan, self._pending_local_keys)
761 received_payloads = self._data_plane.exchange_prepared(prepared_exchange)
762 # Reader-only ranks must finish routing before leaving this round.
763 if not self._topology.is_constructor:
764 return None
765 return self._construct_received_payloads(plan, received_payloads)
767 def _construct_received_payloads(
768 self,
769 plan: DistributedPackingPlan,
770 received_payloads: dict[SampleKey, Any],
771 ) -> Any:
772 """Construct one local batch; the caller must be a Data Constructor."""
773 local_batch_plan = plan.local_batch_for(self._topology.data_rank)
774 local_batch = self._data_constructor.construct(local_batch_plan, received_payloads)
775 # None is reserved for distributed end-of-stream.
776 if local_batch is None:
777 raise ValueError("The Data Constructor must return a non-None batch.")
778 return local_batch
780 def _next_plan_control(self) -> DistributedPackingPlan | None:
781 """Pop one micro-batch plan, planning a new window only when the cache is empty."""
782 if self._window_plans:
783 return self._window_plans.popleft()
784 planning_reader = self._planning_reader()
785 local_input = planning_reader.planning_input() if planning_reader is not None else _ReaderPlanInput()
786 if self._broadcast_plan:
787 gathered_inputs = self._data_plane.gather_object_to_planner(local_input)
788 plan = None
789 if self._topology.global_rank == self._data_plane.planner_rank:
790 plan = self._build_plan_control(gathered_inputs)
791 plan = self._data_plane.broadcast_from_planner(plan)
792 else:
793 plan = self._build_plan_control(self._data_plane.all_gather_object(local_input))
794 if isinstance(plan, tuple):
795 self._window_plans.extend(plan)
796 return self._window_plans.popleft()
797 return plan
799 def _prepare_local_metadata(self, window_micro_batch_num: int) -> None:
800 """Read one explicitly sized source window before constructing its batches."""
801 if self._topology.global_rank not in self._dataset_reader_ranks:
802 return
803 planning_reader = self._planning_reader()
804 planning_reader.prepare_next_step(window_micro_batch_num)
806 def _build_plan_control(
807 self, gathered_inputs: tuple[_ReaderPlanInput, ...],
808 ) -> _PlanControl:
809 """Return the native window's plans; None means end of stream."""
810 # Collective rank order can differ from the mesh's logical DP order.
811 inputs_by_rank = dict(zip(self._data_plane.ranks, gathered_inputs))
812 reader_inputs = [
813 inputs_by_rank[rank] for rank in self._topology.constructor_ranks
814 if rank in self._dataset_reader_ranks
815 ]
816 return self._plan_native_batch(reader_inputs)
818 def _native_window_micro_batch_num(self, reader_inputs: list[_ReaderPlanInput]) -> int:
819 """Validate complete synchronized source rounds before constructing reference bins."""
820 if all(not item.original_batches for item in reader_inputs):
821 return 0
822 if any(not item.original_batches for item in reader_inputs):
823 raise ValueError("Native BatchSampler ranks exhausted at different forward/backward rounds.")
824 counts = {len(item.original_batches) for item in reader_inputs}
825 if len(counts) != 1 or any(
826 len(batch) != self._planner.local_batch_size
827 for item in reader_inputs for batch in item.original_batches
828 ):
829 raise ValueError("Native BatchSampler ranks must provide equal complete micro-batch windows.")
830 return next(iter(counts))
832 def _plan_native_batch(
833 self, reader_inputs: list[_ReaderPlanInput],
834 ) -> _PlanControl:
835 """Plan the native sampler window, returning one plan per training round."""
836 window_micro_batch_num = self._native_window_micro_batch_num(reader_inputs)
837 if not window_micro_batch_num:
838 return None
839 reference_bins = tuple(
840 keys for batches in zip(*(item.original_batches for item in reader_inputs))
841 for batch in batches for keys in batch
842 )
843 metadata_by_key = {metadata.key: metadata for item in reader_inputs for metadata in item.metadata}
844 samples = tuple(metadata_by_key[key] for keys in reference_bins for key in keys)
845 plan = self._planner.plan if window_micro_batch_num == 1 else self._planner.plan_window
846 plans = plan(samples, reference_bins=reference_bins, step=self._step)
847 self._capture_balance_log(samples, reference_bins, plans)
848 return plans
850 def _capture_balance_log(
851 self, samples: tuple[SampleMetadata, ...], reference_bins: tuple[tuple[SampleKey, ...], ...],
852 plans: DistributedPackingPlan | tuple[DistributedPackingPlan, ...],
853 ) -> None:
854 """Freeze scalar window statistics before the producer advances to another plan."""
855 if (
856 not self._log_balance
857 or self._topology.global_rank != self._data_plane.planner_rank
858 or not balance_logger.isEnabledFor(logging.INFO)
859 ):
860 return
861 plans = plans if isinstance(plans, tuple) else (plans,)
862 metadata = {item.key: item for item in samples}
863 costs = self._planner.last_sample_costs
865 def describe(keys: tuple[SampleKey, ...]) -> dict[str, Any]:
866 """Combine evaluated document costs with this layout's constructor padding.
868 Args:
869 keys: Sample occurrences assigned to one pack.
870 """
871 tokens = sum(metadata[key].pack_tokens for key in keys)
872 padding = max(0, self._planner.seq_len - tokens)
873 return {
874 "seq_len": tokens, "samples": len(keys), "padding": padding,
875 "cost": math.fsum(costs[key].llm for key in keys) + self._planner.padding_cost(padding).llm,
876 }
878 local_size = self._planner.local_batch_size
879 bin_count = self._planner.distributed_bin_count
880 rounds = []
881 for offset, plan in enumerate(plans):
882 original = reference_bins[offset * bin_count:(offset + 1) * bin_count]
883 rounds.append({
884 "micro_batch": plan.step % self._accumulation_steps + 1,
885 "bins_before": tuple(
886 tuple(describe(keys) for keys in original[start:start + local_size])
887 for start in range(0, bin_count, local_size)
888 ),
889 "bins_after": tuple(
890 tuple(describe(packing_bin.sample_keys) for packing_bin in local_batch)
891 for local_batch in plan.local_batches
892 ),
893 })
894 stats = {
895 **self._planner.last_balance_decision,
896 "objective": self._planner.objective, "epoch": self._epoch,
897 "planner_rank": self._data_plane.planner_rank,
898 "constructor_ranks": self._topology.constructor_ranks,
899 "group_ranks": tuple(range(self._topology.data_parallel_size)), "rounds": rounds,
900 }
901 self._balance_logs.append(format_balance_stats(stats, self._step // self._accumulation_steps + 1))
903 def _prepare_outgoing(
904 self,
905 plan: DistributedPackingPlan,
906 local_selected_keys: set[SampleKey],
907 ) -> PreparedPayloadExchange:
908 """Group source-owned payloads by destination Constructor for transport encoding."""
909 payloads_by_target_rank: dict[int, list[tuple[SampleKey, Any]]] = {}
910 if local_selected_keys:
911 payloads = self._dataset_reader.selected_payloads(local_selected_keys)
912 target_by_key = {
913 key: self._topology.constructor_ranks[data_rank]
914 for data_rank, local_batch in enumerate(plan.local_batches)
915 for packing_bin in local_batch
916 for key in packing_bin.sample_keys
917 }
918 for key, payload in payloads:
919 payloads_by_target_rank.setdefault(target_by_key[key], []).append((key, payload))
920 return self._data_plane.prepare_exchange(payloads_by_target_rank)
922 def _planning_reader(self) -> Any | None:
923 """Return this rank's online or metadata-only Dataset Reader."""
924 if self._dataset_reader is not None:
925 return self._dataset_reader
926 return self._metadata_reader
930__all__ = ["DistributedDataLoader"]