Coverage for / home / jenkins / .local / lib / python3.10 / site-packages / hyper_parallel / core / dtensor / tensor_redistribution.py: 77%
315 statements
« prev ^ index » next coverage.py v7.13.1, created at 2026-08-25 04:27 +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 if from_layout.has_uneven_shard or to_layout.has_uneven_shard:
287 # FSDP uneven shard/unshard uses its private padded collectives.
288 raise NotImplementedError(
289 "DTensor redistribution does not support uneven chunk-sharded layouts."
290 )
291 return self._redistribution_normal(x, to_layout)
293 def _redistribution_normal(self, input_x: DTensor, to_layout: Layout) -> DTensor:
294 """Run the existing redistribution path for normal layouts."""
295 from_layout = input_x.layout
296 x = input_x
297 if not self.is_init:
298 self.rank_id = platform.get_rank()
299 self.is_init = True
300 key = from_layout.compact_str + to_layout.compact_str + str(self.rank_id)
301 if key in self._transform_cache:
302 x = x.to_local()
303 transform_operator_list = self._transform_cache[key]
304 for transform_operator in transform_operator_list:
305 x = self._construct_op_operator[transform_operator[0]](x, *transform_operator[1])
306 return DTensor.from_local(x, to_layout.mesh, to_layout.alias_placements)
308 full_shape = x.shape
309 key_and_shape = key + str(full_shape)
310 x = x.to_local()
311 if key_and_shape in self._transform_cache:
312 transform_operator_list = self._transform_cache[key_and_shape]
313 for transform_operator in transform_operator_list:
314 x = self._construct_op_operator[transform_operator[0]](x, *transform_operator[1])
315 return DTensor.from_local(x, to_layout.mesh, to_layout.alias_placements)
317 rank_list = from_layout.rank_list
318 if self._apply_eazy_redistribute(from_layout, to_layout):
319 if from_layout.is_partial():
320 from_layout.reset_partial()
321 x = self._redistribution_without_shape(x, from_layout, to_layout, key, rank_list)
322 else:
323 transform_operator_list = self._infer_transform_operator_list(from_layout, to_layout,
324 full_shape, key_and_shape, rank_list)
325 for transform_operator in transform_operator_list:
326 x = self._construct_op_operator[transform_operator[0]](x, *transform_operator[1])
327 return DTensor.from_local(x, to_layout.mesh, to_layout.alias_placements)
329 @staticmethod
330 def ragged_to_normal(input_x: DTensor, to_layout: Layout) -> DTensor:
331 """Gather one flat RaggedShard into its Replicate normal view."""
332 from_layout = input_x.layout
333 info = from_layout.ragged_shard
334 if not isinstance(info, RaggedShardInfo):
335 raise ValueError("ragged_to_normal requires a RaggedShard source layout")
336 global_shape = tuple(input_x.shape)
337 output_splits = _compute_ragged_splits(global_shape, from_layout)
338 local_tensor = input_x.to_local()
339 if len(output_splits) == 1:
340 gathered = local_tensor
341 else:
342 group = from_layout.mesh.get_group(info.mesh_dim)
343 gathered = platform.differentiable_variable_all_gather(
344 local_tensor,
345 output_splits,
346 group,
347 )
348 return DTensor.from_local_with_layout(
349 gathered.reshape(global_shape),
350 to_layout,
351 shape=global_shape,
352 )
354 @staticmethod
355 def normal_to_ragged(input_x: DTensor, to_layout: Layout) -> DTensor:
356 """Slice a normal-view tensor into the target flat RaggedShard."""
357 global_shape = tuple(input_x.shape)
358 local_slice = _compute_ragged_slice(global_shape, to_layout)
359 flat_tensor = input_x.to_local().reshape((-1,))
360 local_tensor = flat_tensor[
361 local_slice.flat_start:local_slice.flat_end
362 ].clone()
363 return DTensor.from_local_with_layout(
364 local_tensor,
365 to_layout,
366 shape=global_shape,
367 )
369 @staticmethod
370 def ragged_to_ragged_via_replicate(
371 input_x: DTensor,
372 replicate_layout: Layout,
373 to_layout: Layout,
374 ) -> DTensor:
375 """Redistribute different RaggedShard dims through a Replicate layout."""
376 from_layout = input_x.layout
377 if (
378 from_layout.mesh_shape != to_layout.mesh_shape
379 or from_layout.rank_list != to_layout.rank_list
380 ):
381 raise ValueError("ragged_to_ragged only supports changes on the same device mesh")
382 replicated = TensorRedistribution.ragged_to_normal(
383 input_x,
384 replicate_layout,
385 )
386 return TensorRedistribution.normal_to_ragged(replicated, to_layout)
388 @staticmethod
389 def ragged_to_ragged(input_x: DTensor, to_layout: Layout) -> DTensor:
390 """Redistribute a local-units-only RaggedShard change with variable all-to-all."""
391 from_layout = input_x.layout
392 source_info = from_layout.ragged_shard
393 if (
394 from_layout.mesh_shape != to_layout.mesh_shape
395 or from_layout.rank_list != to_layout.rank_list
396 ):
397 raise ValueError("ragged_to_ragged only supports local_units changes on the same device mesh")
399 global_shape = tuple(input_x.shape)
400 input_splits, output_splits = _compute_ragged_all_to_all_splits(
401 global_shape,
402 from_layout,
403 to_layout,
404 )
405 flat_input = input_x.to_local().reshape((-1,))
407 if flat_input.shape[0] != sum(input_splits):
408 raise ValueError(
409 "RaggedShard source storage does not match all-to-all splits, "
410 f"got local_numel={flat_input.shape[0]}, input_splits={input_splits!r}"
411 )
412 if len(input_splits) == 1:
413 flat_output = flat_input
414 else:
415 group = from_layout.mesh.get_group(source_info.mesh_dim)
416 flat_output = platform.differentiable_all_to_all_single(
417 flat_input,
418 input_splits,
419 output_splits,
420 group,
421 )
422 if flat_output.shape[0] != sum(output_splits):
423 raise ValueError(
424 "RaggedShard target storage does not match all-to-all splits, "
425 f"got local_numel={flat_output.shape[0]}, output_splits={output_splits!r}"
426 )
427 return DTensor.from_local_with_layout(
428 flat_output,
429 to_layout,
430 shape=global_shape,
431 )
433 def _infer_transform_operator_list(self, from_layout, to_layout, from_full_shape, key, rank_list):
434 """infer transform operator list"""
435 from_layout_tuple, to_layout_tuple = \
436 _construct_layout_tuple_for_transform_operator_list(from_layout, to_layout, from_full_shape)
437 self._transform_cache[key] = \
438 platform.get_tensor_transform().transform_tensor_sharding(from_layout_tuple, to_layout_tuple,
439 rank_list, False, self.rank_id)
440 return self._transform_cache[key]
442 @staticmethod
443 def _allreduce_along_dev_dim(x, op, layout, dev_dim):
444 """Do allreduce at specified axis along dev_dim."""
445 logger.debug(
446 "differentiable_all_reduce: input_shape=%s, op=%s, dev_dim=%s",
447 tuple(x.shape), op, dev_dim,
448 )
449 group = layout.get_comm_group_by_axis(dev_dim)
450 zero_dim = x.dim() == 0
451 if zero_dim:
452 x = x.unsqueeze(0)
453 if op == 'avg':
454 dev_num = layout.mesh_shape[layout.alias_name.index(dev_dim)]
455 x = platform.differentiable_all_reduce(x, 'sum', group)
456 x = x / dev_num
457 elif op == 'all':
458 x_int32 = platform.tensor_type_cast(x.bool(), 'int32') # True→1, False→0
459 x = platform.differentiable_all_reduce(x_int32, 'all', group)
460 x = x.bool()
461 else:
462 x = platform.differentiable_all_reduce(x, op, group)
463 if zero_dim:
464 x = x.squeeze(0)
465 return x
467 @staticmethod
468 def _reduce_scatter_along_dev_dim_with_axis(x, axis, op, layout, dev_dim):
469 """Do reduce_scatter at specified axis along dev_dim."""
470 dev_num = layout.mesh_shape[layout.alias_name.index(dev_dim)]
471 logger.debug(
472 "differentiable_reduce_scatter: input_shape=%s, axis=%d, "
473 "op=%s, dev_dim=%s, dev_num=%d",
474 tuple(x.shape), axis, op, dev_dim, dev_num,
475 )
476 group = layout.get_comm_group_by_axis(dev_dim)
477 output_tensor = platform.differentiable_reduce_scatter(x, dev_num, axis, op, group)
478 return output_tensor
480 def reduce_partial(self, input_x, to_layout):
481 """Reduce partial status."""
482 from_layout = input_x.layout
483 x = input_x
484 if from_layout is None or not from_layout.is_partial():
485 return x
487 x = x.to_local()
488 if from_layout.mesh_shape != to_layout.mesh_shape:
489 raise ValueError(f"For reduce partial, mesh_shape between from_layout and to_layout must be the same, "
490 f"but got {from_layout.mesh_shape} and {to_layout.mesh_shape}")
491 if to_layout.is_partial():
492 raise ValueError(f"For reduce partial, to_layout must be non-partial status, but got to_layout.partial: "
493 f"{to_layout.partial}")
495 dev_map_order = {}
496 for dev_axis in to_layout.alias_tensor_map:
497 if isinstance(dev_axis, tuple):
498 for i, sub_dev_axis in enumerate(dev_axis):
499 dev_map_order[sub_dev_axis] = i
500 else:
501 dev_map_order[dev_axis] = 0
503 pending_reduce_op_list = [] # List[Tuple[comm_op, op, dev_dim, reduce_dim]]
504 for dev_axis_index, op in enumerate(from_layout.partial):
505 if op is None:
506 continue
507 dev_axis = from_layout.alias_name[dev_axis_index]
508 apply_shard_dim = to_layout.get_dev_axis_apply_shard_axis(dev_axis)
509 comm_op = "ReduceScatter" if apply_shard_dim is not None else "AllReduce"
510 pending_reduce_op_list.append((comm_op, op, dev_axis, apply_shard_dim))
512 # sort reduce op
513 # 1. ReduceScatter is executed before AllReduce
514 # 2. If multiple split, the dev axis split outer will be execute first.
515 # e.g. ("cp", "tp"), will execute reduce_scatter along "cp" before "tp"
516 # 3. Lower dev_id execute before higher dev_id
517 def _reduce_pair_sort_key(reduce_pair):
518 return (reduce_pair[0] != "ReduceScatter",
519 dev_map_order.get(reduce_pair[2], 0),
520 to_layout.mesh.axis_id(reduce_pair[2]))
522 sorted_pending_reduce_op_list = sorted(pending_reduce_op_list, key=_reduce_pair_sort_key)
524 output_alias_tensor_map = list(from_layout.alias_tensor_map)
525 for reduce_op_pair in sorted_pending_reduce_op_list:
526 comm_op = reduce_op_pair[0]
527 op = reduce_op_pair[1]
528 dev_axis = reduce_op_pair[2]
529 if comm_op == "AllReduce":
530 x = TensorRedistribution._allreduce_along_dev_dim(x, op, from_layout, dev_axis)
531 elif comm_op == "ReduceScatter":
532 reduce_axis = reduce_op_pair[3]
533 x = self._reduce_scatter_along_dev_dim_with_axis(x, reduce_axis, op, from_layout, dev_axis)
534 if output_alias_tensor_map[reduce_axis] == "None":
535 output_alias_tensor_map[reduce_axis] = dev_axis
536 elif isinstance(output_alias_tensor_map[reduce_axis], tuple):
537 output_alias_tensor_map[reduce_axis] += (dev_axis,)
538 else:
539 output_alias_tensor_map[reduce_axis] = (output_alias_tensor_map[reduce_axis], dev_axis)
541 output_layout = from_layout(*output_alias_tensor_map)
542 output_layout.reset_partial()
543 return DTensor.from_local(x, output_layout.mesh, output_layout.alias_placements)
546_tensor_redistribution = TensorRedistribution()