Coverage for /home/jenkins/.local/lib/python3.10/site-packages/hyper_parallel/distributed/context_parallel/indexer

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"""No-grad query-task redistribution for compressed Indexer selection.

17Main attention, teacher and KL retain the original sequence ownership. Selection

18exchanges low-rank query residuals, merge weights and original rotary embeddings;

19receivers project Indexer Q and return integer TopK results. The canonical Index K

20bank remains differentiable and is never repacked by this communication layer.

21"""

22from __future__ import annotations

24from bisect import bisect_right

25from collections.abc import Callable, Iterator

26from dataclasses import dataclass

27from functools import lru_cache

28from heapq import heapify, heappop, heappush

29from itertools import accumulate

30from typing import Any

32import torch # pylint: disable=forbidden-backend-import

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

36def resolve_indexer_balance_block_size(mode: str, block_size: int | None) -> int | None:

37 """Resolve the greedy-pool query block option, defaulting to 2048 rows.

39 None represents an omitted option so none/zigzag remain valid without it.

40 Explicit values require greedy/hybrid mode and must be positive integers.

41 """

42 if mode not in ("none", "zigzag", "greedy", "hybrid"):

43 raise ValueError("indexer_load_balance must be 'none', 'zigzag', 'greedy' or 'hybrid'")

44 if block_size is not None:

45 if mode not in ("greedy", "hybrid"):

46 raise ValueError("indexer_balance_block_size requires indexer_load_balance='greedy' or 'hybrid'")

47 if not isinstance(block_size, int) or isinstance(block_size, bool) or block_size < 1:

48 raise ValueError("indexer_balance_block_size must be a positive integer")

49 return (2048 if block_size is None else block_size) if mode in ("greedy", "hybrid") else None

52def validate_indexer_hybrid_threshold(mode: str, min_length: int | None) -> None:

53 """Require an explicit raw-token document threshold only for hybrid scheduling.

55 The threshold selects complete zigzag fragments independently of CP size and

56 the greedy block size. No performance threshold is chosen implicitly.

58 Args:

59 mode: Selected Indexer load-balancing strategy.

60 min_length: Explicit positive raw-token document threshold for hybrid.

61 """

62 if mode != "hybrid":

63 if min_length is not None:

64 raise ValueError("indexer_zigzag_min_length requires indexer_load_balance='hybrid'")

65 elif not isinstance(min_length, int) or isinstance(min_length, bool) or min_length < 1:

66 raise ValueError(

67 "indexer_load_balance='hybrid' requires an explicit positive integer indexer_zigzag_min_length"

68 )

71@dataclass(frozen=True, order=True)

72class IndexerTask:

73 """A contiguous query interval in the original document coordinates."""

75 document: int

76 start: int

77 end: int

79 @property

80 def length(self) -> int:

81 """Return the number of real query rows, excluding any synthetic padding."""

82 return self.end - self.start

85def _merge_tasks(tasks: list[IndexerTask]) -> tuple[IndexerTask, ...]:

86 merged = []

87 for task in sorted(tasks):

88 if merged and merged[-1].document == task.document and merged[-1].end == task.start:

89 merged[-1] = IndexerTask(task.document, merged[-1].start, task.end)

90 elif task.length:

91 merged.append(task)

92 return tuple(merged)

95def _work_prefix(length: int, ratio: int) -> int:

96 groups, tail = divmod(length, ratio)

97 return ratio * groups * (groups - 1) // 2 + groups * (tail + 1)

100def _task_work(task: IndexerTask, ratio: int) -> int:

101 return _work_prefix(task.end, ratio) - _work_prefix(task.start, ratio)

102

103

104def _zigzag_tasks(lengths: tuple[int, ...], size: int) -> list[list[IndexerTask]]:

105 """Assign paired document chunks and evenly distribute their leading remainders."""

106 assignments = [[] for _ in range(size)]

107 remainders = []

108 for document, length in enumerate(lengths):

109 chunk, remainder = divmod(length, 2 * size)

110 if remainder:

111 remainders.append(IndexerTask(document, 0, remainder))

112 if chunk:

113 for rank in range(size):

114 for slot in (rank, 2 * size - 1 - rank):

115 start = remainder + slot * chunk

