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

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 routing of aligned image features to decoder token owners.""" 

16 

17from __future__ import annotations 

18 

19from enum import Enum 

20from typing import Any, Mapping, Sequence 

21 

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

23 

24from hyper_parallel.core.dtensor._utils import differentiable_all_to_all_single, differentiable_variable_all_gather 

25from hyper_parallel.distributed.encoder_parallel.contracts import ( 

26 EncoderFeature, 

27 EncoderWorkItem, 

28) 

29 

30 

31class FeatureExchangeMode(Enum): 

32 """Collective selected from static decoder topology.""" 

33 

34 LOCAL = "local" 

35 TP_VARIABLE_ALL_GATHER = "tp_variable_all_gather" 

36 TP_VARIABLE_A2A = "tp_variable_a2a" 

37 PARENT_VARIABLE_A2A = "parent_variable_a2a" 

38 

39 

40def select_feature_exchange(mesh_context: Any) -> FeatureExchangeMode: 

41 """Choose a uniform exchange algorithm from decoder CP/TP/SP topology. 

42 

43 Args: 

44 mesh_context: Decoder CP, TP, and sequence-parallel coordinates. 

45 """ 

46 if mesh_context.cp_size > 1: 

47 return FeatureExchangeMode.PARENT_VARIABLE_A2A 

48 if mesh_context.tp_size == 1: 

49 return FeatureExchangeMode.LOCAL 

50 if mesh_context.sequence_parallel: 

51 return FeatureExchangeMode.TP_VARIABLE_A2A 

52 return FeatureExchangeMode.TP_VARIABLE_ALL_GATHER 

53 

54 

55def cat_or_empty(rows: Sequence[torch.Tensor], *, anchor: torch.Tensor, 

56 feature_width: int) -> torch.Tensor: 

57 """Concatenate rows, preserving an autograd edge for an empty sender. 

58 

59 Args: 

60 rows: Feature tensors to concatenate. 

61 anchor: Tensor supplying the zero-row gradient edge. 

62 feature_width: Width of an empty feature tensor. 

63 """ 

64 if rows: 

65 return torch.cat(tuple(rows), dim=0) 

66 return anchor.sum().reshape(1, 1).expand(0, feature_width) 

67 

68 

69def _ordered_source_records(buckets: Sequence[Sequence[EncoderWorkItem]], encoder_cp_size: int 

70 ) -> dict[int, tuple[EncoderWorkItem, ...]]: 

71 return {bucket_rank * encoder_cp_size: tuple(bucket) for bucket_rank, bucket in enumerate(buckets)} 

72 

73 

74def _route_tp_all_gather(*, local_by_id: Mapping[int, torch.Tensor], 

75 source_records: Mapping[int, tuple[EncoderWorkItem, ...]], 

76 prototype: torch.Tensor, parent_rank: int, 

77 exchange_mesh: Any) -> tuple[torch.Tensor, tuple[EncoderWorkItem, ...]]: 

78 """Gather each producer's complete records to every decoder TP rank.""" 

79 group_size = exchange_mesh.size() 

80 local_items = source_records.get(parent_rank, ()) 

81 local_rows = cat_or_empty([local_by_id[item.item_id] for item in local_items], 

82 anchor=prototype, feature_width=int(prototype.shape[-1])) 

83 splits = [sum(item.output_rows for item in source_records.get(rank, ())) for rank in range(group_size)] 

84 rows = differentiable_variable_all_gather(local_rows, output_splits=splits, 

85 group=exchange_mesh.get_group()) 

86 received_items = tuple(item for rank in range(group_size) for item in source_records.get(rank, ())) 

87 return rows, received_items 

88 

89 

90def _route_variable_a2a(*, mode: FeatureExchangeMode, 

91 local_by_id: Mapping[int, torch.Tensor], 

92 source_records: Mapping[int, tuple[EncoderWorkItem, ...]], 

93 prototype: torch.Tensor, parent_rank: int, 

94 destinations: Mapping[int, tuple[int, ...]], 

95 encoder_cp_size: int, exchange_mesh: Any 

96 ) -> tuple[torch.Tensor, tuple[EncoderWorkItem, ...]]: 

