Coverage for / home / jenkins / .local / lib / python3.10 / site-packages / hyper_parallel / data / parallel / batch_sampler.py: 91%
206 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-agnostic DP index slicing and resumable micro-batch sampling."""
17from __future__ import annotations
19import copy
20import operator
21import os
22import random
23from collections.abc import Iterator, Mapping, Sequence
24from typing import Any, Literal, Union, cast
26import torch
28from hyper_parallel.data.dataset_logging import get_dataset_logger
30IndexMapping = Union[Mapping[int, int], Sequence[int]]
31SamplerType = Literal["single", "cyclic"]
33logger = get_dataset_logger(__name__)
36def sampler_state_for_load(state_dict: Mapping[str, Any]) -> dict[str, Any]:
37 """Build a loader-owned checkpoint template without inventing missing progress.
39 Args:
40 state_dict: Fresh fixed or balanced loader state.
42 Returns:
43 Loading-only placeholders for the shared cursor and legacy Reader paths.
44 """
45 state = copy.deepcopy(dict(state_dict))
46 balanced = "_num_yielded" not in state and "_snapshot" not in state
47 if not balanced and get_sampler_checkpoint_state(state) is None:
48 return state
49 cursor = dict.fromkeys(("consumed_samples", "epoch", "global_batch_size", "dp_world_size", "optimizer_step_offset"))
50 state["_index_sampler_state"] = dict(cursor)
51 state["_iterator_finished"] = None
52 # Missing torchdata fields must not make a cursor-only checkpoint look like
53 # the target's freshly initialized worker snapshot.
54 state["_num_yielded"] = None
55 snapshot = state.setdefault("_snapshot", {}).setdefault("_main_snapshot", {})
56 snapshot["_num_workers"] = None
57 if isinstance(snapshot.get("_index_sampler_state"), dict):
58 snapshot["_index_sampler_state"]["optimizer_step_offset"] = None
59 state["global_rank"] = None
60 payload = dict.fromkeys(("dataset_size", "seed"))
61 for key in ("dataset_reader", "metadata_reader", "sidecar_reader"):
62 state[key] = {"sampler": dict(cursor), "sample_loader": dict(payload)}
63 if balanced:
64 state.update(epoch=0, step=0, stopped=False, direct_sample_loader=dict(payload))
65 state["dataset_state"] = dict(payload)
66 return state
69def has_torchdata_state(state_dict: Mapping[str, Any]) -> bool:
70 """Return whether a loaded checkpoint contains real torchdata iterator state."""
71 snapshot = state_dict.get("_snapshot")
72 return state_dict.get("_num_yielded") is not None or (
73 isinstance(snapshot, Mapping)
74 and snapshot.get("_main_snapshot", {}).get("_num_workers") is not None
75 )
78def get_sampler_checkpoint_state(state_dict: Mapping[str, Any]) -> Mapping[str, Any] | None:
79 """Read the shared native cursor, accepting legacy Reader-only checkpoints.
81 Args:
82 state_dict: Saved loader state or a template filled by checkpoint loading.
84 Returns:
85 Saved sampler state, or None when no cursor was loaded.
86 """
87 candidates = [state_dict.get("_index_sampler_state")]
88 for key in ("dataset_reader", "metadata_reader", "sidecar_reader"):
89 reader = state_dict.get(key)
90 if isinstance(reader, Mapping):
91 candidates.append(reader.get("sampler"))
92 return next((state for state in candidates
93 if isinstance(state, Mapping) and state.get("consumed_samples") is not None), None)
96def sampler_optimizer_step_progress(state: Mapping[str, Any]) -> int:
97 """Return samples accumulated since the saved optimizer-step boundary.
99 Args:
100 state: Native sampler cursor; old checkpoints use an offset of zero.
102 Returns:
103 Zero at a complete optimizer step, otherwise the partial sample count.
104 """
105 consumed, batch_size = state["consumed_samples"], state["global_batch_size"]
106 offset = state.get("optimizer_step_offset")
107 offset = 0 if offset is None else offset
108 if any(isinstance(value, bool) or not isinstance(value, int) for value in (consumed, batch_size, offset)):
109 raise ValueError("Sampler cursor, global_batch_size and optimizer_step_offset must be integers")
110 if not 0 <= offset <= consumed or offset >= batch_size:
111 raise ValueError("Sampler optimizer_step_offset must be within the cursor and global_batch_size")
112 return (consumed - offset) % batch_size
115def _validate_positive_integer(value: int, name: str) -> None:
116 """Validate an integer boundary shared by sampler options."""
117 if value <= 0:
118 raise ValueError(f"{name} must be a positive integer")
121class _DatasetBatchSampler:
122 """Share validation, DP slicing, epoch, and checkpoint state.
124 BatchSampler splits one global batch across DP ranks
125 → each rank receives its logical Dataset index
126 → Dataset.__getitem__(index)
127 → shuffle_index maps the logical index to a random packed sample
128 """
130 def __init__(
131 self,
132 *,
133 total_samples: int,
134 consumed_samples: int,
135 micro_batch_size: int,
136 global_batch_size: int,
137 dp_rank: int,
138 dp_world_size: int,
139 drop_last: bool,
140 index_mapping: IndexMapping | None,
141 ) -> None:
142 """Validate and store DP slicing and checkpoint state."""
143 _validate_positive_integer(total_samples, "total_samples")
144 _validate_positive_integer(micro_batch_size, "micro_batch_size")
145 _validate_positive_integer(dp_world_size, "dp_world_size")
146 self._validate_consumed_samples(consumed_samples, total_samples)
147 if not 0 <= dp_rank < dp_world_size:
148 raise ValueError(f"dp_rank must be in [0, {dp_world_size}), but got {dp_rank!r}")
150 self.total_samples = total_samples
151 self.consumed_samples = consumed_samples
152 self.micro_batch_size = micro_batch_size
153 self.dp_rank = dp_rank
154 self.dp_world_size = dp_world_size
155 self.drop_last = drop_last
156 self.index_mapping = index_mapping
157 self.global_micro_batch_size = micro_batch_size * dp_world_size
158 _validate_positive_integer(global_batch_size, "global_batch_size")
159 if global_batch_size % self.global_micro_batch_size != 0:
160 raise ValueError("global_batch_size must be divisible by micro_batch_size * dp_world_size")
162 self.global_batch_size = global_batch_size
163 self._optimizer_step_offset = 0
164 self.epoch = 0
165 self._resume_at_source_batch = False
167 @staticmethod
168 def _validate_consumed_samples(consumed_samples: int, total_samples: int) -> None:
169 """Allow sequential sampling to stop exactly at the epoch boundary."""
170 if not 0 <= consumed_samples <= total_samples:
171 raise ValueError("consumed_samples must be in [0, total_samples]")
173 def __iter__(self) -> Iterator[list[int]]:
174 """Yield sequential rank-local indices for ``single`` mode."""
175 while self.consumed_samples < self.total_samples:
176 block_start = self.consumed_samples
177 block_stop = min(block_start + self.global_micro_batch_size, self.total_samples)
178 block_size = block_stop - block_start
179 if self.drop_last and block_size < self.global_micro_batch_size:
180 self.consumed_samples = self.total_samples
181 break
183 local_start = block_start + self.dp_rank * self.micro_batch_size
184 local_stop = min(local_start + self.micro_batch_size, block_stop)
185 self.consumed_samples = block_stop
186 if local_start >= block_stop:
187 continue
189 local_indices = [
190 self._resolve_index(index)
191 for index in range(local_start, local_stop)
192 ]
193 yield local_indices
195 def __len__(self) -> int:
196 """Return the remaining number of micro-batches for this DP rank."""
197 remaining_samples = self.total_samples - self.consumed_samples
198 full_batches, partial_samples = divmod(remaining_samples, self.global_micro_batch_size)
199 if self.drop_last or partial_samples <= self.dp_rank * self.micro_batch_size:
200 return full_batches
201 return full_batches + 1
203 def _resolve_index(self, logical_index: int) -> int:
204 """Apply an optional logical-to-physical sample index mapping."""
205 if self.index_mapping is None:
206 return logical_index
208 resolved_index = self.index_mapping[logical_index]
209 if isinstance(resolved_index, bool):
210 raise TypeError(f"index_mapping[{logical_index}] must be an integer")
212 physical_index = operator.index(resolved_index)
213 return physical_index
215 def state_dict(self) -> dict[str, int]:
216 """Return the Dataset epoch and position needed to resume iteration."""
217 sampler_state = {
218 "consumed_samples": self.consumed_samples,
219 "epoch": self.epoch,
220 "global_batch_size": self.global_batch_size,
221 "dp_world_size": self.dp_world_size,
222 }
223 if self._optimizer_step_offset:
224 sampler_state["optimizer_step_offset"] = self._optimizer_step_offset
225 return sampler_state
227 def enable_source_batch_resume(self) -> None:
228 """Restore this sampler at a source-batch boundary for Online batching."""
229 self._resume_at_source_batch = True
231 def load_state_dict(self, state_dict: Mapping[str, Any]) -> None:
232 """Restore the global sample position from a checkpoint.
234 Args:
235 state_dict: Sampler state produced by :meth:`state_dict`.
237 Raises:
238 ValueError: If the checkpoint does not contain a valid sample position.
239 """
240 consumed_samples = state_dict["consumed_samples"]
241 self._validate_consumed_samples(consumed_samples, self.total_samples)
242 progress = sampler_optimizer_step_progress(state_dict)
243 repartition = (state_dict["global_batch_size"] != self.global_batch_size
244 or state_dict["dp_world_size"] != self.dp_world_size)
245 if progress and (repartition or not self._resume_at_source_batch):
246 raise ValueError("Sampler restoration requires a complete global optimizer batch")
247 offset = consumed_samples % self.global_batch_size if repartition else (
248 state_dict.get("optimizer_step_offset") or 0
249 )
250 if (self._resume_at_source_batch and consumed_samples != self.total_samples
251 and (consumed_samples - offset) % self.global_micro_batch_size):
252 raise ValueError("Online source sampler state must align to a distributed micro-batch")
254 self._optimizer_step_offset = offset
255 self.consumed_samples = consumed_samples
256 self.epoch = state_dict["epoch"]
258 def set_epoch(self, epoch: int) -> None:
259 """Start a new epoch while preserving progress in the restored epoch.
261 Args:
262 epoch: Zero-based training epoch.
264 Raises:
265 ValueError: If ``epoch`` is not a non-negative integer.
266 """
267 if isinstance(epoch, bool) or not isinstance(epoch, int) or epoch < 0:
268 raise ValueError("epoch must be a non-negative integer")
270 if epoch != self.epoch:
271 self.consumed_samples = 0
272 self._optimizer_step_offset = 0
273 self.epoch = epoch
276class _CyclicDatasetBatchSampler(_DatasetBatchSampler):
277 """Yield one deterministic shuffled epoch for ``cyclic`` mode."""
279 def __init__(self, *, data_sharding: bool, seed: int, **sampler_options: Any) -> None:
280 """Store the cyclic data-sharding policy."""
281 super().__init__(**sampler_options)
282 if not self.drop_last:
283 raise ValueError("cyclic sampling requires drop_last=True")
285 self.data_sharding = data_sharding
286 self.seed = seed
288 def load_state_dict(self, state_dict: Mapping[str, Any]) -> None:
289 """Restore a global cursor, repartitioning only global shuffle order."""
290 if self.data_sharding and state_dict["dp_world_size"] != self.dp_world_size:
291 raise ValueError("elastic DP resume requires data_sharding=False")
293 super().load_state_dict(state_dict)
295 def __len__(self) -> int:
296 """Return the number of local micro-batches remaining in this epoch."""
297 active_samples = self._active_samples()
298 remaining_samples = active_samples - self.consumed_samples
299 return remaining_samples // self.global_micro_batch_size
301 def __iter__(self) -> Iterator[list[int]]:
302 """Build the current epoch permutation and DP slices."""
303 active_samples = self._active_samples()
304 if (self.consumed_samples - self._optimizer_step_offset) % self.global_micro_batch_size != 0:
305 raise ValueError("cyclic consumed_samples must align to a distributed micro-batch")
307 logger.debug(
308 "Starting cyclic Dataset epoch=%d at consumed_samples=%d",
309 self.epoch,
310 self.consumed_samples,
311 )
313 if self.data_sharding:
314 # Shuffle one contiguous bucket per DP rank.
315 bucket_size = active_samples // self.dp_world_size
316 bucket_offset = self.consumed_samples // self.dp_world_size
317 bucket_start = self.dp_rank * bucket_size
318 random_offsets = list(range(bucket_size))
319 random.Random(self.seed + self.epoch).shuffle(random_offsets)
320 rank_indices = [bucket_start + offset for offset in random_offsets[bucket_offset:]]
321 else:
322 # Shuffle before dropping the tail so source-file tail samples remain eligible.
323 # Legacy checkpoints retain their cursor but not their old truncated permutation.
324 shuffled_indices = list(range(self.total_samples))
325 random.Random(self.seed + self.epoch).shuffle(shuffled_indices)
326 active_indices = shuffled_indices[:active_samples]
327 active_indices = active_indices[self.consumed_samples:]
328 rank_indices = active_indices[self.dp_rank::self.dp_world_size]
330 local_batch = []
331 for index in rank_indices:
332 local_batch.append(index)
333 if len(local_batch) == self.micro_batch_size:
334 self.consumed_samples += self.global_micro_batch_size
335 yield local_batch
336 local_batch = []
338 def _active_samples(self) -> int:
339 """Return the number of samples fitting complete global micro-batches."""
340 remaining = self.total_samples - self.consumed_samples
341 active_samples = self.total_samples - remaining % self.global_micro_batch_size
342 if active_samples <= 0:
343 raise ValueError("cyclic Dataset must contain one complete global micro-batch")
345 return active_samples
348def _resolve_index_mapping(
349 index_mapping: IndexMapping | None,
350 data_rearrange_map: IndexMapping | str | os.PathLike[str] | None,
351) -> IndexMapping | None:
352 """Resolve the optional rearrangement-map configuration."""
353 if index_mapping is not None and data_rearrange_map is not None:
354 raise ValueError("configure only one of index_mapping and data_rearrange_map")
356 mapping_source = data_rearrange_map if data_rearrange_map is not None else index_mapping
357 if isinstance(mapping_source, (str, os.PathLike)):
358 mapping_path = os.fspath(mapping_source)
359 loaded_mapping = torch.load(f=mapping_path)
360 resolved_mapping = cast(IndexMapping, loaded_mapping)
361 return resolved_mapping
362 return mapping_source
365def build_dataset_batch_sampler(
366 *,
367 total_samples: int,
368 micro_batch_size: int,
369 global_batch_size: int,
370 dp_rank: int,
371 dp_world_size: int,
372 consumed_samples: int = 0,
373 drop_last: bool = True,
374 index_mapping: IndexMapping | None = None,
375 data_rearrange_map: IndexMapping | str | os.PathLike[str] | None = None,
376 sampler_type: SamplerType = "single",
377 data_sharding: bool = False,
378 seed: int = 0,
379) -> _DatasetBatchSampler:
380 """Build a Dataset batch sampler.
382 The sampler groups logical indices into global DP blocks, gives each DP rank
383 one contiguous micro-batch, and checkpoints the next global sample position.
384 It does not know how samples are stored or transformed.
386 Supported scenarios:
387 - ``single``: Sample sequential Dataset indices.
388 - ``single`` with ``data_rearrange_map``: Resolve indices through an
389 in-memory mapping or a mapping checkpoint path.
390 - ``cyclic``: Use epoch-based ``randperm`` order, with
391 ``data_sharding=True/False`` and resumable sampler state.
393 Args:
394 total_samples: Number of logical samples in the Dataset.
395 micro_batch_size: Samples consumed by one DP rank per micro-step.
396 global_batch_size: Fixed samples consumed globally by one optimizer step.
397 dp_rank: Rank inside the data-parallel group.
398 dp_world_size: Number of ranks in the data-parallel group.
399 consumed_samples: Global logical sample position restored from a checkpoint.
400 drop_last: Whether to omit an incomplete global DP block.
401 index_mapping: Optional logical-to-physical sample index mapping.
402 data_rearrange_map: Mapping object or checkpoint path.
403 sampler_type: ``single`` for sequential sampling or ``cyclic`` for a
404 deterministic shuffled epoch.
405 data_sharding: Whether cyclic sampling shuffles an independent
406 contiguous bucket on each DP rank.
407 seed: Base seed used by cyclic epoch shuffling.
409 Returns:
410 A resumable iterable of rank-local index lists.
411 """
412 resolved_mapping = _resolve_index_mapping(index_mapping, data_rearrange_map)
413 logger.debug(
414 "Building Dataset sampler: type=%s, total_samples=%d, consumed_samples=%d, micro_batch_size=%d, "
415 "dp_rank=%d, dp_world_size=%d, drop_last=%s, data_sharding=%s, mapped=%s",
416 sampler_type, total_samples, consumed_samples, micro_batch_size, dp_rank, dp_world_size, drop_last,
417 data_sharding, resolved_mapping is not None,
418 )
419 sampler_options = {
420 "total_samples": total_samples,
421 "consumed_samples": consumed_samples,
422 "micro_batch_size": micro_batch_size,
423 "global_batch_size": global_batch_size,
424 "dp_rank": dp_rank,
425 "dp_world_size": dp_world_size,
426 "drop_last": drop_last,
427 "index_mapping": resolved_mapping,
428 }
430 if sampler_type == "single":
431 # Scenario 1 — single without a mapping: sequential Dataset indices.
432 # Scenario 2 — single + data_rearrange_map: the same logical sequence
433 # is resolved through an in-memory mapping or mapping checkpoint.
434 batch_sampler = _DatasetBatchSampler(**sampler_options)
435 return batch_sampler
437 if sampler_type == "cyclic":
438 # Scenario 3 — cyclic: use epoch-based randperm order,
439 # support data_sharding=True/False, and checkpoint consumed_samples.
440 if resolved_mapping is not None:
441 raise ValueError("cyclic sampling does not support a rearrangement map")
443 batch_sampler = _CyclicDatasetBatchSampler(
444 data_sharding=data_sharding,
445 seed=seed,
446 **sampler_options,
447 )
448 return batch_sampler
450 raise ValueError("sampler_type must be one of: single, cyclic")