Coverage for  / home / jenkins / .local / lib / python3.10 / site-packages / hyper_parallel / core / dtensor / layout.py: 96%

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"""layout""" 

16# pylint: disable=C9006,C9007,C9008 

17 

18import copy 

19import functools 

20from typing import Any, NamedTuple, Optional, Sequence 

21 

22import numpy as np 

23import torch.distributed as dist 

24 

25from hyper_parallel.core.dtensor.placement_types import ( 

26 Partial, 

27 Placement, 

28 RaggedShard, 

29 Replicate, 

30 Shard, 

31 StridedShard, 

32) 

33from hyper_parallel.core.dtensor.device_mesh import DeviceMesh, _create_device_mesh 

34 

35 

36class RaggedShardInfo(NamedTuple): 

37 """RaggedShard placement and its corresponding mesh dimension.""" 

38 

39 mesh_dim: int 

40 placement: RaggedShard 

41 

42 

43def _extract_ragged_shard(placements: Sequence[Placement]) -> Optional[RaggedShardInfo]: 

44 """Extract the single RaggedShard placement from a placement sequence.""" 

45 ragged_shard = None 

46 for mesh_dim, placement in enumerate(placements): 

47 if not placement.is_ragged_shard(): 

48 continue 

49 if ragged_shard is not None: 

50 raise ValueError( 

51 "Layout supports at most one RaggedShard placement, " 

52 f"but got placements={tuple(placements)!r}" 

53 ) 

54 ragged_shard = RaggedShardInfo(mesh_dim, placement) 

55 return ragged_shard 

56 

57 

58def _replace_ragged_with_replicate(placements: Sequence[Placement]) -> tuple[Placement, ...]: 

59 """Return placements with RaggedShard represented as Replicate.""" 

60 return tuple( 

61 Replicate() if placement.is_ragged_shard() else placement 

62 for placement in placements 

63 ) 

64 

65 

66def infer_balanced_chunk_range( 

67 tensor_size: int, 

68 shard_count: int, 

69 shard_rank: int, 

70) -> tuple[int, int]: 

71 """Return one balanced half-open shard range.""" 

72 shard_size, remainder = divmod(tensor_size, shard_count) 

73 chunk_start = shard_rank * shard_size + min(shard_rank, remainder) 

74 chunk_end = chunk_start + shard_size + int(shard_rank < remainder) 

75 return chunk_start, chunk_end 

76 

77 

78def infer_ceil_chunk_range( 

79 tensor_size: int, 

80 shard_count: int, 

81 shard_rank: int, 

82) -> tuple[int, int]: 

83 """Return one ceil-chunk range, including empty trailing ranks.""" 

84 chunk_size = (tensor_size + shard_count - 1) // shard_count 

85 chunk_start = min(shard_rank * chunk_size, tensor_size) 

86 chunk_end = min(chunk_start + chunk_size, tensor_size) 

87 return chunk_start, chunk_end 

88 

89 

90def _calculate_mesh_coordinate_for_rank( 

91 mesh_shape: Sequence[int], rank_id: int 

92) -> list[int]: 

93 """Calculate the device-mesh coordinate for a linear rank.""" 

94 coordinate = [0] * len(mesh_shape) 

95 remaining_rank = rank_id 

96 for mesh_dim in range(len(mesh_shape) - 1, -1, -1): 

97 coordinate[mesh_dim] = remaining_rank % mesh_shape[mesh_dim] 

98 remaining_rank //= mesh_shape[mesh_dim] 

99 return coordinate 

100 

101 

102def _infer_tensor_dimension_slice( 

103 global_size: int, 

104 tensor_mapping: int | Sequence[int], 

105 mesh_shape: Sequence[int], 

106 mesh_coordinate: Sequence[int], 

107 uneven_shard_mesh_dims: set[int], 

108) -> tuple[int, int]: 

109 """Infer the half-open slice for one tensor dimension.""" 

110 mapped_dims = (tensor_mapping,) if isinstance(tensor_mapping, int) else tensor_mapping 

111 slice_start, slice_end = 0, global_size 

112 for mapped_mesh_dim in mapped_dims: 

113 if mapped_mesh_dim == -1: 

114 continue 

115 mesh_dim = len(mesh_shape) - mapped_mesh_dim - 1 

116 infer_chunk_range = ( 

117 infer_ceil_chunk_range 

118 if mesh_dim in uneven_shard_mesh_dims 

119 else infer_balanced_chunk_range 

120 ) 

