Coverage for  / home / jenkins / .local / lib / python3.10 / site-packages / hyper_parallel / core / pipeline_parallel / _p2p.py: 97%

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"""Pipeline edge groups and batched P2P communicator initialization.""" 

16import torch.distributed as dist 

17 

18_P2P_MULTI_STREAM_GROUPS = {} 

19 

20 

21def _build_p2p_edge_rank_lists(pp_rank_list: list[int], include_wrap: bool = False) -> list[tuple[int, int]]: 

22 """Build normalized two-rank groups for adjacent pipeline ranks.""" 

23 if not isinstance(pp_rank_list, (list, tuple)): 

24 raise ValueError( 

25 f"pp_rank_list must be a list or tuple of integer ranks, but got {type(pp_rank_list)}." 

26 ) 

27 if any(not isinstance(rank, int) or isinstance(rank, bool) for rank in pp_rank_list): 

28 raise ValueError(f"pp_rank_list must contain only integer ranks, but got {pp_rank_list}.") 

29 if len(set(pp_rank_list)) != len(pp_rank_list): 

30 raise ValueError(f"pp_rank_list must not contain duplicate ranks, but got {pp_rank_list}.") 

31 if len(pp_rank_list) < 2: 

32 return [] 

33 

34 edge_rank_lists = { 

35 tuple(sorted((src_rank, dst_rank))) 

36 for src_rank, dst_rank in zip(pp_rank_list, pp_rank_list[1:]) 

37 } 

38 if include_wrap and len(pp_rank_list) > 2: 

39 edge_rank_lists.add(tuple(sorted((pp_rank_list[-1], pp_rank_list[0])))) 

40 return sorted(edge_rank_lists) 

41 

42 

43def create_p2p_multi_stream_groups( 

44 pp_rank_list: list[int], 

45 include_wrap: bool = False, 

46) -> dict[int, dist.ProcessGroup]: 

47 """Create adjacent two-rank PP groups for independent communication streams. 

48 

49 Args: 

50 pp_rank_list: Ordered global ranks in one pipeline-parallel group. 

51 include_wrap: Whether to include the last-to-first interleaved edge. 

52 

53 Returns: 

54 A mapping from adjacent peer rank to its process-group handle. 

55 """ 

56 current_rank = dist.get_rank() 

57 world_size = dist.get_world_size() 

58 gathered_pp_rank_lists = [None] * world_size 

59 dist.all_gather_object(gathered_pp_rank_lists, list(pp_rank_list)) 

60 edge_rank_lists = sorted({ 

61 edge_ranks 

62 for ranks in gathered_pp_rank_lists 

63 for edge_ranks in _build_p2p_edge_rank_lists(ranks, include_wrap) 

64 }) 

65 

66 # Overlapping groups can deadlock with local synchronization when 

67 # neighboring ranks enter different edge creations first. Every rank 

68 # therefore creates the global edge set in the same order. 

69 local_groups = {} 

70 world_rank_list = tuple(range(world_size)) 

71 for edge_ranks in edge_rank_lists: 

72 group_key = str(edge_ranks) 

73 group = _P2P_MULTI_STREAM_GROUPS.get(group_key) 

74 if group is None: 

75 group = ( 

76 dist.group.WORLD 

77 if edge_ranks == world_rank_list 

78 else dist.new_group(ranks=list(edge_ranks)) 

79 ) 

80 _P2P_MULTI_STREAM_GROUPS[group_key] = group 

81 if current_rank not in edge_ranks: 

82 continue 

83 peer_rank = edge_ranks[0] if edge_ranks[1] == current_rank else edge_ranks[1] 

84 local_groups[peer_rank] = group 

85 return local_groups 

86 

87 

88def prepare_batch_p2p_group(group: dist.ProcessGroup = None) -> None: 

89 """Synchronize a group before its first subset batched P2P call. 

90 

91 PyTorch requires every rank in a process group to participate when 

92 ``batch_isend_irecv`` is the first collective on that group. A barrier 

93 at the common pipeline run boundary initializes the communicator 

94 before ranks reach peer operations at different times. 

95 

96 Args: 

97 group: The process group used by the batched P2P operations. 

98 ``None`` uses the default group. 

99 """ 

100 dist.barrier(group=group)