Coverage for /home/jenkins/.local/lib/python3.10/site-packages/hyper_parallel/data/parallel/batch

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

7# http://www.apache.org/licenses/LICENSE-2.0

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 # pylint: disable=forbidden-backend-import

38import torch.distributed as dist # pylint: disable=forbidden-backend-import

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)

101

102 padded = _pad_batch_for_cp(batch, pad_len)

103

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

114

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

121

122

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

143

144

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).

148

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

160

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))

170

171 if max_local_packs == 0:

172 max_local_packs = 1

173

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

186

187

188# ── DataLoader batch distribution (merged from components/datasets/parallel/batch_parallel.py, 05 §11.2) ──

189

190

191class CPBatchSharder:

192 """Select contiguous CP token intervals with optional multimodal field rules."""

193

194 def __init__(

195 self,

196 parallel_context: DataLoaderParallelContext,

197 token_fields: frozenset[str] | None = None,

198 ) -> None:

199 """Initialize context-parallel batch sharding.

200

201 Args:

202 parallel_context: DataLoader and TP/CP topology information.

203 token_fields: Optional sequence fields to slice.

204 When supplied, fields not listed here remain complete on each CP rank.

205 """

206 self.parallel_context = parallel_context

207 self.token_fields = token_fields

208

209 def shard(

210 self,

211 canonical_batch: Mapping[str, Any] | None,

212 ) -> dict[str, Any] | None:

213 """Shard input IDs and labels along the sequence dim.

214

215 Args:

216 canonical_batch: Normalized complete batch on TP rank zero, or ``None`` elsewhere.

217

218 Returns:

219 CP-local input IDs and labels on TP rank zero.

220 """

221 if canonical_batch is None:

222 cp_batch = None

223 return cp_batch

224

225 if self.token_fields is not None:

226 cp_batch = self._shard_configured_fields(canonical_batch)

227 return cp_batch

228

229

230 input_ids = canonical_batch["input_ids"]

231 cp_batch = {

232 "input_ids": input_ids,

233 "labels": canonical_batch["labels"],

234 }

235

236 cp_size = self.parallel_context.cp_world_size

237 if cp_size <= 1:

238 return cp_batch

239

240 seq_len = input_ids.shape[1]

241 if seq_len % cp_size != 0:

242 raise ValueError(f"seq_len ({seq_len}) must be divisible by cp_size ({cp_size})")

243

244 cp_rank = self.parallel_context.cp_rank

245 for field, value in cp_batch.items():

246 local_value = torch.chunk(value, cp_size, dim=1)[cp_rank]

247 cp_batch[field] = local_value.contiguous()

248

249 return cp_batch

250

251 @staticmethod

252 def _extend_cu_seq_lens(cp_batch: dict[str, Any], sequence_length: int, pad_length: int) -> None:

253 """Keep packed sequence boundaries aligned with token padding."""

254 if "cu_seq_lens" not in cp_batch:

255 return

256 boundaries = cp_batch["cu_seq_lens"].reshape(-1)

257 if int(boundaries[0]) != 0 or int(boundaries[-1]) != sequence_length:

258 raise ValueError("cu_seq_lens must cover the complete CP input sequence")

259 if pad_length:

260 boundaries = torch.cat((boundaries, boundaries[-1:] + pad_length))

261 cp_batch["cu_seq_lens"] = boundaries

262

263 def _shard_configured_fields(self, canonical_batch: Mapping[str, Any]) -> dict[str, Any]:

264 """Slice aligned token fields and preserve global boundaries and media.

265

266 Example:

267 CP=2, TP=4: collate supplies [1, 24], then each CP rank

268 receives [1, 12]. A length of 18 is rejected here.

269 """

270 cp_batch = dict(canonical_batch)

271 cp_size = self.parallel_context.cp_world_size

272 alignment = cp_size * self.parallel_context.tp_world_size

273 if cp_size == 1 and alignment == 1:

274 return cp_batch

275

276 input_ids = cp_batch["input_ids"]

277 if not torch.is_tensor(input_ids) or input_ids.ndim != 2 or input_ids.shape[0] != 1:

278 raise ValueError("Configured CP sharding requires input_ids with shape [1, sequence]")

279

280 sequence_length = input_ids.shape[1]

281 if sequence_length % alignment:

282 raise ValueError(

283 f"sequence_length ({sequence_length}) must be divisible by tp_size * cp_size ({alignment}); "

284 "align token fields in collate_fn before get_batch"

285 )

286

287 if "cu_seq_lens" in cp_batch:

288 boundaries = cp_batch["cu_seq_lens"].reshape(-1)

289 if boundaries.numel() < 2 or int(boundaries[0]) != 0 or int(boundaries[-1]) != sequence_length:

290 raise ValueError("cu_seq_lens must cover the complete CP input sequence")

291

292 cp_batch["cu_seq_lens"] = boundaries

293

294 chunk_length = sequence_length // cp_size

295 local_start = self.parallel_context.cp_rank * chunk_length

296 local_end = local_start + chunk_length

297 for field in self.token_fields:

298 if field not in cp_batch:

299 continue

300

301 value = cp_batch[field]

302 if not torch.is_tensor(value) or value.ndim != 2 or value.shape != input_ids.shape:

303 raise ValueError(f"CP token field {field!r} must match input_ids shape")

304

305 cp_batch[field] = value[:, local_start:local_end].contiguous()

306

307 return cp_batch

308

309

310class TPBatchBroadcaster:

311 """Broadcast CP-local batch fields within one TP group."""

312

313 def __init__(

314 self,

315 parallel_context: DataLoaderParallelContext,

316 device: Any,

317 ) -> None:

318 """Initialize tensor-parallel batch transport.

319

320 Args:

321 parallel_context: DataLoader and TP/CP topology information.

322 device: Device receiving the broadcast batch.

323 """