121 chunk_start, chunk_end = infer_chunk_range( 

122 slice_end - slice_start, 

123 mesh_shape[mesh_dim], 

124 mesh_coordinate[mesh_dim], 

125 ) 

126 slice_start += chunk_start 

127 slice_end = slice_start + chunk_end - chunk_start 

128 return slice_start, slice_end 

129 

130 

131def _infer_slice_area_by_rank( 

132 mesh_shape: tuple[int, ...], 

133 tensor_map: Sequence, 

134 rank_id: int, 

135 full_shape: Sequence[int], 

136 uneven_shard_mesh_dims: set[int], 

137) -> tuple[tuple[int, int], ...]: 

138 """Return one rank's slice using placement-specific shard geometry.""" 

139 mesh_coordinate = _calculate_mesh_coordinate_for_rank(mesh_shape, rank_id) 

140 return tuple( 

141 _infer_tensor_dimension_slice( 

142 global_size, 

143 tensor_map[tensor_dim], 

144 mesh_shape, 

145 mesh_coordinate, 

146 uneven_shard_mesh_dims, 

147 ) 

148 for tensor_dim, global_size in enumerate(full_shape) 

149 ) 

150 

151 

152def infer_slice_area_by_rank( 

153 mesh_shape: tuple[int, ...], 

154 tensor_map: Sequence, 

155 rank_id: int, 

156 full_shape: Sequence[int], 

157 uneven_shard_mesh_dims: Optional[Sequence[int]] = None, 

158) -> tuple[tuple[int, int], ...]: 

159 """Return one rank's global slice using placement-specific shard semantics.""" 

160 return _infer_slice_area_by_rank( 

161 mesh_shape, 

162 tensor_map, 

163 rank_id, 

164 full_shape, 

165 set(uneven_shard_mesh_dims or ()), 

166 ) 

167 

168 

169def infer_slice_area_by_layout( 

170 layout: "Layout", 

171 rank_id: int, 

172 full_shape: Sequence[int], 

173) -> tuple[tuple[int, int], ...]: 

174 """Return one rank's slice using ``uneven_shard`` placement markers.""" 

175 return infer_slice_area_by_rank( 

176 layout.mesh_shape, 

177 layout.tensor_map, 

178 rank_id, 

179 full_shape, 

180 layout.uneven_shard_mesh_dims, 

181 ) 

182 

183 

184def _get_slice_tensor_by_layout(global_tensor, layout): 

185 """Transfer global tensor to local tensor by layout""" 

186 inner_rank_id = layout.rank_list.index(layout.mesh.rank) 

187 slice_area = infer_slice_area_by_layout(layout, inner_rank_id, global_tensor.shape) 

188 

189 def get_slice_data(full_data, offset): 

190 area = () 

191 for begin, end in offset: 

192 area += (slice(begin, end),) 

193 return full_data[area].clone() 

194 

195 local_tensor = get_slice_data(global_tensor, slice_area) 

196 return local_tensor 

197 

198 

199def _infer_slice_shape_by_layout(global_shape, layout): 

200 """Infer slice shape from global_shape and layout""" 

201 inner_rank_id = layout.rank_list.index(layout.mesh.rank) 

202 slice_area = infer_slice_area_by_layout(layout, inner_rank_id, global_shape) 

203 return [end - start for start, end in slice_area] 

204 

205 

206class Layout: 

207 """ 

208 Topological abstraction describing cluster devices for tensor slice placement on the cluster. 

209 

210 Note: 

211 - It is valid only in semi auto parallel or auto parallel mode. 

212 - The multiplication result of the `mesh_shape` must be equal to the device count in a pipeline stage. 

213 - When the layout function is invoked to constructs a sharding strategy, each alias name is only allowed to be 

214 used once to shard a tensor. 

215 

216 Args: 

217 mesh_shape (tuple): Describe the shape of devices arrangement, its element type is int. 

218 alias_name (tuple): The alias name for each axis of mesh_shape, its length shoits element type is string. 

219 When using "interleaved_parallel" as an alias name, the tensor would be split into multiple 

220 copies on the corresponding partition dimension on a single card. 

221 rank_list (tuple, optional): Data is allocated to the device according to rank_list. Default: ``None``. 

222 

223 Raises: 

224 TypeError: `mesh_shape` is not a tuple type. 

225 TypeError: `alias_name` is not a tuple type. 

226 TypeError: 'rank_list' is not a list type. 

227 ValueError: `mesh_shape` length is not equal to `alias_name` length. 

228 TypeError: The element of `mesh_shape` is not int type. 

229 TypeError: The element of `alias_name` is not a str type. 

230 TypeError: The element of `rank_list` is not int type. 

231 ValueError: The element of `alias_name` is an empty str. 

232 ValueError: The element of `alias_name` is "None". 

233 ValueError: `alias_name` contains repeated element. 

234 

235 Supported Platforms: 

236 ``Ascend`` 

237 

238 Examples: 

239 >>> layout = Layout((2, 2, 2), ("dp", "sp", "mp")) 

240 >>> layout0 = layout("dp", "mp") 

241 >>> print(layout0.to_dict()) 

242 {"mesh_shape": (2, 2, 2), "tensor_map": (2, 0), "interleaved_parallel": False, 

243 'alias_name': {'dp', 'sp', 'mp'}, "rank_list": [0, 1, 2, 3, 4, 5, 6, 7]} 

244 >>> layout = Layout((2, 2, 2), ("dp", "sp", "interleaved_parallel")) 

245 >>> layout1 = layout(("dp", "interleaved_parallel"), "sp") 

246 """ 

