Coverage for  / home / jenkins / .local / lib / python3.10 / site-packages / hyper_parallel / distributed_data / locality.py: 92%

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"""Create independent node-local data exchange groups.""" 

16# This distributed-data package is intentionally PyTorch-only. 

17 

18from __future__ import annotations 

19 

20import os 

21import socket 

22from typing import Any 

23 

24import torch.distributed as dist # pylint: disable=forbidden-backend-import 

25 

26from hyper_parallel.distributed_data.topology import DataTopology 

27from hyper_parallel.distributed_data.transport import DataGroups, all_gather_control_object 

28 

29 

30def _node_rank_groups( 

31 topology: DataTopology, 

32 node_ids: dict[int, str], 

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

34 buckets: dict[str, list[int]] = {} 

35 for rank in topology.rank_list: 

36 buckets.setdefault(node_ids[rank], []).append(rank) 

37 return tuple(sorted(tuple(sorted(ranks)) for ranks in buckets.values())) 

38 

39 

40def _fixed_rank_groups( 

41 topology: DataTopology, 

42 group_size: int, 

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

44 """Partition the root mesh rank order into fixed-size balancing groups. 

45 

46 Args: 

47 topology: Root mesh topology whose rank order defines the groups. 

48 group_size: Number of ranks in every balancing group. 

49 

50 Returns: 

51 Fixed-size rank groups in root mesh order. 

52 

53 Raises: 

54 ValueError: If the mesh size is not divisible by ``group_size``. 

55 """ 

56 rank_list = topology.rank_list 

57 if len(rank_list) % group_size: 

58 raise ValueError( 

59 f"balance_group_size={group_size} must divide the root mesh size {len(rank_list)}." 

60 ) 

61 return tuple( 

62 tuple(rank_list[start:start + group_size]) 

63 for start in range(0, len(rank_list), group_size) 

64 ) 

65 

66 

67def _resolve_local_topology( 

68 mesh: Any, 

69 rank: int, 

70 world_size: int, 

71 dp_dim_names: tuple[str, ...] | None, 

72) -> DataTopology: 

73 """Resolve the WORLD-covering topology for independent DP readers.""" 

74 topology = DataTopology.from_mesh(mesh, global_rank=rank, dp_dim_names=dp_dim_names) 

75 if set(topology.rank_list) != set(range(world_size)): 

76 raise ValueError("Local balancing requires a root mesh containing every WORLD rank.") 

77 if any(len(group) != 1 for group in topology.model_parallel_rank_groups): 

78 raise ValueError("Local balancing currently requires CP/TP/PP/Ulysses and other model parallel sizes = 1.") 

79 return topology 

80 

81 

82def _validate_communication_config(backend: str, device: Any, distributed: bool) -> None: 

83 """Validate the node-local collective backend and its rank-local device.""" 

84 if backend not in ("gloo", "hccl"): 

85 raise ValueError("communication_backend must be 'gloo' or 'hccl'.") 

86 if backend != "hccl": 

87 return 

88 if distributed and (device is None or getattr(device, "type", None) != "npu"): 

89 raise ValueError("communication_backend='hccl' requires an NPU communication_device.") 

90 

91 

92def _gather_startup_statuses( 

93 status: Any, 

94 *, 

95 distributed: bool, 

96 backend: str, 

97 device: Any, 

98) -> list[Any]: 

99 """Gather startup state on an explicit Gloo group or HCCL WORLD.""" 

100 if not distributed: 

101 return [status] 

102 startup_group = None 

103 if backend == "gloo": 

104 startup_group = dist.new_group(ranks=list(range(dist.get_world_size())), backend="gloo") 

105 return list(all_gather_control_object(status, group=startup_group, device=device, backend=backend)) 

106 

107 

108def _create_locality_groups( 

109 mesh: Any, 

110 *, 

111 dp_dim_names: tuple[str, ...] | None = None, 

112 balance_group_size: int | None = None, 

113 communication_backend: str = "hccl", 

114 communication_device: Any = None, 

115) -> tuple[DataTopology, DataGroups]: 

116 """Resolve locality once and create all process groups in WORLD rank order. 

117 

118 Args: 

119 mesh: Named root mesh covering WORLD, with no nontrivial model parallelism. 

120 dp_dim_names: Root mesh data-parallel dimensions. 

121 balance_group_size: Optional fixed number of ranks per balancing group. 

122 

123 Returns: 

124 Training topology and this rank's independent data exchange groups. 

125 

126 Note: 

127 WORLD communication is used only during startup. No training mesh or 

128 FSDP process group is modified. All ranks must call this factory together. 

129 """ 

130 distributed = dist.is_available() and dist.is_initialized() 

131 _validate_communication_config(communication_backend, communication_device, distributed) 

132 rank = dist.get_rank() if distributed else 0 

133 world_size = dist.get_world_size() if distributed else 1 

134 topology = _resolve_local_topology(mesh, rank, world_size, dp_dim_names) 

135 if balance_group_size is None: 

136 selected_node = os.environ.get("GROUP_RANK") or socket.gethostname() 

137 statuses = _gather_startup_statuses( 

138 (rank, selected_node), 

139 distributed=distributed, 

140 backend=communication_backend, 

141 device=communication_device, 

142 ) 

143 rank_groups = _node_rank_groups(topology, dict(statuses)) 

144 else: 

145 rank_groups = _fixed_rank_groups(topology, balance_group_size) 

146 own_groups = None 

147 for ranks in rank_groups: 

148 control_group = None 

149 payload_group = None 

150 if distributed and len(ranks) > 1: 

151 control_group = dist.new_group(ranks=list(ranks), backend=communication_backend) 

152 payload_group = control_group 

153 if communication_backend == "hccl": 

154 payload_group = dist.new_group(ranks=list(ranks), backend=communication_backend) 

155 if rank in ranks: 

156 own_groups = DataGroups(ranks, control_group, payload_group, None, min(ranks), distributed) 

157 return topology, own_groups