324 self.parallel_context = parallel_context

325 self.device = device

326

327 def broadcast(

328 self,

329 cp_local_batch: Mapping[str, Any] | None,

330 cu_seq_lens: Any = None,

331 *,

332 broadcast_all_fields: bool = False,

333 ) -> dict[str, Any]:

334 """Broadcast CP-local fields and global boundaries across TP ranks.

335

336 Args:

337 cp_local_batch: CP-local fields on source owners, or ``None`` elsewhere.

338 When ``tp_source_rank=None``, every rank already has its fields

339 and this method only normalizes device and dtype.

340 cu_seq_lens: Global sequence boundaries, unsharded across CP ranks.

341 broadcast_all_fields: Preserve heterogeneous Omni fields and dtypes.

342

343 Returns:

344 Local batch populated on every TP rank.

345 """

346 if broadcast_all_fields:

347 parallel_batch = self._broadcast_mapping(cp_local_batch)

348 return parallel_batch

349

350 tp_rank = self.parallel_context.tp_rank

351 source_rank = self.parallel_context.tp_source_rank

352 if source_rank is None or tp_rank == source_rank:

353 parallel_batch = {

354 field: value.to(self.device, dtype=torch.int64, non_blocking=True)

355 for field, value in cp_local_batch.items()

356 }

357 local_cu_seq_lens = None

358 if cu_seq_lens is not None:

359 local_cu_seq_lens = cu_seq_lens.to(self.device, dtype=torch.int32, non_blocking=True)

360 else:

361 parallel_batch = None

362 local_cu_seq_lens = None

363

364 tp_size = self.parallel_context.tp_world_size

365 if source_rank is None or tp_size <= 1:

366 parallel_batch["cu_seq_lens"] = local_cu_seq_lens

367 return parallel_batch

368

369 # Dynamic Online batches do not have a fixed local sequence length.

370 # Broadcast compact shape metadata before receivers allocate tensors.

371 if tp_rank == source_rank:

372 batch_size, seq_len = parallel_batch["input_ids"].shape

373 num_boundaries = 0 if local_cu_seq_lens is None else local_cu_seq_lens.numel()

374 batch_meta = torch.tensor(

375 [batch_size, seq_len, num_boundaries],

376 dtype=torch.int64,

377 device=self.device,

378 )

379 else:

380 batch_meta = torch.empty((3,), dtype=torch.int64, device=self.device)

381

382 tp_group = self.parallel_context.tp_group

383 dist.broadcast(batch_meta, group=tp_group, group_src=source_rank)

384

385 # Non-source TP ranks allocate the same device shapes and dtypes.

386 if tp_rank != source_rank:

387 batch_size, seq_len, num_boundaries = [int(value) for value in batch_meta.tolist()]

388 shape = (batch_size, seq_len)

389 parallel_batch = {

390 "input_ids": torch.empty(shape, dtype=torch.int64, device=self.device),

391 "labels": torch.empty(shape, dtype=torch.int64, device=self.device),

392 }

393 if num_boundaries > 0:

394 local_cu_seq_lens = torch.empty((num_boundaries,), dtype=torch.int32, device=self.device)

395

396 # Broadcast only model fields that cannot be regenerated locally.

397 for field in ("input_ids", "labels"):

398 tensor = parallel_batch.get(field)

399 if tensor is None:

400 raise ValueError(f"parallel batch requires {field!r}")

401

402 dist.broadcast(tensor, group=tp_group, group_src=source_rank)

403

404 # Packed boundaries remain global across CP ranks and variable in size.

405 if local_cu_seq_lens is not None:

406 dist.broadcast(local_cu_seq_lens, group=tp_group, group_src=source_rank)

407

408 parallel_batch["cu_seq_lens"] = local_cu_seq_lens

409

410 return parallel_batch

411

412 def _broadcast_mapping(self, cp_local_batch: Mapping[str, Any] | None) -> dict[str, Any]:

413 """Broadcast variable-shaped Omni tensors and scalar metadata."""

414 source_rank = self.parallel_context.tp_source_rank

415 if source_rank is None or self.parallel_context.tp_world_size == 1:

416 parallel_batch = self._move_mapping_to_device(cp_local_batch)

417 return parallel_batch

418

419 is_source = self.parallel_context.tp_rank == source_rank

420 if is_source:

421 parallel_batch = self._move_mapping_to_device(cp_local_batch)

422 metadata = []

423 for name, value in parallel_batch.items():

424 if torch.is_tensor(value):

425 metadata.append((name, tuple(value.shape), value.dtype, None))

426 else:

427 metadata.append((name, None, None, value))

428 else:

429 parallel_batch = {}

430 metadata = None

431 object_list = [metadata]

432 dist.broadcast_object_list(

433 object_list, group=self.parallel_context.tp_group, group_src=source_rank, device=self.device

434 )

435 for name, shape, dtype, value in object_list[0]:

436 if shape is None:

437 parallel_batch[name] = value

438 continue

439 if not is_source:

440 parallel_batch[name] = torch.empty(shape, dtype=dtype, device=self.device)

441 dist.broadcast(parallel_batch[name], group=self.parallel_context.tp_group, group_src=source_rank)

442 return parallel_batch

443

444 def _move_mapping_to_device(self, batch: Mapping[str, Any]) -> dict[str, Any]:

445 """Transfer Omni tensors without changing their native dtypes."""

446 return {

447 field: value.to(self.device, non_blocking=True) if torch.is_tensor(value) else value

448 for field, value in batch.items()

449 }

Coverage for / home / jenkins / .local / lib / python3.10 / site-packages / hyper_parallel / data / parallel / batch_parallel.py: 82%

222 statements