247 

248 def __init__(self, mesh_shape, alias_name, rank_list=None, init_backend=True): 

249 self._alias_name = alias_name 

250 self._tensor_map = None 

251 if not rank_list: 

252 self._rank_list = tuple(range(np.prod(np.array(mesh_shape)))) 

253 else: 

254 self._rank_list = tuple(rank_list) 

255 self._partial = [None] * len(mesh_shape) # partial status for each dev dim 

256 self._support_partial_op = ['sum', 'max', 'min', 'avg', 'prod', 'all', None] 

257 self._alias_tensor_map = None 

258 self._tensor_shape = None 

259 self._tensor_stride = None 

260 self._tensor_dtype = None 

261 self._mesh = _create_device_mesh("npu", mesh_shape, mesh_dim_names=alias_name, rank_list=self._rank_list, 

262 init_backend=init_backend) 

263 self._placements = None 

264 self.partial_ops = {} # Initialized in _build_dim_map_from_placements() 

265 self._ragged_shard = None 

266 self._compact_str = self._to_compact_string() 

267 

268 @classmethod 

269 def from_device_mesh(cls, device_mesh: DeviceMesh) -> 'Layout': 

270 """ 

271 Create a Layout from an existing DeviceMesh. 

272 

273 Args: 

274 device_mesh (DeviceMesh): The device mesh to create layout from. 

275 

276 Returns: 

277 Layout: A new Layout instance initialized with the properties of the provided device mesh. 

278 

279 Examples: 

280 >>> from hyper_parallel.core.dtensor.layout import Layout, DeviceMesh 

281 >>> device_mesh = DeviceMesh("npu", (2, 2), mesh_dim_names=("dp", "mp")) 

282 >>> layout = Layout.from_device_mesh(device_mesh) 

283 """ 

284 obj = cls.__new__(cls) 

285 obj._mesh = device_mesh 

286 obj._alias_name = device_mesh.mesh_dim_names 

287 obj._rank_list = device_mesh.rank_list 

288 obj._tensor_map = None 

289 obj._partial = [None] * len(device_mesh.mesh_shape) 

290 obj._support_partial_op = ['sum', 'max', 'min', 'avg', 'prod', 'all', None] 

291 obj._alias_tensor_map = None 

292 obj._tensor_shape = None 

293 obj._tensor_stride = None 

294 obj._tensor_dtype = None 

295 obj._placements = None 

296 obj._ragged_shard = None 

297 obj._compact_str = obj._to_compact_string() 

298 return obj 

299 

300 def __call__(self, *alias_tensor_map): 

301 obj = copy.deepcopy(self) 

302 

303 # Clear the inherited partial status. 

304 # When creating a new layout mapping configuration via __call__, 

305 # it should not inherit the dynamic execution state (Partial) of the original layout. 

306 # If the user intends to create a Partial placement, it will be parsed from alias_tensor_map. 

307 obj._partial = [None] * len(obj.mesh_shape) 

308 

309 if len(alias_tensor_map) == 1 and isinstance(alias_tensor_map[0], (list, tuple)): 

310 if len(alias_tensor_map[0]) > 0 and isinstance(alias_tensor_map[0][0], Placement): 

311 return self._process_placement_layout(obj, alias_tensor_map[0]) 

312 

313 if len(alias_tensor_map) > 0 and isinstance(alias_tensor_map[0], Placement): 

314 return self._process_placement_layout(obj, alias_tensor_map) 

315 

316 return self._process_alias_layout(obj, alias_tensor_map) 

