Coverage for / home / jenkins / .local / lib / python3.10 / site-packages / hyper_parallel / data / batching / get_batch.py: 22%
194 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"""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 math import lcm
22from typing import Any
24import torch
26from hyper_parallel.data.batching.runtime_input import (
27 AttentionRuntime,
28 IndexedBoundaryResolver,
29 OnlineBoundaryResolver,
30 RuntimeInputAdapter,
31)
32from hyper_parallel.data.constants import (
33 OMNI_CP_PAD_VALUES,
34 OMNI_CP_TOKEN_FIELDS,
35 OMNI_INTERNAL_FIELDS,
36 OMNI_LOSS_INPUT_FIELDS,
37)
38from hyper_parallel.data.dataset_logging import get_dataset_logger
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 """Shard Omni token fields over CP and replicate modality fields 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 data_config: Mapping[str, Any] | None = None,
60 ) -> None:
61 """Create the DataLoader ownership and TP/CP transport context.
63 Args:
64 mesh_context: Trainer mesh exposing TP, CP, and PP sizes.
65 device: Destination model device.
66 pp_shared_data: Whether pipeline stages share the source batch.
67 runtime_input_adapter: Provider of additional forward inputs.
68 data_config: Dataset cache settings used by DataLoader ownership.
69 """
70 if runtime_input_adapter is not None and not isinstance(runtime_input_adapter, RuntimeInputAdapter):
71 raise TypeError("OmniParallelBatch runtime_input_adapter must be a RuntimeInputAdapter")
72 if int(getattr(mesh_context, "pp_size", 1)) != 1 or pp_shared_data:
73 raise NotImplementedError("OmniParallelBatch does not support pipeline parallelism or pp_shared_data")
74 source_config = data_config or {}
75 self.parallel_context = create_dataloader_parallel_context(
76 mesh_context,
77 data_index_cache=bool(source_config.get("data_index_cache", False)),
78 shared_storage=not bool(source_config.get("no_shared_storage", False)),
79 )
80 if not self.parallel_context.distributed_enabled and (
81 int(getattr(mesh_context, "tp_size", 1)) > 1 or int(getattr(mesh_context, "cp_size", 1)) > 1
82 ):
83 raise ValueError("OmniParallelBatch requires an initialized process group for TP or CP")
84 self.device = device
85 self.sequence_parallel = bool(mesh_context.sequence_parallel)
86 self.runtime_input_adapter = runtime_input_adapter
87 token_pad_values = {field: OMNI_CP_PAD_VALUES.get(field, 0) for field in OMNI_CP_TOKEN_FIELDS}
88 cp_size = int(mesh_context.cp_size)
89 tp_parts = int(mesh_context.tp_size) if mesh_context.sequence_parallel else 1
90 alignment = lcm(2 * cp_size, cp_size * tp_parts)
91 self.cp_sharder = CPBatchSharder(
92 self.parallel_context, token_pad_values=token_pad_values,
93 required_global_alignment=alignment,
94 )
95 self.tp_broadcaster = TPBatchBroadcaster(self.parallel_context, device)
97 def __call__(
98 self,
99 data_iterator: Any,
100 *,
101 external_batch: Mapping[str, Any] | None = None,
102 ) -> tuple[dict[str, Any], dict[str, Any]]:
103 """Read on TP0, slice token fields for CP, then broadcast within TP."""
104 local_batch = None
105 if self.parallel_context.build_on_rank():
106 source_batch = external_batch if external_batch is not None else next(data_iterator)
107 if not isinstance(source_batch, Mapping):
108 raise TypeError("Omni DataLoader must yield a mapping batch")
109 normalized_batch = self._normalize_source_batch(source_batch)
110 local_batch = self.cp_sharder.shard(normalized_batch)
112 device_batch = self.tp_broadcaster.broadcast(local_batch, broadcast_all_fields=True)
113 runtime_inputs = self._build_runtime_inputs(device_batch)
114 model_inputs, loss_inputs = self._split_model_and_loss_inputs(device_batch)
115 model_inputs.update(runtime_inputs)
117 return model_inputs, loss_inputs
119 @staticmethod
120 def _normalize_source_batch(source_batch: Mapping[str, Any]) -> dict[str, Any]:
121 """Validate required fields and derive the default loss mask."""
122 batch = dict(source_batch)
123 if "input_ids" not in batch:
124 raise ValueError("Omni batch must contain input_ids")
125 if "labels" not in batch:
126 raise ValueError("Omni batch must contain labels")
127 if "loss_mask" not in batch:
128 batch["loss_mask"] = batch["labels"] >= 0
129 if "cu_seq_lens" in batch:
130 batch["cu_seq_lens"] = batch["cu_seq_lens"].reshape(-1)
131 return batch
133 def _build_runtime_inputs(self, batch: Mapping[str, Any]) -> dict[str, Any]:
134 """Call the configured model hook after TP/CP batch transport."""
135 if self.runtime_input_adapter is None:
136 return {}
137 runtime_inputs = self.runtime_input_adapter.build(
138 batch=batch, parallel_context=self.parallel_context,
139 options={"sequence_parallel": self.sequence_parallel},
140 )
141 collisions = set(batch).intersection(runtime_inputs)
142 if collisions:
143 raise ValueError(
144 "runtime inputs cannot replace Omni model inputs: "
145 f"{sorted(collisions)}"
146 )
147 return runtime_inputs
149 @staticmethod
150 def _split_model_and_loss_inputs(
151 batch: Mapping[str, Any],
152 ) -> tuple[dict[str, Any], dict[str, Any]]:
153 """Keep processor fields as model inputs and isolate internal metadata."""
154 model_inputs = {}
155 loss_inputs = {}
156 for field, value in batch.items():
157 if field not in OMNI_INTERNAL_FIELDS:
158 model_inputs[field] = value
159 if field in OMNI_LOSS_INPUT_FIELDS:
160 loss_inputs[field] = value
161 return model_inputs, loss_inputs
164class TextParallelBatch:
165 """Define the DataLoader-to-forward-batch processing boundary.
167 The runtime reads data only on TP rank zero at every CP coordinate, keeps
168 global sequence boundaries, shards token fields for CP, broadcasts the
169 local fields across TP, and then builds common attention inputs.
170 """
172 def __init__(
173 self,
174 mesh_context: Any,
175 device: Any,
176 tokenizer: Any,
177 data_config: Mapping[str, Any],
178 pp_shared_data: bool,
179 *,
180 source_type: str,
181 attention_mode: str = "dense",
182 causal: bool = True,
183 sliding_window: int | None = None,
184 reset_position_ids: bool = False,
185 reset_attention_mask: bool = False,
186 eod_mask_loss: bool = False,
187 runtime_input_adapter: RuntimeInputAdapter | None = None,
188 ) -> None:
189 """Initialize the batch runtime and its parallel execution context.
191 Args:
192 mesh_context: Trainer mesh used to create DataLoader parallel state.
193 device: Device receiving the DataLoader batch.
194 tokenizer: Tokenizer providing EOD and padding token semantics.
195 data_config: Dataset options used to build runtime LTR fields.
196 pp_shared_data: Whether pipeline stages share the prepared batch.
197 source_type: ``online`` or ``indexed`` DataLoader batch contract.
198 attention_mode: ``dense`` or ``compressed`` attention representation.
199 causal: Whether attention uses left-to-right causal semantics.
200 sliding_window: Sliding-window size, or ``None`` for full attention.
201 reset_position_ids: Whether positions restart at sequence boundaries.
202 reset_attention_mask: Whether EOD starts an independent attention sequence.
203 eod_mask_loss: Whether EOD tokens are excluded from the loss.
204 runtime_input_adapter: Optional model-owned forward-input extension.
205 """
206 # pylint: disable=too-many-locals
208 self.parallel_context: DataLoaderParallelContext = create_dataloader_parallel_context(
209 mesh_context,
210 data_index_cache=bool(data_config.get("data_index_cache", False)),
211 shared_storage=not bool(data_config.get("no_shared_storage", False)),
212 )
213 self.device = device
214 self.tokenizer = tokenizer
215 self.data_config = dict(data_config)
216 self.source_type = source_type
217 self.causal = causal
218 self.sliding_window = sliding_window
219 if runtime_input_adapter is not None and not isinstance(
220 runtime_input_adapter, RuntimeInputAdapter
221 ):
222 raise TypeError("TextParallelBatch runtime_input_adapter must be a RuntimeInputAdapter")
223 self.runtime_input_adapter = runtime_input_adapter
224 self._batch_flow_logged = False
226 if source_type == "online":
227 self.boundary_resolver = OnlineBoundaryResolver()
228 self.source_input_field = "input_ids"
229 elif source_type == "indexed":
230 eod_token_id = getattr(tokenizer, "eod", None)
231 self.boundary_resolver = IndexedBoundaryResolver(eod_token_id)
232 self.source_input_field = "tokens"
233 else:
234 raise ValueError(f"Unsupported batch source type: {source_type!r}")
236 self.pp_shared_data = pp_shared_data
237 self.labels_are_shifted = bool(self.data_config.get("labels_are_shifted", True))
238 create_attention_mask = bool(
239 self.data_config.get("create_attention_mask_in_dataloader", attention_mode == "dense")
240 )
241 self.cp_sharder = CPBatchSharder(self.parallel_context)
242 self.tp_broadcaster = TPBatchBroadcaster(self.parallel_context, device)
244 self.reset_position_ids = reset_position_ids or bool(self.data_config.get("reset_position_ids", False))
245 resolved_reset_attention_mask = (
246 reset_attention_mask
247 or bool(self.data_config.get("reset_attention_mask", False))
248 or source_type == "online"
249 )
250 self.eod_mask_loss = eod_mask_loss or bool(self.data_config.get("eod_mask_loss", False))
251 self.attention_runtime = AttentionRuntime(
252 mode=attention_mode,
253 create_mask=create_attention_mask,
254 reset_mask=resolved_reset_attention_mask,
255 sliding_window=sliding_window,
256 )
258 def __call__(
259 self,
260 data_iterator: Any,
261 ) -> tuple[Mapping[str, Any], Mapping[str, Any]]:
262 """Build model and loss inputs for one forward-backward step.
264 Args:
265 data_iterator: Iterator yielding one collated Indexed or Online batch.
267 Returns:
268 Model inputs and loss inputs for one forward-backward step.
269 """
270 source_batch = self._read_source_batch(data_iterator)
271 canonical_batch = self._normalize_source_batch(source_batch)
272 cu_seq_lens = self._resolve_sequence_boundaries(canonical_batch)
273 cp_local_batch = self.cp_sharder.shard(canonical_batch)
274 parallel_batch = self.tp_broadcaster.broadcast(
275 cp_local_batch,
276 cu_seq_lens,
277 )
278 self._log_batch_flow(canonical_batch, parallel_batch)
280 position_ids = self._build_local_position_ids(
281 parallel_batch["input_ids"],
282 parallel_batch["cu_seq_lens"],
283 )
284 parallel_batch["position_ids"] = position_ids
286 loss_mask = self._build_loss_mask(parallel_batch)
287 parallel_batch["loss_mask"] = loss_mask
288 runtime_inputs = self._build_runtime_inputs(parallel_batch)
290 model_inputs, loss_inputs = self._split_model_and_loss_inputs(
291 parallel_batch,
292 runtime_inputs,
293 )
295 return model_inputs, loss_inputs
297 def _build_runtime_inputs(
298 self,
299 parallel_batch: Mapping[str, Any],
300 ) -> dict[str, Any]:
301 """Build common attention inputs plus a model-owned extension."""
302 runtime_inputs = self.attention_runtime.build(
303 batch=parallel_batch,
304 parallel_context=self.parallel_context,
305 )
306 if self.runtime_input_adapter is None:
307 return runtime_inputs
309 model_runtime_inputs = self.runtime_input_adapter.build(
310 batch=parallel_batch,
311 parallel_context=self.parallel_context,
312 options={
313 "source_type": self.source_type,
314 "attention_mode": self.attention_runtime.mode,
315 "causal": self.causal,
316 "sliding_window": self.sliding_window,
317 },
318 )
319 collisions = set(runtime_inputs).intersection(model_runtime_inputs)
320 if collisions:
321 raise ValueError(
322 "model runtime inputs cannot replace common attention inputs: "
323 f"{sorted(collisions)}"
324 )
325 runtime_inputs.update(model_runtime_inputs)
326 return runtime_inputs
328 def _log_batch_flow(
329 self,
330 canonical_batch: Mapping[str, Any] | None,
331 parallel_batch: Mapping[str, Any],
332 ) -> None:
333 """Log the resolved source and parallel batch shapes once."""
334 if not self._batch_flow_logged:
335 source_shape = None if canonical_batch is None else tuple(canonical_batch["input_ids"].shape)
336 local_shape = tuple(parallel_batch["input_ids"].shape)
337 num_boundaries = 0 if parallel_batch["cu_seq_lens"] is None else parallel_batch["cu_seq_lens"].numel()
338 logger.debug(
339 "Parallel batch flow: source=%s, tp_rank=%d/%d, cp_rank=%d/%d, "
340 "source_owner=%s, source_shape=%s, local_shape=%s, global_boundaries=%d",
341 self.source_type,
342 self.parallel_context.tp_rank,
343 self.parallel_context.tp_world_size,
344 self.parallel_context.cp_rank,
345 self.parallel_context.cp_world_size,
346 canonical_batch is not None,
347 source_shape,
348 local_shape,
349 num_boundaries,
350 enabled=True,
351 )
352 self._batch_flow_logged = True
354 def _read_source_batch(self, data_iterator: Any) -> Mapping[str, Any] | None:
355 """Read one complete batch on TP rank zero of each CP coordinate."""
356 if self.parallel_context.build_on_rank():
357 source_batch = next(data_iterator)
358 else:
359 source_batch = None
361 return source_batch
363 def _normalize_source_batch(
364 self,
365 source_batch: Mapping[str, Any] | None,
366 ) -> Mapping[str, Any] | None:
367 """Normalize Indexed and Online token fields before parallel processing."""
368 if source_batch is None:
369 return None
371 canonical_batch = dict(source_batch)
372 if self.source_input_field != "input_ids":
373 canonical_batch["input_ids"] = canonical_batch.pop(self.source_input_field)
375 return canonical_batch
377 def _resolve_sequence_boundaries(
378 self,
379 canonical_batch: Mapping[str, Any] | None,
380 ) -> Any:
381 """Resolve global leading-zero cumulative sequence boundaries."""
382 if canonical_batch is None:
383 cu_seq_lens = None
384 else:
385 cu_seq_lens = self.boundary_resolver.resolve(canonical_batch)
387 return cu_seq_lens
389 def _build_local_position_ids(
390 self,
391 input_ids: Any,
392 cu_seq_lens: Any,
393 ) -> Any:
394 """Build position IDs for the current CP slice on every TP rank.
396 Packed boundaries define the position semantics. CP only determines
397 which global sequence interval is materialized by this rank.
398 """
399 # pylint: disable=too-many-locals
400 batch_size, local_seq_len = input_ids.shape
401 cp_size = self.parallel_context.cp_world_size
402 global_seq_len = local_seq_len * cp_size
403 cp_rank = self.parallel_context.cp_rank
404 cp_seq_start = cp_rank * local_seq_len
405 cp_seq_end = cp_seq_start + local_seq_len
407 # This produces the same values as building global [B, S] positions
408 # and then taking the CP slice, without allocating the global tensor.
409 position_ids = torch.arange(
410 cp_seq_start, cp_seq_end, dtype=torch.int64, device=input_ids.device
411 )
412 local_position_ids = position_ids.unsqueeze(0).expand(batch_size, -1)
414 if self.reset_position_ids:
415 local_position_ids = local_position_ids.clone()
416 boundaries = cu_seq_lens.tolist()
417 for seq_start, seq_end in zip(boundaries[:-1], boundaries[1:]):
418 seq_len = seq_end - seq_start
419 batch_idx, seq_start_in_batch = divmod(seq_start, global_seq_len)
420 seq_end_in_batch = seq_start_in_batch + seq_len
421 # Intersect the sequence with the current CP slice.
422 slice_start = max(seq_start_in_batch, cp_seq_start)
423 slice_end = min(seq_end_in_batch, cp_seq_end)
424 if slice_start >= slice_end:
425 continue
427 local_start = slice_start - cp_seq_start
428 local_end = slice_end - cp_seq_start
429 local_position_ids[batch_idx, local_start:local_end] -= seq_start_in_batch
431 local_position_ids = local_position_ids.contiguous()
433 return local_position_ids
435 def _build_loss_mask(self, parallel_batch: Mapping[str, Any]) -> Any:
436 """Build the local loss mask from labels and input IDs."""
437 loss_mask = (parallel_batch["labels"] >= 0).to(dtype=torch.int64)
439 if self.eod_mask_loss:
440 eod_token_id = getattr(self.tokenizer, "eod", None)
441 loss_mask = loss_mask.masked_fill(parallel_batch["input_ids"] == eod_token_id, 0)
443 return loss_mask
445 def _split_model_and_loss_inputs(
446 self,
447 parallel_batch: Mapping[str, Any],
448 runtime_inputs: Mapping[str, Any] | None = None,
449 ) -> tuple[Mapping[str, Any], Mapping[str, Any]]:
450 """Split forward fields from loss and token-accounting fields."""
451 model_inputs = {
452 "input_ids": parallel_batch["input_ids"],
453 "labels": parallel_batch["labels"],
454 "position_ids": parallel_batch["position_ids"],
455 }
456 if self.labels_are_shifted:
457 model_inputs["shift_labels"] = parallel_batch["labels"]
459 runtime_inputs = runtime_inputs or {}
460 collisions = set(model_inputs).intersection(runtime_inputs)
461 if collisions:
462 raise ValueError(
463 "runtime inputs cannot replace framework-owned model inputs: "
464 f"{sorted(collisions)}"
465 )
466 model_inputs.update(runtime_inputs)
468 loss_inputs = {
469 "labels": parallel_batch["labels"],
470 "loss_mask": parallel_batch["loss_mask"],
471 }
472 if self.labels_are_shifted:
473 loss_inputs["shift_labels"] = parallel_batch["labels"]
475 return model_inputs, loss_inputs