Coverage for / home / jenkins / .local / lib / python3.10 / site-packages / hyper_parallel / data / batching / get_batch.py: 91%
209 statements
« prev ^ index » next coverage.py v7.13.1, created at 2026-11-15 07:18 +0800
1# Copyright 2026 Huawei Technologies Co., Ltd
2#
3# Licensed under the Apache License, Version 2.0 (the "License");
4# you may not use this file except in compliance with the License.
5# You may obtain a copy of the License at
6#
7# http://www.apache.org/licenses/LICENSE-2.0
8#
9# Unless required by applicable law or agreed to in writing, software
10# distributed under the License is distributed on an "AS IS" BASIS,
11# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12# See the License for the specific language governing permissions and
13# limitations under the License.
14# ============================================================================
15"""Build one model and loss batch from a DataLoader iterator."""
16# pylint: disable=forbidden-backend-import
18from __future__ import annotations
20from collections.abc import Mapping
21from typing import Any
23import torch
25from hyper_parallel.data.batching.runtime_input import (
26 AttentionRuntime,
27 IndexedBoundaryResolver,
28 OnlineBoundaryResolver,
29 RuntimeInputAdapter,
30)
31from hyper_parallel.data.constants import (
32 OMNI_CP_TOKEN_FIELDS,
33 OMNI_INTERNAL_FIELDS,
34 OMNI_LOSS_INPUT_FIELDS,
35)
36from hyper_parallel.data.dataset_logging import get_dataset_logger
37from hyper_parallel.data.indexed.indexed_data_config import resolve_dataset_type
38from hyper_parallel.data.indexed.indexed_pretrain_dataset import SFTFromMRDataset
39from hyper_parallel.data.parallel import (
40 CPBatchSharder,
41 DataLoaderParallelContext,
42 TPBatchBroadcaster,
43 create_dataloader_parallel_context,
44)
46logger = get_dataset_logger(__name__)
49class OmniParallelBatch:
50 """Validate collate-aligned Omni batches, shard over CP, and broadcast over TP."""
52 def __init__(
53 self,
54 *,
55 mesh_context: Any,
56 device: Any,
57 pp_shared_data: bool = False,
58 runtime_input_adapter: RuntimeInputAdapter | None = None,
59 dataloader_context: DataLoaderParallelContext | None = None,
60 data_config: Mapping[str, Any] | None = None,
61 ) -> None:
62 """Create the DataLoader ownership and TP/CP transport context.
64 Args:
65 mesh_context: Trainer mesh exposing TP, CP, and PP sizes.
66 device: Destination model device.
67 pp_shared_data: Whether pipeline stages share the source batch.
68 runtime_input_adapter: Provider of additional forward inputs.
69 data_config: Dataset cache settings used by DataLoader ownership.
70 dataloader_context: Shared ownership and transport policy from the Trainer.
71 """
72 # if encoder_dp:
73 # raise NotImplementedError("OmniParallelBatch encoder_dp image-bucket distribution is not implemented")
75 if runtime_input_adapter is not None and not isinstance(runtime_input_adapter, RuntimeInputAdapter):
76 raise TypeError("OmniParallelBatch runtime_input_adapter must be a RuntimeInputAdapter")
78 if int(getattr(mesh_context, "pp_size", 1)) != 1 or pp_shared_data:
79 raise NotImplementedError("OmniParallelBatch does not support pipeline parallelism or pp_shared_data")
81 source_config = data_config or {}
82 self.parallel_context = dataloader_context or create_dataloader_parallel_context(
83 mesh_context,
84 data_index_cache=bool(source_config.get("data_index_cache", False)),
85 shared_storage=not bool(source_config.get("no_shared_storage", False)),
86 )
87 if not self.parallel_context.distributed_enabled and (
88 int(getattr(mesh_context, "tp_size", 1)) > 1 or int(getattr(mesh_context, "cp_size", 1)) > 1
89 ):
90 raise ValueError("OmniParallelBatch requires an initialized process group for TP or CP")
92 self.device = device
93 self.sequence_parallel = bool(mesh_context.sequence_parallel)
94 self.runtime_input_adapter = runtime_input_adapter
95 self.cp_sharder = CPBatchSharder(self.parallel_context, token_fields=OMNI_CP_TOKEN_FIELDS)
96 self.tp_broadcaster = TPBatchBroadcaster(self.parallel_context, device)
97 self.batch_transport = None
98 self._batch_index = 0
100 def __call__(
101 self,
102 data_iterator: Any,
103 *,
104 external_batch: Mapping[str, Any] | None = None,
105 ) -> tuple[dict[str, Any], dict[str, Any]]:
106 """Read on source owners, slice for CP, and distribute only if needed."""
107 self._batch_index += 1
108 local_batch = None
109 if self.parallel_context.build_on_rank():
110 logger.debug("Omni get_batch: batch=%d source_read_begin", self._batch_index, enabled=True)
111 source_batch = external_batch if external_batch is not None else next(data_iterator)
112 if not isinstance(source_batch, Mapping):
113 raise TypeError("Omni DataLoader must yield a mapping batch")
115 normalized_batch = self._normalize_source_batch(source_batch)
116 local_batch = self.cp_sharder.shard(normalized_batch)
117 logger.debug(
118 "Omni get_batch: batch=%d source_ready global_shape=%s local_shape=%s tp=%d cp=%d cp_rank=%d",
119 self._batch_index, tuple(normalized_batch["input_ids"].shape),
120 tuple(local_batch["input_ids"].shape), self.parallel_context.tp_world_size,
121 self.parallel_context.cp_world_size, self.parallel_context.cp_rank, enabled=True,
122 )
124 transport_inputs = {}
125 if self.batch_transport is None:
126 device_batch = self.tp_broadcaster.broadcast(local_batch, broadcast_all_fields=True)
127 else:
128 device_batch, transport_inputs = self.batch_transport(local_batch, self.tp_broadcaster)
129 runtime_inputs = self._build_runtime_inputs(device_batch)
130 model_inputs, loss_inputs = self._split_model_and_loss_inputs(device_batch)
131 if set(transport_inputs).intersection((*device_batch, *runtime_inputs)):
132 raise ValueError("batch transport cannot replace Omni model or runtime inputs")
133 model_inputs.update(runtime_inputs)
134 model_inputs.update(transport_inputs)
135 # Shapes are host metadata; inspecting device tensor values here can stall collectives.
136 logger.debug(
137 "Omni get_batch: batch=%d transport_ready fields=%s runtime_fields=%s",
138 self._batch_index,
139 {field: tuple(value.shape) for field, value in device_batch.items() if torch.is_tensor(value)},
140 list(runtime_inputs), enabled=True,
141 )
143 return model_inputs, loss_inputs
145 @staticmethod
146 def _normalize_source_batch(source_batch: Mapping[str, Any]) -> dict[str, Any]:
147 """Validate required fields and derive the default loss mask."""
148 batch = dict(source_batch)
149 if "input_ids" not in batch:
150 raise ValueError("Omni batch must contain input_ids")
152 if "labels" not in batch:
153 raise ValueError("Omni batch must contain labels")
155 if "loss_mask" not in batch:
156 batch["loss_mask"] = batch["labels"] >= 0
157 if "cu_seq_lens" in batch:
158 batch["cu_seq_lens"] = batch["cu_seq_lens"].reshape(-1)
159 return batch
161 def _build_runtime_inputs(self, batch: Mapping[str, Any]) -> dict[str, Any]:
162 """Call the configured model hook after TP/CP batch transport."""
163 if self.runtime_input_adapter is None:
164 return {}
166 runtime_inputs = self.runtime_input_adapter.build(
167 batch=batch, parallel_context=self.parallel_context,
168 options={"sequence_parallel": self.sequence_parallel},
169 )
170 collisions = set(batch).intersection(runtime_inputs)
171 if collisions:
172 raise ValueError(
173 "runtime inputs cannot replace Omni model inputs: "
174 f"{sorted(collisions)}"
175 )
177 return runtime_inputs
179 @staticmethod
180 def _split_model_and_loss_inputs(
181 batch: Mapping[str, Any],
182 ) -> tuple[dict[str, Any], dict[str, Any]]:
183 """Keep processor fields as model inputs and isolate internal metadata."""
184 model_inputs = {}
185 loss_inputs = {}
186 for field, value in batch.items():
187 if field not in OMNI_INTERNAL_FIELDS:
188 model_inputs[field] = value
189 if field in OMNI_LOSS_INPUT_FIELDS:
190 loss_inputs[field] = value
191 return model_inputs, loss_inputs
194class TextParallelBatch:
195 """Define the DataLoader-to-forward-batch processing boundary.
197 The runtime reads data only on TP rank zero at every CP coordinate, keeps
198 global sequence boundaries, shards token fields for CP, broadcasts the
199 local fields across TP, and then builds common attention inputs.
200 """
202 def __init__(
203 self,
204 mesh_context: Any,
205 device: Any,
206 tokenizer: Any,
207 data_config: Mapping[str, Any],
208 pp_shared_data: bool,
209 *,
210 source_type: str,
211 attention_mode: str = "dense",
212 causal: bool = True,
213 sliding_window: int | None = None,
214 reset_position_ids: bool = False,
215 reset_attention_mask: bool = False,
216 eod_mask_loss: bool = False,
217 runtime_input_adapter: RuntimeInputAdapter | None = None,
218 dataloader_context: DataLoaderParallelContext | None = None,
219 ) -> None:
220 """Initialize the batch runtime and its parallel execution context.
222 Args:
223 mesh_context: Trainer mesh used to create DataLoader parallel state.
224 device: Device receiving the DataLoader batch.
225 tokenizer: Tokenizer providing EOD and padding token semantics.
226 data_config: Dataset options used to build runtime LTR fields.
227 pp_shared_data: Whether pipeline stages share the prepared batch.
228 source_type: ``online``, ``indexed``, or ``packed_text`` batch contract.
229 Packed text provides input_ids, shifted labels, and explicit cu_seq_lens.
230 attention_mode: ``dense`` or ``compressed`` attention representation.
231 causal: Whether attention uses left-to-right causal semantics.
232 sliding_window: Sliding-window size, or ``None`` for full attention.
233 reset_position_ids: Whether positions restart at sequence boundaries.
234 reset_attention_mask: Whether EOD starts an independent attention sequence.
235 eod_mask_loss: Whether EOD tokens are excluded from the loss.
236 runtime_input_adapter: Optional model-owned forward-input extension.
237 dataloader_context: Shared ownership and transport policy from the Trainer.
238 """
239 # pylint: disable=too-many-locals
241 self.parallel_context = dataloader_context or create_dataloader_parallel_context(
242 mesh_context,
243 data_index_cache=bool(data_config.get("data_index_cache", False)),
244 shared_storage=not bool(data_config.get("no_shared_storage", False)),
245 )
246 self.device = device
247 self.tokenizer = tokenizer
248 self.data_config = dict(data_config)
249 self.source_type = source_type
250 dataset_spec = data_config.get("low_level_dataset")
251 self.is_sft_dataset = (source_type == "indexed" and dataset_spec is not None
252 and issubclass(resolve_dataset_type(dataset_spec), SFTFromMRDataset))
253 self.causal = causal
254 self.sliding_window = sliding_window
255 if runtime_input_adapter is not None and not isinstance(
256 runtime_input_adapter, RuntimeInputAdapter
257 ):
258 raise TypeError("TextParallelBatch runtime_input_adapter must be a RuntimeInputAdapter")
260 self.runtime_input_adapter = runtime_input_adapter
261 self._batch_flow_logged = False
263 self.boundary_resolver = self._build_boundary_resolver(source_type, tokenizer)
265 self.pp_shared_data = pp_shared_data
266 self.labels_are_shifted = bool(self.data_config.get("labels_are_shifted", True))
267 create_attention_mask = bool(
268 self.data_config.get("create_attention_mask_in_dataloader", attention_mode == "dense")
269 )
270 self.eod_mask_loss = eod_mask_loss or bool(self.data_config.get("eod_mask_loss", False))
271 self._validate_attention_inputs(attention_mode, create_attention_mask)
272 self.cp_sharder = CPBatchSharder(self.parallel_context)
273 self.tp_broadcaster = TPBatchBroadcaster(self.parallel_context, device)
275 self.reset_position_ids = (
276 reset_position_ids or bool(self.data_config.get("reset_position_ids", False))
277 or source_type == "packed_text"
278 )
279 resolved_reset_attention_mask = (
280 reset_attention_mask
281 or bool(self.data_config.get("reset_attention_mask", False))
282 or source_type in {"online", "packed_text"}
283 or self.is_sft_dataset
284 )
285 self.attention_runtime = AttentionRuntime(
286 mode=attention_mode,
287 create_mask=create_attention_mask,
288 reset_mask=resolved_reset_attention_mask,
289 sliding_window=sliding_window,
290 )
292 @staticmethod
293 def _build_boundary_resolver(source_type: str, tokenizer: Any) -> OnlineBoundaryResolver | IndexedBoundaryResolver:
294 """Select sequence-boundary handling for the source's input_ids batch."""
295 if source_type in {"online", "packed_text"}:
296 return OnlineBoundaryResolver()
297 if source_type == "indexed":
298 return IndexedBoundaryResolver(getattr(tokenizer, "eod", None))
299 raise ValueError(f"Unsupported batch source type: {source_type!r}")
301 def _validate_attention_inputs(self, attention_mode: str, create_attention_mask: bool) -> None:
302 """Require shifted labels and document isolation for prepared text records."""
303 if self.source_type == "packed_text":
304 if not self.labels_are_shifted:
305 raise ValueError("Packed text batches require shifted labels; labels_are_shifted must be True")
306 if not create_attention_mask or (attention_mode == "compressed" and self.runtime_input_adapter is None):
307 raise ValueError(
308 "Packed text batches require document-boundary attention: enable "
309 "create_attention_mask_in_dataloader and supply a runtime_input_adapter for compressed attention"
310 )
311 if self.is_sft_dataset:
312 if not self.labels_are_shifted or self.eod_mask_loss or not self.causal:
313 raise ValueError("Prepared records require shifted labels, causal attention and eod_mask_loss=False")
314 if (attention_mode == "compressed" or not create_attention_mask) and self.runtime_input_adapter is None:
315 raise ValueError("Prepared records without a dense mask require a boundary-aware runtime adapter")
317 def __call__(
318 self,
319 data_iterator: Any,
320 ) -> tuple[Mapping[str, Any], Mapping[str, Any]]:
321 """Build model and loss inputs for one forward-backward step.
323 Args:
324 data_iterator: Iterator yielding one collated Indexed or Online batch.
326 Returns:
327 Model inputs and loss inputs for one forward-backward step.
328 """
329 canonical_batch = self._read_source_batch(data_iterator)
330 cu_seq_lens = self._resolve_sequence_boundaries(canonical_batch)
331 cp_local_batch = self.cp_sharder.shard(canonical_batch)
332 parallel_batch = self.tp_broadcaster.broadcast(
333 cp_local_batch,
334 cu_seq_lens,
335 )
336 self._log_batch_flow(canonical_batch, parallel_batch)
338 position_ids = self._build_local_position_ids(
339 parallel_batch["input_ids"],
340 parallel_batch["cu_seq_lens"],
341 )
342 parallel_batch["position_ids"] = position_ids
344 loss_mask = self._build_loss_mask(parallel_batch)
345 parallel_batch["loss_mask"] = loss_mask
346 runtime_inputs = self._build_runtime_inputs(parallel_batch)
348 model_inputs, loss_inputs = self._split_model_and_loss_inputs(
349 parallel_batch,
350 runtime_inputs,
351 )
353 return model_inputs, loss_inputs
355 def _build_runtime_inputs(
356 self,
357 parallel_batch: Mapping[str, Any],
358 ) -> dict[str, Any]:
359 """Build common attention inputs plus a model-owned extension."""
360 runtime_inputs = self.attention_runtime.build(
361 batch=parallel_batch,
362 parallel_context=self.parallel_context,
363 )
364 if self.runtime_input_adapter is None:
365 return runtime_inputs
367 model_runtime_inputs = self.runtime_input_adapter.build(
368 batch=parallel_batch,
369 parallel_context=self.parallel_context,
370 options={
371 "source_type": self.source_type,
372 "attention_mode": self.attention_runtime.mode,
373 "causal": self.causal,
374 "sliding_window": self.sliding_window,
375 },
376 )
377 collisions = set(runtime_inputs).intersection(model_runtime_inputs)
378 if collisions:
379 raise ValueError(
380 "model runtime inputs cannot replace common attention inputs: "
381 f"{sorted(collisions)}"
382 )
384 runtime_inputs.update(model_runtime_inputs)
385 return runtime_inputs
387 def _log_batch_flow(
388 self,
389 canonical_batch: Mapping[str, Any] | None,
390 parallel_batch: Mapping[str, Any],
391 ) -> None:
392 """Log the resolved source and parallel batch shapes once."""
393 if not self._batch_flow_logged:
394 source_shape = None if canonical_batch is None else tuple(canonical_batch["input_ids"].shape)
395 local_shape = tuple(parallel_batch["input_ids"].shape)
396 num_boundaries = 0 if parallel_batch["cu_seq_lens"] is None else parallel_batch["cu_seq_lens"].numel()
397 logger.debug(
398 "Parallel batch flow: source=%s, tp_rank=%d/%d, cp_rank=%d/%d, "
399 "source_owner=%s, source_shape=%s, local_shape=%s, global_boundaries=%d",
400 self.source_type,
401 self.parallel_context.tp_rank,
402 self.parallel_context.tp_world_size,
403 self.parallel_context.cp_rank,
404 self.parallel_context.cp_world_size,
405 canonical_batch is not None,
406 source_shape,
407 local_shape,
408 num_boundaries,
409 enabled=True,
410 )
411 self._batch_flow_logged = True
413 def _read_source_batch(self, data_iterator: Any) -> Mapping[str, Any] | None:
414 """Advance the source only on ranks selected by the shared context."""
415 if self.parallel_context.build_on_rank():
416 source_batch = next(data_iterator)
417 else:
418 source_batch = None
420 return source_batch
422 def _resolve_sequence_boundaries(
423 self,
424 canonical_batch: Mapping[str, Any] | None,
425 ) -> Any:
426 """Resolve global leading-zero cumulative sequence boundaries."""
427 if canonical_batch is None:
428 cu_seq_lens = None
429 else:
430 cu_seq_lens = self.boundary_resolver.resolve(canonical_batch)
432 return cu_seq_lens
434 def _build_local_position_ids(
435 self,
436 input_ids: Any,
437 cu_seq_lens: Any,
438 ) -> Any:
439 """Build position IDs for the current CP slice on every TP rank.
441 Packed boundaries define the position semantics. CP only determines
442 which global sequence interval is materialized by this rank.
443 """
444 # pylint: disable=too-many-locals
445 batch_size, local_seq_len = input_ids.shape
446 cp_size = self.parallel_context.cp_world_size
447 global_seq_len = local_seq_len * cp_size
448 cp_rank = self.parallel_context.cp_rank
449 cp_seq_start = cp_rank * local_seq_len
450 cp_seq_end = cp_seq_start + local_seq_len
452 # This produces the same values as building global [B, S] positions
453 # and then taking the CP slice, without allocating the global tensor.
454 position_ids = torch.arange(
455 cp_seq_start, cp_seq_end, dtype=torch.int64, device=input_ids.device
456 )
457 local_position_ids = position_ids.unsqueeze(0).expand(batch_size, -1)
459 if self.reset_position_ids:
460 local_position_ids = local_position_ids.clone()
461 boundaries = cu_seq_lens.tolist()
462 for seq_start, seq_end in zip(boundaries[:-1], boundaries[1:]):
463 seq_len = seq_end - seq_start
464 batch_idx, seq_start_in_batch = divmod(seq_start, global_seq_len)
465 seq_end_in_batch = seq_start_in_batch + seq_len
466 # Intersect the sequence with the current CP slice.
467 slice_start = max(seq_start_in_batch, cp_seq_start)
468 slice_end = min(seq_end_in_batch, cp_seq_end)
469 if slice_start >= slice_end:
470 continue
472 local_start = slice_start - cp_seq_start
473 local_end = slice_end - cp_seq_start
474 local_position_ids[batch_idx, local_start:local_end] -= seq_start_in_batch
476 local_position_ids = local_position_ids.contiguous()
478 return local_position_ids
480 def _build_loss_mask(self, parallel_batch: Mapping[str, Any]) -> Any:
481 """Build the local loss mask from labels and input IDs."""
482 loss_mask = (parallel_batch["labels"] >= 0).to(dtype=torch.int64)
484 if self.eod_mask_loss:
485 eod_token_id = getattr(self.tokenizer, "eod", None)
486 loss_mask = loss_mask.masked_fill(parallel_batch["input_ids"] == eod_token_id, 0)
488 return loss_mask
490 def _split_model_and_loss_inputs(
491 self,
492 parallel_batch: Mapping[str, Any],
493 runtime_inputs: Mapping[str, Any] | None = None,
494 ) -> tuple[Mapping[str, Any], Mapping[str, Any]]:
495 """Split forward fields from loss and token-accounting fields."""
496 model_inputs = {
497 "input_ids": parallel_batch["input_ids"],
498 "labels": parallel_batch["labels"],
499 "position_ids": parallel_batch["position_ids"],
500 }
501 if self.labels_are_shifted:
502 model_inputs["shift_labels"] = parallel_batch["labels"]
504 runtime_inputs = runtime_inputs or {}
505 collisions = set(model_inputs).intersection(runtime_inputs)
506 if collisions:
507 raise ValueError(
508 "runtime inputs cannot replace framework-owned model inputs: "
509 f"{sorted(collisions)}"
510 )
512 model_inputs.update(runtime_inputs)
514 loss_inputs = {
515 "labels": parallel_batch["labels"],
516 "loss_mask": parallel_batch["loss_mask"],
517 }
518 if self.labels_are_shifted:
519 loss_inputs["shift_labels"] = parallel_batch["labels"]
521 return model_inputs, loss_inputs