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

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

16 

17import copy 

18import functools 

19from typing import Any, NamedTuple, Optional, Sequence 

20 

21import numpy as np 

22 

23 

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 

34 

35platform = get_platform() 

36 

37 

38class RaggedShardInfo(NamedTuple): 

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

40 

41 mesh_dim: int 

42 placement: RaggedShard 

43 

44 

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 

58 

59 

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 ) 

66 

67 

68def infer_balanced_chunk_range( 

69 tensor_size: int, 

70 shard_count: int, 

71 shard_rank: int, 

72) -> tuple[int, int]: 

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

74 shard_size, remainder = divmod(tensor_size, shard_count) 

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

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

77 return chunk_start, chunk_end 

78 

79 

80def infer_ceil_chunk_range( 

81 tensor_size: int, 

82 shard_count: int, 

83 shard_rank: int, 

84) -> tuple[int, int]: 

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

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

87 chunk_start = min(shard_rank * chunk_size, tensor_size) 

88 chunk_end = min(chunk_start + chunk_size, tensor_size) 

89 return chunk_start, chunk_end 

90 

91 

92def _infer_slice_area_by_rank( 

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

94 tensor_map: Sequence, 

95 rank_id: int, 

96 full_shape: Sequence[int], 

97 uneven_shard_mesh_dims: set[int], 

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

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

100 mesh_coordinate = [0] * len(mesh_shape) 

101 remaining_rank = rank_id 

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

103 mesh_coordinate[mesh_dim] = remaining_rank % mesh_shape[mesh_dim] 

104 remaining_rank //= mesh_shape[mesh_dim] 

105 

106 slice_area = [] 

107 for tensor_dim, global_size in enumerate(full_shape): 

108 tensor_mapping = tensor_map[tensor_dim] 

109 if isinstance(tensor_mapping, int): 

110 tensor_mapping = (tensor_mapping,) 

111 

112 slice_start = 0 

113 slice_end = global_size 

114 for mapped_mesh_dim in tensor_mapping: 

115 if mapped_mesh_dim == -1: 

116 continue 

117 mesh_dim = len(mesh_shape) - mapped_mesh_dim - 1 

118 shard_count = mesh_shape[mesh_dim] 

119 shard_rank = mesh_coordinate[mesh_dim] 

120 current_size = slice_end - slice_start 

121 infer_chunk_range = ( 

122 infer_ceil_chunk_range 

123 if mesh_dim in uneven_shard_mesh_dims 

124 else infer_balanced_chunk_range 

125 ) 

126 chunk_start, chunk_end = infer_chunk_range( 

127 current_size, 

128 shard_count, 

129 shard_rank, 

130 ) 

131 slice_end = slice_start + chunk_end 

132 slice_start += chunk_start 

133 slice_area.append((slice_start, slice_end)) 

134 return tuple(slice_area) 

135 

136 

137def infer_slice_area_by_rank( 

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

139 tensor_map: Sequence, 

140 rank_id: int, 

141 full_shape: Sequence[int], 

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

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

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

145 return _infer_slice_area_by_rank( 

146 mesh_shape, 

147 tensor_map, 

148 rank_id, 

149 full_shape, 

150 set(uneven_shard_mesh_dims or ()), 

151 ) 

152 

153 

154def infer_slice_area_by_layout( 

155 layout: "Layout", 

156 rank_id: int, 

157 full_shape: Sequence[int], 

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

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

160 return infer_slice_area_by_rank( 

161 layout.mesh_shape, 

162 layout.tensor_map, 

163 rank_id, 

164 full_shape, 

165 layout.uneven_shard_mesh_dims, 

166 ) 

167 

168 

169def _get_slice_tensor_by_layout(global_tensor, layout): 

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

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

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

173 

174 def get_slice_data(full_data, offset): 

175 area = () 

176 for begin, end in offset: 

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

178 return full_data[area].clone() 

179 

180 local_tensor = get_slice_data(global_tensor, slice_area) 

181 return local_tensor 

182 

183 

184def _infer_slice_shape_by_layout(global_shape, layout): 

185 """Infer slice shape from global_shape and 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_shape) 

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

189 

190 

191class Layout: 

192 """ 

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

194 

195 Note: 

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

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

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

199 used once to shard a tensor. 

200 

201 Args: 

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

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

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

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

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

207 

208 Raises: 

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

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

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

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

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

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

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

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

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

218 ValueError: `alias_name` contains repeated element. 

219 

220 Supported Platforms: 

221 ``Ascend`` 

222 

223 Examples: 

224 >>> from mindspore.parallel import Layout 

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

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

227 >>> print(layout0.to_dict()) 

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

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

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

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

232 """ 

233 

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

235 self._alias_name = alias_name 

236 self._tensor_map = None 

237 if not rank_list: 

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

239 else: 

240 self._rank_list = tuple(rank_list) 

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

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

243 self._alias_tensor_map = None 

244 self._tensor_shape = None 

245 self._tensor_stride = None 

246 self._tensor_dtype = None 

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

248 init_backend=init_backend) 

249 self._placements = None 

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

251 self._ragged_shard = None 

252 self._compact_str = self._to_compact_string() 

253 

254 @classmethod 

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

256 """ 

257 Create a Layout from an existing DeviceMesh. 

258 

259 Args: 

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

261 

262 Returns: 

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

264 

265 Examples: 

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

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

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

269 """ 

270 obj = cls.__new__(cls) 

271 obj._mesh = device_mesh 

272 obj._alias_name = device_mesh.mesh_dim_names 

273 obj._rank_list = device_mesh.rank_list 

274 obj._tensor_map = None 

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

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

277 obj._alias_tensor_map = None 

278 obj._tensor_shape = None 

279 obj._tensor_stride = None 

280 obj._tensor_dtype = None 

281 obj._placements = None 

282 obj._ragged_shard = None 

283 obj._compact_str = obj._to_compact_string() 

284 return obj 

285 

286 def __call__(self, *alias_tensor_map): 

287 obj = copy.deepcopy(self) 

288 

289 # Clear the inherited partial status. 

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

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

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

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

294 

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

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

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

298 

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

300 return self._process_placement_layout(obj, alias_tensor_map) 

301 

302 return self._process_alias_layout(obj, alias_tensor_map) 

303 

304 def __deepcopy__(self, memo): 

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

306 cls = self.__class__ 

307 result = cls.__new__(cls) 

308 memo[id(self)] = result 

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

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

311 return result 

312 

313 @staticmethod 

314 def _process_placement_layout(obj, placements): 

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

316 obj.set_placements(placements) 

317 return copy.deepcopy(obj) 

318 

319 @staticmethod 

320 def _process_alias_layout(obj, alias_tensor_map): 

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

322 obj.set_alias_tensor_map(alias_tensor_map) 

323 tensor_map = () 

324 writed_map = () 

325 for ele in alias_tensor_map: 

326 if isinstance(ele, tuple): 

327 ele_map = () 

328 for item in ele: 

329 if item == "None": 

330 ele_map += (-1,) 

331 continue 

332 if item not in obj.alias_name: 

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

334 if item in writed_map: 

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

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

337 writed_map += (item,) 

338 tensor_map += (ele_map,) 

339 continue 

340 if ele == "None": 

341 tensor_map += (-1,) 

342 continue 

343 if ele not in obj.alias_name: 

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

345 if ele in writed_map: 

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

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

348 writed_map += (ele,) 

349 obj.set_tensor_map(tensor_map) 

350 obj.tensor_map_to_placement() 

351 obj.update_compact_str() 

352 return copy.deepcopy(obj) 

353 

354 def to_dict(self): 

355 """ 

356 Transform layout to a dictionary. 

357 """ 

358 if self._mesh.mesh_shape is None: 

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

360 if self._tensor_map is None: 

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

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

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

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

365 "rank_list": self._rank_list} 

366 

367 def placement_to_tensor_map(self, dim): 

368 """ 

369 Transform placement to tensor map. 

370 

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

372 Replicate, Partial) 

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

374 

375 Args: 

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

377 

378 Returns: 

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

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

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

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

383 

384 Raises: 

385 ValueError: If `dim` is negative. 

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

387 """ 

388 if dim < 0: 

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

390 if dim == 0: 

391 return self._handle_zero_dim_placement() 

392 

393 dim_map = self._build_dim_map_from_placements(dim) 

394 tensor_map = self._convert_dim_map_to_tensor_map(dim_map) 

395 self.set_tensor_map(tuple(tensor_map)) 

396 self._alias_tensor_map = self._build_readable_tensor_map() 

397 self.update_compact_str() 

398 return tensor_map 

399 

400 def _handle_zero_dim_placement(self): 

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

402 self.set_tensor_map(()) 

403 self._alias_tensor_map = () 

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

405 if isinstance(placement, Partial): 

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

407 return [] 

408 

409 def _build_dim_map_from_placements(self, dim): 

410 """Build dimension map from placements.""" 

411 dim_map = [-1] * dim 

412 self.partial_ops = {} 

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

414 if isinstance(placement, Shard): 

415 shard_dim = placement.dim 

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

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

418 if shard_dim < 0: 

419 shard_dim += dim 

420 if dim_map[shard_dim] == -1: 

421 dim_map[shard_dim] = [mesh_idx] 

422 else: 

423 dim_map[shard_dim].append(mesh_idx) 

424 elif isinstance(placement, Partial): 

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

426 self._validate_strided_shard_split_factor(dim_map) 

427 self._reorder_dim_map_for_strided_shard(dim_map) 

428 return dim_map 

429 

430 @staticmethod 

431 def _placement_split_factor(placement): 

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

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

434 

435 @staticmethod 

436 def _build_order_positions(shard_order): 

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

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

439 

440 def _compute_expected_split_factors(self, shard_axes, shard_order): 

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

442 order_positions = self._build_order_positions(shard_order) 

443 expected_split_factors = {} 

444 for mesh_idx in shard_axes: 

445 split_factor = 1 

446 for right_mesh_idx in shard_axes: 

447 if right_mesh_idx <= mesh_idx: 

448 continue 

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

450 split_factor *= self.mesh_shape[right_mesh_idx] 

451 expected_split_factors[mesh_idx] = split_factor 

452 return expected_split_factors 

453 

454 def _get_effective_shard_axes(self, shard_axes): 

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

456 return sorted( 

457 shard_axes, 

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

459 ) 

460 

461 def _reorder_dim_map_for_strided_shard(self, dim_map): 

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

463 for i, shard_axes in enumerate(dim_map): 

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

465 continue 

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

467 

468 def _validate_strided_shard_split_factor(self, dim_map): 

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

470 for shard_axes in dim_map: 

471 if shard_axes == -1: 

472 continue 

473 shard_order = self._get_effective_shard_axes(shard_axes) 

474 expected_split_factors = self._compute_expected_split_factors( 

475 shard_axes, shard_order 

476 ) 

477 for mesh_idx in shard_axes: 

478 placement = self.placements[mesh_idx] 

479 actual_split_factor = self._placement_split_factor(placement) 

480 expected_split_factor = expected_split_factors[mesh_idx] 

481 if actual_split_factor != expected_split_factor: 

482 raise ValueError( 

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

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

485 ) 

486 

487 @staticmethod 

488 def _extract_reduce_op(placement): 

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

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

491 if isinstance(op_name, str): 

492 op_name = op_name.lower() 

493 return op_name 

494 

495 def _convert_dim_map_to_tensor_map(self, dim_map): 

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

497 device_dim_count = len(self.mesh_shape) 

498 tensor_map = [] 

499 for mesh_idx in dim_map: 

500 if mesh_idx == -1: 

501 tensor_map.append(-1) 

502 continue 

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

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

505 return tensor_map 

506 

507 def _build_readable_tensor_map(self): 

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

509 mesh_dim_names = self._mesh.mesh_dim_names 

510 has_names = mesh_dim_names is not None 

511 

512 def _map_dim(dim): 

513 """convert dimension index to dimension name.""" 

514 if dim == -1: 

515 return "None" 

516 if not has_names: 

517 return f"dim_{dim}" 

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

519 

520 readable_map = [] 

521 for item in self._tensor_map: 

522 if isinstance(item, tuple): 

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

524 readable_map.append(mapped_tuple) 

525 else: 

526 readable_map.append(_map_dim(item)) 

527 return tuple(readable_map) 

528 

529 def tensor_map_to_placement(self): 

530 """ 

531 Transform tensor map to placement. 

532 

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

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

535 `placement_to_tensor_map`. 

536 

537 Returns: 

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

539 dimension of the device mesh. 

540 

541 Raises: 

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

543 """ 

544 if self._tensor_map is None: 

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

546 uneven_shard_dims = { 

547 mesh_idx: placement.dim 

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

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

550 } 

551 mesh_ndim = len(self.mesh_shape) 

552 placements = [Replicate()] * mesh_ndim 

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

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

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

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

557 shard_axes = sorted(mesh_indices) 

558 expected_split_factors = self._compute_expected_split_factors( 

559 shard_axes, mesh_indices 

560 ) 

561 for mesh_idx in shard_axes: 

562 split_factor = expected_split_factors[mesh_idx] 

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

564 placement = ( 

565 StridedShard( 

566 dim=tensor_dim, 

567 split_factor=split_factor, 

568 uneven_shard=uneven_shard, 

569 ) 

570 if split_factor > 1 

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

572 ) 

573 placements[mesh_idx] = placement 

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

575 if op is not None: 

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

577 if self._ragged_shard is not None: 

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

579 self.set_placements(placements) 

580 self._alias_tensor_map = self._build_readable_tensor_map() 

581 self.update_compact_str() 

582 return placements 

583 

584 def __setstate__(self, state): 

585 self.__dict__.update(state) 

586 self.update_mesh(init_backend=False) 

587 

588 @property 

589 def mesh(self): 

590 """ 

591 Get the device mesh associated with this layout. 

592 

593 Returns: 

594 DeviceMesh: The device mesh describing the device topology. 

595 """ 

596 return self._mesh 

597 

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

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

600 

601 Args: 

602 init_backend (bool): Whether to initialize communication backend 

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

604 avoid creating process groups with a stale rank_list from the 

605 sender side. Default ``True``. 

606 """ 

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

608 rank_list=self.rank_list, init_backend=init_backend) 

609 

610 @property 

611 def rank_list(self): 

612 """ 

613 Get the list of ranks participating in this layout. 

614 

615 Returns: 

616 tuple[int]: The rank list. 

617 """ 

618 return self._rank_list 

619 

620 @rank_list.setter 

621 def rank_list(self, val): 

622 self._rank_list = val 

623 

624 @property 

625 def mesh_shape(self): 

626 """mesh shape""" 

627 return self._mesh.mesh_shape 

628 

629 @property 

630 def alias_name(self): 

631 """alias name""" 

632 return self._mesh.mesh_dim_names 

633 

634 @property 

635 def alias_tensor_map(self): 

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

637 return self._alias_tensor_map 

638 

639 @property 

640 def alias_placements(self): 

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

642 

643 alias_tensor_map preserves multi-axis ordering information 

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

645 that Placement objects cannot represent, since both map to 

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

647 

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

649 also carry Partial status which alias_tensor_map cannot encode. 

650 

651 Use this property when constructing DTensors from an existing Layout 

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

653 """ 

654 if self._ragged_shard is not None: 

655 return self._placements 

656 if self._alias_tensor_map is not None and any( 

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

658 ): 

659 return self._alias_tensor_map 

660 return self._placements 

661 

662 def set_alias_tensor_map(self, alias_tensor_map): 

663 """Set alias_tensor_map""" 

664 self._alias_tensor_map = alias_tensor_map 

665 

666 @property 

667 def placements(self): 

668 """placements""" 

669 return self._placements 

670 

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

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

673 self._placements = placements 

674 self._ragged_shard = ( 

675 None if placements is None else _extract_ragged_shard(placements) 

676 ) 

677 

678 @property 

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

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

681 return self._tensor_shape 

682 

683 @property 

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

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

686 return self._tensor_stride 

687 

688 @property 

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

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

691 return self._tensor_dtype 

692 

693 @property 

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

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

696 return tuple( 

697 mesh_dim 

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

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

700 ) 

701 

702 @property 

703 def has_uneven_shard(self) -> bool: 

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

705 return bool(self.uneven_shard_mesh_dims) 

706 

707 def set_tensor_meta( 

708 self, 

709 shape: Sequence[int], 

710 stride: Sequence[int], 

711 dtype: Any, 

712 ) -> None: 

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

714 

715 Args: 

716 shape: Logical global tensor shape. 

717 stride: Logical global tensor stride. 

718 dtype: Logical tensor dtype. 

719 """ 

720 self._tensor_shape = tuple(shape) 

721 self._tensor_stride = tuple(stride) 

722 self._tensor_dtype = dtype 

723 self.update_compact_str() 

724 

725 @property 

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

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

728 if self._placements is None: 

729 return None 

730 return _replace_ragged_with_replicate(self._placements) 

731 

732 @property 

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

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

735 return self._ragged_shard 

736 

737 @property 

738 def tensor_map(self): 

739 """tensor map""" 

740 return self._tensor_map 

741 

742 def set_tensor_map(self, tensor_map): 

743 """Set tensor_map.""" 

744 self._tensor_map = tensor_map 

745 

746 @property 

747 def partial(self): 

748 """partial status""" 

749 return self._partial 

750 

751 def set_partial_by_dev_axis(self, axis, op): 

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

753 if op not in self._support_partial_op: 

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

755 if self.is_dev_axis_apply_shard(axis): 

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

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

758 self.tensor_map_to_placement() 

759 self.update_compact_str() 

760 

761 def get_partial_by_dev_id(self, axis): 

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

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

764 

765 def is_dev_axis_apply_shard(self, axis): 

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

767 axis_id = self._mesh.axis_id(axis) 

768 

769 def flatten(input_x): 

770 flatten_res = [] 

771 for item in input_x: 

772 if isinstance(item, tuple): 

773 flatten_res.extend(flatten(item)) 

774 else: 

775 flatten_res.append(item) 

776 return flatten_res 

777 

778 flatten_tensor_map = flatten(self.tensor_map) 

779 return axis_id in flatten_tensor_map 

780 

781 def get_dev_axis_apply_shard_axis(self, axis): 

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

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

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

785 return dim 

786 return None 

787 

788 def reset_partial(self): 

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

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

791 self.tensor_map_to_placement() 

792 self.update_compact_str() 

793 

794 def is_partial(self): 

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

796 return any(self.partial) 

797 

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

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

800 

801 Args: 

802 tensor_dim: Tensor dimension index to check. 

803 

804 Returns: 

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

806 """ 

807 alias_tm = self.alias_tensor_map 

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

809 return 1 

810 dim_entry = alias_tm[tensor_dim] 

811 if dim_entry == 'None': 

812 return 1 

813 if isinstance(dim_entry, str): 

814 return self.mesh.get_device_num_along_axis(dim_entry) 

815 if isinstance(dim_entry, tuple): 

816 total = 1 

817 for axis in dim_entry: 

818 if axis != 'None': 

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

820 return total 

821 return 1 

822 

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

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

825 

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

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

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

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

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

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

832 

833 Args: 

834 tensor_dim: Tensor dimension index to query. 

835 

836 Returns: 

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

838 """ 

839 alias_tm = self.alias_tensor_map 

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

841 return 0 

842 dim_entry = alias_tm[tensor_dim] 

843 if dim_entry == 'None': 

844 return 0 

845 rank = platform.get_rank() 

846 if isinstance(dim_entry, tuple): 

847 non_none = [ax for ax in dim_entry if ax != 'None'] 

848 if not non_none: 

849 return 0 

850 global_id = 0 

851 for ax in non_none: 

852 rank_list = self.mesh.get_rank_list_along_axis(ax) 

853 local_id = rank_list.index(rank) if rank in rank_list else 0 

854 ax_size = self.mesh.get_device_num_along_axis(ax) 

855 global_id = global_id * ax_size + local_id