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 

20import torch 

21import torch.distributed as dist 

22 

23from hyper_parallel.core.dtensor.device_mesh import DeviceMesh 

24 

25Tensor = torch.Tensor 

26 

27 

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

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

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

31 return 

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

33 mesh._mesh_shape, 

34 mesh.mesh_dim_names, 

35 mesh._rank_list, 

36 ) 

37 

38 

39def mesh_scatter( 

40 output: Tensor, 

41 scatter_list: Sequence[Tensor], 

42 mesh: DeviceMesh, 

43 mesh_dim: int, 

44 *, 

45 group_src: int = 0, 

46) -> Tensor: 

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

48 _ensure_mesh_process_groups(mesh) 

49 group = mesh.get_group(mesh_dim) 

50 contiguous_list = [ 

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

52 for chunk in scatter_list 

53 ] 

54 src = dist.get_global_rank(group, group_src) 

55 if dist.get_group_rank(group, dist.get_rank()) == group_src: 

56 dist.scatter(output, list(contiguous_list), src=src, group=group) 

57 else: 

58 dist.scatter(output, None, src=src, group=group) 

59 return output 

60 

61 

62def mesh_scatter_ragged( 

63 output: Tensor, 

64 scatter_list: Optional[Sequence[Tensor]], 

65 mesh: DeviceMesh, 

66 mesh_dim: int, 

67 *, 

68 group_src: int = 0, 

69) -> Tensor: 

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

71 

72 Args: 

73 output: Preallocated receive buffer for the current rank. 

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

75 may pass ``None``. 

76 mesh: Device mesh containing the communication group. 

77 mesh_dim: Mesh dimension along which to scatter. 

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

79 

80 Returns: 

81 The populated current-rank output buffer. 

82 

83 Raises: 

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

85 """ 

86 _ensure_mesh_process_groups(mesh) 

87 group = mesh.get_group(mesh_dim) 

88 group_size = mesh.size(mesh_dim) 

89 if group_src < 0 or group_src >= group_size: 

90 raise ValueError( 

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

92 ) 

93 

94 group_rank = dist.get_group_rank(group, dist.get_rank()) 

95 source_global_rank = dist.get_global_rank(group, group_src) 

96 if group_rank == group_src: 

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

98 raise ValueError( 

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

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

101 ) 

102 output.copy_(scatter_list[group_src]) 

103 works = [] 

104 for destination_group_rank, chunk in enumerate(scatter_list): 

105 if destination_group_rank == group_src: 

106 continue 

107 destination_global_rank = dist.get_global_rank( 

108 group, destination_group_rank 

109 ) 

110 works.append( 

111 dist.isend( 

112 chunk.contiguous(), 

113 dst=destination_global_rank, 

114 group=group, 

115 ) 

116 ) 

117 for work in works: 

118 work.wait() 

119 return output 

120 

121 work = dist.irecv( 

122 output, 

123 src=source_global_rank, 

124 group=group, 

125 ) 

126 work.wait() 

127 return output 

128 

129 

130def mesh_broadcast( 

131 tensor: Tensor, 

132 mesh: DeviceMesh, 

133 mesh_dim: int, 

134 *, 

135 group_src: int = 0, 

136) -> Tensor: 

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

138 _ensure_mesh_process_groups(mesh) 

139 group = mesh.get_group(mesh_dim) 

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

141 tensor = tensor.contiguous() 

142 dist.broadcast(tensor, src=dist.get_global_rank(group, group_src), group=group) 

143 return tensor