Coverage for  / home / jenkins / .local / lib / python3.10 / site-packages / hyper_parallel / core / utils / shape_utils.py: 100%

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

16Utility functions for distributed tensor operations. 

17 

18This module provides helper functions for computing local shapes, global offsets, 

19and other layout-related calculations in distributed settings. 

20""" 

21from typing import Sequence 

22 

23from hyper_parallel.core.dtensor.layout import ( 

24 Layout, 

25 infer_balanced_chunk_range, 

26 infer_ceil_chunk_range, 

27) 

28 

29 

30def compute_local_shape_and_global_offset(global_shape, device_mesh, placement): 

31 """ 

32 Compute local shard shape and its global offset. 

33 

34 Args: 

35 global_shape: Shape of the global tensor. 

36 device_mesh: Device mesh for distributed execution. 

37 placement: Sharding placements for each mesh dimension. Supports 

38 Placement objects or alias strings. 

39 

40 Returns: 

41 The local shape owned by the current rank. 

42 """ 

43 from hyper_parallel.core.dtensor.dtensor import _is_alias_placements # pylint: disable=C0415 

44 total_layout = Layout.from_device_mesh(device_mesh) 

45 if _is_alias_placements(placement): 

46 layout = total_layout(*placement) 

47 else: 

48 layout = total_layout(placement) 

49 layout.placement_to_tensor_map(len(global_shape)) 

50 local_shape = list(global_shape) 

51 for tensor_dim, mapped_axes in enumerate(layout.alias_tensor_map): 

52 if isinstance(mapped_axes, str): 

53 mapped_axes = (mapped_axes,) 

54 for mapped_axis in mapped_axes: 

55 if mapped_axis == "None": 

56 continue 

57 shard_count = layout.mesh.get_device_num_along_axis(mapped_axis) 

58 if local_shape[tensor_dim] % shard_count == 0: 

59 local_shape[tensor_dim] //= shard_count 

60 continue 

61 chunk_start, chunk_end = infer_balanced_chunk_range( 

62 local_shape[tensor_dim], 

63 shard_count, 

64 layout.mesh.get_local_rank(mapped_axis), 

65 ) 

66 local_shape[tensor_dim] = chunk_end - chunk_start 

67 return local_shape 

68 

69 

70def compute_local_shape_and_global_offset_by_ceil_chunk( 

71 global_shape: Sequence[int], 

72 shard_dim: int, 

73 shard_count: int, 

74 shard_rank: int, 

75) -> tuple[list[int], list[int]]: 

76 """Return one FSDP local shape and offset using ceil-chunk geometry. 

77 

78 Unlike balanced Shard geometry, ceil-chunk keeps a fixed maximum chunk 

79 size and represents ranks beyond the last chunk with an empty shard. 

80 

81 Args: 

82 global_shape: Shape before applying the FSDP shard. 

83 shard_dim: Tensor dimension partitioned by FSDP. 

84 shard_count: Number of ranks in the FSDP shard mesh. 

85 shard_rank: Current rank within the FSDP shard mesh. 

86 

87 Returns: 

88 The local shape and its global offset relative to ``global_shape``. 

89 """ 

90 local_shape = list(global_shape) 

91 global_offset = [0] * len(local_shape) 

92 chunk_start, chunk_end = infer_ceil_chunk_range( 

93 local_shape[shard_dim], 

94 shard_count, 

95 shard_rank, 

96 ) 

97 local_shape[shard_dim] = chunk_end - chunk_start 

98 global_offset[shard_dim] = chunk_start 

99 return local_shape, global_offset