Coverage for / home / jenkins / .local / lib / python3.10 / site-packages / hyper_parallel / core / dtensor / tensor_redistribution.py: 80%
381 statements
« prev ^ index » next coverage.py v7.13.1, created at 2026-11-15 07:18 +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, infer_slice_area_by_rank
30from hyper_parallel.core.dtensor.placement_types import Replicate
31from hyper_parallel.core.dtensor.redistribute_infer import RedistributionOperatorInfer
33logger = logging.getLogger(__name__)
36def _construct_layout_tuple_for_transform_operator_list(from_layout, to_layout, from_full_shape):
37 """_construct_layout_tuple_for_transform_operator_list"""
38 from_layout_dict = from_layout.to_dict()
39 to_layout_dict = to_layout.to_dict()
40 from_layout_tuple = (
41 from_layout_dict["mesh_shape"], from_layout_dict["tensor_map"], list(from_full_shape)
42 )
43 # NOTE: consider reshape scenario when to_full_shape differs from from_full_shape
44 to_layout_tuple = (
45 to_layout_dict["mesh_shape"], to_layout_dict["tensor_map"], list(from_full_shape)
46 )
47 return from_layout_tuple, to_layout_tuple
50class TensorRedistribution:
51 """
52 TensorRedistribution.
53 """
54 def __init__(self):
55 self.is_init = False
56 self.rank_id = None # current rank_id (global)
57 self._transform_cache = {}
58 self._construct_op_operator = {
59 "Reshape": self._construct_reshape,
60 "AllConcat": self._construct_all_concat,
61 "StridedSlice": self._construct_strided_slice,
62 "all_concat": TensorRedistribution._construct_all_concat_new,
63 "all_split": self._construct_all_split,
64 "all_to_all": self._construct_all_to_all
65 }
67 @staticmethod
68 def _construct_reshape(x, *args):
69 """args: (*shape)"""
70 return x.view(args)
72 @staticmethod
73 def _construct_all_concat(x, *args):
74 """args: (*rank_list, concat_dim)"""
75 rank_list = args[0:-1]
76 concat_dim = args[-1]
77 group = _utils.create_group(rank_list)
78 concat_size = len(rank_list)
79 logger.debug(
80 "differentiable_all_gather_concat: input_shape=%s, concat_dim=%d, "
81 "concat_size=%d, rank_list=%s",
82 tuple(x.shape), concat_dim, concat_size, rank_list,
83 )
84 return _utils.differentiable_all_gather_concat(x, group, concat_size, concat_dim, rank_list)
87 @staticmethod
88 def _construct_strided_slice(x, *args):
89 """args: (begin, end, strides)"""
90 dims = len(args) // 3
91 return _utils.construct_strided_slice(x, args[0: dims], args[dims: 2 * dims], args[2 * dims:])
93 @staticmethod
94 def _construct_all_concat_new(x, *args):
95 """args: (concat_dim, concat_size, group)"""
96 rank_list = args[2]
97 concat_dim = args[0]
98 concat_size = args[1]
99 group = _utils.create_group(rank_list)
100 logger.debug(
101 "differentiable_all_gather_concat: input_shape=%s, concat_dim=%d, "
102 "concat_size=%d, rank_list=%s",
103 tuple(x.shape), concat_dim, concat_size, rank_list,
104 )
105 return _utils.differentiable_all_gather_concat(x, group, concat_size, concat_dim, rank_list)
107 def _construct_all_split(self, x, *args):
108 """args: (split_dim, split_size, group)"""
109 rank_list = list(args[2])
110 split_dim = args[0]
111 split_size = args[1]
112 idx = rank_list.index(self.rank_id)
113 return torch.chunk(x, split_size, dim=split_dim)[idx]
115 @staticmethod
116 def _construct_all_to_all(x, *args):
117 """args: (split_dim, concat_dim, permute_size, group)"""
118 split_dim, concat_dim, split_count, rank_list = args
119 group = _utils.create_group(rank_list)
120 logger.debug(
121 "differentiable_all_to_all: input_shape=%s, split_dim=%d, "
122 "concat_dim=%d, split_count=%d, rank_list=%s",
123 tuple(x.shape), split_dim, concat_dim, split_count, rank_list,
124 )
125 original_shape = x.shape
127 dim_size = original_shape[split_dim]
128 if dim_size % split_count != 0:
129 raise ValueError(f"Dimension {split_dim} with size {dim_size} "
130 f"cannot be evenly split into {split_count} parts")
132 split_size = dim_size // split_count
133 final_shape = list(original_shape)
134 if split_dim != concat_dim:
135 final_shape[split_dim] = split_size
136 final_shape[concat_dim] = final_shape[concat_dim] * split_count
137 final_shape = tuple(final_shape)
139 pre_special_handle = all(original_shape[i] == 1 for i in range(split_dim))
140 if pre_special_handle:
141 reshape_shape = (split_count * split_size,) + original_shape[split_dim + 1:]
142 x_reshaped = x.view(reshape_shape)
143 else:
144 reshape_dims = list(original_shape)
145 reshape_dims[split_dim] = split_count
146 reshape_dims.insert(split_dim + 1, split_size)
148 trans_dims = list(range(len(reshape_dims)))
149 trans_dims.remove(split_dim)
150 trans_dims.insert(0, split_dim)
152 x_reshaped = x.reshape(reshape_dims).permute(trans_dims).contiguous()
154 reshape_shape = list(x_reshaped.shape)
155 reshape_shape[0] = reshape_shape[0] * reshape_shape[1]
156 reshape_shape.pop(1)
157 reshape_shape = tuple(reshape_shape)
158 x_reshaped = x_reshaped.reshape(reshape_shape)
159 x_reshaped = x_reshaped.contiguous()
160 output_tensor = _utils.differentiable_all_to_all(
161 input_data=x_reshaped,
162 output_shape=reshape_shape,
163 group=group
164 )
166 post_special_handle = all(final_shape[i] == 1 for i in range(concat_dim))
167 if post_special_handle:
168 return output_tensor.view(final_shape)
170 # When pre_special_handle collapsed leading size-1 dims, the A2A was executed
171 # in a reduced-rank space where the effective concat axis is shifted left by
172 # split_dim positions. Use recon_concat_dim for all post-A2A reshaping so
173 # that split_count is merged into the correct dimension.
174 recon_concat_dim = (concat_dim - split_dim) if pre_special_handle else concat_dim
176 output_reshape = list(output_tensor.shape)
177 output_reshape[0] = split_count
178 output_reshape.insert(1, output_tensor.shape[0] // split_count)
180 out_trans_dims = list(range(len(output_reshape)))
181 first_dim = out_trans_dims.pop(0)
182 if recon_concat_dim >= len(out_trans_dims):
183 out_trans_dims.append(first_dim)
184 else:
185 out_trans_dims.insert(recon_concat_dim, first_dim)
187 final_output = output_tensor.reshape(output_reshape).permute(out_trans_dims).contiguous()
189 final_reshape = list(final_output.shape)
190 if recon_concat_dim < len(final_reshape) - 1:
191 final_reshape[recon_concat_dim] = (
192 final_reshape[recon_concat_dim] * final_reshape[recon_concat_dim + 1]
193 )
194 final_reshape.pop(recon_concat_dim + 1)
196 result = final_output.reshape(final_reshape)
197 if pre_special_handle:
198 result = result.view(final_shape)
199 return result
201 @staticmethod
202 def _apply_eazy_redistribute(src_layout, dst_layout):
203 """_apply_eazy_redistribute"""
204 if (src_layout.mesh_shape != dst_layout.mesh_shape or
205 src_layout.rank_list != dst_layout.rank_list):
206 return False
208 tensor_map_size = len(src_layout.tensor_map)
209 if len(dst_layout.tensor_map) != tensor_map_size:
210 return False
211 return True
213 def _redistribution_without_shape(self, local_x, src_layout, dst_layout, key, rank_list):
214 """_redistribution_without_shape"""
215 inferrer = RedistributionOperatorInfer(
216 dev_mat=src_layout.mesh_shape,
217 in_tensor_map=list(src_layout.tensor_map),
218 out_tensor_map=list(dst_layout.tensor_map)
219 )
220 op_list = inferrer.infer_ops_list(self.rank_id, rank_list)
221 self._transform_cache[key] = op_list
222 for op in op_list:
223 local_x = self._construct_op_operator[op[0]](local_x, *op[1])
224 return local_x
226 @staticmethod
227 def _to_normal_layout(layout: Layout, tensor_dim: int) -> Layout:
228 """Build the normal view consumed by the legacy redistribution path."""
229 normal_layout = Layout.from_device_mesh(layout.mesh)
230 normal_layout.set_placements(layout.normal_placements)
231 normal_layout.placement_to_tensor_map(tensor_dim)
232 return normal_layout
234 def _redistribute_ragged(self, input_x: DTensor, to_layout: Layout) -> DTensor:
235 """Adapt RaggedShard layouts to the supported redistribution primitives."""
236 from_layout = input_x.layout
237 source_info = from_layout.ragged_shard
238 target_info = to_layout.ragged_shard
239 source_is_ragged = isinstance(source_info, RaggedShardInfo)
240 target_is_ragged = isinstance(target_info, RaggedShardInfo)
241 tensor_dim = len(input_x.shape)
243 if source_is_ragged and target_is_ragged:
244 source_normal_layout = self._to_normal_layout(from_layout, tensor_dim)
245 target_normal_layout = self._to_normal_layout(to_layout, tensor_dim)
246 if (
247 source_info.mesh_dim == target_info.mesh_dim
248 and source_info.placement.dims == target_info.placement.dims
249 and source_normal_layout == target_normal_layout
250 ):
251 return self.ragged_to_ragged(input_x, to_layout)
252 return self.ragged_to_ragged_via_replicate(
253 input_x,
254 source_normal_layout,
255 to_layout,
256 )
258 if source_is_ragged:
259 source_normal_layout = self._to_normal_layout(from_layout, tensor_dim)
260 normal = self.ragged_to_normal(input_x, source_normal_layout)
261 if normal.layout == to_layout:
262 return normal
263 return self._redistribution_normal(normal, to_layout)
265 target_normal_layout = self._to_normal_layout(to_layout, tensor_dim)
266 normal = input_x
267 if from_layout != target_normal_layout:
268 normal = self._redistribution_normal(normal, target_normal_layout)
269 return self.normal_to_ragged(normal, to_layout)
271 def redistribution(self, input_x: DTensor, to_layout: Layout) -> DTensor:
272 """Redistribute a tensor to the requested layout.
274 Args:
275 input_x: Distributed input tensor.
276 to_layout: Requested output layout on the same rank list.
277 """
278 x_layout = input_x.layout
279 if (x_layout.has_uneven_shard
280 and not isinstance(x_layout.ragged_shard, RaggedShardInfo)
281 and not isinstance(to_layout.ragged_shard, RaggedShardInfo)):
282 # Validate the export-only path before any collective or layout rewrite.
283 if x_layout.rank_list != to_layout.rank_list:
284 raise ValueError("Uneven redistribution requires the same rank list")
285 return self._gather_uneven_to_replicate(input_x, to_layout)
286 x = input_x
287 if input_x.layout.is_partial():
288 # Solve partial status first
289 if input_x.layout.mesh_shape == to_layout.mesh_shape:
290 x = self.reduce_partial(input_x, to_layout)
291 else:
292 x = self.reduce_partial(input_x, x_layout)
294 from_layout = x.layout
295 if from_layout.rank_list != to_layout.rank_list:
296 raise ValueError(f"The from_layout rank list: {from_layout.rank_list} is not equal to "
297 f"to_layout rank list: {to_layout.rank_list}")
298 if isinstance(from_layout.ragged_shard, RaggedShardInfo) or isinstance(
299 to_layout.ragged_shard, RaggedShardInfo
300 ):
301 return self._redistribute_ragged(x, to_layout)
302 if from_layout.has_uneven_shard or to_layout.has_uneven_shard:
303 return self._gather_uneven_to_replicate(x, to_layout)
304 return self._redistribution_normal(x, to_layout)
306 @staticmethod
307 @torch.no_grad()
308 def _padded_all_gather(local_tensor, sizes, tensor_dim, rank_list):
309 """Gather export buffers; return independent storage without padding."""
310 if len(rank_list) == 1:
311 return local_tensor.detach().clone()
312 max_size = max(sizes)
313 if tensor_dim != 0 and any(size != max_size for size in sizes):
314 raise NotImplementedError("Uneven export communication only supports dim=0")
315 padded = local_tensor.detach().contiguous()
316 if tensor_dim == 0 and local_tensor.shape[0] != max_size:
317 padded_shape = (max_size, *local_tensor.shape[1:])
318 padded = local_tensor.new_zeros(padded_shape)
319 padded[:local_tensor.shape[0]].copy_(local_tensor)
320 if padded.numel() == 0:
321 output_shape = list(padded.shape)
322 output_shape[tensor_dim] = sum(sizes)
323 return padded.new_empty(output_shape)
324 group = _utils.create_group(rank_list)
325 group_ranks = dist.get_process_group_ranks(group)
326 gathered = [torch.empty_like(padded) for _ in group_ranks]
327 dist.all_gather(gathered, padded, group=group)
328 group_indices = {rank: index for index, rank in enumerate(group_ranks)}
329 # cat copies every trimmed shard into owned storage. Returning a view
330 # (or contiguous(), which may be a no-op) would retain padded buffers.
331 return torch.cat([
332 gathered[group_indices[rank]].narrow(tensor_dim, 0, size)
333 for rank, size in zip(rank_list, sizes)
334 ], dim=tensor_dim)
336 @staticmethod
337 def _validate_uneven_local_shape(input_x: DTensor) -> None:
338 """Reject padded storage masquerading as a true local shard."""
339 source = input_x.layout
340 area = infer_slice_area_by_rank(
341 source.mesh_shape, source.tensor_map, source.rank_list.index(source.mesh.rank),
342 tuple(input_x.shape), source.uneven_shard_mesh_dims,
343 )
344 expected_shape = tuple(end - start for start, end in area)
345 local_shape = tuple(input_x.to_local().shape)
346 if local_shape != expected_shape:
347 raise ValueError(f"Uneven shard local shape {local_shape} does not match logical slice {expected_shape}")
349 @staticmethod
350 def _gather_uneven_to_replicate(input_x: DTensor, to_layout: Layout) -> DTensor:
351 """Reconstruct detached FSDP weights, including empty dim-zero shards.
353 Undo each tensor dimension's sharding from innermost to outermost.
354 This respects both ceil-chunk FSDP shards and balanced TP shards,
355 including the effective ordering encoded for StridedShard layouts.
356 This is an export-only path, not a differentiable redistribution.
357 Uneven nonzero axes and Partial placements remain unsupported.
358 """
359 source = input_x.layout
360 if (source.mesh_shape != to_layout.mesh_shape
361 or not all(isinstance(placement, Replicate) for placement in to_layout.placements)):
362 raise NotImplementedError(
363 "Uneven chunk-sharded redistribution only supports a fully replicated target on the same mesh."
364 )
365 full_shape = tuple(input_x.shape)
366 rank = source.mesh.rank
367 tensor_map = list(source.tensor_map)
368 uneven_dims = source.uneven_shard_mesh_dims
369 local_tensor = input_x.to_local()
370 if source.is_partial():
371 raise NotImplementedError("Uneven FSDP weight export does not support Partial placements")
372 if any(source.placements[axis].dim != 0 for axis in uneven_dims):
373 raise NotImplementedError("Uneven FSDP weight export only supports dim=0")
374 if torch.is_grad_enabled() and local_tensor.requires_grad:
375 raise RuntimeError("Uneven FSDP weight export is not differentiable; use torch.no_grad() or detach()")
376 TensorRedistribution._validate_uneven_local_shape(input_x)
377 for tensor_dim, mapping in enumerate(source.tensor_map):
378 axes = list(mapping if isinstance(mapping, tuple) else (mapping,))
379 axes = [axis for axis in axes if axis != -1]
380 while axes:
381 mesh_dim = len(source.mesh_shape) - 1 - axes[-1]
382 ranks = source.mesh.get_devices_for_axis(mesh_dim, rank)
383 areas = [
384 infer_slice_area_by_rank(
385 source.mesh_shape, tensor_map, source.rank_list.index(peer), full_shape, uneven_dims,
386 ) for peer in ranks
387 ]
388 sizes = [area[tensor_dim][1] - area[tensor_dim][0] for area in areas]
389 local_tensor = TensorRedistribution._padded_all_gather(local_tensor, sizes, tensor_dim, ranks)
390 axes.pop()
391 tensor_map[tensor_dim] = tuple(axes) if axes else -1
392 return DTensor.from_local_with_layout(local_tensor, to_layout, shape=full_shape)
394 def _redistribution_normal(self, input_x: DTensor, to_layout: Layout) -> DTensor:
395 """Run the existing redistribution path for normal layouts."""
396 from_layout = input_x.layout
397 x = input_x
398 if not self.is_init:
399 self.rank_id = dist.get_rank()
400 self.is_init = True
401 key = from_layout.compact_str + to_layout.compact_str + str(self.rank_id)
402 if key in self._transform_cache:
403 x = x.to_local()
404 transform_operator_list = self._transform_cache[key]
405 for transform_operator in transform_operator_list:
406 x = self._construct_op_operator[transform_operator[0]](x, *transform_operator[1])
407 return DTensor.from_local(x, to_layout.mesh, to_layout.alias_placements)
409 full_shape = x.shape
410 key_and_shape = key + str(full_shape)
411 x = x.to_local()
412 if key_and_shape in self._transform_cache:
413 transform_operator_list = self._transform_cache[key_and_shape]
414 for transform_operator in transform_operator_list:
415 x = self._construct_op_operator[transform_operator[0]](x, *transform_operator[1])
416 return DTensor.from_local(x, to_layout.mesh, to_layout.alias_placements)
418 rank_list = from_layout.rank_list
419 if self._apply_eazy_redistribute(from_layout, to_layout):
420 if from_layout.is_partial():
421 from_layout.reset_partial()
422 x = self._redistribution_without_shape(x, from_layout, to_layout, key, rank_list)
423 else:
424 transform_operator_list = self._infer_transform_operator_list(from_layout, to_layout,
425 full_shape, key_and_shape, rank_list)
426 for transform_operator in transform_operator_list:
427 x = self._construct_op_operator[transform_operator[0]](x, *transform_operator[1])
428 return DTensor.from_local(x, to_layout.mesh, to_layout.alias_placements)
430 @staticmethod
431 def ragged_to_normal(input_x: DTensor, to_layout: Layout) -> DTensor:
432 """Gather one flat RaggedShard into its Replicate normal view."""
433 from_layout = input_x.layout
434 info = from_layout.ragged_shard
435 if not isinstance(info, RaggedShardInfo):
436 raise ValueError("ragged_to_normal requires a RaggedShard source layout")
437 global_shape = tuple(input_x.shape)
438 output_splits = _compute_ragged_splits(global_shape, from_layout)
439 local_tensor = input_x.to_local()
440 if len(output_splits) == 1:
441 gathered = local_tensor
442 else:
443 group = from_layout.mesh.get_group(info.mesh_dim)
444 gathered = _utils.differentiable_variable_all_gather(
445 local_tensor,
446 output_splits,
447 group,
448 )
449 return DTensor.from_local_with_layout(
450 gathered.reshape(global_shape),
451 to_layout,
452 shape=global_shape,
453 )
455 @staticmethod
456 def normal_to_ragged(input_x: DTensor, to_layout: Layout) -> DTensor:
457 """Slice a normal-view tensor into the target flat RaggedShard."""
458 global_shape = tuple(input_x.shape)
459 local_slice = _compute_ragged_slice(global_shape, to_layout)
460 flat_tensor = input_x.to_local().reshape((-1,))
461 local_tensor = flat_tensor[
462 local_slice.flat_start:local_slice.flat_end
463 ].clone()
464 return DTensor.from_local_with_layout(
465 local_tensor,
466 to_layout,
467 shape=global_shape,
468 )
470 @staticmethod
471 def ragged_to_ragged_via_replicate(
472 input_x: DTensor,
473 replicate_layout: Layout,
474 to_layout: Layout,
475 ) -> DTensor:
476 """Redistribute different RaggedShard dims through a Replicate layout."""
477 from_layout = input_x.layout
478 if (
479 from_layout.mesh_shape != to_layout.mesh_shape
480 or from_layout.rank_list != to_layout.rank_list
481 ):
482 raise ValueError("ragged_to_ragged only supports changes on the same device mesh")
483 replicated = TensorRedistribution.ragged_to_normal(
484 input_x,
485 replicate_layout,
486 )
487 return TensorRedistribution.normal_to_ragged(replicated, to_layout)
489 @staticmethod
490 def ragged_to_ragged(input_x: DTensor, to_layout: Layout) -> DTensor:
491 """Redistribute a local-units-only RaggedShard change with variable all-to-all."""
492 from_layout = input_x.layout
493 source_info = from_layout.ragged_shard
494 if (
495 from_layout.mesh_shape != to_layout.mesh_shape
496 or from_layout.rank_list != to_layout.rank_list
497 ):
498 raise ValueError("ragged_to_ragged only supports local_units changes on the same device mesh")
500 global_shape = tuple(input_x.shape)
501 input_splits, output_splits = _compute_ragged_all_to_all_splits(
502 global_shape,
503 from_layout,
504 to_layout,
505 )
506 flat_input = input_x.to_local().reshape((-1,))
508 if flat_input.shape[0] != sum(input_splits):
509 raise ValueError(
510 "RaggedShard source storage does not match all-to-all splits, "
511 f"got local_numel={flat_input.shape[0]}, input_splits={input_splits!r}"
512 )
513 if len(input_splits) == 1:
514 flat_output = flat_input
515 else:
516 group = from_layout.mesh.get_group(source_info.mesh_dim)
517 flat_output = _utils.differentiable_all_to_all_single(
518 flat_input,
519 input_splits,
520 output_splits,
521 group,
522 )
523 if flat_output.shape[0] != sum(output_splits):
524 raise ValueError(
525 "RaggedShard target storage does not match all-to-all splits, "
526 f"got local_numel={flat_output.shape[0]}, output_splits={output_splits!r}"
527 )
528 return DTensor.from_local_with_layout(
529 flat_output,
530 to_layout,
531 shape=global_shape,
532 )
534 def _infer_transform_operator_list(self, from_layout, to_layout, from_full_shape, key, rank_list):
535 """infer transform operator list"""
536 from_layout_tuple, to_layout_tuple = \
537 _construct_layout_tuple_for_transform_operator_list(from_layout, to_layout, from_full_shape)
538 self._transform_cache[key] = \
539 _utils.get_tensor_transform().transform_tensor_sharding(from_layout_tuple, to_layout_tuple,
540 rank_list, False, self.rank_id)
541 return self._transform_cache[key]
543 @staticmethod
544 def _allreduce_along_dev_dim(x, op, layout, dev_dim):
545 """Do allreduce at specified axis along dev_dim."""
546 logger.debug(
547 "differentiable_all_reduce: input_shape=%s, op=%s, dev_dim=%s",
548 tuple(x.shape), op, dev_dim,
549 )
550 group = layout.get_comm_group_by_axis(dev_dim)
551 zero_dim = x.dim() == 0
552 if zero_dim:
553 x = x.unsqueeze(0)
554 if op == 'avg':
555 dev_num = layout.mesh_shape[layout.alias_name.index(dev_dim)]
556 x = _utils.differentiable_all_reduce(x, 'sum', group)
557 x = x / dev_num
558 elif op == 'all':
559 x_int32 = _utils.tensor_type_cast(x.bool(), 'int32') # True→1, False→0
560 x = _utils.differentiable_all_reduce(x_int32, 'all', group)
561 x = x.bool()
562 else:
563 x = _utils.differentiable_all_reduce(x, op, group)
564 if zero_dim:
565 x = x.squeeze(0)
566 return x
568 @staticmethod
569 def _reduce_scatter_along_dev_dim_with_axis(x, axis, op, layout, dev_dim):
570 """Do reduce_scatter at specified axis along dev_dim."""
571 dev_num = layout.mesh_shape[layout.alias_name.index(dev_dim)]
572 logger.debug(
573 "differentiable_reduce_scatter: input_shape=%s, axis=%d, "
574 "op=%s, dev_dim=%s, dev_num=%d",
575 tuple(x.shape), axis, op, dev_dim, dev_num,
576 )
577 group = layout.get_comm_group_by_axis(dev_dim)
578 output_tensor = _utils.differentiable_reduce_scatter(x, dev_num, axis, op, group)
579 return output_tensor
581 def reduce_partial(self, input_x, to_layout):
582 """Reduce partial status."""
583 from_layout = input_x.layout
584 x = input_x
585 if from_layout is None or not from_layout.is_partial():
586 return x
588 x = x.to_local()
589 if from_layout.mesh_shape != to_layout.mesh_shape:
590 raise ValueError(f"For reduce partial, mesh_shape between from_layout and to_layout must be the same, "
591 f"but got {from_layout.mesh_shape} and {to_layout.mesh_shape}")
592 if to_layout.is_partial():
593 raise ValueError(f"For reduce partial, to_layout must be non-partial status, but got to_layout.partial: "
594 f"{to_layout.partial}")
596 dev_map_order = {}
597 for dev_axis in to_layout.alias_tensor_map:
598 if isinstance(dev_axis, tuple):
599 for i, sub_dev_axis in enumerate(dev_axis):
600 dev_map_order[sub_dev_axis] = i
601 else:
602 dev_map_order[dev_axis] = 0
604 pending_reduce_op_list = [] # List[Tuple[comm_op, op, dev_dim, reduce_dim]]
605 for dev_axis_index, op in enumerate(from_layout.partial):
606 if op is None:
607 continue
608 dev_axis = from_layout.alias_name[dev_axis_index]
609 apply_shard_dim = to_layout.get_dev_axis_apply_shard_axis(dev_axis)
610 comm_op = "ReduceScatter" if apply_shard_dim is not None else "AllReduce"
611 pending_reduce_op_list.append((comm_op, op, dev_axis, apply_shard_dim))
613 # sort reduce op
614 # 1. ReduceScatter is executed before AllReduce
615 # 2. If multiple split, the dev axis split outer will be execute first.
616 # e.g. ("cp", "tp"), will execute reduce_scatter along "cp" before "tp"
617 # 3. Lower dev_id execute before higher dev_id
618 def _reduce_pair_sort_key(reduce_pair):
619 return (reduce_pair[0] != "ReduceScatter",
620 dev_map_order.get(reduce_pair[2], 0),
621 to_layout.mesh.axis_id(reduce_pair[2]))
623 sorted_pending_reduce_op_list = sorted(pending_reduce_op_list, key=_reduce_pair_sort_key)
625 output_alias_tensor_map = list(from_layout.alias_tensor_map)
626 for reduce_op_pair in sorted_pending_reduce_op_list:
627 comm_op = reduce_op_pair[0]
628 op = reduce_op_pair[1]
629 dev_axis = reduce_op_pair[2]
630 if comm_op == "AllReduce":
631 x = TensorRedistribution._allreduce_along_dev_dim(x, op, from_layout, dev_axis)
632 elif comm_op == "ReduceScatter":
633 reduce_axis = reduce_op_pair[3]
634 x = self._reduce_scatter_along_dev_dim_with_axis(x, reduce_axis, op, from_layout, dev_axis)
635 if output_alias_tensor_map[reduce_axis] == "None":
636 output_alias_tensor_map[reduce_axis] = dev_axis
637 elif isinstance(output_alias_tensor_map[reduce_axis], tuple):
638 output_alias_tensor_map[reduce_axis] += (dev_axis,)
639 else:
640 output_alias_tensor_map[reduce_axis] = (output_alias_tensor_map[reduce_axis], dev_axis)
642 output_layout = from_layout(*output_alias_tensor_map)
643 output_layout.reset_partial()
644 return DTensor.from_local(x, output_layout.mesh, output_layout.alias_placements)
647_tensor_redistribution = TensorRedistribution()