116 assignments[rank].append(IndexerTask(document, start, start + chunk))

117 # All ranks already own equally many paired rows. Distribute the real

118 # remainder rows evenly without inserting tokens into any document.

119 capacity = sum(task.length for task in remainders) // size

120 rank, used = 0, 0

121 for task in remainders:

122 start = task.start

123 while start < task.end:

124 if used == capacity:

125 rank += 1

126 used = 0

127 end = min(task.end, start + capacity - used)

128 assignments[rank].append(IndexerTask(task.document, start, end))

129 used += end - start

130 start = end

131 return assignments

132

133

134def _greedy_tasks(lengths: tuple[int, ...], size: int, ratio: int, block: int) -> list[list[IndexerTask]]:

135 """Keep the original cost/count/rank ordering without scanning every rank."""

136 tasks = [IndexerTask(document, start, min(start + block, length))

137 for document, length in enumerate(lengths) for start in range(0, length, block)]

138 assignments = [[] for _ in range(size)]

139 _assign_greedy_tasks(assignments, tasks, ratio, 2 * (sum(lengths) // size))

140 return assignments

141

142

143def _assign_greedy_tasks(assignments: list[list[IndexerTask]], tasks: list[IndexerTask],

144 ratio: int, capacity: int) -> None:

145 """Add pooled tasks to the existing per-worker cost and row budgets."""

146 available = [(sum(_task_work(task, ratio) for task in assigned), sum(task.length for task in assigned), rank)

147 for rank, assigned in enumerate(assignments)]

148 heapify(available)

149 # Reuse the cost computed for sorting when updating the selected worker.

150 for negative_work, task in sorted((-_task_work(task, ratio), task) for task in tasks):

151 skipped = []

152 for _ in range(len(available)):

153 cost, count, rank = heappop(available)

154 if count + task.length <= capacity:

155 break

156 skipped.append((cost, count, rank))

157 else:

158 raise ValueError(f"No Indexer worker has capacity for {task.length} query rows (limit: {capacity})")

159 assignments[rank].append(task)

160 heappush(available, (cost - negative_work, count + task.length, rank))

161 for item in skipped:

162 heappush(available, item)

163

164

165def _hybrid_tasks(lengths: tuple[int, ...], size: int, ratio: int, block: int,

166 min_length: int) -> list[list[IndexerTask]]:

167 """Keep eligible zigzag pairs intact and greedily allocate every remaining row.

168

169 Documents shorter than 2*size contain no complete pair and enter the greedy pool.

170 """

171 assignments, pooled = [[] for _ in range(size)], []

172 for document, length in enumerate(lengths):

173 chunk, remainder = divmod(length, 2 * size)

174 if length >= min_length and chunk:

175 for rank in range(size):

176 for slot in (rank, 2 * size - 1 - rank):

177 start = remainder + slot * chunk

178 assignments[rank].append(IndexerTask(document, start, start + chunk))

179 pool_end = remainder

180 else:

181 pool_end = length

182 pooled.extend(IndexerTask(document, start, min(start + block, pool_end))

183 for start in range(0, pool_end, block))

184 _assign_greedy_tasks(assignments, pooled, ratio, 2 * (sum(lengths) // size))

185 return assignments

186

187

188def _validate_task_plan(lengths: tuple[int, ...], size: int, ratio: int, mode: str,

189 block_size: int | None) -> int | None:

190 """Keep input validation identical for direct and all-worker planning."""

191 if size < 1 or ratio < 1 or not lengths or min(lengths) < 1 or sum(lengths) % size:

192 raise ValueError("Indexer balance needs positive document lengths/ratio and an equal contiguous CP partition")

193 block = resolve_indexer_balance_block_size(mode, block_size)

194 if mode not in ("zigzag", "greedy", "hybrid"):

195 raise ValueError("Indexer task mode must be 'zigzag', 'greedy' or 'hybrid'")

196 local = sum(lengths) // size

197 if block is not None and block > local:

198 raise ValueError(f"Indexer greedy block size {block} exceeds the local query count {local}")

199 return block

200

201

202@lru_cache(maxsize=128)

203def indexer_task_plan(lengths: tuple[int, ...], size: int, ratio: int, mode: str,

204 block_size: int | None = None, zigzag_min_length: int | None = None) -> tuple:

205 """Assign real queries using zigzag, greedy or a shared-budget hybrid plan.

206

207 Args:

208 lengths: Global document lengths in packed order.

209 size: CP group size; total packed tokens must divide this value.

210 ratio: Compression ratio used by the unchanged Indexer kernel.

211 mode: 'zigzag' pairs document tails and heads after removing the leading

212 remainder; 'greedy' weights blocks by the number of causal Q/K pairs;

213 'hybrid' keeps long-document zigzag fragments and greedily places

214 their leading remainders together with all shorter documents.

215 block_size: Maximum original query rows per greedy task, before adjacent

216 tasks assigned to one rank are merged. Greedy/hybrid default to 2048.

217 Explicit values require greedy/hybrid and must be positive and at most local Q

218 to guarantee its per-worker row bound. Omit this option for zigzag.

219 zigzag_min_length: Hybrid-only minimum document length in original tokens.

220 Must be specified explicitly. Eligible fragments remain whole,

221 independently of block_size. Documents shorter than 2*size enter the

222 greedy pool because they have no complete zigzag fragments.

223

224 Returns:

225 A sorted, disjoint task list for each rank. Greedy/hybrid cap each worker at

226 twice the original local query count; zigzag preserves equal row counts.

227 """

228 block = _validate_task_plan(lengths, size, ratio, mode, block_size)

229 validate_indexer_hybrid_threshold(mode, zigzag_min_length)

230 if mode == "zigzag":

231 assignments = _zigzag_tasks(lengths, size)

232 elif mode == "hybrid":

233 assignments = _hybrid_tasks(lengths, size, ratio, block, zigzag_min_length)

234 else:

235 assignments = _greedy_tasks(lengths, size, ratio, block)

236 return tuple(_merge_tasks(tasks) for tasks in assignments)

237

238

239def _segments(tasks: tuple, lengths: tuple[int, ...], ratio: int) -> tuple:

240 key_bases = (0, *accumulate(length // ratio for length in lengths))

241 result, cursor = [], 0

242 for task in tasks:

243 base = key_bases[task.document]

244 result.append((cursor, cursor + task.length, base, base + task.end // ratio, task.end % ratio))

245 cursor += task.length

246 return tuple(result)

247

248

249def _send_ranges(tasks: tuple, boundaries: tuple[int, ...], owner: int, local: int) -> tuple:

250 """Find only tasks intersecting this owner in each sorted worker assignment."""

251 first, last = owner * local, (owner + 1) * local

252 start = bisect_right(tasks, first, key=lambda task: boundaries[task.document] + task.end)

253 result = []

254 for index in range(start, len(tasks)):

255 task = tasks[index]

256 begin, end = boundaries[task.document] + task.start, boundaries[task.document] + task.end

257 if begin >= last:

258 break

259 begin, end = max(first, begin), min(last, end)

260 result.append((begin - first, end - begin))

261 return tuple(result)

262

263

264def _receive_ranges(tasks: tuple, boundaries: tuple[int, ...], size: int, local: int) -> tuple:

265 """Index a worker's tasks by owner once, retaining its original packed offsets."""

266 owners, cursor = [[] for _ in range(size)], 0

267 for task in tasks:

268 begin, end = boundaries[task.document] + task.start, boundaries[task.document] + task.end

269 for owner in range(begin // local, (end - 1) // local + 1):

270 first, last = max(begin, owner * local), min(end, (owner + 1) * local)

271 owners[owner].append((first - owner * local, cursor + first - begin, last - first))

272 cursor += task.length

273 return tuple(tuple(rows) for rows in owners)

274

275

276def _append_ordered_task(tasks: list[IndexerTask], document: int, start: int, end: int) -> None:

277 """Merge an already ordered interval without sorting the worker's tasks."""

278 if start >= end:

279 return

280 if tasks and tasks[-1].document == document and tasks[-1].end == start:

281 tasks[-1] = IndexerTask(document, tasks[-1].start, end)

282 else:

283 tasks.append(IndexerTask(document, start, end))

284

285

286def _zigzag_rank_tasks(lengths: tuple[int, ...], size: int, rank: int) -> tuple[IndexerTask, ...]:

287 """Construct only this worker's paired chunks and share of leading remainders."""

288 capacity = sum(length % (2 * size) for length in lengths) // size

289 first, last = rank * capacity, (rank + 1) * capacity

290 pooled, tasks = 0, []

291 for document, length in enumerate(lengths):

292 chunk, remainder = divmod(length, 2 * size)

293 _append_ordered_task(tasks, document, max(0, first - pooled), min(remainder, last - pooled))

294 pooled += remainder

295 if chunk:

296 for slot in (rank, 2 * size - 1 - rank):

297 start = remainder + slot * chunk

298 _append_ordered_task(tasks, document, start, start + chunk)

299 return tuple(tasks)

300

301

302def _zigzag_intersections(length: int, size: int, pooled: int, capacity: int,

303 begin: int, end: int) -> Iterator[tuple[int, int, int]]:

304 """Yield each destination intersecting an owner's interval in one document."""

305 chunk, remainder = divmod(length, 2 * size)

306 # Positive remainders imply positive capacity after validating total % size.

307 while begin < min(end, remainder):

308 peer = (pooled + begin) // capacity

309 stop = min(end, remainder, (peer + 1) * capacity - pooled)

310 yield peer, begin, stop - begin

311 begin = stop

312 while begin < end:

313 slot = (begin - remainder) // chunk

314 peer = min(slot, 2 * size - 1 - slot)

315 stop = min(end, remainder + (slot + 1) * chunk)

316 yield peer, begin, stop - begin

317 begin = stop

318

319

320def _zigzag_send_ranges(lengths: tuple[int, ...], size: int, rank: int, boundaries: tuple) -> tuple:

321 """Invert zigzag placement for local rows instead of visiting all workers."""

322 local = boundaries[-1] // size

323 pooled = (0, *accumulate(length % (2 * size) for length in lengths))

324 capacity = pooled[-1] // size

325 first, last = rank * local, (rank + 1) * local

326 sends = [[] for _ in range(size)]

327 document = bisect_right(boundaries, first) - 1

328 while document < len(lengths) and boundaries[document] < last:

329 for peer, start, count in _zigzag_intersections(

330 lengths[document], size, pooled[document], capacity,

331 max(0, first - boundaries[document]), min(lengths[document], last - boundaries[document])):

332 sends[peer].append((boundaries[document] + start - first, count))

333 document += 1

334 return tuple(tuple(ranges) for ranges in sends)

335

336

337@dataclass(frozen=True)

338class IndexerRoute:

339 """Device-independent row ordering and matching variable collective splits."""

340

341 segments: tuple

342 input_splits: tuple[int, ...]

343 output_splits: tuple[int, ...]

344 send_rows: tuple[int, ...]

345 receive_rows: tuple[int, ...]

346 local_sources: tuple[int, ...]

347 local_targets: tuple[int, ...]

348 query_count: int

349

350

351@lru_cache(maxsize=128)

352def indexer_route(lengths: tuple[int, ...], size: int, rank: int, ratio: int, mode: str,

353 block_size: int | None = None, zigzag_min_length: int | None = None) -> IndexerRoute:

354 """Keep self-owned tasks local and communicate the remaining rows in owner order.

355

356 Args:

357 lengths: Shared document lengths in raw tokens.

358 size: Number of ranks in the CP group.

359 rank: Current rank within the CP group.

360 ratio: Compression ratio used by the Indexer.

361 mode: Zigzag, greedy or hybrid task allocation.

362 block_size: Maximum initial greedy-pool task size in Q rows.

363 zigzag_min_length: Hybrid-only long-document threshold.

364

365 Returns:

366 Immutable input/result route with a local-task bypass.

367 """

368 if not 0 <= rank < size:

369 raise ValueError("Indexer balance rank must be inside its CP group")

370 if mode == "zigzag":

371 _validate_task_plan(lengths, size, ratio, mode, block_size)

372 validate_indexer_hybrid_threshold(mode, zigzag_min_length)

373 boundaries = (0, *accumulate(lengths))

374 tasks = _zigzag_rank_tasks(lengths, size, rank)

375 sends = _zigzag_send_ranges(lengths, size, rank, boundaries)

376 receives = _receive_ranges(tasks, boundaries, size, boundaries[-1] // size)

377 return _route_from_ranges(tasks, lengths, rank, ratio, sends, receives)

378 assignments = indexer_task_plan(lengths, size, ratio, mode, block_size,

379 zigzag_min_length)

380 return _route_from_tasks(assignments, lengths, rank, ratio)

381

382

383def _route_from_tasks(assignments: tuple, lengths: tuple[int, ...], rank: int, ratio: int) -> IndexerRoute:

384 size = len(assignments)

385 boundaries, local = (0, *accumulate(lengths)), sum(lengths) // size

386 receives = _receive_ranges(assignments[rank], boundaries, size, local)

387 sends = tuple(_send_ranges(tasks, boundaries, rank, local) if peer != rank else ()

388 for peer, tasks in enumerate(assignments))

389 return _route_from_ranges(assignments[rank], lengths, rank, ratio, sends, receives)

390

391

392def _route_from_ranges(tasks: tuple, lengths: tuple[int, ...], rank: int, ratio: int,

393 sends: tuple, receives: tuple) -> IndexerRoute:

394 """Preserve peer ordering and local bypass for either task-planning path."""

395 send_rows, receive_rows, inputs, outputs = [], [], [], []

396 for peer, sent in enumerate(sends):

397 sent = sent if peer != rank else ()

398 received = receives[peer] if peer != rank else ()

399 inputs.append(sum(count for _, count in sent))

400 outputs.append(sum(count for _, _, count in received))

401 send_rows.extend(row for start, count in sent for row in range(start, start + count))

402 receive_rows.extend(row for _, start, count in received for row in range(start, start + count))

403 sources = tuple(row for start, _, count in receives[rank] for row in range(start, start + count))

404 targets = tuple(row for _, start, count in receives[rank] for row in range(start, start + count))

405 return IndexerRoute(_segments(tasks, lengths, ratio), tuple(inputs), tuple(outputs),

406 tuple(send_rows), tuple(receive_rows), sources, targets,

407 sum(task.length for task in tasks))

408

409

410@lru_cache(maxsize=32)

411@torch.inference_mode(False)

412def _device_rows(route: IndexerRoute, device: torch.device) -> tuple:

413 return tuple(torch.tensor(rows, device=device, dtype=torch.long) for rows in

414 (route.send_rows, route.receive_rows, route.local_sources, route.local_targets))

415

416

417def _exchange_rows(tensor: torch.Tensor, route: IndexerRoute, group: Any, rows: tuple) -> torch.Tensor:

418 """Exchange remote rows and place local bypass rows in worker task order."""

419 send_rows, receive_rows, local_sources, local_targets = rows

420 send = tensor.index_select(0, send_rows).contiguous()

421 received = tensor.new_empty((sum(route.output_splits), *tensor.shape[1:]))

422 work = dist.all_to_all_single(received, send, list(route.output_splits), list(route.input_splits),

423 group=group, async_op=True)

424 output = tensor.new_empty((route.query_count, *tensor.shape[1:]))

425 output.index_copy_(0, local_targets, tensor.index_select(0, local_sources))

426 work.wait()

427 output.index_copy_(0, receive_rows, received)

428 return output

429

430

431def _exchange_inputs(inputs: tuple[torch.Tensor, ...], route: IndexerRoute, group: Any, rows: tuple) -> tuple:

432 """Pack fields of equal dtype together without rounding the original RoPE values."""

433 by_dtype = {}

434 for index, tensor in enumerate(inputs):

435 by_dtype.setdefault(tensor.dtype, []).append(index)

436 outputs = [None] * len(inputs)

437 for indices in by_dtype.values():

438 widths = [inputs[index].shape[-1] for index in indices]

439 packed = torch.cat([inputs[index][0] for index in indices], dim=-1)

440 received = _exchange_rows(packed, route, group, rows)

441 for index, field in zip(indices, received.split(widths, dim=-1)):

442 outputs[index] = field.unsqueeze(0).contiguous()

443 return tuple(outputs)

444

445

446@torch.no_grad()

447def balanced_indexer_select(query_residual: torch.Tensor, weights: torch.Tensor,

448 position_embeddings: tuple[torch.Tensor, torch.Tensor], selector: Callable,

449 lengths: tuple[int, ...], ratio: int, mode: str, group: Any,

450 block_size: int | None = None, zigzag_min_length: int | None = None) -> torch.Tensor:

451 """Redistribute low-rank selection inputs and return TopK to the original owners.

452

453 Args:

454 query_residual: Local low-rank query residual [1, local tokens, q_lora_rank].

455 weights: Local merge weights [1, local tokens, heads].

456 position_embeddings: Original per-token cos/sin, each [1, local tokens, rotary/2].

457 All payload dtypes are preserved; tensors of equal dtype share a collective.

458 selector: Callable (qr, weights, (cos, sin), segments) that projects Indexer Q

459 with the current layer parameters and selects against the planned Index K bank.

460 It runs without autograd; owner-side Q and weights retain their KL gradients.

461 lengths: Global original document lengths; no per-document CP padding is added.

462 ratio: Compression ratio used by this selection layer.

463 mode: 'zigzag', 'greedy' or 'hybrid', identical on all CP ranks.

464 group: Original sequence-ordered CP process group.

465 block_size: Greedy task length in original query rows; see indexer_task_plan.

466 zigzag_min_length: Hybrid-only minimum original document length.

467 """

468 size, rank = dist.get_world_size(group), dist.get_rank(group)

469 inputs = (query_residual, weights, *position_embeddings)

470 if (len(position_embeddings) != 2 or query_residual.ndim != 3 or query_residual.shape[0] != 1

471 or query_residual.shape[1] * size != sum(lengths)):

472 raise ValueError("Indexer balance requires batch one and equal contiguous source query shards")

473 if any(tensor.ndim != 3 or tensor.shape[:2] != query_residual.shape[:2]

474 or tensor.device != query_residual.device for tensor in inputs):

475 raise ValueError("Indexer qr, weights and rotary embeddings must have matching batch/token axes and device")

476 route = indexer_route(lengths, size, rank, ratio, mode, block_size,

477 zigzag_min_length)

478 rows = _device_rows(route, query_residual.device)

479 balanced_qr, balanced_weights, cos, sin = _exchange_inputs(inputs, route, group, rows)

480 selected = selector(balanced_qr, balanced_weights, (cos, sin), route.segments)[0]

481 send_rows, receive_rows, local_sources, local_targets = rows

482 send = selected.index_select(0, receive_rows).contiguous()

483 received = selected.new_empty((sum(route.input_splits), selected.shape[-1]))

484 work = dist.all_to_all_single(received, send, list(route.input_splits), list(route.output_splits),

485 group=group, async_op=True)

486 output = selected.new_empty((query_residual.shape[1], selected.shape[-1]))

487 output.index_copy_(0, local_sources, selected.index_select(0, local_targets))

488 work.wait()

489 output.index_copy_(0, send_rows, received)

490 return output.unsqueeze(0)

491

492

493@dataclass

494class _PendingRows:

495 """Own communication buffers until the receive is safe to consume."""

496

497 send: torch.Tensor | None

498 received: torch.Tensor | None

499 output: torch.Tensor | None

500 receive_rows: torch.Tensor

501 work: Any

502 consumed: bool = False

503

504 @torch.no_grad()

505 def wait(self) -> torch.Tensor:

506 """Wait once, then restore received rows and release transfer storage."""

507 if self.consumed:

508 raise RuntimeError("Indexer communication handle cannot be consumed twice")

509 self.consumed = True

510 self.work.wait()

511 output = self.output

512 output.index_copy_(0, self.receive_rows, self.received)

513 self.send = self.received = self.output = self.work = None

514 return output

515

516

517@dataclass

518class IndexerInputExchange:

519 """One balanced query dispatch with original routing and dtype groups."""

520

521 route: IndexerRoute

522 rows: tuple

523 group: Any

524 owner_queries: int

525 transfers: tuple

526 fields: tuple

527 widths: tuple

528 consumed: bool = False

529

530 @torch.no_grad()

531 def wait(self) -> tuple:

532 """Restore qr, weights and original RoPE after all input launches."""

533 if self.consumed:

534 raise RuntimeError("Indexer input exchange cannot be consumed twice")

535 self.consumed = True

536 outputs = [None] * len(self.widths)

537 for transfer, indices in zip(self.transfers, self.fields):

538 merged = transfer.wait()

539 for index, value in zip(indices, merged.split([self.widths[i] for i in indices], dim=-1)):

540 outputs[index] = value.unsqueeze(0).contiguous()

541 self.transfers = ()

542 return tuple(outputs)

543

544 @torch.no_grad()

545 def launch_result(self, selected: torch.Tensor) -> _PendingRows:

546 """Return selection IDs while owner-side differentiable projections run."""

547 send_rows, receive_rows, local_sources, local_targets = self.rows

548 send = selected.index_select(0, receive_rows).contiguous()

549 received = selected.new_empty((sum(self.route.input_splits), selected.shape[-1]))

550 work = dist.all_to_all_single(

551 received, send, list(self.route.input_splits), list(self.route.output_splits),

552 group=self.group, async_op=True,

553 )

554 output = selected.new_empty((self.owner_queries, selected.shape[-1]))

555 output.index_copy_(0, local_sources, selected.index_select(0, local_targets))

556 return _PendingRows(send, received, output, send_rows, work)

557

558

559@torch.no_grad()

560def launch_indexer_inputs(

561 query_residual: torch.Tensor, weights: torch.Tensor, embeddings: tuple,

562 lengths: tuple[int, ...], ratio: int, mode: str, group: Any,

563 block_size: int | None = None, zigzag_min_length: int | None = None,

564) -> IndexerInputExchange:

565 """Launch all selection input fields without waiting for worker readiness.

566

567 Args:

568 query_residual: Owner qr, shaped [1, local Q, q_lora_rank].

569 weights: Owner Indexer merge weights.

570 embeddings: Original cos/sin; retain their dtypes and positions.

571 lengths: Global document lengths for contiguous CP ownership.

572 ratio: Key compression ratio.

573 mode: Existing zigzag, greedy or hybrid task assignment.

574 group: CP process group.

575 block_size: Greedy task length in original query rows.

576 zigzag_min_length: Hybrid long-document threshold.

577

578 Returns:

579 Single-use handle retaining transfer storage until wait.

580 """

581 size, rank = dist.get_world_size(group), dist.get_rank(group)

582 inputs = (query_residual, weights, *embeddings)

583 if (len(embeddings) != 2 or query_residual.ndim != 3 or query_residual.shape[0] != 1

584 or query_residual.shape[1] * size != sum(lengths)):

585 raise ValueError("Indexer balance requires batch one and equal contiguous source query shards")

586 if any(value.ndim != 3 or value.shape[:2] != query_residual.shape[:2]

587 or value.device != query_residual.device for value in inputs):

588 raise ValueError("Indexer qr, weights and rotary embeddings must share token axes and device")

589 route = indexer_route(lengths, size, rank, ratio, mode, block_size,

590 zigzag_min_length)

591 rows = _device_rows(route, query_residual.device)

592 send_rows, receive_rows, local_sources, local_targets = rows

593 grouped = {}

594 for index, value in enumerate(inputs):

595 grouped.setdefault(value.dtype, []).append(index)

596 transfers, fields = [], []

597 for indices in grouped.values():

598 packed = torch.cat([inputs[index][0] for index in indices], dim=-1)

599 send = packed.index_select(0, send_rows).contiguous()

600 received = packed.new_empty((sum(route.output_splits), packed.shape[-1]))

601 work = dist.all_to_all_single(

602 received, send, list(route.output_splits), list(route.input_splits), group=group, async_op=True,

603 )

604 output = packed.new_empty((route.query_count, packed.shape[-1]))

605 output.index_copy_(0, local_targets, packed.index_select(0, local_sources))

606 transfers.append(_PendingRows(send, received, output, receive_rows, work))

607 fields.append(tuple(indices))

608 return IndexerInputExchange(route, rows, group, query_residual.shape[1], tuple(transfers),

609 tuple(fields), tuple(value.shape[-1] for value in inputs))

Coverage for / home / jenkins / .local / lib / python3.10 / site-packages / hyper_parallel / distributed / context_parallel / indexer_balance.py: 73%

361 statements