Coverage for  / home / jenkins / .local / lib / python3.10 / site-packages / hyper_parallel / distributed_data / topology.py: 89%

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"""Derive Dataset Reader and Data Constructor ownership from a named mesh.""" 

16# This distributed-data package is intentionally Torch-only. 

17 

18from __future__ import annotations 

19 

20import math 

21from dataclasses import dataclass 

22from typing import Any 

23 

24import torch.distributed as dist # pylint: disable=forbidden-backend-import 

25 

26_DEFAULT_DP_DIM_NAMES = ("dp_replicate", "dp_shard", "dp") 

27 

28 

29def _flatten_coordinate(coordinate: tuple[int, ...], shape: tuple[int, ...]) -> int: 

30 flat_index = 0 

31 for value, size in zip(coordinate, shape, strict=True): 

32 flat_index = flat_index * size + value 

33 return flat_index 

34 

35 

36def _unflatten_coordinate(flat_index: int, shape: tuple[int, ...]) -> tuple[int, ...]: 

37 coordinate = [0] * len(shape) 

38 remaining = flat_index 

39 for dim_index in range(len(shape) - 1, -1, -1): 

40 coordinate[dim_index] = remaining % shape[dim_index] 

41 remaining //= shape[dim_index] 

42 return tuple(coordinate) 

43 

44 

45def _mesh_layout(mesh: Any) -> tuple[tuple[int, ...], tuple[int, ...]]: 

46 mesh_shape = getattr(mesh, "mesh_shape", None) 

47 rank_list = getattr(mesh, "rank_list", None) 

48 if mesh_shape is not None and rank_list is not None: 

49 return tuple(int(size) for size in mesh_shape), tuple(int(rank) for rank in rank_list) 

50 

51 mesh_tensor = getattr(mesh, "mesh", None) 

52 if mesh_tensor is None: 

53 raise ValueError("mesh must expose mesh_shape/rank_list or a Torch mesh tensor.") 

54 return ( 

55 tuple(int(size) for size in mesh_tensor.shape), 

56 tuple(int(rank) for rank in mesh_tensor.reshape(-1).tolist()), 

57 ) 

58 

59 

60@dataclass(frozen=True) 

61class DataTopology: 

62 """DP constructors and their model-parallel consumer groups.""" 

63 

64 mesh_shape: tuple[int, ...] 

65 mesh_dim_names: tuple[str, ...] 

66 rank_list: tuple[int, ...] 

67 global_rank: int 

68 dp_dim_names: tuple[str, ...] 

69 coordinate: tuple[int, ...] 

70 data_rank: int 

71 data_parallel_size: int 

72 constructor_rank: int 

73 constructor_ranks: tuple[int, ...] 

74 model_parallel_rank_groups: tuple[tuple[int, ...], ...] 

75 model_parallel_ranks: tuple[int, ...] 

76 

77 @classmethod 

78 def from_mesh( 

79 cls, 

80 mesh: Any, 

81 *, 

82 global_rank: int | None = None, 

83 dp_dim_names: tuple[str, ...] | None = None, 

84 ) -> "DataTopology": 

85 """Build topology from a named HyperParallel or Torch DeviceMesh. 

86 

87 Args: 

88 mesh: Root training mesh with named dimensions. 

89 global_rank: Current global rank. Defaults to torch.distributed or 

90 rank zero in a standalone process. 

91 dp_dim_names: Mesh dimensions defining independent data consumers. 

92 

93 Returns: 

94 Derived data topology. 

95 """ 

96 mesh_dim_names = getattr(mesh, "mesh_dim_names", None) 

97 if not mesh_dim_names: 

98 raise ValueError("Distributed data loading requires a mesh with named dimensions.") 

99 rank = global_rank 

100 if rank is None: 

101 rank = dist.get_rank() if dist.is_available() and dist.is_initialized() else 0 

102 mesh_shape, rank_list = _mesh_layout(mesh) 

103 return cls.from_layout( 

104 mesh_shape, 

105 tuple(mesh_dim_names), 

106 rank_list, 

107 rank, 

108 dp_dim_names=dp_dim_names, 

109 ) 

110 

111 @classmethod 

112 def from_layout( 

113 cls, 

114 mesh_shape: tuple[int, ...], 

115 mesh_dim_names: tuple[str, ...], 

116 rank_list: tuple[int, ...], 

117 global_rank: int, 

118 *, 

119 dp_dim_names: tuple[str, ...] | None = None, 

120 ) -> "DataTopology": 

121 """Build topology from a raw named layout. 

122 

123 Args: 

124 mesh_shape: Size of every named mesh dimension. 

125 mesh_dim_names: Unique mesh dimension names. 

126 rank_list: Global ranks in row-major mesh order. 

127 global_rank: Rank for which ownership is derived. 

128 dp_dim_names: Dimensions that form the DP coordinate. 

129 

130 Returns: 

131 Derived data topology. 

132 """ 

133 cls._validate_layout(mesh_shape, mesh_dim_names, rank_list, global_rank) 

134 selected_dp_names = dp_dim_names 

135 if selected_dp_names is None: 

136 selected_dp_names = tuple(name for name in mesh_dim_names if name in _DEFAULT_DP_DIM_NAMES) 

