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

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"""Reusable sequence transport over an existing HyperParallel CP mesh.""" 

16 

17from __future__ import annotations 

18 

19from dataclasses import dataclass 

20from typing import Any 

21 

22from torch import Tensor 

23 

24from hyper_parallel.distributed.context_parallel.collectives import ( 

25 AsyncCPCollective, async_cp_allgather_launch, ulysses_head_to_seq, ulysses_seq_to_head, 

26) 

27 

28 

29@dataclass(frozen=True) 

30class SequencePartition: 

31 """Contiguous, equal-length sequence shards in process-group rank order.""" 

32 

33 lengths: tuple[int, ...] 

34 rank: int 

35 

36 @property 

37 def offset(self) -> int: 

38 """Global sequence coordinate of the first local token.""" 

39 return sum(self.lengths[:self.rank]) 

40 

41 @property 

42 def global_length(self) -> int: 

43 """Number of tokens represented by this partition.""" 

44 return sum(self.lengths) 

45 

46 

47@dataclass 

48class PendingSequence: 

49 """Per-call pending output; consumers obtain tensors only through wait().""" 

50 

51 collective: AsyncCPCollective 

52 partition: SequencePartition 

53 

54 def wait(self) -> tuple[Tensor, SequencePartition]: 

55 """Wait before exposing the gathered tensor and its coordinate receipt.""" 

56 return self.collective.wait(), self.partition 

57 

58 

59class CPContext: 

60 """Transport only: model families retain projection, mask and attention math.""" 

61 

62 def __init__(self, mesh: Any) -> None: 

63 """Borrow a one-dimensional CP mesh without creating process groups.""" 

64 if mesh is None or mesh.ndim != 1: 

65 raise ValueError("CPContext requires an existing one-dimensional HP mesh") 

66 self.mesh = mesh 

67 self.size = mesh.size() 

68 self.rank = mesh.get_local_rank() 

69 

70 def partition(self, local_length: int) -> SequencePartition: 

71 """Describe equal contiguous source shards without reading activations. 

72 

73 Args: 

74 local_length: Sequence length on each rank of the CP group. 

75 """ 

76 if local_length < 0: 

77 raise ValueError("local sequence length cannot be negative") 

78 return SequencePartition((local_length,) * self.size, self.rank) 

79 

80 def launch_gather_sequence(self, tensor: Tensor, *, dim: int = 1) -> PendingSequence: 

81 """Launch differentiable sequence gathering with explicit readiness. 

82 

83 Args: 

84 tensor: Local activation shard, with equal lengths on all CP ranks. 

85 dim: Sequence dimension; rank order defines global token coordinates. 

86 """ 

87 partition = self.partition(tensor.shape[dim]) 

88 return PendingSequence(async_cp_allgather_launch(tensor, dim, self.mesh), partition) 

89 

90 def gather_sequence(self, tensor: Tensor, *, dim: int = 1) -> tuple[Tensor, SequencePartition]: 

91 """Gather and wait synchronously; backward sums consumers into source shards. 

92 

93 Args: 

94 tensor: Local sequence shard. 

95 dim: Sequence dimension. 

96 """ 

97 return self.launch_gather_sequence(tensor, dim=dim).wait() 

98 

99 def sequence_to_head(self, tensor: Tensor, *, seq_dim: int = 1, head_dim: int = 2) -> Tensor: 

100 """Exchange sequence-sharded activations into head-sharded activations. 

101 

102 Args: 

103 tensor: Local sequence shard with divisible head count. 

104 seq_dim: Sequence dimension. 

105 head_dim: Head dimension. 

106 """ 

107 return ulysses_seq_to_head(tensor, seq_dim=seq_dim, head_dim=head_dim, cp_mesh=self.mesh) 

108 

109 def head_to_sequence(self, tensor: Tensor, *, seq_dim: int = 1, head_dim: int = 2) -> Tensor: 

110 """Apply the inverse head-to-sequence exchange. 

111 

112 Args: 

113 tensor: Local head shard with divisible sequence length. 

114 seq_dim: Sequence dimension. 

115 head_dim: Head dimension. 

116 """ 

117 return ulysses_head_to_seq(tensor, seq_dim=seq_dim, head_dim=head_dim, cp_mesh=self.mesh)