Coverage for / home / jenkins / .local / lib / python3.10 / site-packages / hyper_parallel / data / batching / build_collate_fn.py: 79%
175 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 final batch collators for Indexed and Online sources."""
17from __future__ import annotations
19import logging
20from collections.abc import Callable, Mapping, Sequence
21from dataclasses import dataclass
22from functools import partial
23from typing import Any
25import torch # pylint: disable=forbidden-backend-import
26from torch.utils.data import default_collate # pylint: disable=forbidden-backend-import
28from hyper_parallel.data.constants import (
29 IGNORE_INDEX,
30 OMNI_CONCAT_FIELDS,
31 OMNI_CP_PAD_VALUES,
32 OMNI_CP_TOKEN_FIELDS,
33)
34from hyper_parallel.data.dataset_logging import get_dataset_logger
36logger = get_dataset_logger(__name__)
39def _get_token_alignment_size(mesh_context: Any | None) -> int:
40 """Return TP X CP, the required final batch token alignment."""
41 if mesh_context is None:
42 return 1
44 cp_size = int(getattr(mesh_context, "cp_size", 1))
45 tp_size = int(getattr(mesh_context, "tp_size", 1))
46 if cp_size <= 0 or tp_size <= 0:
47 raise ValueError("mesh_context tp_size and cp_size must be positive")
49 alignment_size = cp_size * tp_size
50 return alignment_size
53def _validate_padding_config(alignment_size: int, pad_to_length: int | None) -> None:
54 """Check the shared text/Omni padding contract once at construction."""
55 if alignment_size <= 0:
56 raise ValueError("alignment_size must be positive")
58 if pad_to_length is None:
59 return
61 if isinstance(pad_to_length, bool) or not isinstance(pad_to_length, int) or pad_to_length <= 0:
62 raise ValueError("pad_to_length must be a positive integer or None")
64 if pad_to_length % alignment_size:
65 raise ValueError(
66 f"pad_to_length={pad_to_length} must be divisible by TP * CP alignment_size={alignment_size}"
67 )
70def _pad_batch_tokens(
71 batch: Mapping[str, Any],
72 alignment_size: int,
73 token_pad_values: Mapping[str, int],
74 pad_to_length: int | None = None,
75) -> dict[str, Any]:
76 """Pad to a fixed length or alignment while preserving real boundaries and media.
78 Example:
79 Length 18, TP=4, CP=2: token fields grow to 24, labels pad with
80 -100, and cu_seq_lens [0, 10, 18] becomes [0, 10, 18, 24].
81 With pad_to_length=32, the same boundaries become [0, 10, 18, 32].
82 """
83 padded_batch = dict(batch)
84 input_ids = padded_batch["input_ids"]
85 if not torch.is_tensor(input_ids) or input_ids.ndim != 2:
86 raise ValueError("Batch alignment requires input_ids with shape [batch, sequence]")
88 sequence_length = input_ids.shape[-1]
89 if pad_to_length is None:
90 pad_length = (-sequence_length) % alignment_size
91 else:
92 if sequence_length > pad_to_length:
93 raise ValueError(
94 f"Packed sequence length {sequence_length} exceeds pad_to_length={pad_to_length}; "
95 "increase pad_to_length or reduce the packing token budget and sample lengths"
96 )
98 pad_length = pad_to_length - sequence_length
100 if not pad_length:
101 return padded_batch
103 for field, pad_value in token_pad_values.items():
104 if field not in padded_batch:
105 continue
107 value = padded_batch[field]
108 if not torch.is_tensor(value) or value.shape != input_ids.shape:
109 raise ValueError(f"Batch token field {field!r} must match input_ids shape")
111 padded_batch[field] = torch.nn.functional.pad(value, (0, pad_length), value=pad_value)
113 if "cu_seq_lens" in padded_batch:
114 boundaries = padded_batch["cu_seq_lens"].reshape(-1)
115 if (input_ids.shape[0] != 1 or boundaries.numel() < 2
116 or int(boundaries[0]) != 0 or int(boundaries[-1]) != sequence_length):
117 raise ValueError("cu_seq_lens must cover one complete packed input row")
119 # A separate ignored sequence keeps attention metadata consistent with
120 # the physical tail without changing any real sample boundary.
121 padded_batch["cu_seq_lens"] = torch.cat((boundaries, boundaries[-1:] + pad_length))
123 return padded_batch
126@dataclass
127class TextPackingCollator:
128 """Pack unpadded Online text samples and emit ``cu_seq_lens``.
130 Only ``input_ids`` and pre-shifted ``labels`` are packed here. Unified
131 ``get_batch`` constructs loss, position, mask, and CP runtime fields.
132 """
134 alignment_size: int = 1
135 pad_to_length: int | None = None
137 def __post_init__(self) -> None:
138 """Validate fixed padding before any samples are read."""
139 _validate_padding_config(self.alignment_size, self.pad_to_length)
141 def __call__(self, model_samples: Sequence[Mapping[str, Any]]) -> Mapping[str, Any]:
142 """Pack samples into one ``[1, packed_length]`` forward-backward batch.
144 Args:
145 model_samples: Online samples selected for one forward-backward step.
147 Returns:
148 Packed ``input_ids`` and ``labels`` plus int32 ``cu_seq_lens``.
150 Raises:
151 ValueError: If no samples are provided.
152 """
153 if not model_samples:
154 raise ValueError("model_samples must contain at least one Online sample")
156 packed_batch = {}
157 for field in ("input_ids", "labels"):
158 values = [model_sample[field] for model_sample in model_samples]
159 packed_batch[field] = torch.cat(values, dim=-1).unsqueeze(0)
161 seq_lens = model_samples[0]["input_ids"].new_tensor(
162 [model_sample["input_ids"].shape[-1] for model_sample in model_samples]
163 )
164 zero = seq_lens.new_zeros(1)
165 seq_ends = seq_lens.cumsum(dim=0)
166 cu_seq_lens = torch.cat((zero, seq_ends))
167 cu_seq_lens = cu_seq_lens.to(torch.int32)
168 packed_batch["cu_seq_lens"] = cu_seq_lens
170 aligned_batch = _pad_batch_tokens(
171 packed_batch, self.alignment_size, {"input_ids": 0, "labels": IGNORE_INDEX}, self.pad_to_length,
172 )
173 return aligned_batch
176@dataclass
177class OmniCollator:
178 """Collate Omni samples, then pad token fields to a fixed length or TP × CP."""
180 alignment_size: int = 1
181 pad_to_length: int | None = None
183 def __post_init__(self) -> None:
184 """Validate fixed padding before any samples are read."""
185 _validate_padding_config(self.alignment_size, self.pad_to_length)
187 def __call__(self, model_samples: Sequence[Mapping[str, Any]]) -> Mapping[str, Any]:
188 """Stack token fields, concatenate media, and pad only the final batch.
190 Example:
191 One packed sample of length 18, TP=4, CP=2 produces [1, 24]
192 token fields. Image tensors and global image starts are unchanged.
193 """
194 if not model_samples:
195 raise ValueError("model_samples must contain at least one Omni sample")
197 sample_fields = set()
198 for model_sample in model_samples:
199 sample_fields.update(model_sample)
201 batch = {}
202 for field in sample_fields:
203 values = []
204 for model_sample in model_samples:
205 if field in model_sample:
206 values.append(model_sample[field])
208 if field in OMNI_CONCAT_FIELDS:
209 batch[field] = torch.cat(values, dim=0)
210 else:
211 batch[field] = default_collate(values)
213 token_pad_values = {field: OMNI_CP_PAD_VALUES.get(field, 0) for field in OMNI_CP_TOKEN_FIELDS}
214 aligned_batch = _pad_batch_tokens(batch, self.alignment_size, token_pad_values, self.pad_to_length)
215 if logger.isEnabledFor(logging.DEBUG) and batch["input_ids"].device.type == "cpu":
216 self._log_batch(aligned_batch, batch["input_ids"].shape[-1])
217 return aligned_batch
220 def _log_batch(self, batch: Mapping[str, Any], real_length: int) -> None:
221 """Inspect CPU batches before transport without synchronizing accelerator streams."""
222 input_ids = batch["input_ids"]
223 labels = batch["labels"]
224 boundaries = batch.get("cu_seq_lens")
225 boundaries = [] if boundaries is None else boundaries.reshape(-1).tolist()
226 sequence_length = input_ids.shape[-1]
227 issues = []
228 if labels.shape != input_ids.shape:
229 issues.append("labels shape differs from input_ids")
230 if sequence_length == 0:
231 issues.append("empty input_ids")
232 if boundaries and (boundaries[0] != 0 or boundaries[-1] != sequence_length
233 or any(end <= start for start, end in zip(boundaries[:-1], boundaries[1:]))):
234 issues.append("cu_seq_lens do not cover strictly increasing packed intervals")
235 padding = sequence_length - real_length
236 supervised = int((labels >= 0).sum())
237 if not supervised:
238 issues.append("no supervised tokens")
239 if padding and bool(torch.any(labels[..., real_length:] != IGNORE_INDEX)):
240 issues.append("padding tail contains supervised labels")
241 lengths = [end - start for start, end in zip(boundaries[:-1], boundaries[1:])]
242 logger.debug(
243 "Omni collate: input_shape=%s label_shape=%s real_tokens=%d padding_tokens=%d "
244 "alignment=%d cu_seq_lens=%s segment_lengths=%s supervised_tokens=%d issues=%s",
245 tuple(input_ids.shape), tuple(labels.shape), real_length, padding, self.alignment_size,
246 boundaries, lengths, supervised, issues, enabled=True,
247 )
248 if issues:
249 logger.warning("Omni collate data anomalies: %s", issues)
252def build_default_collate_fn() -> Callable[[list[Any]], Any]:
253 """Build ordinary PyTorch collation for already prepared samples.
255 Returns:
256 PyTorch default collation.
257 """
258 collate_fn = default_collate
260 return collate_fn
263# Indexed samples are already prepared and use the same ordinary collation.
264build_indexed_collate_fn = build_default_collate_fn
267def _sft_row_boundaries(boundaries: Any, sequence_length: int) -> list[int]:
268 """Remove boundary sentinels and validate one complete SFT row."""
269 row = torch.as_tensor(boundaries)
270 if row.ndim != 1 or row.dtype not in (torch.int32, torch.int64):
271 raise ValueError("Each cu_seq_lens row must be a one-dimensional integer array")
272 row = row[row != -1]
273 if (row.numel() < 2 or row[0] != 0 or row[-1] != sequence_length
274 or torch.any(row[1:] <= row[:-1])):
275 raise ValueError("Document boundaries must strictly increase from zero to row length")
276 return row.tolist()
279def _collate_sft_samples(batch: list[Any], *, alignment_size: int) -> Any:
280 """Stack SFT rows and retain per-row boundaries without sentinel padding.
282 Records from SFTFromMRDataset have 1D input_ids/labels; already grouped
283 records may have [rows, S]. Variable-length rows are padded, never truncated.
284 """
285 if not isinstance(batch[0], Mapping) or "cu_seq_lens" not in batch[0]:
286 return default_collate(batch)
288 token_rows, label_rows, boundary_rows = [], [], []
289 for sample in batch:
290 tokens, labels = torch.as_tensor(sample["input_ids"]), torch.as_tensor(sample["labels"])
291 if tokens.ndim not in (1, 2) or tokens.shape != labels.shape or not tokens.numel():
292 raise ValueError("SFT tokens and labels must have matching nonempty [S] or [rows, S] shapes")
293 boundaries = sample["cu_seq_lens"]
294 if tokens.ndim == 1:
295 tokens, labels, boundaries = tokens.unsqueeze(0), labels.unsqueeze(0), [boundaries]
296 if len(boundaries) != len(tokens):
297 raise ValueError("cu_seq_lens must have one boundary array per token row")
298 for row in boundaries:
299 boundary_rows.append(_sft_row_boundaries(row, tokens.shape[-1]))
300 token_rows.extend(tokens.unbind(0))
301 label_rows.extend(labels.unbind(0))
303 sequence_length = max(row.numel() for row in token_rows)
304 sequence_length += (-sequence_length) % alignment_size
305 packed_keys = {"input_ids", "labels", "cu_seq_lens"}
306 result = default_collate([
307 {key: value for key, value in sample.items() if key not in packed_keys} for sample in batch
308 ])
309 for key, rows, value in (("input_ids", token_rows, 0), ("labels", label_rows, IGNORE_INDEX)):
310 result[key] = torch.stack([
311 torch.nn.functional.pad(row, (0, sequence_length - row.numel()), value=value) for row in rows
312 ])
313 for row in boundary_rows:
314 if row[-1] < sequence_length:
315 # Isolate ignored tail padding from every real document.
316 row.append(sequence_length)
317 result["cu_seq_lens"] = boundary_rows
318 return result
321def build_sft_collate_fn(mesh_context: Any | None = None) -> Callable[[list[Any]], Any]:
322 """Build SFT collation preserving [B, S] and each row's document boundaries.
324 Args:
325 mesh_context: Optional topology for TP x CP sequence-length alignment.
327 Returns:
328 A CPU collator accepting 1D records or already grouped 2D records.
329 Ordinary records without cu_seq_lens use default collation.
330 """
331 return partial(_collate_sft_samples, alignment_size=_get_token_alignment_size(mesh_context))
334def build_omni_collate_fn(
335 mesh_context: Any | None = None,
336 pad_to_length: int | None = None,
337) -> OmniCollator:
338 """Build Omni collation with optional fixed padding after packing.
340 Args:
341 mesh_context: Runtime TP/CP topology injected by the Trainer.
342 pad_to_length: Exact global token length, divisible by TP × CP. None
343 pads only for alignment. Longer packed sequences raise ValueError.
345 Returns:
346 A collator that pads token lengths and preserves modality metadata.
348 Raises:
349 ValueError: If pad_to_length is not a positive aligned integer or None.
350 """
351 alignment_size = _get_token_alignment_size(mesh_context)
352 collate_fn = OmniCollator(alignment_size=alignment_size, pad_to_length=pad_to_length)
353 return collate_fn
356def build_online_text_collate_fn(
357 mesh_context: Any | None = None,
358 pad_to_length: int | None = None,
359) -> TextPackingCollator:
360 """Build Online text collation shared by fixed N and dynamic K batching.
362 Packing concatenates ``input_ids`` and ``labels`` and emits ``cu_seq_lens``.
363 Tail padding defaults to alignment from the runtime topology. A fixed
364 pad_to_length can instead specify the exact global packed token length.
365 The unified alignment size is ``cp_size * tp_size``, derived directly
366 from the topology without a separate SP switch. Real sample boundaries
367 remain unchanged and one synthetic final boundary covers the physical
368 padding segment.
370 Args:
371 mesh_context: Runtime TP/CP topology injected by the Trainer.
372 pad_to_length: Exact global token length, divisible by TP × CP. None
373 pads only for alignment. Longer packed sequences raise ValueError.
375 Returns:
376 A collator producing one forward-backward batch.
378 Raises:
379 ValueError: If pad_to_length is not a positive aligned integer or None.
380 """
381 alignment_size = _get_token_alignment_size(mesh_context)
382 collate_fn = TextPackingCollator(alignment_size=alignment_size, pad_to_length=pad_to_length)
384 return collate_fn