From e7a3651537d72463ba97c729636a38262397163c Mon Sep 17 00:00:00 2001 From: Xiaoyu Zhang Date: Wed, 26 Aug 2026 18:37:25 -0700 Subject: [PATCH 1/7] Add canonical component manifests Resolve task roles, module paths, and HuggingFace source ownership into one immutable manifest. Switch inspection and optimization-role lookup to the shared metadata without changing graph or weight behavior. Signed-off-by: Xiaoyu Zhang --- src/mobius/_builder.py | 6 +- src/mobius/_component_manifest.py | 135 +++++++++++++++++++++++++ src/mobius/_component_manifest_test.py | 94 +++++++++++++++++ src/mobius/_inspect.py | 28 +++-- src/mobius/tasks/_base.py | 22 +++- 5 files changed, 268 insertions(+), 17 deletions(-) create mode 100644 src/mobius/_component_manifest.py create mode 100644 src/mobius/_component_manifest_test.py diff --git a/src/mobius/_builder.py b/src/mobius/_builder.py index c8b766973..69f1a6589 100644 --- a/src/mobius/_builder.py +++ b/src/mobius/_builder.py @@ -162,12 +162,16 @@ def build_from_module( if prune_prefill_prefix: task = _enable_prefill_prefix_pruning_task(task) resolved_task = get_task(task) + component_manifest = resolved_task.component_manifest() capabilities = ep_registry.require(execution_provider) with build_context(capabilities, dtype): package = resolved_task.build(module, config) for name, model in package.items(): - role = resolved_task.model_roles.get(name) or _MODEL_ROLE_MAP.get(name, "decoder") + descriptor = component_manifest.get(name) + role = ( + descriptor.role if descriptor is not None else _MODEL_ROLE_MAP.get(name, "decoder") + ) optimize_model( model, ep=execution_provider, diff --git a/src/mobius/_component_manifest.py b/src/mobius/_component_manifest.py new file mode 100644 index 000000000..4139f6967 --- /dev/null +++ b/src/mobius/_component_manifest.py @@ -0,0 +1,135 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT License. + +"""Canonical metadata for the components of a model package.""" + +from __future__ import annotations + +__all__ = [ + "ComponentDescriptor", + "ComponentManifest", + "get_hf_component_sources", + "resolve_component_manifest", +] + +import dataclasses +from collections.abc import Iterator, Mapping +from types import MappingProxyType +from typing import TYPE_CHECKING + +if TYPE_CHECKING: + pass + + +@dataclasses.dataclass(frozen=True) +class ComponentDescriptor: + """One package component and all metadata needed to address it. + + Attributes: + name: Key used by :class:`~mobius.ModelPackage`. + module_path: Dotted path from the top-level Mobius module to the + sub-module that constructs this component. The empty string means + the top-level module itself. + role: Optimization role such as ``decoder``, ``encoder``, ``embedding`` + or ``glue``. + source_paths: Runtime HuggingFace ``named_modules()`` paths whose + weights belong to this component. + """ + + name: str + module_path: str + role: str + source_paths: tuple[str, ...] = () + + def __post_init__(self) -> None: + if not self.name: + raise ValueError("component name must not be empty") + if not self.role: + raise ValueError(f"component {self.name!r} must declare a role") + if any(not path for path in self.source_paths): + raise ValueError( + f"component {self.name!r} source_paths must not contain empty paths" + ) + + +@dataclasses.dataclass(frozen=True) +class ComponentManifest(Mapping[str, ComponentDescriptor]): + """Ordered, immutable component metadata keyed by package component name.""" + + components: tuple[ComponentDescriptor, ...] + _by_name: Mapping[str, ComponentDescriptor] = dataclasses.field( + init=False, + repr=False, + compare=False, + ) + + def __post_init__(self) -> None: + by_name: dict[str, ComponentDescriptor] = {} + for component in self.components: + if component.name in by_name: + raise ValueError( + f"component manifest declares {component.name!r} more than once" + ) + by_name[component.name] = component + object.__setattr__(self, "_by_name", MappingProxyType(by_name)) + + def __getitem__(self, name: str) -> ComponentDescriptor: + return self._by_name[name] + + def __iter__(self) -> Iterator[str]: + return iter(self._by_name) + + def __len__(self) -> int: + return len(self._by_name) + + @property + def names(self) -> tuple[str, ...]: + """Component names in task declaration order.""" + return tuple(self._by_name) + + +def get_hf_component_sources( + module_class: type, + model_type: str, + hf_config: object, +) -> dict[str, tuple[str, ...]]: + """Read runtime HuggingFace component paths from a registered model class.""" + resolver = getattr(module_class, "get_hf_component_sources", None) + if resolver is not None: + resolved = resolver(model_type=model_type, hf_config=hf_config) + else: + resolved = getattr(module_class, "HF_COMPONENT_SOURCES", {}) + return {name: tuple(paths) for name, paths in resolved.items()} + + +def resolve_component_manifest( + task: object, + *, + module_class: type | None = None, + model_type: str | None = None, + hf_config: object | None = None, +) -> ComponentManifest: + """Combine task roles/paths and model source ownership into one manifest.""" + roles = dict(getattr(task, "model_roles", {}) or {}) + component_spec = getattr(task, "components", None) + module_paths = dict(component_spec.items()) if component_spec is not None else {} + + component_sources: dict[str, tuple[str, ...]] = {} + if module_class is not None and model_type is not None and hf_config is not None: + component_sources = get_hf_component_sources( + module_class, + model_type, + hf_config, + ) + + ordered_names = tuple(dict.fromkeys((*roles, *module_paths))) + descriptors = tuple( + ComponentDescriptor( + name=name, + module_path=module_paths.get(name, "" if name == "model" else name), + role=roles.get(name, "decoder"), + source_paths=component_sources.get(name, ()), + ) + for name in ordered_names + ) + return ComponentManifest(descriptors) diff --git a/src/mobius/_component_manifest_test.py b/src/mobius/_component_manifest_test.py new file mode 100644 index 000000000..6765445a5 --- /dev/null +++ b/src/mobius/_component_manifest_test.py @@ -0,0 +1,94 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT License. + +"""Tests for canonical component manifest resolution.""" + +from __future__ import annotations + +from typing import ClassVar + +import pytest + +from mobius._component_manifest import ( + ComponentDescriptor, + ComponentManifest, + resolve_component_manifest, +) +from mobius.tasks import ComponentSpec + + +class _Task: + model_roles: ClassVar[dict[str, str]] = { + "decoder": "decoder", + "vision_encoder": "encoder", + "embedding": "embedding", + } + components = ComponentSpec( + decoder="language", + vision_encoder="vision.tower", + embedding="embedding", + ) + + +class _Model: + HF_COMPONENT_SOURCES: ClassVar[dict[str, tuple[str, ...]]] = { + "decoder": ("model.language_model.layers", "lm_head"), + "vision_encoder": ("model.vision_tower", "model.projector"), + "embedding": ("model.language_model.embed_tokens",), + } + + +def test_manifest_combines_task_and_model_metadata(): + manifest = resolve_component_manifest( + _Task(), + module_class=_Model, + model_type="test", + hf_config=object(), + ) + + assert manifest.names == ("decoder", "vision_encoder", "embedding") + assert manifest["decoder"] == ComponentDescriptor( + name="decoder", + module_path="language", + role="decoder", + source_paths=("model.language_model.layers", "lm_head"), + ) + assert manifest["vision_encoder"].module_path == "vision.tower" + assert manifest["vision_encoder"].role == "encoder" + + +def test_dynamic_source_resolver_is_authoritative(): + class _DynamicModel: + @classmethod + def get_hf_component_sources(cls, *, model_type, hf_config): + assert model_type == "dynamic" + assert hf_config == "config" + return {"decoder": ("resolved.decoder",)} + + manifest = resolve_component_manifest( + _Task(), + module_class=_DynamicModel, + model_type="dynamic", + hf_config="config", + ) + + assert manifest["decoder"].source_paths == ("resolved.decoder",) + assert manifest["vision_encoder"].source_paths == () + + +def test_single_component_uses_root_module_path(): + class _SingleTask: + model_roles: ClassVar[dict[str, str]] = {"model": "encoder"} + components = None + + manifest = resolve_component_manifest(_SingleTask()) + + assert manifest["model"].module_path == "" + assert manifest["model"].role == "encoder" + + +def test_duplicate_component_names_are_rejected(): + component = ComponentDescriptor("decoder", "decoder", "decoder") + + with pytest.raises(ValueError, match="more than once"): + ComponentManifest((component, component)) diff --git a/src/mobius/_inspect.py b/src/mobius/_inspect.py index abe7ae0f1..9b023c260 100644 --- a/src/mobius/_inspect.py +++ b/src/mobius/_inspect.py @@ -114,10 +114,9 @@ def _get_hf_component_sources( hf_config: object, ) -> dict[str, tuple[str, ...]]: """Read runtime HuggingFace component paths from a registered model class.""" - resolver = getattr(module_class, "get_hf_component_sources", None) - if resolver is not None: - return resolver(model_type=model_type, hf_config=hf_config) - return getattr(module_class, "HF_COMPONENT_SOURCES", {}) + from mobius._component_manifest import get_hf_component_sources + + return get_hf_component_sources(module_class, model_type, hf_config) def inspect_components( @@ -151,23 +150,22 @@ def inspect_components( model_id, task, trust_remote_code ) task_obj = get_task(resolved_task) - roles = task_obj.model_roles or {} - - # Runtime HF paths are owned by the registered model class. Most classes - # declare a fixed ``HF_COMPONENT_SOURCES`` mapping; classes shared by - # several HF layouts can resolve paths from the already-loaded config. - component_sources: dict[str, tuple[str, ...]] = {} + module_class = None if model_type is not None and hf_config is not None and model_type in registry: module_class = registry.get(model_type) - component_sources = _get_hf_component_sources(module_class, model_type, hf_config) + manifest = task_obj.component_manifest( + module_class=module_class, + model_type=model_type, + hf_config=hf_config, + ) components = [ ComponentInfo( - name=name, - role=role, - source_paths=tuple(component_sources.get(name, ())), + name=component.name, + role=component.role, + source_paths=component.source_paths, ) - for name, role in roles.items() + for component in manifest.values() ] logger.debug( "inspect_components(%s): task=%s components=%s", diff --git a/src/mobius/tasks/_base.py b/src/mobius/tasks/_base.py index 3547314d1..f847b5d72 100644 --- a/src/mobius/tasks/_base.py +++ b/src/mobius/tasks/_base.py @@ -6,7 +6,7 @@ from __future__ import annotations from abc import ABC, abstractmethod -from typing import ClassVar +from typing import TYPE_CHECKING, ClassVar import onnx_ir as ir from onnxscript import GraphBuilder, nn @@ -16,6 +16,9 @@ from mobius._constants import OPSET_VERSION from mobius._model_package import ModelPackage +if TYPE_CHECKING: + from mobius._component_manifest import ComponentManifest + class ComponentSpec: """Declares which sub-module attributes a multi-component task requires. @@ -159,6 +162,23 @@ def build(self, module, config): #: exist on the module before building begins. components: ClassVar[ComponentSpec | None] = None + def component_manifest( + self, + *, + module_class: type | None = None, + model_type: str | None = None, + hf_config: object | None = None, + ) -> ComponentManifest: + """Resolve canonical metadata for every component produced by this task.""" + from mobius._component_manifest import resolve_component_manifest + + return resolve_component_manifest( + self, + module_class=module_class, + model_type=model_type, + hf_config=hf_config, + ) + def _validate_components(self, module: nn.Module) -> None: """Validate that *module* exposes all attributes declared in :attr:`components`. From 6d03d13685866e07c50cb54b3b222f966ae480d4 Mon Sep 17 00:00:00 2001 From: Xiaoyu Zhang Date: Wed, 26 Aug 2026 18:40:48 -0700 Subject: [PATCH 2/7] Map component modules to source names Derive candidate HuggingFace module names from component-local paths so later loader stages can apply exact and regex quantization exclusions consistently. Signed-off-by: Xiaoyu Zhang --- src/mobius/_component_manifest.py | 28 ++++++++++++++++++++++++++ src/mobius/_component_manifest_test.py | 19 +++++++++++++++++ 2 files changed, 47 insertions(+) diff --git a/src/mobius/_component_manifest.py b/src/mobius/_component_manifest.py index 4139f6967..5d3509614 100644 --- a/src/mobius/_component_manifest.py +++ b/src/mobius/_component_manifest.py @@ -51,6 +51,34 @@ def __post_init__(self) -> None: f"component {self.name!r} source_paths must not contain empty paths" ) + def source_module_names(self, local_module_path: str) -> tuple[str, ...]: + """Candidate HuggingFace names for a component-local module path. + + Source roots and Mobius paths commonly share an anchor segment even + when their prefixes differ. For example, source root + ``model.language_model.layers`` and local path + ``model.layers.0.self_attn.q_proj`` share ``layers`` and resolve to + ``model.language_model.layers.0.self_attn.q_proj``. + """ + if not local_module_path: + return self.source_paths + + local_parts = local_module_path.split(".") + candidates = [local_module_path] + for source_path in self.source_paths: + source_parts = source_path.split(".") + anchor = source_parts[-1] + anchor_indices = [ + index for index, part in enumerate(local_parts) if part == anchor + ] + if anchor_indices: + for index in anchor_indices: + suffix = local_parts[index + 1 :] + candidates.append(".".join((*source_parts, *suffix))) + else: + candidates.append(".".join((*source_parts, *local_parts))) + return tuple(dict.fromkeys(candidates)) + @dataclasses.dataclass(frozen=True) class ComponentManifest(Mapping[str, ComponentDescriptor]): diff --git a/src/mobius/_component_manifest_test.py b/src/mobius/_component_manifest_test.py index 6765445a5..b3383a0a0 100644 --- a/src/mobius/_component_manifest_test.py +++ b/src/mobius/_component_manifest_test.py @@ -76,6 +76,25 @@ def get_hf_component_sources(cls, *, model_type, hf_config): assert manifest["vision_encoder"].source_paths == () +def test_descriptor_maps_local_path_to_huggingface_source_name(): + descriptor = ComponentDescriptor( + name="decoder", + module_path="decoder", + role="decoder", + source_paths=("model.language_model.layers", "lm_head"), + ) + + assert descriptor.source_module_names("model.layers.0.per_layer_input_gate") == ( + "model.layers.0.per_layer_input_gate", + "model.language_model.layers.0.per_layer_input_gate", + "lm_head.model.layers.0.per_layer_input_gate", + ) + assert descriptor.source_module_names("lm_head") == ( + "lm_head", + "model.language_model.layers.lm_head", + ) + + def test_single_component_uses_root_module_path(): class _SingleTask: model_roles: ClassVar[dict[str, str]] = {"model": "encoder"} From 604e5c13c355106c533e42daa26a10bd89432936 Mon Sep 17 00:00:00 2001 From: Xiaoyu Zhang Date: Wed, 26 Aug 2026 19:20:39 -0700 Subject: [PATCH 3/7] Avoid synthetic component source paths Only derive HuggingFace module candidates when a declared source root shares an anchor with the component-local path, preventing unrelated roots such as lm_head from matching every decoder module. Signed-off-by: Xiaoyu Zhang --- src/mobius/_component_manifest.py | 2 -- src/mobius/_component_manifest_test.py | 8 ++------ 2 files changed, 2 insertions(+), 8 deletions(-) diff --git a/src/mobius/_component_manifest.py b/src/mobius/_component_manifest.py index 5d3509614..e1d80f8e6 100644 --- a/src/mobius/_component_manifest.py +++ b/src/mobius/_component_manifest.py @@ -75,8 +75,6 @@ def source_module_names(self, local_module_path: str) -> tuple[str, ...]: for index in anchor_indices: suffix = local_parts[index + 1 :] candidates.append(".".join((*source_parts, *suffix))) - else: - candidates.append(".".join((*source_parts, *local_parts))) return tuple(dict.fromkeys(candidates)) diff --git a/src/mobius/_component_manifest_test.py b/src/mobius/_component_manifest_test.py index b3383a0a0..f0a43d4f1 100644 --- a/src/mobius/_component_manifest_test.py +++ b/src/mobius/_component_manifest_test.py @@ -84,15 +84,11 @@ def test_descriptor_maps_local_path_to_huggingface_source_name(): source_paths=("model.language_model.layers", "lm_head"), ) - assert descriptor.source_module_names("model.layers.0.per_layer_input_gate") == ( + assert descriptor.source_module_names("model.layers.0.per_layer_input_gate" ) == ( "model.layers.0.per_layer_input_gate", "model.language_model.layers.0.per_layer_input_gate", - "lm_head.model.layers.0.per_layer_input_gate", - ) - assert descriptor.source_module_names("lm_head") == ( - "lm_head", - "model.language_model.layers.lm_head", ) + assert descriptor.source_module_names("lm_head") == ("lm_head",) def test_single_component_uses_root_module_path(): From b44036a161bc9dc45ba9363ef02878ccf8f0b017 Mon Sep 17 00:00:00 2001 From: Xiaoyu Zhang Date: Wed, 26 Aug 2026 19:41:20 -0700 Subject: [PATCH 4/7] Support explicit component source aliases Allow model declarations to map component-local module prefixes to HuggingFace source prefixes when structural anchor inference is insufficient. Signed-off-by: Xiaoyu Zhang --- src/mobius/_component_manifest.py | 20 ++++++++++++++++++++ src/mobius/_component_manifest_test.py | 12 +++++++++++- 2 files changed, 31 insertions(+), 1 deletion(-) diff --git a/src/mobius/_component_manifest.py b/src/mobius/_component_manifest.py index e1d80f8e6..fb95acbc4 100644 --- a/src/mobius/_component_manifest.py +++ b/src/mobius/_component_manifest.py @@ -34,12 +34,15 @@ class ComponentDescriptor: or ``glue``. source_paths: Runtime HuggingFace ``named_modules()`` paths whose weights belong to this component. + source_path_aliases: Pairs of ``(local_prefix, source_prefix)`` for + component paths that cannot be aligned by a shared anchor segment. """ name: str module_path: str role: str source_paths: tuple[str, ...] = () + source_path_aliases: tuple[tuple[str, str], ...] = () def __post_init__(self) -> None: if not self.name: @@ -50,6 +53,11 @@ def __post_init__(self) -> None: raise ValueError( f"component {self.name!r} source_paths must not contain empty paths" ) + if any(not local or not source for local, source in self.source_path_aliases): + raise ValueError( + f"component {self.name!r} source_path_aliases must contain " + "non-empty local/source prefixes" + ) def source_module_names(self, local_module_path: str) -> tuple[str, ...]: """Candidate HuggingFace names for a component-local module path. @@ -65,6 +73,12 @@ def source_module_names(self, local_module_path: str) -> tuple[str, ...]: local_parts = local_module_path.split(".") candidates = [local_module_path] + for local_prefix, source_prefix in self.source_path_aliases: + if local_module_path == local_prefix: + candidates.append(source_prefix) + elif local_module_path.startswith(f"{local_prefix}."): + suffix = local_module_path[len(local_prefix) + 1 :] + candidates.append(f"{source_prefix}.{suffix}") for source_path in self.source_paths: source_parts = source_path.split(".") anchor = source_parts[-1] @@ -141,12 +155,17 @@ def resolve_component_manifest( module_paths = dict(component_spec.items()) if component_spec is not None else {} component_sources: dict[str, tuple[str, ...]] = {} + component_aliases: dict[str, tuple[tuple[str, str], ...]] = {} if module_class is not None and model_type is not None and hf_config is not None: component_sources = get_hf_component_sources( module_class, model_type, hf_config, ) + raw_aliases = getattr(module_class, "HF_COMPONENT_MODULE_ALIASES", {}) + component_aliases = { + name: tuple(aliases.items()) for name, aliases in raw_aliases.items() + } ordered_names = tuple(dict.fromkeys((*roles, *module_paths))) descriptors = tuple( @@ -155,6 +174,7 @@ def resolve_component_manifest( module_path=module_paths.get(name, "" if name == "model" else name), role=roles.get(name, "decoder"), source_paths=component_sources.get(name, ()), + source_path_aliases=component_aliases.get(name, ()), ) for name in ordered_names ) diff --git a/src/mobius/_component_manifest_test.py b/src/mobius/_component_manifest_test.py index f0a43d4f1..f81df3bce 100644 --- a/src/mobius/_component_manifest_test.py +++ b/src/mobius/_component_manifest_test.py @@ -36,6 +36,12 @@ class _Model: "vision_encoder": ("model.vision_tower", "model.projector"), "embedding": ("model.language_model.embed_tokens",), } + HF_COMPONENT_MODULE_ALIASES: ClassVar[dict[str, dict[str, str]]] = { + "vision_encoder": { + "encoder": "model.vision_tower", + "projector": "model.projector", + } + } def test_manifest_combines_task_and_model_metadata(): @@ -55,6 +61,10 @@ def test_manifest_combines_task_and_model_metadata(): ) assert manifest["vision_encoder"].module_path == "vision.tower" assert manifest["vision_encoder"].role == "encoder" + assert manifest["vision_encoder"].source_module_names("encoder.layers.0.q_proj") == ( + "encoder.layers.0.q_proj", + "model.vision_tower.layers.0.q_proj", + ) def test_dynamic_source_resolver_is_authoritative(): @@ -84,7 +94,7 @@ def test_descriptor_maps_local_path_to_huggingface_source_name(): source_paths=("model.language_model.layers", "lm_head"), ) - assert descriptor.source_module_names("model.layers.0.per_layer_input_gate" ) == ( + assert descriptor.source_module_names("model.layers.0.per_layer_input_gate") == ( "model.layers.0.per_layer_input_gate", "model.language_model.layers.0.per_layer_input_gate", ) From 1f47daad54dc954f79733045f5395a0a9ef2520c Mon Sep 17 00:00:00 2001 From: Xiaoyu Zhang Date: Thu, 27 Aug 2026 14:05:07 -0700 Subject: [PATCH 5/7] Clarify component manifest field semantics Rename module_path to module_attribute_path and document task-defined optimization roles so callers cannot confuse Python module ownership with package or checkpoint prefixes. Signed-off-by: Xiaoyu Zhang --- src/mobius/_component_manifest.py | 17 +++++++++++------ src/mobius/_component_manifest_test.py | 8 ++++---- 2 files changed, 15 insertions(+), 10 deletions(-) diff --git a/src/mobius/_component_manifest.py b/src/mobius/_component_manifest.py index fb95acbc4..61babcfb7 100644 --- a/src/mobius/_component_manifest.py +++ b/src/mobius/_component_manifest.py @@ -27,11 +27,13 @@ class ComponentDescriptor: Attributes: name: Key used by :class:`~mobius.ModelPackage`. - module_path: Dotted path from the top-level Mobius module to the + module_attribute_path: Dotted Python attribute path from the root + :class:`onnxscript.nn.Module` passed to ``task.build()`` to the sub-module that constructs this component. The empty string means - the top-level module itself. - role: Optimization role such as ``decoder``, ``encoder``, ``embedding`` - or ``glue``. + the root module itself. This is not a package key or checkpoint + prefix. + role: Task-defined optimization category. Current roles include + ``decoder``, ``encoder``, ``vision``, ``embedding``, and ``glue``. source_paths: Runtime HuggingFace ``named_modules()`` paths whose weights belong to this component. source_path_aliases: Pairs of ``(local_prefix, source_prefix)`` for @@ -39,7 +41,7 @@ class ComponentDescriptor: """ name: str - module_path: str + module_attribute_path: str role: str source_paths: tuple[str, ...] = () source_path_aliases: tuple[tuple[str, str], ...] = () @@ -171,7 +173,10 @@ def resolve_component_manifest( descriptors = tuple( ComponentDescriptor( name=name, - module_path=module_paths.get(name, "" if name == "model" else name), + module_attribute_path=module_paths.get( + name, + "" if name == "model" else name, + ), role=roles.get(name, "decoder"), source_paths=component_sources.get(name, ()), source_path_aliases=component_aliases.get(name, ()), diff --git a/src/mobius/_component_manifest_test.py b/src/mobius/_component_manifest_test.py index f81df3bce..23985f691 100644 --- a/src/mobius/_component_manifest_test.py +++ b/src/mobius/_component_manifest_test.py @@ -55,11 +55,11 @@ def test_manifest_combines_task_and_model_metadata(): assert manifest.names == ("decoder", "vision_encoder", "embedding") assert manifest["decoder"] == ComponentDescriptor( name="decoder", - module_path="language", + module_attribute_path="language", role="decoder", source_paths=("model.language_model.layers", "lm_head"), ) - assert manifest["vision_encoder"].module_path == "vision.tower" + assert manifest["vision_encoder"].module_attribute_path == "vision.tower" assert manifest["vision_encoder"].role == "encoder" assert manifest["vision_encoder"].source_module_names("encoder.layers.0.q_proj") == ( "encoder.layers.0.q_proj", @@ -89,7 +89,7 @@ def get_hf_component_sources(cls, *, model_type, hf_config): def test_descriptor_maps_local_path_to_huggingface_source_name(): descriptor = ComponentDescriptor( name="decoder", - module_path="decoder", + module_attribute_path="decoder", role="decoder", source_paths=("model.language_model.layers", "lm_head"), ) @@ -108,7 +108,7 @@ class _SingleTask: manifest = resolve_component_manifest(_SingleTask()) - assert manifest["model"].module_path == "" + assert manifest["model"].module_attribute_path == "" assert manifest["model"].role == "encoder" From 5cc1af13356990a151edb53b308e3ba2017c7a13 Mon Sep 17 00:00:00 2001 From: Xiaoyu Zhang Date: Wed, 26 Aug 2026 18:45:14 -0700 Subject: [PATCH 6/7] Add typed checkpoint weight records Group affine packed sidecars into component-owned logical records and introduce a quantization codec registry over the existing Olive, GPTQ, and AWQ normalization helpers. The active loader remains unchanged. Signed-off-by: Xiaoyu Zhang --- src/mobius/weights/__init__.py | 28 ++++ src/mobius/weights/_codecs.py | 199 ++++++++++++++++++++++++++++ src/mobius/weights/_codecs_test.py | 132 ++++++++++++++++++ src/mobius/weights/_records.py | 121 +++++++++++++++++ src/mobius/weights/_records_test.py | 47 +++++++ 5 files changed, 527 insertions(+) create mode 100644 src/mobius/weights/__init__.py create mode 100644 src/mobius/weights/_codecs.py create mode 100644 src/mobius/weights/_codecs_test.py create mode 100644 src/mobius/weights/_records.py create mode 100644 src/mobius/weights/_records_test.py diff --git a/src/mobius/weights/__init__.py b/src/mobius/weights/__init__.py new file mode 100644 index 000000000..a3dd7744d --- /dev/null +++ b/src/mobius/weights/__init__.py @@ -0,0 +1,28 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT License. + +"""Typed checkpoint records and quantization format codecs.""" + +from __future__ import annotations + +from mobius.weights._codecs import ( + QuantizationCodec, + QuantizationCodecRegistry, + codec_registry, +) +from mobius.weights._records import ( + FloatWeight, + PackedWeight, + WeightBundle, + WeightRecord, +) + +__all__ = [ + "FloatWeight", + "PackedWeight", + "QuantizationCodec", + "QuantizationCodecRegistry", + "WeightBundle", + "WeightRecord", + "codec_registry", +] diff --git a/src/mobius/weights/_codecs.py b/src/mobius/weights/_codecs.py new file mode 100644 index 000000000..9a796b6a9 --- /dev/null +++ b/src/mobius/weights/_codecs.py @@ -0,0 +1,199 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT License. + +"""Quantization format codecs for existing packed checkpoint weights.""" + +from __future__ import annotations + +__all__ = [ + "QuantizationCodec", + "QuantizationCodecRegistry", + "codec_registry", +] + +from collections.abc import Mapping +from typing import Protocol + +import torch + +from mobius._component_manifest import ComponentDescriptor +from mobius._configs import QuantizationConfig +from mobius._weight_utils import preprocess_quantized_weights +from mobius.weights._records import ( + FloatWeight, + PackedWeight, + WeightBundle, + WeightRecord, +) + + +class QuantizationCodec(Protocol): + """Groups and normalizes one producer's existing packed weight layout.""" + + method: str + + def group( + self, + component: ComponentDescriptor, + state_dict: Mapping[str, torch.Tensor], + config: QuantizationConfig, + ) -> WeightBundle: + """Group checkpoint sidecars into typed logical records.""" + ... + + def normalize( + self, + record: WeightRecord, + config: QuantizationConfig, + ) -> dict[str, torch.Tensor]: + """Convert one packed record to Mobius's canonical parameter layout.""" + ... + + +class QuantizationCodecRegistry: + """Registry keyed by serialized ``quant_method``.""" + + def __init__(self) -> None: + self._codecs: dict[str, QuantizationCodec] = {} + + def register(self, codec: QuantizationCodec) -> None: + if not codec.method: + raise ValueError("quantization codec method must not be empty") + if codec.method in self._codecs: + raise ValueError(f"quantization codec {codec.method!r} is already registered") + self._codecs[codec.method] = codec + + def get(self, method: str) -> QuantizationCodec: + try: + return self._codecs[method] + except KeyError: + raise KeyError( + f"No quantization codec registered for {method!r}. " + f"Available methods: {sorted(self._codecs)}" + ) from None + + def __contains__(self, method: str) -> bool: + return method in self._codecs + + +class _LegacyAffineCodec: + """Typed facade over the existing Olive/GPTQ/AWQ normalization helpers.""" + + def __init__(self, method: str): + self.method = method + + @staticmethod + def _logical_name(qweight_key: str) -> tuple[str, str, str, str | None]: + if qweight_key.endswith("_qweight"): + stem = qweight_key[: -len("_qweight")] + logical_name = stem if stem.endswith(".weight") else f"{stem}.weight" + return ( + logical_name, + f"{stem}_scales", + qweight_key, + f"{stem}_qzeros", + ) + if qweight_key.endswith(".qweight"): + stem = qweight_key[: -len(".qweight")] + return ( + f"{stem}.weight", + f"{stem}.scales", + qweight_key, + f"{stem}.qzeros", + ) + raise ValueError(f"{qweight_key!r} is not a packed qweight key") + + def group( + self, + component: ComponentDescriptor, + state_dict: Mapping[str, torch.Tensor], + config: QuantizationConfig, + ) -> WeightBundle: + if config.quant_method != self.method: + raise ValueError( + f"codec {self.method!r} cannot group quant_method {config.quant_method!r}" + ) + + records: dict[str, WeightRecord] = {} + consumed: set[str] = set() + qweight_keys = sorted( + key for key in state_dict if key.endswith(("_qweight", ".qweight")) + ) + for qweight_key in qweight_keys: + logical_name, scales_key, _, zero_points_key = self._logical_name(qweight_key) + if scales_key not in state_dict: + raise ValueError( + f"Packed weight {qweight_key!r} is missing scales {scales_key!r}" + ) + zero_points = state_dict.get(zero_points_key) + if not config.sym and zero_points is None: + raise ValueError( + f"Asymmetric packed weight {qweight_key!r} is missing " + f"zero points {zero_points_key!r}" + ) + storage = PackedWeight( + qweight=state_dict[qweight_key], + scales=state_dict[scales_key], + zero_points=zero_points, + qweight_key=qweight_key, + scales_key=scales_key, + zero_points_key=zero_points_key if zero_points is not None else None, + method=self.method, + ) + if logical_name in records: + raise ValueError( + f"Checkpoint declares logical weight {logical_name!r} more than once" + ) + records[logical_name] = WeightRecord( + name=logical_name, + component=component.name, + storage=storage, + ) + consumed.update(storage.source_keys) + + orphan_sidecars = sorted( + key + for key in state_dict + if key.endswith(("_scales", "_qzeros", ".scales", ".qzeros")) + and key not in consumed + ) + if orphan_sidecars: + raise ValueError( + f"Packed checkpoint sidecars have no matching qweight: {orphan_sidecars}" + ) + + for key, value in state_dict.items(): + if key in consumed: + continue + if key in records: + raise ValueError(f"Float and packed checkpoint values both target {key!r}") + records[key] = WeightRecord( + name=key, + component=component.name, + storage=FloatWeight(value=value, source_key=key), + ) + return WeightBundle(component=component, records=records) + + def normalize( + self, + record: WeightRecord, + config: QuantizationConfig, + ) -> dict[str, torch.Tensor]: + if not isinstance(record.storage, PackedWeight): + raise TypeError(f"weight {record.name!r} is not packed") + if record.storage.method != self.method: + raise ValueError( + f"weight {record.name!r} uses {record.storage.method!r}, " + f"not codec {self.method!r}" + ) + return preprocess_quantized_weights( + record.storage.as_state_dict(), + config, + tie_embeddings=False, + qmoe_target_path=None, + ) + + +codec_registry = QuantizationCodecRegistry() +for _method in ("olive", "gptq", "awq"): + codec_registry.register(_LegacyAffineCodec(_method)) diff --git a/src/mobius/weights/_codecs_test.py b/src/mobius/weights/_codecs_test.py new file mode 100644 index 000000000..338fbbcf3 --- /dev/null +++ b/src/mobius/weights/_codecs_test.py @@ -0,0 +1,132 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT License. + +"""Tests for typed quantization format codecs.""" + +from __future__ import annotations + +import pytest +import torch + +from mobius._component_manifest import ComponentDescriptor +from mobius._configs import QuantizationConfig +from mobius.weights import PackedWeight, QuantizationCodecRegistry, codec_registry + + +def _component() -> ComponentDescriptor: + return ComponentDescriptor( + name="decoder", + module_path="decoder", + role="decoder", + source_paths=("model.layers",), + ) + + +def _config(method: str = "olive", *, sym: bool = True) -> QuantizationConfig: + return QuantizationConfig( + bits=4, + group_size=16, + quant_method=method, + sym=sym, + ) + + +def test_groups_olive_sidecars_into_one_record(): + state_dict = { + "model.q_proj.weight_qweight": torch.zeros(32, 32, dtype=torch.uint8), + "model.q_proj.weight_scales": torch.ones(32, 4), + "model.norm.weight": torch.ones(64), + } + + bundle = codec_registry.get("olive").group( + _component(), + state_dict, + _config(), + ) + + record = bundle["model.q_proj.weight"] + assert isinstance(record.storage, PackedWeight) + assert record.source_keys == ( + "model.q_proj.weight_qweight", + "model.q_proj.weight_scales", + ) + assert bundle["model.norm.weight"].is_quantized is False + + +def test_groups_gptq_dotted_sidecars(): + state_dict = { + "model.q_proj.qweight": torch.zeros(8, 32, dtype=torch.int32), + "model.q_proj.scales": torch.ones(4, 32), + } + + bundle = codec_registry.get("gptq").group( + _component(), + state_dict, + _config("gptq"), + ) + + assert bundle["model.q_proj.weight"].is_quantized is True + + +def test_rejects_missing_scales(): + state_dict = { + "model.q_proj.weight_qweight": torch.zeros(32, 32, dtype=torch.uint8), + } + + with pytest.raises(ValueError, match="missing scales"): + codec_registry.get("olive").group( + _component(), + state_dict, + _config(), + ) + + +def test_rejects_missing_asymmetric_zero_points(): + state_dict = { + "model.q_proj.weight_qweight": torch.zeros(32, 32, dtype=torch.uint8), + "model.q_proj.weight_scales": torch.ones(32, 4), + } + + with pytest.raises(ValueError, match="missing zero points"): + codec_registry.get("olive").group( + _component(), + state_dict, + _config(sym=False), + ) + + +def test_rejects_orphan_sidecars(): + with pytest.raises(ValueError, match="no matching qweight"): + codec_registry.get("olive").group( + _component(), + {"model.q_proj.weight_scales": torch.ones(32, 4)}, + _config(), + ) + + +def test_compatibility_normalizer_uses_existing_packer(): + state_dict = { + "model.q_proj.weight_qweight": torch.zeros(32, 32, dtype=torch.uint8), + "model.q_proj.weight_scales": torch.ones(32, 4), + } + codec = codec_registry.get("olive") + record = codec.group(_component(), state_dict, _config())["model.q_proj.weight"] + + normalized = codec.normalize(record, _config()) + + assert normalized["model.q_proj.weight"].shape == (32, 4, 8) + assert normalized["model.q_proj.scales"].shape == (32, 4) + + +def test_registry_rejects_duplicate_method(): + registry = QuantizationCodecRegistry() + codec = codec_registry.get("olive") + registry.register(codec) + + with pytest.raises(ValueError, match="already registered"): + registry.register(codec) + + +def test_registry_reports_unknown_method(): + with pytest.raises(KeyError, match="Available methods"): + QuantizationCodecRegistry().get("unknown") diff --git a/src/mobius/weights/_records.py b/src/mobius/weights/_records.py new file mode 100644 index 000000000..cc274f576 --- /dev/null +++ b/src/mobius/weights/_records.py @@ -0,0 +1,121 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT License. + +"""Typed logical weights used between checkpoint readers and model adapters.""" + +from __future__ import annotations + +__all__ = [ + "FloatWeight", + "PackedWeight", + "WeightBundle", + "WeightRecord", +] + +import dataclasses +from collections.abc import Iterator, Mapping +from types import MappingProxyType + +import torch + +from mobius._component_manifest import ComponentDescriptor + + +@dataclasses.dataclass(frozen=True) +class FloatWeight: + """One ordinary floating-point checkpoint tensor.""" + + value: torch.Tensor + source_key: str + + +@dataclasses.dataclass(frozen=True) +class PackedWeight: + """One logical affine-quantized weight grouped from checkpoint sidecars.""" + + qweight: torch.Tensor + scales: torch.Tensor + zero_points: torch.Tensor | None + qweight_key: str + scales_key: str + zero_points_key: str | None + method: str + + @property + def source_keys(self) -> tuple[str, ...]: + """Checkpoint keys consumed by this logical packed weight.""" + keys = [self.qweight_key, self.scales_key] + if self.zero_points_key is not None: + keys.append(self.zero_points_key) + return tuple(keys) + + def as_state_dict(self) -> dict[str, torch.Tensor]: + """Reconstruct the source sidecars for a compatibility codec.""" + tensors = { + self.qweight_key: self.qweight, + self.scales_key: self.scales, + } + if self.zero_points_key is not None and self.zero_points is not None: + tensors[self.zero_points_key] = self.zero_points + return tensors + + +WeightStorage = FloatWeight | PackedWeight + + +@dataclasses.dataclass(frozen=True) +class WeightRecord: + """A named logical weight owned by exactly one package component.""" + + name: str + component: str + storage: WeightStorage + + @property + def is_quantized(self) -> bool: + """Whether this record stores an existing packed weight.""" + return isinstance(self.storage, PackedWeight) + + @property + def source_keys(self) -> tuple[str, ...]: + """Checkpoint keys represented by this record.""" + if isinstance(self.storage, FloatWeight): + return (self.storage.source_key,) + return self.storage.source_keys + + +@dataclasses.dataclass(frozen=True) +class WeightBundle(Mapping[str, WeightRecord]): + """Immutable records routed to one component descriptor.""" + + component: ComponentDescriptor + records: Mapping[str, WeightRecord] + + def __post_init__(self) -> None: + for name, record in self.records.items(): + if name != record.name: + raise ValueError( + f"weight bundle key {name!r} does not match record name {record.name!r}" + ) + if record.component != self.component.name: + raise ValueError( + f"weight {name!r} belongs to component " + f"{record.component!r}, expected {self.component.name!r}" + ) + object.__setattr__(self, "records", MappingProxyType(dict(self.records))) + + def __getitem__(self, name: str) -> WeightRecord: + return self.records[name] + + def __iter__(self) -> Iterator[str]: + return iter(self.records) + + def __len__(self) -> int: + return len(self.records) + + @property + def source_keys(self) -> frozenset[str]: + """All checkpoint keys represented by this bundle.""" + return frozenset( + source_key for record in self.records.values() for source_key in record.source_keys + ) diff --git a/src/mobius/weights/_records_test.py b/src/mobius/weights/_records_test.py new file mode 100644 index 000000000..ddf1dc10c --- /dev/null +++ b/src/mobius/weights/_records_test.py @@ -0,0 +1,47 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT License. + +"""Tests for typed logical checkpoint records.""" + +from __future__ import annotations + +import pytest +import torch + +from mobius._component_manifest import ComponentDescriptor +from mobius.weights import FloatWeight, WeightBundle, WeightRecord + + +def _component() -> ComponentDescriptor: + return ComponentDescriptor( + name="decoder", + module_path="decoder", + role="decoder", + source_paths=("model.layers",), + ) + + +def test_bundle_tracks_source_keys(): + record = WeightRecord( + name="model.norm.weight", + component="decoder", + storage=FloatWeight( + value=torch.ones(4), + source_key="model.norm.weight", + ), + ) + bundle = WeightBundle(_component(), {record.name: record}) + + assert bundle.source_keys == frozenset({"model.norm.weight"}) + assert bundle["model.norm.weight"] is record + + +def test_bundle_rejects_wrong_component(): + record = WeightRecord( + name="model.norm.weight", + component="vision_encoder", + storage=FloatWeight(torch.ones(4), "model.norm.weight"), + ) + + with pytest.raises(ValueError, match="belongs to component"): + WeightBundle(_component(), {record.name: record}) From 15c0dfc609c84e1c14d200864b3a0fc05239d7d7 Mon Sep 17 00:00:00 2001 From: Xiaoyu Zhang Date: Thu, 27 Aug 2026 14:06:06 -0700 Subject: [PATCH 7/7] Use explicit component module attribute paths Update typed weight fixtures for the clarified ComponentDescriptor API. Signed-off-by: Xiaoyu Zhang --- src/mobius/weights/_codecs_test.py | 2 +- src/mobius/weights/_records_test.py | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/src/mobius/weights/_codecs_test.py b/src/mobius/weights/_codecs_test.py index 338fbbcf3..e65ffaf0b 100644 --- a/src/mobius/weights/_codecs_test.py +++ b/src/mobius/weights/_codecs_test.py @@ -16,7 +16,7 @@ def _component() -> ComponentDescriptor: return ComponentDescriptor( name="decoder", - module_path="decoder", + module_attribute_path="decoder", role="decoder", source_paths=("model.layers",), ) diff --git a/src/mobius/weights/_records_test.py b/src/mobius/weights/_records_test.py index ddf1dc10c..f15aae40a 100644 --- a/src/mobius/weights/_records_test.py +++ b/src/mobius/weights/_records_test.py @@ -15,7 +15,7 @@ def _component() -> ComponentDescriptor: return ComponentDescriptor( name="decoder", - module_path="decoder", + module_attribute_path="decoder", role="decoder", source_paths=("model.layers",), )