Coverage for / home / jenkins / .local / lib / python3.10 / site-packages / hyper_parallel / data / parallel / batch_sampler.py: 19%
162 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"""Dataset-agnostic DP index slicing and resumable micro-batch sampling."""
17from __future__ import annotations
19import operator
20import os
21import random
22from collections.abc import Iterator, Mapping, Sequence
23from typing import Any, Literal, Union, cast
25import torch
27from hyper_parallel.data.dataset_logging import get_dataset_logger
29IndexMapping = Union[Mapping[int, int], Sequence[int]]
30SamplerType = Literal["single", "cyclic"]
32logger = get_dataset_logger(__name__)
35def _validate_positive_integer(value: int, name: str) -> None:
36 """Validate an integer boundary shared by sampler options."""
37 if value <= 0:
38 raise ValueError(f"{name} must be a positive integer")
41class _DatasetBatchSampler:
42 """Share validation, DP slicing, epoch, and checkpoint state.
44 BatchSampler splits one global batch across DP ranks
45 → each rank receives its logical Dataset index
46 → Dataset.__getitem__(index)
47 → shuffle_index maps the logical index to a random packed sample
48 """
50 def __init__(
51 self,
52 *,
53 total_samples: int,
54 consumed_samples: int,
55 micro_batch_size: int,
56 global_batch_size: int,
57 dp_rank: int,
58 dp_world_size: int,
59 drop_last: bool,
60 index_mapping: IndexMapping | None,
61 ) -> None:
62 """Validate and store DP slicing and checkpoint state."""
63 _validate_positive_integer(total_samples, "total_samples")
64 _validate_positive_integer(micro_batch_size, "micro_batch_size")
65 _validate_positive_integer(dp_world_size, "dp_world_size")
66 self._validate_consumed_samples(consumed_samples, total_samples)
67 if not 0 <= dp_rank < dp_world_size:
68 raise ValueError(f"dp_rank must be in [0, {dp_world_size}), but got {dp_rank!r}")
70 self.total_samples = total_samples
71 self.consumed_samples = consumed_samples
72 self.micro_batch_size = micro_batch_size
73 self.dp_rank = dp_rank
74 self.dp_world_size = dp_world_size
75 self.drop_last = drop_last
76 self.index_mapping = index_mapping
77 self.global_micro_batch_size = micro_batch_size * dp_world_size
78 _validate_positive_integer(global_batch_size, "global_batch_size")
79 if global_batch_size % self.global_micro_batch_size != 0:
80 raise ValueError("global_batch_size must be divisible by micro_batch_size * dp_world_size")
82 self.global_batch_size = global_batch_size
83 self.epoch = 0
84 self._resume_at_source_batch = False
86 @staticmethod
87 def _validate_consumed_samples(consumed_samples: int, total_samples: int) -> None:
88 """Allow sequential sampling to stop exactly at the epoch boundary."""
89 if not 0 <= consumed_samples <= total_samples:
90 raise ValueError("consumed_samples must be in [0, total_samples]")
92 def __iter__(self) -> Iterator[list[int]]:
93 """Yield sequential rank-local indices for ``single`` mode."""
94 while self.consumed_samples < self.total_samples:
95 block_start = self.consumed_samples
96 block_stop = min(block_start + self.global_micro_batch_size, self.total_samples)
97 block_size = block_stop - block_start
98 if self.drop_last and block_size < self.global_micro_batch_size:
99 self.consumed_samples = self.total_samples
100 break
102 local_start = block_start + self.dp_rank * self.micro_batch_size
103 local_stop = min(local_start + self.micro_batch_size, block_stop)
104 self.consumed_samples = block_stop
105 if local_start >= block_stop:
106 continue
108 local_indices = [
109 self._resolve_index(index)
110 for index in range(local_start, local_stop)
111 ]
112 yield local_indices
114 def __len__(self) -> int:
115 """Return the remaining number of micro-batches for this DP rank."""
116 remaining_samples = self.total_samples - self.consumed_samples
117 full_batches, partial_samples = divmod(remaining_samples, self.global_micro_batch_size)
118 if self.drop_last or partial_samples <= self.dp_rank * self.micro_batch_size:
119 return full_batches
120 return full_batches + 1
122 def _resolve_index(self, logical_index: int) -> int:
123 """Apply an optional logical-to-physical sample index mapping."""
124 if self.index_mapping is None:
125 return logical_index
127 resolved_index = self.index_mapping[logical_index]
128 if isinstance(resolved_index, bool):
129 raise TypeError(f"index_mapping[{logical_index}] must be an integer")
131 physical_index = operator.index(resolved_index)
132 return physical_index
134 def state_dict(self) -> dict[str, int]:
135 """Return the Dataset epoch and position needed to resume iteration."""
136 sampler_state = {
137 "consumed_samples": self.consumed_samples,
138 "epoch": self.epoch,
139 "global_batch_size": self.global_batch_size,
140 "dp_world_size": self.dp_world_size,
141 }
142 return sampler_state
144 def enable_source_batch_resume(self) -> None:
145 """Restore this sampler at a source-batch boundary for Online batching."""
146 self._resume_at_source_batch = True
148 def load_state_dict(self, state_dict: Mapping[str, Any]) -> None:
149 """Restore the global sample position from a checkpoint.
151 Args:
152 state_dict: Sampler state produced by :meth:`state_dict`.
154 Raises:
155 ValueError: If the checkpoint does not contain a valid sample position.
156 """
157 consumed_samples = state_dict["consumed_samples"]
158 self._validate_consumed_samples(consumed_samples, self.total_samples)
159 if state_dict["global_batch_size"] != self.global_batch_size:
160 if self._resume_at_source_batch:
161 raise ValueError("Online dataloader resume requires an unchanged global_batch_size")
163 raise ValueError("elastic DP resume requires an unchanged global_batch_size")
165 if self._resume_at_source_batch:
166 if state_dict["dp_world_size"] != self.dp_world_size:
167 raise ValueError("Online dataloader resume does not support DP world-size changes")
169 if (
170 consumed_samples != self.total_samples
171 and consumed_samples % self.global_micro_batch_size != 0
172 ):
173 raise ValueError("Online source sampler state must align to a distributed micro-batch")
175 elif consumed_samples % self.global_batch_size != 0:
176 raise ValueError("sampler state must align to a global optimizer batch")
178 self.consumed_samples = consumed_samples
179 self.epoch = state_dict["epoch"]
181 def set_epoch(self, epoch: int) -> None:
182 """Start a new epoch while preserving progress in the restored epoch.
184 Args:
185 epoch: Zero-based training epoch.
187 Raises:
188 ValueError: If ``epoch`` is not a non-negative integer.
189 """
190 if isinstance(epoch, bool) or not isinstance(epoch, int) or epoch < 0:
191 raise ValueError("epoch must be a non-negative integer")
193 if epoch != self.epoch:
194 self.consumed_samples = 0
195 self.epoch = epoch
198class _CyclicDatasetBatchSampler(_DatasetBatchSampler):
199 """Yield one deterministic shuffled epoch for ``cyclic`` mode."""
201 def __init__(self, *, data_sharding: bool, seed: int, **sampler_options: Any) -> None:
202 """Store the cyclic data-sharding policy."""
203 super().__init__(**sampler_options)
204 if not self.drop_last:
205 raise ValueError("cyclic sampling requires drop_last=True")
207 self.data_sharding = data_sharding
208 self.seed = seed
210 def load_state_dict(self, state_dict: Mapping[str, Any]) -> None:
211 """Restore a fixed-GBS cursor, repartitioning only global shuffle order."""
212 if self.data_sharding and state_dict["dp_world_size"] != self.dp_world_size:
213 raise ValueError("elastic DP resume requires data_sharding=False")
215 super().load_state_dict(state_dict)
217 def __len__(self) -> int:
218 """Return the number of local micro-batches remaining in this epoch."""
219 active_samples = self._active_samples()
220 remaining_samples = active_samples - self.consumed_samples
221 return remaining_samples // self.global_micro_batch_size
223 def __iter__(self) -> Iterator[list[int]]:
224 """Build the current epoch permutation and DP slices."""
225 active_samples = self._active_samples()
226 if self.consumed_samples % self.global_micro_batch_size != 0:
227 raise ValueError("cyclic consumed_samples must align to a distributed micro-batch")
229 logger.debug(
230 "Starting cyclic Dataset epoch=%d at consumed_samples=%d",
231 self.epoch,
232 self.consumed_samples,
233 )
235 if self.data_sharding:
236 # Shuffle one contiguous bucket per DP rank.
237 bucket_size = active_samples // self.dp_world_size
238 bucket_offset = self.consumed_samples // self.dp_world_size
239 bucket_start = self.dp_rank * bucket_size
240 random_offsets = list(range(bucket_size))
241 random.Random(self.seed + self.epoch).shuffle(random_offsets)
242 rank_indices = [bucket_start + offset for offset in random_offsets[bucket_offset:]]
243 else:
244 # Shuffle the complete global-batch region, then stride by DP rank.
245 shuffled_indices = list(range(active_samples))
246 random.Random(self.seed + self.epoch).shuffle(shuffled_indices)
247 active_indices = shuffled_indices[self.consumed_samples:]
248 rank_indices = active_indices[self.dp_rank::self.dp_world_size]
250 local_batch = []
251 for index in rank_indices:
252 local_batch.append(index)
253 if len(local_batch) == self.micro_batch_size:
254 self.consumed_samples += self.global_micro_batch_size
255 yield local_batch
256 local_batch = []
258 def _active_samples(self) -> int:
259 """Return the complete global micro-batch region reused every epoch."""
260 active_samples = self.total_samples - self.total_samples % self.global_micro_batch_size
261 if active_samples <= 0:
262 raise ValueError("cyclic Dataset must contain one complete global micro-batch")
264 return active_samples
267def _resolve_index_mapping(
268 index_mapping: IndexMapping | None,
269 data_rearrange_map: IndexMapping | str | os.PathLike[str] | None,
270) -> IndexMapping | None:
271 """Resolve the optional rearrangement-map configuration."""
272 if index_mapping is not None and data_rearrange_map is not None:
273 raise ValueError("configure only one of index_mapping and data_rearrange_map")
275 mapping_source = data_rearrange_map if data_rearrange_map is not None else index_mapping
276 if isinstance(mapping_source, (str, os.PathLike)):
277 mapping_path = os.fspath(mapping_source)
278 loaded_mapping = torch.load(f=mapping_path)
279 resolved_mapping = cast(IndexMapping, loaded_mapping)
280 return resolved_mapping
281 return mapping_source
284def build_dataset_batch_sampler(
285 *,
286 total_samples: int,
287 micro_batch_size: int,
288 global_batch_size: int,
289 dp_rank: int,
290 dp_world_size: int,
291 consumed_samples: int = 0,
292 drop_last: bool = True,
293 index_mapping: IndexMapping | None = None,
294 data_rearrange_map: IndexMapping | str | os.PathLike[str] | None = None,
295 sampler_type: SamplerType = "single",
296 data_sharding: bool = False,
297 seed: int = 0,
298) -> _DatasetBatchSampler:
299 """Build a Dataset batch sampler.
301 The sampler groups logical indices into global DP blocks, gives each DP rank
302 one contiguous micro-batch, and checkpoints the next global sample position.
303 It does not know how samples are stored or transformed.
305 Supported scenarios:
306 - ``single``: Sample sequential Dataset indices.
307 - ``single`` with ``data_rearrange_map``: Resolve indices through an
308 in-memory mapping or a mapping checkpoint path.
309 - ``cyclic``: Use epoch-based ``randperm`` order, with
310 ``data_sharding=True/False`` and resumable sampler state.
312 Args:
313 total_samples: Number of logical samples in the Dataset.
314 micro_batch_size: Samples consumed by one DP rank per micro-step.
315 global_batch_size: Fixed samples consumed globally by one optimizer step.
316 dp_rank: Rank inside the data-parallel group.
317 dp_world_size: Number of ranks in the data-parallel group.
318 consumed_samples: Global logical sample position restored from a checkpoint.
319 drop_last: Whether to omit an incomplete global DP block.
320 index_mapping: Optional logical-to-physical sample index mapping.
321 data_rearrange_map: Mapping object or checkpoint path.
322 sampler_type: ``single`` for sequential sampling or ``cyclic`` for a
323 deterministic shuffled epoch.
324 data_sharding: Whether cyclic sampling shuffles an independent
325 contiguous bucket on each DP rank.
326 seed: Base seed used by cyclic epoch shuffling.
328 Returns:
329 A resumable iterable of rank-local index lists.
330 """
331 resolved_mapping = _resolve_index_mapping(index_mapping, data_rearrange_map)
332 logger.debug(
333 "Building Dataset sampler: type=%s, total_samples=%d, consumed_samples=%d, micro_batch_size=%d, "
334 "dp_rank=%d, dp_world_size=%d, drop_last=%s, data_sharding=%s, mapped=%s",
335 sampler_type, total_samples, consumed_samples, micro_batch_size, dp_rank, dp_world_size, drop_last,
336 data_sharding, resolved_mapping is not None,
337 )
338 sampler_options = {
339 "total_samples": total_samples,
340 "consumed_samples": consumed_samples,
341 "micro_batch_size": micro_batch_size,
342 "global_batch_size": global_batch_size,
343 "dp_rank": dp_rank,
344 "dp_world_size": dp_world_size,
345 "drop_last": drop_last,
346 "index_mapping": resolved_mapping,
347 }
349 if sampler_type == "single":
350 # Scenario 1 — single without a mapping: sequential Dataset indices.
351 # Scenario 2 — single + data_rearrange_map: the same logical sequence
352 # is resolved through an in-memory mapping or mapping checkpoint.
353 batch_sampler = _DatasetBatchSampler(**sampler_options)
354 return batch_sampler
356 if sampler_type == "cyclic":
357 # Scenario 3 — cyclic: use epoch-based randperm order,
358 # support data_sharding=True/False, and checkpoint consumed_samples.
359 if resolved_mapping is not None:
360 raise ValueError("cyclic sampling does not support a rearrangement map")
362 batch_sampler = _CyclicDatasetBatchSampler(
363 data_sharding=data_sharding,
364 seed=seed,
365 **sampler_options,
366 )
367 return batch_sampler
369 raise ValueError("sampler_type must be one of: single, cyclic")