Coverage for  / home / jenkins / .local / lib / python3.10 / site-packages / hyper_parallel / trainer / runtime / distributed.py: 68%

1# Copyright 2025 Bytedance Ltd. and/or its affiliates 

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"""Collective and process-group lifecycle helpers for the trainer runtime. 

15 

16``get_world_size_safe`` / ``get_global_rank_safe`` / ``get_local_rank_safe`` 

17merged in from the former ``components/distributed/init_utils.py`` in stage 7 

18(05 §15.11 step 3); ``initialize_distributed`` / 

19``create_distributed_setup_from_config`` / ``destroy_process_group`` come from 

20the former ``components/distributed/infrastructure.py`` (05 §15.11 step 3, 

21plan §1195-1196). The mesh objects themselves are AutoModels-side in 

22``hyper_parallel.distributed.mesh``. 

23""" 

24 

25__all__ = ["all_gather", "all_reduce"] 

26 

27import logging 

28import os 

29from typing import TYPE_CHECKING, Any, List, Literal, Optional, Union 

30 

31import torch 

32from torch import distributed as dist 

33 

34from hyper_parallel.distributed.mesh import ( 

35 DistributedSetup, 

36 MeshContext, 

37 ParallelTopologySpec, 

38 resolve_parallel_layout, 

39) 

40from hyper_parallel.trainer.runtime.device import get_device_type, get_torch_device 

41 

42logger = logging.getLogger(__name__) 

43 

44 

45if TYPE_CHECKING: 

46 from torch.distributed import ProcessGroup 

47 

48 

49def all_gather(tensor: "torch.Tensor", world_size: int) -> "torch.Tensor": 

50 """ 

51 Gathers the tensor from all ranks and concats them along the first dim. 

52 """ 

53 output_tensor = torch.empty(world_size * tensor.numel(), dtype=tensor.dtype, device=get_device_type()) 

54 dist.all_gather_into_tensor(output_tensor, tensor) 

55 return output_tensor.view(-1, *tensor.size()[1:]) 

56 

57 

58def all_reduce( 

59 data: Union[int, float, List[Union[int, float]], "torch.Tensor"], 

60 op: Literal["mean", "sum", "max", "min"] = "mean", 

61 group: Optional["ProcessGroup"] = None, 

62) -> Union[int, float, List[Union[int, float]]]: 

63 """ 

64 Performs all reduce in the given process group. 

65 """ 

66 if not dist.is_initialized(): 

67 raise RuntimeError("Distributed environment is not initialized.") 

68 

69 if not isinstance(data, torch.Tensor): 

70 data = torch.tensor(data, dtype=torch.float, device=get_device_type()) 

71 

72 reduce_ops = { 

73 "mean": dist.ReduceOp.SUM, 

74 "sum": dist.ReduceOp.SUM, 

75 "max": dist.ReduceOp.MAX, 

76 "min": dist.ReduceOp.MIN, 

77 } 

78 reduce_op = reduce_ops.get(op) 

79 if reduce_op is None: 

80 raise ValueError("op must be one of: mean, sum, max, min") 

81 dist.all_reduce(data, op=reduce_op, group=group) 

82 if op == "mean": # ReduceOp.AVG is not supported by the NPU backend 

83 data /= dist.get_world_size(group=group) 

84 

85 if data.numel() == 1: 

86 return data.item() 

87 return data.tolist() 

88 

89 

90def get_world_size_safe() -> int: 

91 """Return dist.get_world_size() if initialized, else 1.""" 

92 if dist.is_initialized(): 

93 return dist.get_world_size() 

94 return 1 

95 

96 

97def get_global_rank_safe() -> int: 

98 """Return dist.get_rank() if initialized, else 0.""" 

99 if dist.is_initialized(): 

100 return dist.get_rank() 

101 return 0 

102 

103 

104def get_local_rank_safe() -> int: 

105 """Return dist.get_node_local_rank() if initialized, else 0.""" 

106 if dist.is_initialized(): 

107 return dist.get_node_local_rank() 

108 return 0 

109 

110 

111def initialize_distributed(backend: str = "nccl") -> Any: 

112 """Initialize torch.distributed process group. 

113 

114 Calls dist.init_process_group if not already initialized. 

115 Falls back to gloo when the requested backend is unavailable (e.g. CPU). 

116 Sets the current accelerator device from ``LOCAL_RANK`` before process 

117 group initialization. 

118 """ 

119 device_type = get_device_type() 

120 local_rank = int(os.environ.get("LOCAL_RANK", 0)) 

121 if device_type in ("cuda", "npu"): 

122 get_torch_device().set_device(local_rank) 

123 

124 if not dist.is_initialized(): 

125 effective_backend = backend 

126 if backend == "nccl" and device_type == "npu": 

127 effective_backend = "hccl" 

