Coverage for  / home / jenkins / .local / lib / python3.10 / site-packages / hyper_parallel / core / dtensor / _ragged_utils.py: 75%

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"""Flat-storage geometry and distribution helpers for RaggedShard.""" 

16from math import prod 

17from typing import NamedTuple, Optional, Sequence 

18 

19import numpy as np 

20import torch 

21 

22from hyper_parallel.core.dtensor._collective_utils import mesh_scatter_ragged 

23from hyper_parallel.core.dtensor.layout import Layout, RaggedShardInfo 

24 

25Tensor = torch.Tensor 

26 

27 

28def _layout_has_ragged_shard(layout: object) -> bool: 

29 """Return whether a concrete Layout carries RaggedShard metadata.""" 

30 return isinstance(getattr(layout, "ragged_shard", None), RaggedShardInfo) 

31 

32 

33class _RaggedSlice(NamedTuple): 

34 """Flat interval owned by one rank in a RaggedShard layout.""" 

35 

36 flat_start: int 

37 flat_end: int 

38 

39 @property 

40 def local_numel(self) -> int: 

41 """Return the number of flat elements in the interval.""" 

42 return self.flat_end - self.flat_start 

43 

44 

45def _normalize_global_shape(shape: Sequence[int]) -> tuple[int, ...]: 

46 """Normalize a concrete global tensor shape.""" 

47 normalized = tuple(shape) 

48 if any( 

49 not isinstance(size, (int, np.integer)) or isinstance(size, bool) or size < 0 

50 for size in normalized 

51 ): 

52 raise ValueError( 

53 f"DTensor global shape must contain non-negative integers, got {normalized!r}" 

54 ) 

55 return tuple(int(size) for size in normalized) 

56 

57 

58def _compute_ragged_slice( 

59 global_shape: Sequence[int], 

60 layout: Layout, 

61 local_rank: Optional[int] = None, 

62) -> _RaggedSlice: 

63 """Compute one rank's flat RaggedShard interval. 

64 

65 Phase one supports exactly one RaggedShard and Replicate placements on all 

66 other mesh dimensions. 

67 """ 

68 info = layout.ragged_shard 

69 if info is None: 

70 raise ValueError("RaggedShard slice computation requires a ragged layout") 

71 

72 for mesh_dim, placement in enumerate(layout.placements): 

73 if mesh_dim == info.mesh_dim: 

74 continue 

75 if not placement.is_replicate(): 

76 raise NotImplementedError( 

77 "RaggedShard phase one only supports Replicate on other mesh dimensions, " 

78 f"got mesh_dim={mesh_dim}, placement={placement!r}" 

79 ) 

80 

81 shape = _normalize_global_shape(global_shape) 

82 ragged = info.placement 

83 prefix_ndim = len(ragged.dims) 

84 if prefix_ndim > len(shape): 

85 raise ValueError( 

86 f"RaggedShard dims {ragged.dims!r} exceed global shape rank {len(shape)}" 

87 ) 

88 

89 mesh_dim_size = layout.mesh.size(info.mesh_dim) 

90 if len(ragged.local_units) != mesh_dim_size: 

91 raise ValueError( 

92 "RaggedShard len(local_units) must equal mesh.size(mesh_dim), " 

93 f"got len(local_units)={len(ragged.local_units)}, mesh_dim_size={mesh_dim_size}" 

94 ) 

95 

96 prefix_cells = prod(shape[:prefix_ndim]) 

97 total_units = sum(ragged.local_units) 

98 if prefix_cells % total_units != 0: 

99 raise ValueError( 

100 "RaggedShard prefix cell count must be divisible by sum(local_units), " 

101 f"got prefix_cells={prefix_cells}, local_units={ragged.local_units!r}" 

102 ) 

103 

104 if local_rank is None: 

105 local_rank = layout.mesh.get_local_rank(info.mesh_dim) 

106 if local_rank < 0 or local_rank >= mesh_dim_size: 

107 raise ValueError( 

108 f"RaggedShard local rank must be in [0, {mesh_dim_size}), got {local_rank}" 

109 ) 

110 cells_per_unit = prefix_cells // total_units 

111 prefix_start = sum(ragged.local_units[:local_rank]) * cells_per_unit 

112 local_prefix_cells = ragged.local_units[local_rank] * cells_per_unit 

113 suffix_numel = prod(shape[prefix_ndim:]) 

114 flat_start = prefix_start * suffix_numel 

115 flat_end = flat_start + local_prefix_cells * suffix_numel 

116 return _RaggedSlice(flat_start, flat_end) 

117 

118 