317 

318 def __deepcopy__(self, memo): 

319 """Deep copy layout without rebuilding the underlying device mesh.""" 

320 cls = self.__class__ 

321 result = cls.__new__(cls) 

322 memo[id(self)] = result 

323 for k, v in self.__dict__.items(): 

324 setattr(result, k, copy.deepcopy(v, memo)) 

325 return result 

326 

327 @staticmethod 

328 def _process_placement_layout(obj, placements): 

329 """Process layout defined by Placement types.""" 

330 obj.set_placements(placements) 

331 return copy.deepcopy(obj) 

332 

333 @staticmethod 

334 def _process_alias_layout(obj, alias_tensor_map): 

335 """Process layout defined by alias strings.""" 

336 obj.set_alias_tensor_map(alias_tensor_map) 

337 tensor_map = () 

338 writed_map = () 

339 for ele in alias_tensor_map: 

340 if isinstance(ele, tuple): 

341 ele_map = () 

342 for item in ele: 

343 if item == "None": 

344 ele_map += (-1,) 

345 continue 

346 if item not in obj.alias_name: 

347 raise ValueError(f'The axis {item} is not found in {obj.alias_name}') 

348 if item in writed_map: 

349 raise ValueError(f'The axis {item} has been set more than one in {obj.alias_name}') 

350 ele_map += (len(obj.alias_name) - 1 - obj.alias_name.index(item),) 

351 writed_map += (item,) 

352 tensor_map += (ele_map,) 

353 continue 

354 if ele == "None": 

355 tensor_map += (-1,) 

356 continue 

357 if ele not in obj.alias_name: 

358 raise ValueError(f'The axis {ele} is not found in {obj.alias_name}') 

359 if ele in writed_map: 

360 raise ValueError(f'The axis {ele} has been set more than one in {obj.alias_name}') 

361 tensor_map += (len(obj.alias_name) - 1 - obj.alias_name.index(ele),) 

362 writed_map += (ele,) 

363 obj.set_tensor_map(tensor_map) 

364 obj.tensor_map_to_placement() 

365 obj.update_compact_str() 

366 return copy.deepcopy(obj) 

367 

368 def to_dict(self): 

369 """ 

370 Transform layout to a dictionary. 

371 """ 

372 if self._mesh.mesh_shape is None: 

373 raise ValueError("The device_shape of layout is None") 

374 if self._tensor_map is None: 

375 raise ValueError("The tensor_map of layout is None") 

376 interleaved_parallel = "interleaved_parallel" in self._mesh.mesh_dim_names 

377 return {"mesh_shape": self._mesh.mesh_shape, "tensor_map": self._tensor_map, 

378 "interleaved_parallel": interleaved_parallel, "alias_name": self._mesh.mesh_dim_names, 

379 "rank_list": self._rank_list} 

380 

381 def placement_to_tensor_map(self, dim): 

382 """ 

383 Transform placement to tensor map. 

384 

385 This method converts the `placements` configuration (consisting of Shard, StridedShard, 

386 Replicate, Partial) 

387 into a `tensor_map` representation used for distributed tensor operations. 

388 

389 Args: 

390 dim (int): The dimension of the tensor. Must be a positive integer. 

391 

392 Returns: 

393 tuple: A tuple representing the tensor map, where each element corresponds to a tensor dimension. 

394 A value of -1 indicates the dimension is not sharded, an integer indicates the mesh 

395 dimension index along which the tensor dimension is sharded, and a tuple indicates 

396 that the same tensor dimension is sharded multiple times in order. 

397 

398 Raises: 

399 ValueError: If `dim` is negative. 

400 ValueError: If a shard dimension in `placements` is out of bounds for the given tensor dimension. 

401 """ 

402 if dim < 0: 

403 raise ValueError(f"Tensor dimension must be positive, but got {dim}") 

404 if dim == 0: 

405 return self._handle_zero_dim_placement() 

406 

407 dim_map = self._build_dim_map_from_placements(dim) 

408 tensor_map = self._convert_dim_map_to_tensor_map(dim_map) 

409 self.set_tensor_map(tuple(tensor_map)) 

410 self._alias_tensor_map = self._build_readable_tensor_map() 

411 self.update_compact_str() 

412 return tensor_map 

413 

414 def _handle_zero_dim_placement(self): 

415 """Handle the special case of zero-dimensional tensor.""" 

416 self.set_tensor_map(()) 

417 self._alias_tensor_map = () 

418 for mesh_idx, placement in enumerate(self.normal_placements): 

