Coverage for / home / jenkins / .local / lib / python3.10 / site-packages / hyper_parallel / distributed / context_parallel / collectives.py: 35%
269 statements
« prev ^ index » next coverage.py v7.13.1, created at 2026-11-15 07:18 +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# ============================================================================
16"""context_parallel.collectives: CP collective primitives (05 §4.4.2 / §6.3.4).
18- ``flex_cp_allgather``: K/V all-gather along the CP dim (reuses
19 cp_mesh.get_group(); new_group is forbidden);
20- ``ulysses_seq_to_head`` / ``ulysses_head_to_seq``: differentiable Pure
21 Ulysses all-to-all layout transforms;
22- async collective launches (``AsyncCPCollective`` /
23 ``async_cp_allgather_launch`` / ``async_ulysses_seq_to_head_launch``);
24- hybrid (Ulysses + all-gather) CP sub-mesh construction and attention
25 orchestration (``_build_hybrid_cp_submeshes`` / ``hybrid_cp_attention``).
27Note (G5): the seq_len % (2*cp) padding constraint originates from the
28zigzag/ring load-balancing scheme; this design uses all-gather K/V +
29contiguous chunk (D-01'' rejected ring), so each rank's Q chunk is
30equal-length and FLOPs are naturally balanced -- the constraint is redundant
31but harmless: the implementation is kept and documented here.
33Split out of components/distributed/cp_utils.py in stage 4e; the attention
34adaptations live in context_parallel/attention.py and the batch sharding in
35hyper_parallel/data/parallel/batch_parallel.py.
36"""
38# The legacy header checker does not recognize the existing multi-year copyright.
39# pylint: disable=missing-apache-license-header
41from dataclasses import dataclass
42from typing import Any, Callable, Optional
43import torch # pylint: disable=forbidden-backend-import
44import torch.distributed as dist # pylint: disable=forbidden-backend-import
45from torch import Tensor # pylint: disable=forbidden-backend-import
46from torch.distributed.nn.functional import all_gather as differentiable_all_gather # pylint: disable=forbidden-backend-import
47from hyper_parallel.core.dtensor.device_mesh import DeviceMesh
48from hyper_parallel.core.utils import communication
51_ULYSSES_WRAPPED_FLAG = "_hyper_ulysses_wrapped"
54_HYBRID_MESH_CACHE = {}
57class _SeqAllToAll(torch.autograd.Function):
58 """Autograd-aware exchange from one tensor dimension to another."""
60 @staticmethod
61 def forward(
62 ctx: Any,
63 group: Any,
64 tensor: Tensor,
65 scatter_dim: int,
66 gather_dim: int,
67 ) -> Tensor:
68 """Exchange tensor shards between the sequence and head dimensions."""
69 ctx.group = group
70 ctx.scatter_dim = scatter_dim
71 ctx.gather_dim = gather_dim
72 world_size = dist.get_world_size(group)
73 if tensor.size(scatter_dim) % world_size:
74 raise ValueError(
75 f"dimension {scatter_dim} size {tensor.size(scatter_dim)} "
76 f"is not divisible by CP size {world_size}")
77 inputs = [
78 part.contiguous()
79 for part in tensor.chunk(world_size, dim=scatter_dim)
80 ]
81 outputs = [torch.empty_like(inputs[0]) for _ in range(world_size)]
82 dist.all_to_all(outputs, inputs, group=group)
83 return torch.cat(outputs, dim=gather_dim).contiguous()
85 @staticmethod
86 def backward(
87 ctx: Any,
88 grad_output: Tensor,
89 ) -> tuple[None, Tensor, None, None]:
90 """Apply the inverse all-to-all to the output gradient."""
91 grad_input = _SeqAllToAll.apply(
92 ctx.group, grad_output, ctx.gather_dim, ctx.scatter_dim)
93 return None, grad_input, None, None
96@dataclass(frozen=True)
97class _UlyssesContext:
98 """Communication context derived from the framework-owned CP mesh."""
100 cp_mesh: Any
102 @property
103 def group(self) -> Any:
104 """Process group of the CP mesh."""
105 return self.cp_mesh.get_group()
107 @property
108 def size(self) -> int:
109 """Number of ranks in the CP mesh."""
110 return self.cp_mesh.size()
112 @property
113 def rank(self) -> int:
114 """Local rank of this process within the CP mesh."""
115 return self.cp_mesh.get_local_rank()
118def _sequence_to_head(tensor, context):
119 """[B, S/CP, H, D] -> [B, S, H/CP, D]."""
120 return _SeqAllToAll.apply(context.group, tensor, 2, 1)
123def _head_to_sequence(tensor, context):
124 """[B, S, H/CP, D] -> [B, S/CP, H, D]."""
125 return _SeqAllToAll.apply(context.group, tensor, 1, 2)
128def _gather_sequence(tensor, context):
129 """Gather sequence shards and sum their gradients during backward."""
130 return torch.cat(
131 differentiable_all_gather(
132 tensor.contiguous(), group=context.group), dim=1)
135def _slice_sequence(tensor, dim, context):
136 if tensor.size(dim) % context.size:
137 raise ValueError(
138 f"sequence length {tensor.size(dim)} is not divisible by "
139 f"CP size {context.size}")
140 length = tensor.size(dim) // context.size
141 return tensor.narrow(
142 dim, context.rank * length, length).contiguous()
145def _global_seq_len(actual, length, device):
146 if actual is None:
147 return [length]
148 if isinstance(actual, Tensor):
149 return actual.to(device=device, dtype=torch.int32)
150 return actual
153def _slice_sink(tensor, context):
154 if tensor is None or tensor.size(2) == 1:
155 return tensor
156 if tensor.size(2) % context.size:
157 raise ValueError(
158 f"parameter-sink heads {tensor.size(2)} are not divisible by "
159 f"CP size {context.size}")
160 count = tensor.size(2) // context.size
161 start = context.rank * count
162 return tensor[:, :, start:start + count].contiguous()
165def _normalize_dim(dim: int, ndim: int) -> int:
166 """Normalize and validate a tensor dimension."""
167 normalized = dim + ndim if dim < 0 else dim
168 if normalized < 0 or normalized >= ndim:
169 raise ValueError(f"dimension {dim} is out of range for a {ndim}D tensor")
170 return normalized
173def _reconstruct_all_to_all(output: Tensor, concat_dim: int) -> Tensor:
174 """Move the source-rank axis next to ``concat_dim`` and merge it."""
175 rank_axis_position = concat_dim + 1
176 permutation = (
177 list(range(1, rank_axis_position))
178 + [0]
179 + list(range(rank_axis_position, output.dim()))
180 )
181 reconstructed = output.permute(permutation).contiguous()
182 shape = list(reconstructed.shape)
183 merged = shape[concat_dim] * shape[concat_dim + 1]
184 return reconstructed.reshape(shape[:concat_dim] + [merged] + shape[concat_dim + 2:])
187def _move_dim_to_front(tensor: Tensor, dim: int) -> Tensor:
188 """Move one dimension to the leading communication dimension."""
189 dim = _normalize_dim(dim, tensor.dim())
190 if dim == 0:
191 return tensor.contiguous()
192 return tensor.movedim(dim, 0).contiguous()
195def _move_dim_from_front(tensor: Tensor, dim: int) -> Tensor:
196 """Restore a leading communication dimension to its original position."""
197 dim = _normalize_dim(dim, tensor.dim())
198 if dim == 0:
199 return tensor.contiguous()
200 return tensor.movedim(0, dim).contiguous()
203def _prepare_ulysses_send(
204 tensor: Tensor,
205 scatter_dim: int,
206 world_size: int,
207) -> Tensor:
208 """Split one tensor dimension and move the destination-rank axis first."""
209 ndim = tensor.dim()
210 scatter_dim = _normalize_dim(scatter_dim, ndim)
211 shape = list(tensor.shape)
212 scatter_size = shape[scatter_dim]
213 if scatter_size % world_size:
214 raise ValueError(
215 f"tensor dimension {scatter_dim} size ({scatter_size}) must be "
216 f"divisible by Ulysses degree ({world_size})"
217 )
218 split_shape = (
219 shape[:scatter_dim]
220 + [world_size, scatter_size // world_size]
221 + shape[scatter_dim + 1:]
222 )
223 permutation = (
224 [scatter_dim]
225 + list(range(scatter_dim))
226 + list(range(scatter_dim + 1, ndim + 1))
227 )
228 return tensor.contiguous().reshape(split_shape).permute(permutation).contiguous()
231class _AsyncAllGatherWait(torch.autograd.Function):
232 """Attach a pre-launched all-gather to autograd at its wait point."""
234 @staticmethod
235 def forward(
236 ctx: Any,
237 tensor: Tensor,
238 work: Any,
239 output: Tensor,
240 group: Any,
241 world_size: int,
242 gather_dim: int,
243 ) -> Tensor:
244 """Wait for a launched all-gather and expose its autograd edge."""
245 del tensor
246 ctx.group = group
247 ctx.world_size = world_size
248 ctx.gather_dim = gather_dim
249 work.wait()
250 return _move_dim_from_front(output, gather_dim)
252 @staticmethod
253 def backward(
254 ctx: Any,
255 grad_output: Tensor,
256 ) -> tuple[Tensor, None, None, None, None, None]:
257 """SUM consumer gradients directly into each owner's sequence shard."""
258 send = _move_dim_to_front(grad_output, ctx.gather_dim)
259 output = grad_output.new_empty((send.shape[0] // ctx.world_size, *send.shape[1:]))
260 work = dist.reduce_scatter_tensor(
261 output, send, op=dist.ReduceOp.SUM, group=ctx.group, async_op=True,
262 )
263 work.wait()
264 local = _move_dim_from_front(output, ctx.gather_dim)
265 return (
266 local,
267 None,
268 None,
269 None,
270 None,
271 None,
272 )
275class _AsyncUlyssesWait(torch.autograd.Function):
276 """Attach a pre-launched sequence-to-head A2A to autograd."""
278 @staticmethod
279 def forward(
280 ctx: Any,
281 tensor: Tensor,
282 work: Any,
283 output: Tensor,
284 group: Any,
285 world_size: int,
286 seq_dim: int,
287 head_dim: int,
288 ) -> Tensor:
289 """Wait for a launched A2A and expose its autograd edge."""
290 del tensor
291 ctx.group = group
292 ctx.world_size = world_size
293 ctx.seq_dim = seq_dim
294 ctx.head_dim = head_dim
295 work.wait()
296 return _reconstruct_all_to_all(output, seq_dim)
298 @staticmethod
299 def backward(
300 ctx: Any,
301 grad_output: Tensor,
302 ) -> tuple[Tensor, None, None, None, None, None, None]:
303 """Apply the inverse A2A to restore the input-gradient layout."""
304 send = _prepare_ulysses_send(
305 grad_output,
306 scatter_dim=ctx.seq_dim,
307 world_size=ctx.world_size,
308 )
309 output = torch.empty(
310 list(send.shape), device=send.device, dtype=send.dtype
311 )
312 work = dist.all_to_all_single(
313 output, send, group=ctx.group, async_op=True
314 )
315 work.wait()
316 grad_input = _reconstruct_all_to_all(output, ctx.head_dim)
317 return grad_input, None, None, None, None, None, None
320@dataclass
321class AsyncCPCollective:
322 """A local-tensor CP collective launched before its consumer."""
324 tensor: Tensor
325 work: Any
326 output: Tensor
327 group: Any
328 world_size: int
329 kind: str
330 seq_dim: int
331 head_dim: Optional[int] = None
332 waited: bool = False
334 def wait(self) -> Tensor:
335 """Materialize the collective result exactly once."""
336 if self.waited:
337 raise RuntimeError("an async CP collective handle cannot be waited twice")
338 self.waited = True
339 if self.work is None:
340 return self.tensor
341 if self.kind == "allgather":
342 return _AsyncAllGatherWait.apply(
343 self.tensor,
344 self.work,
345 self.output,
346 self.group,
347 self.world_size,
348 self.seq_dim,
349 )
350 if self.kind == "ulysses":
351 return _AsyncUlyssesWait.apply(
352 self.tensor,
353 self.work,
354 self.output,
355 self.group,
356 self.world_size,
357 self.seq_dim,
358 self.head_dim,
359 )
360 raise RuntimeError(f"unknown async CP collective kind {self.kind!r}")
363def async_cp_allgather_launch(
364 tensor: Tensor,
365 gather_dim: int,
366 cp_mesh: DeviceMesh,
367) -> AsyncCPCollective:
368 """Launch an asynchronous all-gather along gather_dim."""
369 world_size = cp_mesh.size()
370 if world_size <= 1:
371 return AsyncCPCollective(
372 tensor, None, tensor, None, world_size, "allgather", gather_dim
373 )
374 gather_dim = _normalize_dim(gather_dim, tensor.dim())
375 send = _move_dim_to_front(tensor.detach(), gather_dim)
376 output_shape = list(send.shape)
377 output_shape[0] *= world_size
378 group = cp_mesh.get_group()
379 output = torch.empty(
380 output_shape, device=send.device, dtype=send.dtype
381 )
382 work = dist.all_gather_into_tensor(
383 output, send, group=group, async_op=True
384 )
385 return AsyncCPCollective(
386 tensor, work, output, group, world_size, "allgather", gather_dim
387 )
390def async_ulysses_seq_to_head_launch(
391 tensor: Tensor,
392 seq_dim: int,
393 head_dim: int,
394 cp_mesh: DeviceMesh,
395) -> AsyncCPCollective:
396 """Launch an asynchronous Ulysses sequence-to-head A2A."""
397 world_size = cp_mesh.size()
398 if world_size <= 1:
399 return AsyncCPCollective(
400 tensor,
401 None,
402 tensor,
403 None,
404 world_size,
405 "ulysses",
406 seq_dim,
407 head_dim,
408 )
409 seq_dim = _normalize_dim(seq_dim, tensor.dim())
410 head_dim = _normalize_dim(head_dim, tensor.dim())
411 if seq_dim == head_dim:
412 raise ValueError("Ulysses seq_dim and head_dim must be different")
413 send = _prepare_ulysses_send(
414 tensor.detach(), scatter_dim=head_dim, world_size=world_size
415 )
416 group = cp_mesh.get_group()
417 output = torch.empty(
418 list(send.shape), device=send.device, dtype=send.dtype
419 )
420 work = dist.all_to_all_single(
421 output, send, group=group, async_op=True
422 )
423 return AsyncCPCollective(
424 tensor,
425 work,
426 output,
427 group,
428 world_size,
429 "ulysses",
430 seq_dim,
431 head_dim,
432 )
435def _ulysses_all_to_all(
436 tensor: Tensor,
437 *,
438 scatter_dim: int,
439 gather_dim: int,
440 cp_mesh: DeviceMesh,
441) -> Tensor:
442 """Split one tensor dimension and gather another through all-to-all."""
443 cp_size = cp_mesh.size()
444 if cp_size <= 1:
445 return tensor
447 ndim = tensor.dim()
448 scatter_dim = _normalize_dim(scatter_dim, ndim)
449 gather_dim = _normalize_dim(gather_dim, ndim)
450 if scatter_dim == gather_dim:
451 raise ValueError("Ulysses scatter_dim and gather_dim must be different")
453 shape = list(tensor.shape)
454 scatter_size = shape[scatter_dim]
455 if scatter_size % cp_size != 0:
456 raise ValueError(
457 f"tensor dimension {scatter_dim} size ({scatter_size}) must be "
458 f"divisible by Ulysses degree ({cp_size})"
459 )
461 split_shape = (
462 shape[:scatter_dim]
463 + [cp_size, scatter_size // cp_size]
464 + shape[scatter_dim + 1:]
465 )
466 split_ndim = ndim + 1
467 permutation = (
468 [scatter_dim]
469 + list(range(scatter_dim))
470 + list(range(scatter_dim + 1, split_ndim))
471 )
472 send = tensor.contiguous().reshape(split_shape).permute(permutation).contiguous()
473 received = communication.differentiable_all_to_all(
474 send, list(send.shape), cp_mesh.get_group())
475 return _reconstruct_all_to_all(received, gather_dim)
478def ulysses_seq_to_head(
479 tensor: Tensor,
480 seq_dim: int,
481 head_dim: int,
482 cp_mesh: DeviceMesh,
483) -> Tensor:
484 """Convert a sequence shard into a full-sequence head shard.
486 Args:
487 tensor: Local Q, K, or V tensor.
488 seq_dim: Sequence dimension to gather across the CP group.
489 head_dim: Head dimension to split across the CP group.
490 cp_mesh: One-dimensional CP DeviceMesh.
492 Returns:
493 Tensor with the sequence dimension gathered and head dimension sharded.
494 """
495 return _ulysses_all_to_all(
496 tensor,
497 scatter_dim=head_dim,
498 gather_dim=seq_dim,
499 cp_mesh=cp_mesh,
500 )
503def ulysses_head_to_seq(
504 tensor: Tensor,
505 seq_dim: int,
506 head_dim: int,
507 cp_mesh: DeviceMesh,
508) -> Tensor:
509 """Convert a full-sequence head shard back into a sequence shard.
511 Args:
512 tensor: Local attention output in Ulysses head-sharded layout.
513 seq_dim: Full sequence dimension to split across the CP group.
514 head_dim: Head dimension to gather across the CP group.
515 cp_mesh: One-dimensional CP DeviceMesh.
517 Returns:
518 Tensor restored to the original local sequence-sharded layout.
519 """
520 return _ulysses_all_to_all(
521 tensor,
522 scatter_dim=seq_dim,
523 gather_dim=head_dim,
524 cp_mesh=cp_mesh,
525 )
528def _validate_ulysses_head_counts(
529 query: Tensor, key: Tensor, value: Tensor,
530 ulysses_degree: int, context: str) -> None:
531 """Require every Q/K/V head dimension to be shardable by Ulysses."""
532 head_counts = {
533 "query": query.shape[1],
534 "key": key.shape[1],
535 "value": value.shape[1],
536 }
537 invalid = {
538 name: count
539 for name, count in head_counts.items()
540 if count % ulysses_degree
541 }
542 if invalid:
543 details = ", ".join(
544 f"{name}_heads={count}" for name, count in invalid.items()
545 )
546 raise ValueError(
547 f"{context} requires every Q/K/V head count to be divisible by "
548 f"ulysses_degree ({ulysses_degree}), got {details}"
549 )
552class _AllGatherAlongDim(torch.autograd.Function):
553 """all-gather along cp_dim + backward reduce-scatter semantics (sum across
554 ranks, then take this rank's chunk)."""
556 @staticmethod
557 def forward(
558 ctx: Any,
559 t: Tensor,
560 cp_dim: int,
561 group: Any,
562 cp_size: int,
563 ) -> Tensor:
564 """All-gather the tensor along cp_dim in CP-rank order."""
565 ctx.cp_dim = cp_dim
566 ctx.group = group
567 ctx.cp_size = cp_size
568 world_t = [torch.empty_like(t) for _ in range(cp_size)]
569 dist.all_gather(world_t, t.contiguous(), group=group)
570 # cat in cp_rank order: [chunk_rank0, chunk_rank1, ...]
571 return torch.cat(world_t, dim=cp_dim)
573 @staticmethod
574 def backward(
575 ctx: Any,
576 grad_output: Tensor,
577 ) -> tuple[Tensor, None, None, None]:
578 """Sum the gradient across ranks and take this rank's chunk."""
579 # reduce-scatter: sum the gradient across ranks, take this rank's chunk
580 grad = grad_output.contiguous().clone()
581 dist.all_reduce(grad, group=ctx.group)
582 rank = dist.get_rank(ctx.group)
583 local = torch.chunk(grad, ctx.cp_size, dim=ctx.cp_dim)[rank]
584 return local.contiguous(), None, None, None
587def flex_cp_allgather(
588 k: Tensor,
589 v: Tensor,
590 cp_dim: int,
591 cp_mesh: DeviceMesh,
592) -> tuple[Tensor, Tensor]:
593 """All-gather K/V along CP dimension for context parallel attention.
595 Forward: all-gather K/V along cp_dim (each rank ends up holding the full K/V).
596 Backward: reduce-scatter semantics (gradients summed across ranks, then this
597 rank's chunk is taken -- implemented explicitly by the _AllGatherAlongDim
598 autograd.Function, since plain ``dist.all_gather`` has no autograd kernel).
600 Args:
601 k, v: [B, N, S_local, H] (cp_dim=2 is the sequence dim).
602 cp_dim: the gather dimension.
603 cp_mesh: DeviceMesh of the CP dim. The communication group is taken from
604 ``cp_mesh.get_group()`` -- already created and cached at DeviceMesh
605 construction; **dist.new_group must NOT be called here** (otherwise
606 every forward leaks one process group, and the semantics would be
607 misaligned).
608 """
609 cp_size = cp_mesh.size()
610 if cp_size <= 1:
611 return k, v
612 group = cp_mesh.get_group()
613 return (_AllGatherAlongDim.apply(k, cp_dim, group, cp_size),
614 _AllGatherAlongDim.apply(v, cp_dim, group, cp_size))
617def _build_hybrid_cp_submeshes(cp_mesh, ulysses_degree: int):
618 """Build cached communication submeshes for local-tensor Hybrid CP."""
619 if cp_mesh is None or cp_mesh.size() <= 1:
620 raise ValueError("Hybrid CP requires an active CP mesh")
621 if isinstance(ulysses_degree, bool) or not isinstance(ulysses_degree, int):
622 raise TypeError(
623 "Hybrid CP requires integer ulysses_degree, got "
624 f"{type(ulysses_degree).__name__}"
625 )
626 cp_size = cp_mesh.size()
627 if not 1 < ulysses_degree < cp_size:
628 raise ValueError(
629 "Hybrid CP requires 1 < ulysses_degree < cp_size, got "
630 f"ulysses_degree={ulysses_degree}, cp_size={cp_size}"
631 )
632 if cp_size % ulysses_degree:
633 raise ValueError(
634 f"cp_size ({cp_size}) must be divisible by ulysses_degree "
635 f"({ulysses_degree})"
636 )
638 cache_key = (cp_mesh.to_hash(), ulysses_degree)
639 cached = _HYBRID_MESH_CACHE.get(cache_key)
640 if cached is not None:
641 return cached
643 colossal_degree = cp_size // ulysses_degree
644 if cp_mesh.ndim == 2:
645 hybrid_mesh = cp_mesh
646 elif (cp_mesh.root_mesh is not None
647 and cp_mesh.mesh_dim_names == ("cp",)):
648 hybrid_mesh = cp_mesh._unflatten( # pylint: disable=protected-access
649 "cp",
650 (colossal_degree, ulysses_degree),
651 ("colossal", "ulysses"),
652 )
653 else:
654 ranks = list(cp_mesh.rank_list)
655 hybrid_mesh = DeviceMesh(
656 cp_mesh.device_type,
657 [
658 ranks[index * ulysses_degree:(index + 1) * ulysses_degree]
659 for index in range(colossal_degree)
660 ],
661 mesh_dim_names=("colossal", "ulysses"),
662 )
663 dim_names = hybrid_mesh.mesh_dim_names
664 if not dim_names or len(dim_names) != 2:
665 raise ValueError("Hybrid CP requires a named two-dimensional submesh")
666 colossal_mesh = hybrid_mesh[dim_names[0]]
667 ulysses_mesh = hybrid_mesh[dim_names[1]]
668 result = (ulysses_mesh, colossal_mesh)
669 _HYBRID_MESH_CACHE[cache_key] = result
670 return result
673def hybrid_cp_attention(
674 attention_fn: Callable[[Tensor, Tensor, Tensor, dict[str, Any]], Any],
675 query: Tensor, key: Tensor, value: Tensor,
676 attention_kwargs: dict[str, Any], cp_mesh: Any,
677 ulysses_degree: int) -> Any:
678 """Run local-tensor Hybrid communication around an attention callable.
680 ``attention_kwargs`` are forwarded unchanged. Mask construction and other
681 model-input semantics belong to the caller's input-preparation contract.
682 """
683 _validate_ulysses_head_counts(
684 query, key, value, ulysses_degree, "Hybrid CP"
685 )
686 ulysses_mesh, colossal_mesh = (
687 _build_hybrid_cp_submeshes(cp_mesh, ulysses_degree)
688 )
689 local_shape = tuple(query.shape)
690 query, key, value = (
691 ulysses_seq_to_head(tensor, 2, 1, ulysses_mesh)
692 for tensor in (query, key, value)
693 )
694 key, value = flex_cp_allgather(key, value, 2, colossal_mesh)
695 output = attention_fn(query, key, value, attention_kwargs)
696 output = ulysses_head_to_seq(output, 2, 1, ulysses_mesh)
697 if tuple(output.shape) != local_shape:
698 raise RuntimeError(
699 f"Hybrid CP output shape {tuple(output.shape)} does not restore "
700 f"the local query shape {local_shape}"
701 )
702 return output