119def _compute_ragged_splits( 

120 global_shape: Sequence[int], 

121 layout: Layout, 

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

123 """Return flat element counts contributed by all ranks on the ragged axis.""" 

124 info = layout.ragged_shard 

125 if info is None: 

126 raise ValueError("RaggedShard split computation requires a ragged layout") 

127 return tuple( 

128 _compute_ragged_slice(global_shape, layout, local_rank=rank).local_numel 

129 for rank in range(layout.mesh.size(info.mesh_dim)) 

130 ) 

131 

132 

133def _interval_overlap_size(first: _RaggedSlice, second: _RaggedSlice) -> int: 

134 """Return the size of two half-open flat intervals' intersection.""" 

135 return max(0, min(first.flat_end, second.flat_end) - max(first.flat_start, second.flat_start)) 

136 

137 

138def _compute_ragged_all_to_all_splits( 

139 global_shape: Sequence[int], 

140 from_layout: Layout, 

141 to_layout: Layout, 

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

143 """Compute variable all-to-all splits for a local-units-only change.""" 

144 source_info = from_layout.ragged_shard 

145 target_info = to_layout.ragged_shard 

146 if source_info is None or target_info is None: 

147 raise ValueError("RaggedShard all-to-all split computation requires ragged source and target layouts") 

148 if source_info.mesh_dim != target_info.mesh_dim: 

149 raise ValueError( 

150 "RaggedShard all-to-all requires the same ragged mesh dimension, " 

151 f"got source={source_info.mesh_dim}, target={target_info.mesh_dim}" 

152 ) 

153 if source_info.placement.dims != target_info.placement.dims: 

154 raise ValueError( 

155 "RaggedShard all-to-all only supports local_units changes; dims must stay unchanged, " 

156 f"got source={source_info.placement.dims!r}, target={target_info.placement.dims!r}" 

157 ) 

158 

159 mesh_dim = source_info.mesh_dim 

160 source_rank = from_layout.mesh.get_local_rank(mesh_dim) 

161 target_rank = to_layout.mesh.get_local_rank(mesh_dim) 

162 source_interval = _compute_ragged_slice(global_shape, from_layout, source_rank) 

163 target_interval = _compute_ragged_slice(global_shape, to_layout, target_rank) 

164 input_splits = tuple( 

165 _interval_overlap_size( 

166 source_interval, 

167 _compute_ragged_slice(global_shape, to_layout, rank), 

168 ) 

169 for rank in range(to_layout.mesh.size(mesh_dim)) 

170 ) 

171 output_splits = tuple( 

172 _interval_overlap_size( 

173 _compute_ragged_slice(global_shape, from_layout, rank), 

174 target_interval, 

175 ) 

176 for rank in range(from_layout.mesh.size(mesh_dim)) 

177 ) 

178 return input_splits, output_splits 

179 

180 

181def _slice_ragged_tensor(tensor: Tensor, layout: Layout) -> Tensor: 

182 """Return an independent flat local shard from a complete global tensor.""" 

183 if hasattr(tensor, "is_contiguous") and not tensor.is_contiguous(): 

184 raise ValueError("distribute_tensor with RaggedShard requires a contiguous tensor") 

185 ragged_slice = _compute_ragged_slice(tuple(tensor.shape), layout) 

186 flat_tensor = tensor.reshape((-1,)) 

187 return flat_tensor[ragged_slice.flat_start:ragged_slice.flat_end].clone() 

188 

189 

190def _scatter_ragged_tensor( 

191 tensor: Tensor, 

192 layout: Layout, 

193 src_data_rank: int, 

194) -> Tensor: 

195 """Distribute variable flat shards from one group-relative source rank.""" 

196 if hasattr(tensor, "is_contiguous") and not tensor.is_contiguous(): 

197 raise ValueError("distribute_tensor with RaggedShard requires a contiguous tensor") 

198 info = layout.ragged_shard 

199 ragged_slice = _compute_ragged_slice(tuple(tensor.shape), layout) 

200 flat_tensor = tensor.reshape((-1,)) 

201 output = torch.empty( 

202 (ragged_slice.local_numel,), 

203 dtype=tensor.dtype, 

204 device=getattr(tensor, "device", None), 

205 ) 

206 

207 scatter_list = None 

208 if layout.mesh.get_local_rank(info.mesh_dim) == src_data_rank: 

209 scatter_list = [] 

210 for destination_rank in range(len(info.placement.local_units)): 

211 destination_slice = _compute_ragged_slice( 

212 tuple(tensor.shape), 

213 layout, 

214 local_rank=destination_rank, 

215 ) 

216 scatter_list.append( 

217 flat_tensor[destination_slice.flat_start:destination_slice.flat_end] 

218 ) 

219 

220 return mesh_scatter_ragged( 

221 output, 

222 scatter_list, 

223 layout.mesh, 

224 info.mesh_dim, 

225 group_src=src_data_rank, 

226 )