419 if isinstance(placement, Partial): 

420 self._partial[mesh_idx] = self._extract_reduce_op(placement) 

421 return [] 

422 

423 def _build_dim_map_from_placements(self, dim): 

424 """Build dimension map from placements.""" 

425 dim_map = [-1] * dim 

426 self.partial_ops = {} 

427 for mesh_idx, placement in enumerate(self.normal_placements): 

428 if isinstance(placement, Shard): 

429 shard_dim = placement.dim 

430 if shard_dim < -dim or shard_dim >= dim: 

431 raise ValueError(f"Shard dimension {shard_dim} is out of bounds for tensor of dimension {dim}") 

432 if shard_dim < 0: 

433 shard_dim += dim 

434 if dim_map[shard_dim] == -1: 

435 dim_map[shard_dim] = [mesh_idx] 

436 else: 

437 dim_map[shard_dim].append(mesh_idx) 

438 elif isinstance(placement, Partial): 

439 self._partial[mesh_idx] = self._extract_reduce_op(placement) 

440 self._validate_strided_shard_split_factor(dim_map) 

441 self._reorder_dim_map_for_strided_shard(dim_map) 

442 return dim_map 

443 

444 @staticmethod 

445 def _placement_split_factor(placement): 

446 """Return the effective split factor carried by a placement.""" 

447 return placement.split_factor if isinstance(placement, StridedShard) else 1 

448 

449 @staticmethod 

450 def _build_order_positions(shard_order): 

451 """Build a mesh axis to order position mapping.""" 

452 return {mesh_idx: order_idx for order_idx, mesh_idx in enumerate(shard_order)} 

453 

454 def _compute_expected_split_factors(self, shard_axes, shard_order): 

455 """Infer the split_factor each mesh axis should carry for the given sharding order.""" 

456 order_positions = self._build_order_positions(shard_order) 

457 expected_split_factors = {} 

458 for mesh_idx in shard_axes: 

459 split_factor = 1 

460 for right_mesh_idx in shard_axes: 

461 if right_mesh_idx <= mesh_idx: 

462 continue 

463 if order_positions[right_mesh_idx] < order_positions[mesh_idx]: 

464 split_factor *= self.mesh_shape[right_mesh_idx] 

465 expected_split_factors[mesh_idx] = split_factor 

466 return expected_split_factors 

467 

468 def _get_effective_shard_axes(self, shard_axes): 

469 """Return shard axes ordered by their effective sharding order.""" 

470 return sorted( 

471 shard_axes, 

472 key=lambda mesh_idx: self._placement_split_factor(self.placements[mesh_idx]), 

473 ) 

474 

475 def _reorder_dim_map_for_strided_shard(self, dim_map): 

476 """Reorder dim_map entries to reflect the effective sharding order.""" 

477 for i, shard_axes in enumerate(dim_map): 

478 if shard_axes == -1 or len(shard_axes) <= 1: 

479 continue 

480 dim_map[i] = self._get_effective_shard_axes(shard_axes) 

481 

482 def _validate_strided_shard_split_factor(self, dim_map): 

483 """Validate that split factors match the effective sharding order.""" 

484 for shard_axes in dim_map: 

485 if shard_axes == -1: 

486 continue 

487 shard_order = self._get_effective_shard_axes(shard_axes) 

488 expected_split_factors = self._compute_expected_split_factors( 

489 shard_axes, shard_order 

490 ) 

491 for mesh_idx in shard_axes: 

492 placement = self.placements[mesh_idx] 

493 actual_split_factor = self._placement_split_factor(placement) 

494 expected_split_factor = expected_split_factors[mesh_idx] 

495 if actual_split_factor != expected_split_factor: 

496 raise ValueError( 

497 f"StridedShard split_factor mismatch on mesh axis {mesh_idx}: " 

498 f"expected {expected_split_factor}, got {actual_split_factor}." 

499 ) 

500 

501 @staticmethod 

502 def _extract_reduce_op(placement): 

503 """Extract reduce operation name from Partial placement.""" 

504 op_name = getattr(placement, "reduce_op", "sum") 

505 if isinstance(op_name, str): 

506 op_name = op_name.lower() 

507 return op_name 

508 

509 def _convert_dim_map_to_tensor_map(self, dim_map): 

510 """Convert dimension map to tensor map format.""" 

511 device_dim_count = len(self.mesh_shape) 

512 tensor_map = [] 

513 for mesh_idx in dim_map: 

514 if mesh_idx == -1: 

515 tensor_map.append(-1) 

