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

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"""Pure topology, token ownership, and image work planning."""

17from __future__ import annotations

19from dataclasses import dataclass

20from math import lcm

21from typing import Any, Mapping, Sequence

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

25from hyper_parallel.distributed.encoder_parallel.contracts import EncoderSegment, EncoderWorkItem

28def required_omni_global_alignment(mesh_context: Any) -> int:

29 """Return the shared compact-attention and decoder-SP alignment.

31 Args:

32 mesh_context: Decoder CP, TP, and sequence-parallel topology.

33 """

34 cp_size = int(mesh_context.cp_size)

35 tp_parts = int(mesh_context.tp_size) if mesh_context.sequence_parallel else 1

36 return lcm(2 * cp_size, cp_size * tp_parts)

39@dataclass(frozen=True)

40class DecoderSequenceLayout:

41 """The exact decoder hidden-token interval on one CP/TP rank."""

43 sequence_parallel: bool

44 cp_rank: int

45 cp_size: int

46 tp_rank: int

47 tp_size: int

48 cp_local_length: int

49 global_sequence_length: int

50 cp_local_start: int

51 cp_local_stop: int

52 global_start: int

53 global_stop: int

55 @classmethod

56 def build(cls, *, mesh_context: Any, cp_local_length: int,

57 cp_rank: int | None = None, tp_rank: int | None = None) -> "DecoderSequenceLayout":

58 """Derive the token interval after embedding's TP reduce-scatter."""

59 cp_size = int(mesh_context.cp_size)

60 tp_size = int(mesh_context.tp_size)

61 cp_rank = int(mesh_context.cp_rank if cp_rank is None else cp_rank)

62 tp_rank = int(mesh_context.tp_rank if tp_rank is None else tp_rank)

63 sp = bool(mesh_context.sequence_parallel)

64 if not 0 <= cp_rank < cp_size or not 0 <= tp_rank < tp_size:

65 raise ValueError(

66 "decoder CP/TP coordinate is outside topology; "

67 f"cp_rank={cp_rank}; cp_size={cp_size}; tp_rank={tp_rank}; tp_size={tp_size}"

68 )

69 if cp_local_length < 1 or (sp and cp_local_length % tp_size):

70 raise ValueError(

71 "invalid decoder sequence layout; "

72 f"sequence_parallel={sp}; cp_local_length={cp_local_length}; tp_size={tp_size}; "

73 "constraint=cp_local_length % tp_size == 0 when SP is enabled"

74 )

75 local_length = cp_local_length // (tp_size if sp else 1)

76 local_start = (tp_rank if sp else 0) * local_length

77 global_start = cp_rank * cp_local_length + local_start

78 return cls(sp, cp_rank, cp_size, tp_rank, tp_size, cp_local_length,

79 cp_size * cp_local_length, local_start, local_start + local_length,

80 global_start, global_start + local_length)

82 def localize(self, value: torch.Tensor, *, sequence_dim: int = 1) -> torch.Tensor:

83 """Slice a CP-local auxiliary tensor to this hidden-token interval.

85 Args:

86 value: Auxiliary tensor replicated across decoder TP ranks.

87 sequence_dim: Dimension containing CP-local token rows.

88 """

89 if value.shape[sequence_dim] != self.cp_local_length:

90 raise ValueError(

91 "CP-local token tensor length mismatch; "

92 f"actual={value.shape[sequence_dim]}; expected={self.cp_local_length}; "

93 f"cp_rank={self.cp_rank}; tp_rank={self.tp_rank}"

94 )

95 return value.narrow(sequence_dim, self.cp_local_start,

96 self.cp_local_stop - self.cp_local_start).contiguous()

98 def validate_hidden(self, hidden: torch.Tensor, *, sequence_dim: int = 1) -> None:

99 """Fail before merging against an unexpected decoder hidden layout.

100

101 Args:

102 hidden: Decoder hidden tensor after embedding redistribution.

103 sequence_dim: Dimension containing decoder token rows.

104 """

105 expected = self.global_stop - self.global_start

106 if hidden.shape[sequence_dim] != expected:

