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