516 continue 

517 mapped_axes = tuple(device_dim_count - 1 - axis for axis in mesh_idx) 

518 tensor_map.append(mapped_axes[0] if len(mapped_axes) == 1 else mapped_axes) 

519 return tensor_map 

520 

521 def _build_readable_tensor_map(self): 

522 """Build human-readable alias tensor map from tensor_map.""" 

523 mesh_dim_names = self._mesh.mesh_dim_names 

524 has_names = mesh_dim_names is not None 

525 

526 def _map_dim(dim): 

527 """convert dimension index to dimension name.""" 

528 if dim == -1: 

529 return "None" 

530 if not has_names: 

531 return f"dim_{dim}" 

532 return mesh_dim_names[len(mesh_dim_names) - 1 - dim] 

533 

534 readable_map = [] 

535 for item in self._tensor_map: 

536 if isinstance(item, tuple): 

537 mapped_tuple = tuple(_map_dim(dim) for dim in item) 

538 readable_map.append(mapped_tuple) 

539 else: 

540 readable_map.append(_map_dim(item)) 

541 return tuple(readable_map) 

542 

543 def tensor_map_to_placement(self): 

544 """ 

545 Transform tensor map to placement. 

546 

547 This method converts the existing `tensor_map` and `partial` status into a list of `Placement` objects 

548 (Shard, StridedShard, Replicate, Partial). This is the inverse operation of 

549 `placement_to_tensor_map`. 

550 

551 Returns: 

552 list[Placement]: A list of Placement objects describing the distribution strategy for each 

553 dimension of the device mesh. 

554 

555 Raises: 

556 ValueError: If `tensor_map` is not configured (None). 

557 """ 

558 if self._tensor_map is None: 

559 raise ValueError("The tensor_map is None, cannot transform to placements.") 

560 uneven_shard_dims = { 

561 mesh_idx: placement.dim 

562 for mesh_idx, placement in enumerate(self.placements or ()) 

563 if isinstance(placement, Shard) and placement.uneven_shard 

564 } 

565 mesh_ndim = len(self.mesh_shape) 

566 placements = [Replicate()] * mesh_ndim 

567 for tensor_dim, mapping in enumerate(self._tensor_map): 

568 mapping_list = mapping if isinstance(mapping, tuple) else (mapping,) 

569 valid_mapping = [map_val for map_val in mapping_list if map_val != -1] 

570 mesh_indices = [mesh_ndim - 1 - map_val for map_val in valid_mapping] 

571 shard_axes = sorted(mesh_indices) 

572 expected_split_factors = self._compute_expected_split_factors( 

573 shard_axes, mesh_indices 

574 ) 

575 for mesh_idx in shard_axes: 

576 split_factor = expected_split_factors[mesh_idx] 

577 uneven_shard = uneven_shard_dims.get(mesh_idx) == tensor_dim 

578 placement = ( 

579 StridedShard( 

580 dim=tensor_dim, 

581 split_factor=split_factor, 

582 uneven_shard=uneven_shard, 

583 ) 

584 if split_factor > 1 

585 else Shard(dim=tensor_dim, uneven_shard=uneven_shard) 

586 ) 

587 placements[mesh_idx] = placement 

588 for mesh_idx, op in enumerate(self.partial): 

589 if op is not None: 

590 placements[mesh_idx] = Partial(reduce_op=op) 

591 if self._ragged_shard is not None: 

592 placements[self._ragged_shard.mesh_dim] = self._ragged_shard.placement 

593 self.set_placements(placements) 

594 self._alias_tensor_map = self._build_readable_tensor_map() 

595 self.update_compact_str() 

596 return placements 

597 

598 def __setstate__(self, state): 

599 self.__dict__.update(state) 

600 self.update_mesh(init_backend=False) 

601 

602 @property 

603 def mesh(self): 

604 """ 

605 Get the device mesh associated with this layout. 

606 

607 Returns: 

608 DeviceMesh: The device mesh describing the device topology. 

609 """ 

610 return self._mesh 

611 

612 def update_mesh(self, init_backend: bool = True): 

613 """Recreate the internal DeviceMesh from current layout properties. 

614 

615 Args: 

616 init_backend (bool): Whether to initialize communication backend 

617 (process groups). Set to ``False`` during deserialization to 

618 avoid creating process groups with a stale rank_list from the 

619 sender side. Default ``True``. 

620 """ 

621 self._mesh = _create_device_mesh("npu", self.mesh_shape, mesh_dim_names=self.alias_name, 

622 rank_list=self.rank_list, init_backend=init_backend) 

