Coverage for / home / jenkins / .local / lib / python3.10 / site-packages / hyper_parallel / components / modules / module_family.py: 92%
248 statements
« prev ^ index » next coverage.py v7.13.1, created at 2026-11-15 07:18 +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"""Shared construction of concrete family members for direct and source models."""
17from __future__ import annotations
19import inspect
20from collections.abc import Mapping
21from copy import deepcopy
22from dataclasses import dataclass, replace
23from types import MappingProxyType
24from typing import Any, Callable
26from torch import nn
28from hyper_parallel.components.modules.family import (
29 ModuleFamily,
30 SelectorSpec,
31 VariantRule,
32 _FAMILIES,
33 _validate_deferred_arguments,
34 _select_variant,
35 get_family_selector_specs,
36 get_module_family,
37 register_module_family,
38 resolve_module_family_target,
39)
42__all__ = [
43 "ModuleFamily", "SelectorSpec", "VariantRule", "FamilyRequest", "ResolvedFamily",
44 "FamilyBuildInfo", "FamilyBoundary", "register_module_family", "get_module_family",
45 "get_family_selector_specs", "resolve_module_family_target", "resolve_family", "build_family",
46 "get_family_build_info", "collect_family_modules", "validate_family_context",
47 "FamilyShardingContext",
48]
51@dataclass(frozen=True)
52class FamilyShardingContext:
53 """Logical computation axes and global shapes; no communication resources."""
55 mesh_dim_names: tuple[str, ...]
56 axis_sizes: Mapping[str, int]
57 parameter_shapes: Mapping[str, tuple[int, ...]]
58 sequence_parallel: bool = False
59 loss_parallel: bool = False
61 def __post_init__(self) -> None:
62 object.__setattr__(self, "axis_sizes", MappingProxyType(dict(self.axis_sizes)))
63 object.__setattr__(self, "parameter_shapes", MappingProxyType(dict(self.parameter_shapes)))
66@dataclass(frozen=True)
67class FamilyRequest:
68 """A public family and its explicit construction arguments, before selection."""
70 target: Callable
71 kwargs: Mapping[str, Any]
73 def __post_init__(self) -> None:
74 family = get_module_family(self.target)
75 if family is None or family.default_variant is None:
76 raise ValueError("FamilyRequest requires a family with a concrete default_variant")
77 _validate_deferred_arguments(family, self.kwargs)
78 object.__setattr__(self, "kwargs", MappingProxyType(deepcopy(dict(self.kwargs))))
81@dataclass(frozen=True)
82class ResolvedFamily:
83 """One validated implementation choice without parameters or communication resources."""
85 family: ModuleFamily
86 variant: str
87 constructor: Callable
88 arguments: Mapping[str, Any]
89 selectors: Mapping[str, Any]
90 axis_sizes: Mapping[str, int]
92 @property
93 def sharding(self) -> Callable | None:
94 """Return the selected family's layout provider."""
95 return self.family.sharding
98@dataclass(frozen=True)
99class _FamilyLayout:
100 """Borrowed mesh and completed component layout, outside serialized requests."""
102 mesh: Any
103 sharding: Any
104 parallel: Any
105 parameter_shapes: Mapping[str, tuple[int, ...]]
106 source_shard_info: Mapping[str, Any]
107 applied: bool = True
110@dataclass(frozen=True)
111class FamilyBuildInfo:
112 """Construction provenance and layout records; execution resources stay on the module."""
114 request: FamilyRequest
115 resolved: ResolvedFamily
116 config_snapshot: Any
117 origin: str
118 layout: _FamilyLayout | None = None
121@dataclass(frozen=True)
122class FamilyBoundary:
123 """One concrete family instance with all its paths in the final model tree."""
125 module: nn.Module
126 fqns: tuple[str, ...]
127 info: FamilyBuildInfo
130def _axis_sizes(context: Mapping[str, Any]) -> Mapping[str, int]:
131 """Read and validate logical parallel degrees from the construction context."""
132 result = {}
133 for axis in ("tp", "cp", "ep", "pp"):
134 size = context.get(f"{axis}_size")
135 if size is None:
136 if context.get(axis):
137 raise ValueError(f"{axis}_size is required when {axis} is active")
138 size = 1
139 actual_type = type(size)
140 if actual_type is not int or size < 1:
141 raise ValueError(f"{axis}_size must be a positive integer")
142 result[axis] = size
143 return MappingProxyType(result)
146def _mesh_axis_size(mesh: Any, axis: str) -> int:
147 """Validate declared sizes against the actual borrowed axis views."""
148 device_mesh = mesh.device_mesh
149 region = getattr(mesh, "fsdp_moe_mesh", None) if axis == "ep" else None
150 if region is None:
151 region = device_mesh
152 names = tuple(getattr(region, "mesh_dim_names", ()) or ())
153 actual = region[axis].size() if axis in names else 1
154 size = getattr(mesh, f"{axis}_size", actual)
155 actual_type = type(size)
156 if actual_type is not int or size < 1:
157 raise ValueError(f"{axis}_size must be a positive integer")
158 if axis in names and size != actual:
159 raise ValueError(f"{axis}_size={size} disagrees with the mesh size {actual}")
160 if axis not in names and size > 1 and axis != "pp":
161 raise ValueError(f"{axis}_size={size} requires an axis named {axis!r}; available mesh axes: {names!r}")
162 return size
165def _context_from_mesh(mesh: Any, context: Mapping[str, Any] | None) -> Mapping[str, Any]:
166 """Read actual job axes without copying mesh resources into a request."""
167 result = dict(context or {})
168 if mesh is None:
169 mesh = result.get("mesh")
170 if getattr(mesh, "device_mesh", None) is None:
171 return result
172 device_mesh = getattr(mesh, "device_mesh", None)
173 if device_mesh is None:
174 raise ValueError("mesh requires a prebuilt device_mesh")
175 previous = result.get("mesh")
176 if previous is not None and not _same_mesh(previous, mesh):
177 raise ValueError("context and mesh belong to different jobs")
178 for axis in ("tp", "cp", "ep", "pp"):
179 size = _mesh_axis_size(mesh, axis)
180 if f"{axis}_size" in result and result[f"{axis}_size"] != size:
181 raise ValueError(f"context and mesh disagree on {axis}_size")
182 result[f"{axis}_size"] = size
183 result[axis] = size > 1
184 result["mesh"] = mesh
185 result["sequence_parallel"] = bool(getattr(mesh, "sequence_parallel", False))
186 result["loss_parallel"] = bool(getattr(mesh, "loss_parallel", False))
187 return result
190def resolve_family(request: FamilyRequest, *, context: Mapping[str, Any] | None = None) -> ResolvedFamily:
191 """Resolve exactly one allowlisted member using the current job context.
193 Args:
194 request: Original family and explicit arguments.
195 context: Read-only job topology used for automatic CP/EP selection.
197 Returns:
198 Concrete construction information without allocating model resources.
199 """
200 family = get_module_family(request.target)
201 context = context or {}
202 axes = _axis_sizes(context)
203 args = deepcopy(dict(request.kwargs))
204 parallel = family.parallel_axis is not None and axes[family.parallel_axis] > 1
205 constructor, variant, selectors = _select_variant(family, args, parallel=parallel)
206 if constructor is family.baseline:
207 raise ValueError("a concrete family member cannot be its public factory")
208 return ResolvedFamily(
209 family, variant, constructor, MappingProxyType(args), MappingProxyType(selectors), axes,
210 )
213def _constructor_arguments(
214 member: ResolvedFamily, config: Any, module: nn.Module | None, fqn: str, context: Mapping[str, Any],
215) -> dict[str, Any]:
216 """Bind source or config construction arguments accepted by the selected member."""
217 signature = inspect.signature(member.constructor)
218 required = "config" if module is None else "module"
219 if required not in signature.parameters:
220 raise TypeError(f"member {member.variant!r} must explicitly accept {required}")
221 kwargs = dict(member.arguments)
222 unknown = set(kwargs) - set(signature.parameters)
223 if unknown:
224 raise TypeError(f"member {member.variant!r} does not accept {sorted(unknown)}")
225 runtime = {"config": config, "module": module, "module_fqn": fqn, "context": context}
226 kwargs.update({name: value for name, value in runtime.items() if name in signature.parameters})
227 signature.bind(**kwargs)
228 return kwargs
231def _normalize_family_request(target, kwargs):
232 """Resolve a public name without selecting a topology-dependent member."""
233 if isinstance(target, str):
234 matches = [family.baseline for family in _FAMILIES.values() if family.name == target]
235 if len(matches) != 1:
236 raise ValueError(f"family name {target!r} must identify exactly one registration")
237 target = matches[0]
238 if isinstance(target, FamilyRequest):
239 if kwargs:
240 raise ValueError("request arguments cannot be combined with additional keywords")
241 return target
242 return FamilyRequest(target, kwargs)
245def _build_family_member(
246 request: FamilyRequest,
247 member: ResolvedFamily,
248 *,
249 config: Any = None,
250 module: nn.Module | None = None,
251 module_fqn: str = "",
252 context: Mapping[str, Any],
253) -> nn.Module:
254 """Construct an already resolved choice without parsing the same request again."""
255 context = MappingProxyType(dict(context))
256 if member.family.validate is not None:
257 member.family.validate(config=config, module=module, context=context, member=member)
258 config = deepcopy(config)
259 result = member.constructor(**_constructor_arguments(member, config, module, module_fqn, context))
260 if not isinstance(result, nn.Module):
261 raise TypeError("a family constructor must return nn.Module")
262 info = FamilyBuildInfo(request, member, deepcopy(config), "direct" if module is None else "replacement")
263 result._hp_family_build_info = info # pylint: disable=protected-access
264 if config is not None and callable(getattr(result, "reset_parameters", None)):
265 result._hp_reset_after_materialization = True # pylint: disable=protected-access
266 return result
269def build_family(
270 target: Callable | str | FamilyRequest,
271 *,
272 config: Any = None,
273 mesh: Any = None,
274 module: nn.Module | None = None,
275 module_fqn: str = "",
276 context: Mapping[str, Any] | None = None,
277 **kwargs: Any,
278) -> nn.Module:
279 """Construct a selected member and retain its request for model preparation.
281 Args:
282 target: Public family callable, registered name, or preserved request.
283 config: Complete logical model configuration for direct construction.
284 mesh: Prebuilt job mesh context; select, shard and bind this component when supplied.
285 Active CP requires a 'cp' axis in device_mesh. Active EP requires an 'ep' axis
286 in fsdp_moe_mesh, or device_mesh when no expert view is supplied.
287 module: Existing source module; mutually exclusive with config.
288 module_fqn: Source path when converting an existing module.
289 context: Optional job constraints; without mesh, may supply logical topology for selection.
290 **kwargs: Explicit selectors and ordinary public factory arguments.
292 Returns:
293 The concrete member with construction provenance. Logical/meta builds finish physical work at runtime.
294 """
295 request = _normalize_family_request(target, kwargs)
296 arguments = dict(request.kwargs)
297 if "config" in arguments:
298 if config is not None:
299 raise ValueError("config was supplied twice")
300 config = arguments.pop("config")
301 request = FamilyRequest(request.target, arguments)
302 if (config is None) == (module is None):
303 raise ValueError("provide exactly one of config or module")
304 context = _context_from_mesh(mesh, context)
305 member = resolve_family(request, context=context)
306 result = _build_family_member(request, member, config=config, module=module,
307 module_fqn=module_fqn, context=context)
308 if mesh is not None:
309 # Layout discovery imports this module; defer its runtime dependency to avoid a cycle.
310 from hyper_parallel.distributed._builder.family_layout import _prepare_family_layout # pylint: disable=C0415,R0401
311 _prepare_family_layout(result, mesh)
312 return result
315def get_family_build_info(module: nn.Module) -> FamilyBuildInfo | None:
316 """Return only framework-created family provenance, or None for ordinary modules.
318 Args:
319 module: Module whose construction provenance is queried.
320 """
321 value = getattr(module, "_hp_family_build_info", None)
322 return value if isinstance(value, FamilyBuildInfo) else None
325def _family_parameter_shapes(model: nn.Module) -> dict[str, tuple[int, ...]]:
326 """Return logical parameter shapes even after component construction localized weights."""
327 by_identity = {}
328 for boundary in collect_family_modules(model):
329 layout = boundary.info.layout
330 if layout is not None:
331 for name, shape in layout.parameter_shapes.items():
332 identity = id(boundary.module.get_parameter(name))
333 if by_identity.setdefault(identity, shape) != shape:
334 raise ValueError(f"shared family parameter {name!r} has conflicting logical shapes")
335 return {name: by_identity.get(id(value), tuple(value.shape))
336 for name, value in model.named_parameters(remove_duplicate=False)}
339def _module_aliases(model: nn.Module) -> dict[int, tuple[nn.Module, tuple[str, ...]]]:
340 """Inventory registered module identities and all paths, including the root."""
341 found: dict[int, tuple[nn.Module, list[str]]] = {}
342 for fqn, module in model.named_modules(remove_duplicate=False):
343 found.setdefault(id(module), (module, []))[1].append(fqn)
344 return {identity: (module, tuple(paths)) for identity, (module, paths) in found.items()}
347def collect_family_modules(model: nn.Module) -> tuple[FamilyBoundary, ...]:
348 """Discover final family boundaries by object identity, including root and aliases.
350 Args:
351 model: Final or source model tree for this build.
352 """
353 boundaries = []
354 for module, paths in _module_aliases(model).values():
355 info = get_family_build_info(module)
356 if info is not None:
357 boundaries.append(FamilyBoundary(module, paths, info))
358 return tuple(boundaries)
361def validate_family_context(model: nn.Module, context: Mapping[str, Any]) -> None:
362 """Validate that model preparation has finalized members for the job context.
364 Args:
365 model: Final or source model tree for this build.
366 context: Read-only job context and effective model facts.
367 """
368 axes = _axis_sizes(_context_from_mesh(None, context))
369 for boundary in collect_family_modules(model):
370 if boundary.info.layout is not None and boundary.info.layout.applied:
371 _context_from_mesh(boundary.info.layout.mesh, context)
372 member = boundary.info.resolved
373 for axis in ("cp", "ep"):
374 if member.axis_sizes[axis] != axes[axis]:
375 raise ValueError(
376 f"family at {boundary.fqns} was constructed with {axis}={member.axis_sizes[axis]}, "
377 f"job requires {axes[axis]}; call prepare_model_components with the job context before layout"
378 )
381def _validate_applied_layout(cached: _FamilyLayout, spec: Any, fqn: str) -> None:
382 """Reject post-construction changes to a component's declared computation layout."""
383 original = cached.sharding
384 if any(name not in original.params or named != original.params[name] for name, named in spec.params.items()):
385 raise ValueError(f"layout at {fqn!r} conflicts with its constructed family layout")
386 for field_name in ("in_src", "in_dst", "out_src", "out_dst", "out_names", "is_boundary"):
387 if getattr(spec, field_name) != getattr(original, field_name):
388 raise ValueError(f"{field_name} at {fqn!r} conflicts with its constructed family layout")
391def _record_family_layouts(model, plan, mesh, shapes, source_shard_info, *, applied=True):
392 """Retain relative declarations once for constructor and model-level application."""
393 for boundary in collect_family_modules(model):
394 fqn = boundary.fqns[0]
395 if fqn not in plan.modules or boundary.info.resolved.sharding is None:
396 continue
397 existing = boundary.info.layout
398 if existing is not None and existing.applied:
399 continue
400 prefix = fqn + "." if fqn else ""
401 names = plan.modules[fqn].params
402 relative_shapes = {name: shapes[prefix + name] for name in names}
403 source = {name: source_shard_info[prefix + name] for name in names
404 if source_shard_info is not None and prefix + name in source_shard_info}
405 layout = _FamilyLayout(mesh, deepcopy(plan.modules[fqn]), plan.parallel_modules.get(fqn),
406 MappingProxyType(relative_shapes), MappingProxyType(source), applied)
407 boundary.module._hp_family_build_info = replace(boundary.info, layout=layout) # pylint: disable=protected-access
410def _family_source_shard_info(model: nn.Module) -> dict[str, Any]:
411 """Gather constructed component source layouts under their final model paths."""
412 result = {}
413 for boundary in collect_family_modules(model):
414 layout = boundary.info.layout
415 if layout is not None and layout.applied:
416 for fqn in boundary.fqns:
417 result.update({f"{fqn}.{name}" if fqn else name: value
418 for name, value in layout.source_shard_info.items()})
419 return result
422def _same_mesh(first: Any, second: Any) -> bool:
423 """Compare borrowed job views without copying or reconstructing their groups."""
424 return (getattr(first, "device_mesh", first) is getattr(second, "device_mesh", second)
425 and getattr(first, "fsdp_moe_mesh", None) is getattr(second, "fsdp_moe_mesh", None))
428def _family_mesh(model: nn.Module, mesh: Any) -> Any:
429 """Reuse the complete context when a caller passes its dense DeviceMesh view."""
430 layouts = [boundary.info.layout for boundary in collect_family_modules(model)
431 if boundary.info.layout is not None and boundary.info.layout.applied]
432 if not hasattr(mesh, "device_mesh"):
433 mesh = next((layout.mesh for layout in layouts if layout.mesh.device_mesh is mesh), mesh)
434 if any(not _same_mesh(layout.mesh, mesh) for layout in layouts):
435 raise ValueError("constructed families cannot be applied with a different job mesh")
436 return mesh