Coverage for / home / jenkins / .local / lib / python3.10 / site-packages / hyper_parallel / data / parallel / batch_parallel.py: 63%
216 statements
« prev ^ index » next coverage.py v7.13.1, created at 2026-10-25 05:45 +0800
1# Copyright 2025-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# pylint: disable=missing-apache-license-header,forbidden-backend-import
17"""data.parallel.batch_parallel: CP-aware batch sharding.
19``shard_batch_for_cp``: data-pipeline CP sharding (aligned with the THD
20contract of the 02 collater); ``_shard_seq_lens_for_cp`` recomputes
21seq_lens/seq_lens_padded per CP rank. ``CPBatchSharder`` slices
22``input_ids``/``labels`` contiguously along the sequence dimension and
23``TPBatchBroadcaster`` moves the CP-local tensors and global sequence
24boundaries to the target device, broadcasting metadata and tensor fields
25to peer TP ranks (the canonical batch exists only on TP rank zero for the
26active CP coordinate).
28Split out of components/distributed/cp_utils.py in stage 4e; merged with
29components/datasets/parallel/batch_parallel.py in stage 6 (05 §11.2).
30"""
32from __future__ import annotations
34from collections.abc import Mapping
35from typing import Any
37import torch
38import torch.distributed as dist
40from hyper_parallel.core.dtensor.device_mesh import DeviceMesh
41from hyper_parallel.data.parallel.dataloader_parallel import DataLoaderParallelContext
44def _pad_batch_for_cp(batch: dict[str, Any], pad_len: int) -> dict[str, Any]:
45 """Pad sequence-dimension tensor fields for context parallelism."""
46 padded = dict(batch)
47 if pad_len <= 0:
48 return padded
50 pad_values = {"labels": -100, "input_ids": 0, "attention_mask": 0}
51 for key, value in batch.items():
52 if key == "qkv_format" or not isinstance(value, torch.Tensor) or value.ndim < 1:
53 continue
54 if key in ("seq_lens", "seq_lens_padded"):
55 continue # recomputed separately, not padded
56 if key == "position_ids":
57 # position_ids increment-pad: continue incrementing from the last value
58 last = value[..., -1:].to(torch.long)
59 increment = torch.arange(1, pad_len + 1, device=value.device,
60 dtype=value.dtype)
61 increment = increment.reshape(*([1] * (value.ndim - 1)), pad_len)
62 pad_block = increment.expand(*value.shape[:-1], pad_len) + last
63 else:
64 shape = list(value.shape)
65 shape[-1] = pad_len
66 pad_block = torch.full(shape, pad_values.get(key, 0),
67 dtype=value.dtype, device=value.device)
68 padded[key] = torch.cat([value, pad_block], dim=-1)
69 return padded
72def shard_batch_for_cp(batch: dict[str, Any], cp_mesh: DeviceMesh) -> dict[str, Any]:
73 """Shard the sequence-dim tensors of a batch along the CP mesh
74 (05 §6.3.4 canonical).
76 Contract (aligned with the 02 collater output):
77 - input_ids/labels/position_ids: [B, S] int64
78 - seq_lens / seq_lens_padded: [B, max_num_packs] int64, padded with the
79 -1000 sentinel
80 - qkv_format: "thd" (passthrough)
82 Sharding strategy: pad to a multiple of 2*cp, then slice the token
83 interval [cp_rank*chunk, (cp_rank+1)*chunk); the seq_lens family is
84 recomputed separately (_shard_seq_lens_for_cp).
86 Args:
87 batch: Canonical token batch to shard.
88 cp_mesh: Context-parallel process mesh.
89 """
90 cp_size = cp_mesh.size()
91 if cp_size <= 1:
92 return batch
94 cp_rank = cp_mesh.get_local_rank()
95 seq_len = batch["input_ids"].shape[1]
96 pad_len = (-seq_len) % (cp_size * 2)
97 chunk = (seq_len + pad_len) // cp_size
98 lo = cp_rank * chunk
99 hi = lo + chunk
100 slc = slice(lo, hi)
102 padded = _pad_batch_for_cp(batch, pad_len)
104 out = {}
105 for k, v in padded.items():
106 if k in ("seq_lens", "seq_lens_padded"):
107 continue
108 if k == "qkv_format":
109 out[k] = v
110 elif isinstance(v, torch.Tensor) and v.ndim >= 1:
111 out[k] = v[..., slc]
112 else:
113 out[k] = v
115 if "seq_lens" in batch and "seq_lens_padded" in batch:
116 out["seq_lens"], out["seq_lens_padded"] = _shard_seq_lens_for_cp(
117 batch["seq_lens"], batch["seq_lens_padded"],
118 cp_rank=cp_rank, chunk=chunk,
119 )
120 return out
123def _shard_seq_lens_row(row_lens, row_padded, lo, hi):
124 """Recompute one sample's pack lengths inside a CP shard."""
125 local_lens, local_padded = [], []
126 offset = 0
127 for raw_len, raw_pad in zip(row_lens, row_padded):
128 if raw_len == -1000:
129 break
130 pack_start = offset
131 pack_end = offset + raw_pad
132 offset = pack_end
133 inter_start = max(pack_start, lo)
134 inter_end = min(pack_end, hi)
135 if inter_start >= inter_end:
136 continue
137 local_actual = max(min(pack_start + raw_len, hi) - max(pack_start, lo), 0)
138 local_pad = inter_end - inter_start
139 if local_actual > 0 or local_pad > 0:
140 local_lens.append(local_actual)
141 local_padded.append(local_pad)
142 return local_lens, local_padded
145def _shard_seq_lens_for_cp(seq_lens, seq_lens_padded, *, cp_rank: int, chunk: int):
146 """Recompute seq_lens/seq_lens_padded per CP shard (preserving the -1000
147 sentinel semantics).
149 Walk the cumulative pack offsets of each sample (accumulated by
150 seq_lens_padded); for each pack:
151 - fully inside [lo, hi): kept as-is;
152 - crossing the boundary: truncated to [lo, hi), with the actual /
153 padding-inclusive lengths recomputed after truncation;
154 - fully outside: skipped.
155 The output is shifted to the local coordinate system; when
156 max_local_packs=0 it is set to 1 to avoid an empty tensor.
157 """
158 lo = cp_rank * chunk
159 hi = lo + chunk
161 local_lens_b, local_lens_padded_b = [], []
162 max_local_packs = 0
163 for b in range(seq_lens.shape[0]):
164 local_lens, local_padded = _shard_seq_lens_row(
165 seq_lens[b].tolist(), seq_lens_padded[b].tolist(), lo, hi,
166 )
167 local_lens_b.append(local_lens)
168 local_lens_padded_b.append(local_padded)
169 max_local_packs = max(max_local_packs, len(local_lens))
171 if max_local_packs == 0:
172 max_local_packs = 1
174 out_lens = torch.full((seq_lens.shape[0], max_local_packs), -1000,
175 dtype=seq_lens.dtype, device=seq_lens.device)
176 out_padded = torch.full((seq_lens.shape[0], max_local_packs), -1000,
177 dtype=seq_lens_padded.dtype, device=seq_lens.device)
178 for b in range(seq_lens.shape[0]):
179 n = len(local_lens_b[b])
180 if n > 0:
181 out_lens[b, :n] = torch.tensor(
182 local_lens_b[b], dtype=seq_lens.dtype, device=seq_lens.device)
183 out_padded[b, :n] = torch.tensor(
184 local_lens_padded_b[b], dtype=seq_lens_padded.dtype, device=seq_lens.device)
185 return out_lens, out_padded
188# ── DataLoader batch distribution (merged from components/datasets/parallel/batch_parallel.py, 05 §11.2) ──
191class CPBatchSharder:
192 """Select contiguous CP token intervals with optional multimodal field rules."""
194 def __init__(
195 self,
196 parallel_context: DataLoaderParallelContext,
197 token_pad_values: Mapping[str, int] | None = None,
198 required_global_alignment: int | None = None,
199 ) -> None:
200 """Initialize context-parallel batch sharding.
202 Args:
203 parallel_context: DataLoader and TP/CP topology information.
204 token_pad_values: Optional sequence fields and their padding values.
205 When supplied, fields not listed here remain complete on each CP rank.
206 """
207 self.parallel_context = parallel_context
208 self.token_pad_values = token_pad_values
209 if required_global_alignment is not None and required_global_alignment < 1:
210 raise ValueError("required_global_alignment must be positive")
211 self.required_global_alignment = required_global_alignment
213 def shard(
214 self,
215 canonical_batch: Mapping[str, Any] | None,
216 ) -> dict[str, Any] | None:
217 """Shard input IDs and labels along the sequence dim.
219 Args:
220 canonical_batch: Normalized complete batch on TP rank zero, or ``None`` elsewhere.
222 Returns:
223 CP-local input IDs and labels on TP rank zero.
224 """
225 if canonical_batch is None:
226 cp_batch = None
227 return cp_batch
229 if self.token_pad_values is not None:
230 return self._shard_configured_fields(canonical_batch)
232 input_ids = canonical_batch["input_ids"]
233 cp_batch = {
234 "input_ids": input_ids,
235 "labels": canonical_batch["labels"],
236 }
238 cp_size = self.parallel_context.cp_world_size
239 if cp_size <= 1:
240 return cp_batch
242 seq_len = input_ids.shape[1]
243 if seq_len % cp_size != 0:
244 raise ValueError(f"seq_len ({seq_len}) must be divisible by cp_size ({cp_size})")
246 cp_rank = self.parallel_context.cp_rank
247 for field, value in cp_batch.items():
248 local_value = torch.chunk(value, cp_size, dim=1)[cp_rank]
249 cp_batch[field] = local_value.contiguous()
251 return cp_batch
253 @staticmethod
254 def _extend_cu_seq_lens(cp_batch: dict[str, Any], sequence_length: int, pad_length: int) -> None:
255 """Keep packed sequence boundaries aligned with token padding."""
256 if "cu_seq_lens" not in cp_batch:
257 return
258 boundaries = cp_batch["cu_seq_lens"].reshape(-1)
259 if int(boundaries[0]) != 0 or int(boundaries[-1]) != sequence_length:
260 raise ValueError("cu_seq_lens must cover the complete CP input sequence")
261 if pad_length:
262 boundaries = torch.cat((boundaries, boundaries[-1:] + pad_length))
263 cp_batch["cu_seq_lens"] = boundaries
265 def _shard_configured_fields(self, canonical_batch: Mapping[str, Any]) -> dict[str, Any]:
266 """Shard configured token fields and preserve all modality metadata."""
267 cp_batch = dict(canonical_batch)
268 cp_size = self.parallel_context.cp_world_size
269 input_ids = cp_batch["input_ids"]
270 if not torch.is_tensor(input_ids) or input_ids.ndim != 2 or input_ids.shape[0] != 1:
271 raise ValueError("Configured CP sharding requires input_ids with shape [1, sequence]")
273 sequence_length = input_ids.shape[1]
274 alignment = self.required_global_alignment or 2 * cp_size
275 pad_length = (-sequence_length) % alignment
276 self._extend_cu_seq_lens(cp_batch, sequence_length, pad_length)
278 if cp_size == 1 and pad_length == 0:
279 return cp_batch
280 chunk_length = (sequence_length + pad_length) // cp_size
281 local_start = self.parallel_context.cp_rank * chunk_length
282 local_end = local_start + chunk_length
283 for field, pad_value in self.token_pad_values.items():
284 if field not in cp_batch:
285 continue
287 value = cp_batch[field]
288 if not torch.is_tensor(value) or value.ndim != 2 or value.shape != input_ids.shape:
289 raise ValueError(f"CP token field {field!r} must match input_ids shape")
291 padded = torch.nn.functional.pad(value, (0, pad_length), value=pad_value)
292 cp_batch[field] = padded[:, local_start:local_end].contiguous()
294 return cp_batch
297class TPBatchBroadcaster:
298 """Broadcast CP-local batch fields within one TP group."""
300 def __init__(
301 self,
302 parallel_context: DataLoaderParallelContext,
303 device: Any,
304 ) -> None:
305 """Initialize tensor-parallel batch transport.
307 Args:
308 parallel_context: DataLoader and TP/CP topology information.
309 device: Device receiving the broadcast batch.
310 """
311 self.parallel_context = parallel_context
312 self.device = device
314 def broadcast(
315 self,
316 cp_local_batch: Mapping[str, Any] | None,
317 cu_seq_lens: Any = None,
318 *,
319 broadcast_all_fields: bool = False,
320 ) -> dict[str, Any]:
321 """Broadcast CP-local fields and global boundaries across TP ranks.
323 Args:
324 cp_local_batch: CP-local fields on TP rank zero, or ``None`` elsewhere.
325 cu_seq_lens: Global sequence boundaries, unsharded across CP ranks.
326 broadcast_all_fields: Preserve heterogeneous Omni fields and dtypes.
328 Returns:
329 Local batch populated on every TP rank.
330 """
331 if broadcast_all_fields:
332 return self._broadcast_mapping(cp_local_batch)
333 tp_rank = self.parallel_context.tp_rank
334 # TP rank zero owns the CP-local DataLoader fields and moves them to
335 # the target device before communication.
336 if tp_rank == 0:
337 parallel_batch = {
338 field: value.to(self.device, dtype=torch.int64, non_blocking=True)
339 for field, value in cp_local_batch.items()
340 }
341 local_cu_seq_lens = None
342 if cu_seq_lens is not None:
343 local_cu_seq_lens = cu_seq_lens.to(self.device, dtype=torch.int32, non_blocking=True)
344 else:
345 parallel_batch = None
346 local_cu_seq_lens = None
348 tp_size = self.parallel_context.tp_world_size
349 if tp_size <= 1:
350 parallel_batch["cu_seq_lens"] = local_cu_seq_lens
351 return parallel_batch
353 # Dynamic Online batches do not have a fixed local sequence length.
354 # Broadcast compact shape metadata before receivers allocate tensors.
355 if tp_rank == 0:
356 batch_size, seq_len = parallel_batch["input_ids"].shape
357 num_boundaries = 0 if local_cu_seq_lens is None else local_cu_seq_lens.numel()
358 batch_meta = torch.tensor(
359 [batch_size, seq_len, num_boundaries],
360 dtype=torch.int64,
361 device=self.device,
362 )
363 else:
364 batch_meta = torch.empty((3,), dtype=torch.int64, device=self.device)
366 tp_group = self.parallel_context.tp_group
367 dist.broadcast(batch_meta, group=tp_group, group_src=0)
369 # Non-source TP ranks allocate the same device shapes and dtypes.
370 if tp_rank != 0:
371 batch_size, seq_len, num_boundaries = [int(value) for value in batch_meta.tolist()]
372 shape = (batch_size, seq_len)
373 parallel_batch = {
374 "input_ids": torch.empty(shape, dtype=torch.int64, device=self.device),
375 "labels": torch.empty(shape, dtype=torch.int64, device=self.device),
376 }
377 if num_boundaries > 0:
378 local_cu_seq_lens = torch.empty((num_boundaries,), dtype=torch.int32, device=self.device)
380 # Broadcast only model fields that cannot be regenerated locally.
381 for field in ("input_ids", "labels"):
382 tensor = parallel_batch.get(field)
383 if tensor is None:
384 raise ValueError(f"parallel batch requires {field!r}")
385 dist.broadcast(tensor, group=tp_group, group_src=0)
387 # Packed boundaries remain global across CP ranks and variable in size.
388 if local_cu_seq_lens is not None:
389 dist.broadcast(local_cu_seq_lens, group=tp_group, group_src=0)
391 parallel_batch["cu_seq_lens"] = local_cu_seq_lens
393 return parallel_batch
395 def _broadcast_mapping(self, cp_local_batch: Mapping[str, Any] | None) -> dict[str, Any]:
396 """Broadcast variable-shaped Omni tensors and scalar metadata."""
397 if self.parallel_context.tp_world_size == 1:
398 return self._move_mapping_to_device(cp_local_batch)
399 is_source = self.parallel_context.tp_rank == 0
400 if is_source:
401 parallel_batch = self._move_mapping_to_device(cp_local_batch)
402 metadata = []
403 for name, value in parallel_batch.items():
404 if torch.is_tensor(value):
405 metadata.append((name, tuple(value.shape), value.dtype, None))
406 else:
407 metadata.append((name, None, None, value))
408 else:
409 parallel_batch = {}
410 metadata = None
411 object_list = [metadata]
412 dist.broadcast_object_list(
413 object_list, group=self.parallel_context.tp_group, group_src=0, device=self.device
414 )
415 for name, shape, dtype, value in object_list[0]:
416 if shape is None:
417 parallel_batch[name] = value
418 continue
419 if not is_source:
420 parallel_batch[name] = torch.empty(shape, dtype=dtype, device=self.device)
421 dist.broadcast(parallel_batch[name], group=self.parallel_context.tp_group, group_src=0)
422 return parallel_batch
424 def _move_mapping_to_device(self, batch: Mapping[str, Any]) -> dict[str, Any]:
425 """Transfer Omni tensors without changing their native dtypes."""
426 return {
427 field: value.to(self.device, non_blocking=True) if torch.is_tensor(value) else value
428 for field, value in batch.items()
429 }