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