128 logger.info("Using HCCL instead of NCCL on the NPU backend.") 

129 elif backend == "nccl" and device_type == "cpu": 

130 effective_backend = "gloo" 

131 logger.warning( 

132 "CUDA not available; falling back to '%s' backend for distributed.", 

133 effective_backend, 

134 ) 

135 dist.init_process_group(backend=effective_backend) 

136 

137 return dist 

138 

139 

140def _populate_mesh_ranks(mesh_context: MeshContext, dp_shard_size: int) -> None: 

141 """Populate local mesh ranks after device-mesh construction.""" 

142 def _local_rank(dim: str) -> int: 

143 if ( 

144 mesh_context.device_mesh is None 

145 or dim not in (mesh_context.device_mesh.mesh_dim_names or ()) 

146 ): 

147 return 0 

148 return mesh_context.device_mesh.get_local_rank(dim) 

149 

150 device_mesh_dim_names = ( 

151 mesh_context.device_mesh.mesh_dim_names 

152 if mesh_context.device_mesh is not None 

153 else () 

154 ) 

155 if "dp" in device_mesh_dim_names: 

156 mesh_context.dp_rank = _local_rank("dp") 

157 else: 

158 dp_shard_rank = _local_rank("dp_shard") 

159 dp_replicate_rank = _local_rank("dp_replicate") 

160 mesh_context.dp_rank = dp_replicate_rank * dp_shard_size + dp_shard_rank 

161 mesh_context.tp_rank = _local_rank("tp") 

162 mesh_context.cp_rank = _local_rank("cp") 

163 if mesh_context.pp_size > 1 and dist.is_initialized(): 

164 stage_world_size = mesh_context.dp_size * mesh_context.cp_size * mesh_context.tp_size 

165 mesh_context.pp_rank = dist.get_rank() // stage_world_size 

166 else: 

167 mesh_context.pp_rank = 0 

168 mesh_context.ep_rank = ( 

169 mesh_context.fsdp_moe_mesh.get_local_rank("ep") 

170 if mesh_context.fsdp_moe_mesh is not None 

171 else 0 

172 ) 

173 

174 

175def create_distributed_setup_from_config(cfg: Any) -> DistributedSetup: 

176 """Create DistributedSetup and build the configured mesh domains.""" 

177 accel = cfg.accelerator if cfg is not None and hasattr(cfg, "accelerator") else None 

178 if accel is None: 

179 return DistributedSetup(mesh_context=MeshContext()) 

180 

181 topology_spec = ParallelTopologySpec.from_config(cfg) 

182 world_size = dist.get_world_size() if dist.is_initialized() else 1 

183 layout = resolve_parallel_layout(topology_spec, world_size) 

184 mesh_context = MeshContext.from_layout(layout) 

185 if dist.is_initialized(): 

186 mesh_context.build_meshs(get_device_type(), world_size) 

187 

188 _populate_mesh_ranks(mesh_context, topology_spec.dp_shard_size) 

189 

190 return DistributedSetup.from_config( 

191 cfg, 

192 mesh_context=mesh_context, 

193 strategy_enabled=dist.is_initialized(), 

194 ) 

195 

196 

197def destroy_process_group() -> None: 

198 """Destroy the process group and clear caches tied to the distributed runtime.""" 

199 try: 

200 if dist.is_initialized(): 

201 dist.destroy_process_group() 

202 finally: 

203 # Import here to avoid adding the cache modules to the auto-model import graph during package initialization. 

204 from hyper_parallel.distributed.context_parallel.collectives import ( # pylint: disable=C0415 

205 _HYBRID_MESH_CACHE, 

206 ) 

207 from hyper_parallel.core.dtensor.device_mesh import _DEVICE_MESH_MAP # pylint: disable=C0415 

208 from hyper_parallel.core.dtensor.dtensor import _LAYOUT_CACHE # pylint: disable=C0415 

209 from hyper_parallel.core.dtensor.tensor_redistribution import _tensor_redistribution # pylint: disable=C0415 

210 from hyper_parallel.core.utils.communication import EXISTING_COMM_GROUPS # pylint: disable=C0415 

211 from hyper_parallel.core.pipeline_parallel._p2p import ( # pylint: disable=C0415 

212 _P2P_MULTI_STREAM_GROUPS, 

213 ) 

214 

215 EXISTING_COMM_GROUPS.clear() 

216 _P2P_MULTI_STREAM_GROUPS.clear() 

217 _DEVICE_MESH_MAP.clear() 

218 _LAYOUT_CACHE.clear() 

219 _HYBRID_MESH_CACHE.clear() 

220 _tensor_redistribution._transform_cache.clear() # pylint: disable=protected-access 

221 _tensor_redistribution.is_init = False 

222 _tensor_redistribution.rank_id = None