Coverage for / home / jenkins / .local / lib / python3.10 / site-packages / hyper_parallel / distributed / _builder / local_region.py: 0%
78 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"""local_region: a DTensor -> local -> DTensor local-compute-region wrapper.
17Built on the ``core.shard.custom_shard`` skeleton, with three enhancements for
18the 05 dual-mode DTensor design (the MoE local_map / CP attention internal
19regions in validate mode):
211. **Named-parameter binding**: ``in_placements`` is a
22 ``dict[str, placements]``, aligned with the dict contract of
23 ``ModuleShardingSpec.in_dst``; positional args are mapped to parameter
24 names via ``inspect.signature``, and kwargs are natively supported (HF
25 forwards are predominantly kwargs-based).
262. **Tolerant passthrough**: inputs that are not DTensors are passed through
27 as-is (the production path where parameters are already unwrapped); outputs
28 that are already DTensors are not re-wrapped; when no input is a DTensor,
29 outputs are not wrapped either.
31**Why there is no backward stitching (important)**: hyper_parallel's DTensor
32is an in-house **forward-only** placement/dispatch system; the backward pass
33does not go through DTensor (there is no DTensor autograd). Therefore this
34function only performs forward unwrap/wrap and contains -- and needs -- no
35autograd.Function stitching or gradient-placement declarations (unlike
36PyTorch ``local_map`` / Titan ``LocalMapConfig.in_grad_placements``, which
37exist because torch DTensor has backward semantics). Backward inside the
38region is plain autograd on local tensors, with gradients landing directly on
39the local parameter shards, consistent with production mode.
41Relationship with production mode: production's forward wrapper
42(``_wrap_local_region_forward``) permanently unwraps parameters to plain
43tensors at build time, and boundary communication is executed by
44``PrecompiledBoundary``; it does not use this function. This function serves
45**validate mode** (parameters stay DTensors and the region boundary needs
46DTensor-contract stitching) and standalone use.
47"""
49import functools
50import inspect
51from typing import Any, Callable, Dict, Optional, Sequence
53import torch
55from hyper_parallel.core.dtensor.dtensor import DTensor
56from hyper_parallel.core.dtensor.layout import DeviceMesh
57from hyper_parallel.core.dtensor.placement_types import Placement
59# Placements of a single tensor: tuple[Placement, ...], aligned with the mesh dims
60Placements = Sequence[Placement]
63def _bind_arg_names(func: Callable) -> Dict[str, int]:
64 """Map positional-argument positions to parameter names (for name-based
65 lookup under the dict contract).
67 When the signature cannot be introspected (C extensions, etc.), an empty
68 dict is returned -- in that case only kwargs-passed arguments can be
69 matched by in_placements, and all positional arguments pass through.
70 """
71 try:
72 sig = inspect.signature(func)
73 except (TypeError, ValueError):
74 return {}
75 return {
76 name: idx
77 for idx, (name, p) in enumerate(sig.parameters.items())
78 if p.kind in (p.POSITIONAL_ONLY, p.POSITIONAL_OR_KEYWORD)
79 }
82def _normalize_out_placements(out_placements, num_outputs: int):
83 """Normalize out_placements into a per-output
84 tuple[tuple[Placement, ...] | None, ...].
86 Accepted forms:
87 - flat single-output form ``(Partial(), Replicate())`` (all elements
88 are Placements);
89 - per-output form ``((Partial(), Replicate()), None, (Shard(1), Replicate()))``
90 (any element is a tuple or None; length must equal the output count).
91 """
92 if len(out_placements) == 0:
93 raise ValueError("out_placements must not be empty")
94 per_output = any(p is None or isinstance(p, tuple) for p in out_placements)
95 if per_output:
96 if len(out_placements) != num_outputs:
97 raise ValueError(
98 f"out_placements count {len(out_placements)} does not match "
99 f"output count {num_outputs}!"
100 )
101 return tuple(out_placements)
102 if num_outputs != 1:
103 raise ValueError(
104 f"flat out_placements only valid for single-output functions, "
105 f"got {num_outputs} outputs!"
106 )
107 return (tuple(out_placements),)
110def _unwrap_region_inputs(
111 args: Sequence[Any],
112 kwargs: Dict[str, Any],
113 name_to_idx: Dict[str, int],
114 in_placements: Dict[str, Optional[Placements]],
115 device_mesh: DeviceMesh,
116 redistribute_inputs: bool,
117) -> tuple[list[Any], bool]:
118 """Convert declared DTensor inputs to local tensors."""
119 local_args = list(args)
120 saw_dtensor = False
121 for name, placements in in_placements.items():
122 if name in kwargs:
123 from_kwargs, index = True, None
124 value = kwargs[name]
125 else:
126 from_kwargs = False
127 index = name_to_idx.get(name)
128 if index is None or index >= len(local_args):
129 continue
130 value = local_args[index]
131 if not isinstance(value, DTensor):
132 continue
133 saw_dtensor = True
134 dtensor = value
135 if (
136 redistribute_inputs
137 and placements is not None
138 and tuple(dtensor.placements) != tuple(placements)
139 ):
140 dtensor = dtensor.redistribute(device_mesh, placements)
141 local_value = dtensor.to_local()
142 if from_kwargs:
143 kwargs[name] = local_value
144 else:
145 local_args[index] = local_value
146 return local_args, saw_dtensor
149def _wrap_region_outputs(
150 output: Any,
151 out_placements: Sequence[Optional[Placements]],
152 device_mesh: DeviceMesh,
153) -> Any:
154 """Wrap tensor outputs into DTensors using declared placements."""
155 single = not isinstance(output, tuple)
156 output_items = (output,) if single else output
157 placements_items = _normalize_out_placements(out_placements, len(output_items))
158 wrapped_output = []
159 for item, placements in zip(output_items, placements_items):
160 if isinstance(item, DTensor):
161 wrapped_output.append(item)
162 elif isinstance(item, torch.Tensor):
163 if placements is None:
164 raise TypeError("Tensor output requires non-None out_placements entry!")
165 wrapped_output.append(DTensor.from_local(item, device_mesh, tuple(placements)))
166 else:
167 wrapped_output.append(item)
168 return wrapped_output[0] if single else tuple(wrapped_output)
171def local_region(
172 func: Optional[Callable] = None,
173 *,
174 device_mesh: DeviceMesh,
175 in_placements: Optional[Dict[str, Optional[Placements]]] = None,
176 out_placements: Optional[Sequence[Optional[Placements]]] = None,
177 redistribute_inputs: bool = False,
178) -> Callable:
179 """Wrap func into a DTensor -> local -> DTensor local compute region (forward).
181 Args:
182 func: the function to wrap (a forward or any callable). Can also be
183 used as a decorator factory.
184 device_mesh: the mesh used for DTensor construction / redistribution.
185 in_placements: ``{arg_name: placements}`` -- the expected placements
186 of each DTensor input at the region entry. Inputs left as None
187 (value) or not listed are not redistributed; non-DTensor inputs
188 always pass through.
189 out_placements: placement declarations for the outputs at the region
190 exit. A single output may be written flat as
191 ``(Partial(), Replicate())``; multiple outputs are written
192 position-by-position, with None as the placeholder for non-tensor
193 outputs. When None, outputs are not wrapped (returned as-is).
194 redistribute_inputs: whether to redistribute inputs to the placements
195 declared in in_placements at the entry. In dual-mode scenarios the
196 boundary communication is already done by PrecompiledBoundary, so
197 pass False (default); pass True for standalone use.
199 Returns:
200 The wrapped function, with the same signature as func.
202 Example:
203 >>> # validate-mode MoE module: boundary DTensor contract preserved,
204 >>> # internal local all-to-all
205 >>> wrapped = local_region(
206 ... moe.forward, device_mesh=mesh,
207 ... in_placements={"hidden_states": (Replicate(), Replicate())},
208 ... out_placements=(Partial(), Replicate()),
209 ... )
211 >>> # decorator form (standalone use, entry redistributes by itself)
212 >>> @local_region(device_mesh=mesh,
213 ... in_placements={"x": (Shard(0),)},
214 ... out_placements=((Shard(0),),),
215 ... redistribute_inputs=True)
216 ... def my_fn(x, bias=None):
217 ... return x + bias
218 """
219 def decorator(fn: Callable) -> Callable:
220 """Build the local-region wrapper around one concrete function."""
221 name_to_idx = None # lazily cached signature mapping
223 @functools.wraps(fn)
224 def wrapped(*args: Any, **kwargs: Any) -> Any:
225 """Unwrap DTensor inputs to local, run fn, and re-wrap outputs."""
226 nonlocal name_to_idx
227 if name_to_idx is None:
228 name_to_idx = _bind_arg_names(fn)
230 args, saw_dtensor = _unwrap_region_inputs(
231 args, kwargs, name_to_idx, in_placements or {}, device_mesh, redistribute_inputs
232 )
234 out = fn(*args, **kwargs)
236 if not saw_dtensor or out_placements is None:
237 return out
239 return _wrap_region_outputs(out, out_placements, device_mesh)
241 return wrapped
243 if func is not None:
244 return decorator(func)
245 return decorator