Coverage for / home / jenkins / .local / lib / python3.10 / site-packages / hyper_parallel / core / dtensor / layout.py: 86%
478 statements
« prev ^ index » next coverage.py v7.13.1, created at 2026-08-21 04:29 +0800
1# Copyright 2025 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"""layout"""
17import copy
18import functools
19from typing import NamedTuple, Optional, Sequence
21import numpy as np
24from hyper_parallel.core.dtensor.placement_types import (
25 Partial,
26 Placement,
27 RaggedShard,
28 Replicate,
29 Shard,
30 StridedShard,
31)
32from hyper_parallel.core.dtensor.device_mesh import DeviceMesh, _create_device_mesh
33from hyper_parallel.platform import get_platform
35platform = get_platform()
38class RaggedShardInfo(NamedTuple):
39 """RaggedShard placement and its corresponding mesh dimension."""
41 mesh_dim: int
42 placement: RaggedShard
45def _extract_ragged_shard(placements: Sequence[Placement]) -> Optional[RaggedShardInfo]:
46 """Extract the single RaggedShard placement from a placement sequence."""
47 ragged_shard = None
48 for mesh_dim, placement in enumerate(placements):
49 if not placement.is_ragged_shard():
50 continue
51 if ragged_shard is not None:
52 raise ValueError(
53 "Layout supports at most one RaggedShard placement, "
54 f"but got placements={tuple(placements)!r}"
55 )
56 ragged_shard = RaggedShardInfo(mesh_dim, placement)
57 return ragged_shard
60def _replace_ragged_with_replicate(placements: Sequence[Placement]) -> tuple[Placement, ...]:
61 """Return placements with RaggedShard represented as Replicate."""
62 return tuple(
63 Replicate() if placement.is_ragged_shard() else placement
64 for placement in placements
65 )
68def _infer_slice_area_by_rank(mesh_shape, tensor_map, rank_id: int, full_shape: tuple): # -> tuple[tuple[int]]:
69 """Return the range of each axis from full tensor for slice in current rank."""
71 def _get_dev_num_alone_dim(mesh_shape, dim):
72 """_get_dev_num_alone_dim."""
73 return mesh_shape[-dim - 1] if dim != -1 else 1
75 def _rank_id_to_dev_id_list(mesh_shape, rank_id):
76 """Infer dev id list by rank_id and mesh_shape"""
77 dims = len(mesh_shape)
78 dev_id_list = [0] * dims
79 for i in range(dims - 1, -1, -1):
80 dev_id_list[i] = rank_id % mesh_shape[i]
81 rank_id = rank_id // mesh_shape[i]
82 return dev_id_list
84 dev_id_list = _rank_id_to_dev_id_list(mesh_shape, rank_id)
86 dims = len(full_shape)
87 area = []
88 for axis in range(dims):
89 mapping = tensor_map[axis]
90 if isinstance(mapping, int):
91 mapping = (mapping,)
92 split_num = 1
93 for dim in mapping:
94 split_num *= _get_dev_num_alone_dim(mesh_shape, dim)
96 slice_id = 0
97 coef = 1
98 for dim in reversed(mapping):
99 if dim == -1:
100 continue
101 slice_id += dev_id_list[-dim - 1] * coef
102 coef *= _get_dev_num_alone_dim(mesh_shape, dim)
103 slice_size = full_shape[axis] // split_num
104 start = slice_id * slice_size
105 end = start + slice_size
106 area.append((start, end))
107 return area
110def _get_slice_tensor_by_layout(global_tensor, layout):
111 """Transfer global tensor to local tensor by layout"""
112 inner_rank_id = layout.rank_list.index(layout.mesh.rank)
113 slice_area = _infer_slice_area_by_rank(layout.mesh_shape, layout.tensor_map, inner_rank_id, global_tensor.shape)
115 def get_slice_data(full_data, offset):
116 area = ()
117 for begin, end in offset:
118 area += (slice(begin, end),)
119 return full_data[area].clone()
121 local_tensor = get_slice_data(global_tensor, slice_area)
122 return local_tensor
125def _infer_slice_shape_by_layout(global_shape, layout):
126 """Infer slice shape from global_shape and layout"""
127 slice_shape = list(global_shape)
128 alias_tensor_map = layout.alias_tensor_map
129 for i in range(len(global_shape)):
130 axis_name = alias_tensor_map[i]
131 if isinstance(axis_name, str):
132 axis_name = (axis_name,)
133 for sub_axis_name in axis_name:
134 if sub_axis_name != "None":
135 slice_shape[i] = slice_shape[i] // layout.mesh.get_device_num_along_axis(sub_axis_name)
136 return slice_shape
139class Layout:
140 """
141 Topological abstraction describing cluster devices for tensor slice placement on the cluster.
143 Note:
144 - It is valid only in semi auto parallel or auto parallel mode.
145 - The multiplication result of the `mesh_shape` must be equal to the device count in a pipeline stage.
146 - When the layout function is invoked to constructs a sharding strategy, each alias name is only allowed to be
147 used once to shard a tensor.
149 Args:
150 mesh_shape (tuple): Describe the shape of devices arrangement, its element type is int.
151 alias_name (tuple): The alias name for each axis of mesh_shape, its length shoits element type is string.
152 When using "interleaved_parallel" as an alias name, the tensor would be split into multiple
153 copies on the corresponding partition dimension on a single card.
154 rank_list (tuple, optional): Data is allocated to the device according to rank_list. Default: ``None``.
156 Raises:
157 TypeError: `mesh_shape` is not a tuple type.
158 TypeError: `alias_name` is not a tuple type.
159 TypeError: 'rank_list' is not a list type.
160 ValueError: `mesh_shape` length is not equal to `alias_name` length.
161 TypeError: The element of `mesh_shape` is not int type.
162 TypeError: The element of `alias_name` is not a str type.
163 TypeError: The element of `rank_list` is not int type.
164 ValueError: The element of `alias_name` is an empty str.
165 ValueError: The element of `alias_name` is "None".
166 ValueError: `alias_name` contains repeated element.
168 Supported Platforms:
169 ``Ascend``
171 Examples:
172 >>> from mindspore.parallel import Layout
173 >>> layout = Layout((2, 2, 2), ("dp", "sp", "mp"))
174 >>> layout0 = layout("dp", "mp")
175 >>> print(layout0.to_dict())
176 {"mesh_shape": (2, 2, 2), "tensor_map": (2, 0), "interleaved_parallel": False,
177 'alias_name': {'dp', 'sp', 'mp'}, "rank_list": [0, 1, 2, 3, 4, 5, 6, 7]}
178 >>> layout = Layout((2, 2, 2), ("dp", "sp", "interleaved_parallel"))
179 >>> layout1 = layout(("dp", "interleaved_parallel"), "sp")
180 """
182 def __init__(self, mesh_shape, alias_name, rank_list=None, init_backend=True):
183 self._alias_name = alias_name
184 self._tensor_map = None
185 if not rank_list:
186 self._rank_list = tuple(range(np.prod(np.array(mesh_shape))))
187 else:
188 self._rank_list = tuple(rank_list)
189 self._partial = [None] * len(mesh_shape) # partial status for each dev dim
190 self._support_partial_op = ['sum', 'max', 'min', 'avg', 'prod', 'all', None]
191 self._alias_tensor_map = None
192 self._mesh = _create_device_mesh("npu", mesh_shape, mesh_dim_names=alias_name, rank_list=self._rank_list,
193 init_backend=init_backend)
194 self._compact_str = self._to_compact_string()
195 self._placements = None
196 self.partial_ops = {} # Initialized in _build_dim_map_from_placements()
197 self._ragged_shard = None
199 @classmethod
200 def from_device_mesh(cls, device_mesh: DeviceMesh) -> 'Layout':
201 """
202 Create a Layout from an existing DeviceMesh.
204 Args:
205 device_mesh (DeviceMesh): The device mesh to create layout from.
207 Returns:
208 Layout: A new Layout instance initialized with the properties of the provided device mesh.
210 Examples:
211 >>> from hyper_parallel.core.dtensor.layout import Layout, DeviceMesh
212 >>> device_mesh = DeviceMesh("npu", (2, 2), mesh_dim_names=("dp", "mp"))
213 >>> layout = Layout.from_device_mesh(device_mesh)
214 """
215 obj = cls.__new__(cls)
216 obj._mesh = device_mesh
217 obj._alias_name = device_mesh.mesh_dim_names
218 obj._rank_list = device_mesh.rank_list
219 obj._tensor_map = None
220 obj._partial = [None] * len(device_mesh.mesh_shape)
221 obj._support_partial_op = ['sum', 'max', 'min', 'avg', 'prod', 'all', None]
222 obj._alias_tensor_map = None
223 obj._placements = None
224 obj._ragged_shard = None
225 obj._compact_str = obj._to_compact_string()
226 return obj
228 def __call__(self, *alias_tensor_map):
229 obj = copy.deepcopy(self)
231 # Clear the inherited partial status.
232 # When creating a new layout mapping configuration via __call__,
233 # it should not inherit the dynamic execution state (Partial) of the original layout.
234 # If the user intends to create a Partial placement, it will be parsed from alias_tensor_map.
235 obj._partial = [None] * len(obj.mesh_shape)
237 if len(alias_tensor_map) == 1 and isinstance(alias_tensor_map[0], (list, tuple)):
238 if len(alias_tensor_map[0]) > 0 and isinstance(alias_tensor_map[0][0], Placement):
239 return self._process_placement_layout(obj, alias_tensor_map[0])
241 if len(alias_tensor_map) > 0 and isinstance(alias_tensor_map[0], Placement):
242 return self._process_placement_layout(obj, alias_tensor_map)
244 return self._process_alias_layout(obj, alias_tensor_map)
246 def __deepcopy__(self, memo):
247 """Deep copy layout without rebuilding the underlying device mesh."""
248 cls = self.__class__
249 result = cls.__new__(cls)
250 memo[id(self)] = result
251 for k, v in self.__dict__.items():
252 setattr(result, k, copy.deepcopy(v, memo))
253 return result
255 @staticmethod
256 def _process_placement_layout(obj, placements):
257 """Process layout defined by Placement types."""
258 obj.set_placements(placements)
259 return copy.deepcopy(obj)
261 @staticmethod
262 def _process_alias_layout(obj, alias_tensor_map):
263 """Process layout defined by alias strings."""
264 obj.set_alias_tensor_map(alias_tensor_map)
265 tensor_map = ()
266 writed_map = ()
267 for ele in alias_tensor_map:
268 if isinstance(ele, tuple):
269 ele_map = ()
270 for item in ele:
271 if item == "None":
272 ele_map += (-1,)
273 continue
274 if item not in obj.alias_name:
275 raise ValueError(f'The axis {item} is not found in {obj.alias_name}')
276 if item in writed_map:
277 raise ValueError(f'The axis {item} has been set more than one in {obj.alias_name}')
278 ele_map += (len(obj.alias_name) - 1 - obj.alias_name.index(item),)
279 writed_map += (item,)
280 tensor_map += (ele_map,)
281 continue
282 if ele == "None":
283 tensor_map += (-1,)
284 continue
285 if ele not in obj.alias_name:
286 raise ValueError(f'The axis {ele} is not found in {obj.alias_name}')
287 if ele in writed_map:
288 raise ValueError(f'The axis {ele} has been set more than one in {obj.alias_name}')
289 tensor_map += (len(obj.alias_name) - 1 - obj.alias_name.index(ele),)
290 writed_map += (ele,)
291 obj.set_tensor_map(tensor_map)
292 obj.tensor_map_to_placement()
293 obj.update_compact_str()
294 return copy.deepcopy(obj)
296 def to_dict(self):
297 """
298 Transform layout to a dictionary.
299 """
300 if self._mesh.mesh_shape is None:
301 raise ValueError("The device_shape of layout is None")
302 if self._tensor_map is None:
303 raise ValueError("The tensor_map of layout is None")
304 interleaved_parallel = "interleaved_parallel" in self._mesh.mesh_dim_names
305 return {"mesh_shape": self._mesh.mesh_shape, "tensor_map": self._tensor_map,
306 "interleaved_parallel": interleaved_parallel, "alias_name": self._mesh.mesh_dim_names,
307 "rank_list": self._rank_list}
309 def placement_to_tensor_map(self, dim):
310 """
311 Transform placement to tensor map.
313 This method converts the `placements` configuration (consisting of Shard, StridedShard,
314 Replicate, Partial)
315 into a `tensor_map` representation used for distributed tensor operations.
317 Args:
318 dim (int): The dimension of the tensor. Must be a positive integer.
320 Returns:
321 tuple: A tuple representing the tensor map, where each element corresponds to a tensor dimension.
322 A value of -1 indicates the dimension is not sharded, an integer indicates the mesh
323 dimension index along which the tensor dimension is sharded, and a tuple indicates
324 that the same tensor dimension is sharded multiple times in order.
326 Raises:
327 ValueError: If `dim` is negative.
328 ValueError: If a shard dimension in `placements` is out of bounds for the given tensor dimension.
329 """
330 if dim < 0:
331 raise ValueError(f"Tensor dimension must be positive, but got {dim}")
332 if dim == 0:
333 return self._handle_zero_dim_placement()
335 dim_map = self._build_dim_map_from_placements(dim)
336 tensor_map = self._convert_dim_map_to_tensor_map(dim_map)
337 self.set_tensor_map(tuple(tensor_map))
338 self._alias_tensor_map = self._build_readable_tensor_map()
339 self.update_compact_str()
340 return tensor_map
342 def _handle_zero_dim_placement(self):
343 """Handle the special case of zero-dimensional tensor."""
344 self.set_tensor_map(())
345 self._alias_tensor_map = ()
346 for mesh_idx, placement in enumerate(self.normal_placements):
347 if isinstance(placement, Partial):
348 self._partial[mesh_idx] = self._extract_reduce_op(placement)
349 return []
351 def _build_dim_map_from_placements(self, dim):
352 """Build dimension map from placements."""
353 dim_map = [-1] * dim
354 self.partial_ops = {}
355 for mesh_idx, placement in enumerate(self.normal_placements):
356 if isinstance(placement, Shard):
357 shard_dim = placement.dim
358 if shard_dim < -dim or shard_dim >= dim:
359 raise ValueError(f"Shard dimension {shard_dim} is out of bounds for tensor of dimension {dim}")
360 if shard_dim < 0:
361 shard_dim += dim
362 if dim_map[shard_dim] == -1:
363 dim_map[shard_dim] = [mesh_idx]
364 else:
365 dim_map[shard_dim].append(mesh_idx)
366 elif isinstance(placement, Partial):
367 self._partial[mesh_idx] = self._extract_reduce_op(placement)
368 self._validate_strided_shard_split_factor(dim_map)
369 self._reorder_dim_map_for_strided_shard(dim_map)
370 return dim_map
372 @staticmethod
373 def _placement_split_factor(placement):
374 """Return the effective split factor carried by a placement."""
375 return placement.split_factor if isinstance(placement, StridedShard) else 1
377 @staticmethod
378 def _build_order_positions(shard_order):
379 """Build a mesh axis to order position mapping."""
380 return {mesh_idx: order_idx for order_idx, mesh_idx in enumerate(shard_order)}
382 def _compute_expected_split_factors(self, shard_axes, shard_order):
383 """Infer the split_factor each mesh axis should carry for the given sharding order."""
384 order_positions = self._build_order_positions(shard_order)
385 expected_split_factors = {}
386 for mesh_idx in shard_axes:
387 split_factor = 1
388 for right_mesh_idx in shard_axes:
389 if right_mesh_idx <= mesh_idx:
390 continue
391 if order_positions[right_mesh_idx] < order_positions[mesh_idx]:
392 split_factor *= self.mesh_shape[right_mesh_idx]
393 expected_split_factors[mesh_idx] = split_factor
394 return expected_split_factors
396 def _get_effective_shard_axes(self, shard_axes):
397 """Return shard axes ordered by their effective sharding order."""
398 return sorted(
399 shard_axes,
400 key=lambda mesh_idx: self._placement_split_factor(self.placements[mesh_idx]),
401 )
403 def _reorder_dim_map_for_strided_shard(self, dim_map):
404 """Reorder dim_map entries to reflect the effective sharding order."""
405 for i, shard_axes in enumerate(dim_map):
406 if shard_axes == -1 or len(shard_axes) <= 1:
407 continue
408 dim_map[i] = self._get_effective_shard_axes(shard_axes)
410 def _validate_strided_shard_split_factor(self, dim_map):
411 """Validate that split factors match the effective sharding order."""
412 for shard_axes in dim_map:
413 if shard_axes == -1:
414 continue
415 shard_order = self._get_effective_shard_axes(shard_axes)
416 expected_split_factors = self._compute_expected_split_factors(
417 shard_axes, shard_order
418 )
419 for mesh_idx in shard_axes:
420 placement = self.placements[mesh_idx]
421 actual_split_factor = self._placement_split_factor(placement)
422 expected_split_factor = expected_split_factors[mesh_idx]
423 if actual_split_factor != expected_split_factor:
424 raise ValueError(
425 f"StridedShard split_factor mismatch on mesh axis {mesh_idx}: "
426 f"expected {expected_split_factor}, got {actual_split_factor}."
427 )
429 @staticmethod
430 def _extract_reduce_op(placement):
431 """Extract reduce operation name from Partial placement."""
432 op_name = getattr(placement, "reduce_op", "sum")
433 if isinstance(op_name, str):
434 op_name = op_name.lower()
435 return op_name
437 def _convert_dim_map_to_tensor_map(self, dim_map):
438 """Convert dimension map to tensor map format."""
439 device_dim_count = len(self.mesh_shape)
440 tensor_map = []
441 for mesh_idx in dim_map:
442 if mesh_idx == -1:
443 tensor_map.append(-1)
444 continue
445 mapped_axes = tuple(device_dim_count - 1 - axis for axis in mesh_idx)
446 tensor_map.append(mapped_axes[0] if len(mapped_axes) == 1 else mapped_axes)
447 return tensor_map
449 def _build_readable_tensor_map(self):
450 """Build human-readable alias tensor map from tensor_map."""
451 mesh_dim_names = self._mesh.mesh_dim_names
452 has_names = mesh_dim_names is not None
454 def _map_dim(dim):
455 """convert dimension index to dimension name."""
456 if dim == -1:
457 return "None"
458 if not has_names:
459 return f"dim_{dim}"
460 return mesh_dim_names[len(mesh_dim_names) - 1 - dim]
462 readable_map = []
463 for item in self._tensor_map:
464 if isinstance(item, tuple):
465 mapped_tuple = tuple(_map_dim(dim) for dim in item)
466 readable_map.append(mapped_tuple)
467 else:
468 readable_map.append(_map_dim(item))
469 return tuple(readable_map)
471 def tensor_map_to_placement(self):
472 """
473 Transform tensor map to placement.
475 This method converts the existing `tensor_map` and `partial` status into a list of `Placement` objects
476 (Shard, StridedShard, Replicate, Partial). This is the inverse operation of
477 `placement_to_tensor_map`.
479 Returns:
480 list[Placement]: A list of Placement objects describing the distribution strategy for each
481 dimension of the device mesh.
483 Raises:
484 ValueError: If `tensor_map` is not configured (None).
485 """
486 if self._tensor_map is None:
487 raise ValueError("The tensor_map is None, cannot transform to placements.")
488 mesh_ndim = len(self.mesh_shape)
489 placements = [Replicate()] * mesh_ndim
490 for tensor_dim, mapping in enumerate(self._tensor_map):
491 mapping_list = mapping if isinstance(mapping, tuple) else (mapping,)
492 valid_mapping = [map_val for map_val in mapping_list if map_val != -1]
493 mesh_indices = [mesh_ndim - 1 - map_val for map_val in valid_mapping]
494 shard_axes = sorted(mesh_indices)
495 expected_split_factors = self._compute_expected_split_factors(
496 shard_axes, mesh_indices
497 )
498 for mesh_idx in shard_axes:
499 split_factor = expected_split_factors[mesh_idx]
500 placement = (
501 StridedShard(dim=tensor_dim, split_factor=split_factor)
502 if split_factor > 1
503 else Shard(dim=tensor_dim)
504 )
505 placements[mesh_idx] = placement
506 for mesh_idx, op in enumerate(self.partial):
507 if op is not None:
508 placements[mesh_idx] = Partial(reduce_op=op)
509 if self._ragged_shard is not None:
510 placements[self._ragged_shard.mesh_dim] = self._ragged_shard.placement
511 self.set_placements(placements)
512 self._alias_tensor_map = self._build_readable_tensor_map()
513 self.update_compact_str()
514 return placements
516 def __setstate__(self, state):
517 self.__dict__.update(state)
518 self.update_mesh(init_backend=False)
520 @property
521 def mesh(self):
522 """
523 Get the device mesh associated with this layout.
525 Returns:
526 DeviceMesh: The device mesh describing the device topology.
527 """
528 return self._mesh
530 def update_mesh(self, init_backend: bool = True):
531 """Recreate the internal DeviceMesh from current layout properties.
533 Args:
534 init_backend (bool): Whether to initialize communication backend
535 (process groups). Set to ``False`` during deserialization to
536 avoid creating process groups with a stale rank_list from the
537 sender side. Default ``True``.
538 """
539 self._mesh = _create_device_mesh("npu", self.mesh_shape, mesh_dim_names=self.alias_name,
540 rank_list=self.rank_list, init_backend=init_backend)
542 @property
543 def rank_list(self):
544 """
545 Get the list of ranks participating in this layout.
547 Returns:
548 tuple[int]: The rank list.
549 """
550 return self._rank_list
552 @rank_list.setter
553 def rank_list(self, val):
554 self._rank_list = val
556 @property
557 def mesh_shape(self):
558 """mesh shape"""
559 return self._mesh.mesh_shape
561 @property
562 def alias_name(self):
563 """alias name"""
564 return self._mesh.mesh_dim_names
566 @property
567 def alias_tensor_map(self):
568 """Return the human-readable alias tensor map for this layout."""
569 return self._alias_tensor_map
571 @property
572 def alias_placements(self):
573 """Return alias_tensor_map when it contains multi-axis tuples, otherwise placements.
575 alias_tensor_map preserves multi-axis ordering information
576 (e.g., (("dp", "tp"), "None") vs (("tp", "dp"), "None"))
577 that Placement objects cannot represent, since both map to
578 [Shard(0), Shard(0)].
580 For single-axis layouts, Placement objects are preferred because they
581 also carry Partial status which alias_tensor_map cannot encode.
583 Use this property when constructing DTensors from an existing Layout
584 to avoid the lossy Placement round-trip for multi-axis cases.
585 """
586 if self._ragged_shard is not None:
587 return self._placements
588 if self._alias_tensor_map is not None and any(
589 isinstance(item, tuple) for item in self._alias_tensor_map
590 ):
591 return self._alias_tensor_map
592 return self._placements
594 def set_alias_tensor_map(self, alias_tensor_map):
595 """Set alias_tensor_map"""
596 self._alias_tensor_map = alias_tensor_map
598 @property
599 def placements(self):
600 """placements"""
601 return self._placements
603 def set_placements(self, placements: Optional[Sequence[Placement]]) -> None:
604 """Set placements and retain the RaggedShard omitted from tensor_map."""
605 self._placements = placements
606 self._ragged_shard = (
607 None if placements is None else _extract_ragged_shard(placements)
608 )
610 @property
611 def normal_placements(self) -> Optional[tuple[Placement, ...]]:
612 """Return placements with RaggedShard represented as Replicate."""
613 if self._placements is None:
614 return None
615 return _replace_ragged_with_replicate(self._placements)
617 @property
618 def ragged_shard(self) -> Optional[RaggedShardInfo]:
619 """Return the RaggedShard placement and its mesh dimension, if present."""
620 return self._ragged_shard
622 @property
623 def tensor_map(self):
624 """tensor map"""
625 return self._tensor_map
627 def set_tensor_map(self, tensor_map):
628 """Set tensor_map."""
629 self._tensor_map = tensor_map
631 @property
632 def partial(self):
633 """partial status"""
634 return self._partial
636 def set_partial_by_dev_axis(self, axis, op):
637 """Set the partial status for the specified dev ID, means pending to do reduce by op."""
638 if op not in self._support_partial_op:
639 raise ValueError(f"Partial op must be one of {self._support_partial_op}, but got {op}")
640 if self.is_dev_axis_apply_shard(axis):
641 raise ValueError("Partial dim must be replicate.")
642 self._partial[self._mesh.axis_index(axis)] = op
643 self.tensor_map_to_placement()
644 self.update_compact_str()
646 def get_partial_by_dev_id(self, axis):
647 """Get the partial status for the specified dev id"""
648 return self.partial[self._mesh.axis_index(axis)]
650 def is_dev_axis_apply_shard(self, axis):
651 """Return true if device axis is applying shard"""
652 axis_id = self._mesh.axis_id(axis)
654 def flatten(input_x):
655 flatten_res = []
656 for item in input_x:
657 if isinstance(item, tuple):
658 flatten_res.extend(flatten(item))
659 else:
660 flatten_res.append(item)
661 return flatten_res
663 flatten_tensor_map = flatten(self.tensor_map)
664 return axis_id in flatten_tensor_map
666 def get_dev_axis_apply_shard_axis(self, axis):
667 """Return the axis which be split by axis. If axis not be apply to shard, return None."""
668 for dim, dim_map in enumerate(self.alias_tensor_map):
669 if (isinstance(dim_map, tuple) and axis in dim_map) or axis == dim_map:
670 return dim
671 return None
673 def reset_partial(self):
674 """Clear all partial statuses and regenerate placements from the tensor map."""
675 self._partial = [None] * len(self.mesh_shape)
676 self.tensor_map_to_placement()
677 self.update_compact_str()
679 def is_partial(self):
680 """Return true if any dim in mesh_shape is partial"""
681 return any(self.partial)
683 def get_dim_split_num(self, tensor_dim: int) -> int:
684 """Return the total shard count for ``tensor_dim`` via alias_tensor_map.
686 Args:
687 tensor_dim: Tensor dimension index to check.
689 Returns:
690 Number of shards (1 if not sharded or no alias_tensor_map set).
691 """
692 alias_tm = self.alias_tensor_map
693 if alias_tm is None or tensor_dim >= len(alias_tm):
694 return 1
695 dim_entry = alias_tm[tensor_dim]
696 if dim_entry == 'None':
697 return 1
698 if isinstance(dim_entry, str):
699 return self.mesh.get_device_num_along_axis(dim_entry)
700 if isinstance(dim_entry, tuple):
701 total = 1
702 for axis in dim_entry:
703 if axis != 'None':
704 total *= self.mesh.get_device_num_along_axis(axis)
705 return total
706 return 1
708 def get_split_id(self, tensor_dim: int) -> int:
709 """Return this rank's global position among all shards of ``tensor_dim``.
711 For a single sharding axis, returns the rank's position within that axis group.
712 For multiple sharding axes (e.g. dp+cp both sharding T1), computes the combined
713 global position as a mixed-radix number ordered by the axis tuple:
714 global_id = ax0_pos * ax1_size * ... + ax1_pos * ax2_size * ... + axN_pos
715 This matches MindFormers' ``offset_id = dp_rank * (cp*tp) + cp_rank * tp + tp_rank``
716 for combined sequence-parallel sharding across dp, cp, and tp dimensions.
718 Args:
719 tensor_dim: Tensor dimension index to query.
721 Returns:
722 Split index for this rank (0 if not sharded or rank not in rank list).
723 """
724 alias_tm = self.alias_tensor_map
725 if alias_tm is None or tensor_dim >= len(alias_tm):
726 return 0
727 dim_entry = alias_tm[tensor_dim]
728 if dim_entry == 'None':
729 return 0
730 rank = platform.get_rank()
731 if isinstance(dim_entry, tuple):
732 non_none = [ax for ax in dim_entry if ax != 'None']
733 if not non_none:
734 return 0
735 global_id = 0
736 for ax in non_none:
737 rank_list = self.mesh.get_rank_list_along_axis(ax)
738 local_id = rank_list.index(rank) if rank in rank_list else 0
739 ax_size = self.mesh.get_device_num_along_axis(ax)
740 global_id = global_id * ax_size + local_id
741 return global_id
742 if isinstance(dim_entry, str):
743 rank_list = self.mesh.get_rank_list_along_axis(dim_entry)
744 return rank_list.index(rank) if rank in rank_list else 0
745 return 0
747 def get_global_shape(self, slice_shape):
748 """get global shape"""
749 return self._mesh.get_global_shape(slice_shape, self._tensor_map)
751 def get_devices_for_axis(self, axis, rank):
752 """
753 Get the repeat rank list when the axis is not shard.
755 Args:
756 layout (Layout): Layout
757 axis (str): Axis name.
758 rank (int): Global rank
760 Returns:
761 list: reduce rank list
762 """
763 return self._mesh.get_devices_for_axis(axis, rank)
765 def get_comm_group_by_axis(self, axis):
766 """Return the communication group for the specified mesh axis via the underlying DeviceMesh."""
767 return self._mesh.get_comm_group_by_axis(axis)
769 def repeat_num(self):
770 """
771 Number of repeated placements.
772 For example:
773 layout = Layout((2, 4), ("dp", "mp"))
774 x_layout = layout("dp", "None")
775 The repeat_num is equal to all device num 8 divided by device num corresponding to used axis 2, that is 4.
776 """
777 if self._tensor_map is None:
778 raise ValueError(f"The tensor_map is None, the mesh_shape is {self._mesh.mesh_shape},"
779 f" alias_name is {self._mesh.mesh_dim_names}")
781 all_device_num = functools.reduce(lambda x, y: x * y, self._mesh.mesh_shape)
782 used_dev_num = 1
783 for ele in self._tensor_map:
784 if isinstance(ele, tuple):
785 for item in ele:
786 if item >= 0:
787 used_dev_num *= self._mesh.mesh_shape[len(self._mesh.mesh_shape) - item - 1]
788 continue
789 if ele >= 0:
790 used_dev_num *= self._mesh.mesh_shape[len(self._mesh.mesh_shape) - ele - 1]
792 return all_device_num // used_dev_num
794 def _to_compact_string(self):
795 """
796 generate dict key
798 Returns:
799 str: string for compact
800 """
801 mesh_key = self._mesh.to_hash()
802 hash_key = (self._tensor_map, self.partial)
803 hash_key += mesh_key
804 return str(hash_key)
806 @property
807 def compact_str(self):
808 """Return the cached compact string representation of this layout."""
809 return self._compact_str
811 def update_compact_str(self):
812 """Recompute and store the compact string representation of this layout."""
813 self._compact_str = self._to_compact_string()
815 def to_string(self):
816 """
817 layout dump
819 Returns:
820 str: layout string
821 """
822 device_info = f"Mesh shape: {self._mesh.mesh_shape}"
823 alias_info = f"Alias Names: {self._mesh.mesh_dim_names}"
824 rank_info = f"Rank List: {self._rank_list}"
825 partial_info = f"Partial: {self.partial}"
827 if self._tensor_map is None:
828 tensor_info = "Tensor Map: Not configured"
829 else:
830 readable_map = []
831 for item in self._tensor_map:
832 if isinstance(item, tuple):
833 # handle nested tuple
834 mapped_tuple = tuple(
835 self._mesh.mesh_dim_names[len(self._mesh.mesh_dim_names) - 1 - dim] if dim != -1 else "None"
836 for dim in item
837 )
838 readable_map.append(mapped_tuple)
839 else:
840 readable_map.append(
841 self._mesh.mesh_dim_names[len(self._mesh.mesh_dim_names) - 1 - item] if item != -1 else "None"
842 )
844 tensor_info = f"Tensor Map: {tuple(readable_map)}"
846 interleaved = "Yes" if "interleaved_parallel" in self._mesh.mesh_dim_names else "No"
847 interleaved_info = f"Interleaved Parallel: {interleaved}"
849 return (
850 f"Layout Configuration:\n"
851 f" {device_info}\n"
852 f" {alias_info}\n"
853 f" {partial_info}\n"
854 f" {tensor_info}\n"
855 f" {interleaved_info}\n"
856 f" {rank_info}"
857 )
859 def __str__(self):
860 """__str__"""
861 return self.to_string()
863 def __repr__(self):
864 """__repr__"""
865 return f"<Layout at {hex(id(self))}>"
867 def __eq__(self, other):
868 """
869 __eq__
870 """
871 if not isinstance(other, Layout):
872 return False
874 same_layout_attrs = (
875 self.mesh_shape,
876 self.alias_name,
877 self.partial,
878 self.rank_list,
879 ) == (
880 other.mesh_shape,
881 other.alias_name,
882 other.partial,
883 other.rank_list,
884 )
885 if not same_layout_attrs:
886 return False
888 if self._tensor_map is None or other.tensor_map is None:
889 return self._tensor_map is other.tensor_map
890 return self._tensor_map == other.tensor_map