137 if len(set(selected_dp_names)) != len(selected_dp_names): 

138 raise ValueError(f"dp_dim_names must be unique, but got {selected_dp_names}.") 

139 if any(name not in mesh_dim_names for name in selected_dp_names): 

140 raise ValueError(f"dp_dim_names {selected_dp_names} must be present in {mesh_dim_names}.") 

141 

142 coordinate = _unflatten_coordinate(rank_list.index(global_rank), mesh_shape) 

143 dp_indices = tuple(mesh_dim_names.index(name) for name in selected_dp_names) 

144 dp_shape = tuple(mesh_shape[index] for index in dp_indices) 

145 dp_coordinate = tuple(coordinate[index] for index in dp_indices) 

146 data_parallel_size = math.prod(dp_shape) if dp_shape else 1 

147 data_rank = _flatten_coordinate(dp_coordinate, dp_shape) if dp_shape else 0 

148 constructor_ranks = cls._build_constructor_ranks(mesh_shape, rank_list, dp_indices, dp_shape) 

149 model_groups = cls._build_model_groups(mesh_shape, rank_list, dp_indices, dp_shape) 

150 return cls( 

151 mesh_shape=mesh_shape, 

152 mesh_dim_names=mesh_dim_names, 

153 rank_list=rank_list, 

154 global_rank=global_rank, 

155 dp_dim_names=selected_dp_names, 

156 coordinate=coordinate, 

157 data_rank=data_rank, 

158 data_parallel_size=data_parallel_size, 

159 constructor_rank=constructor_ranks[data_rank], 

160 constructor_ranks=constructor_ranks, 

161 model_parallel_rank_groups=model_groups, 

162 model_parallel_ranks=model_groups[data_rank], 

163 ) 

164 

165 @staticmethod 

166 def _validate_layout( 

167 mesh_shape: tuple[int, ...], 

168 mesh_dim_names: tuple[str, ...], 

169 rank_list: tuple[int, ...], 

170 global_rank: int, 

171 ) -> None: 

172 """Require a named mesh filled by unique ranks and containing this process.""" 

173 invalid_sizes = any( 

174 not isinstance(size, int) or isinstance(size, bool) or size < 1 

175 for size in mesh_shape 

176 ) 

177 if not mesh_shape or invalid_sizes: 

178 raise ValueError(f"mesh_shape must contain positive integers, but got {mesh_shape}.") 

179 if len(mesh_shape) != len(mesh_dim_names): 

180 raise ValueError("mesh_shape and mesh_dim_names must have the same length.") 

181 if len(set(mesh_dim_names)) != len(mesh_dim_names): 

182 raise ValueError(f"mesh_dim_names must be unique, but got {mesh_dim_names}.") 

183 if math.prod(mesh_shape) != len(rank_list) or len(set(rank_list)) != len(rank_list): 

184 raise ValueError("rank_list must uniquely fill every mesh position.") 

185 if global_rank not in rank_list: 

186 raise ValueError(f"global_rank {global_rank} is not present in rank_list {rank_list}.") 

187 

188 @staticmethod 

189 def _build_constructor_ranks( 

190 mesh_shape: tuple[int, ...], 

191 rank_list: tuple[int, ...], 

192 dp_indices: tuple[int, ...], 

193 dp_shape: tuple[int, ...], 

194 ) -> tuple[int, ...]: 

195 """Choose the zero model-parallel coordinate for each DP coordinate.""" 

196 data_parallel_size = math.prod(dp_shape) if dp_shape else 1 

197 constructors = [] 

198 for data_rank in range(data_parallel_size): 

199 dp_coordinate = _unflatten_coordinate(data_rank, dp_shape) if dp_shape else () 

200 coordinate = [0] * len(mesh_shape) 

201 for dim_index, value in zip(dp_indices, dp_coordinate, strict=True): 

202 coordinate[dim_index] = value 

203 constructors.append(rank_list[_flatten_coordinate(tuple(coordinate), mesh_shape)]) 

204 return tuple(constructors) 

205 

206 @staticmethod 

207 def _build_model_groups( 

208 mesh_shape: tuple[int, ...], 

209 rank_list: tuple[int, ...], 

210 dp_indices: tuple[int, ...], 

211 dp_shape: tuple[int, ...], 

212 ) -> tuple[tuple[int, ...], ...]: 

213 """Group ranks that share one DP coordinate across model-parallel dimensions.""" 

214 data_parallel_size = math.prod(dp_shape) if dp_shape else 1 

215 groups: list[list[int]] = [[] for _ in range(data_parallel_size)] 

216 for flat_index, candidate_rank in enumerate(rank_list): 

217 coordinate = _unflatten_coordinate(flat_index, mesh_shape) 

218 dp_coordinate = tuple(coordinate[index] for index in dp_indices) 

219 data_rank = _flatten_coordinate(dp_coordinate, dp_shape) if dp_shape else 0 

220 groups[data_rank].append(candidate_rank) 

221 return tuple(tuple(group) for group in groups) 

222 

223 @property 

224 def is_constructor(self) -> bool: 

225 """Return whether this rank constructs its DP coordinate's local batch.""" 

226 return self.global_rank == self.constructor_rank