Coverage for / home / jenkins / .local / lib / python3.10 / site-packages / hyper_parallel / core / dtensor / tensor_redistribution.py: 77%
313 statements
« prev ^ index » next coverage.py v7.13.1, created at 2026-08-22 04:23 +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._ragged_utils import (
19 _compute_ragged_all_to_all_splits,
20 _compute_ragged_slice,
21 _compute_ragged_splits,
22)
23from hyper_parallel.core.dtensor.dtensor import DTensor
24from hyper_parallel.core.dtensor.layout import Layout, RaggedShardInfo
25from hyper_parallel.core.dtensor.redistribute_infer import RedistributionOperatorInfer
26from hyper_parallel.platform import get_platform
27platform = get_platform()
29logger = logging.getLogger(__name__)
32def _construct_layout_tuple_for_transform_operator_list(from_layout, to_layout, from_full_shape):
33 """_construct_layout_tuple_for_transform_operator_list"""
34 from_layout_dict = from_layout.to_dict()
35 to_layout_dict = to_layout.to_dict()
36 from_layout_tuple = (
37 from_layout_dict["mesh_shape"], from_layout_dict["tensor_map"], list(from_full_shape)
38 )
39 # NOTE: consider reshape scenario when to_full_shape differs from from_full_shape
40 to_layout_tuple = (
41 to_layout_dict["mesh_shape"], to_layout_dict["tensor_map"], list(from_full_shape)
42 )
43 return from_layout_tuple, to_layout_tuple
46class TensorRedistribution:
47 """
48 TensorRedistribution.
49 """
50 def __init__(self):
51 self.is_init = False
52 self.rank_id = None # current rank_id (global)
53 self._transform_cache = {}
54 self._construct_op_operator = {
55 "Reshape": self._construct_reshape,
56 "AllConcat": self._construct_all_concat,
57 "StridedSlice": self._construct_strided_slice,
58 "all_concat": TensorRedistribution._construct_all_concat_new,
59 "all_split": self._construct_all_split,
60 "all_to_all": self._construct_all_to_all
61 }
63 @staticmethod
64 def _construct_reshape(x, *args):
65 """args: (*shape)"""
66 return x.view(args)
68 @staticmethod
69 def _construct_all_concat(x, *args):
70 """args: (*rank_list, concat_dim)"""
71 rank_list = args[0:-1]
72 concat_dim = args[-1]
73 group = platform.create_group(rank_list)
74 concat_size = len(rank_list)
75 logger.debug(
76 "differentiable_all_gather_concat: input_shape=%s, concat_dim=%d, "
77 "concat_size=%d, rank_list=%s",
78 tuple(x.shape), concat_dim, concat_size, rank_list,
79 )
80 return platform.differentiable_all_gather_concat(x, group, concat_size, concat_dim, rank_list)
83 @staticmethod
84 def _construct_strided_slice(x, *args):
85 """args: (begin, end, strides)"""
86 dims = len(args) // 3
87 return platform.construct_strided_slice(x, args[0: dims], args[dims: 2 * dims], args[2 * dims:])
89 @staticmethod
90 def _construct_all_concat_new(x, *args):
91 """args: (concat_dim, concat_size, group)"""
92 rank_list = args[2]
93 concat_dim = args[0]
94 concat_size = args[1]
95 group = platform.create_group(rank_list)
96 logger.debug(
97 "differentiable_all_gather_concat: input_shape=%s, concat_dim=%d, "
98 "concat_size=%d, rank_list=%s",
99 tuple(x.shape), concat_dim, concat_size, rank_list,
100 )
101 return platform.differentiable_all_gather_concat(x, group, concat_size, concat_dim, rank_list)
103 def _construct_all_split(self, x, *args):
104 """args: (split_dim, split_size, group)"""
105 rank_list = list(args[2])
106 split_dim = args[0]
107 split_size = args[1]
108 idx = rank_list.index(self.rank_id)
109 return platform.chunk(x, split_dim, split_size, idx)
111 @staticmethod
112 def _construct_all_to_all(x, *args):
113 """args: (split_dim, concat_dim, permute_size, group)"""
114 split_dim, concat_dim, split_count, rank_list = args
115 group = platform.create_group(rank_list)
116 logger.debug(
117 "differentiable_all_to_all: input_shape=%s, split_dim=%d, "
118 "concat_dim=%d, split_count=%d, rank_list=%s",
119 tuple(x.shape), split_dim, concat_dim, split_count, rank_list,
120 )
121 original_shape = x.shape
123 dim_size = original_shape[split_dim]
124 if dim_size % split_count != 0:
125 raise ValueError(f"Dimension {split_dim} with size {dim_size} "
126 f"cannot be evenly split into {split_count} parts")
128 split_size = dim_size // split_count
129 final_shape = list(original_shape)
130 if split_dim != concat_dim:
131 final_shape[split_dim] = split_size
132 final_shape[concat_dim] = final_shape[concat_dim] * split_count
133 final_shape = tuple(final_shape)
135 pre_special_handle = all(original_shape[i] == 1 for i in range(split_dim))
136 if pre_special_handle:
137 reshape_shape = (split_count * split_size,) + original_shape[split_dim + 1:]
138 x_reshaped = x.view(reshape_shape)
139 else:
140 reshape_dims = list(original_shape)
141 reshape_dims[split_dim] = split_count
142 reshape_dims.insert(split_dim + 1, split_size)
144 trans_dims = list(range(len(reshape_dims)))
145 trans_dims.remove(split_dim)
146 trans_dims.insert(0, split_dim)
148 x_reshaped = x.reshape(reshape_dims).permute(trans_dims).contiguous()
150 reshape_shape = list(x_reshaped.shape)
151 reshape_shape[0] = reshape_shape[0] * reshape_shape[1]
152 reshape_shape.pop(1)
153 reshape_shape = tuple(reshape_shape)
154 x_reshaped = x_reshaped.reshape(reshape_shape)
155 x_reshaped = x_reshaped.contiguous()
156 output_tensor = platform.differentiable_all_to_all(
157 input_data=x_reshaped,
158 output_shape=reshape_shape,
159 group=group
160 )
162 post_special_handle = all(final_shape[i] == 1 for i in range(concat_dim))
163 if post_special_handle:
164 return output_tensor.view(final_shape)
166 # When pre_special_handle collapsed leading size-1 dims, the A2A was executed
167 # in a reduced-rank space where the effective concat axis is shifted left by
168 # split_dim positions. Use recon_concat_dim for all post-A2A reshaping so
169 # that split_count is merged into the correct dimension.
170 recon_concat_dim = (concat_dim - split_dim) if pre_special_handle else concat_dim
172 output_reshape = list(output_tensor.shape)
173 output_reshape[0] = split_count
174 output_reshape.insert(1, output_tensor.shape[0] // split_count)
176 out_trans_dims = list(range(len(output_reshape)))
177 first_dim = out_trans_dims.pop(0)
178 if recon_concat_dim >= len(out_trans_dims):
179 out_trans_dims.append(first_dim)
180 else:
181 out_trans_dims.insert(recon_concat_dim, first_dim)
183 final_output = output_tensor.reshape(output_reshape).permute(out_trans_dims).contiguous()
185 final_reshape = list(final_output.shape)
186 if recon_concat_dim < len(final_reshape) - 1:
187 final_reshape[recon_concat_dim] = (
188 final_reshape[recon_concat_dim] * final_reshape[recon_concat_dim + 1]
189 )
190 final_reshape.pop(recon_concat_dim + 1)
192 result = final_output.reshape(final_reshape)
193 if pre_special_handle:
194 result = result.view(final_shape)
195 return result
197 @staticmethod
198 def _apply_eazy_redistribute(src_layout, dst_layout):
199 """_apply_eazy_redistribute"""
200 if (src_layout.mesh_shape != dst_layout.mesh_shape or
201 src_layout.rank_list != dst_layout.rank_list):
202 return False
204 tensor_map_size = len(src_layout.tensor_map)
205 if len(dst_layout.tensor_map) != tensor_map_size:
206 return False
207 return True
209 def _redistribution_without_shape(self, local_x, src_layout, dst_layout, key, rank_list):
210 """_redistribution_without_shape"""
211 inferrer = RedistributionOperatorInfer(
212 dev_mat=src_layout.mesh_shape,
213 in_tensor_map=list(src_layout.tensor_map),
214 out_tensor_map=list(dst_layout.tensor_map)
215 )
216 op_list = inferrer.infer_ops_list(self.rank_id, rank_list)
217 self._transform_cache[key] = op_list
218 for op in op_list:
219 local_x = self._construct_op_operator[op[0]](local_x, *op[1])
220 return local_x
222 @staticmethod
223 def _to_normal_layout(layout: Layout, tensor_dim: int) -> Layout:
224 """Build the normal view consumed by the legacy redistribution path."""
225 normal_layout = Layout.from_device_mesh(layout.mesh)
226 normal_layout.set_placements(layout.normal_placements)
227 normal_layout.placement_to_tensor_map(tensor_dim)
228 return normal_layout
230 def _redistribute_ragged(self, input_x: DTensor, to_layout: Layout) -> DTensor:
231 """Adapt RaggedShard layouts to the supported redistribution primitives."""
232 from_layout = input_x.layout
233 source_info = from_layout.ragged_shard
234 target_info = to_layout.ragged_shard
235 source_is_ragged = isinstance(source_info, RaggedShardInfo)
236 target_is_ragged = isinstance(target_info, RaggedShardInfo)
237 tensor_dim = len(input_x.shape)
239 if source_is_ragged and target_is_ragged:
240 source_normal_layout = self._to_normal_layout(from_layout, tensor_dim)
241 target_normal_layout = self._to_normal_layout(to_layout, tensor_dim)
242 if (
243 source_info.mesh_dim == target_info.mesh_dim
244 and source_info.placement.dims == target_info.placement.dims
245 and source_normal_layout == target_normal_layout
246 ):
247 return self.ragged_to_ragged(input_x, to_layout)
248 return self.ragged_to_ragged_via_replicate(
249 input_x,
250 source_normal_layout,
251 to_layout,
252 )
254 if source_is_ragged:
255 source_normal_layout = self._to_normal_layout(from_layout, tensor_dim)
256 normal = self.ragged_to_normal(input_x, source_normal_layout)
257 if normal.layout == to_layout:
258 return normal
259 return self._redistribution_normal(normal, to_layout)
261 target_normal_layout = self._to_normal_layout(to_layout, tensor_dim)
262 normal = input_x
263 if from_layout != target_normal_layout:
264 normal = self._redistribution_normal(normal, target_normal_layout)
265 return self.normal_to_ragged(normal, to_layout)
267 def redistribution(self, input_x, to_layout):
268 """tensor redistribution"""
269 x_layout = input_x.layout
270 x = input_x
271 if input_x.layout.is_partial():
272 # Solve partial status first
273 if input_x.layout.mesh_shape == to_layout.mesh_shape:
274 x = self.reduce_partial(input_x, to_layout)
275 else:
276 x = self.reduce_partial(input_x, x_layout)
278 from_layout = x.layout
279 if from_layout.rank_list != to_layout.rank_list:
280 raise ValueError(f"The from_layout rank list: {from_layout.rank_list} is not equal to "
281 f"to_layout rank list: {to_layout.rank_list}")
282 if isinstance(from_layout.ragged_shard, RaggedShardInfo) or isinstance(
283 to_layout.ragged_shard, RaggedShardInfo
284 ):
285 return self._redistribute_ragged(x, to_layout)
286 return self._redistribution_normal(x, to_layout)
288 def _redistribution_normal(self, input_x: DTensor, to_layout: Layout) -> DTensor:
289 """Run the existing redistribution path for normal layouts."""
290 from_layout = input_x.layout
291 x = input_x
292 if not self.is_init:
293 self.rank_id = platform.get_rank()
294 self.is_init = True
295 key = from_layout.compact_str + to_layout.compact_str + str(self.rank_id)
296 if key in self._transform_cache:
297 x = x.to_local()
298 transform_operator_list = self._transform_cache[key]
299 for transform_operator in transform_operator_list:
300 x = self._construct_op_operator[transform_operator[0]](x, *transform_operator[1])
301 return DTensor.from_local(x, to_layout.mesh, to_layout.alias_placements)
303 full_shape = x.shape
304 key_and_shape = key + str(full_shape)
305 x = x.to_local()
306 if key_and_shape in self._transform_cache:
307 transform_operator_list = self._transform_cache[key_and_shape]
308 for transform_operator in transform_operator_list:
309 x = self._construct_op_operator[transform_operator[0]](x, *transform_operator[1])
310 return DTensor.from_local(x, to_layout.mesh, to_layout.alias_placements)
312 rank_list = from_layout.rank_list
313 if self._apply_eazy_redistribute(from_layout, to_layout):
314 if from_layout.is_partial():
315 from_layout.reset_partial()
316 x = self._redistribution_without_shape(x, from_layout, to_layout, key, rank_list)
317 else:
318 transform_operator_list = self._infer_transform_operator_list(from_layout, to_layout,
319 full_shape, key_and_shape, rank_list)
320 for transform_operator in transform_operator_list:
321 x = self._construct_op_operator[transform_operator[0]](x, *transform_operator[1])
322 return DTensor.from_local(x, to_layout.mesh, to_layout.alias_placements)
324 @staticmethod
325 def ragged_to_normal(input_x: DTensor, to_layout: Layout) -> DTensor:
326 """Gather one flat RaggedShard into its Replicate normal view."""
327 from_layout = input_x.layout
328 info = from_layout.ragged_shard
329 if not isinstance(info, RaggedShardInfo):
330 raise ValueError("ragged_to_normal requires a RaggedShard source layout")
331 global_shape = tuple(input_x.shape)
332 output_splits = _compute_ragged_splits(global_shape, from_layout)
333 local_tensor = input_x.to_local()
334 if len(output_splits) == 1:
335 gathered = local_tensor
336 else:
337 group = from_layout.mesh.get_group(info.mesh_dim)
338 gathered = platform.differentiable_variable_all_gather(
339 local_tensor,
340 output_splits,
341 group,
342 )
343 return DTensor.from_local_with_layout(
344 gathered.reshape(global_shape),
345 to_layout,
346 shape=global_shape,
347 )
349 @staticmethod
350 def normal_to_ragged(input_x: DTensor, to_layout: Layout) -> DTensor:
351 """Slice a normal-view tensor into the target flat RaggedShard."""
352 global_shape = tuple(input_x.shape)
353 local_slice = _compute_ragged_slice(global_shape, to_layout)
354 flat_tensor = input_x.to_local().reshape((-1,))
355 local_tensor = flat_tensor[
356 local_slice.flat_start:local_slice.flat_end
357 ].clone()
358 return DTensor.from_local_with_layout(
359 local_tensor,
360 to_layout,
361 shape=global_shape,
362 )
364 @staticmethod
365 def ragged_to_ragged_via_replicate(
366 input_x: DTensor,
367 replicate_layout: Layout,
368 to_layout: Layout,
369 ) -> DTensor:
370 """Redistribute different RaggedShard dims through a Replicate layout."""
371 from_layout = input_x.layout
372 if (
373 from_layout.mesh_shape != to_layout.mesh_shape
374 or from_layout.rank_list != to_layout.rank_list
375 ):
376 raise ValueError("ragged_to_ragged only supports changes on the same device mesh")
377 replicated = TensorRedistribution.ragged_to_normal(
378 input_x,
379 replicate_layout,
380 )
381 return TensorRedistribution.normal_to_ragged(replicated, to_layout)
383 @staticmethod
384 def ragged_to_ragged(input_x: DTensor, to_layout: Layout) -> DTensor:
385 """Redistribute a local-units-only RaggedShard change with variable all-to-all."""
386 from_layout = input_x.layout
387 source_info = from_layout.ragged_shard
388 if (
389 from_layout.mesh_shape != to_layout.mesh_shape
390 or from_layout.rank_list != to_layout.rank_list
391 ):
392 raise ValueError("ragged_to_ragged only supports local_units changes on the same device mesh")
394 global_shape = tuple(input_x.shape)
395 input_splits, output_splits = _compute_ragged_all_to_all_splits(
396 global_shape,
397 from_layout,
398 to_layout,
399 )
400 flat_input = input_x.to_local().reshape((-1,))
402 if flat_input.shape[0] != sum(input_splits):
403 raise ValueError(
404 "RaggedShard source storage does not match all-to-all splits, "
405 f"got local_numel={flat_input.shape[0]}, input_splits={input_splits!r}"
406 )
407 if len(input_splits) == 1:
408 flat_output = flat_input
409 else:
410 group = from_layout.mesh.get_group(source_info.mesh_dim)
411 flat_output = platform.differentiable_all_to_all_single(
412 flat_input,
413 input_splits,
414 output_splits,
415 group,
416 )
417 if flat_output.shape[0] != sum(output_splits):
418 raise ValueError(
419 "RaggedShard target storage does not match all-to-all splits, "
420 f"got local_numel={flat_output.shape[0]}, output_splits={output_splits!r}"
421 )
422 return DTensor.from_local_with_layout(
423 flat_output,
424 to_layout,
425 shape=global_shape,
426 )
428 def _infer_transform_operator_list(self, from_layout, to_layout, from_full_shape, key, rank_list):
429 """infer transform operator list"""
430 from_layout_tuple, to_layout_tuple = \
431 _construct_layout_tuple_for_transform_operator_list(from_layout, to_layout, from_full_shape)
432 self._transform_cache[key] = \
433 platform.get_tensor_transform().transform_tensor_sharding(from_layout_tuple, to_layout_tuple,
434 rank_list, False, self.rank_id)
435 return self._transform_cache[key]
437 @staticmethod
438 def _allreduce_along_dev_dim(x, op, layout, dev_dim):
439 """Do allreduce at specified axis along dev_dim."""
440 logger.debug(
441 "differentiable_all_reduce: input_shape=%s, op=%s, dev_dim=%s",
442 tuple(x.shape), op, dev_dim,
443 )
444 group = layout.get_comm_group_by_axis(dev_dim)
445 zero_dim = x.dim() == 0
446 if zero_dim:
447 x = x.unsqueeze(0)
448 if op == 'avg':
449 dev_num = layout.mesh_shape[layout.alias_name.index(dev_dim)]
450 x = platform.differentiable_all_reduce(x, 'sum', group)
451 x = x / dev_num
452 elif op == 'all':
453 x_int32 = platform.tensor_type_cast(x.bool(), 'int32') # True→1, False→0
454 x = platform.differentiable_all_reduce(x_int32, 'all', group)
455 x = x.bool()
456 else:
457 x = platform.differentiable_all_reduce(x, op, group)
458 if zero_dim:
459 x = x.squeeze(0)
460 return x
462 @staticmethod
463 def _reduce_scatter_along_dev_dim_with_axis(x, axis, op, layout, dev_dim):
464 """Do reduce_scatter at specified axis along dev_dim."""
465 dev_num = layout.mesh_shape[layout.alias_name.index(dev_dim)]
466 logger.debug(
467 "differentiable_reduce_scatter: input_shape=%s, axis=%d, "
468 "op=%s, dev_dim=%s, dev_num=%d",
469 tuple(x.shape), axis, op, dev_dim, dev_num,
470 )
471 group = layout.get_comm_group_by_axis(dev_dim)
472 output_tensor = platform.differentiable_reduce_scatter(x, dev_num, axis, op, group)
473 return output_tensor
475 def reduce_partial(self, input_x, to_layout):
476 """Reduce partial status."""
477 from_layout = input_x.layout
478 x = input_x
479 if from_layout is None or not from_layout.is_partial():
480 return x
482 x = x.to_local()
483 if from_layout.mesh_shape != to_layout.mesh_shape:
484 raise ValueError(f"For reduce partial, mesh_shape between from_layout and to_layout must be the same, "
485 f"but got {from_layout.mesh_shape} and {to_layout.mesh_shape}")
486 if to_layout.is_partial():
487 raise ValueError(f"For reduce partial, to_layout must be non-partial status, but got to_layout.partial: "
488 f"{to_layout.partial}")
490 dev_map_order = {}
491 for dev_axis in to_layout.alias_tensor_map:
492 if isinstance(dev_axis, tuple):
493 for i, sub_dev_axis in enumerate(dev_axis):
494 dev_map_order[sub_dev_axis] = i
495 else:
496 dev_map_order[dev_axis] = 0
498 pending_reduce_op_list = [] # List[Tuple[comm_op, op, dev_dim, reduce_dim]]
499 for dev_axis_index, op in enumerate(from_layout.partial):
500 if op is None:
501 continue
502 dev_axis = from_layout.alias_name[dev_axis_index]
503 apply_shard_dim = to_layout.get_dev_axis_apply_shard_axis(dev_axis)
504 comm_op = "ReduceScatter" if apply_shard_dim is not None else "AllReduce"
505 pending_reduce_op_list.append((comm_op, op, dev_axis, apply_shard_dim))
507 # sort reduce op
508 # 1. ReduceScatter is executed before AllReduce
509 # 2. If multiple split, the dev axis split outer will be execute first.
510 # e.g. ("cp", "tp"), will execute reduce_scatter along "cp" before "tp"
511 # 3. Lower dev_id execute before higher dev_id
512 def _reduce_pair_sort_key(reduce_pair):
513 return (reduce_pair[0] != "ReduceScatter",
514 dev_map_order.get(reduce_pair[2], 0),
515 to_layout.mesh.axis_id(reduce_pair[2]))
517 sorted_pending_reduce_op_list = sorted(pending_reduce_op_list, key=_reduce_pair_sort_key)
519 output_alias_tensor_map = list(from_layout.alias_tensor_map)
520 for reduce_op_pair in sorted_pending_reduce_op_list:
521 comm_op = reduce_op_pair[0]
522 op = reduce_op_pair[1]
523 dev_axis = reduce_op_pair[2]
524 if comm_op == "AllReduce":
525 x = TensorRedistribution._allreduce_along_dev_dim(x, op, from_layout, dev_axis)
526 elif comm_op == "ReduceScatter":
527 reduce_axis = reduce_op_pair[3]
528 x = self._reduce_scatter_along_dev_dim_with_axis(x, reduce_axis, op, from_layout, dev_axis)
529 if output_alias_tensor_map[reduce_axis] == "None":
530 output_alias_tensor_map[reduce_axis] = dev_axis
531 elif isinstance(output_alias_tensor_map[reduce_axis], tuple):
532 output_alias_tensor_map[reduce_axis] += (dev_axis,)
533 else:
534 output_alias_tensor_map[reduce_axis] = (output_alias_tensor_map[reduce_axis], dev_axis)
536 output_layout = from_layout(*output_alias_tensor_map)
537 output_layout.reset_partial()
538 return DTensor.from_local(x, output_layout.mesh, output_layout.alias_placements)
541_tensor_redistribution = TensorRedistribution()