Coverage for / home / jenkins / .local / lib / python3.10 / site-packages / hyper_parallel / distributed / context_parallel / sequence_dispatch.py: 96%
117 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"""Differentiable all-to-all-v dispatch of sequence bank intervals."""
17from __future__ import annotations
19from bisect import bisect_right
20from dataclasses import dataclass
21from functools import lru_cache
22from typing import Any
24import torch # pylint: disable=forbidden-backend-import
25import torch.distributed as dist # pylint: disable=forbidden-backend-import
27from hyper_parallel.distributed.context_parallel.compressed_document import (
28 compressed_document_offsets, compressed_owner_offsets,
29)
32@dataclass(frozen=True)
33class SequenceDispatchRoute:
34 """Row splits for compact intervals assembled from contiguous owner shards.
36 Global start/end bound the requested ranges, including gaps when present.
37 Self rows stay local. Remote rows are ordered by owner rank, so the local
38 shard fits between the received prefix and suffix without tensor indices.
39 Backward reverses the exchange and SUMs overlapping consumer contributions.
40 """
42 local_length: int
43 global_start: int
44 global_end: int
45 local_output_start: int
46 send_offsets: tuple[int, ...]
47 input_splits: tuple[int, ...]
48 output_splits: tuple[int, ...]
49 needs_exchange: bool
51 def owner_ranges(self) -> tuple[tuple[int, int], ...]:
52 """Return send slices in destination-rank order, excluding self rows."""
53 return tuple(zip(self.send_offsets, self.input_splits))
56def _validate_interval_owners(owners: tuple[int, ...], rank: int, intervals: tuple) -> None:
57 """Keep each complete owner shard in its consumer bank, including empty shards."""
58 if (len(owners) != len(intervals) + 1 or not intervals or owners[0] != 0
59 or not 0 <= rank < len(intervals) or any(end < start for start, end in zip(owners, owners[1:]))):
60 raise ValueError("Dispatch requires ordered owner offsets and a valid rank")
61 for peer, (start, end) in enumerate(intervals):
62 if not 0 <= start <= owners[peer] <= owners[peer + 1] <= end <= owners[-1]:
63 raise ValueError("Dispatch intervals must cover their complete owner shard within the global bank")
66def interval_route(
67 local_length: int, rank: int, intervals: tuple[tuple[int, int], ...],
68 *, owner_offsets: tuple[int, ...] | None = None,
69) -> SequenceDispatchRoute:
70 """Derive owner/consumer intersections; all ranks agree whether to exchange.
72 Args:
73 local_length: Actual row count owned by the current rank, possibly zero.
74 rank: Current rank in sequence-shard order.
75 intervals: One consumer interval per rank, containing its complete owner shard.
76 owner_offsets: Canonical source boundaries; omitted for equal-sized raw shards.
78 Returns:
79 Per-peer row splits and the unchanged owner's position in its consumer bank.
80 """
81 owners = (tuple(peer * local_length for peer in range(len(intervals) + 1))
82 if owner_offsets is None else owner_offsets)
83 _validate_interval_owners(owners, rank, intervals)
84 if owners[rank + 1] - owners[rank] != local_length:
85 raise ValueError("Local bank length differs from the supplied owner offsets")
86 local_start = owners[rank]
87 global_start, global_end = intervals[rank]
88 offsets, sends, receives = [], [], []
89 for peer, (start, end) in enumerate(intervals):
90 first = max(local_start, start)
91 count = max(0, min(owners[rank + 1], end) - first) if peer != rank else 0
92 offsets.append(first - local_start if count else 0)
93 sends.append(count)
94 receives.append(max(0, min(owners[peer + 1], global_end) - max(owners[peer], global_start))
95 if peer != rank else 0)
96 needs_exchange = any(start != owners[peer] or end != owners[peer + 1]
97 for peer, (start, end) in enumerate(intervals))
98 return SequenceDispatchRoute(local_length, global_start, global_end, local_start - global_start,
99 tuple(offsets), tuple(sends), tuple(receives), needs_exchange)
102class _SequenceDispatchWait(torch.autograd.Function):
103 """Attach remote bank rows to their owners without retaining send buffers."""
105 @staticmethod
106 def forward(ctx: Any, local: torch.Tensor, received: torch.Tensor, handle: Any,
107 route: SequenceDispatchRoute, group: Any) -> torch.Tensor:
108 """Wait before exposing received data and insert the uncommunicated local rows.
110 Args:
111 ctx: Autograd context retaining the route and process group.
112 local: Owner shard with its sequence dimension moved to axis zero.
113 received: Remote rows whose writes finish when handle completes.
114 handle: Outstanding all-to-all-v work handle.
115 route: Per-peer splits and local placement for this bank.
116 group: Framework-owned CP process group.
117 """
118 ctx.route = route
119 ctx.group = group
120 handle.wait()
121 before = route.local_output_start
122 return torch.cat((received[:before], local, received[before:]), dim=0)
124 @staticmethod
125 def backward(ctx: Any, grad: torch.Tensor) -> tuple:
126 """Reverse dispatch, SUMming consumers in rank order, including empty ranks.
128 Args:
129 ctx: Saved dispatch route and CP process group.
130 grad: Consumer gradient in the assembled bank's row order.
131 """
132 route = ctx.route
133 first, last = route.local_output_start, route.local_output_start + route.local_length
134 send = torch.cat((grad[:first], grad[last:]), dim=0)
135 received = grad.new_empty((sum(route.input_splits), *grad.shape[1:]))
136 handle = dist.all_to_all_single(
137 received, send, output_split_sizes=list(route.input_splits),
138 input_split_sizes=list(route.output_splits), group=ctx.group, async_op=True,
139 )
140 local = grad[first:last].clone(memory_format=torch.contiguous_format)
141 handle.wait()
142 cursor = 0
143 for offset, count in route.owner_ranges():
144 if count:
145 local.narrow(0, offset, count).add_(received.narrow(0, cursor, count))
146 cursor += count
147 return local, None, None, None, None
150@dataclass
151class AsyncSequenceDispatch:
152 """Single-use launch handle owning send storage until the communication completes."""
154 local: torch.Tensor
155 send: torch.Tensor | None
156 received: torch.Tensor | None
157 handle: Any
158 route: SequenceDispatchRoute
159 group: Any
160 sequence_dim: int
161 waited: bool = False
163 def wait(self) -> torch.Tensor:
164 """Finish the exchange and return the compact bank in its original layout."""
165 if self.waited:
166 raise RuntimeError("a sequence dispatch handle cannot be waited twice")
167 self.waited = True
168 output = self.local
169 if self.received is not None:
170 output = _SequenceDispatchWait.apply(
171 self.local, self.received, self.handle, self.route, self.group,
172 )
173 self.send = None
174 self.received = None
175 return output.movedim(0, self.sequence_dim).contiguous()
178def async_sequence_dispatch(tensor: torch.Tensor, sequence_dim: int,
179 route: SequenceDispatchRoute, group: Any) -> AsyncSequenceDispatch:
180 """Launch variable row exchange with an autograd edge attached at wait.
182 Args:
183 tensor: Contiguous local owner shard of a raw, compressed or Index K bank.
184 sequence_dim: Sequence dimension; feature dimensions are transported unchanged.
185 route: Immutable per-batch routing agreed by every member of the CP group.
186 group: Existing framework-owned CP process group, in sequence-shard rank order.
187 """
188 if not -tensor.ndim <= sequence_dim < tensor.ndim:
189 raise ValueError("sequence dispatch dimension is outside the input tensor dimensions")
190 sequence_dim %= tensor.ndim
191 if tensor.shape[sequence_dim] != route.local_length:
192 raise ValueError("sequence dispatch owner length does not match the bank's physical layout")
193 local = tensor.movedim(sequence_dim, 0).contiguous()
194 if not route.needs_exchange:
195 return AsyncSequenceDispatch(local, None, None, None, route, group, sequence_dim)
196 parts = [local.detach().narrow(0, offset, count)
197 for offset, count in route.owner_ranges() if count]
198 send = torch.cat(parts, dim=0) if parts else local.detach()[:0]
199 received = local.new_empty((sum(route.output_splits), *local.shape[1:]))
200 handle = dist.all_to_all_single(
201 received, send, output_split_sizes=list(route.output_splits),
202 input_split_sizes=list(route.input_splits), group=group, async_op=True,
203 )
204 return AsyncSequenceDispatch(local, send, received, handle, route, group, sequence_dim)
207@lru_cache(maxsize=128)
208def sliding_window_route(boundaries: tuple[int, ...], local_length: int, size: int,
209 rank: int, window: int) -> SequenceDispatchRoute:
210 """Plan the same-document left halo; every rank agrees whether to exchange."""
211 if local_length <= 0 or size <= 0 or not 0 <= rank < size or window <= 0:
212 raise ValueError("SWA dispatch requires positive dimensions and a valid CP rank")
213 if (len(boundaries) < 2 or boundaries[0] != 0 or boundaries[-1] != local_length * size
214 or any(end <= start for start, end in zip(boundaries, boundaries[1:]))):
215 raise ValueError("SWA document boundaries must span the full contiguous CP sequence")
216 intervals = tuple(
217 (max(boundaries[bisect_right(boundaries, peer * local_length) - 1], peer * local_length - window + 1),
218 (peer + 1) * local_length)
219 for peer in range(size)
220 )
221 return interval_route(local_length, rank, intervals)
224@lru_cache(maxsize=128)
225def document_prefix_route(boundaries: tuple[int, ...], local_length: int, size: int,
226 rank: int, ratio: int) -> SequenceDispatchRoute:
227 """Route natural-document main KV prefixes using complete-block ownership."""
228 if ratio not in (1, 2) or local_length <= 0 or size <= 0 or not 0 <= rank < size:
229 raise ValueError("Main KV dispatch requires a supported ratio and equal positive raw CP shards")
230 if (len(boundaries) < 2 or boundaries[0] != 0 or boundaries[-1] != local_length * size
231 or any(end <= start for start, end in zip(boundaries, boundaries[1:]))):
232 raise ValueError("Main KV document boundaries must span the complete contiguous CP sequence")
233 keys = compressed_document_offsets(boundaries, ratio)
234 owners = compressed_owner_offsets(boundaries, ratio, local_length)
235 intervals = tuple(
236 (keys[bisect_right(boundaries, peer * local_length) - 1], owners[peer + 1])
237 for peer in range(size)
238 )
239 return interval_route(owners[rank + 1] - owners[rank], rank, intervals, owner_offsets=owners)