623 

624 @property 

625 def rank_list(self): 

626 """ 

627 Get the list of ranks participating in this layout. 

628 

629 Returns: 

630 tuple[int]: The rank list. 

631 """ 

632 return self._rank_list 

633 

634 @rank_list.setter 

635 def rank_list(self, val): 

636 self._rank_list = val 

637 

638 @property 

639 def mesh_shape(self): 

640 """mesh shape""" 

641 return self._mesh.mesh_shape 

642 

643 @property 

644 def alias_name(self): 

645 """alias name""" 

646 return self._mesh.mesh_dim_names 

647 

648 @property 

649 def alias_tensor_map(self): 

650 """Return the human-readable alias tensor map for this layout.""" 

651 return self._alias_tensor_map 

652 

653 @property 

654 def alias_placements(self): 

655 """Return alias_tensor_map when it contains multi-axis tuples, otherwise placements. 

656 

657 alias_tensor_map preserves multi-axis ordering information 

658 (e.g., (("dp", "tp"), "None") vs (("tp", "dp"), "None")) 

659 that Placement objects cannot represent, since both map to 

660 [Shard(0), Shard(0)]. 

661 

662 For single-axis layouts, Placement objects are preferred because they 

663 also carry Partial status which alias_tensor_map cannot encode. 

664 

665 Use this property when constructing DTensors from an existing Layout 

666 to avoid the lossy Placement round-trip for multi-axis cases. 

667 """ 

668 if self._ragged_shard is not None: 

669 return self._placements 

670 if self._alias_tensor_map is not None and any( 

671 isinstance(item, tuple) for item in self._alias_tensor_map 

672 ): 

673 return self._alias_tensor_map 

674 return self._placements 

675 

676 def set_alias_tensor_map(self, alias_tensor_map): 

677 """Set alias_tensor_map""" 

678 self._alias_tensor_map = alias_tensor_map 

679 

680 @property 

681 def placements(self): 

682 """placements""" 

683 return self._placements 

684 

685 def set_placements(self, placements: Optional[Sequence[Placement]]) -> None: 

686 """Set placements and retain the RaggedShard omitted from tensor_map.""" 

687 self._placements = placements 

688 self._ragged_shard = ( 

689 None if placements is None else _extract_ragged_shard(placements) 

690 ) 

691 

692 @property 

693 def tensor_shape(self) -> Optional[tuple[int, ...]]: 

694 """Return the explicit logical global shape, if present.""" 

695 return self._tensor_shape 

696 

697 @property 

698 def tensor_stride(self) -> Optional[tuple[int, ...]]: 

699 """Return the explicit logical global stride, if present.""" 

700 return self._tensor_stride 

701 

702 @property 

703 def tensor_dtype(self) -> Optional[Any]: 

704 """Return the explicit logical dtype, if present.""" 

705 return self._tensor_dtype 

706 

707 @property 

708 def uneven_shard_mesh_dims(self) -> tuple[int, ...]: 

709 """Return mesh dimensions carrying FSDP ceil-chunk placements.""" 

710 return tuple( 

711 mesh_dim 

712 for mesh_dim, placement in enumerate(self._placements or ()) 

713 if isinstance(placement, Shard) and placement.uneven_shard 

714 ) 

715 

716 @property 

717 def has_uneven_shard(self) -> bool: 

718 """Return whether any placement uses FSDP ceil-chunk geometry.""" 

719 return bool(self.uneven_shard_mesh_dims) 

720 

721 def set_tensor_meta( 

722 self, 

723 shape: Sequence[int], 

724 stride: Sequence[int], 

725 dtype: Any, 

726 ) -> None: 

727 """Set logical tensor metadata independently from local shard storage. 

728 

729 Args: 

730 shape: Logical global tensor shape. 

731 stride: Logical global tensor stride. 

732 dtype: Logical tensor dtype. 

733 """ 

734 self._tensor_shape = tuple(shape) 

735 self._tensor_stride = tuple(stride) 

736 self._tensor_dtype = dtype 

737 self.update_compact_str() 

738 

739 @property 

740 def normal_placements(self) -> Optional[tuple[Placement, ...]]: 

741 """Return placements with RaggedShard represented as Replicate.""" 

742 if self._placements is None: 

743 return None 

744 return _replace_ragged_with_replicate(self._placements) 

745 

746 @property 

747 def ragged_shard(self) -> Optional[RaggedShardInfo]: 

748 """Return the RaggedShard placement and its mesh dimension, if present.""" 

