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

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"""Mesh-scoped collectives for :func:`distribute_tensor` (PyTorch DTensor parity).""" 

16from __future__ import annotations 

17 

18from typing import Optional, Sequence 

19 

20from hyper_parallel.core.dtensor.device_mesh import DeviceMesh 

21from hyper_parallel.platform import get_platform 

22 

23platform = get_platform() 

24Tensor = platform.Tensor 

25 

26 

27def _ensure_mesh_process_groups(mesh: DeviceMesh) -> None: 

28 """Lazily create per-axis process groups when mesh was built with ``init_backend=False``.""" 

29 if hasattr(mesh, "_dim_group_names") and mesh._dim_group_names is not None: 

30 return 

31 mesh._dim_group_names = DeviceMesh._init_process_groups( # pylint: disable=protected-access 

32 mesh._mesh_shape, 

33 mesh.mesh_dim_names, 

34 mesh._rank_list, 

35 ) 

36 

37 

38def mesh_scatter( 

39 output: Tensor, 

40 scatter_list: Sequence[Tensor], 

41 mesh: DeviceMesh, 

42 mesh_dim: int, 

43 *, 

44 group_src: int = 0, 

45) -> Tensor: 

46 """Scatter tensor chunks along one mesh dimension (PyTorch ``mesh_scatter`` parity).""" 

47 _ensure_mesh_process_groups(mesh) 

48 group = mesh.get_group(mesh_dim) 

49 contiguous_list = [ 

50 chunk.contiguous() if hasattr(chunk, "is_contiguous") and not chunk.is_contiguous() else chunk 

51 for chunk in scatter_list 

52 ] 

53 if platform.get_group_rank(group) == group_src: 

54 platform.scatter(output, list(contiguous_list), group=group, group_src=group_src) 

55 else: 

56 platform.scatter(output, None, group=group, group_src=group_src) 

57 return output 

58 

59 

60def mesh_scatter_ragged( 

61 output: Tensor, 

62 scatter_list: Optional[Sequence[Tensor]], 

63 mesh: DeviceMesh, 

64 mesh_dim: int, 

65 *, 

66 group_src: int = 0, 

67) -> Tensor: 

68 """Scatter variable-length flat tensors with point-to-point communication. 

69 

70 Args: 

71 output: Preallocated receive buffer for the current rank. 

72 scatter_list: Source-rank tensors ordered by group rank. Non-source ranks 

73 may pass ``None``. 

74 mesh: Device mesh containing the communication group. 

75 mesh_dim: Mesh dimension along which to scatter. 

76 group_src: Source rank relative to the mesh-dimension group. 

77 

78 Returns: 

79 The populated current-rank output buffer. 

80 

81 Raises: 

82 ValueError: If the source rank or source scatter list is invalid. 

83 """ 

84 _ensure_mesh_process_groups(mesh) 

85 group = mesh.get_group(mesh_dim) 

86 group_size = mesh.size(mesh_dim) 

87 if group_src < 0 or group_src >= group_size: 

88 raise ValueError( 

89 f"group_src must be in [0, {group_size}), but got {group_src}" 

90 ) 

91 

92 group_rank = platform.get_group_rank(group) 

93 source_global_rank = platform.get_global_rank(group, group_src) 

94 if group_rank == group_src: 

95 if scatter_list is None or len(scatter_list) != group_size: 

96 raise ValueError( 

97 "source scatter_list length must equal the mesh dimension size, " 

98 f"got scatter_list={scatter_list!r}, group_size={group_size}" 

99 ) 

100 output.copy_(scatter_list[group_src]) 

101 works = [] 

102 for destination_group_rank, chunk in enumerate(scatter_list): 

103 if destination_group_rank == group_src: 

104 continue 

105 destination_global_rank = platform.get_global_rank( 

106 group, destination_group_rank 

107 ) 

108 works.append( 

109 platform.isend( 

110 chunk.contiguous(), 

111 dst=destination_global_rank, 

112 group=group, 

113 ) 

114 ) 

115 for work in works: 

116 work.wait() 

117 return output 

118 

119 work = platform.irecv( 

120 output, 

121 src=source_global_rank, 

122 group=group, 

123 ) 

124 work.wait() 

125 return output 

126 

127 

128def mesh_broadcast( 

129 tensor: Tensor, 

130 mesh: DeviceMesh, 

131 mesh_dim: int, 

132 *, 

133 group_src: int = 0, 

134) -> Tensor: 

135 """Broadcast a tensor along one mesh dimension (PyTorch ``mesh_broadcast`` parity).""" 

136 _ensure_mesh_process_groups(mesh) 

137 group = mesh.get_group(mesh_dim) 

138 if hasattr(tensor, "is_contiguous") and not tensor.is_contiguous(): 

139 tensor = tensor.contiguous() 

140 platform.broadcast(tensor, group=group, group_src=group_src) 

141 return tensor