107 raise ValueError(

108 "decoder hidden does not match planned CP/SP shard; "

109 f"actual={hidden.shape[sequence_dim]}; expected={expected}; "

110 f"global_interval=({self.global_start},{self.global_stop}); "

111 f"cp_rank={self.cp_rank}; tp_rank={self.tp_rank}"

112 )

113

114

115def assign_lpt(items: Sequence[EncoderWorkItem], encoder_dp_size: int,

116 encoder_cp_size: int) -> tuple[tuple[EncoderWorkItem, ...], ...]:

117 """Assign image work by stable largest-processing-time-first scheduling.

118

119 Args:

120 items: Images with valid patch counts and stable IDs.

121 encoder_dp_size: Number of visual work buckets.

122 encoder_cp_size: Sequence-sharding degree of each bucket.

123 """

124 if encoder_dp_size < 1 or encoder_cp_size < 1:

125 raise ValueError("encoder DP and CP sizes must be positive")

126 buckets: list[list[EncoderWorkItem]] = [[] for _ in range(encoder_dp_size)]

127 loads = [0] * encoder_dp_size

128 for item in sorted(items, key=lambda value: (-value.sequence_length ** 2, value.batch_index, value.item_id)):

129 bucket = min(range(encoder_dp_size), key=lambda rank: (loads[rank], rank))

130 buckets[bucket].append(item)

131 loads[bucket] += item.sequence_length ** 2 // encoder_cp_size

132 return tuple(tuple(bucket) for bucket in buckets)

133

134

135def build_segments(items: Sequence[EncoderWorkItem], encoder_cp_rank: int,

136 encoder_cp_size: int) -> tuple[EncoderSegment, ...]:

137 """Build equal-size padded sequence shards for one Encoder-CP rank.

138

139 Args:

140 items: Images assigned to this Encoder-DP bucket.

141 encoder_cp_rank: Sequence-shard coordinate within the bucket.

142 encoder_cp_size: Number of shards per image.

143 """

144 if not 0 <= encoder_cp_rank < encoder_cp_size:

145 raise ValueError("invalid encoder CP rank")

146 segments = []

147 offset = 0

148 for item in items:

149 if item.sequence_length < 1:

150 raise ValueError("image patch count must be positive")

151 local_length = (item.sequence_length + encoder_cp_size - 1) // encoder_cp_size

152 start = encoder_cp_rank * local_length

153 valid = max(0, min(local_length, item.sequence_length - start))

154 segments.append(EncoderSegment(item, start, valid, local_length, offset))

155 offset += local_length

156 if not segments:

157 segments.append(EncoderSegment(None, 0, 1, 1, 0))

158 return tuple(segments)

159

160

161def build_destinations(items: Sequence[EncoderWorkItem], layouts: Mapping[int, DecoderSequenceLayout]

162 ) -> dict[int, tuple[int, ...]]:

163 """Find every CP/TP rank whose decoder hidden intersects each image span.

164

165 Args:

166 items: Images carrying their global decoder token spans.

167 layouts: Decoder token interval for every rank in the parent group.

168 """

169 lengths = {layout.global_sequence_length for layout in layouts.values()}

170 if len(lengths) != 1:

171 raise ValueError("inconsistent decoder global lengths")

172 global_length = next(iter(lengths))

173 result = {}

174 for item in items:

175 start = item.decoder_token_start

176 stop = start + item.decoder_token_length

177 if item.decoder_token_length <= 0 or start < 0 or stop > global_length:

178 raise ValueError(

179 "image span is outside the decoder domain; "

180 f"item_id={item.item_id}; span=({start},{stop}); global_sequence_length={global_length}"

181 )

182 owners = tuple(rank for rank, layout in layouts.items()

183 if max(start, layout.global_start) < min(stop, layout.global_stop))

184 if not owners:

185 raise ValueError(f"image span has no decoder token owner; item_id={item.item_id}")

186 result[item.item_id] = owners

187 return result

Coverage for / home / jenkins / .local / lib / python3.10 / site-packages / hyper_parallel / distributed / encoder_parallel / planning.py: 92%

88 statements