Diff Coverage

        mesh_shape=(layout.dp_size, topology_spec.cp_size, topology_spec.tp_size),
    )
    mesh_context = MeshContext.from_layout(layout)
    mesh_context.device_mesh = topology
    if topology_spec.ep_size > 1:
        mesh_context.fsdp_moe_mesh = LogicalDeviceMesh(
            mesh_dim_names=("edp", "ep"),
            mesh_shape=(topology.size() // topology_spec.ep_size, topology_spec.ep_size),
        )
    return DistributedSetup.from_config(

    load_base_model: bool,
    session: CodegenBuildSession,
) -> Any:
    """Apply native replacement declarations while recording their result."""
    specs = list(getattr(distributed_setup, "module_replacements", None) or ())
    original_compile = module_build.compile_module_replacements

    def _observe_compile(*args: Any, **kwargs: Any) -> Any:
        plan = original_compile(*args, **kwargs)
        specs[:] = plan.specs
        return plan

    module_build.compile_module_replacements = _observe_compile
    try:
        with _observe_replacement_installations(specs) as applications:
            model, _ = model_builder._maybe_apply_module_replacements(  # pylint: disable=protected-access
                model, distributed_setup, low_precision_config, load_base_model,
            )
    except replacement_api.ModuleReplacementMatchError as exc:
        session.record_replacements(

            applications=(), unmatched=exc.unmatched,
        )
        raise
    finally:
        module_build.compile_module_replacements = original_compile
    session.record_replacements(
        declarations=replacement_api.describe_replacement_declarations(specs),
        applications=applications, unmatched=(),
    )

def _observe_replacement_installations(specs: Sequence[Any]) -> Iterator[list[dict[str, Any]]]:
    """Record metadata only after the shared installer commits a prepared candidate."""
    applications: list[dict[str, Any]] = []
    spec_indexes = {id(spec): index for index, spec in enumerate(specs)}
    pending = {}
    original_prepare = replacement_api._prepare_replacement_target  # pylint: disable=protected-access
    original_install = replacement_api._install_module_candidates  # pylint: disable=protected-access

    def _observe_prepare(target: Any, *args: Any, **kwargs: Any) -> Any:
        candidate, transforms = original_prepare(target, *args, **kwargs)
        if id(target.spec) not in spec_indexes:
            spec_indexes.update({id(spec): index for index, spec in enumerate(specs)})
        index = spec_indexes.setdefault(id(target.spec), len(spec_indexes))
        pending[(target.module_fqns, id(candidate))] = {
            "rule_id": replacement_api._replacement_rule_id(target.spec, index),  # pylint: disable=protected-access
            "module_fqns": list(target.module_fqns),
            "primary_fqn": target.module_fqns[0],
            "source_type": replacement_api._symbol_path(type(target.source)),

            "source_type": replacement_api._symbol_path(type(target.source)),
            "replacement_type": replacement_api._symbol_path(type(candidate)),
            "weight_transforms": [replacement_api._symbol_path(type(transform)) for transform in transforms],
        }
        return candidate, transforms

    def _observe_install(root: Any, candidates: Any) -> Any:
        candidates = tuple(candidates)
        result = original_install(root, candidates)
        for fqns, _, candidate in candidates:
            record = pending.pop((fqns, id(candidate)), None)
            if record is not None:
                applications.append(record)
        return result

    replacement_api._prepare_replacement_target = _observe_prepare  # pylint: disable=protected-access
    replacement_api._install_module_candidates = _observe_install  # pylint: disable=protected-access
    try:
        yield applications
    finally:
        replacement_api._prepare_replacement_target = original_prepare  # pylint: disable=protected-access
        replacement_api._install_module_candidates = original_install  # pylint: disable=protected-access


def _observe_plan(
    model: Any,

        return
    if mesh.device_mesh is None:
        session.mark_not_required("device_mesh_unavailable")
        return
    if (is_hf_model and not model_sharding_requested
            and not model_builder._has_module_owned_layouts(model)):  # pylint: disable=protected-access
        if session.has_replacement_applications:
            session.observe_plan(None, None, mesh)
        else:
            session.mark_not_required("pure_data_parallel")
        return
    plan = sharding_planner.plan(

            swap_inputs,
            allow_cache_entry_mutation,
            activation_swap,
        )
        factory, context = auto_model._resolve_module_construction(  # pylint: disable=protected-access
            kwargs, hf_config, distributed_setup, model_init_dtype or torch_dtype,
        )
        init_kwargs = auto_model._model_init_kwargs(  # pylint: disable=protected-access
            dict(kwargs),

            load_base_model=False,
        )
        if "trust_remote_code" in kwargs:
            init_kwargs["trust_remote_code"] = kwargs["trust_remote_code"]
        _prepare_model_source_compat(context["model_spec"])
        with _codegen_model_init_context():
            if factory is not None:
                if model_args:
                    raise ValueError("model_factory uses config and context, not additional positional model arguments")
                model = auto_model.construct_family_model(factory, hf_config, context=context)
            else:
                _, model = auto_model._init_model(  # pylint: disable=protected-access
                    cls, None, hf_config, attn_implementation, torch_dtype, is_hf_model,
                    *model_args, **init_kwargs,
                )
                auto_model.attach_build_context(model, context)
        _record_source(
            session,
            model,
            hf_config,

            model_builder._apply_parameter_freezing(  # pylint: disable=protected-access
                model,
                freeze_config,
            )
        model_builder._apply_model_method_replacements(  # pylint: disable=protected-access
            model, distributed_setup, getattr(distributed_setup, "low_precision_config", None),
        )
        _observe_plan(model, mesh, sharding_planner, is_hf_model, session)
        snapshot = _publish_session(session, artifact_dir, result)

    """Check complete, unique whitelist entries before registration."""
    defaults = _validate_selector_specs(family)
    _validate_variant_paths(family)
    if family.parallel_axis not in (None, "cp", "ep"):
        raise ValueError("parallel_axis must be cp, ep or None")
    if family.default_variant is not None:
        if family.default_variant not in family.variants:
            raise ValueError("default_variant must name a declared concrete variant")
        if "parallel" in family.selector_specs:
            raise ValueError("parallel is derived from context, not a public selector")
        defaults = {**defaults, "parallel": False}
    conditions = [defaults]  # The default combination is an explicit baseline allowance.
    for rule in family.variant_rules:
        if not isinstance(rule, VariantRule) or not isinstance(rule.when, Mapping):
            raise TypeError(f"ModuleFamily {family.name}: expected VariantRule with a mapping condition")

        if set(rule.when) != set(defaults):
            raise ValueError(f"ModuleFamily {family.name}: each rule must contain exactly {sorted(defaults)}")
        for name, value in rule.when.items():
            if name == "parallel" and family.default_variant is not None:
                actual_type = type(value)
                if actual_type is not bool:
                    raise TypeError("derived parallel conditions must be bool")
            else:
                _validate_selector_value(family, name, value)
        if rule.variant not in family.variants:
            raise ValueError(f"ModuleFamily {family.name}: undeclared variant {rule.variant!r}")

    selectors = {name: args.pop(name, spec.default) for name, spec in family.selector_specs.items()}
    for name, value in selectors.items():
        _validate_selector_value(family, name, value)
    defaults = {name: spec.default for name, spec in family.selector_specs.items()}
    if family.default_variant is not None:
        selectors["parallel"] = parallel
        defaults["parallel"] = False
    if selectors == defaults:
        variant = family.default_variant
    else:
        matches = [rule.variant for rule in family.variant_rules if selectors == dict(rule.when)]
        if len(matches) != 1:
            legal = [defaults, *(dict(rule.when) for rule in family.variant_rules)]
            raise ValueError(
                f"ModuleFamily {family.name}: unsupported selector combination {selectors}; legal combinations: {legal}"
            )
        variant = matches[0]
    constructor = family.baseline
    if variant is not None:
        module_path, _, name = family.variants[variant].rpartition(".")
        constructor = getattr(import_module(module_path), name)
    return constructor, variant, selectors


def resolve_module_family_target(target: Callable, arguments: Mapping) -> tuple[Callable, dict]:
    """Select a registered implementation and validate its ordinary keywords.

    args = normalize_family_selectors(target, arguments)
    family = get_module_family(target)
    if family is None:
        return target, args
    if family.default_variant is not None:
        _validate_deferred_arguments(family, args)
        return target, args
    target, _, _ = _select_variant(family, args)
    target_signature = signature(target)
    allowed = {
        name
        for name, parameter in target_signature.parameters.items()

def _validate_deferred_arguments(family: ModuleFamily, args: Mapping[str, Any]) -> None:
    """Validate a new-family request without choosing a topology-dependent member."""
    allowed = set(signature(family.baseline).parameters) - {"module", "module_fqn", "context", "mesh"}
    unknown = set(args) - allowed
    if unknown:
        raise TypeError(f"ModuleFamily {family.name}: unknown replacement arguments: {sorted(unknown)}")
    for name in family.selector_specs:
        if name in args:
            _validate_selector_value(family, name, args[name])

    sequence_parallel: bool = False
    loss_parallel: bool = False

    def __post_init__(self) -> None:
        object.__setattr__(self, "axis_sizes", MappingProxyType(dict(self.axis_sizes)))
        object.__setattr__(self, "parameter_shapes", MappingProxyType(dict(self.parameter_shapes)))


@dataclass(frozen=True)
class FamilyRequest:

    target: Callable
    kwargs: Mapping[str, Any]

    def __post_init__(self) -> None:
        family = get_module_family(self.target)
        if family is None or family.default_variant is None:
            raise ValueError("FamilyRequest requires a family with a concrete default_variant")
        _validate_deferred_arguments(family, self.kwargs)
        object.__setattr__(self, "kwargs", MappingProxyType(deepcopy(dict(self.kwargs))))


@dataclass(frozen=True)
class ResolvedFamily:

    @property
    def sharding(self) -> Callable | None:
        """Return the selected family's layout provider."""
        return self.family.sharding


@dataclass(frozen=True)
class _FamilyLayout:

def _axis_sizes(context: Mapping[str, Any]) -> Mapping[str, int]:
    """Read and validate logical parallel degrees from the construction context."""
    result = {}
    for axis in ("tp", "cp", "ep", "pp"):
        size = context.get(f"{axis}_size")
        if size is None:
            if context.get(axis):
                raise ValueError(f"{axis}_size is required when {axis} is active")
            size = 1
        actual_type = type(size)
        if actual_type is not int or size < 1:
            raise ValueError(f"{axis}_size must be a positive integer")
        result[axis] = size
    return MappingProxyType(result)


def _mesh_axis_size(mesh: Any, axis: str) -> int:
    """Validate declared sizes against the actual borrowed axis views."""
    device_mesh = mesh.device_mesh
    region = getattr(mesh, "fsdp_moe_mesh", None) if axis == "ep" else None
    if region is None:
        region = device_mesh
    names = tuple(getattr(region, "mesh_dim_names", ()) or ())
    actual = region[axis].size() if axis in names else 1
    size = getattr(mesh, f"{axis}_size", actual)
    actual_type = type(size)
    if actual_type is not int or size < 1:
        raise ValueError(f"{axis}_size must be a positive integer")
    if axis in names and size != actual:
        raise ValueError(f"{axis}_size={size} disagrees with the mesh size {actual}")
    if axis not in names and size > 1 and axis != "pp":
        raise ValueError(f"{axis}_size={size} requires an axis named {axis!r}; available mesh axes: {names!r}")
    return size


def _context_from_mesh(mesh: Any, context: Mapping[str, Any] | None) -> Mapping[str, Any]:
    """Read actual job axes without copying mesh resources into a request."""
    result = dict(context or {})
    if mesh is None:
        mesh = result.get("mesh")
        if getattr(mesh, "device_mesh", None) is None:
            return result
    device_mesh = getattr(mesh, "device_mesh", None)
    if device_mesh is None:
        raise ValueError("mesh requires a prebuilt device_mesh")
    previous = result.get("mesh")
    if previous is not None and not _same_mesh(previous, mesh):
        raise ValueError("context and mesh belong to different jobs")
    for axis in ("tp", "cp", "ep", "pp"):
        size = _mesh_axis_size(mesh, axis)
        if f"{axis}_size" in result and result[f"{axis}_size"] != size:
            raise ValueError(f"context and mesh disagree on {axis}_size")
        result[f"{axis}_size"] = size
        result[axis] = size > 1
    result["mesh"] = mesh
    result["sequence_parallel"] = bool(getattr(mesh, "sequence_parallel", False))
    result["loss_parallel"] = bool(getattr(mesh, "loss_parallel", False))
    return result


def resolve_family(request: FamilyRequest, *, context: Mapping[str, Any] | None = None) -> ResolvedFamily:
    """Resolve exactly one allowlisted member using the current job context.

    Returns:
        Concrete construction information without allocating model resources.
    """
    family = get_module_family(request.target)
    context = context or {}
    axes = _axis_sizes(context)
    args = deepcopy(dict(request.kwargs))
    parallel = family.parallel_axis is not None and axes[family.parallel_axis] > 1
    constructor, variant, selectors = _select_variant(family, args, parallel=parallel)
    if constructor is family.baseline:
        raise ValueError("a concrete family member cannot be its public factory")
    return ResolvedFamily(
        family, variant, constructor, MappingProxyType(args), MappingProxyType(selectors), axes,
    )

def _constructor_arguments(
    member: ResolvedFamily, config: Any, module: nn.Module | None, fqn: str, context: Mapping[str, Any],
) -> dict[str, Any]:
    """Bind source or config construction arguments accepted by the selected member."""
    signature = inspect.signature(member.constructor)
    required = "config" if module is None else "module"
    if required not in signature.parameters:
        raise TypeError(f"member {member.variant!r} must explicitly accept {required}")
    kwargs = dict(member.arguments)
    unknown = set(kwargs) - set(signature.parameters)
    if unknown:
        raise TypeError(f"member {member.variant!r} does not accept {sorted(unknown)}")
    runtime = {"config": config, "module": module, "module_fqn": fqn, "context": context}
    kwargs.update({name: value for name, value in runtime.items() if name in signature.parameters})
    signature.bind(**kwargs)
    return kwargs


def _normalize_family_request(target, kwargs):
    """Resolve a public name without selecting a topology-dependent member."""
    if isinstance(target, str):
        matches = [family.baseline for family in _FAMILIES.values() if family.name == target]
        if len(matches) != 1:
            raise ValueError(f"family name {target!r} must identify exactly one registration")
        target = matches[0]
    if isinstance(target, FamilyRequest):
        if kwargs:
            raise ValueError("request arguments cannot be combined with additional keywords")
        return target
    return FamilyRequest(target, kwargs)


def _build_family_member(
    request: FamilyRequest,

    module_fqn: str = "",
    context: Mapping[str, Any],
) -> nn.Module:
    """Construct an already resolved choice without parsing the same request again."""
    context = MappingProxyType(dict(context))
    if member.family.validate is not None:
        member.family.validate(config=config, module=module, context=context, member=member)
    config = deepcopy(config)
    result = member.constructor(**_constructor_arguments(member, config, module, module_fqn, context))
    if not isinstance(result, nn.Module):
        raise TypeError("a family constructor must return nn.Module")
    info = FamilyBuildInfo(request, member, deepcopy(config), "direct" if module is None else "replacement")
    result._hp_family_build_info = info  # pylint: disable=protected-access
    if config is not None and callable(getattr(result, "reset_parameters", None)):
        result._hp_reset_after_materialization = True  # pylint: disable=protected-access
    return result


def build_family(
    target: Callable | str | FamilyRequest,

    Returns:
        The concrete member with construction provenance. Logical/meta builds finish physical work at runtime.
    """
    request = _normalize_family_request(target, kwargs)
    arguments = dict(request.kwargs)
    if "config" in arguments:
        if config is not None:
            raise ValueError("config was supplied twice")
        config = arguments.pop("config")
        request = FamilyRequest(request.target, arguments)
    if (config is None) == (module is None):
        raise ValueError("provide exactly one of config or module")
    context = _context_from_mesh(mesh, context)
    member = resolve_family(request, context=context)
    result = _build_family_member(request, member, config=config, module=module,
                                 module_fqn=module_fqn, context=context)
    if mesh is not None:
        # Layout discovery imports this module; defer its runtime dependency to avoid a cycle.
        from hyper_parallel.distributed._builder.family_layout import _prepare_family_layout  # pylint: disable=C0415,R0401
        _prepare_family_layout(result, mesh)
    return result


def get_family_build_info(module: nn.Module) -> FamilyBuildInfo | None:
    """Return only framework-created family provenance, or None for ordinary modules.

    Args:
        module: Module whose construction provenance is queried.
    """
    value = getattr(module, "_hp_family_build_info", None)
    return value if isinstance(value, FamilyBuildInfo) else None


def _family_parameter_shapes(model: nn.Module) -> dict[str, tuple[int, ...]]:
    """Return logical parameter shapes even after component construction localized weights."""
    by_identity = {}
    for boundary in collect_family_modules(model):
        layout = boundary.info.layout
        if layout is not None:
            for name, shape in layout.parameter_shapes.items():
                identity = id(boundary.module.get_parameter(name))
                if by_identity.setdefault(identity, shape) != shape:
                    raise ValueError(f"shared family parameter {name!r} has conflicting logical shapes")
    return {name: by_identity.get(id(value), tuple(value.shape))
            for name, value in model.named_parameters(remove_duplicate=False)}


def _module_aliases(model: nn.Module) -> dict[int, tuple[nn.Module, tuple[str, ...]]]:

def _module_aliases(model: nn.Module) -> dict[int, tuple[nn.Module, tuple[str, ...]]]:
    """Inventory registered module identities and all paths, including the root."""
    found: dict[int, tuple[nn.Module, list[str]]] = {}
    for fqn, module in model.named_modules(remove_duplicate=False):
        found.setdefault(id(module), (module, []))[1].append(fqn)
    return {identity: (module, tuple(paths)) for identity, (module, paths) in found.items()}


def collect_family_modules(model: nn.Module) -> tuple[FamilyBoundary, ...]:
    """Discover final family boundaries by object identity, including root and aliases.

    Args:
        model: Final or source model tree for this build.
    """
    boundaries = []
    for module, paths in _module_aliases(model).values():
        info = get_family_build_info(module)
        if info is not None:
            boundaries.append(FamilyBoundary(module, paths, info))
    return tuple(boundaries)


def validate_family_context(model: nn.Module, context: Mapping[str, Any]) -> None:
    """Validate that model preparation has finalized members for the job context.

    Args:
        model: Final or source model tree for this build.
        context: Read-only job context and effective model facts.
    """
    axes = _axis_sizes(_context_from_mesh(None, context))
    for boundary in collect_family_modules(model):
        if boundary.info.layout is not None and boundary.info.layout.applied:
            _context_from_mesh(boundary.info.layout.mesh, context)
        member = boundary.info.resolved
        for axis in ("cp", "ep"):
            if member.axis_sizes[axis] != axes[axis]:
                raise ValueError(
                    f"family at {boundary.fqns} was constructed with {axis}={member.axis_sizes[axis]}, "
                    f"job requires {axes[axis]}; call prepare_model_components with the job context before layout"
                )

def _validate_applied_layout(cached: _FamilyLayout, spec: Any, fqn: str) -> None:
    """Reject post-construction changes to a component's declared computation layout."""
    original = cached.sharding
    if any(name not in original.params or named != original.params[name] for name, named in spec.params.items()):
        raise ValueError(f"layout at {fqn!r} conflicts with its constructed family layout")
    for field_name in ("in_src", "in_dst", "out_src", "out_dst", "out_names", "is_boundary"):
        if getattr(spec, field_name) != getattr(original, field_name):
            raise ValueError(f"{field_name} at {fqn!r} conflicts with its constructed family layout")


def _record_family_layouts(model, plan, mesh, shapes, source_shard_info, *, applied=True):
    """Retain relative declarations once for constructor and model-level application."""
    for boundary in collect_family_modules(model):
        fqn = boundary.fqns[0]
        if fqn not in plan.modules or boundary.info.resolved.sharding is None:
            continue
        existing = boundary.info.layout
        if existing is not None and existing.applied:
            continue
        prefix = fqn + "." if fqn else ""
        names = plan.modules[fqn].params
        relative_shapes = {name: shapes[prefix + name] for name in names}
        source = {name: source_shard_info[prefix + name] for name in names
                  if source_shard_info is not None and prefix + name in source_shard_info}
        layout = _FamilyLayout(mesh, deepcopy(plan.modules[fqn]), plan.parallel_modules.get(fqn),
                              MappingProxyType(relative_shapes), MappingProxyType(source), applied)
        boundary.module._hp_family_build_info = replace(boundary.info, layout=layout)  # pylint: disable=protected-access


def _family_source_shard_info(model: nn.Module) -> dict[str, Any]:
    """Gather constructed component source layouts under their final model paths."""
    result = {}
    for boundary in collect_family_modules(model):
        layout = boundary.info.layout
        if layout is not None and layout.applied:
            for fqn in boundary.fqns:
                result.update({f"{fqn}.{name}" if fqn else name: value
                               for name, value in layout.source_shard_info.items()})
    return result


def _same_mesh(first: Any, second: Any) -> bool:
    """Compare borrowed job views without copying or reconstructing their groups."""
    return (getattr(first, "device_mesh", first) is getattr(second, "device_mesh", second)
            and getattr(first, "fsdp_moe_mesh", None) is getattr(second, "fsdp_moe_mesh", None))


def _family_mesh(model: nn.Module, mesh: Any) -> Any:

def _family_mesh(model: nn.Module, mesh: Any) -> Any:
    """Reuse the complete context when a caller passes its dense DeviceMesh view."""
    layouts = [boundary.info.layout for boundary in collect_family_modules(model)
               if boundary.info.layout is not None and boundary.info.layout.applied]
    if not hasattr(mesh, "device_mesh"):
        mesh = next((layout.mesh for layout in layouts if layout.mesh.device_mesh is mesh), mesh)
    if any(not _same_mesh(layout.mesh, mesh) for layout in layouts):
        raise ValueError("constructed families cannot be applied with a different job mesh")
    return mesh

        )
    else:
        scattered = state_dict

    model.load_state_dict(scattered, strict=options.strict, assign=assign)


# =====================================================================
# Optimizer state dict helpers

    raw_sd: Dict[str, Any],
    saved_id_to_fqn: Dict[int, str],
) -> List[Dict[str, Any]]:
    """Build result param_groups with FQN keys from saved id-based groups."""
    del optimizer
    result_param_groups: List[Dict[str, Any]] = []
    for saved_group in raw_sd["param_groups"]:
        pg: Dict[str, Any] = {}
        for k, v in saved_group.items():

    Returns:
        Tuple of (fqn_to_saved_id, dtensor_state_ids).
    """
    del param_by_fqn
    target_raw_sd = optimizer.state_dict()

    target_id_to_param: Dict[int, nn.Parameter] = {}
    for runtime_group, saved_group in zip(optimizer.param_groups, target_raw_sd["param_groups"]):

def _optimizer_tensor_result(tensor: DTensor, stores_dtensor: bool) -> Union[torch.Tensor, DTensor]:
    """Keep global layout only when the target optimizer stores DTensor values."""
    return tensor if stores_dtensor else tensor.to_local()


def _convert_input_tensor_to_target(
    tensor: torch.Tensor,

    target_device = param.to_local().device

    if isinstance(tensor, DTensor):
        if tensor.device_mesh == mesh and tensor.placements == placements:
            return _optimizer_tensor_result(tensor, stores_dtensor)
        redistributed = tensor.redistribute(mesh, placements)
        return _optimizer_tensor_result(redistributed, stores_dtensor)

    if full_state_dict:
        if cpu_offload:
            tensor = tensor.to(target_device)

    if full_state_dict:
        if cpu_offload:
            tensor = tensor.to(target_device)
        dt = distribute_tensor(tensor, mesh, placements)
        return _optimizer_tensor_result(dt, stores_dtensor)

    if cpu_offload and tensor.device != target_device:
        tensor = tensor.to(target_device)

    cpu_offload: bool,
) -> Optional[Tuple[str, str, torch.Tensor]]:
    """Parse a state.* metadata entry and return (fqn, state_key, tensor) or None."""
    # DCP re-exports these state helpers; defer metadata to avoid that package import cycle.
    from hyper_parallel.core.distributed_checkpoint.metadata import (  # pylint: disable=import-outside-toplevel
        TensorStorageMetadata,
    )

    remainder = meta_key[len("state."):]

    meta_val: Any,
) -> Optional[Tuple[int, str, torch.Tensor]]:
    """Parse a param_group.* metadata entry and return (group_idx, field_name, tensor) or None."""
    # DCP re-exports these state helpers; defer metadata to avoid that package import cycle.
    from hyper_parallel.core.distributed_checkpoint.metadata import (  # pylint: disable=import-outside-toplevel
        TensorStorageMetadata,
    )

    remainder = meta_key[len("param_group."):]

        state.<FQN>.<state_key>  -> TensorStorageMetadata
        param_group.<FQN>.<field> -> TensorStorageMetadata or BytesStorageMetadata
    """
    # DCP re-exports these state helpers; defer metadata to avoid that package import cycle.
    from hyper_parallel.core.distributed_checkpoint.metadata import (  # pylint: disable=import-outside-toplevel
        TensorStorageMetadata,
    )

    flat: Dict[str, Any] = {}

    """Require an explicit forward replacement for CP attention boundaries."""
    if cp_mesh is None or cp_mesh.size() <= 1:
        return
    for fqn, spec in plan.modules.items():
        if fqn in plan.parallel_modules:
            continue
        if not spec.is_boundary or not getattr(spec, "_needs_cp_attn", False):
            continue
        if spec.inner_wrapper is not None or spec.local_compute_fn is not None:
            continue

    from hyper_parallel.distributed.expert_parallel.recipes import (  # pylint: disable=C0415
        EP_ARCHETYPE_SUGGESTIONS,
    )
    for fqn, spec in plan.modules.items():
        if fqn in plan.parallel_modules:
            continue
        if not spec.is_boundary or not getattr(spec, "_ep_size", 0):
            continue
        if spec.local_compute_fn is not None or spec.region_dispatch is False:
            continue

        else create_tp_collective_lowerer(mesh, mesh_dim_names)
    )
    for module_fqn, spec in sorted(
            plan.modules.items(), key=_boundary_post_order_key):
        if module_fqn in plan.parallel_modules:
            continue
        if not spec.is_boundary:
            continue
        module = _resolve_module(model, module_fqn)
        family = get_family_build_info(module)
        if family is not None and family.layout is not None:
            continue
        boundary = PrecompiledBoundary(
            spec,
            mesh,
            mesh_dim_names,

    aliases: list[tuple[str, str]] = field(default_factory=list)


def _full(fqn, name):
    return f"{fqn}.{name}" if fqn else name


def _is_child(fqn, parent):
    return fqn.startswith(parent + ".") if parent else bool(fqn)


def _module_parallel(module, member):
    provider = getattr(module, "parallel_module_spec", None)
    declaration = provider() if provider is not None else None
    if declaration is not None and not isinstance(declaration, ParallelModuleSpec):
        raise TypeError("parallel_module_spec() must return ParallelModuleSpec")
    if member is not None and member.selectors["parallel"] and declaration is None:
        raise ValueError("a parallel family member must declare its explicit execution protocol")
    return declaration


def _add_provider(entries, model, fqn, sharding, parallel):
    module = model.get_submodule(fqn)
    identity = id(module)
    if identity in entries:
        raise ValueError(f"multiple layout/communication owners for module {fqn!r}")
    if not isinstance(sharding, ModuleShardingSpec):
        raise TypeError("layout providers must return ModuleShardingSpec")
    sharding = copy.deepcopy(sharding)
    _normalize_out_fields(sharding)
    entries[identity] = (fqn, module, sharding, parallel)


def _native_providers(entries, model, spec, context, boundaries):
    """Locate model-owned sharding providers for native parallel components."""
    if spec is None or spec.native_parallel is None:
        return
    requests = spec.native_parallel(model=model, context=context)
    all_modules = tuple(model.named_modules(remove_duplicate=False))
    family_subtrees = {id(child) for boundary in boundaries for child in boundary.module.modules()}
    for request in requests:
        if not isinstance(request, NativeParallelBindingRequest):
            raise TypeError("native_parallel must return NativeParallelBindingRequest objects")
        matched = set()
        for pattern in request.match:
            hits = [(fqn, module) for fqn, module in all_modules if fnmatch.fnmatchcase(fqn, pattern)]
            if not hits:
                raise ValueError(f"native binding pattern {pattern!r} matched no module")
            for fqn, module in hits:
                if id(module) in matched:
                    continue
                if _overlaps_family(module, family_subtrees):
                    raise ValueError(f"native binding at {fqn!r} overlaps a family-owned subtree")
                matched.add(id(module))
                layout = request.sharding(module=module, module_fqn=fqn, member=None, context=context)
                _add_provider(entries, model, fqn, layout, request.parallel)


def _overlaps_family(module, family_subtrees):
    return any(id(child) in family_subtrees for child in module.modules())


def _validate_ports(module, spec):
    """Validate named input and output layout declarations against the module interface."""
    injection_fields = ("local_compute_fn", "inner_wrapper", "inner_target", "region_dispatch")
    if any(getattr(spec, name) is not None for name in injection_fields):
        raise ValueError("module-owned layouts cannot install legacy compute injection")
    if spec.params is None or spec.in_src is None or spec.in_dst is None:
        raise ValueError("module-owned providers must declare params, in_src and in_dst explicitly")
    if set(spec.in_src) != set(spec.in_dst):
        raise ValueError("module-owned in_src and in_dst must declare the same input names")
    signature = inspect.signature(module.forward)
    accepts_kwargs = any(item.kind == inspect.Parameter.VAR_KEYWORD for item in signature.parameters.values())
    unknown = set(spec.in_src) - set(signature.parameters)
    if unknown and not accepts_kwargs:
        raise ValueError(f"layout declares unknown forward inputs {sorted(unknown)}")
    if spec.out_names is not None and len(set(spec.out_names)) != len(spec.out_names):
        raise ValueError("output names must be unique")


def _validate_placements(spec, context):
    """Reject unknown mesh axes and incompatible parameter placement declarations."""
    mappings = [*spec.params.values(), *spec.in_src.values(), *spec.in_dst.values()]
    mappings.extend((spec.out_src or {}).values())
    mappings.extend((spec.out_dst or {}).values())
    for named in mappings:
        for axis, placement in named.items():
            if axis not in context.axis_sizes:
                raise ValueError(f"module layouts cannot declare storage/unknown axis {axis!r}")
            if not hasattr(placement, "is_shard"):
                raise TypeError("module providers must return concrete Placement objects")
    if any(isinstance(placement, Partial) for named in spec.params.values() for placement in named.values()):
        raise ValueError("parameter source layouts cannot be Partial")


def _claim_layout_parameters(fqn, parameters, spec, children, owner_by_id, canonical_by_id):
    """Assign each Parameter identity once and verify every alias has the same layout."""
    declared = set()
    for name, placement in tuple(spec.params.items()):
        parameter = parameters.get(name)
        if parameter is None:
            raise ValueError(f"layout parameter {_full(fqn, name)!r} does not exist")
        full_name = _full(fqn, name)
        if any(full_name.startswith(child + ".") for child in children):
            raise ValueError(f"parent layout duplicates nested module ownership at {full_name!r}")
        identity = id(parameter)
        declared.add(identity)
        if identity in owner_by_id:
            if owner_by_id[identity] != placement:
                raise ValueError(f"shared parameter {full_name!r} has conflicting layouts")
            del spec.params[name]
        else:
            owner_by_id[identity] = placement
            canonical_by_id[identity] = full_name
    return declared


def _assign_parameters(result, entries, model):
    """Assign one owned layout per logical parameter while retaining shared aliases."""
    all_parameters = dict(model.named_parameters(remove_duplicate=False))
    owner_by_id = {}
    canonical_by_id = {}
    for fqn, module, spec, _ in entries.values():
        parameters = dict(module.named_parameters(remove_duplicate=False))
        children = [name for name, *_ in entries.values() if _is_child(name, fqn)]
        expected = {id(parameter) for name, parameter in parameters.items()
                    if not any(_full(fqn, name).startswith(child + ".") for child in children)}
        declared = _claim_layout_parameters(fqn, parameters, spec, children, owner_by_id, canonical_by_id)
        if expected - declared:
            missing = [name for name, parameter in parameters.items() if id(parameter) in expected - declared]
            raise ValueError(f"module-owned layout {fqn!r} omits parameters {missing}")
        result.expected_ids.update(expected)
        result.specs[fqn] = spec
        result.required_by_fqn[fqn] = {id(parameters[name]) for name in spec.params}
    for name, parameter in all_parameters.items():
        identity = id(parameter)
        if identity not in owner_by_id:
            continue
        result.owned_names.add(name)
        canonical = canonical_by_id[identity]
        if canonical != name:
            result.aliases.append((canonical, name))


def collect_owned_layouts(model: nn.Module, context: Any, spec: Any = None) -> OwnedLayouts:
    """Discover model/native/family providers independently of generic TP templates.