Coverage for / home / jenkins / .local / lib / python3.10 / site-packages / hyper_parallel / distributed / context_parallel / compressed_document.py: 85%
124 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"""Document-local compression blocks and compact bank ownership for contiguous CP."""
17from __future__ import annotations
19from bisect import bisect_right
20from dataclasses import dataclass
21from functools import lru_cache
22from itertools import accumulate
23from typing import Any
25import torch # pylint: disable=forbidden-backend-import
26from torch.nn import functional # pylint: disable=forbidden-backend-import
29@dataclass(frozen=True)
30class CompressionBoundaryRoute:
31 """Projected rows needed to complete a block owned by its final-token rank."""
33 global_start: int
34 send_offsets: tuple[int, ...]
35 input_splits: tuple[int, ...]
36 output_splits: tuple[int, ...]
37 needs_exchange: bool
40def _boundary_route(local_length: int, rank: int, intervals: tuple) -> CompressionBoundaryRoute:
41 """Derive left-boundary exchanges; all ranks enter the same collective."""
42 local_start = rank * local_length
43 global_start, global_end = intervals[rank]
44 offsets, sends, receives = [], [], []
45 for peer, (start, end) in enumerate(intervals):
46 first = max(local_start, start)
47 count = max(0, min(local_start + local_length, end) - first) if peer != rank else 0
48 offsets.append(first - local_start if count else 0)
49 sends.append(count)
50 peer_start = peer * local_length
51 receives.append(max(0, min(peer_start + local_length, global_end) - max(peer_start, global_start))
52 if peer != rank else 0)
53 needs_exchange = any(start < peer * local_length for peer, (start, _) in enumerate(intervals))
54 return CompressionBoundaryRoute(global_start, tuple(offsets), tuple(sends), tuple(receives), needs_exchange)
57@lru_cache(maxsize=128)
58def compressed_document_offsets(boundaries: tuple[int, ...], ratio: int) -> tuple[int, ...]:
59 """Return canonical K boundaries, dropping only incomplete document tails.
61 Args:
62 boundaries: Cumulative raw document boundaries, beginning at zero.
63 ratio: Positive number of raw tokens per compressed block.
65 Returns:
66 Cumulative complete-block counts in document order, beginning at zero.
67 """
68 if ratio < 1 or len(boundaries) < 2 or boundaries[0] != 0 or any(
69 end <= start for start, end in zip(boundaries, boundaries[1:])):
70 raise ValueError("Compression requires a positive ratio and strictly increasing document boundaries")
71 return (0, *accumulate((end-start)//ratio for start, end in zip(boundaries, boundaries[1:])))
74@lru_cache(maxsize=128)
75def compressed_owner_offsets(boundaries: tuple[int, ...], ratio: int, local_length: int) -> tuple[int, ...]:
76 """Assign each complete block to the raw rank containing its last token.
78 Args:
79 boundaries: Cumulative raw document boundaries, beginning at zero.
80 ratio: Positive number of raw tokens per compressed block.
81 local_length: Equal positive raw-token count on each CP rank.
83 Returns:
84 Cumulative compressed-bank offsets for each rank and the final bank end.
85 """
86 keys = compressed_document_offsets(boundaries, ratio)
87 if local_length <= 0 or boundaries[-1] % local_length:
88 raise ValueError("Compression requires equal positive raw CP shard lengths")
89 offsets = []
90 for position in range(0, boundaries[-1], local_length):
91 document = bisect_right(boundaries, position) - 1
92 offsets.append(keys[document] + (position-boundaries[document])//ratio)
93 return (*offsets, keys[-1])
96@dataclass(frozen=True)
97class DocumentCompressionPlan:
98 """Immutable block ranges; never enumerate the full sequence on the CPU."""
100 ratio: int
101 local_length: int
102 rank: int
103 counts: tuple[int, ...]
104 local_blocks: tuple[tuple[int, int], ...]
105 route: CompressionBoundaryRoute
107 def indices(self, device: torch.device) -> tuple[torch.Tensor, torch.Tensor]:
108 """Reuse device indices for local blocks and their original rotary positions.
110 Args:
111 device: Device on which the row indices are consumed.
113 Returns:
114 Flattened pooling rows and block-start rows, relative to exchanged input.
115 """
116 return _compression_indices(self, device)
118 def exchange(self, tensor: torch.Tensor, context: Any) -> torch.Tensor:
119 """Complete boundary blocks with the existing differentiable A2AV protocol.
121 Args:
122 tensor: Local projected values with raw tokens on dimension one.
123 context: CP boundary-dispatch context, or None when no exchange is needed.
125 Returns:
126 Required preceding rows followed by the local rows, preserving autograd.
127 """
128 if not self.route.needs_exchange:
129 return tensor
130 if context is None or context.exchange_boundary is None:
131 raise ValueError("natural document compression across CP requires boundary dispatch")
132 return context.exchange_boundary(tensor, self.route)
134 def pool(self, key: torch.Tensor, gate: torch.Tensor, context: Any) -> torch.Tensor:
135 """Pool complete same-document groups using the model's original FP32 softmax.
137 Args:
138 key: Local projected values shaped [batch, raw tokens, channels].
139 gate: Pooling logits with the same shape as key.
140 context: CP boundary-dispatch context, or None when no exchange is needed.
142 Returns:
143 Pooled values shaped [batch, locally owned complete blocks, channels].
144 """
145 rows = self.indices(key.device)[0]
146 projected = self.exchange(torch.cat((key, gate), dim=-1), context)
147 blocks = projected.index_select(1, rows).reshape(
148 key.shape[0], self.counts[self.rank], self.ratio, projected.shape[-1],
149 )
150 values, scores = blocks.chunk(2, dim=-1)
151 return (values * scores.float().softmax(dim=2).to(values.dtype)).sum(dim=2)
153 def positions(self, cos: torch.Tensor, sin: torch.Tensor, context: Any) -> tuple:
154 """Take original block-start RoPE; retain FP32 boundary metadata precision.
156 Args:
157 cos: Local raw-token rotary cosine coefficients.
158 sin: Local raw-token rotary sine coefficients with the same shape as cos.
159 context: CP boundary-dispatch context, or None when no exchange is needed.
161 Returns:
162 Cosine and sine coefficients for each locally owned block's first token.
163 """
164 positions = self.indices(cos.device)[1]
165 embeddings = self.exchange(torch.cat((cos, sin), dim=-1), context)
166 return embeddings.index_select(1, positions).chunk(2, dim=-1)
168 def launch_bank(self, tensor: torch.Tensor, context: Any) -> CompactBankGather:
169 """All-gather equal-capacity containers and exclude every padding slot.
171 Args:
172 tensor: Locally owned compressed bank with blocks on dimension one.
173 context: CP context providing an asynchronous sequence all-gather.
175 Returns:
176 Handle whose wait method returns the global bank without padding slots.
177 """
178 capacity = max(self.counts)
179 if not capacity:
180 return CompactBankGather(None, tensor, None)
181 padded = tensor if tensor.shape[1] == capacity else functional.pad(
182 tensor, (0, 0, 0, capacity - tensor.shape[1]),
183 )
184 slots = _bank_slots(self.counts, tensor.device)
185 return CompactBankGather(context.launch(padded, 1), tensor, slots)
188@dataclass
189class CompactBankGather:
190 """Complete one bank gather with at most one compacting index selection."""
192 handle: Any
193 tensor: torch.Tensor
194 slots: torch.Tensor | None
196 def wait(self) -> torch.Tensor:
197 """Keep compaction in autograd so all consumers return gradients to owners.
199 Returns:
200 Complete compressed bank in canonical document order, without padding.
201 """
202 output = self.tensor if self.handle is None else self.handle.wait()
203 return output if self.slots is None else output.index_select(1, self.slots)
206@lru_cache(maxsize=128)
207def document_compression_plan(boundaries: tuple[int, ...], ratio: int, local_length: int,
208 size: int, rank: int) -> DocumentCompressionPlan:
209 """Assign complete blocks to the rank owning their final token.
211 Work on document/rank intersections rather than all global token rows.
212 The packed-sequence boundary validator owns interval validation.
214 Args:
215 boundaries: Cumulative raw document boundaries covering the global sequence.
216 ratio: Number of raw tokens per compressed block, greater than one.
217 local_length: Equal positive raw-token count on each CP rank.
218 size: Number of contiguous CP shards.
219 rank: Local rank within the CP group.
221 Returns:
222 Complete-block ownership, local block ranges, and boundary-exchange routing.
223 """
224 if ratio <= 1 or local_length <= 0 or not 0 <= rank < size:
225 raise ValueError("document block planning requires ratio > 1 and a valid contiguous CP shard")
226 if boundaries[0] != 0 or boundaries[-1] != size * local_length:
227 raise ValueError("document boundaries must cover the complete contiguous CP sequence")
228 owners = compressed_owner_offsets(boundaries, ratio, local_length)
229 counts = tuple(end-start for start, end in zip(owners, owners[1:]))
230 first_rows = [peer * local_length for peer in range(size)]
231 local_blocks = []
232 for start, end in zip(boundaries, boundaries[1:]):
233 blocks = (end - start) // ratio
234 if not blocks:
235 continue
236 first_owner = (start + ratio - 1) // local_length
237 last_owner = (start + blocks * ratio - 1) // local_length
238 for peer in range(first_owner, last_owner + 1):
239 first = max(0, (peer * local_length - start) // ratio)
240 last = min(blocks, ((peer + 1) * local_length - start) // ratio)
241 if last > first:
242 first_rows[peer] = min(first_rows[peer], start + first * ratio)
243 if peer == rank:
244 local_blocks.append((start + first * ratio, last - first))
245 intervals = tuple((first, (peer + 1) * local_length) for peer, first in enumerate(first_rows))
246 return DocumentCompressionPlan(ratio, local_length, rank, tuple(counts), tuple(local_blocks),
247 _boundary_route(local_length, rank, intervals))
250@lru_cache(maxsize=8)
251@torch.inference_mode(False)
252def _compression_indices(plan: DocumentCompressionPlan, device: torch.device) -> tuple:
253 """Reuse block rows; bound the cache since compact AG maps scale with global K."""
254 starts = [torch.arange(start, start + count * plan.ratio, plan.ratio, device=device)
255 for start, count in plan.local_blocks]
256 starts = torch.cat(starts) if starts else torch.empty(0, device=device, dtype=torch.long)
257 starts = starts - plan.route.global_start
258 rows = (starts.unsqueeze(1) + torch.arange(plan.ratio, device=device)).flatten()
259 return rows.long(), starts.long()
262@lru_cache(maxsize=8)
263@torch.inference_mode(False)
264def _bank_slots(counts: tuple[int, ...], device: torch.device) -> torch.Tensor | None:
265 """Materialize global padding maps only for the all-gather reference path."""
266 if len(set(counts)) == 1:
267 return None
268 capacity = max(counts)
269 return torch.cat([torch.arange(peer*capacity, peer*capacity+count, device=device)
270 for peer, count in enumerate(counts)])