Coverage for / home / jenkins / .local / lib / python3.10 / site-packages / hyper_parallel / distributed_data / api.py: 87%
299 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"""Public builder for sample-balanced distributed dynamic packing.
17Source modes and buffering
18--------------------------
19Online mode without ``batch_sampler`` requires ``external_step_source``.
20The source emits one complete local step of raw-sample bins. HP derives
21metadata, freezes the sample union,
22balances it across Data Constructors, and routes payloads when necessary.
23Metadata mode requires ``batch_sampler`` to define step/sample boundaries.
24It looks up ``metadata[index]`` before any Dataset read; target constructors
25then directly read their assigned indices from the shared Dataset and skip
26payload A2A. Metadata-only streaming selection is not supported.
28With ``batch_sampler``, native HP sampling owns step selection:
29``balance_window_size`` consecutive sampler yields per DP Constructor are
30planned together, stopping at optimizer and epoch boundaries. Each iterator
31yield still returns exactly one forward/backward round. The current window
32and at most one future window are double-buffered on Host. Device staging
33independently retains up to ``device_prefetch_depth`` future micro-batches;
34model-parallel input broadcast runs when the current batch is consumed.
35Checkpoints advance only at complete optimizer-step boundaries.
36Dataset outputs stay whole by default. Online ``split_fn`` decomposes each
37selected output into independent sequences; ``pack_fn`` rebuilds the planned
38bins without changing the original sampler window's sample membership.
39Set ``config.repack_sequences=True`` to use built-in text callbacks instead
40of supplying split_fn/pack_fn. ``config.pad_token_id`` controls source PAD
41trimming and new PAD; ``config.log_balance`` enables balancing reports.
43Checkpoint and distributed contracts
44-----------------------------------
45Native sampler checkpoint replay requires a deterministic Dataset stream
46for a given epoch; arbitrary worker-side RNG state is not captured. Metadata and
47Dataset lengths must agree across all Data Constructor ranks; the caller
48must ensure this consistency because the builder only validates local inputs. Metadata
49entries must describe deterministic, rank-independent Dataset outputs.
50An external_step_source uses pure DP, node-local
51communication and buffered H2D. Every step evaluates a candidate; sample
52exchange occurs only when its objective improves by more than
53min_balance_gain. HCCL is the default backend; Gloo transports control
54and CPU payloads on the host. HCCL transports control and payload
55tensors on the rank-local NPU.
56Checkpoint/resume remains available only on the sampler path, at complete
57optimizer-step boundaries. Save training state after the optimizer update;
58an incomplete step retains its starting cursor and replays in full.
59The caller must keep topology, configuration and custom policies
60consistent across ranks and compatible with restored checkpoints.
61"""
62# This distributed-data package is intentionally PyTorch-only.
64from __future__ import annotations
66import math
67from collections.abc import Callable, Mapping, Sequence
68from dataclasses import dataclass
69from functools import partial
70from multiprocessing.context import BaseContext
71from typing import Any, Literal
73import torch # pylint: disable=forbidden-backend-import
75from hyper_parallel.distributed_data.batch_sampler import (
76 BatchSamplerReader,
77 preserve_sample,
78 validate_native_sampler,
79)
80from hyper_parallel.distributed_data.data_constructor import (
81 PackingDataConstructor,
82 default_collate_fn,
83 default_pack_fn,
84)
85from hyper_parallel.distributed_data.device_prefetch import _create_device_prefetcher, _resolve_device
86from hyper_parallel.distributed_data.distributed_dataloader import DistributedDataLoader
87from hyper_parallel.distributed_data.balancing_algorithm import BalancingAlgorithm
88from hyper_parallel.distributed_data.cost_model import CausalAttentionCostModel, CostModel
89from hyper_parallel.distributed_data.model_batch import TensorFieldSpec
90from hyper_parallel.distributed_data.indexed_text import (
91 _validate_pad_token_id,
92 collate_indexed_text_sequences,
93 pack_indexed_text_samples,
94 split_packed_text_sample,
95)
96from hyper_parallel.distributed_data.packed_balancing import LocalBalancingDataLoader, build_local_balancing_dataloader
97from hyper_parallel.distributed_data.external_step import ExternalStepSource
98from hyper_parallel.distributed_data.planner import DynamicPackingPlanner, OversizedPolicy
99from hyper_parallel.distributed_data.schema import PackingConstraints, SampleMetadata
100from hyper_parallel.distributed_data.metadata import PlannedSampleLoader
101from hyper_parallel.distributed_data.dataset_reader import _validate_worker_options
102from hyper_parallel.distributed_data.topology import DataTopology
103from hyper_parallel.distributed_data.transport import (
104 DataPlaneTransport,
105 ModelParallelTransport,
106 create_data_groups,
107)
109_DEFAULT_DATALOADER_KWARGS = {
110 "num_workers": 0,
111 "pin_memory": False,
112 "prefetch_factor": None,
113 "persistent_workers": False,
114}
115_OPTIONAL_DATALOADER_KWARGS = frozenset({
116 "in_order",
117 "multiprocessing_context",
118 "pin_memory_device",
119 "timeout",
120 "worker_init_fn",
121})
122_SUPPORTED_DATALOADER_KWARGS = frozenset(_DEFAULT_DATALOADER_KWARGS) | _OPTIONAL_DATALOADER_KWARGS
123_MANAGED_DATALOADER_KWARGS = frozenset({
124 "batch_sampler",
125 "batch_size",
126 "collate_fn",
127 "dataset",
128 "drop_last",
129 "generator",
130 "sampler",
131 "shuffle",
132})
135@dataclass(frozen=True)
136class DistributedDatasetConfig:
137 """Configure sample planning, dynamic packing, and data-plane communication.
139 ``seq_len`` and ``local_batch_size`` are the only required sizing inputs.
140 One iterator yield contains ``local_batch_size`` packed sequences on each
141 DP rank. Native sampler global batch size defines optimizer boundaries.
143 Args:
144 seq_len: Maximum token count in one non-oversized packed sequence.
145 local_batch_size: Packed sequences constructed per DP rank and yield.
146 balance_window_size: Maximum native sampler micro-batches planned together.
147 Native prefetch buffers the current window and at most one future window,
148 each clipped at the optimizer-step boundary.
149 Defaults to one. Windows stop at optimizer and epoch boundaries;
150 gradient accumulation is inferred from the native sampler's global
151 batch size. External-step sources support only one.
152 device_prefetch_depth: Future native NPU/CUDA micro-batches copied ahead
153 on Constructors. Model peers receive each batch only at delivery.
154 Defaults to one. Zero stages only on
155 demand while retaining Host window prefetch. Must be a non-negative
156 integer and identical across ranks. Independent of balance_window_size;
157 ignored on CPU. External-step sources support only the default value.
158 dp_dim_names: Named mesh dimensions that define DP coordinates.
159 balance_group_size: Optional number of DP ranks in each independent
160 local-balancing group. ``None`` preserves launcher/node grouping.
161 This applies only to external-step
162 loading; native ``batch_sampler`` loading remains global.
163 dataset_reader_ranks: Optional Dataset Reader ranks. They read raw
164 samples online or metadata only in metadata mode. Defaults to the
165 Data Constructor ranks.
166 dataset_already_sharded: Legacy sharding configuration flag.
167 Native BatchSampler requires False because it owns DP slicing;
168 external sources manage their own partitioning.
169 planner_rank: Optional centralized Planner rank. Defaults to the lowest
170 Data Constructor rank.
171 oversized_policy: ``error`` by default; ``single`` explicitly permits
172 one oversized sample to occupy a bin alone.
173 drop_last: Whether to drop a tail with fewer than one sample per global
174 packing bin. Only ``True`` is supported in this first version.
175 seed: DataLoader worker seed. Sampling order is configured on the sampler
176 or external step producer.
177 min_balance_gain: Minimum relative improvement in the algorithm's
178 objective. A candidate is accepted only when its gain is strictly
179 greater; zero rejects equal or worse placements.
180 packing_budgets: Optional per-packed-sequence additive hard limits.
181 Each configured stage must
182 occur in every sample's packing_costs. These limits are independent
183 of cost-model balancing scores.
184 communication_backend: Backend for HP data-plane collectives. ``hccl``
185 is the default for NPU training and serializes control objects into
186 accelerator tensors. ``gloo`` keeps control and CPU payload
187 communication on Gloo. HCCL requires an NPU communication device.
188 metadata_mode: Whether metadata is precomputed and available by Dataset
189 index. ``False`` means the Dataset Reader loads samples and calls
190 ``metadata_fn`` during native sampler planning.
191 broadcast_plan: Whether one Planner computes and broadcasts each plan.
192 Defaults to ``True``. If ``False``, every data-plane rank computes
193 the plan from gathered metadata. Configuration and deterministic
194 cost/assignment policies must produce identical plans across ranks.
195 log_balance: Whether to emit INFO-level balancing reports. Defaults to
196 ``False``; enabling INFO logging alone does not enable these reports.
197 repack_sequences: Use built-in text splitting and packing on the native
198 online sampler path. Defaults to ``False``. Records must contain
199 input_ids, pre-shifted labels and explicit cu_seq_lens. Cannot be
200 combined with custom split_fn or pack_fn. Unless overridden, use
201 the text collator and causal Attention square cost including new PAD.
202 pad_token_id: Optional non-negative source PAD ID for built-in repacking.
203 Only trailing matching tokens with labels=-100 are removed. The same
204 ID fills new PAD. ``None`` preserves source tokens and fills with 0.
205 Has no effect when repack_sequences is false.
206 """
208 seq_len: int
209 local_batch_size: int
210 dp_dim_names: tuple[str, ...] | None = None
211 dataset_reader_ranks: tuple[int, ...] | None = None
212 dataset_already_sharded: bool = False
213 planner_rank: int | None = None
214 oversized_policy: OversizedPolicy = "error"
215 drop_last: bool = True
216 seed: int = 1234
217 min_balance_gain: float = 0.0
218 packing_budgets: dict[str, float] | None = None
219 communication_backend: Literal["gloo", "hccl"] = "hccl"
220 balance_group_size: int | None = None
221 metadata_mode: bool = False
222 broadcast_plan: bool = True
223 balance_window_size: int = 1
224 log_balance: bool = False
225 repack_sequences: bool = False
226 pad_token_id: int | None = None
227 device_prefetch_depth: int = 1
229 def __post_init__(self) -> None:
230 """Validate topology-independent configuration boundaries."""
231 self._validate_integer_fields()
232 self._validate_policy_and_flags()
233 if self.pad_token_id is not None:
234 _validate_pad_token_id(self.pad_token_id)
235 if self.repack_sequences and self.metadata_mode:
236 raise ValueError("repack_sequences requires online metadata (metadata_mode=False).")
237 self._validate_rank_tuple(self.dataset_reader_ranks, "dataset_reader_ranks")
238 self._validate_name_tuple(self.dp_dim_names, "dp_dim_names")
239 self._validate_planner_rank()
240 PackingConstraints(self.seq_len, self.oversized_policy, self.packing_budgets)
242 def _validate_integer_fields(self) -> None:
243 """Validate window sizes, batch dimensions and non-negative integer options."""
244 for name in ("seq_len", "local_batch_size", "balance_window_size"):
245 value = getattr(self, name)
246 if not isinstance(value, int) or isinstance(value, bool) or value < 1:
247 raise ValueError(f"{name} must be a positive integer, but got {value!r}.")
248 for name in ("seed", "device_prefetch_depth"):
249 value = getattr(self, name)
250 if not isinstance(value, int) or isinstance(value, bool) or value < 0:
251 raise ValueError(f"{name} must be a non-negative integer, but got {value!r}.")
252 if (
253 self.balance_group_size is not None
254 and (
255 not isinstance(self.balance_group_size, int)
256 or isinstance(self.balance_group_size, bool)
257 or self.balance_group_size < 1
258 )
259 ):
260 raise ValueError("balance_group_size must be a positive integer or None.")
262 def _validate_policy_and_flags(self) -> None:
263 """Validate policy names, boolean flags and the relative-gain threshold."""
264 if self.oversized_policy not in ("error", "single"):
265 raise ValueError("oversized_policy must be 'error' or 'single'.")
266 for name in (
267 "dataset_already_sharded",
268 "drop_last",
269 "metadata_mode",
270 "broadcast_plan",
271 "log_balance",
272 "repack_sequences",
273 ):
274 if not isinstance(getattr(self, name), bool):
275 raise ValueError(f"{name} must be boolean.")
276 if (
277 not isinstance(self.min_balance_gain, (int, float))
278 or isinstance(self.min_balance_gain, bool)
279 or not math.isfinite(self.min_balance_gain)
280 or not 0.0 <= self.min_balance_gain < 1.0
281 ):
282 raise ValueError("min_balance_gain must be in [0, 1).")
283 if self.communication_backend not in ("gloo", "hccl"):
284 raise ValueError("communication_backend must be 'gloo' or 'hccl'.")
285 if not self.drop_last:
286 raise ValueError(
287 "Dynamic distributed packing currently requires drop_last=True so every DP rank receives the "
288 "same number of non-empty packing bins."
289 )
291 def _validate_planner_rank(self) -> None:
292 if self.planner_rank is not None and (
293 not isinstance(self.planner_rank, int)
294 or isinstance(self.planner_rank, bool)
295 or self.planner_rank < 0
296 ):
297 raise ValueError("planner_rank must be a non-negative integer or None.")
299 @staticmethod
300 def _validate_rank_tuple(value: tuple[int, ...] | None, name: str) -> None:
301 if value is None:
302 return
303 if not isinstance(value, tuple) or not value:
304 raise ValueError(f"{name} must be a non-empty tuple or None.")
305 if any(not isinstance(rank, int) or isinstance(rank, bool) or rank < 0 for rank in value):
306 raise ValueError(f"{name} must contain non-negative integer ranks.")
307 if len(value) != len(set(value)):
308 raise ValueError(f"{name} must not contain duplicate ranks.")
310 @staticmethod
311 def _validate_name_tuple(value: tuple[str, ...] | None, name: str) -> None:
312 if value is None:
313 return
314 if not isinstance(value, tuple) or not value or any(not isinstance(item, str) or not item for item in value):
315 raise ValueError(f"{name} must be a non-empty tuple of strings or None.")
316 if len(value) != len(set(value)):
317 raise ValueError(f"{name} must not contain duplicate names.")
320def _resolve_service_ranks(
321 topology: DataTopology,
322 config: DistributedDatasetConfig,
323) -> tuple[tuple[int, ...], int]:
324 """Resolve Reader and planner ownership within the root mesh."""
325 dataset_reader_ranks = config.dataset_reader_ranks or topology.constructor_ranks
326 unknown_ranks = set(dataset_reader_ranks) - set(topology.rank_list)
327 if unknown_ranks:
328 raise ValueError(f"dataset_reader_ranks are outside the root mesh: {sorted(unknown_ranks)}.")
329 planner_rank = min(topology.constructor_ranks) if config.planner_rank is None else config.planner_rank
330 data_plane_ranks = set(dataset_reader_ranks) | set(topology.constructor_ranks)
331 if planner_rank not in data_plane_ranks:
332 raise ValueError(f"planner_rank {planner_rank} must be a Dataset Reader or Data Constructor rank.")
333 return dataset_reader_ranks, planner_rank
336def _normalize_dataloader_kwargs(
337 dataloader_kwargs: Mapping[str, Any] | None,
338) -> dict[str, Any]:
339 """Resolve validated DataLoader execution options."""
340 if dataloader_kwargs is None:
341 supplied_options = {}
342 elif not isinstance(dataloader_kwargs, Mapping):
343 raise ValueError("dataloader_kwargs must be a mapping or None.")
344 else:
345 try:
346 supplied_options = dict(dataloader_kwargs)
347 except Exception as exc:
348 raise ValueError(f"dataloader_kwargs cannot be copied: {exc}") from exc
349 _validate_dataloader_option_keys(supplied_options)
350 options = dict(_DEFAULT_DATALOADER_KWARGS)
351 options.update(supplied_options)
352 _validate_worker_options(options)
353 _validate_dataloader_timeout(options)
354 worker_init_fn = options.get("worker_init_fn")
355 if worker_init_fn is not None and not callable(worker_init_fn):
356 raise ValueError("dataloader_kwargs worker_init_fn must be callable or None.")
357 _validate_multiprocessing_context(options)
358 pin_memory_device = options.get("pin_memory_device", "")
359 if not isinstance(pin_memory_device, str):
360 raise ValueError("dataloader_kwargs pin_memory_device must be a string.")
361 return options
364def _validate_dataloader_option_keys(supplied_options: Mapping[str, Any]) -> None:
365 """Reject unsupported keys and overrides of internally managed DataLoader options."""
366 non_string_keys = [key for key in supplied_options if not isinstance(key, str)]
367 if non_string_keys:
368 raise ValueError(f"dataloader_kwargs keys must be strings, but got {non_string_keys!r}.")
369 managed_keys = sorted(set(supplied_options) & _MANAGED_DATALOADER_KWARGS)
370 if managed_keys:
371 raise ValueError(
372 f"dataloader_kwargs cannot override distributed sampling, batching, collation, or deterministic "
373 f"worker-seeding options "
374 f"{managed_keys}."
375 )
376 unsupported_keys = sorted(set(supplied_options) - _SUPPORTED_DATALOADER_KWARGS)
377 if unsupported_keys:
378 raise ValueError(f"dataloader_kwargs contains unsupported options {unsupported_keys}.")
381def _validate_dataloader_timeout(options: Mapping[str, Any]) -> None:
382 """Require a finite timeout compatible with the configured worker count."""
383 num_workers = options["num_workers"]
384 timeout = options.get("timeout", 0)
385 try:
386 normalized_timeout = float(timeout)
387 except (TypeError, ValueError, OverflowError):
388 normalized_timeout = math.nan
389 if (
390 not isinstance(timeout, (int, float))
391 or isinstance(timeout, bool)
392 or not math.isfinite(normalized_timeout)
393 or normalized_timeout < 0
394 ):
395 raise ValueError(f"dataloader_kwargs timeout must be finite and non-negative, but got {timeout!r}.")
396 if num_workers == 0 and timeout != 0:
397 raise ValueError("dataloader_kwargs timeout must be zero when num_workers=0.")
400def _validate_multiprocessing_context(options: Mapping[str, Any]) -> None:
401 """Check that the worker context uses a supported multiprocessing start method."""
402 multiprocessing_context = options.get("multiprocessing_context")
403 if multiprocessing_context is None:
404 return
405 if options["num_workers"] == 0:
406 raise ValueError("dataloader_kwargs multiprocessing_context requires num_workers > 0.")
407 if not isinstance(multiprocessing_context, (str, BaseContext)):
408 raise ValueError("dataloader_kwargs multiprocessing_context is invalid.")
409 start_method = multiprocessing_context
410 if isinstance(multiprocessing_context, BaseContext):
411 try:
412 start_method = multiprocessing_context.get_start_method()
413 except Exception as exc:
414 raise ValueError(
415 f"dataloader_kwargs multiprocessing_context cannot report its start method: {exc}"
416 ) from exc
417 if start_method not in torch.multiprocessing.get_all_start_methods():
418 raise ValueError(f"dataloader_kwargs multiprocessing_context={start_method!r} is not supported.")
421def _normalize_communication_device(communication_device: Any) -> torch.device | None:
422 """Normalize a rank-local A2A device."""
423 if communication_device is None:
424 return None
425 try:
426 return torch.device(communication_device)
427 except Exception as exc:
428 raise ValueError(f"communication_device is invalid: {communication_device!r}.") from exc
431class _BuildState:
432 """Validated rank-local components needed to assemble a distributed DataLoader."""
434 def __init__(
435 self,
436 dataset: Any | None,
437 mesh: Any,
438 config: DistributedDatasetConfig,
439 metadata_fn: Callable[[Any], SampleMetadata] | None,
440 metadata: Sequence[SampleMetadata] | None,
441 dataloader_kwargs: Mapping[str, Any] | None,
442 pack_fn: Callable[[Sequence[Any], int], Any] | None,
443 collate_fn: Callable[[Sequence[Any]], Any] | None,
444 communication_device: Any,
445 batch_sampler: Any = None,
446 *,
447 split_fn: Callable[[Any], Sequence[Any]] | None = None,
448 model_config: Any = None,
449 cost_model: CostModel | None = None,
450 balancing_algorithm: BalancingAlgorithm | None = None,
451 ) -> None:
452 """Validate inputs and construct all rank-local components.
454 Args:
455 dataset: Mapping Dataset on each Reader rank.
456 mesh: Named root device mesh.
457 config: Explicit loading mode, topology and packing options.
458 metadata_fn: Derives metadata from loaded samples in online mode.
459 metadata: Indexed precomputed metadata in metadata mode.
460 dataloader_kwargs: Rank-local sample loading options.
461 pack_fn: Optional packed-sequence constructor.
462 collate_fn: Optional local-batch collator.
463 communication_device: Rank-local control and payload device.
464 batch_sampler: Native sampler defining each local batch.
465 split_fn: Optional decomposition of each selected record into independent sequences.
466 model_config: Backbone dimensions for the default cost model.
467 cost_model: Optional workload estimator.
468 balancing_algorithm: Optional assignment and objective policy.
469 """
470 if not isinstance(config, DistributedDatasetConfig):
471 raise ValueError(f"config must be DistributedDatasetConfig, but got {type(config)}.")
472 _validate_builder_callbacks(metadata_fn, metadata, pack_fn, collate_fn)
473 split_fn, pack_fn, collate_fn = _resolve_repacking_callbacks(
474 config, batch_sampler, split_fn, pack_fn, collate_fn,
475 )
476 if split_fn is not None and (not callable(split_fn) or config.metadata_mode):
477 raise ValueError("split_fn requires online metadata and must be callable.")
478 normalized_options = _normalize_dataloader_kwargs(dataloader_kwargs)
479 loader_options = {name: normalized_options[name] for name in _DEFAULT_DATALOADER_KWARGS}
480 loader_options["dataloader_kwargs"] = {
481 name: value
482 for name, value in normalized_options.items()
483 if name not in _DEFAULT_DATALOADER_KWARGS
484 }
485 self.communication_device = _normalize_communication_device(communication_device)
486 if batch_sampler is None:
487 if config.metadata_mode:
488 raise ValueError("Metadata mode requires batch_sampler to define step/sample boundaries.")
489 raise ValueError("Online mode requires batch_sampler or external_step_source to define complete steps.")
490 effective_pack_fn, effective_collate_fn = _resolve_constructor_callbacks(
491 pack_fn, collate_fn, split_fn,
492 )
493 self.topology = DataTopology.from_mesh(mesh, dp_dim_names=config.dp_dim_names)
494 self.dataset_reader_ranks, self.planner_rank = _resolve_service_ranks(self.topology, config)
495 self.dataset_reader: BatchSamplerReader | None = None
496 self.metadata_reader: BatchSamplerReader | None = None
497 self.direct_sample_loader: PlannedSampleLoader | None = None
498 if config.metadata_mode:
499 if metadata_fn is not None:
500 raise ValueError("metadata_mode=True requires precomputed metadata, not metadata_fn.")
501 if metadata is None:
502 raise ValueError("metadata_mode=True requires metadata.")
503 elif metadata is not None:
504 raise ValueError("metadata_mode=False requires metadata_fn, not precomputed metadata.")
505 elif metadata_fn is None:
506 raise ValueError("metadata_mode=False requires metadata_fn.")
507 self._configure_batch_sampler_sources(
508 dataset, metadata_fn, metadata, config, batch_sampler, loader_options, split_fn,
509 )
510 self.planner, self.constructor = self._build_components(
511 config,
512 effective_pack_fn,
513 effective_collate_fn,
514 model_config,
515 cost_model,
516 balancing_algorithm,
517 )
519 @property
520 def is_reader(self) -> bool:
521 """Return whether this rank owns a Dataset Reader."""
522 return self.topology.global_rank in self.dataset_reader_ranks
524 def _configure_batch_sampler_sources(
525 self,
526 dataset: Any,
527 metadata_fn: Callable[[Any], SampleMetadata] | None,
528 metadata: Sequence[SampleMetadata] | None,
529 config: DistributedDatasetConfig,
530 batch_sampler: Any,
531 loader_options: dict[str, Any],
532 split_fn: Callable[[Any], Sequence[Any]] | None,
533 ) -> None:
534 """Use native DP sampler owners as Readers without applying a second stride."""
535 if self.dataset_reader_ranks != self.topology.constructor_ranks:
536 raise ValueError("Native batch_sampler currently requires Dataset Readers on the Data Constructor ranks.")
537 if config.dataset_already_sharded:
538 raise ValueError(
539 "Native batch_sampler owns DP slicing; disable dataset_already_sharded."
540 )
541 validate_native_sampler(
542 batch_sampler, data_rank=self.topology.data_rank,
543 dp_size=self.topology.data_parallel_size, local_batch_size=config.local_batch_size,
544 )
545 if not self.is_reader:
546 return
547 if dataset is None:
548 raise ValueError("Native batch_sampler requires a mapping Dataset on each Data Constructor.")
549 sample_loader = PlannedSampleLoader(dataset, seed=config.seed, **loader_options)
550 sample_loader.set_epoch(batch_sampler.epoch)
551 if config.metadata_mode:
552 if not callable(getattr(metadata, "__getitem__", None)) or len(metadata) != len(dataset):
553 raise ValueError("Native batch_sampler metadata must align with the mapping Dataset.")
554 self.direct_sample_loader = sample_loader
555 reader = BatchSamplerReader(
556 batch_sampler, reader_rank=self.topology.global_rank,
557 metadata=metadata, metadata_fn=metadata_fn,
558 sample_loader=None if config.metadata_mode else sample_loader,
559 split_fn=split_fn,
560 )
561 if config.metadata_mode:
562 self.metadata_reader = reader
563 else:
564 self.dataset_reader = reader
566 def _build_components(
567 self,
568 config: DistributedDatasetConfig,
569 effective_pack_fn: Callable[[Sequence[Any], int], Any],
570 effective_collate_fn: Callable[[Sequence[Any]], Any],
571 model_config: Any,
572 cost_model: CostModel | None,
573 balancing_algorithm: BalancingAlgorithm | None,
574 ) -> tuple[DynamicPackingPlanner, PackingDataConstructor]:
575 """Create the planner and data constructor."""
576 if cost_model is None and config.repack_sequences:
577 cost_model = CausalAttentionCostModel()
578 planner = DynamicPackingPlanner(
579 data_parallel_size=self.topology.data_parallel_size,
580 seq_len=config.seq_len,
581 local_batch_size=config.local_batch_size,
582 oversized_policy=config.oversized_policy,
583 packing_budgets=config.packing_budgets,
584 min_balance_gain=config.min_balance_gain,
585 model_config=model_config,
586 cost_model=cost_model,
587 balancing_algorithm=balancing_algorithm,
588 )
589 return planner, PackingDataConstructor(effective_pack_fn, effective_collate_fn, seq_len=config.seq_len)
592def _validate_builder_callbacks(
593 metadata_fn: Callable[[Any], SampleMetadata] | None,
594 metadata: Sequence[SampleMetadata] | None,
595 pack_fn: Callable[[Sequence[Any], int], Any] | None,
596 collate_fn: Callable[[Sequence[Any]], Any] | None,
597) -> None:
598 """Validate callback types and mutually exclusive metadata providers."""
599 if metadata_fn is not None and not callable(metadata_fn):
600 raise ValueError("metadata_fn must be callable or None.")
601 if metadata_fn is not None and metadata is not None:
602 raise ValueError("Provide either metadata or metadata_fn, but not both.")
603 if pack_fn is not None and not callable(pack_fn):
604 raise ValueError("pack_fn must be callable or None.")
605 if collate_fn is not None and not callable(collate_fn):
606 raise ValueError("collate_fn must be callable or None.")
609def _resolve_repacking_callbacks(
610 config: DistributedDatasetConfig,
611 batch_sampler: Any,
612 split_fn: Callable[[Any], Sequence[Any]] | None,
613 pack_fn: Callable[[Sequence[Any], int], Any] | None,
614 collate_fn: Callable[[Sequence[Any]], Any] | None,
615) -> tuple[
616 Callable[[Any], Sequence[Any]] | None,
617 Callable[[Sequence[Any], int], Any] | None,
618 Callable[[Sequence[Any]], Any] | None,
619]:
620 """Select built-in text callbacks without mixing packing contracts."""
621 if not config.repack_sequences:
622 return split_fn, pack_fn, collate_fn
623 if batch_sampler is None:
624 raise ValueError("repack_sequences requires a native batch_sampler.")
625 if split_fn is not None or pack_fn is not None:
626 raise ValueError("repack_sequences cannot be combined with custom split_fn or pack_fn.")
627 return (
628 partial(split_packed_text_sample, pad_token_id=config.pad_token_id),
629 partial(pack_indexed_text_samples, pad_token_id=0 if config.pad_token_id is None else config.pad_token_id),
630 collate_indexed_text_sequences if collate_fn is None else collate_fn,
631 )
634def _resolve_constructor_callbacks(
635 pack_fn: Callable[[Sequence[Any], int], Any] | None,
636 collate_fn: Callable[[Sequence[Any]], Any] | None,
637 split_fn: Callable[[Any], Sequence[Any]] | None = None,
638) -> tuple[Callable[[Sequence[Any], int], Any], Callable[[Sequence[Any]], Any]]:
639 """Resolve constructor callbacks for a native BatchSampler source."""
640 if split_fn is not None and pack_fn is None:
641 raise ValueError("Native split_fn requires pack_fn to rebuild the planned sequences.")
642 if split_fn is None and pack_fn is not None:
643 raise ValueError("Native batch_sampler preserves Dataset outputs; pack_fn must be omitted.")
644 effective_pack_fn = preserve_sample if split_fn is None else pack_fn
645 effective_collate_fn = default_collate_fn if collate_fn is None else collate_fn
646 return effective_pack_fn, effective_collate_fn
649def build_distributed_dataloader(
650 dataset: Any | None,
651 mesh: Any,
652 config: DistributedDatasetConfig,
653 *,
654 metadata_fn: Callable[[Any], SampleMetadata] | None = None,
655 metadata: Sequence[SampleMetadata] | None = None,
656 dataloader_kwargs: Mapping[str, Any] | None = None,
657 pack_fn: Callable[[Sequence[Any], int], Any] | None = None,
658 split_fn: Callable[[Any], Sequence[Any]] | None = None,
659 collate_fn: Callable[[Sequence[Any]], Any] | None = None,
660 device: Any = None,
661 model_batch_spec: Any = None,
662 batch_sampler: Any = None,
663 external_step_source: ExternalStepSource | None = None,
664 model_config: Any = None,
665 cost_model: CostModel | None = None,
666 balancing_algorithm: BalancingAlgorithm | None = None,
667) -> DistributedDataLoader | LocalBalancingDataLoader:
668 """Build a sample-balanced loader using the source and checkpoint contracts above.
670 Args:
671 dataset: Shared mapping Dataset on native Data Constructors. External sources
672 own loading and may pass None; other ranks may also pass None.
673 mesh: Named root HyperParallel or native DeviceMesh.
674 config: Packing, repacking, logging, service-rank and communication options.
675 metadata_fn: Derive metadata from each native Dataset output or external raw
676 sample. With splitting, receives each independent sequence instead.
677 metadata: Explicit indexable metadata on native Dataset Reader ranks in metadata
678 mode. Must implement __len__ and __getitem__; metadata[index] must describe
679 the corresponding deterministic, rank-independent dataset[index].
680 dataloader_kwargs: Native worker options: num_workers=0, pin_memory=False,
681 prefetch_factor=None, persistent_workers=False by default. Sampling,
682 batching, shuffling and DataLoader collation cannot be overridden.
683 pack_fn: Build one packed sequence; the external default preserves a bin as a
684 raw-sample tuple. Omit when config.repack_sequences selects text packing.
685 split_fn: Online native callback returning independent sequences in original
686 order. Requires pack_fn. Each original record must fit one bin and its
687 sequences retain their fallback grouping. Preserve tokens, targets and
688 attention isolation; contextual slices are not independent. Omit with
689 config.repack_sequences; unsupported in metadata mode.
690 collate_fn: Collate local_batch_size planned bins. Defaults to text collation
691 with config.repack_sequences, otherwise a tuple; an explicit callback wins.
692 device: Rank-local training device. HCCL uses its NPU for control, payload and
693 final H2D. Gloo communicates on Host and defaults to CPU; pass an explicit
694 NPU/CUDA device to enable H2D with Gloo.
695 model_batch_spec: Identical layout on every MP rank: nested dict/list/tuple
696 leaves are TensorFieldSpec or static scalars. Dynamic values must be tensors;
697 fields and dtypes are not broadcast. Defaults to collate_fn.model_batch_spec
698 or the built-in text layout. Required for custom MP batches, unnecessary for DP.
699 batch_sampler: Native HP BatchSampler supplied on every rank; only Constructors
700 advance it. Readers must coincide with Constructors and drop_last must be
701 true. Dataset outputs stay whole unless split_fn/pack_fn or online text
702 repacking is enabled. Metadata indexes outputs, not underlying documents.
703 Checkpoint via this loader, not the sampler's speculative cursor.
704 external_step_source: Online pure-DP iterable emitting one complete local step
705 of local_batch_size raw-sample bins. Optional set_epoch is forwarded;
706 exhaustion ends iteration. The caller owns recovery and step limits.
707 model_config: Backbone dimensions for the default cost model when repacking
708 is disabled and cost_model is omitted.
709 cost_model: Optional workload callback. Text repacking defaults to
710 CausalAttentionCostModel including new PAD; otherwise DefaultCostModel.
711 balancing_algorithm: Optional assignment/objective policy receiving scored
712 samples. HyperParallel enforces capacity and the gain gate.
714 Returns:
715 Collective iterator yielding one constructed local batch per micro-batch.
716 """
717 if external_step_source is not None:
718 if model_batch_spec is not None:
719 raise ValueError("model_batch_spec is only supported with a native batch_sampler.")
720 if split_fn is not None:
721 raise ValueError("split_fn is only supported with a native batch_sampler.")
722 if config.metadata_mode:
723 raise ValueError("external_step_source requires metadata_mode=False for online metadata.")
724 if batch_sampler is not None or metadata is not None:
725 raise ValueError("external_step_source cannot be combined with a sampler or metadata sequence.")
726 if dataloader_kwargs:
727 raise ValueError("Configure source worker options on external_step_source, not the balancing wrapper.")
728 return build_local_balancing_dataloader(
729 external_step_source, mesh, config, metadata_fn=metadata_fn,
730 pack_fn=default_pack_fn if pack_fn is None else pack_fn,
731 collate_fn=default_collate_fn if collate_fn is None else collate_fn,
732 model_config=model_config, cost_model=cost_model,
733 balancing_algorithm=balancing_algorithm, device=device,
734 )
735 return _build_distributed_dataloader_impl(
736 dataset, mesh, config, metadata_fn=metadata_fn, metadata=metadata,
737 dataloader_kwargs=dataloader_kwargs, pack_fn=pack_fn, collate_fn=collate_fn,
738 split_fn=split_fn, communication_device=device, model_batch_spec=model_batch_spec,
739 batch_sampler=batch_sampler, model_config=model_config,
740 cost_model=cost_model, balancing_algorithm=balancing_algorithm,
741 )
744def _build_distributed_dataloader_impl(
745 dataset: Any | None,
746 mesh: Any,
747 config: DistributedDatasetConfig,
748 *,
749 metadata_fn: Callable[[Any], SampleMetadata] | None,
750 metadata: Sequence[SampleMetadata] | None,
751 dataloader_kwargs: Mapping[str, Any] | None,
752 pack_fn: Callable[[Sequence[Any], int], Any] | None,
753 collate_fn: Callable[[Sequence[Any]], Any] | None,
754 communication_device: Any,
755 model_batch_spec: Any = None,
756 batch_sampler: Any = None,
757 split_fn: Callable[[Any], Sequence[Any]] | None = None,
758 model_config: Any = None,
759 cost_model: CostModel | None = None,
760 balancing_algorithm: BalancingAlgorithm | None = None,
761) -> DistributedDataLoader:
762 """Build native sampler readers, planning, transport and device staging for this rank."""
763 if communication_device is None:
764 communication_device = _resolve_device(communication_backend=getattr(config, "communication_backend", "hccl"))
765 if getattr(config, "balance_group_size", None) is not None:
766 raise ValueError("balance_group_size is only supported for local balancing paths.")
767 state = _BuildState(
768 dataset,
769 mesh,
770 config,
771 metadata_fn,
772 metadata,
773 dataloader_kwargs,
774 pack_fn,
775 collate_fn,
776 communication_device,
777 batch_sampler,
778 split_fn=split_fn,
779 model_config=model_config,
780 cost_model=cost_model,
781 balancing_algorithm=balancing_algorithm,
782 )
783 device = state.communication_device
784 # Validate device/stream support on every rank before any group is created.
785 device_prefetch = _create_device_prefetcher(device)
786 communication_device = device if config.communication_backend == "hccl" else None
787 if model_batch_spec is None:
788 model_batch_spec = getattr(collate_fn, "model_batch_spec", None)
789 if model_batch_spec is None and (
790 collate_fn is collate_indexed_text_sequences or (config.repack_sequences and collate_fn is None)
791 ):
792 shape = (config.local_batch_size, config.seq_len)
793 model_batch_spec = {
794 "input_ids": TensorFieldSpec(shape, torch.int64),
795 "labels": TensorFieldSpec(shape, torch.int64),
796 "cu_seq_lens": TensorFieldSpec((None,), torch.int32),
797 }
798 groups = create_data_groups(
799 state.topology,
800 state.dataset_reader_ranks,
801 state.planner_rank,
802 cpu_backend=config.communication_backend,
803 payload_backend=config.communication_backend,
804 communication_device=communication_device,
805 enable_payload_exchange=not config.metadata_mode,
806 model_device=device,
807 )
809 accumulation_steps = batch_sampler.global_batch_size // batch_sampler.global_micro_batch_size
810 return DistributedDataLoader(
811 accumulation_steps=accumulation_steps,
812 balance_window_size=config.balance_window_size,
813 device_prefetch_depth=config.device_prefetch_depth,
814 device=device,
815 initial_epoch=batch_sampler.epoch,
816 initial_consumed_samples=batch_sampler.consumed_samples,
817 topology=state.topology,
818 dataset_reader_ranks=state.dataset_reader_ranks,
819 dataset_reader=state.dataset_reader,
820 metadata_reader=state.metadata_reader,
821 direct_sample_loader=state.direct_sample_loader,
822 metadata_mode=config.metadata_mode,
823 broadcast_plan=config.broadcast_plan,
824 log_balance=config.log_balance,
825 planner=state.planner,
826 data_constructor=state.constructor,
827 data_plane=DataPlaneTransport(
828 groups,
829 state.topology.global_rank,
830 communication_device=communication_device,
831 ),
832 model_transport=ModelParallelTransport(
833 state.topology,
834 groups,
835 model_batch_spec=model_batch_spec,
836 device=device,
837 communication_device=communication_device,
838 ),
839 device_prefetch=device_prefetch if state.topology.is_constructor else None,
840 )
843__all__ = ["DistributedDatasetConfig", "build_distributed_dataloader"]