Coverage for / home / jenkins / .local / lib / python3.10 / site-packages / hyper_parallel / distributed / _builder / tp_collective_lowering.py: 87%
78 statements
« prev ^ index » next coverage.py v7.13.1, created at 2026-10-25 05:45 +0800
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"""Lower TP placement transitions to differentiable local-tensor collectives."""
17import logging
18from dataclasses import dataclass
19from typing import Any, Optional, Sequence
21import torch
22import torch.distributed as dist
24from hyper_parallel.core.dtensor.placement_types import Partial, Placement, Replicate, Shard
25from hyper_parallel.core.utils import communication
27logger = logging.getLogger(__name__)
30def classify_tp_transition(
31 src: Sequence[Placement],
32 dst: Sequence[Placement],
33 mesh_dim_names: Sequence[str],
34) -> Optional[dict]:
35 """Classify one placement transition for the TP collective path.
37 Single source of truth shared by the runtime lowerer
38 (``TPCollectiveLowerer``) and the codegen boundary-form classifier
39 (``hyper_parallel.codegen.plan.boundary_forms``), so the emit-time form
40 decision and the install-time lowering can never disagree about what a
41 transition means.
43 Args:
44 src: Source placements aligned to ``mesh_dim_names``.
45 dst: Target placements aligned to ``mesh_dim_names``.
46 mesh_dim_names: Axis names the placements resolve against.
48 Returns:
49 ``{"kind": "identity"}`` when no communication is needed;
50 ``{"kind": "all_gather", "tensor_dim": d}`` for Shard -> Replicate on
51 the tp axis; ``{"kind": "all_reduce", "reduce_op": "sum"}`` for
52 Partial -> Replicate; ``{"kind": "reduce_scatter", "tensor_dim": d,
53 "reduce_op": "sum"}`` for Partial -> Shard; ``None`` when the
54 transition is not lowerable on this path (a non-tp axis differs, the
55 placement pair is unsupported, or the reduce op is not ``sum``).
56 """
57 if tuple(src) == tuple(dst):
58 return {"kind": "identity"}
59 if "tp" not in mesh_dim_names:
60 return None
61 tp_axis = tuple(mesh_dim_names).index("tp")
62 if any(
63 axis != tp_axis and src_placement != dst_placement
64 for axis, (src_placement, dst_placement) in enumerate(zip(src, dst))
65 ):
66 return None
68 src_tp = src[tp_axis]
69 dst_tp = dst[tp_axis]
70 if isinstance(src_tp, Shard) and isinstance(dst_tp, Replicate):
71 return {"kind": "all_gather", "tensor_dim": src_tp.dim}
72 if isinstance(src_tp, Partial) and isinstance(dst_tp, Replicate):
73 if src_tp.reduce_op != "sum":
74 return None
75 return {"kind": "all_reduce", "reduce_op": "sum"}
76 if isinstance(src_tp, Partial) and isinstance(dst_tp, Shard):
77 if src_tp.reduce_op != "sum":
78 return None
79 return {
80 "kind": "reduce_scatter",
81 "tensor_dim": dst_tp.dim,
82 "reduce_op": "sum",
83 }
84 return None
87@dataclass
88class TPExecutionOp:
89 """Executable local-tensor operation for one TP placement transition."""
91 kind: str
92 group: object
93 group_size: int
94 group_rank: int
95 tensor_dim: Optional[int] = None
96 reduce_op: str = "sum"
98 def execute(self, tensor: Any) -> Any:
99 """Execute the differentiable collective selected during lowering."""
100 if self.kind == "all_gather":
101 return communication.differentiable_all_gather_concat(
102 tensor,
103 self.group,
104 self.group_size,
105 self.tensor_dim,
106 )
107 if self.kind == "all_reduce":
108 return communication.differentiable_all_reduce(
109 tensor,
110 self.reduce_op,
111 self.group,
112 )
113 if self.kind == "reduce_scatter":
114 return communication.differentiable_reduce_scatter(
115 tensor,
116 self.group_size,
117 self.tensor_dim,
118 self.reduce_op,
119 self.group,
120 )
121 if self.kind == "all_reduce_shard":
122 reduced = communication.differentiable_all_reduce(
123 tensor,
124 self.reduce_op,
125 self.group,
126 )
127 return torch.chunk(
128 reduced,
129 self.group_size,
130 dim=self.tensor_dim,
131 )[self.group_rank]
132 raise ValueError(f"Unsupported TP execution operation: {self.kind!r}")
135@dataclass(frozen=True)
136class TPCollectiveLowerer:
137 """Lower supported TP-only placement differences to execution operations."""
139 mesh_dim_names: tuple[str, ...]
140 group: object
141 group_size: int
142 group_rank: int
143 backend: str
145 def __call__(
146 self,
147 src: Sequence[Placement],
148 dst: Sequence[Placement],
149 ) -> Optional[TPExecutionOp]:
150 """Return an execution operation or ``None`` for the generic fallback."""
151 result = classify_tp_transition(src, dst, self.mesh_dim_names)
152 if result is None or result["kind"] == "identity":
153 return None
154 return self.execution_op(
155 result["kind"],
156 tensor_dim=result.get("tensor_dim"),
157 reduce_op=result.get("reduce_op", "sum"),
158 )
160 def execution_op(
161 self,
162 kind: str,
163 tensor_dim: Optional[int] = None,
164 reduce_op: str = "sum",
165 ) -> TPExecutionOp:
166 """Build an execution op for one bare-operator call by kind.
168 The codegen bare-operator runtime (``runtime.TPOperators``) calls this
169 with the kind the generated forward named, so a statically emitted
170 ``reduce_scatter`` runs the exact same gloo remap (all_reduce + local
171 chunk) the placement-driven ``__call__`` path applies.
172 """
173 # Gloo has no reduce_scatter: lower to all_reduce + local chunk, the
174 # same remap the bare-operator runtime (``TPOperators``) applies.
175 if kind == "reduce_scatter" and "gloo" in self.backend:
176 kind = "all_reduce_shard"
177 return TPExecutionOp(
178 kind=kind,
179 group=self.group,
180 group_size=self.group_size,
181 group_rank=self.group_rank,
182 tensor_dim=tensor_dim,
183 reduce_op=reduce_op,
184 )
187def create_tp_collective_lowerer(
188 mesh: Any,
189 mesh_dim_names: Sequence[str],
190 *,
191 collective_backend: Optional[str] = None,
192) -> Optional[TPCollectiveLowerer]:
193 """Resolve immutable TP runtime metadata and create the boundary lowerer."""
194 mesh_dim_names = tuple(mesh_dim_names)
195 if "tp" not in mesh_dim_names:
196 return None
198 tp_mesh = mesh if mesh_dim_names == ("tp",) else mesh["tp"]
199 group = tp_mesh.get_group()
200 mesh_ranks = tuple(tp_mesh.rank_list)
201 group_ranks = tuple(dist.get_process_group_ranks(group))
202 if mesh_ranks != group_ranks:
203 logger.warning(
204 "TP mesh rank order %s differs from process group rank order %s; "
205 "falling back to DTensor redistribution.",
206 mesh_ranks,
207 group_ranks,
208 )
209 return None
211 backend = collective_backend or dist.get_backend(group)
212 return TPCollectiveLowerer(
213 mesh_dim_names=mesh_dim_names,
214 group=group,
215 group_size=tp_mesh.size(),
216 group_rank=tp_mesh.get_local_rank(),
217 backend=str(backend).lower(),
218 )
221__all__ = [
222 "TPCollectiveLowerer",
223 "TPExecutionOp",
224 "classify_tp_transition",
225 "create_tp_collective_lowerer",
226]