749 return self._ragged_shard 

750 

751 @property 

752 def tensor_map(self): 

753 """tensor map""" 

754 return self._tensor_map 

755 

756 def set_tensor_map(self, tensor_map): 

757 """Set tensor_map.""" 

758 self._tensor_map = tensor_map 

759 

760 @property 

761 def partial(self): 

762 """partial status""" 

763 return self._partial 

764 

765 def set_partial_by_dev_axis(self, axis, op): 

766 """Set the partial status for the specified dev ID, means pending to do reduce by op.""" 

767 if op not in self._support_partial_op: 

768 raise ValueError(f"Partial op must be one of {self._support_partial_op}, but got {op}") 

769 if self.is_dev_axis_apply_shard(axis): 

770 raise ValueError("Partial dim must be replicate.") 

771 self._partial[self._mesh.axis_index(axis)] = op 

772 self.tensor_map_to_placement() 

773 self.update_compact_str() 

774 

775 def get_partial_by_dev_id(self, axis): 

776 """Get the partial status for the specified dev id""" 

777 return self.partial[self._mesh.axis_index(axis)] 

778 

779 def is_dev_axis_apply_shard(self, axis): 

780 """Return true if device axis is applying shard""" 

781 axis_id = self._mesh.axis_id(axis) 

782 

783 def flatten(input_x): 

784 flatten_res = [] 

785 for item in input_x: 

786 if isinstance(item, tuple): 

787 flatten_res.extend(flatten(item)) 

788 else: 

789 flatten_res.append(item) 

790 return flatten_res 

791 

792 flatten_tensor_map = flatten(self.tensor_map) 

793 return axis_id in flatten_tensor_map 

794 

795 def get_dev_axis_apply_shard_axis(self, axis): 

796 """Return the axis which be split by axis. If axis not be apply to shard, return None.""" 

797 for dim, dim_map in enumerate(self.alias_tensor_map): 

798 if (isinstance(dim_map, tuple) and axis in dim_map) or axis == dim_map: 

799 return dim 

800 return None 

801 

802 def reset_partial(self): 

803 """Clear all partial statuses and regenerate placements from the tensor map.""" 

804 self._partial = [None] * len(self.mesh_shape) 

805 self.tensor_map_to_placement() 

806 self.update_compact_str() 

807 

808 def is_partial(self): 

809 """Return true if any dim in mesh_shape is partial""" 

810 return any(self.partial) 

811 

812 def get_dim_split_num(self, tensor_dim: int) -> int: 

813 """Return the total shard count for ``tensor_dim`` via alias_tensor_map. 

814 

815 Args: 

816 tensor_dim: Tensor dimension index to check. 

817 

818 Returns: 

819 Number of shards (1 if not sharded or no alias_tensor_map set). 

820 """ 

821 alias_tm = self.alias_tensor_map 

822 if alias_tm is None or tensor_dim >= len(alias_tm): 

823 return 1 

824 dim_entry = alias_tm[tensor_dim] 

825 if dim_entry == 'None': 

826 return 1 

827 if isinstance(dim_entry, str): 

828 return self.mesh.get_device_num_along_axis(dim_entry) 

829 if isinstance(dim_entry, tuple): 

830 total = 1 

831 for axis in dim_entry: 

832 if axis != 'None': 

833 total *= self.mesh.get_device_num_along_axis(axis) 

834 return total 

835 return 1 

836 

837 def get_split_id(self, tensor_dim: int) -> int: 

838 """Return this rank's global position among all shards of ``tensor_dim``. 

839 

840 For a single sharding axis, returns the rank's position within that axis group. 

841 For multiple sharding axes (e.g. dp+cp both sharding T1), computes the combined 

842 global position as a mixed-radix number ordered by the axis tuple: 

843 global_id = ax0_pos * ax1_size * ... + ax1_pos * ax2_size * ... + axN_pos 

844 This matches MindFormers' ``offset_id = dp_rank * (cp*tp) + cp_rank * tp + tp_rank`` 

845 for combined sequence-parallel sharding across dp, cp, and tp dimensions. 

846 

847 Args: 

848 tensor_dim: Tensor dimension index to query. 

849 

850 Returns: 

851 Split index for this rank (0 if not sharded or rank not in rank list). 

852 """ 

853 alias_tm = self.alias_tensor_map 

854 if alias_tm is None or tensor_dim >= len(alias_tm): 

855 return 0 

856 dim_entry = alias_tm[tensor_dim] 

857 if dim_entry == 'None': 

858 return 0