97 """Send each full record only to its exact decoder token owners.""" 

98 group_size = exchange_mesh.size() 

99 local_rank = exchange_mesh.get_local_rank() 

100 if group_size != len(source_records) * encoder_cp_size: 

101 raise ValueError(f"{mode.value} exchange size disagrees with encoder parent") 

102 width = int(prototype.shape[-1]) 

103 send_chunks = [] 

104 input_splits = [] 

105 for destination_rank in range(group_size): 

106 items = [item for item in source_records.get(parent_rank, ()) 

107 if destination_rank in destinations[item.item_id]] 

108 chunk = cat_or_empty([local_by_id[item.item_id] for item in items], 

109 anchor=prototype, feature_width=width) 

110 send_chunks.append(chunk) 

111 input_splits.append(chunk.shape[0]) 

112 packed = cat_or_empty(send_chunks, anchor=prototype, feature_width=width) 

113 output_splits = [] 

114 received_items = [] 

115 for source_rank in range(group_size): 

116 items = [item for item in source_records.get(source_rank, ()) 

117 if local_rank in destinations[item.item_id]] 

118 output_splits.append(sum(item.output_rows for item in items)) 

119 received_items.extend(items) 

120 rows = differentiable_all_to_all_single(packed, input_splits=input_splits, 

121 output_splits=output_splits, group=exchange_mesh.get_group()) 

122 return rows, tuple(received_items) 

123 

124 

125def route_features(*, mode: FeatureExchangeMode, producer_records: Sequence[EncoderFeature], 

126 prototype: torch.Tensor, parent_rank: int, destinations: Mapping[int, tuple[int, ...]], 

127 buckets: Sequence[Sequence[EncoderWorkItem]], encoder_cp_size: int, 

128 exchange_mesh: Any) -> tuple[torch.Tensor, tuple[EncoderWorkItem, ...]]: 

129 """Exchange full feature records using a deterministic plan shared by all ranks.""" 

130 width = int(prototype.shape[-1]) 

131 local_by_id = {item.item_id: feature for item, feature in producer_records} 

132 source_records = _ordered_source_records(buckets, encoder_cp_size) 

133 if mode is FeatureExchangeMode.LOCAL: 

134 received_items = tuple(item for item, _ in producer_records) 

135 rows = cat_or_empty([feature for _, feature in producer_records], anchor=prototype, feature_width=width) 

136 return rows, received_items 

137 

138 if mode is FeatureExchangeMode.TP_VARIABLE_ALL_GATHER: 

139 return _route_tp_all_gather(local_by_id=local_by_id, source_records=source_records, 

140 prototype=prototype, parent_rank=parent_rank, 

141 exchange_mesh=exchange_mesh) 

142 if mode in (FeatureExchangeMode.TP_VARIABLE_A2A, FeatureExchangeMode.PARENT_VARIABLE_A2A): 

143 return _route_variable_a2a( 

144 mode=mode, local_by_id=local_by_id, source_records=source_records, 

145 prototype=prototype, parent_rank=parent_rank, destinations=destinations, 

146 encoder_cp_size=encoder_cp_size, exchange_mesh=exchange_mesh, 

147 ) 

148 raise ValueError(f"unsupported feature exchange mode: {mode!r}") 

149 

150 

151def unpack_records(rows: torch.Tensor, received_items: Sequence[EncoderWorkItem] 

152 ) -> tuple[EncoderFeature, ...]: 

153 """Split transported rows using the shared plan's image order. 

154 

155 Args: 

156 rows: Received aligned feature rows. 

157 received_items: Images in source-rank then bucket order. 

158 """ 

159 records = [] 

160 offset = 0 

161 for item in received_items: 

162 stop = offset + item.output_rows 

163 records.append((item, rows[offset:stop])) 

164 offset = stop 

165 if offset != rows.shape[0]: 

166 raise ValueError( 

167 "transported rows disagree with planned images; " 

168 f"rows={rows.shape[0]}; described={offset}" 

169 ) 

170 return tuple(sorted(records, key=lambda record: ( 

171 record[0].batch_index, record[0].decoder_token_start, record[0].item_id 

172 )))