Coverage for / home / jenkins / .local / lib / python3.10 / site-packages / hyper_parallel / core / dtensor / tensor_redistribution.py: 71%
222 statements
« prev ^ index » next coverage.py v7.13.1, created at 2026-07-06 05:41 +0800
1# Copyright 2025-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"""tensor_redistribution"""
17from hyper_parallel.core.dtensor.dtensor import DTensor
18from hyper_parallel.core.dtensor.redistribute_infer import RedistributionOperatorInfer
19from hyper_parallel.platform import get_platform
20platform = get_platform()
23def _construct_layout_tuple_for_transform_operator_list(from_layout, to_layout, from_full_shape):
24 """_construct_layout_tuple_for_transform_operator_list"""
25 from_layout_dict = from_layout.to_dict()
26 to_layout_dict = to_layout.to_dict()
27 from_layout_tuple = (
28 from_layout_dict["mesh_shape"], from_layout_dict["tensor_map"], list(from_full_shape)
29 )
30 # NOTE: consider reshape scenario when to_full_shape differs from from_full_shape
31 to_layout_tuple = (
32 to_layout_dict["mesh_shape"], to_layout_dict["tensor_map"], list(from_full_shape)
33 )
34 return from_layout_tuple, to_layout_tuple
37class TensorRedistribution:
38 """
39 TensorRedistribution.
40 """
41 def __init__(self):
42 self.is_init = False
43 self.rank_id = None # current rank_id (global)
44 self._transform_cache = {}
45 self._construct_op_operator = {
46 "Reshape": self._construct_reshape,
47 "AllConcat": self._construct_all_concat,
48 "StridedSlice": self._construct_strided_slice,
49 "all_concat": TensorRedistribution._construct_all_concat_new,
50 "all_split": self._construct_all_split,
51 "all_to_all": self._construct_all_to_all
52 }
54 @staticmethod
55 def _construct_reshape(x, *args):
56 """args: (*shape)"""
57 return x.view(args)
59 @staticmethod
60 def _construct_all_concat(x, *args):
61 """args: (*rank_list, concat_dim)"""
62 rank_list = args[0:-1]
63 concat_dim = args[-1]
64 group = platform.create_group(rank_list)
65 concat_size = len(rank_list)
66 return platform.differentiable_all_gather_concat(x, group, concat_size, concat_dim, rank_list)
69 @staticmethod
70 def _construct_strided_slice(x, *args):
71 """args: (begin, end, strides)"""
72 dims = len(args) // 3
73 return platform.construct_strided_slice(x, args[0: dims], args[dims: 2 * dims], args[2 * dims:])
75 @staticmethod
76 def _construct_all_concat_new(x, *args):
77 """args: (concat_dim, concat_size, group)"""
78 rank_list = args[2]
79 concat_dim = args[0]
80 concat_size = args[1]
81 group = platform.create_group(rank_list)
82 return platform.differentiable_all_gather_concat(x, group, concat_size, concat_dim, rank_list)
84 def _construct_all_split(self, x, *args):
85 """args: (split_dim, split_size, group)"""
86 rank_list = list(args[2])
87 split_dim = args[0]
88 split_size = args[1]
89 idx = rank_list.index(self.rank_id)
90 return platform.chunk(x, split_dim, split_size, idx)
92 @staticmethod
93 def _construct_all_to_all(x, *args):
94 """args: (split_dim, concat_dim, permute_size, group)"""
95 split_dim, concat_dim, split_count, rank_list = args
96 group = platform.create_group(rank_list)
97 original_shape = x.shape
99 dim_size = original_shape[split_dim]
100 if dim_size % split_count != 0:
101 raise ValueError(f"Dimension {split_dim} with size {dim_size} "
102 f"cannot be evenly split into {split_count} parts")
104 split_size = dim_size // split_count
105 final_shape = list(original_shape)
106 if split_dim != concat_dim:
107 final_shape[split_dim] = split_size
108 final_shape[concat_dim] = final_shape[concat_dim] * split_count
109 final_shape = tuple(final_shape)
111 pre_special_handle = all(original_shape[i] == 1 for i in range(split_dim))
112 if pre_special_handle:
113 reshape_shape = (split_count * split_size,) + original_shape[split_dim + 1:]
114 x_reshaped = x.view(reshape_shape)
115 else:
116 reshape_dims = list(original_shape)
117 reshape_dims[split_dim] = split_count
118 reshape_dims.insert(split_dim + 1, split_size)
120 trans_dims = list(range(len(reshape_dims)))
121 trans_dims.remove(split_dim)
122 trans_dims.insert(0, split_dim)
124 x_reshaped = x.reshape(reshape_dims).permute(trans_dims).contiguous()
126 reshape_shape = list(x_reshaped.shape)
127 reshape_shape[0] = reshape_shape[0] * reshape_shape[1]
128 reshape_shape.pop(1)
129 reshape_shape = tuple(reshape_shape)
130 x_reshaped = x_reshaped.reshape(reshape_shape)
131 x_reshaped = x_reshaped.contiguous()
132 output_tensor = platform.differentiable_all_to_all(
133 input_data=x_reshaped,
134 output_shape=reshape_shape,
135 group=group
136 )
138 post_special_handle = all(final_shape[i] == 1 for i in range(concat_dim))
139 if post_special_handle:
140 return output_tensor.view(final_shape)
142 # When pre_special_handle collapsed leading size-1 dims, the A2A was executed
143 # in a reduced-rank space where the effective concat axis is shifted left by
144 # split_dim positions. Use recon_concat_dim for all post-A2A reshaping so
145 # that split_count is merged into the correct dimension.
146 recon_concat_dim = (concat_dim - split_dim) if pre_special_handle else concat_dim
148 output_reshape = list(output_tensor.shape)
149 output_reshape[0] = split_count
150 output_reshape.insert(1, output_tensor.shape[0] // split_count)
152 out_trans_dims = list(range(len(output_reshape)))
153 first_dim = out_trans_dims.pop(0)
154 if recon_concat_dim >= len(out_trans_dims):
155 out_trans_dims.append(first_dim)
156 else:
157 out_trans_dims.insert(recon_concat_dim, first_dim)
159 final_output = output_tensor.reshape(output_reshape).permute(out_trans_dims).contiguous()
161 final_reshape = list(final_output.shape)
162 if recon_concat_dim < len(final_reshape) - 1:
163 final_reshape[recon_concat_dim] = (
164 final_reshape[recon_concat_dim] * final_reshape[recon_concat_dim + 1]
165 )
166 final_reshape.pop(recon_concat_dim + 1)
168 result = final_output.reshape(final_reshape)
169 if pre_special_handle:
170 result = result.view(final_shape)
171 return result
173 @staticmethod
174 def _apply_eazy_redistribute(src_layout, dst_layout):
175 """_apply_eazy_redistribute"""
176 if (src_layout.mesh_shape != dst_layout.mesh_shape or
177 src_layout.rank_list != dst_layout.rank_list):
178 return False
180 tensor_map_size = len(src_layout.tensor_map)
181 if len(dst_layout.tensor_map) != tensor_map_size:
182 return False
183 return True
185 def _redistribution_without_shape(self, local_x, src_layout, dst_layout, key, rank_list):
186 """_redistribution_without_shape"""
187 inferrer = RedistributionOperatorInfer(
188 dev_mat=src_layout.mesh_shape,
189 in_tensor_map=list(src_layout.tensor_map),
190 out_tensor_map=list(dst_layout.tensor_map)
191 )
192 op_list = inferrer.infer_ops_list(self.rank_id, rank_list)
193 self._transform_cache[key] = op_list
194 for op in op_list:
195 local_x = self._construct_op_operator[op[0]](local_x, *op[1])
196 return local_x
198 def redistribution(self, input_x, to_layout):
199 """tensor redistribution"""
200 x_layout = input_x.layout
201 x = input_x
202 if input_x.layout.is_partial():
203 # Solve partial status first
204 if input_x.layout.mesh_shape == to_layout.mesh_shape:
205 x = self.reduce_partial(input_x, to_layout)
206 else:
207 x = self.reduce_partial(input_x, x_layout)
209 from_layout = x.layout
210 if not self.is_init:
211 self.rank_id = platform.get_rank()
212 self.is_init = True
213 if from_layout.rank_list != to_layout.rank_list:
214 raise ValueError(f"The from_layout rank list: {from_layout.rank_list} is not equal to "
215 f"to_layout rank list: {to_layout.rank_list}")
216 key = from_layout.compact_str + to_layout.compact_str + str(self.rank_id)
217 if key in self._transform_cache:
218 x = x.to_local()
219 transform_operator_list = self._transform_cache[key]
220 for transform_operator in transform_operator_list:
221 x = self._construct_op_operator[transform_operator[0]](x, *transform_operator[1])
222 return DTensor.from_local(x, to_layout.mesh, to_layout.alias_placements)
224 full_shape = x.shape
225 key_and_shape = key + str(full_shape)
226 x = x.to_local()
227 if key_and_shape in self._transform_cache:
228 transform_operator_list = self._transform_cache[key_and_shape]
229 for transform_operator in transform_operator_list:
230 x = self._construct_op_operator[transform_operator[0]](x, *transform_operator[1])
231 return DTensor.from_local(x, to_layout.mesh, to_layout.alias_placements)
233 rank_list = from_layout.rank_list
234 if self._apply_eazy_redistribute(from_layout, to_layout):
235 if from_layout.is_partial():
236 from_layout.reset_partial()
237 x = self._redistribution_without_shape(x, from_layout, to_layout, key, rank_list)
238 else:
239 transform_operator_list = self._infer_transform_operator_list(from_layout, to_layout,
240 full_shape, key_and_shape, rank_list)
241 for transform_operator in transform_operator_list:
242 x = self._construct_op_operator[transform_operator[0]](x, *transform_operator[1])
243 return DTensor.from_local(x, to_layout.mesh, to_layout.alias_placements)
245 def _infer_transform_operator_list(self, from_layout, to_layout, from_full_shape, key, rank_list):
246 """infer transform operator list"""
247 from_layout_tuple, to_layout_tuple = \
248 _construct_layout_tuple_for_transform_operator_list(from_layout, to_layout, from_full_shape)
249 self._transform_cache[key] = \
250 platform.get_tensor_transform().transform_tensor_sharding(from_layout_tuple, to_layout_tuple,
251 rank_list, False, self.rank_id)
252 return self._transform_cache[key]
254 @staticmethod
255 def _allreduce_along_dev_dim(x, op, layout, dev_dim):
256 """Do allreduce at specified axis along dev_dim."""
257 group = layout.get_comm_group_by_axis(dev_dim)
258 zero_dim = x.dim() == 0
259 if zero_dim:
260 x = x.unsqueeze(0)
261 if op == 'avg':
262 dev_num = layout.mesh_shape[layout.alias_name.index(dev_dim)]
263 x = platform.differentiable_all_reduce(x, 'sum', group)
264 x = x / dev_num
265 elif op == 'all':
266 x_int32 = platform.tensor_type_cast(x.bool(), 'int32') # True→1, False→0
267 x = platform.differentiable_all_reduce(x_int32, 'all', group)
268 x = x.bool()
269 else:
270 x = platform.differentiable_all_reduce(x, op, group)
271 if zero_dim:
272 x = x.squeeze(0)
273 return x
275 @staticmethod
276 def _reduce_scatter_along_dev_dim_with_axis(x, axis, op, layout, dev_dim):
277 """Do reduce_scatter at specified axis along dev_dim."""
278 dev_num = layout.mesh_shape[layout.alias_name.index(dev_dim)]
279 group = layout.get_comm_group_by_axis(dev_dim)
280 output_tensor = platform.differentiable_reduce_scatter(x, dev_num, axis, op, group)
281 return output_tensor
283 def reduce_partial(self, input_x, to_layout):
284 """Reduce partial status."""
285 from_layout = input_x.layout
286 x = input_x
287 if from_layout is None or not from_layout.is_partial():
288 return x
290 x = x.to_local()
291 if from_layout.mesh_shape != to_layout.mesh_shape:
292 raise ValueError(f"For reduce partial, mesh_shape between from_layout and to_layout must be the same, "
293 f"but got {from_layout.mesh_shape} and {to_layout.mesh_shape}")
294 if to_layout.is_partial():
295 raise ValueError(f"For reduce partial, to_layout must be non-partial status, but got to_layout.partial: "
296 f"{to_layout.partial}")
298 dev_map_order = {}
299 for dev_axis in to_layout.alias_tensor_map:
300 if isinstance(dev_axis, tuple):
301 for i, sub_dev_axis in enumerate(dev_axis):
302 dev_map_order[sub_dev_axis] = i
303 else:
304 dev_map_order[dev_axis] = 0
306 pending_reduce_op_list = [] # List[Tuple[comm_op, op, dev_dim, reduce_dim]]
307 for dev_axis_index, op in enumerate(from_layout.partial):
308 if op is None:
309 continue
310 dev_axis = from_layout.alias_name[dev_axis_index]
311 apply_shard_dim = to_layout.get_dev_axis_apply_shard_axis(dev_axis)
312 comm_op = "ReduceScatter" if apply_shard_dim is not None else "AllReduce"
313 pending_reduce_op_list.append((comm_op, op, dev_axis, apply_shard_dim))
315 # sort reduce op
316 # 1. ReduceScatter is executed before AllReduce
317 # 2. If multiple split, the dev axis split outer will be execute first.
318 # e.g. ("cp", "tp"), will execute reduce_scatter along "cp" before "tp"
319 # 3. Lower dev_id execute before higher dev_id
320 def _reduce_pair_sort_key(reduce_pair):
321 return (reduce_pair[0] != "ReduceScatter",
322 dev_map_order.get(reduce_pair[2], 0),
323 to_layout.mesh.axis_id(reduce_pair[2]))
325 sorted_pending_reduce_op_list = sorted(pending_reduce_op_list, key=_reduce_pair_sort_key)
327 output_alias_tensor_map = list(from_layout.alias_tensor_map)
328 for reduce_op_pair in sorted_pending_reduce_op_list:
329 comm_op = reduce_op_pair[0]
330 op = reduce_op_pair[1]
331 dev_axis = reduce_op_pair[2]
332 if comm_op == "AllReduce":
333 x = TensorRedistribution._allreduce_along_dev_dim(x, op, from_layout, dev_axis)
334 elif comm_op == "ReduceScatter":
335 reduce_axis = reduce_op_pair[3]
336 x = self._reduce_scatter_along_dev_dim_with_axis(x, reduce_axis, op, from_layout, dev_axis)
337 if output_alias_tensor_map[reduce_axis] == "None":
338 output_alias_tensor_map[reduce_axis] = dev_axis
339 elif isinstance(output_alias_tensor_map[reduce_axis], tuple):
340 output_alias_tensor_map[reduce_axis] += (dev_axis,)
341 else:
342 output_alias_tensor_map[reduce_axis] = (output_alias_tensor_map[reduce_axis], dev_axis)
344 output_layout = from_layout(*output_alias_tensor_map)
345 output_layout.reset_partial()
346 return DTensor.from_local(x, output_layout.mesh, output_layout.alias_placements)
349_tensor_redistribution = TensorRedistribution()