Coverage for / home / jenkins / .local / lib / python3.10 / site-packages / hyper_parallel / distributed / _builder / precompiled_boundary.py: 85%
126 statements
« prev ^ index » next coverage.py v7.13.1, created at 2026-10-25 05:45 +0800
1# Copyright 2025-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"""precompiled_boundary: compile-time communication planning (05 §4.3).
17RedistOp / PrecompiledBoundary compile the placement differences of
18in_src→in_dst and out_src→out_dst into sequences of RedistOps. An optional
19lowerer may attach local-tensor execution operations while the boundary keeps
20the generic DTensor redistribution fallback.
22API adaptation (differences from the 05 doc pseudocode; actual in-house
23DTensor signatures):
24- ``DTensor.from_local(local, mesh, placements)``: no run_check parameter;
25- ``dt.redistribute(mesh, placements)``: mesh is the first argument, no async_op.
26"""
28import logging
29from dataclasses import dataclass
30from typing import Any, Callable, Optional, Protocol, Sequence, Tuple
32import torch
34from hyper_parallel.core.dtensor.dtensor import DTensor
35from hyper_parallel.core.dtensor.placement_types import Partial, Placement, Shard
36from hyper_parallel.distributed.recipe_spec import resolve_placements
38logger = logging.getLogger(__name__)
41def _classify_collective(src, dst) -> str:
42 """Derive the communication type from placements (a debug/profiling label,
43 not a communication-path selector).
45 Only differing dimensions are compared. Identity placement dimensions do
46 not participate in classification, so a Shard-to-Replicate transition is
47 classified as all_gather.
48 """
49 if tuple(src) == tuple(dst):
50 return "identity"
51 diff_src = tuple(s for s, d in zip(src, dst) if s != d)
52 diff_dst = tuple(d for s, d in zip(src, dst) if s != d)
54 has_shard_src = any(isinstance(p, Shard) for p in diff_src)
55 has_partial_src = any(isinstance(p, Partial) for p in diff_src)
56 has_shard_dst = any(isinstance(p, Shard) for p in diff_dst)
57 all_replicate_dst = all(
58 not isinstance(p, (Shard, Partial)) for p in diff_dst
59 )
61 if has_partial_src and has_shard_dst:
62 return "reduce_scatter"
63 if has_partial_src and all_replicate_dst:
64 return "all_reduce"
65 if has_shard_src and all_replicate_dst:
66 return "all_gather"
67 return "redistribute"
70def _get_arg(args, kwargs, name, idx, default=None):
71 if name in kwargs:
72 return kwargs[name]
73 if idx is not None and idx < len(args):
74 return args[idx]
75 return default
78def _set_arg(args, kwargs, name, idx, value):
79 if name in kwargs:
80 kwargs[name] = value
81 return args, kwargs
82 if idx is not None and idx < len(args):
83 args = list(args)
84 args[idx] = value
85 return tuple(args), kwargs
86 kwargs[name] = value
87 return args, kwargs
90class BoundaryExecutionOp(Protocol):
91 """Operation attached by a boundary lowerer and executed on local tensors."""
93 def execute(self, tensor: Any) -> Any:
94 """Execute one pre-lowered placement transition."""
97BoundaryOpLowerer = Callable[
98 [Sequence[Placement], Sequence[Placement]],
99 Optional[BoundaryExecutionOp],
100]
103@dataclass
104class RedistOp:
105 """A single precompiled redistribute operation (05 §4.3.1).
107 ``collective_type`` is a debug/profiling label. Validation and transitions
108 rejected by the injected lowerer use ``DTensor.redistribute()``; accepted
109 transitions use ``execution_op`` directly on local tensors.
110 """
111 arg_name: str
112 arg_index: Optional[int]
113 mesh: object # DeviceMesh
114 src_placements: Tuple[Placement, ...]
115 dst_placements: Tuple[Placement, ...]
116 collective_type: str
117 execution_op: Optional[BoundaryExecutionOp] = None
119 def execute(self, tensor: torch.Tensor, *, as_dtensor: bool = False) -> Any:
120 """Execute the communication.
122 Args:
123 tensor: input local tensor (or DTensor).
124 as_dtensor: True → return a DTensor (validate mode), False → return a local tensor.
125 """
126 if self.collective_type == "identity":
127 if isinstance(tensor, DTensor):
128 # validate (as_dtensor=True) keeps the DTensor; production
129 # returns local — an identity op's input may come from a
130 # from_local re-wrap in a local region (MoE/CP wrapper), so
131 # the boundary exit must unwrap it.
132 return tensor if as_dtensor else tensor.to_local()
133 if as_dtensor:
134 return DTensor.from_local(
135 tensor, self.mesh, tuple(self.src_placements))
136 return tensor
138 if not as_dtensor and self.execution_op is not None:
139 local_tensor = tensor.to_local() if isinstance(tensor, DTensor) else tensor
140 return self.execution_op.execute(local_tensor)
142 # Unified path: zero-copy wrap → redistribute → optional to_local
143 if isinstance(tensor, DTensor):
144 dt = tensor
145 else:
146 dt = DTensor.from_local(tensor, self.mesh, tuple(self.src_placements))
147 dt = dt.redistribute(self.mesh, tuple(self.dst_placements))
148 return dt if as_dtensor else dt.to_local()
151class PrecompiledBoundary:
152 """Compile-time communication plan (05 §4.3.3): two RedistOp sequences, in_plan/out_plan."""
154 def __init__(
155 self,
156 spec: Any,
157 mesh: Any,
158 mesh_dim_names: Sequence[str],
159 *,
160 op_lowerer: Optional[BoundaryOpLowerer] = None,
161 ) -> None:
162 """Compile the in/out RedistOp plans for one module boundary.
164 Args:
165 spec: ModuleShardingSpec carrying the in_src/in_dst/out_src/out_dst
166 placement contracts.
167 mesh: DeviceMesh the redistributions execute on.
168 mesh_dim_names: Ordered mesh dimension names used to resolve
169 placement shorthands.
170 op_lowerer: Optional lowerer that turns a (src, dst) placement
171 transition into a local-tensor execution op.
172 """
173 self.spec = spec
174 self.mesh = mesh
175 self.mesh_dim_names = tuple(mesh_dim_names)
176 self._op_lowerer = op_lowerer
177 self.in_plan = self._compile_input_plan(spec, mesh, self.mesh_dim_names)
178 self.out_plan = self._compile_output_plan(spec, mesh, self.mesh_dim_names)
180 def _lower_execution_op(self, src, dst) -> Optional[BoundaryExecutionOp]:
181 """Delegate optional execution lowering without owning backend details."""
182 return self._op_lowerer(src, dst) if self._op_lowerer is not None else None
184 # ── Compilation ─────────────────────────────────────────────────────
186 def _compile_input_plan(self, spec, mesh, mesh_dim_names):
187 """Compile the input communication plan from in_src → in_dst (identity
188 dimensions naturally compile to pass-through ops)."""
189 plan = []
190 in_src = spec.in_src or {} # tolerate None: a hand-written spec (debug
191 in_dst = spec.in_dst or {} # shortcut) may not be normalized by the plan
192 # (the input-side semantics of "inherit when not declared")
193 all_names = set(in_src.keys()) | set(in_dst.keys())
194 for name in sorted(all_names):
195 src_p = tuple(resolve_placements(
196 in_src.get(name, {}), mesh_dim_names))
197 dst_p = tuple(resolve_placements(
198 in_dst.get(name, {}), mesh_dim_names))
199 plan.append(RedistOp(
200 arg_name=name,
201 arg_index=None,
202 mesh=mesh,
203 src_placements=src_p,
204 dst_placements=dst_p,
205 collective_type=_classify_collective(src_p, dst_p),
206 execution_op=self._lower_execution_op(src_p, dst_p),
207 ))
208 return plan
210 def _compile_output_plan(self, spec, mesh, mesh_dim_names):
211 """Compile the output communication plan from out_src → out_dst (identity
212 skipped, multi-output supported).
214 arg_index source priority: (1) explicit order in spec.out_names;
215 (2) key order of out_src.
216 out_src=None or out_dst=None → nothing is compiled.
217 """
218 if spec.out_src is None or spec.out_dst is None:
219 return []
221 out_names = getattr(spec, "out_names", None) or list(spec.out_src.keys())
222 name_to_idx = {name: i for i, name in enumerate(out_names)}
224 plan = []
225 all_names = set(spec.out_src.keys()) | set(spec.out_dst.keys())
226 for name in sorted(all_names):
227 src_p = tuple(resolve_placements(
228 spec.out_src.get(name, {}), mesh_dim_names))
229 dst_p = tuple(resolve_placements(
230 spec.out_dst.get(name, {}), mesh_dim_names))
231 if src_p == dst_p:
232 continue # identity, no communication needed
233 plan.append(RedistOp(
234 arg_name=name,
235 arg_index=name_to_idx.get(name, 0),
236 mesh=mesh,
237 src_placements=src_p,
238 dst_placements=dst_p,
239 collective_type=_classify_collective(src_p, dst_p),
240 execution_op=self._lower_execution_op(src_p, dst_p),
241 ))
242 return plan
244 # ── Runtime execution ───────────────────────────────────────────────
246 def redistribute_inputs(
247 self,
248 args: Sequence[Any],
249 kwargs: dict,
250 *,
251 as_dtensor: bool = False,
252 ) -> Tuple[Sequence[Any], dict]:
253 """Execute input redistribution. as_dtensor=True → return DTensors (validate mode).
255 When an arg is not found in args/kwargs (None) the op is skipped —
256 e.g. embed's in_src key "input" differs from the actual kwargs name
257 "input_ids" and is identity.
258 """
259 for op in self.in_plan:
260 arg = _get_arg(args, kwargs, op.arg_name, op.arg_index, default=None)
261 if arg is None:
262 continue
263 result = op.execute(arg, as_dtensor=as_dtensor)
264 args, kwargs = _set_arg(args, kwargs, op.arg_name, op.arg_index, result)
265 return args, kwargs
267 def redistribute_outputs(self, outputs: Any, *, as_dtensor_input: bool = False) -> Any:
268 """Execute output redistribution (single Tensor output / multi-output
269 tuple, order preserved, same structure returned).
271 as_dtensor_input=True → inputs are already DTensors (validate mode).
272 """
273 is_sequence = isinstance(outputs, (tuple, list))
274 outputs_list = list(outputs) if is_sequence else [outputs]
275 for op in self.out_plan:
276 idx = op.arg_index if op.arg_index is not None else 0
277 if idx >= len(outputs_list):
278 logger.warning(
279 "PrecompiledBoundary: out_plan expects output '%s' at index %d, "
280 "but module returned only %d outputs. Skipping.",
281 op.arg_name, idx, len(outputs_list),
282 )
283 continue
284 tensor = outputs_list[idx]
285 if tensor is None:
286 continue
287 # as_dtensor_input=True (validate) → keep DTensors for out_dst
288 # validation; otherwise return local (production / final boundary exit).
289 outputs_list[idx] = op.execute(tensor, as_dtensor=as_dtensor_input)
290 if isinstance(outputs, tuple):
291 return tuple(outputs_list)
292 if isinstance(outputs, list):
293 return outputs_list
294 return outputs_list[0]