Coverage for / home / jenkins / .local / lib / python3.10 / site-packages / hyper_parallel / trainer / runtime / distributed.py: 29%
120 statements
« prev ^ index » next coverage.py v7.13.1, created at 2026-10-25 05:45 +0800
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.
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"""
25__all__ = ["all_gather", "all_reduce"]
27import logging
28import os
29from typing import TYPE_CHECKING, Any, List, Literal, Optional, Union
31import torch
32from torch import distributed as dist
34from hyper_parallel.distributed.mesh import (
35 DistributedSetup,
36 MeshContext,
37 _build_device_mesh_from_accelerator,
38)
39from hyper_parallel.trainer.runtime.device import get_device_type, get_torch_device
41logger = logging.getLogger(__name__)
44if TYPE_CHECKING:
45 from torch.distributed import ProcessGroup
48def all_gather(tensor: "torch.Tensor", world_size: int) -> "torch.Tensor":
49 """
50 Gathers the tensor from all ranks and concats them along the first dim.
51 """
52 output_tensor = torch.empty(world_size * tensor.numel(), dtype=tensor.dtype, device=get_device_type())
53 dist.all_gather_into_tensor(output_tensor, tensor)
54 return output_tensor.view(-1, *tensor.size()[1:])
57def all_reduce(
58 data: Union[int, float, List[Union[int, float]], "torch.Tensor"],
59 op: Literal["mean", "sum", "max", "min"] = "mean",
60 group: Optional["ProcessGroup"] = None,
61) -> Union[int, float, List[Union[int, float]]]:
62 """
63 Performs all reduce in the given process group.
64 """
65 if not dist.is_initialized():
66 raise RuntimeError("Distributed environment is not initialized.")
68 if not isinstance(data, torch.Tensor):
69 data = torch.tensor(data, dtype=torch.float, device=get_device_type())
71 reduce_ops = {
72 "mean": dist.ReduceOp.SUM,
73 "sum": dist.ReduceOp.SUM,
74 "max": dist.ReduceOp.MAX,
75 "min": dist.ReduceOp.MIN,
76 }
77 reduce_op = reduce_ops.get(op)
78 if reduce_op is None:
79 raise ValueError("op must be one of: mean, sum, max, min")
80 dist.all_reduce(data, op=reduce_op, group=group)
81 if op == "mean": # ReduceOp.AVG is not supported by the NPU backend
82 data /= dist.get_world_size(group=group)
84 if data.numel() == 1:
85 return data.item()
86 return data.tolist()
89def get_world_size_safe() -> int:
90 """Return dist.get_world_size() if initialized, else 1."""
91 if dist.is_initialized():
92 return dist.get_world_size()
93 return 1
96def get_global_rank_safe() -> int:
97 """Return dist.get_rank() if initialized, else 0."""
98 if dist.is_initialized():
99 return dist.get_rank()
100 return 0
103def get_local_rank_safe() -> int:
104 """Return dist.get_node_local_rank() if initialized, else 0."""
105 if dist.is_initialized():
106 return dist.get_node_local_rank()
107 return 0
110def initialize_distributed(backend: str = "nccl") -> Any:
111 """Initialize torch.distributed process group.
113 Calls dist.init_process_group if not already initialized.
114 Falls back to gloo when the requested backend is unavailable (e.g. CPU).
115 Sets the current accelerator device from ``LOCAL_RANK`` before process
116 group initialization.
117 """
118 device_type = get_device_type()
119 local_rank = int(os.environ.get("LOCAL_RANK", 0))
120 if device_type in ("cuda", "npu"):
121 get_torch_device().set_device(local_rank)
123 if not dist.is_initialized():
124 effective_backend = backend
125 if backend == "nccl" and device_type == "npu":
126 effective_backend = "hccl"
127 logger.info("Using HCCL instead of NCCL on the NPU backend.")
128 elif backend == "nccl" and device_type == "cpu":
129 effective_backend = "gloo"
130 logger.warning(
131 "CUDA not available; falling back to '%s' backend for distributed.",
132 effective_backend,
133 )
134 dist.init_process_group(backend=effective_backend)
136 return dist
139def _resolve_data_parallel_sizes(
140 world_size: int,
141 tp_size: int,
142 cp_size: int,
143 pp_size: int,
144 dp_shard_size: int,
145) -> tuple[int, int]:
146 """Validate parallel sizes and return DP and replicated-DP sizes."""
147 non_dp_size = tp_size * cp_size * pp_size
148 if world_size % non_dp_size != 0:
149 raise ValueError(
150 f"world_size {world_size} is not divisible by non-DP size {non_dp_size}"
151 )
152 dp_size = world_size // non_dp_size
153 fsdp_data_parallel_size = dp_size * cp_size
154 if fsdp_data_parallel_size % dp_shard_size != 0:
155 raise ValueError(
156 "DP+CP size "
157 f"{fsdp_data_parallel_size} is not divisible by FSDP shard size "
158 f"{dp_shard_size}"
159 )
160 return dp_size, fsdp_data_parallel_size // dp_shard_size
163def _populate_mesh_ranks(mesh_context: MeshContext, dp_shard_size: int) -> None:
164 """Populate local mesh ranks after device-mesh construction."""
165 def _local_rank(dim: str) -> int:
166 if (
167 mesh_context.device_mesh is None
168 or dim not in (mesh_context.device_mesh.mesh_dim_names or ())
169 ):
170 return 0
171 return mesh_context.device_mesh.get_local_rank(dim)
173 device_mesh_dim_names = (
174 mesh_context.device_mesh.mesh_dim_names
175 if mesh_context.device_mesh is not None
176 else ()
177 )
178 if "dp" in device_mesh_dim_names:
179 mesh_context.dp_rank = _local_rank("dp")
180 else:
181 dp_shard_rank = _local_rank("dp_shard")
182 dp_replicate_rank = _local_rank("dp_replicate")
183 mesh_context.dp_rank = dp_replicate_rank * dp_shard_size + dp_shard_rank
184 mesh_context.tp_rank = _local_rank("tp")
185 mesh_context.cp_rank = _local_rank("cp")
186 if mesh_context.pp_size > 1 and dist.is_initialized():
187 stage_world_size = mesh_context.dp_size * mesh_context.cp_size * mesh_context.tp_size
188 mesh_context.pp_rank = dist.get_rank() // stage_world_size
189 else:
190 mesh_context.pp_rank = 0
191 mesh_context.ep_rank = (
192 mesh_context.fsdp_moe_mesh.get_local_rank("ep")
193 if mesh_context.fsdp_moe_mesh is not None
194 else 0
195 )
198def create_distributed_setup_from_config(cfg: Any) -> DistributedSetup:
199 """Create DistributedSetup and build the configured mesh domains."""
200 accel = cfg.accelerator if cfg is not None and hasattr(cfg, "accelerator") else None
201 if accel is None:
202 return DistributedSetup(mesh_context=MeshContext())
204 fsdp_config = cfg.fsdp_config
205 dp_shard_size = max(1, fsdp_config.dp_shard_size)
206 edp_shard_size = max(1, fsdp_config.edp_shard_size)
207 tp_size = max(1, accel.tp_size)
208 cp_size = max(1, accel.cp_size)
209 pp_size = max(1, accel.pp_size)
210 ep_size = max(1, accel.ep_size)
212 world_size = dist.get_world_size() if dist.is_initialized() else 1
213 dp_size, dp_replicate_size = _resolve_data_parallel_sizes(
214 world_size,
215 tp_size,
216 cp_size,
217 pp_size,
218 dp_shard_size,
219 )
220 mesh_context = MeshContext(
221 dp_size=dp_size,
222 dp_replicate_size=dp_replicate_size,
223 dp_shard_size=dp_shard_size,
224 edp_shard_size=edp_shard_size,
225 tp_size=tp_size,
226 cp_size=cp_size,
227 pp_size=pp_size,
228 ep_size=ep_size,
229 sequence_parallel=bool(accel.sequence_parallel),
230 loss_parallel=bool(accel.loss_parallel),
231 )
232 if dist.is_initialized():
233 mesh_context, _ = _build_device_mesh_from_accelerator(
234 accel,
235 dp_shard_size,
236 dp_replicate_size,
237 world_size,
238 edp_shard_size,
239 )
241 _populate_mesh_ranks(mesh_context, dp_shard_size)
243 fsdp_enabled = dist.is_initialized() and (
244 dp_shard_size > 1 or dp_replicate_size > 1 or edp_shard_size > 1
245 )
246 strategy_config = fsdp_config if fsdp_enabled else None
247 return DistributedSetup(
248 mesh_context=mesh_context,
249 strategy_config=strategy_config,
250 # plan_overrides/module_replacements stay unset here: the Trainer
251 # desugars the raw YAML entries onto the setup (05 §15.2.6) before
252 # model construction.
253 low_precision_config=getattr(
254 getattr(cfg, "training", None),
255 "low_precision",
256 None,
257 ),
258 fp32_main_params=cfg.optimizer.fp32_main_params,
259 )
262def destroy_process_group() -> None:
263 """Destroy the process group and clear caches tied to the distributed runtime."""
264 try:
265 if dist.is_initialized():
266 dist.destroy_process_group()
267 finally:
268 # Import here to avoid adding the cache modules to the auto-model import graph during package initialization.
269 from hyper_parallel.distributed.context_parallel.collectives import ( # pylint: disable=C0415
270 _HYBRID_MESH_CACHE,
271 )
272 from hyper_parallel.core.dtensor.device_mesh import _DEVICE_MESH_MAP # pylint: disable=C0415
273 from hyper_parallel.core.dtensor.dtensor import _LAYOUT_CACHE # pylint: disable=C0415
274 from hyper_parallel.core.dtensor.tensor_redistribution import _tensor_redistribution # pylint: disable=C0415
275 from hyper_parallel.core.utils.communication import EXISTING_COMM_GROUPS # pylint: disable=C0415
276 from hyper_parallel.core.pipeline_parallel._p2p import ( # pylint: disable=C0415
277 _P2P_MULTI_STREAM_GROUPS,
278 )
280 EXISTING_COMM_GROUPS.clear()
281 _P2P_MULTI_STREAM_GROUPS.clear()
282 _DEVICE_MESH_MAP.clear()
283 _LAYOUT_CACHE.clear()
284 _HYBRID_MESH_CACHE.clear()
285 _tensor_redistribution._transform_cache.clear() # pylint: disable=protected-access
286 _tensor_redistribution.is_init = False
287 _tensor_redistribution.rank_id = None